deploy YOLACT+ fracture detection demo

README.md

@@ -1,14 +1,160 @@
+# YOLACT+ (ResNet-18) on FracAtlas
+
+Instance segmentation pipeline for fracture detection using YOLACT+ with a lightweight ResNet-18 backbone.
+
 ---
-title: FracAtlas YOLACT
-emoji: 🏃
-colorFrom: green
-colorTo: red
-sdk: gradio
-sdk_version: 6.13.0
-app_file: app.py
-pinned: false
-license: apache-2.0
-short_description: Fracture detection and segmentation on X-ray images
+
+## File Overview
+
+```
+fracatlas_yolact/
+├── prepare_dataset.py   # Step 1 — split FracAtlas into train/val/test
+├── dataloader.py        # Dataset class + DataLoader factory
+├── model.py             # YOLACT+ architecture (ResNet-18 backbone)
+├── loss.py              # Focal cls + SmoothL1 box + BCE mask loss
+├── train.py             # Training loop with CSV logging + checkpointing
+├── inference.py         # Single-image / val-set / test-set inference
+└── requirements.txt
+```
+
+---
+
+## Dataset Splits
+
+| Split | Fractured | Non-fractured | Total |
+|-------|-----------|---------------|-------|
+| Train | 500       | 500           | 1000  |
+| Val   | 100       | 100           | 200   |
+| Test  | 100       | 100           | 200   |
+
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+## Quick Start
+
+### 1. Install dependencies
+```bash
+pip install -r requirements.txt
+```
+
+### 2. Prepare dataset
+```bash
+python prepare_dataset.py \
+    --fracatlas_root /path/to/FracAtlas \
+    --output_dir     data/
+```
+
+Expected FracAtlas layout:
+```
+FracAtlas/
+├── images/
+│   ├── Fractured/
+│   └── Non_fractured/
+└── Annotations/COCO JSON/COCO_fracture_all.json
+```
+
+Output:
+```
+data/
+├── train/images/ + annotations.json   (1000 images)
+├── val/images/   + annotations.json   (200 images)
+└── test/images/  + annotations.json   (200 images)
+```
+
+### 3. Train
+```bash
+python train.py \
+    --data_root  data/ \
+    --epochs     100   \
+    --batch_size 8     \
+    --lr         1e-4
+```
+
+Checkpoints → `runs/fracatlas_yolact_resnet18/checkpoints/`
+- `best.pth`      — lowest validation loss
+- `last.pth`      — most recent epoch
+- `epoch_NNN.pth` — every 10 epochs
+
+Training log → `runs/fracatlas_yolact_resnet18/train_log.csv`
+
+### 4. Inference
+
+#### Single image
+```bash
+python inference.py \
+    --mode        single \
+    --weights     runs/fracatlas_yolact_resnet18/checkpoints/best.pth \
+    --image       path/to/xray.jpg \
+    --save_vis    output_vis.jpg
+```
+
+#### Validation set (with metrics)
+```bash
+python inference.py \
+    --mode         val \
+    --weights      runs/fracatlas_yolact_resnet18/checkpoints/best.pth \
+    --data_root    data/ \
+    --save_vis_dir results/vis/val/ \
+    --save_results results/val_metrics.json
+```
+
+#### Test set (with metrics)
+```bash
+python inference.py \
+    --mode         test \
+    --weights      runs/fracatlas_yolact_resnet18/checkpoints/best.pth \
+    --data_root    data/ \
+    --save_vis_dir results/vis/test/ \
+    --save_results results/test_metrics.json
+```
+
+---
+
+## Architecture Summary
+
+```
+Input [B, 3, 550, 550]
+  │
+  ▼
+ResNet-18 Backbone
+  │  C3 [128, 69, 69]
+  │  C4 [256, 35, 35]
+  │  C5 [512, 18, 18]
+  │
+  ▼
+FPN (5 levels P3–P7, each 256 channels)
+  │
+  ├─► ProtoNet (on P3) → Prototypes [B, 32, 138, 138]
+  │
+  └─► PredictionHead (shared across levels)
+        ├─ cls_pred  [B, A, num_classes+1]
+        ├─ box_pred  [B, A, 4]
+        └─ coef_pred [B, A, 32]
+
+Final masks = sigmoid(prototypes ⊗ coefficients)
+```
+
+---
+
+## Key Hyperparameters
+
+| Parameter        | Default | Notes                                      |
+|------------------|---------|--------------------------------------------|
+| `img_size`       | 550     | YOLACT+ standard input resolution          |
+| `num_prototypes` | 32      | Prototype masks                            |
+| `fpn_channels`   | 256     | FPN output channels                        |
+| `lr`             | 1e-4    | AdamW learning rate                        |
+| `lambda_cls`     | 1.0     | Classification (focal) loss weight         |
+| `lambda_box`     | 1.5     | Box regression (SmoothL1) loss weight      |
+| `lambda_mask`    | 6.125   | Mask (BCE) loss weight (YOLACT+ default)   |
+| `score_thresh`   | 0.3     | Detection confidence threshold             |
+| `nms_thresh`     | 0.5     | NMS IoU threshold                          |
+
+---
+
+## Evaluation Metrics
+
+The inference script reports per-class and mean:
+- **Precision** — TP / (TP + FP)
+- **Recall**    — TP / (TP + FN)
+- **F1 Score**  — harmonic mean of precision and recall
+- **Avg IoU**   — average box IoU of matched detections (threshold @ 0.5)

README_hf.md

@@ -0,0 +1,58 @@
+---
+title: FracAtlas YOLACT+
+emoji: 🦴
+colorFrom: red
+colorTo: gray
+sdk: gradio
+sdk_version: 4.44.0
+app_file: app.py
+pinned: false
+license: apache-2.0
+short_description: Fracture detection and segmentation on X-ray images using YOLACT+ with ResNet-18 backbone, trained on the FracAtlas dataset.
+---
+
+# FracAtlas Fracture Detection — YOLACT+ (ResNet-18)
+
+Instance segmentation model for bone fracture detection on X-ray images.
+
+## Model
+
+- **Architecture:** YOLACT+ with ResNet-18 backbone
+- **Neck:** Feature Pyramid Network (FPN, 5 levels)
+- **Prototypes:** 32 mask prototypes
+- **Input size:** 550×550
+
+## Dataset
+
+[FracAtlas](https://figshare.com/articles/dataset/The_dataset/22363012)
+
+| Split | Fractured | Non-fractured | Total |
+|-------|-----------|---------------|-------|
+| Train | 500 | 500 | 1000 |
+| Val   | 100 | 100 | 200  |
+| Test  | 100 | 100 | 200  |
+
+## Training
+
+- **Epochs:** 200
+- **Optimizer:** AdamW (lr=5e-5, weight_decay=5e-4)
+- **Scheduler:** Cosine decay with 5-epoch warmup
+- **Loss:** Focal classification + SmoothL1 box + BCE mask
+
+## Results (Validation Set)
+
+| Metric | Value |
+|--------|-------|
+| Precision | 0.5328 |
+| Recall | 0.5422 |
+| F1 Score | 0.5374 |
+| Avg IoU | 0.9405 |
+
+## Author
+
+Muhammad Adil — MS Data Science, Information Technology University (ITU) Lahore, Pakistan  
+GitHub: [Adil6312](https://github.com/Adil6312)
+
+## License
+
+Apache 2.0

app.py

@@ -0,0 +1,231 @@
+"""
+FracAtlas YOLACT+ Demo
+======================
+Gradio app for fracture detection and segmentation on X-ray images.
+Deployed on Hugging Face Spaces: https://huggingface.co/spaces/MuhammadAdil63/FracAtlas-YOLACT
+"""
+
+import os
+import cv2
+import numpy as np
+import torch
+import torch.nn.functional as F
+import gradio as gr
+from huggingface_hub import hf_hub_download
+import albumentations as A
+from albumentations.pytorch import ToTensorV2
+
+from model import YOLACTPlus
+
+
+# ─── Config ───────────────────────────────────────────────────────────────────
+IMG_SIZE     = 550
+NUM_CLASSES  = 1
+CLASS_NAMES  = ["fracture"]
+SCORE_THRESH = 0.4
+NMS_THRESH   = 0.4
+DEVICE       = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+# HuggingFace repo where best.pth is hosted
+HF_REPO_ID   = "MuhammadAdil63/FracAtlas-YOLACT"   # change if model hosted separately
+CKPT_FILE    = "best.pth"
+
+# Mask overlay colour (red in RGB)
+MASK_COLOR   = (220, 50, 50)
+BOX_COLOR    = (220, 50, 50)
+
+
+# ─── Load model (cached after first load) ────────────────────────────────────
+
+def load_model():
+    print(f"[INFO] Loading model on {DEVICE} ...")
+
+    # Try local first, then download from HF Hub
+    if os.path.isfile(CKPT_FILE):
+        ckpt_path = CKPT_FILE
+        print(f"[INFO] Using local checkpoint: {ckpt_path}")
+    else:
+        print(f"[INFO] Downloading checkpoint from HF Hub ...")
+        ckpt_path = hf_hub_download(repo_id=HF_REPO_ID, filename=CKPT_FILE)
+
+    model = YOLACTPlus(num_classes=NUM_CLASSES, img_size=IMG_SIZE, pretrained=False)
+    ckpt  = torch.load(ckpt_path, map_location=DEVICE)
+    model.load_state_dict(ckpt.get("model", ckpt))
+    model.to(DEVICE).eval()
+    print("[INFO] Model loaded successfully!")
+    return model
+
+
+# Load once at startup
+MODEL = load_model()
+
+
+# ─── Preprocessing ────────────────────────────────────────────────────────────
+
+def preprocess(image_rgb: np.ndarray):
+    transform = A.Compose([
+        A.LongestMaxSize(max_size=IMG_SIZE),
+        A.PadIfNeeded(min_height=IMG_SIZE, min_width=IMG_SIZE, fill=0),
+        A.Normalize(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)),
+        ToTensorV2(),
+    ])
+    out    = transform(image=image_rgb)
+    tensor = out["image"].unsqueeze(0).to(DEVICE)
+    return tensor
+
+
+# ─── Draw predictions ─────────────────────────────────────────────────────────
+
+def draw_predictions(image_rgb: np.ndarray, result: dict) -> np.ndarray:
+    vis = image_rgb.copy().astype(np.float32)
+    H, W = vis.shape[:2]
+
+    boxes  = result["boxes"]    # [N, 4] normalised
+    scores = result["scores"]   # [N]
+    masks  = result["masks"]    # [N, IMG_SIZE, IMG_SIZE]
+
+    for i in range(len(scores)):
+        # Mask overlay
+        mask_np = masks[i].numpy()
+        mask_rs = cv2.resize(mask_np, (W, H), interpolation=cv2.INTER_LINEAR)
+        mask_bin = (mask_rs > 0.5).astype(np.float32)
+        colour = np.array(MASK_COLOR, dtype=np.float32)
+        for c in range(3):
+            vis[:, :, c] = vis[:, :, c] * (1 - 0.45 * mask_bin) + \
+                           colour[c] * (0.45 * mask_bin)
+
+        # Bounding box
+        x1 = int(boxes[i, 0].item() * W)
+        y1 = int(boxes[i, 1].item() * H)
+        x2 = int(boxes[i, 2].item() * W)
+        y2 = int(boxes[i, 3].item() * H)
+        vis_u8 = np.clip(vis, 0, 255).astype(np.uint8)
+        cv2.rectangle(vis_u8, (x1, y1), (x2, y2), BOX_COLOR, 2)
+
+        # Label
+        score = scores[i].item()
+        tag   = f"fracture: {score:.2f}"
+        (tw, th), _ = cv2.getTextSize(tag, cv2.FONT_HERSHEY_SIMPLEX, 0.6, 1)
+        cv2.rectangle(vis_u8, (x1, y1 - th - 8), (x1 + tw + 4, y1), BOX_COLOR, -1)
+        cv2.putText(vis_u8, tag, (x1 + 2, y1 - 4),
+                    cv2.FONT_HERSHEY_SIMPLEX, 0.6, (255, 255, 255), 1, cv2.LINE_AA)
+        vis = vis_u8.astype(np.float32)
+
+    return np.clip(vis, 0, 255).astype(np.uint8)
+
+
+# ─── Main inference function ──────────────────────────────────────────────────
+
+def predict(image: np.ndarray, score_threshold: float, nms_threshold: float):
+    """
+    Gradio inference function.
+    Args:
+        image           : RGB numpy array from Gradio
+        score_threshold : confidence threshold slider
+        nms_threshold   : NMS IoU threshold slider
+    Returns:
+        annotated image, result summary text
+    """
+    if image is None:
+        return None, "No image provided."
+
+    orig_rgb = image.copy()
+    tensor   = preprocess(orig_rgb)
+
+    with torch.no_grad():
+        results = MODEL.predict(tensor, score_threshold, nms_threshold)
+
+    res = results[0]
+    n   = len(res["scores"])
+
+    # Draw
+    annotated = draw_predictions(orig_rgb, res)
+
+    # Summary text
+    if n == 0:
+        summary = "**No fractures detected.**\n\nTry lowering the Score Threshold."
+    else:
+        lines = [f"**{n} fracture(s) detected:**\n"]
+        for i in range(n):
+            score = res["scores"][i].item()
+            box   = res["boxes"][i].tolist()
+            lines.append(
+                f"- Detection {i+1}: confidence **{score:.3f}** | "
+                f"box `[{box[0]:.3f}, {box[1]:.3f}, {box[2]:.3f}, {box[3]:.3f}]`"
+            )
+        summary = "\n".join(lines)
+
+    return annotated, summary
+
+
+# ─── Gradio UI ────────────────────────────────────────────────────────────────
+
+DESCRIPTION = """
+## FracAtlas Fracture Detection — YOLACT+ (ResNet-18)
+
+Upload an X-ray image to detect and segment bone fractures.
+
+**Model:** YOLACT+ with ResNet-18 backbone  
+**Dataset:** [FracAtlas](https://figshare.com/articles/dataset/The_dataset/22363012) — 717 fractured + 3366 non-fractured X-rays  
+**Training:** 200 epochs | AdamW | Cosine LR decay with warmup  
+**Val F1:** 0.537 | **Val Avg IoU:** 0.940
+"""
+
+EXAMPLES = [
+    ["examples/fractured_example.jpg",   0.4, 0.4],
+    ["examples/nonfractured_example.jpg", 0.4, 0.4],
+]
+
+with gr.Blocks(theme=gr.themes.Soft(), title="FracAtlas YOLACT+") as demo:
+
+    gr.Markdown(DESCRIPTION)
+
+    with gr.Row():
+        with gr.Column(scale=1):
+            input_image = gr.Image(
+                label="Input X-ray Image",
+                type="numpy",
+                image_mode="RGB",
+            )
+            with gr.Accordion("Detection Settings", open=False):
+                score_thresh = gr.Slider(
+                    minimum=0.1, maximum=0.9, value=0.4, step=0.05,
+                    label="Score Threshold",
+                    info="Higher = fewer but more confident detections",
+                )
+                nms_thresh = gr.Slider(
+                    minimum=0.1, maximum=0.9, value=0.4, step=0.05,
+                    label="NMS Threshold",
+                    info="Lower = suppress more overlapping boxes",
+                )
+            run_btn = gr.Button("Detect Fractures", variant="primary")
+
+        with gr.Column(scale=1):
+            output_image = gr.Image(
+                label="Predicted Segmentation",
+                type="numpy",
+            )
+            output_text = gr.Markdown(label="Detection Summary")
+
+    run_btn.click(
+        fn=predict,
+        inputs=[input_image, score_thresh, nms_thresh],
+        outputs=[output_image, output_text],
+    )
+
+    # Also run on image upload
+    input_image.change(
+        fn=predict,
+        inputs=[input_image, score_thresh, nms_thresh],
+        outputs=[output_image, output_text],
+    )
+
+    gr.Markdown("""
+---
+**Note:** This is a research demo. Not intended for clinical use.  
+**Author:** Muhammad Adil | MS Data Science, ITU Lahore  
+**GitHub:** [Adil6312](https://github.com/Adil6312)
+    """)
+
+if __name__ == "__main__":
+    demo.launch()

best.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:bac75ce750a77f810c60c9a4a72dd8be8d0b39ebd1969140a643d25de8af578b
+size 229563787

dataloader.py

@@ -0,0 +1,292 @@
+"""
+FracAtlas DataLoader for YOLACT+ (ResNet-18 backbone)
+======================================================
+Provides:
+  - FracAtlasDataset   : torch.utils.data.Dataset over COCO-format splits
+  - detection_collate  : custom collate for variable-size masks/boxes
+  - get_dataloader     : factory function for train / val / test loaders
+"""
+
+import os
+import os
+import cv2
+import numpy as np
+import torch
+import warnings
+warnings.filterwarnings("ignore", message=".*Premature end.*")
+warnings.filterwarnings("ignore", message=".*Corrupt JPEG.*")
+# Suppress OpenCV JPEG warnings
+try:
+    cv2.setLogLevel(0)
+except AttributeError:
+    os.environ["OPENCV_LOG_LEVEL"] = "SILENT"
+from torch.utils.data import Dataset, DataLoader
+from pycocotools.coco import COCO
+from pycocotools import mask as coco_mask
+import albumentations as A
+from albumentations.pytorch import ToTensorV2
+
+
+# ─── Augmentation pipelines ───────────────────────────────────────────────────
+
+def get_train_transforms(img_size: int = 550):
+    return A.Compose(
+        [
+            A.LongestMaxSize(max_size=img_size),
+            A.PadIfNeeded(
+                min_height=img_size,
+                min_width=img_size,
+                fill=0,
+            ),
+            A.HorizontalFlip(p=0.5),
+            A.VerticalFlip(p=0.2),
+            A.RandomBrightnessContrast(brightness_limit=0.2, contrast_limit=0.2, p=0.5),
+            A.GaussNoise(p=0.3),
+            A.Affine(
+                translate_percent=0.05,
+                scale=(0.9, 1.1),
+                rotate=(-10, 10),
+                p=0.4,
+            ),
+            A.Normalize(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)),
+            ToTensorV2(),
+        ],
+        bbox_params=A.BboxParams(
+            format="pascal_voc",
+            label_fields=["class_labels"],
+            min_visibility=0.3,
+        ),
+    )
+
+
+def get_val_transforms(img_size: int = 550):
+    return A.Compose(
+        [
+            A.LongestMaxSize(max_size=img_size),
+            A.PadIfNeeded(
+                min_height=img_size,
+                min_width=img_size,
+                fill=0,
+            ),
+            A.Normalize(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)),
+            ToTensorV2(),
+        ],
+        bbox_params=A.BboxParams(
+            format="pascal_voc",
+            label_fields=["class_labels"],
+            min_visibility=0.3,
+        ),
+    )
+
+
+# ─── Dataset ──────────────────────────────────────────────────────────────────
+
+class FracAtlasDataset(Dataset):
+    """
+    COCO-format dataset for FracAtlas fracture detection.
+
+    Each item returns:
+        image   : FloatTensor [3, H, W]  (normalised)
+        target  : dict with keys
+                    boxes   : FloatTensor [N, 4]  (x1y1x2y2, normalised 0-1)
+                    labels  : LongTensor  [N]
+                    masks   : FloatTensor [N, H, W]  (binary, same spatial size as image)
+                    image_id: int
+    """
+
+    def __init__(
+        self,
+        image_dir: str,
+        ann_file: str,
+        img_size: int = 550,
+        split: str = "train",
+    ):
+        self.image_dir = image_dir
+        self.img_size = img_size
+        self.split = split
+
+        self.coco = COCO(ann_file)
+        self.image_ids = sorted(self.coco.imgs.keys())
+
+        # Build category → 0-indexed label map
+        # NOTE: FracAtlas uses category_id=0 ('fractured') — handle offset
+        cats = self.coco.loadCats(self.coco.getCatIds())
+        self.cat_id_to_label = {c["id"]: i for i, c in enumerate(cats)}
+        # If only one class and its id is 0, map it to label 0
+        if len(cats) == 1 and cats[0]["id"] == 0:
+            self.cat_id_to_label = {0: 0}
+        self.num_classes = len(cats)
+        self.class_names = [c["name"] for c in cats]
+
+        self.transforms = (
+            get_train_transforms(img_size)
+            if split == "train"
+            else get_val_transforms(img_size)
+        )
+
+        print(
+            f"[{split}] {len(self.image_ids)} images | "
+            f"{self.num_classes} classes: {self.class_names}"
+        )
+
+    def __len__(self):
+        return len(self.image_ids)
+
+    def _decode_mask(self, ann: dict, h: int, w: int) -> np.ndarray:
+        """Decode COCO RLE or polygon segmentation to binary mask."""
+        seg = ann.get("segmentation", None)
+        if seg is None:
+            # Fall back: create mask from bbox
+            x, y, bw, bh = [int(v) for v in ann["bbox"]]
+            m = np.zeros((h, w), dtype=np.uint8)
+            m[y : y + bh, x : x + bw] = 1
+            return m
+        if isinstance(seg, dict):  # RLE
+            return coco_mask.decode(seg).astype(np.uint8)
+        else:  # polygon
+            rle = coco_mask.frPyObjects(seg, h, w)
+            merged = coco_mask.merge(rle)
+            return coco_mask.decode(merged).astype(np.uint8)
+
+    def __getitem__(self, idx: int):
+        img_id = self.image_ids[idx]
+        img_info = self.coco.imgs[img_id]
+
+        # ── Load image ────────────────────────────────────────────────────────
+        img_path = os.path.join(self.image_dir, img_info["file_name"])
+        image = cv2.imread(img_path)
+        if image is None:
+            raise FileNotFoundError(f"Cannot read image: {img_path}")
+        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
+        orig_h, orig_w = image.shape[:2]
+
+        # ── Load annotations ─────────────────────────────────────────────────
+        ann_ids = self.coco.getAnnIds(imgIds=img_id)
+        anns = self.coco.loadAnns(ann_ids)
+
+        boxes, class_labels, raw_masks = [], [], []
+        for ann in anns:
+            x, y, bw, bh = ann["bbox"]
+            x1, y1, x2, y2 = x, y, x + bw, y + bh
+            # Clip to image bounds
+            x1 = max(0.0, x1)
+            y1 = max(0.0, y1)
+            x2 = min(float(orig_w), x2)
+            y2 = min(float(orig_h), y2)
+            if x2 <= x1 or y2 <= y1:
+                continue
+            boxes.append([x1, y1, x2, y2])
+            class_labels.append(self.cat_id_to_label[ann["category_id"]])
+            raw_masks.append(self._decode_mask(ann, orig_h, orig_w))
+
+        # Non-fractured images: create a dummy background instance so the
+        # tensor shapes are consistent (YOLACT handles empty targets fine too,
+        # but keeping consistent is safer).
+        if len(boxes) == 0:
+            boxes = [[0.0, 0.0, float(orig_w), float(orig_h)]]
+            class_labels = [0]  # background / non-fractured
+            raw_masks = [np.zeros((orig_h, orig_w), dtype=np.uint8)]
+
+        # ── Albumentations ───────────────────────────────────────────────────
+        transformed = self.transforms(
+            image=image,
+            masks=raw_masks,
+            bboxes=boxes,
+            class_labels=class_labels,
+        )
+        image_t = transformed["image"]          # [3, H, W]
+        boxes_t = transformed["bboxes"]
+        labels_t = transformed["class_labels"]
+        masks_t = transformed["masks"]           # list of H×W arrays
+
+        _, H, W = image_t.shape
+
+        # ── Build target tensors ─────────────────────────────────────────────
+        if len(boxes_t) == 0:
+            # All boxes removed by augmentation (e.g. min_visibility)
+            boxes_out = torch.zeros((0, 4), dtype=torch.float32)
+            labels_out = torch.zeros((0,), dtype=torch.long)
+            masks_out = torch.zeros((0, H, W), dtype=torch.float32)
+        else:
+            boxes_np = np.array(boxes_t, dtype=np.float32)
+            # Normalise to [0, 1]
+            boxes_np[:, [0, 2]] /= W
+            boxes_np[:, [1, 3]] /= H
+            boxes_np = np.clip(boxes_np, 0.0, 1.0)
+
+            boxes_out = torch.from_numpy(boxes_np)
+            labels_out = torch.tensor(labels_t, dtype=torch.long)
+            # Albumentations >=2.x returns masks as tensors; older versions return numpy.
+            def to_float_tensor(m):
+                if isinstance(m, torch.Tensor):
+                    return m.float()
+                return torch.from_numpy(np.array(m, dtype=np.float32))
+            masks_out = torch.stack([to_float_tensor(m) for m in masks_t])
+
+        target = {
+            "boxes": boxes_out,
+            "labels": labels_out,
+            "masks": masks_out,
+            "image_id": img_id,
+        }
+        return image_t, target
+
+
+# ─── Collate ──────────────────────────────────────────────────────────────────
+
+def detection_collate(batch):
+    """
+    Custom collate for object detection.
+    Images are stacked; targets are kept as a list (variable number of instances).
+    """
+    images, targets = zip(*batch)
+    images = torch.stack(images, dim=0)  # [B, 3, H, W]
+    return images, list(targets)
+
+
+# ─── DataLoader factory ───────────────────────────────────────────────────────
+
+def get_dataloader(
+    image_dir: str,
+    ann_file: str,
+    split: str = "train",
+    img_size: int = 550,
+    batch_size: int = 8,
+    num_workers: int = 4,
+    pin_memory: bool = True,
+) -> DataLoader:
+    """
+    Returns a DataLoader for the given split.
+
+    Args:
+        image_dir  : path to images/ folder for this split
+        ann_file   : path to annotations.json for this split
+        split      : "train" | "val" | "test"
+        img_size   : input resolution fed to the network (default 550 for YOLACT+)
+        batch_size : mini-batch size
+        num_workers: parallel data-loading workers
+        pin_memory : pin CPU memory for faster GPU transfer
+
+    Returns:
+        torch.utils.data.DataLoader
+    """
+    dataset = FracAtlasDataset(
+        image_dir=image_dir,
+        ann_file=ann_file,
+        img_size=img_size,
+        split=split,
+    )
+    shuffle = split == "train"
+    # pin_memory only works when CUDA is available
+    import torch
+    use_pin = pin_memory and torch.cuda.is_available()
+    loader = DataLoader(
+        dataset,
+        batch_size=batch_size,
+        shuffle=shuffle,
+        num_workers=num_workers,
+        pin_memory=use_pin,
+        collate_fn=detection_collate,
+        drop_last=(split == "train"),
+    )
+    return loader

model.py

@@ -0,0 +1,380 @@
+"""
+YOLACT+ with ResNet-18 Backbone
+=================================
+A faithful implementation of YOLACT+ adapted for a lightweight ResNet-18 backbone.
+
+Architecture:
+  Backbone  : ResNet-18  (torchvision, ImageNet pre-trained)
+  Neck      : FPN  (Feature Pyramid Network)
+  Head      : PredictionHead  (class + box + mask coefficient)
+  Proto     : ProtoNet  (generates prototype masks)
+  Mask      : linear combination of prototypes × coefficients
+  NMS       : Fast NMS (YOLACT-style)
+"""
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torchvision.models import resnet18, ResNet18_Weights
+from torchvision.ops import nms
+
+
+# ─── Constants ────────────────────────────────────────────────────────────────
+NUM_PROTOTYPES = 32
+FPN_CHANNELS   = 256
+PROTO_CHANNELS = 256
+
+
+# ─── Backbone ─────────────────────────────────────────────────────────────────
+
+class ResNet18Backbone(nn.Module):
+    """
+    ResNet-18 feature extractor.
+    Returns C3, C4, C5 feature maps (strides 8, 16, 32).
+    """
+
+    def __init__(self, pretrained: bool = True):
+        super().__init__()
+        weights = ResNet18_Weights.IMAGENET1K_V1 if pretrained else None
+        base = resnet18(weights=weights)
+
+        self.layer0 = nn.Sequential(base.conv1, base.bn1, base.relu, base.maxpool)
+        self.layer1 = base.layer1   # stride 4,  channels 64
+        self.layer2 = base.layer2   # stride 8,  channels 128  → C3
+        self.layer3 = base.layer3   # stride 16, channels 256  → C4
+        self.layer4 = base.layer4   # stride 32, channels 512  → C5
+
+        self.out_channels = [128, 256, 512]  # C3, C4, C5
+
+    def forward(self, x):
+        x = self.layer0(x)
+        x = self.layer1(x)
+        c3 = self.layer2(x)
+        c4 = self.layer3(c3)
+        c5 = self.layer4(c4)
+        return c3, c4, c5
+
+
+# ─── FPN ──────────────────────────────────────────────────────────────────────
+
+class FPN(nn.Module):
+    """
+    5-level FPN: P3–P7 (P6, P7 generated by strided convolution on P5).
+    """
+
+    def __init__(self, in_channels: list, out_channels: int = FPN_CHANNELS):
+        super().__init__()
+        self.lateral_convs = nn.ModuleList(
+            [nn.Conv2d(c, out_channels, 1) for c in in_channels]
+        )
+        self.output_convs = nn.ModuleList(
+            [nn.Conv2d(out_channels, out_channels, 3, padding=1) for _ in in_channels]
+        )
+        # Extra levels P6, P7
+        self.p6_conv = nn.Conv2d(out_channels, out_channels, 3, stride=2, padding=1)
+        self.p7_conv = nn.Conv2d(out_channels, out_channels, 3, stride=2, padding=1)
+
+    def forward(self, features):
+        c3, c4, c5 = features
+        lat = [l(f) for l, f in zip(self.lateral_convs, [c3, c4, c5])]
+
+        # Top-down pathway
+        lat[1] = lat[1] + F.interpolate(lat[2], size=lat[1].shape[-2:], mode="nearest")
+        lat[0] = lat[0] + F.interpolate(lat[1], size=lat[0].shape[-2:], mode="nearest")
+
+        p3 = self.output_convs[0](lat[0])
+        p4 = self.output_convs[1](lat[1])
+        p5 = self.output_convs[2](lat[2])
+        p6 = self.p6_conv(p5)
+        p7 = self.p7_conv(F.relu(p6))
+
+        return [p3, p4, p5, p6, p7]
+
+
+# ─── ProtoNet ─────────────────────────────────────────────────────────────────
+
+class ProtoNet(nn.Module):
+    """
+    Generates K prototype masks from the P3 feature map.
+    Output: [B, K, H/4, W/4]
+    """
+
+    def __init__(self, in_channels: int = FPN_CHANNELS, num_protos: int = NUM_PROTOTYPES):
+        super().__init__()
+        self.proto_net = nn.Sequential(
+            nn.Conv2d(in_channels, PROTO_CHANNELS, 3, padding=1), nn.ReLU(),
+            nn.Conv2d(PROTO_CHANNELS, PROTO_CHANNELS, 3, padding=1), nn.ReLU(),
+            nn.Conv2d(PROTO_CHANNELS, PROTO_CHANNELS, 3, padding=1), nn.ReLU(),
+            nn.Upsample(scale_factor=2, mode="bilinear", align_corners=False),
+            nn.Conv2d(PROTO_CHANNELS, PROTO_CHANNELS, 3, padding=1), nn.ReLU(),
+            nn.Conv2d(PROTO_CHANNELS, num_protos, 1),
+        )
+
+    def forward(self, p3):
+        return self.proto_net(p3)   # [B, K, H', W']
+
+
+# ─── Anchor generator ─────────────────────────────────────────────────────────
+
+class AnchorGenerator:
+    """
+    Pre-computed anchor boxes for each FPN level.
+    Scales: [24, 48, 96, 192, 384]  (for 550×550 input)
+    Aspect ratios: [1.0, 0.5, 2.0]
+    """
+
+    SCALES       = [24, 48, 96, 192, 384]
+    ASPECT_RATIOS = [1.0, 0.5, 2.0]
+
+    def __init__(self, img_size: int = 550):
+        self.img_size = img_size
+        self.num_anchors_per_cell = len(self.ASPECT_RATIOS)
+
+    def make_anchors(self, feature_sizes: list) -> torch.Tensor:
+        """Returns [total_anchors, 4] in cx/cy/w/h format (normalised 0-1)."""
+        all_anchors = []
+        for lvl, (fh, fw) in enumerate(feature_sizes):
+            scale = self.SCALES[lvl]
+            for row in range(fh):
+                for col in range(fw):
+                    cx = (col + 0.5) / fw
+                    cy = (row + 0.5) / fh
+                    for ar in self.ASPECT_RATIOS:
+                        w = scale * (ar ** 0.5) / self.img_size
+                        h = scale / (ar ** 0.5) / self.img_size
+                        all_anchors.append([cx, cy, w, h])
+        return torch.tensor(all_anchors, dtype=torch.float32)
+
+
+# ─── Prediction Head ──────────────────────────────────────────────────────────
+
+class PredictionHead(nn.Module):
+    """
+    Shared prediction head applied to each FPN level.
+    Outputs:
+        cls_pred  : [B, A, num_classes+1]
+        box_pred  : [B, A, 4]
+        coef_pred : [B, A, K]
+    """
+
+    def __init__(
+        self,
+        in_channels: int = FPN_CHANNELS,
+        num_classes: int = 2,
+        num_anchors: int = 3,
+        num_protos: int = NUM_PROTOTYPES,
+    ):
+        super().__init__()
+        self.num_classes = num_classes
+        self.num_anchors = num_anchors
+
+        self.shared = nn.Sequential(
+            nn.Conv2d(in_channels, in_channels, 3, padding=1), nn.ReLU(),
+            nn.Conv2d(in_channels, in_channels, 3, padding=1), nn.ReLU(),
+            nn.Conv2d(in_channels, in_channels, 3, padding=1), nn.ReLU(),
+            nn.Conv2d(in_channels, in_channels, 3, padding=1), nn.ReLU(),
+        )
+        self.cls_layer  = nn.Conv2d(in_channels, num_anchors * (num_classes + 1), 1)
+        self.box_layer  = nn.Conv2d(in_channels, num_anchors * 4, 1)
+        self.coef_layer = nn.Conv2d(in_channels, num_anchors * num_protos, 1)
+
+    def forward(self, feat):
+        B, _, H, W = feat.shape
+        x = self.shared(feat)
+
+        cls  = self.cls_layer(x)               # [B, A*(C+1), H, W]
+        box  = self.box_layer(x)               # [B, A*4,     H, W]
+        coef = self.coef_layer(x)              # [B, A*K,     H, W]
+
+        # Reshape to [B, H*W*A, ...]
+        A, C, K = self.num_anchors, self.num_classes, NUM_PROTOTYPES
+        cls  = cls.permute(0, 2, 3, 1).contiguous().view(B, -1, C + 1)
+        box  = box.permute(0, 2, 3, 1).contiguous().view(B, -1, 4)
+        coef = coef.permute(0, 2, 3, 1).contiguous().view(B, -1, K)
+        coef = torch.tanh(coef)
+
+        return cls, box, coef
+
+
+# ─── YOLACT+ ──────────────────────────────────────────────────────────────────
+
+class YOLACTPlus(nn.Module):
+    """
+    YOLACT+ with ResNet-18 backbone.
+
+    Args:
+        num_classes : number of foreground classes (background added internally)
+        img_size    : input image resolution (square, default 550)
+        pretrained  : use ImageNet-pretrained ResNet-18
+    """
+
+    def __init__(
+        self,
+        num_classes: int = 1,
+        img_size: int = 550,
+        pretrained: bool = True,
+    ):
+        super().__init__()
+        self.num_classes = num_classes
+        self.img_size = img_size
+
+        self.backbone  = ResNet18Backbone(pretrained=pretrained)
+        self.fpn       = FPN(self.backbone.out_channels)
+        self.proto_net = ProtoNet(FPN_CHANNELS, NUM_PROTOTYPES)
+        self.head      = PredictionHead(
+            FPN_CHANNELS, num_classes, len(AnchorGenerator.ASPECT_RATIOS), NUM_PROTOTYPES
+        )
+        self.anchor_gen = AnchorGenerator(img_size)
+        self._anchors   = None  # cached after first forward pass
+
+    # ── Forward ───────────────────────────────────────────────────────────────
+
+    def forward(self, images: torch.Tensor):
+        """
+        Args:
+            images : [B, 3, H, W]
+        Returns (training mode):
+            {
+              "cls_pred"  : [B, total_anchors, num_classes+1]
+              "box_pred"  : [B, total_anchors, 4]
+              "coef_pred" : [B, total_anchors, K]
+              "proto_out" : [B, K, H', W']
+              "anchors"   : [total_anchors, 4]  (cx/cy/w/h, normalised)
+            }
+        """
+        features = self.backbone(images)
+        fpn_feats = self.fpn(features)
+
+        proto_out = self.proto_net(fpn_feats[0])   # P3 → prototypes
+
+        # Cache anchors (they depend only on feature map sizes)
+        if self._anchors is None or self._anchors.device != images.device:
+            feat_sizes = [(f.shape[2], f.shape[3]) for f in fpn_feats]
+            self._anchors = self.anchor_gen.make_anchors(feat_sizes).to(images.device)
+
+        cls_preds, box_preds, coef_preds = [], [], []
+        for feat in fpn_feats:
+            cls, box, coef = self.head(feat)
+            cls_preds.append(cls)
+            box_preds.append(box)
+            coef_preds.append(coef)
+
+        return {
+            "cls_pred":  torch.cat(cls_preds,  dim=1),
+            "box_pred":  torch.cat(box_preds,  dim=1),
+            "coef_pred": torch.cat(coef_preds, dim=1),
+            "proto_out": proto_out,
+            "anchors":   self._anchors,
+        }
+
+    # ── Post-processing (inference) ───────────────────────────────────────────
+
+    @torch.no_grad()
+    def predict(
+        self,
+        images: torch.Tensor,
+        score_thresh: float = 0.3,
+        nms_thresh: float   = 0.5,
+        top_k: int          = 100,
+    ) -> list:
+        """
+        Run inference and return decoded predictions.
+
+        Returns:
+            List (one per image) of dicts:
+                boxes  : [N, 4]  x1y1x2y2 normalised
+                scores : [N]
+                labels : [N]
+                masks  : [N, H, W]  float binary masks (upsampled to input size)
+        """
+        self.eval()
+        out = self.forward(images)
+
+        cls_pred  = out["cls_pred"]   # [B, A, C+1]
+        box_pred  = out["box_pred"]   # [B, A, 4]
+        coef_pred = out["coef_pred"]  # [B, A, K]
+        proto     = out["proto_out"]  # [B, K, H', W']
+        anchors   = out["anchors"]    # [A, 4]
+
+        results = []
+        B = images.shape[0]
+        for i in range(B):
+            scores_all = torch.softmax(cls_pred[i], dim=-1)        # [A, C+1]
+            scores, labels = scores_all[:, 1:].max(dim=-1)         # foreground only
+
+            keep_mask = scores > score_thresh
+            if keep_mask.sum() == 0:
+                results.append({"boxes": torch.zeros(0, 4), "scores": torch.zeros(0),
+                                 "labels": torch.zeros(0, dtype=torch.long),
+                                 "masks": torch.zeros(0, self.img_size, self.img_size)})
+                continue
+
+            scores  = scores[keep_mask]
+            labels  = labels[keep_mask]
+            boxes_d = box_pred[i][keep_mask]   # deltas
+            coefs   = coef_pred[i][keep_mask]  # [N, K]
+            anch    = anchors[keep_mask]        # [N, 4]
+
+            # Decode box deltas → cx/cy/w/h
+            pred_cx = boxes_d[:, 0] * anch[:, 2] + anch[:, 0]
+            pred_cy = boxes_d[:, 1] * anch[:, 3] + anch[:, 1]
+            pred_w  = torch.exp(boxes_d[:, 2]) * anch[:, 2]
+            pred_h  = torch.exp(boxes_d[:, 3]) * anch[:, 3]
+
+            # → x1y1x2y2
+            x1 = torch.clamp(pred_cx - pred_w / 2, 0, 1)
+            y1 = torch.clamp(pred_cy - pred_h / 2, 0, 1)
+            x2 = torch.clamp(pred_cx + pred_w / 2, 0, 1)
+            y2 = torch.clamp(pred_cy + pred_h / 2, 0, 1)
+            boxes_xyxy = torch.stack([x1, y1, x2, y2], dim=1)
+
+            # ── Filter out oversized boxes ────────────────────────────
+            # Remove boxes whose area exceeds 50% of the image area.
+            # These are almost always spurious full-image anchors.
+            box_w = boxes_xyxy[:, 2] - boxes_xyxy[:, 0]
+            box_h = boxes_xyxy[:, 3] - boxes_xyxy[:, 1]
+            box_area = box_w * box_h
+            size_mask = box_area < 0.50   # keep boxes < 50% of image area
+            boxes_xyxy = boxes_xyxy[size_mask]
+            scores     = scores[size_mask]
+            labels     = labels[size_mask]
+            coefs      = coefs[size_mask]
+
+            if boxes_xyxy.shape[0] == 0:
+                results.append({"boxes": torch.zeros(0, 4), "scores": torch.zeros(0),
+                                 "labels": torch.zeros(0, dtype=torch.long),
+                                 "masks": torch.zeros(0, self.img_size, self.img_size)})
+                continue
+
+            # NMS (pixel-scale for torchvision nms)
+            scale = float(self.img_size)
+            keep = nms(boxes_xyxy * scale, scores, nms_thresh)
+            keep = keep[:top_k]
+
+            boxes_xyxy = boxes_xyxy[keep]
+            scores     = scores[keep]
+            labels     = labels[keep]
+            coefs      = coefs[keep]   # [N, K]
+
+            # Decode masks: proto [K, H', W'], coefs [N, K]
+            proto_i = proto[i]                                    # [K, H', W']
+            K, pH, pW = proto_i.shape
+            proto_flat = proto_i.view(K, -1).T                   # [H'W', K]
+            mask_flat  = torch.sigmoid(proto_flat @ coefs.T)     # [H'W', N]
+            masks_raw  = mask_flat.T.view(len(keep), pH, pW)     # [N, H', W']
+
+            # Upsample to input resolution
+            masks_up = F.interpolate(
+                masks_raw.unsqueeze(0), size=(self.img_size, self.img_size),
+                mode="bilinear", align_corners=False
+            ).squeeze(0)
+            masks_bin = (masks_up > 0.5).float()
+
+            results.append({
+                "boxes":  boxes_xyxy.cpu(),
+                "scores": scores.cpu(),
+                "labels": labels.cpu(),
+                "masks":  masks_bin.cpu(),
+            })
+
+        return results

requirements.txt

@@ -0,0 +1,10 @@
+# YOLACT+ FracAtlas Requirements
+# Install with: pip install -r requirements.txt
+
+torch>=2.0.0
+torchvision>=0.15.0
+albumentations>=1.3.0
+pycocotools>=2.0.6
+opencv-python>=4.7.0
+numpy>=1.24.0
+tqdm>=4.65.0

requirements_hf.txt

@@ -0,0 +1,9 @@
+torch
+torchvision
+albumentations
+pycocotools
+opencv-python-headless
+numpy
+matplotlib
+gradio
+huggingface_hub