feat: integrate real SensiNet mammography model

- Add SensiNet dual-stream architecture (Xception + EfficientNet-B3 with CBAM)
- Replace mock inference with real Bayesian MC-Dropout prediction (10 passes)
- Add /analyze endpoint with SSRF protection (allowlist + private IP blocking)
- Add malignancy_probability to response schema
- Add training and data preparation utility scripts
- Model weights (131MB .pth) gitignored — must be downloaded separately

.env.example

@@ -1,2 +1,6 @@
-MODEL_MODE=mock
-MODEL_VERSION=baseline-mock-v1
+MODEL_MODE=real
+MODEL_VERSION=sensinet-v1
+# Path to .pth weights (defaults to weights/advanced_model_best.pth)
+# MODEL_WEIGHTS=weights/advanced_model_best.pth
+# Comma-separated list of allowed hostnames for the /analyze endpoint (SSRF protection)
+# ALLOWED_IMAGE_HOSTS=your-project.supabase.co

.gitignore

@@ -3,3 +3,4 @@ __pycache__/
 *.pyc
 .env
 .DS_Store
+weights/*.pth

app/architecture.py

@@ -0,0 +1,100 @@
+"""
+SensiNet Dual-Stream Architecture for Mammographic Classification.
+
+Architecture:
+    - Stream 1: Xception (feature-rich legacy backbone)
+    - Stream 2: EfficientNet-B3 (modern efficient backbone)
+    - Fusion: Projected feature maps concatenated and refined via CBAM attention
+    - Output: Single logit (sigmoid → malignancy probability)
+
+Source: Aredeksu/SensiNet-Mammography on Hugging Face (Apache-2.0 license)
+Trained on: CBIS-DDSM dataset
+"""
+
+import timm
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+class ChannelAttention(nn.Module):
+    def __init__(self, in_planes: int, ratio: int = 16):
+        super().__init__()
+        self.avg_pool = nn.AdaptiveAvgPool2d(1)
+        self.max_pool = nn.AdaptiveMaxPool2d(1)
+        self.fc1 = nn.Conv2d(in_planes, in_planes // ratio, 1, bias=False)
+        self.relu1 = nn.ReLU()
+        self.fc2 = nn.Conv2d(in_planes // ratio, in_planes, 1, bias=False)
+        self.sigmoid = nn.Sigmoid()
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        avg_out = self.fc2(self.relu1(self.fc1(self.avg_pool(x))))
+        max_out = self.fc2(self.relu1(self.fc1(self.max_pool(x))))
+        return self.sigmoid(avg_out + max_out)
+
+
+class SpatialAttention(nn.Module):
+    def __init__(self, kernel_size: int = 7):
+        super().__init__()
+        padding = 3 if kernel_size == 7 else 1
+        self.conv1 = nn.Conv2d(2, 1, kernel_size, padding=padding, bias=False)
+        self.sigmoid = nn.Sigmoid()
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        avg_out = torch.mean(x, dim=1, keepdim=True)
+        max_out, _ = torch.max(x, dim=1, keepdim=True)
+        x = torch.cat([avg_out, max_out], dim=1)
+        return self.sigmoid(self.conv1(x))
+
+
+class CBAM(nn.Module):
+    def __init__(self, planes: int):
+        super().__init__()
+        self.ca = ChannelAttention(planes)
+        self.sa = SpatialAttention()
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = x * self.ca(x)
+        x = x * self.sa(x)
+        return x
+
+
+class AdvancedBreastCancerModel(nn.Module):
+    def __init__(self) -> None:
+        super().__init__()
+
+        # Stream 1: Xception → 2048 channels
+        self.stream1 = timm.create_model("xception", pretrained=False, num_classes=0)
+        # Stream 2: EfficientNet-B3 → 1536 channels
+        self.stream2 = timm.create_model("efficientnet_b3", pretrained=False, num_classes=0)
+
+        # Projection layers to 512 each
+        self.proj1 = nn.Conv2d(2048, 512, 1)
+        self.proj2 = nn.Conv2d(1536, 512, 1)
+
+        # Attention fusion (512 + 512 = 1024)
+        self.fusion_attention = CBAM(1024)
+
+        # Classification head
+        self.classifier = nn.Sequential(
+            nn.AdaptiveAvgPool2d(1),
+            nn.Flatten(),
+            nn.Linear(1024, 512),
+            nn.BatchNorm1d(512),
+            nn.ReLU(),
+            nn.Dropout(0.5),
+            nn.Linear(512, 1),
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        f1 = self.stream1.forward_features(x)  # [B, 2048, H1, W1]
+        f2 = self.stream2.forward_features(x)  # [B, 1536, H2, W2]
+
+        if f1.shape[2:] != f2.shape[2:]:
+            f2 = F.interpolate(f2, size=f1.shape[2:], mode="bilinear", align_corners=False)
+
+        p1 = self.proj1(f1)
+        p2 = self.proj2(f2)
+        concat = torch.cat([p1, p2], dim=1)
+        refined = self.fusion_attention(concat)
+        return self.classifier(refined)

app/main.py

@@ -1,10 +1,21 @@
 import io
+import ipaddress
+import os
+import socket
+from urllib.parse import urlparse
 
+import requests as http_requests
 from fastapi import FastAPI, File, HTTPException, UploadFile
 from PIL import Image, UnidentifiedImageError
 
 from app.model import MammogramModel
-from app.schemas import PredictResponse
+from app.schemas import AnalyzeRequest, PredictResponse
+
+# Comma-separated list of allowed URL hostnames (e.g. your Supabase storage host)
+_ALLOWED_HOSTS_ENV = os.getenv("ALLOWED_IMAGE_HOSTS", "")
+ALLOWED_HOSTS: set[str] = {
+    h.strip().lower() for h in _ALLOWED_HOSTS_ENV.split(",") if h.strip()
+}
 
 app = FastAPI(title="Mammogram Inference API", version="0.1.0")
 model = MammogramModel()
@@ -31,3 +42,44 @@ async def predict(file: UploadFile = File(...)) -> PredictResponse:
 
     result = model.predict(image)
     return PredictResponse(**result)
+
+
+def _validate_url(url: str) -> str:
+    """Validate image URL to prevent SSRF attacks."""
+    parsed = urlparse(url)
+    if parsed.scheme not in ("https",):
+        raise HTTPException(status_code=400, detail="Only HTTPS URLs are allowed")
+    hostname = (parsed.hostname or "").lower()
+    if not hostname:
+        raise HTTPException(status_code=400, detail="Invalid URL")
+    if ALLOWED_HOSTS and hostname not in ALLOWED_HOSTS:
+        raise HTTPException(status_code=400, detail="Image host not in allowlist")
+    # Block private/loopback IPs to prevent SSRF
+    try:
+        for info in socket.getaddrinfo(hostname, None):
+            addr = info[4][0]
+            ip = ipaddress.ip_address(addr)
+            if ip.is_private or ip.is_loopback or ip.is_link_local:
+                raise HTTPException(status_code=400, detail="URL resolves to a private address")
+    except socket.gaierror as exc:
+        raise HTTPException(status_code=400, detail="Cannot resolve hostname") from exc
+    return url
+
+
+@app.post("/analyze", response_model=PredictResponse)
+def analyze(body: AnalyzeRequest) -> PredictResponse:
+    """Accept a public image URL, download it, and run inference."""
+    _validate_url(body.image_url)
+    try:
+        resp = http_requests.get(body.image_url, timeout=30)
+        resp.raise_for_status()
+    except http_requests.RequestException as exc:
+        raise HTTPException(status_code=400, detail=f"Failed to fetch image: {exc}") from exc
+
+    try:
+        image = Image.open(io.BytesIO(resp.content))
+    except UnidentifiedImageError as exc:
+        raise HTTPException(status_code=400, detail="URL did not return a valid image") from exc
+
+    result = model.predict(image)
+    return PredictResponse(**result)

app/model.py

@@ -1,50 +1,153 @@
-import hashlib
+import logging
 import os
+from pathlib import Path
 
 import numpy as np
+import torch
+import torch.nn as nn
 from PIL import Image
+from torchvision import transforms
+
+from app.architecture import AdvancedBreastCancerModel
+
+logger = logging.getLogger(__name__)
+
+# ImageNet normalisation (same as SensiNet training pipeline)
+TRANSFORM = transforms.Compose([
+    transforms.Resize((299, 299)),
+    transforms.ToTensor(),
+    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
+])
+
+WEIGHTS_DIR = Path(__file__).resolve().parent.parent / "weights"
+DEFAULT_WEIGHTS = WEIGHTS_DIR / "advanced_model_best.pth"
+
+# Malignancy probability threshold (same as SensiNet default)
+THRESHOLD = 0.40
+
+# Number of Bayesian MC-Dropout forward passes
+MC_PASSES = 10
+
+
+def _prob_to_birads(prob: float) -> int:
+    """Map malignancy probability to BI-RADS category."""
+    if prob < 0.10:
+        return 1  # Negative
+    if prob < 0.25:
+        return 2  # Benign
+    if prob < 0.50:
+        return 3  # Probably benign
+    if prob < 0.75:
+        return 4  # Suspicious
+    return 5  # Highly suggestive of malignancy
+
+
+def _birads_findings(birads: int, prob: float, prediction: str) -> str:
+    templates = {
+        1: "No suspicious findings detected. Mammographic appearance is unremarkable.",
+        2: "Benign-appearing pattern identified. Correlate with prior imaging if available.",
+        3: "Probably benign appearance. Short-interval follow-up may be considered.",
+        4: "Suspicious abnormality pattern detected. Tissue biopsy is recommended.",
+        5: "Highly suggestive of malignancy. Urgent diagnostic workup is recommended.",
+    }
+    base = templates.get(birads, "Analysis complete.")
+    return f"Model prediction: {prediction} (probability {prob:.1%}). {base}"
 
 
 class MammogramModel:
+    """Loads the SensiNet dual-stream model and runs inference."""
+
     def __init__(self) -> None:
-        self.mode = os.getenv("MODEL_MODE", "mock")
-        self.version = os.getenv("MODEL_VERSION", "baseline-mock-v1")
+        self.mode = os.getenv("MODEL_MODE", "real")
+        self.version = os.getenv("MODEL_VERSION", "sensinet-v1")
+        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        self._model: AdvancedBreastCancerModel | None = None
+
+        weights_path = Path(os.getenv("MODEL_WEIGHTS", str(DEFAULT_WEIGHTS)))
+        if weights_path.exists():
+            self._load_model(weights_path)
+        else:
+            logger.warning("Weights not found at %s — falling back to mock mode", weights_path)
+            self.mode = "mock"
+
+    def _load_model(self, weights_path: Path) -> None:
+        logger.info("Loading SensiNet model from %s onto %s …", weights_path, self.device)
+        net = AdvancedBreastCancerModel()
+        state = torch.load(weights_path, map_location=self.device, weights_only=False)
+        net.load_state_dict(state)
+        net.to(self.device)
+        net.eval()
+        self._model = net
+        logger.info("Model loaded successfully.")
+
+    # ------------------------------------------------------------------
 
     def predict(self, image: Image.Image) -> dict:
-        # Placeholder implementation: deterministic mock score from image bytes.
-        # Replace with your real PyTorch model loading and inference later.
-        rgb = image.convert("L")
-        arr = np.array(rgb, dtype=np.float32) / 255.0
-        mean_intensity = float(arr.mean())
+        if self._model is None or self.mode == "mock":
+            return self._mock_predict(image)
+        return self._real_predict(image)
+
+    # ------------------------------------------------------------------
+    # Real inference with Bayesian MC-Dropout
+    # ------------------------------------------------------------------
+
+    def _real_predict(self, image: Image.Image) -> dict:
+        rgb = image.convert("RGB")
+        tensor = TRANSFORM(rgb).unsqueeze(0).to(self.device)
 
+        def enable_dropout(m: nn.Module) -> None:
+            if isinstance(m, (nn.Dropout, nn.Dropout2d)):
+                m.train()
+
+        self._model.apply(enable_dropout)
+
+        mc_predictions: list[float] = []
+        with torch.no_grad():
+            for _ in range(MC_PASSES):
+                logits = self._model(tensor)
+                prob = torch.sigmoid(logits).item()
+                mc_predictions.append(prob)
+
+        self._model.eval()
+
+        prob_malig = float(np.mean(mc_predictions))
+        variance = float(np.var(mc_predictions))
+
+        decision_confidence = max(0.50, 0.99 - (variance * 2.0))
+        if prob_malig < 0.10 or prob_malig > 0.90:
+            decision_confidence = min(0.99, decision_confidence + 0.10)
+
+        prediction = "Malignant" if prob_malig >= THRESHOLD else "Benign"
+        birads = _prob_to_birads(prob_malig)
+
+        return {
+            "birads": birads,
+            "confidence": round(decision_confidence, 3),
+            "malignancy_probability": round(prob_malig, 3),
+            "findings_text": _birads_findings(birads, prob_malig, prediction),
+            "model_version": self.version,
+        }
+
+    # ------------------------------------------------------------------
+    # Deterministic mock fallback (no weights needed)
+    # ------------------------------------------------------------------
+
+    @staticmethod
+    def _mock_predict(image: Image.Image) -> dict:
+        import hashlib
+
+        arr = np.array(image.convert("L"), dtype=np.float32) / 255.0
         digest = hashlib.sha256(arr.tobytes()).hexdigest()
         seed = int(digest[:8], 16)
         rng = np.random.default_rng(seed)
-        jitter = float(rng.uniform(-0.04, 0.04))
-
-        raw = min(max(mean_intensity + jitter, 0.0), 1.0)
-
-        if raw < 0.20:
-            birads = 1
-            findings = "No suspicious findings detected in this preliminary model pass."
-        elif raw < 0.35:
-            birads = 2
-            findings = "Likely benign pattern; correlate with prior imaging."
-        elif raw < 0.55:
-            birads = 3
-            findings = "Probably benign appearance; short-interval follow-up may be considered."
-        elif raw < 0.75:
-            birads = 4
-            findings = "Suspicious abnormality pattern; biopsy correlation recommended."
-        else:
-            birads = 5
-            findings = "Highly suggestive of malignancy pattern; urgent diagnostic follow-up recommended."
+        raw = float(min(max(arr.mean() + rng.uniform(-0.04, 0.04), 0.0), 1.0))
 
-        confidence = max(0.55, min(0.98, 0.55 + abs(raw - 0.5)))
+        birads = _prob_to_birads(raw)
 
         return {
             "birads": birads,
-            "confidence": round(confidence, 3),
-            "findings_text": findings,
-            "model_version": self.version,
+            "confidence": round(max(0.55, min(0.98, 0.55 + abs(raw - 0.5))), 3),
+            "malignancy_probability": round(raw, 3),
+            "findings_text": _birads_findings(birads, raw, "Malignant" if raw >= THRESHOLD else "Benign"),
+            "model_version": "mock-v1",
         }

app/schemas.py

@@ -1,8 +1,13 @@
 from pydantic import BaseModel, Field
 
 
+class AnalyzeRequest(BaseModel):
+    image_url: str
+
+
 class PredictResponse(BaseModel):
     birads: int = Field(ge=0, le=6)
     confidence: float = Field(ge=0.0, le=1.0)
+    malignancy_probability: float = Field(ge=0.0, le=1.0, default=0.0)
     findings_text: str
     model_version: str

prepare_data.py

@@ -0,0 +1,75 @@
+"""
+prepare_data.py  —  organise raw CBIS-DDSM images into train/val folder structure.
+
+If your downloaded images are already in  data/train/benign  etc., skip this.
+
+Usage
+-----
+python prepare_data.py --images /path/to/raw/images --csv /path/to/labels.csv
+
+CSV must have columns:  file_path, pathology
+  pathology values:  BENIGN, MALIGNANT  (or benign, malignant)
+
+Output
+------
+data/
+  train/benign/  train/malignant/
+  val/benign/    val/malignant/
+"""
+
+import argparse
+import os
+import shutil
+import random
+
+TRAIN_RATIO = 0.85
+
+
+def prepare(images_dir: str, csv_path: str, output_dir: str, seed: int = 42) -> None:
+    import csv
+
+    random.seed(seed)
+
+    records: list[tuple[str, str]] = []
+    with open(csv_path, newline="") as f:
+        reader = csv.DictReader(f)
+        for row in reader:
+            # normalise label
+            label = row.get("pathology", row.get("label", "")).strip().lower()
+            if label in ("benign", "benign_without_callback"):
+                label = "benign"
+            elif label in ("malignant",):
+                label = "malignant"
+            else:
+                continue  # skip unknown labels
+
+            img_path = os.path.join(images_dir, row.get("file_path", "").strip())
+            if os.path.isfile(img_path):
+                records.append((img_path, label))
+
+    print(f"Found {len(records)} labelled images")
+    random.shuffle(records)
+
+    split = int(len(records) * TRAIN_RATIO)
+    splits = {"train": records[:split], "val": records[split:]}
+
+    for split_name, items in splits.items():
+        for label in ("benign", "malignant"):
+            os.makedirs(os.path.join(output_dir, split_name, label), exist_ok=True)
+        for src, label in items:
+            fname = os.path.basename(src)
+            dst = os.path.join(output_dir, split_name, label, fname)
+            shutil.copy2(src, dst)
+        counts = {lbl: sum(1 for _, l in items if l == lbl) for lbl in ("benign", "malignant")}
+        print(f"{split_name}: {counts}")
+
+    print(f"Data prepared in {output_dir}/")
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--images", required=True, help="Directory containing raw image files")
+    parser.add_argument("--csv", required=True, help="CSV file with file_path and pathology columns")
+    parser.add_argument("--output", default="data", help="Output directory")
+    args = parser.parse_args()
+    prepare(args.images, args.csv, args.output)

requirements.txt

@@ -4,3 +4,8 @@ python-multipart==0.0.9
 pydantic==2.11.3
 numpy==2.2.4
 pillow==11.1.0
+torch>=2.0.0
+torchvision>=0.15.0
+timm>=0.9.0
+opencv-python-headless>=4.8.0
+requests>=2.28.0

train.py

@@ -0,0 +1,188 @@
+"""
+train.py  —  Fine-tune the SensiNet dual-stream model on a binary mammogram dataset.
+
+Expected dataset layout
+-----------------------
+data/
+  train/
+    benign/        <- benign mammogram images (.jpg / .png / .dcm converted to jpg)
+    malignant/     <- malignant mammogram images
+  val/
+    benign/
+    malignant/
+
+If you only have a flat folder + CSV (CBIS-DDSM style), run prepare_data.py first.
+
+Usage
+-----
+python train.py --data data --output models/advanced_model_best.pth
+
+The saved file is a raw state_dict compatible with MammogramModel._load_model().
+"""
+
+import argparse
+import os
+from pathlib import Path
+
+import torch
+import torch.nn as nn
+from torch.optim import Adam
+from torch.optim.lr_scheduler import ReduceLROnPlateau
+from torch.utils.data import DataLoader
+from torchvision import datasets, transforms
+
+from app.architecture import AdvancedBreastCancerModel
+
+# ── Hyperparameters ────────────────────────────────────────────────────────────
+IMG_SIZE = 299          # Xception / EfficientNet-B3 both happy at 299
+BATCH_SIZE = 16
+EPOCHS_HEAD = 20        # frozen backbone, train classifier + projection layers only
+EPOCHS_FINE = 50        # unfreeze all, lower LR
+LR_HEAD = 1e-3
+LR_FINE = 1e-5
+PATIENCE_EARLY = 10
+PATIENCE_LR = 4
+# ──────────────────────────────────────────────────────────────────────────────
+
+
+def make_loaders(data_dir: str):
+    train_tf = transforms.Compose([
+        transforms.Resize((IMG_SIZE, IMG_SIZE)),
+        transforms.RandomHorizontalFlip(),
+        transforms.RandomRotation(15),
+        transforms.ColorJitter(brightness=0.15, contrast=0.15),
+        transforms.ToTensor(),
+        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
+    ])
+    val_tf = transforms.Compose([
+        transforms.Resize((IMG_SIZE, IMG_SIZE)),
+        transforms.ToTensor(),
+        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
+    ])
+
+    train_ds = datasets.ImageFolder(os.path.join(data_dir, "train"), transform=train_tf)
+    val_ds = datasets.ImageFolder(os.path.join(data_dir, "val"), transform=val_tf)
+
+    # Expect exactly two classes: benign=0, malignant=1
+    print(f"Class mapping: {train_ds.class_to_idx}")
+    assert set(train_ds.class_to_idx.keys()) == {"benign", "malignant"}, (
+        "Dataset must have exactly 'benign' and 'malignant' subdirs"
+    )
+
+    train_loader = DataLoader(train_ds, batch_size=BATCH_SIZE, shuffle=True, num_workers=4, pin_memory=True)
+    val_loader = DataLoader(val_ds, batch_size=BATCH_SIZE, shuffle=False, num_workers=4, pin_memory=True)
+    return train_loader, val_loader, train_ds.class_to_idx
+
+
+def _freeze_backbones(model: AdvancedBreastCancerModel) -> None:
+    for param in model.stream1.parameters():
+        param.requires_grad = False
+    for param in model.stream2.parameters():
+        param.requires_grad = False
+
+
+def _unfreeze_all(model: AdvancedBreastCancerModel) -> None:
+    for param in model.parameters():
+        param.requires_grad = True
+
+
+def run_epoch(model, loader, criterion, optimizer, device, training: bool):
+    model.train() if training else model.eval()
+    total_loss = 0.0
+    correct = 0
+    total = 0
+
+    ctx = torch.enable_grad() if training else torch.no_grad()
+    with ctx:
+        for images, labels in loader:
+            images = images.to(device)
+            # labels: 0=benign, 1=malignant → float for BCEWithLogitsLoss
+            targets = labels.float().to(device)
+
+            logits = model(images).squeeze(1)
+            loss = criterion(logits, targets)
+
+            if training:
+                optimizer.zero_grad()
+                loss.backward()
+                optimizer.step()
+
+            total_loss += loss.item() * images.size(0)
+            preds = (torch.sigmoid(logits) >= 0.40).long()
+            correct += (preds == labels.to(device)).sum().item()
+            total += images.size(0)
+
+    return total_loss / total, correct / total
+
+
+def train(data_dir: str, output_path: str) -> None:
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    print(f"Device: {device}")
+
+    train_loader, val_loader, _ = make_loaders(data_dir)
+
+    model = AdvancedBreastCancerModel().to(device)
+    criterion = nn.BCEWithLogitsLoss()
+
+    best_val_acc = 0.0
+    output_path = Path(output_path)
+    output_path.parent.mkdir(parents=True, exist_ok=True)
+
+    # ── Phase 1: train head only ───────────────────────────────────────────────
+    print("\n=== Phase 1: training classifier head (frozen backbones) ===")
+    _freeze_backbones(model)
+    optimizer = Adam(filter(lambda p: p.requires_grad, model.parameters()), lr=LR_HEAD)
+    scheduler = ReduceLROnPlateau(optimizer, factor=0.5, patience=PATIENCE_LR, min_lr=1e-7, verbose=True)
+    no_improve = 0
+
+    for epoch in range(1, EPOCHS_HEAD + 1):
+        tr_loss, tr_acc = run_epoch(model, train_loader, criterion, optimizer, device, training=True)
+        vl_loss, vl_acc = run_epoch(model, val_loader, criterion, optimizer, device, training=False)
+        scheduler.step(vl_loss)
+        print(f"[P1 {epoch:02d}/{EPOCHS_HEAD}] loss={tr_loss:.4f} acc={tr_acc:.3f} | val_loss={vl_loss:.4f} val_acc={vl_acc:.3f}")
+
+        if vl_acc > best_val_acc:
+            best_val_acc = vl_acc
+            torch.save(model.state_dict(), output_path)
+            print(f"  ✓ Saved (val_acc={best_val_acc:.3f})")
+            no_improve = 0
+        else:
+            no_improve += 1
+            if no_improve >= PATIENCE_EARLY:
+                print("  Early stopping (Phase 1)")
+                break
+
+    # ── Phase 2: fine-tune all layers ─────────────────────────────────────────
+    print("\n=== Phase 2: fine-tuning all layers ===")
+    _unfreeze_all(model)
+    optimizer = Adam(model.parameters(), lr=LR_FINE)
+    scheduler = ReduceLROnPlateau(optimizer, factor=0.5, patience=PATIENCE_LR, min_lr=1e-8, verbose=True)
+    no_improve = 0
+
+    for epoch in range(1, EPOCHS_FINE + 1):
+        tr_loss, tr_acc = run_epoch(model, train_loader, criterion, optimizer, device, training=True)
+        vl_loss, vl_acc = run_epoch(model, val_loader, criterion, optimizer, device, training=False)
+        scheduler.step(vl_loss)
+        print(f"[P2 {epoch:02d}/{EPOCHS_FINE}] loss={tr_loss:.4f} acc={tr_acc:.3f} | val_loss={vl_loss:.4f} val_acc={vl_acc:.3f}")
+
+        if vl_acc > best_val_acc:
+            best_val_acc = vl_acc
+            torch.save(model.state_dict(), output_path)
+            print(f"  ✓ Saved (val_acc={best_val_acc:.3f})")
+            no_improve = 0
+        else:
+            no_improve += 1
+            if no_improve >= PATIENCE_EARLY:
+                print("  Early stopping (Phase 2)")
+                break
+
+    print(f"\nDone. Best val_acc={best_val_acc:.3f}")
+    print(f"Weights → {output_path}")
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="Train SensiNet mammogram classifier")
+    parser.add_argument("--data", default="data", help="Root data dir (must contain train/ and val/)")
+    parser.add_argument("--output", default="weights/advanced_model_best.pth", help="Output weights path")
+    args = parser.parse_args()
+    train(args.data, args.output)