aklbpsd/wealth-slide-classifier

1from fastai.vision.all import *2import gradio as gr3import fitz  # PyMuPDF4import tempfile5import os6from pathlib import Path7 8# Load model9model_path = Path(__file__).parent / "wealth_slide_classifier.pkl"10learn = load_learner(model_path)11 12# Prediction function13def classify_pdf(pdf_file):14    import fitz  # PyMuPDF15    import tempfile16    from PIL import Image17 18    results = []19 20    # Save uploaded PDF temporarily21    with tempfile.NamedTemporaryFile(suffix=".pdf", delete=False) as tmp_pdf:22        tmp_pdf.write(pdf_file)23        tmp_pdf_path = tmp_pdf.name24 25    # Open PDF with PyMuPDF26    doc = fitz.open(tmp_pdf_path)27 28    for i, page in enumerate(doc):29        # Convert page to image30        pix = page.get_pixmap(dpi=200)31        img_path = f"/tmp/slide_{i+1}.png"32        pix.save(img_path)33 34        # Predict using FastAI35        pred, pred_idx, probs = learn.predict(PILImage.create(img_path))36 37        if pred != 'not_useful':38            # Load image before deleting file39            image = Image.open(img_path)40            results.append((f"Slide {i+1} → {pred}", image))41 42        # Delete image file after using it43        os.remove(img_path)44 45    # Clean up temp PDF file46    os.remove(tmp_pdf_path)47 48    if not results:49        return ["No relevant slides found."], None50 51    labels, imgs = zip(*results)52    return list(labels), list(imgs)53 54# Gradio UI55demo = gr.Interface(56    fn=classify_pdf,57    inputs=gr.File(label="Upload Investor Presentation (PDF)", file_types=[".pdf"], type="binary"),58    outputs=[gr.Textbox(label="Predicted Slide Labels"), gr.Gallery(label="Relevant Slides")],59    title="Wealth Slides Classifier",60    description="Upload an investor presentation. The model extracts and classifies key slides (quarterly/yearly info, general info)."61)62 63demo.launch()64