Create app.py

app.py

@@ -0,0 +1,78 @@
+import os
+import gradio as gr
+import pandas as pd
+from PIL import Image
+import fitz  # PyMuPDF
+from transformers import pipeline
+
+# Pick a lightweight doc classifier. Swap to your preferred HF model.
+MODEL_ID = os.getenv("MODEL_ID", "HAMMALE/vit-tiny-classifier-rvlcdip")
+
+clf = pipeline(
+    task="image-classification",
+    model=MODEL_ID,
+    device=0 if os.getenv("CUDA_VISIBLE_DEVICES") not in (None, "", "-1") else -1,
+)
+
+def pdf_to_images(pdf_path: str, max_pages: int = 6, dpi: int = 150):
+    doc = fitz.open(pdf_path)
+    images = []
+    zoom = dpi / 72.0
+    mat = fitz.Matrix(zoom, zoom)
+    for i in range(min(len(doc), max_pages)):
+        page = doc.load_page(i)
+        pix = page.get_pixmap(matrix=mat, alpha=False)
+        img = Image.frombytes("RGB", [pix.width, pix.height], pix.samples)
+        images.append(img)
+    doc.close()
+    return images
+
+def run_infer(file_obj, max_pages: int = 6, top_k: int = 5):
+    path = file_obj.name
+    ext = os.path.splitext(path)[1].lower()
+
+    if ext == ".pdf":
+        images = pdf_to_images(path, max_pages=max_pages)
+        page_labels = [f"page_{i+1}" for i in range(len(images))]
+    else:
+        images = [Image.open(path).convert("RGB")]
+        page_labels = ["image"]
+
+    rows = []
+    # Aggregate by summing scores per label across pages (simple + robust)
+    agg = {}
+
+    for label, img in zip(page_labels, images):
+        preds = clf(img, top_k=top_k)
+        for p in preds:
+            rows.append({"item": label, "label": p["label"], "score": float(p["score"])})
+            agg[p["label"]] = agg.get(p["label"], 0.0) + float(p["score"])
+
+    per_page = pd.DataFrame(rows).sort_values(["item", "score"], ascending=[True, False])
+
+    agg_df = (
+        pd.DataFrame([{"label": k, "score_sum": v} for k, v in agg.items()])
+        .sort_values("score_sum", ascending=False)
+        .head(top_k)
+        .reset_index(drop=True)
+    )
+
+    return per_page, agg_df
+
+demo = gr.Interface(
+    fn=run_infer,
+    inputs=[
+        gr.File(label="Upload PDF/PNG/JPG"),
+        gr.Slider(1, 30, value=6, step=1, label="Max PDF pages"),
+        gr.Slider(1, 20, value=5, step=1, label="Top-K labels"),
+    ],
+    outputs=[
+        gr.Dataframe(label="Per-page predictions"),
+        gr.Dataframe(label="Aggregated across pages (sum of scores)"),
+    ],
+    title="Document Classifier (PDF/PNG)",
+    description=f"Model: {MODEL_ID}. Upload a PDF or image to classify document type.",
+)
+
+if __name__ == "__main__":
+    demo.launch()