aklbpsd/wealth-slide-api

1from fastapi import FastAPI, File, UploadFile2from fastapi.responses import JSONResponse3import fitz  # PyMuPDF4from fastai.vision.all import *5from pathlib import Path6import tempfile7from PIL import Image8import os9import base6410 11app = FastAPI()12 13# Load model at startup14MODEL_PATH = Path(__file__).parent / "wealth_slide_classifier.pkl"15learn = load_learner(MODEL_PATH)16 17@app.get("/")18def root():19    return {"message": "Wealth Slide API is running 🚀"}20 21@app.post("/predict")22async def classify_pdf(file: UploadFile = File(...)):23    # Save uploaded PDF to temp file24    pdf_bytes = await file.read()25    with tempfile.NamedTemporaryFile(delete=False, suffix=".pdf") as tmp_pdf:26        tmp_pdf.write(pdf_bytes)27        tmp_pdf_path = tmp_pdf.name28 29    # Open PDF30    doc = fitz.open(tmp_pdf_path)31    results = []32 33    for i, page in enumerate(doc):34        # Convert page to image35        pix = page.get_pixmap(dpi=200)36        img_path = f"/tmp/slide_{i+1}.png"37        pix.save(img_path)38 39        # Run prediction40        pred, pred_idx, probs = learn.predict(PILImage.create(img_path))41 42        # Filter out 'not_useful' slides43        if pred != "not_useful":44            # Read image as base6445            with open(img_path, "rb") as img_file:46                img_bytes = img_file.read()47                img_b64 = base64.b64encode(img_bytes).decode("utf-8")48 49            results.append({50                "page_num": i + 1,51                "label": pred,52                "image_base64": img_b6453            })54 55        # Clean up the temp image56        os.remove(img_path)57 58    # Clean up the temp PDF59    os.remove(tmp_pdf_path)60 61    return JSONResponse(content={"results": results})62 63