Update app.py

app.py

@@ -2,11 +2,10 @@ import os
 import tempfile
 import zipfile
 from pathlib import Path
-from typing import List, Tuple
+from typing import List, Tuple, Any
 
 import gradio as gr
 import torch
-import torch.nn.functional as F
 
 from model import PairNet, load_weights_if_any
 from utils import load_exam_as_batch, aggregate_predictions, clamp_days, today_plus_days
@@ -27,11 +26,20 @@ def init_model() -> Tuple[torch.nn.Module, str, bool]:
 
 MODEL, LOAD_MSG, HAS_WEIGHTS = init_model()
 
-def predict_on_files(files: List[gr.File]) -> dict:
-    # Collect file paths (also support zip of images)
+def _to_path(maybe_file: Any) -> Path:
+    if isinstance(maybe_file, str):
+        return Path(maybe_file)
+    if isinstance(maybe_file, dict) and "name" in maybe_file:
+        return Path(maybe_file["name"])
+    name = getattr(maybe_file, "name", None)
+    if name:
+        return Path(name)
+    return Path(str(maybe_file))
+
+def predict_on_files(files) -> dict:
     paths: List[str] = []
     for f in files or []:
-        p = Path(f.name)
+        p = _to_path(f)
         if p.suffix.lower() == ".zip":
             with zipfile.ZipFile(p, "r") as z:
                 with tempfile.TemporaryDirectory() as td:
@@ -44,17 +52,14 @@ def predict_on_files(files: List[gr.File]) -> dict:
     if not paths:
         return {"status": "no files received"}
 
-    x = load_exam_as_batch(paths).to(DEVICE)
-    x = x.float()
+    x = load_exam_as_batch(paths).to(DEVICE).float()
 
     if HAS_WEIGHTS:
         with torch.no_grad():
             days_raw, logits = MODEL(x)
-            # Aggregate across frames by mean
             days = days_raw.squeeze(-1).detach().cpu().tolist()
             proba = torch.sigmoid(logits.squeeze(-1)).detach().cpu().tolist()
     else:
-        # Baseline demo: constant mid-gestation-ish guess and neutral probability
         days = [150.0 for _ in range(x.shape[0])]
         proba = [0.5 for _ in range(x.shape[0])]
 
@@ -63,7 +68,7 @@ def predict_on_files(files: List[gr.File]) -> dict:
 
     days_mean, proba_mean = aggregate_predictions(days, proba)
 
-    result = {
+    return {
         "frames": len(paths),
         "per_frame_days": days,
         "per_frame_preterm_proba": proba,
@@ -74,31 +79,32 @@ def predict_on_files(files: List[gr.File]) -> dict:
         "aggregate_preterm_label": "Preterm" if proba_mean >= 0.5 else "Term",
         "weights_message": LOAD_MSG
     }
-    return result
 
 with gr.Blocks(title="PAIR-inspired Delivery Timing Predictor") as demo:
     gr.Markdown(
-        "PAIR-inspired Delivery Timing Predictor
+        """
+PAIR-inspired Delivery Timing Predictor
 
-"
-        "This app is a technical scaffold inspired by the PAIR study. "
-        "It does not include the proprietary model or clinical dataset. "
-        "Not for medical use."
+This app is a technical scaffold inspired by the PAIR study.
+It does not include the proprietary model or clinical dataset.
+Not for medical use.
+        """
     )
 
     with gr.Row():
         with gr.Column():
-            in_files = gr.Files(label="Upload ultrasound images or a ZIP (PNG/JPG/DICOM)", file_count="multiple", type="filepath")
+            in_files = gr.Files(
+                label="Upload ultrasound images or a ZIP (PNG/JPG/DICOM)",
+                file_count="multiple",
+                type="filepath"
+            )
             run_btn = gr.Button("Run prediction", variant="primary")
         with gr.Column():
             status = gr.JSON(label="Outputs")
 
-    note = gr.Markdown(f"Model status: {LOAD_MSG}")
+    gr.Markdown(f"Model status: {LOAD_MSG}")
 
-    def _run(files):
-        return predict_on_files(files)
-
-    run_btn.click(_run, inputs=[in_files], outputs=[status])
+    run_btn.click(fn=predict_on_files, inputs=[in_files], outputs=[status])
 
 if __name__ == "__main__":
     demo.launch()