Upload 5 files

README.md

@@ -1,13 +1,31 @@
----
-title: PAIR
-emoji: 🦀
-colorFrom: pink
-colorTo: green
-sdk: gradio
-sdk_version: 5.42.0
-app_file: app.py
-pinned: false
-short_description: PAIR
----
-
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+# PAIR-inspired Delivery Timing Predictor (Gradio Space)
+
+Important
+This demo is inspired by the PAIR study (Perinatal artificial intelligence in ultrasound) but does not include the proprietary model or the private clinical dataset described in the paper. It is provided only as a technical scaffold and demonstration UI.
+
+What this Space does
+1. Lets you upload ultrasound images (PNG/JPG) or DICOM files and returns:
+   - Predicted days-until-delivery (regression).
+   - Preterm probability and label (binary classification with threshold 0.5).
+2. Aggregates predictions across multiple images from the same exam using a simple mean.
+3. Can load your own PyTorch weights via a Hugging Face repo id or a path to a .pt file.
+
+How to use in your own Space
+1) Create a new Space on Hugging Face with SDK set to Gradio.
+2) Upload the files in this repository.
+3) Optional: place your model weights at weights/pair_v4.pt or set the env var HF_WEIGHTS to point to a Hugging Face model repo or a local .pt file.
+4) Click Run. If weights are missing, the app falls back to a constant baseline just to demonstrate the UI.
+
+Model input
+- One or more 2D ultrasound images or DICOM frames from one exam. The app will convert grayscale to 3-channels when needed.
+
+Model output
+- days_to_delivery: float in [1, 300] (clamped).
+- preterm_proba: float in [0, 1].
+- preterm_label: Term if proba < 0.5, Preterm otherwise.
+- predicted_date: today + days_to_delivery (for demo; in clinical use you would provide the scan date).
+
+Notes
+- This is not a medical device and is not for clinical use.
+- Performance will be meaningless without appropriate training on a suitable dataset; this scaffold is for integration and UI only.
+- To replicate the paper, you would need authorized access to a comparable dataset and a trained model.

app.py

@@ -0,0 +1,104 @@
+import os
+import tempfile
+import zipfile
+from pathlib import Path
+from typing import List, Tuple
+
+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
+
+DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
+
+def init_model() -> Tuple[torch.nn.Module, str, bool]:
+    model = PairNet(pretrained=True).to(DEVICE)
+    model.eval()
+    weights_hint = os.getenv("HF_WEIGHTS", "").strip()
+    ok = False
+    msg = "Running in baseline mode (no weights). Output is a UI demo only."
+    if weights_hint:
+        ok, msg = load_weights_if_any(model, weights_hint)
+    elif Path("weights/pair_v4.pt").exists():
+        ok, msg = load_weights_if_any(model, "weights/pair_v4.pt")
+    return model, msg, ok
+
+MODEL, LOAD_MSG, HAS_WEIGHTS = init_model()
+
+def predict_on_files(files: List[gr.File]) -> dict:
+    # Collect file paths (also support zip of images)
+    paths: List[str] = []
+    for f in files or []:
+        p = Path(f.name)
+        if p.suffix.lower() == ".zip":
+            with zipfile.ZipFile(p, "r") as z:
+                with tempfile.TemporaryDirectory() as td:
+                    z.extractall(td)
+                    for ext in (".png", ".jpg", ".jpeg", ".dcm"):
+                        paths.extend([str(q) for q in Path(td).rglob(f"*{ext}")])
+        else:
+            paths.append(str(p))
+
+    if not paths:
+        return {"status": "no files received"}
+
+    x = load_exam_as_batch(paths).to(DEVICE)
+    x = x.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])]
+
+    days = [clamp_days(float(d)) for d in days]
+    preterm = ["Preterm" if p >= 0.5 else "Term" for p in proba]
+
+    days_mean, proba_mean = aggregate_predictions(days, proba)
+
+    result = {
+        "frames": len(paths),
+        "per_frame_days": days,
+        "per_frame_preterm_proba": proba,
+        "per_frame_preterm_label": preterm,
+        "aggregate_days_mean": days_mean,
+        "aggregate_predicted_date": today_plus_days(days_mean),
+        "aggregate_preterm_proba": proba_mean,
+        "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
+
+"
+        "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")
+            run_btn = gr.Button("Run prediction", variant="primary")
+        with gr.Column():
+            status = gr.JSON(label="Outputs")
+
+    note = gr.Markdown(f"Model status: {LOAD_MSG}")
+
+    def _run(files):
+        return predict_on_files(files)
+
+    run_btn.click(_run, inputs=[in_files], outputs=[status])
+
+if __name__ == "__main__":
+    demo.launch()

model.py

@@ -0,0 +1,62 @@
+import torch
+import torch.nn as nn
+from torchvision import models
+
+class PairNet(nn.Module):
+    """
+    A lightweight backbone + dual-head network:
+      - Regression head for days-to-delivery
+      - Classification head for preterm probability
+    This is a scaffold and not the proprietary model from the paper.
+    """
+    def __init__(self, backbone_name: str = "efficientnet_b0", pretrained: bool = True):
+        super().__init__()
+        if backbone_name == "efficientnet_b0":
+            try:
+                weights = models.EfficientNet_B0_Weights.DEFAULT if pretrained else None
+            except Exception:
+                weights = None
+            backbone = models.efficientnet_b0(weights=weights)
+            in_feats = backbone.classifier[1].in_features
+            backbone.classifier = nn.Identity()
+        else:
+            # Fallback to resnet18
+            try:
+                weights = models.ResNet18_Weights.DEFAULT if pretrained else None
+            except Exception:
+                weights = None
+            backbone = models.resnet18(weights=weights)
+            in_feats = backbone.fc.in_features
+            backbone.fc = nn.Identity()
+
+        self.backbone = backbone
+        self.reg_head = nn.Linear(in_feats, 1)
+        self.cls_head = nn.Linear(in_feats, 1)
+
+    def forward(self, x):
+        feats = self.backbone(x)
+        days = self.reg_head(feats)   # unconstrained
+        logits = self.cls_head(feats) # unconstrained
+        return days, logits
+
+def load_weights_if_any(model: nn.Module, weights_path: str | None):
+    if not weights_path:
+        return False, "No weights path provided"
+    import os
+    if os.path.isfile(weights_path):
+        state = torch.load(weights_path, map_location="cpu")
+        if "state_dict" in state:
+            state = state["state_dict"]
+        missing, unexpected = model.load_state_dict(state, strict=False)
+        return True, f"Loaded local weights. missing={len(missing)} unexpected={len(unexpected)}"
+    # Try huggingface hub repo id
+    try:
+        from huggingface_hub import hf_hub_download
+        fp = hf_hub_download(repo_id=weights_path, filename="pytorch_model.bin", local_dir="weights")
+        state = torch.load(fp, map_location="cpu")
+        if "state_dict" in state:
+            state = state["state_dict"]
+        missing, unexpected = model.load_state_dict(state, strict=False)
+        return True, f"Loaded HF weights from {weights_path}. missing={len(missing)} unexpected={len(unexpected)}"
+    except Exception as e:
+        return False, f"Failed to load weights: {e}"

requirements.txt

@@ -0,0 +1,7 @@
+gradio>=4.44.0
+torch>=2.2.0
+torchvision>=0.17.0
+pillow>=10.0.0
+numpy>=1.26.0
+pydicom>=2.4.4
+huggingface_hub>=0.23.0

utils.py

@@ -0,0 +1,54 @@
+import io
+import os
+import math
+import datetime as dt
+from typing import List, Tuple
+
+import numpy as np
+from PIL import Image
+import torch
+import pydicom
+
+IMG_SIZE = 224
+
+def _read_image(file_path: str) -> Image.Image:
+    ext = os.path.splitext(file_path)[1].lower()
+    if ext == ".dcm":
+        ds = pydicom.dcmread(file_path)
+        arr = ds.pixel_array.astype(np.float32)
+        # normalize to 0..255
+        arr = arr - arr.min()
+        if arr.max() > 0:
+            arr = arr / arr.max()
+        arr = (arr * 255.0).clip(0,255).astype(np.uint8)
+        return Image.fromarray(arr)
+    else:
+        return Image.open(file_path).convert("L")  # grayscale
+
+def _to_tensor(img: Image.Image) -> torch.Tensor:
+    # Resize, center-crop/pad to square, stack to 3 channels, normalize 0..1
+    img = img.resize((IMG_SIZE, IMG_SIZE))
+    arr = np.array(img).astype(np.float32) / 255.0
+    if arr.ndim == 2:
+        arr = np.stack([arr, arr, arr], axis=0)  # 3xHxW
+    elif arr.ndim == 3:
+        arr = arr.transpose(2, 0, 1)  # HWC -> CHW
+    return torch.from_numpy(arr)
+
+def load_exam_as_batch(file_paths: List[str]) -> torch.Tensor:
+    imgs = [_to_tensor(_read_image(p)) for p in file_paths]
+    x = torch.stack(imgs, dim=0)  # Nx3xHxW
+    return x
+
+def aggregate_predictions(days_list: List[float], proba_list: List[float]) -> Tuple[float, float]:
+    if len(days_list) == 0:
+        return 0.0, 0.0
+    return float(np.mean(days_list)), float(np.mean(proba_list))
+
+def clamp_days(d: float) -> float:
+    return float(max(1.0, min(300.0, d)))
+
+def today_plus_days(days: float) -> str:
+    base = dt.date.today()
+    target = base + dt.timedelta(days=int(round(days)))
+    return target.isoformat()