Deploy inference service from GitHub Actions

Dockerfile

@@ -3,23 +3,18 @@ FROM python:3.11-slim
 
 WORKDIR /app
 
-# System deps for OpenCV
 RUN apt-get update && apt-get install -y --no-install-recommends \
     libglib2.0-0 libsm6 libxrender1 libxext6 curl \
     && rm -rf /var/lib/apt/lists/*
 
-# Install Python deps
 COPY inference/requirements.txt requirements.txt
 RUN pip install --no-cache-dir -r requirements.txt
 
-# Copy training package (needed for model classes)
 COPY training/ training/
 COPY inference/ inference/
 
-# Expose port
 EXPOSE 7860
 
-# Health check
 HEALTHCHECK --interval=30s --timeout=10s --retries=3 \
     CMD sh -c 'curl -f "http://localhost:${PORT:-7860}/health" || exit 1'
 

README.md

@@ -12,3 +12,7 @@ license: mit
 # AnemiaScan
 
 Research-only non-invasive anemia screening demo.
+
+Concept, design, build, training, deployment, testing by: Dr Siddalingaiah H S, Professor, Community Medicine, Shridevi Institute of Medical Sciences and Research Hospital, Tumkur, hssling@yahoo.com, 8941087719.
+
+ORCID: 0000-0002-4771-8285

inference/Dockerfile

@@ -3,23 +3,18 @@ FROM python:3.11-slim
 
 WORKDIR /app
 
-# System deps for OpenCV
 RUN apt-get update && apt-get install -y --no-install-recommends \
     libglib2.0-0 libsm6 libxrender1 libxext6 curl \
     && rm -rf /var/lib/apt/lists/*
 
-# Install Python deps
 COPY inference/requirements.txt requirements.txt
 RUN pip install --no-cache-dir -r requirements.txt
 
-# Copy training package (needed for model classes)
 COPY training/ training/
 COPY inference/ inference/
 
-# Expose port
 EXPOSE 7860
 
-# Health check
 HEALTHCHECK --interval=30s --timeout=10s --retries=3 \
     CMD sh -c 'curl -f "http://localhost:${PORT:-7860}/health" || exit 1'
 

inference/__init__.py

@@ -0,0 +1 @@
+# inference/__init__.py

inference/app.py

@@ -1,16 +1,12 @@
-# inference/app.py
-"""
-AnemiaScan Inference Server -- FastAPI + Gradio
-
-Routes:
-  GET  /health          -> {"status": "ok", "models_loaded": [...]}
-  POST /api/predict     -> JSON prediction result
-  GET  /demo            -> Gradio demo interface
-"""
+"""FastAPI inference service with a mounted Gradio demo."""
+from __future__ import annotations
 
+import base64
 import io
 import logging
 import os
+from contextlib import suppress
+from typing import Any
 
 import gradio as gr
 import uvicorn
@@ -20,218 +16,201 @@ from fastapi.responses import RedirectResponse
 from PIL import Image
 
 from inference.gradcam import generate_gradcam
-from inference.model_loader import load_model, preload_all_models
+from inference.model_loader import _MODEL_CACHE, load_model, preload_available_models
 from inference.predict import preprocess_image, run_full_prediction
 
 logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
 log = logging.getLogger(__name__)
 
-app = FastAPI(
+PORT = int(os.getenv("PORT", "7860"))
+W_CONJ = float(os.getenv("ENSEMBLE_W_CONJ", "1.0"))
+W_NAIL = float(os.getenv("ENSEMBLE_W_NAIL", "0.0"))
+
+base_app = FastAPI(
     title="AnemiaScan Inference API",
-    description="Non-invasive anemia screening from conjunctival and nail-bed images",
-    version="1.0.0",
+    description="Non-invasive anemia screening from conjunctiva and nail-bed images",
+    version="0.4.0",
 )
 
-app.add_middleware(
+base_app.add_middleware(
     CORSMiddleware,
-    allow_origins=["*"],  # Netlify frontend
+    allow_origins=["*"],
     allow_methods=["GET", "POST"],
     allow_headers=["*"],
 )
 
-W_CONJ = float(os.getenv("ENSEMBLE_W_CONJ", "0.5"))
-W_NAIL = float(os.getenv("ENSEMBLE_W_NAIL", "0.5"))
-
-SPACE_CSS = """
-:root {
-  --glass-bg: rgba(248, 243, 236, 0.7);
-  --glass-line: rgba(36, 26, 20, 0.1);
-}
-body, .gradio-container {
-  background:
-    radial-gradient(circle at top left, rgba(156, 47, 31, 0.18), transparent 24%),
-    radial-gradient(circle at bottom right, rgba(212, 169, 79, 0.16), transparent 22%),
-    linear-gradient(180deg, #f6efe3 0%, #ecdfcf 100%);
-  font-family: "Aptos", "Segoe UI Variable Text", "Segoe UI", sans-serif;
-}
-.gradio-container {
-  max-width: 1100px !important;
-}
-.gr-box, .block, .gr-panel, .gr-form {
-  border-radius: 24px !important;
-}
-.gradio-container .gr-panel,
-.gradio-container .gr-box,
-.gradio-container .block {
-  background: var(--glass-bg) !important;
-  backdrop-filter: blur(18px) saturate(135%);
-  border: 1px solid var(--glass-line) !important;
-  box-shadow: 0 22px 60px rgba(70, 44, 30, 0.12);
-}
-.gradio-container h1, .gradio-container h2, .gradio-container h3 {
-  letter-spacing: -0.02em;
-}
-footer { display: none !important; }
-@media (prefers-color-scheme: dark) {
-  :root {
-    --glass-bg: rgba(20, 29, 41, 0.72);
-    --glass-line: rgba(237, 241, 247, 0.08);
-  }
-  body, .gradio-container {
-    background:
-      radial-gradient(circle at top left, rgba(255, 122, 89, 0.14), transparent 24%),
-      radial-gradient(circle at bottom right, rgba(127, 81, 255, 0.12), transparent 20%),
-      linear-gradient(180deg, #0e141c 0%, #101922 100%);
-  }
-}
-"""
-
-_MAX_UPLOAD_BYTES = 10 * 1024 * 1024  # 10 MB
-_MAX_IMAGE_PIXELS = 4096 * 4096  # ~16 MP
-
-
-def _open_image(raw: bytes) -> Image.Image:
-    """Open image bytes with size guards to prevent OOM from malicious uploads."""
-    if len(raw) > _MAX_UPLOAD_BYTES:
-        raise HTTPException(status_code=413, detail="Image file too large (max 10 MB).")
-    img = Image.open(io.BytesIO(raw))
-    if img.width * img.height > _MAX_IMAGE_PIXELS:
-        raise HTTPException(status_code=400, detail="Image dimensions too large (max 4096x4096).")
-    return img.convert("RGB")
-
-
-@app.on_event("startup")
-async def startup():
-    log.info("Preloading models ...")
-    preload_all_models()
-    log.info("Models ready.")
-
-
-@app.get("/health")
-def health():
-    from inference.model_loader import _MODEL_CACHE
-
-    return {"status": "ok", "models_loaded": list(_MODEL_CACHE.keys())}
-
-
-@app.get("/")
-def root():
-    return RedirectResponse(url="/demo/")
-
-
-@app.post("/api/predict")
+
+def _read_upload_as_pil(upload: UploadFile | None) -> Image.Image | None:
+    if upload is None:
+        return None
+    return Image.open(io.BytesIO(upload.file.read())).convert("RGB")
+
+
+def _try_load_model(site: str):
+    with suppress(Exception):
+        return load_model(site)
+    return None
+
+
+def _select_inputs_for_mode(
+    mode: str,
+    conj_img: Image.Image | None,
+    nail_img: Image.Image | None,
+) -> tuple[Image.Image | None, Image.Image | None]:
+    normalized = mode.lower().strip()
+    if normalized == "conjunctiva":
+        return conj_img, None
+    if normalized == "nailbed":
+        return None, nail_img
+    return conj_img, nail_img
+
+
+def _augment_with_gradcam(
+    result: dict[str, Any],
+    conj_img: Image.Image | None,
+    nail_img: Image.Image | None,
+    conj_model,
+    nail_model,
+) -> dict[str, Any]:
+    primary_img = conj_img if conj_img is not None else nail_img
+    primary_model = conj_model if conj_img is not None else nail_model
+    if primary_img is None or primary_model is None:
+        result["gradcam_b64"] = None
+        return result
+
+    try:
+        gradcam_tensor = preprocess_image(primary_img)
+        result["gradcam_b64"] = generate_gradcam(primary_model, gradcam_tensor)
+    except Exception as exc:
+        log.warning("Grad-CAM generation failed: %s", exc)
+        result["gradcam_b64"] = None
+    return result
+
+
+@base_app.on_event("startup")
+async def startup_event() -> None:
+    log.info("Preloading available inference models")
+    preload_available_models()
+
+
+@base_app.get("/")
+def root() -> RedirectResponse:
+    return RedirectResponse(url="/demo")
+
+
+@base_app.get("/health")
+def health() -> dict[str, Any]:
+    return {
+        "status": "ok",
+        "models_loaded": sorted(_MODEL_CACHE.keys()),
+        "default_mode": "conjunctiva",
+    }
+
+
+@base_app.post("/api/predict")
 async def predict(
     conjunctiva_image: UploadFile | None = File(default=None),
     nailbed_image: UploadFile | None = File(default=None),
-    model_name: str = Form(default="ensemble"),
-    include_gradcam: str = Form(default="false"),
-):
-    if conjunctiva_image is None and nailbed_image is None:
-        raise HTTPException(status_code=400, detail="Provide at least one image.")
+    model: str = Form(default="ensemble"),
+) -> dict[str, Any]:
+    conj_img = _read_upload_as_pil(conjunctiva_image)
+    nail_img = _read_upload_as_pil(nailbed_image)
 
-    conj_pil = None
-    nail_pil = None
-
-    if conjunctiva_image is not None:
-        raw = await conjunctiva_image.read()
-        conj_pil = _open_image(raw)
+    if conj_img is None and nail_img is None:
+        raise HTTPException(status_code=400, detail="Provide at least one image.")
 
-    if nailbed_image is not None:
-        raw = await nailbed_image.read()
-        nail_pil = _open_image(raw)
+    conj_img, nail_img = _select_inputs_for_mode(model, conj_img, nail_img)
+    conj_model = _try_load_model("conjunctiva") if conj_img is not None else None
+    nail_model = _try_load_model("nailbed") if nail_img is not None else None
 
-    try:
-        conj_model = load_model("conjunctiva") if conj_pil else None
-        nail_model = load_model("nailbed") if nail_pil else None
-    except Exception as e:
-        raise HTTPException(status_code=503, detail=f"Model loading failed: {e}")
+    if conj_img is not None and conj_model is None and nail_img is None:
+        raise HTTPException(status_code=503, detail="Conjunctiva model is unavailable.")
+    if nail_img is not None and nail_model is None and conj_img is None:
+        raise HTTPException(status_code=503, detail="Nail-bed model is unavailable.")
 
     try:
         result = run_full_prediction(
-            conj_img=conj_pil,
-            nail_img=nail_pil,
+            conj_img=conj_img,
+            nail_img=nail_img,
             conj_model=conj_model,
             nail_model=nail_model,
             w_conj=W_CONJ,
             w_nail=W_NAIL,
         )
-    except ValueError as e:
-        raise HTTPException(status_code=400, detail=str(e))
-    except Exception:
-        log.exception("Prediction failed")
-        raise HTTPException(status_code=500, detail="Internal prediction error.")
-
-    if include_gradcam.lower() == "true":
-        try:
-            gc_img = conj_pil if conj_pil else nail_pil
-            gc_site = "conjunctiva" if conj_pil else "nailbed"
-            gc_model = load_model(gc_site)
-            gc_tensor = preprocess_image(gc_img)
-            result["gradcam_b64"] = generate_gradcam(gc_model, gc_tensor)
-        except Exception as e:
-            log.warning(f"Grad-CAM generation failed: {e}")
-            result["gradcam_b64"] = None
+    except ValueError as exc:
+        raise HTTPException(status_code=400, detail=str(exc)) from exc
 
-    return result
+    result["requested_mode"] = model
+    return _augment_with_gradcam(result, conj_img, nail_img, conj_model, nail_model)
 
 
-def gradio_predict(conj_img, nail_img):
-    """Gradio wrapper for the prediction endpoint."""
+def gradio_predict(conj_img, nail_img, mode):
     conj_pil = Image.fromarray(conj_img) if conj_img is not None else None
     nail_pil = Image.fromarray(nail_img) if nail_img is not None else None
 
-    if conj_pil is None and nail_pil is None:
-        return "Please upload at least one image.", {}
-
     try:
-        conj_model = load_model("conjunctiva") if conj_pil else None
-        nail_model = load_model("nailbed") if nail_pil else None
-
+        selected_conj, selected_nail = _select_inputs_for_mode(mode, conj_pil, nail_pil)
+        conj_model = _try_load_model("conjunctiva") if selected_conj is not None else None
+        nail_model = _try_load_model("nailbed") if selected_nail is not None else None
         result = run_full_prediction(
-            conj_img=conj_pil,
-            nail_img=nail_pil,
+            conj_img=selected_conj,
+            nail_img=selected_nail,
             conj_model=conj_model,
             nail_model=nail_model,
             w_conj=W_CONJ,
             w_nail=W_NAIL,
         )
-    except Exception as e:
-        return f"Prediction error: {e}", {}
+        result = _augment_with_gradcam(result, selected_conj, selected_nail, conj_model, nail_model)
+    except Exception as exc:
+        return f"Prediction failed: {exc}", {}, None
 
     summary = (
-        f"**Hb Estimate:** {result['hb_estimate']} g/dL "
-        f"(95% CI: {result['hb_ci_95'][0]}-{result['hb_ci_95'][1]})\n\n"
-        f"**Classification:** {result['classification'].replace('_', ' ').title()}\n\n"
-        f"Warning: {result['disclaimer']}"
+        f"**Hb Estimate:** {result['hb_estimate']} g/dL\n\n"
+        f"**95% CI:** {result['hb_ci_95'][0]} to {result['hb_ci_95'][1]}\n\n"
+        f"**Classification:** {result['classification'].title()}\n\n"
+        f"**Disclaimer:** {result['disclaimer']}"
     )
-    return summary, result["class_probabilities"]
+    gradcam_image = None
+    if result["gradcam_b64"] is not None:
+        gradcam_bytes = io.BytesIO()
+        gradcam_bytes.write(base64.b64decode(result["gradcam_b64"]))
+        gradcam_bytes.seek(0)
+        gradcam_image = Image.open(gradcam_bytes)
+
+    return summary, result["class_probabilities"], gradcam_image
 
 
 demo = gr.Interface(
     fn=gradio_predict,
     inputs=[
-        gr.Image(label="Conjunctiva Image (optional)", type="numpy"),
-        gr.Image(label="Nail-bed Image (optional)", type="numpy"),
+        gr.Image(label="Conjunctiva Image", type="numpy"),
+        gr.Image(label="Nail-bed Image", type="numpy"),
+        gr.Radio(
+            choices=["ensemble", "conjunctiva", "nailbed"],
+            value="conjunctiva",
+            label="Prediction mode",
+        ),
     ],
     outputs=[
         gr.Markdown(label="Result"),
         gr.Label(label="Class Probabilities", num_top_classes=4),
+        gr.Image(label="Grad-CAM", type="pil"),
     ],
-    title="AnemiaScan -- AI Anemia Screening",
+    title="AnemiaScan",
     description=(
-        "Upload a palpebral conjunctiva and/or nail-bed image to estimate hemoglobin "
-        "and classify anemia severity. **Research tool only -- not a medical device.**"
-    ),
-    examples=[],
-    theme=gr.themes.Soft(
-        primary_hue="rose",
-        secondary_hue="amber",
-        neutral_hue="stone",
+        "Upload conjunctiva and/or nail-bed images for research-only Hb estimation.\n\n"
+        "Model summary: EfficientNet-B4 dual-head model, ImageNet-pretrained and fine-tuned on 380x380 RGB images. "
+        "Inference includes MC-dropout uncertainty and Grad-CAM explanations. "
+        "Tracked evaluation metrics include MAE, RMSE, Pearson r, AUC, F1, sensitivity, specificity, and Bland-Altman analysis. "
+        "Current live deployment loads conjunctiva weights first; nail-bed support follows once model weights are uploaded.\n\n"
+        "<sub>Concept, design, build, training, deployment, testing by: Dr Siddalingaiah H S, "
+        "Professor, Community Medicine, Shridevi Institute of Medical Sciences and Research "
+        "Hospital, Tumkur, hssling@yahoo.com, 8941087719. ORCID: 0000-0002-4771-8285.</sub>"
     ),
-    css=SPACE_CSS,
 )
 
-app = gr.mount_gradio_app(app, demo, path="/demo")
+app = gr.mount_gradio_app(base_app, demo, path="/demo")
+
 
 if __name__ == "__main__":
-    uvicorn.run("inference.app:app", host="0.0.0.0", port=8000, reload=False)
+    uvicorn.run("inference.app:app", host="0.0.0.0", port=PORT, reload=False)

inference/gradcam.py

@@ -1,12 +1,8 @@
-# inference/gradcam.py
-"""
-Grad-CAM heatmap generation for inference explanations.
-Returns a base64-encoded PNG overlay for the API response.
-"""
+"""Grad-CAM helper that returns a base64 encoded overlay."""
+from __future__ import annotations
 
 import base64
 import io
-import logging
 
 import cv2
 import numpy as np
@@ -14,7 +10,15 @@ import torch
 import torch.nn.functional as F
 from PIL import Image
 
-log = logging.getLogger(__name__)
+
+def _get_last_conv_layer(model: torch.nn.Module) -> torch.nn.Module:
+    last_conv = None
+    for layer in model.modules():
+        if isinstance(layer, torch.nn.Conv2d):
+            last_conv = layer
+    if last_conv is None:
+        raise RuntimeError("No convolutional layer found for Grad-CAM.")
+    return last_conv
 
 
 def generate_gradcam(
@@ -22,89 +26,59 @@ def generate_gradcam(
     image_tensor: torch.Tensor,
     target_class: int | None = None,
 ) -> str:
-    """
-    Generate a Grad-CAM heatmap for the given image and model.
-
-    Args:
-        model: AnemiaModel instance
-        image_tensor: (1, 3, H, W) preprocessed tensor
-        target_class: class index to visualise (default: argmax of classification head)
-
-    Returns:
-        Base64-encoded PNG string of the heatmap overlay.
-    """
+    """Generate a PNG overlay encoded as base64."""
     model.eval()
     activations: dict[str, torch.Tensor] = {}
     gradients: dict[str, torch.Tensor] = {}
 
     target_layer = _get_last_conv_layer(model)
 
-    def forward_hook(_, __, output):
+    def forward_hook(_module, _inputs, output):
         activations["value"] = output.detach()
 
-    def backward_hook(_, __, grad_output):
+    def backward_hook(_module, _grad_input, grad_output):
         gradients["value"] = grad_output[0].detach()
 
-    fh = target_layer.register_forward_hook(forward_hook)
-    bh = target_layer.register_full_backward_hook(backward_hook)
+    forward_handle = target_layer.register_forward_hook(forward_hook)
+    backward_handle = target_layer.register_full_backward_hook(backward_hook)
 
     try:
-        image_tensor = image_tensor.requires_grad_(True)
-        hb_pred, cls_logits = model(image_tensor)
-
+        input_tensor = image_tensor.clone().detach().requires_grad_(True)
+        _, cls_logits = model(input_tensor)
         if target_class is None:
             target_class = int(cls_logits.argmax(dim=1).item())
 
-        score = cls_logits[0, target_class]
-        model.zero_grad()
+        score = cls_logits[:, target_class].sum()
+        model.zero_grad(set_to_none=True)
         score.backward()
 
-        act = activations["value"].squeeze()  # (C, H, W)
-        grad = gradients["value"].squeeze()  # (C, H, W)
-
-        weights = grad.mean(dim=(1, 2))  # (C,)
-        cam = (weights[:, None, None] * act).sum(dim=0)  # (H, W)
+        activation = activations["value"].squeeze(0)
+        gradient = gradients["value"].squeeze(0)
+        weights = gradient.mean(dim=(1, 2))
+        cam = (weights[:, None, None] * activation).sum(dim=0)
         cam = F.relu(cam)
         cam = cam - cam.min()
         cam = cam / (cam.max() + 1e-8)
         cam_np = cam.cpu().numpy()
 
-        h, w = image_tensor.shape[2], image_tensor.shape[3]
-        cam_resized = cv2.resize(cam_np, (w, h))
-        heatmap = cv2.applyColorMap(np.uint8(255 * cam_resized), cv2.COLORMAP_JET)
+        height, width = input_tensor.shape[2], input_tensor.shape[3]
+        heatmap = cv2.applyColorMap(
+            np.uint8(255 * cv2.resize(cam_np, (width, height))),
+            cv2.COLORMAP_JET,
+        )
         heatmap = cv2.cvtColor(heatmap, cv2.COLOR_BGR2RGB)
 
-        orig = image_tensor.squeeze().detach().cpu().numpy()
         mean = np.array([0.485, 0.456, 0.406])[:, None, None]
         std = np.array([0.229, 0.224, 0.225])[:, None, None]
-        orig = np.clip((orig * std + mean) * 255, 0, 255).astype(np.uint8).transpose(1, 2, 0)
-
-        overlay = cv2.addWeighted(orig, 0.6, heatmap, 0.4, 0)
-        pil_img = Image.fromarray(overlay)
-
-        buf = io.BytesIO()
-        pil_img.save(buf, format="PNG")
-        return base64.b64encode(buf.getvalue()).decode("utf-8")
-
+        original = input_tensor.squeeze(0).detach().cpu().numpy()
+        original = np.clip((original * std + mean) * 255.0, 0, 255).astype(np.uint8)
+        original = original.transpose(1, 2, 0)
+
+        overlay = cv2.addWeighted(original, 0.6, heatmap, 0.4, 0)
+        image = Image.fromarray(overlay)
+        buffer = io.BytesIO()
+        image.save(buffer, format="PNG")
+        return base64.b64encode(buffer.getvalue()).decode("utf-8")
     finally:
-        fh.remove()
-        bh.remove()
-
-
-def _get_last_conv_layer(model: torch.nn.Module) -> torch.nn.Module:
-    """Return the last convolutional layer of the backbone (EfficientNet-B4)."""
-    try:
-        blocks = list(model.backbone.blocks)
-        for block in reversed(blocks):
-            for layer in reversed(list(block.modules())):
-                if isinstance(layer, torch.nn.Conv2d):
-                    return layer
-    except AttributeError:
-        pass
-    last_conv = None
-    for layer in model.modules():
-        if isinstance(layer, torch.nn.Conv2d):
-            last_conv = layer
-    if last_conv is None:
-        raise RuntimeError("No Conv2d layer found in model")
-    return last_conv
+        forward_handle.remove()
+        backward_handle.remove()

inference/model_loader.py

@@ -1,9 +1,5 @@
 # inference/model_loader.py
-"""
-Download and cache model weights from HuggingFace Hub.
-Models are loaded once at startup and cached in memory.
-"""
-
+"""Download and cache model weights from HuggingFace Hub."""
 import logging
 import os
 
@@ -15,42 +11,40 @@ from training.models.efficientnet_b4 import AnemiaModel
 
 log = logging.getLogger(__name__)
 
-_MODEL_CACHE: dict[str, AnemiaModel] = {}
+_MODEL_CACHE: dict = {}
 
 HF_REPOS = {
     "conjunctiva": os.getenv("HF_CONJ_MODEL_REPO", "hssling/anemia-efficientnet-b4-conjunctiva"),
-    "nailbed": os.getenv("HF_NAIL_MODEL_REPO", "hssling/anemia-efficientnet-b4-nailbed"),
+    "nailbed":     os.getenv("HF_NAIL_MODEL_REPO", "hssling/anemia-efficientnet-b4-nailbed"),
 }
-HF_TOKEN = os.getenv("HF_TOKEN") or os.getenv("HUGGING_FACE_HUB_TOKEN")
 
-DEVICE = torch.device("cpu")  # HF Spaces free tier is CPU-only
+DEVICE = torch.device("cpu")
 
 
 def load_model(site: str) -> AnemiaModel:
-    """Load and cache model for a given site ('conjunctiva' or 'nailbed')."""
     if site in _MODEL_CACHE:
         return _MODEL_CACHE[site]
 
     repo_id = HF_REPOS.get(site)
     if repo_id is None:
-        raise ValueError(f"Unknown site: {site!r}. Must be 'conjunctiva' or 'nailbed'.")
+        raise ValueError(f"Unknown site: {site!r}")
 
-    log.info(f"Downloading model weights from {repo_id} ...")
-    ckpt_path = hf_hub_download(repo_id=repo_id, filename="model.safetensors", token=HF_TOKEN)
+    log.info(f"Downloading {site} model from {repo_id} ...")
+    ckpt_path = hf_hub_download(repo_id=repo_id, filename="model.safetensors")
     model = AnemiaModel(pretrained=False)
     state_dict = load_file(ckpt_path, device="cpu")
     model.load_state_dict(state_dict)
     model.to(DEVICE)
     model.eval()
     _MODEL_CACHE[site] = model
-    log.info(f"Model loaded for site: {site}")
+    log.info(f"Model loaded: {site}")
     return model
 
 
-def preload_all_models():
-    """Eagerly load all models at startup to avoid cold-start delays."""
+def preload_available_models():
+    """Load conjunctiva model; skip nailbed if not yet on Hub."""
     for site in HF_REPOS:
         try:
             load_model(site)
         except Exception as e:
-            log.warning(f"Could not preload {site} model: {e}")
+            log.warning(f"Could not preload {site} model (skipping): {e}")

inference/predict.py

@@ -1,70 +1,80 @@
-# inference/predict.py
-"""
-Image preprocessing and MC-Dropout inference pipeline.
-"""
+"""Image preprocessing and MC-dropout inference helpers."""
+from __future__ import annotations
 
-import logging
 from typing import Any
 
 import numpy as np
 import torch
 from PIL import Image
 
-log = logging.getLogger(__name__)
-
 CLASS_NAMES = ["normal", "mild", "moderate", "severe"]
-_IMAGENET_MEAN = torch.tensor([0.485, 0.456, 0.406]).view(3, 1, 1)
-_IMAGENET_STD = torch.tensor([0.229, 0.224, 0.225]).view(3, 1, 1)
+_IMAGENET_MEAN = torch.tensor([0.485, 0.456, 0.406], dtype=torch.float32).view(3, 1, 1)
+_IMAGENET_STD = torch.tensor([0.229, 0.224, 0.225], dtype=torch.float32).view(3, 1, 1)
 
 
 def preprocess_image(img: Image.Image, image_size: int = 380) -> torch.Tensor:
-    """Convert PIL Image to normalised (1, 3, H, W) float tensor."""
-    img = img.convert("RGB").resize((image_size, image_size), Image.BICUBIC)
-    arr = np.array(img, dtype=np.float32) / 255.0
-    tensor = torch.from_numpy(arr).permute(2, 0, 1)  # (3, H, W)
+    """Convert a PIL image to a normalised BCHW tensor."""
+    resized = img.convert("RGB").resize((image_size, image_size), Image.BICUBIC)
+    array = np.asarray(resized, dtype=np.float32) / 255.0
+    tensor = torch.from_numpy(array).permute(2, 0, 1)
     tensor = (tensor - _IMAGENET_MEAN) / _IMAGENET_STD
-    return tensor.unsqueeze(0)  # (1, 3, H, W)
+    return tensor.unsqueeze(0)
+
+
+def _enable_dropout(module: torch.nn.Module) -> None:
+    for child in module.modules():
+        if isinstance(child, torch.nn.Dropout):
+            child.train()
 
 
 def mc_dropout_predict(
     model: torch.nn.Module,
     image_tensor: torch.Tensor,
-    n_samples: int = 30,
+    n_samples: int = 20,
 ) -> dict[str, Any]:
-    """
-    Run MC Dropout inference.
-    Activates dropout at inference time for n_samples forward passes.
-    Returns Hb estimate, 95% CI, class probabilities, and classification label.
-    """
-    model.train()  # activate dropout
-    hb_samples = []
-    cls_samples = []
-
-    try:
-        with torch.no_grad():
-            for _ in range(n_samples):
-                hb_pred, cls_logits = model(image_tensor)
-                hb_samples.append(hb_pred.item())
-                cls_samples.append(torch.softmax(cls_logits, dim=1).squeeze().numpy())
-    finally:
-        model.eval()  # always restore eval mode, even on exception
-
-    hb_arr = np.array(hb_samples)
-    cls_arr = np.array(cls_samples).mean(axis=0)  # (4,)
-
-    hb_mean = float(np.mean(hb_arr))
-    hb_lo = float(np.percentile(hb_arr, 2.5))
-    hb_hi = float(np.percentile(hb_arr, 97.5))
-    pred_class_idx = int(np.argmax(cls_arr))
+    """Run repeated stochastic forward passes and aggregate predictions."""
+    model.eval()
+    _enable_dropout(model)
+
+    hb_samples: list[float] = []
+    cls_samples: list[np.ndarray] = []
+
+    with torch.no_grad():
+        for _ in range(n_samples):
+            hb_pred, cls_logits = model(image_tensor)
+            hb_samples.append(float(hb_pred.squeeze().item()))
+            cls_samples.append(torch.softmax(cls_logits, dim=1).squeeze(0).cpu().numpy())
+
+    hb_array = np.asarray(hb_samples, dtype=np.float32)
+    cls_array = np.asarray(cls_samples, dtype=np.float32).mean(axis=0)
+    class_idx = int(np.argmax(cls_array))
 
     return {
-        "hb_estimate": round(hb_mean, 2),
-        "hb_ci_95": [round(hb_lo, 2), round(hb_hi, 2)],
-        "classification": CLASS_NAMES[pred_class_idx],
+        "hb_estimate": round(float(hb_array.mean()), 2),
+        "hb_ci_95": [
+            round(float(np.percentile(hb_array, 2.5)), 2),
+            round(float(np.percentile(hb_array, 97.5)), 2),
+        ],
+        "classification": CLASS_NAMES[class_idx],
         "class_probabilities": {
-            name: round(float(cls_arr[i]), 4) for i, name in enumerate(CLASS_NAMES)
+            class_name: round(float(cls_array[i]), 4)
+            for i, class_name in enumerate(CLASS_NAMES)
         },
-        "_hb_samples": hb_arr.tolist(),  # kept for ensemble CI computation; stripped before API response
+    }
+
+
+def _combine_weighted_probabilities(
+    first: dict[str, float],
+    second: dict[str, float],
+    first_weight: float,
+    second_weight: float,
+) -> dict[str, float]:
+    return {
+        class_name: round(
+            first_weight * first[class_name] + second_weight * second[class_name],
+            4,
+        )
+        for class_name in CLASS_NAMES
     }
 
 
@@ -76,58 +86,67 @@ def run_full_prediction(
     w_conj: float = 0.5,
     w_nail: float = 0.5,
     image_size: int = 380,
-    n_mc_samples: int = 30,
+    n_mc_samples: int = 20,
 ) -> dict[str, Any]:
-    """
-    Run prediction on available images, ensemble if both present.
-    Fills 'per_model' field with individual model results.
-    """
-    results = {}
+    """Predict from conjunctiva and/or nail-bed images with optional ensembling."""
+    per_model: dict[str, dict[str, Any]] = {}
 
     if conj_img is not None and conj_model is not None:
-        t = preprocess_image(conj_img, image_size)
-        results["conjunctiva"] = mc_dropout_predict(conj_model, t, n_mc_samples)
+        tensor = preprocess_image(conj_img, image_size=image_size)
+        per_model["conjunctiva"] = mc_dropout_predict(conj_model, tensor, n_samples=n_mc_samples)
 
     if nail_img is not None and nail_model is not None:
-        t = preprocess_image(nail_img, image_size)
-        results["nailbed"] = mc_dropout_predict(nail_model, t, n_mc_samples)
-
-    if not results:
-        raise ValueError("No model results — ensure at least one image and model are provided.")
-
-    # Ensemble
-    if "conjunctiva" in results and "nailbed" in results:
-        cls_probs = {
-            k: w_conj * results["conjunctiva"]["class_probabilities"][k]
-            + w_nail * results["nailbed"]["class_probabilities"][k]
-            for k in CLASS_NAMES
-        }
-        best_cls = max(cls_probs, key=cls_probs.get)
-        # Compute CI from combined weighted MC samples (statistically valid)
-        samples_c = np.array(results["conjunctiva"]["_hb_samples"])
-        samples_n = np.array(results["nailbed"]["_hb_samples"])
-        ensemble_samples = w_conj * samples_c + w_nail * samples_n
-        hb_mean = float(np.mean(ensemble_samples))
-        ci_lo = float(np.percentile(ensemble_samples, 2.5))
-        ci_hi = float(np.percentile(ensemble_samples, 97.5))
-        ensemble = {
-            "hb_estimate": round(hb_mean, 2),
-            "hb_ci_95": [round(ci_lo, 2), round(ci_hi, 2)],
-            "classification": best_cls,
-            "class_probabilities": {k: round(v, 4) for k, v in cls_probs.items()},
-        }
-    elif "conjunctiva" in results:
-        ensemble = results["conjunctiva"]
+        tensor = preprocess_image(nail_img, image_size=image_size)
+        per_model["nailbed"] = mc_dropout_predict(nail_model, tensor, n_samples=n_mc_samples)
+
+    if not per_model:
+        raise ValueError("No usable model/image pair was provided.")
+
+    if {"conjunctiva", "nailbed"} <= set(per_model):
+        hb_estimate = round(
+            w_conj * per_model["conjunctiva"]["hb_estimate"]
+            + w_nail * per_model["nailbed"]["hb_estimate"],
+            2,
+        )
+        hb_ci_95 = [
+            round(
+                w_conj * per_model["conjunctiva"]["hb_ci_95"][0]
+                + w_nail * per_model["nailbed"]["hb_ci_95"][0],
+                2,
+            ),
+            round(
+                w_conj * per_model["conjunctiva"]["hb_ci_95"][1]
+                + w_nail * per_model["nailbed"]["hb_ci_95"][1],
+                2,
+            ),
+        ]
+        class_probabilities = _combine_weighted_probabilities(
+            per_model["conjunctiva"]["class_probabilities"],
+            per_model["nailbed"]["class_probabilities"],
+            w_conj,
+            w_nail,
+        )
+        classification = max(class_probabilities, key=class_probabilities.get)
     else:
-        ensemble = results["nailbed"]
-
-    # Strip internal MC samples before returning — not part of the public API contract
-    for r in results.values():
-        r.pop("_hb_samples", None)
+        primary_result = next(iter(per_model.values()))
+        hb_estimate = primary_result["hb_estimate"]
+        hb_ci_95 = primary_result["hb_ci_95"]
+        classification = primary_result["classification"]
+        class_probabilities = primary_result["class_probabilities"]
+
+    warnings: list[str] = []
+    if conj_img is not None and conj_model is None:
+        warnings.append("Conjunctiva image ignored because the conjunctiva model is unavailable.")
+    if nail_img is not None and nail_model is None:
+        warnings.append("Nail-bed image ignored because the nail-bed model is unavailable.")
 
     return {
-        **ensemble,
-        "per_model": results,
-        "model_version": "v1.0.0",
+        "hb_estimate": hb_estimate,
+        "hb_ci_95": hb_ci_95,
+        "classification": classification,
+        "class_probabilities": class_probabilities,
+        "per_model": per_model,
+        "warnings": warnings,
+        "model_version": "v0.4.0",
         "disclaimer": "Research tool only. Not a certified diagnostic device. Clinical confirmation required.",
     }

inference/requirements.txt

@@ -11,3 +11,5 @@ opencv-python-headless>=4.9.0
 albumentations>=1.4.0
 numpy>=1.26.0
 python-multipart>=0.0.9
+scipy>=1.12.0
+scikit-learn>=1.4.0

training/config.yaml

@@ -1,9 +1,7 @@
 # training/config.yaml
-# All hyperparameters and paths. Override via CLI args in train.py.
-
 data:
   hf_dataset_repo: "hssling/anemia-conjunctiva-nailbed"
-  image_size: 380          # EfficientNet-B4 canonical input
+  image_size: 380
   batch_size: 32
   num_workers: 4
 
@@ -19,7 +17,7 @@ model:
       pretrained: true
       unfreeze_last_n_blocks: 3
   dropout_rate: 0.3
-  mc_dropout_samples: 30   # for uncertainty (CI95) estimation
+  mc_dropout_samples: 30
 
 training:
   phase1_epochs: 10
@@ -46,4 +44,4 @@ output:
 
 wandb:
   project: "anemiascan"
-  entity: null    # set to your W&B username if needed
+  entity: null

training/cross_validation.py

@@ -1,23 +1,11 @@
 # training/cross_validation.py
-"""
-5-fold stratified cross-validation runner.
-
-CV is used for metric estimation only.
-Final model is retrained on full train+val after CV.
-
-Usage:
-    python training/cross_validation.py \
-        --model efficientnet_b4 \
-        --config training/config.yaml \
-        --output-dir outputs/cv/
-"""
-
-import argparse
+"""5-fold stratified cross-validation runner."""
 import json
 import logging
 import pathlib
 
 import numpy as np
+import yaml
 from sklearn.model_selection import StratifiedKFold
 
 from training.train import load_config, train_model
@@ -26,33 +14,21 @@ log = logging.getLogger(__name__)
 logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
 
 
-def run_cross_validation(
-    rows: list,
-    model_name: str,
-    config: dict,
-    output_dir: pathlib.Path,
-) -> dict:
-    """
-    Run 5-fold stratified CV. Returns dict with mean +/- std of each metric.
-    """
+def run_cross_validation(rows, model_name, config, output_dir):
     n_folds = config["training"]["n_folds"]
-    fold_metrics = []
-
     strat_labels = [r["anemia_class"] for r in rows]
     skf = StratifiedKFold(
         n_splits=n_folds, shuffle=True, random_state=config["training"]["random_seed"]
     )
-
+    fold_metrics = []
     for fold, (train_idx, val_idx) in enumerate(skf.split(rows, strat_labels)):
         log.info(f"=== Fold {fold + 1}/{n_folds} ===")
-        train_rows = [rows[i] for i in train_idx]
-        val_rows = [rows[i] for i in val_idx]
-        fold_out = output_dir / f"fold_{fold}"
+        fold_out = pathlib.Path(output_dir) / f"fold_{fold}"
         fold_out.mkdir(parents=True, exist_ok=True)
         metrics = train_model(
             model_name=model_name,
-            train_rows=train_rows,
-            val_rows=val_rows,
+            train_rows=[rows[i] for i in train_idx],
+            val_rows=[rows[i] for i in val_idx],
             config=config,
             output_dir=fold_out,
             fold=fold,
@@ -60,36 +36,15 @@ def run_cross_validation(
         )
         fold_metrics.append(metrics)
 
-    all_keys = fold_metrics[0].keys()
-    summary = {}
-    for key in all_keys:
-        vals = [m[key] for m in fold_metrics if isinstance(m.get(key), int | float)]
+    summary = {"n_folds": n_folds, "model": model_name}
+    for key in fold_metrics[0].keys():
+        vals = [m[key] for m in fold_metrics if isinstance(m.get(key), (int, float))]
         if vals:
             summary[f"{key}_mean"] = float(np.mean(vals))
             summary[f"{key}_std"] = float(np.std(vals))
 
-    summary["n_folds"] = n_folds
-    summary["model"] = model_name
-    out_path = output_dir / f"{model_name}_cv_summary.json"
+    out_path = pathlib.Path(output_dir) / f"{model_name}_cv_summary.json"
     with open(out_path, "w") as f:
         json.dump(summary, f, indent=2)
-    log.info(
-        f"CV summary: MAE={summary.get('mae_mean', '?'):.3f} +/- {summary.get('mae_std', '?'):.3f}"
-    )
+    log.info(f"CV MAE={summary.get('mae_mean', '?'):.3f} ± {summary.get('mae_std', '?'):.3f}")
     return summary
-
-
-def main():
-    parser = argparse.ArgumentParser()
-    parser.add_argument("--model", required=True)
-    parser.add_argument("--config", default="training/config.yaml", type=pathlib.Path)
-    parser.add_argument("--output-dir", default="outputs/cv", type=pathlib.Path)
-    args = parser.parse_args()
-
-    load_config(args.config)
-    log.info(f"Cross-validation for {args.model}")
-    log.info("Load your dataset rows and call run_cross_validation(rows, ...)")
-
-
-if __name__ == "__main__":
-    main()

training/evaluation/metrics.py

@@ -1,13 +1,11 @@
 # training/evaluation/metrics.py
 """Evaluation metrics for hemoglobin regression and anemia classification."""
-
 import numpy as np
 from scipy import stats
 from sklearn.metrics import confusion_matrix, f1_score, roc_auc_score
 
 
 def compute_regression_metrics(y_true: np.ndarray, y_pred: np.ndarray) -> dict:
-    """MAE, RMSE, Pearson r for Hb regression."""
     mae = float(np.mean(np.abs(y_true - y_pred)))
     rmse = float(np.sqrt(np.mean((y_true - y_pred) ** 2)))
     r, p_val = stats.pearsonr(y_true, y_pred)
@@ -15,25 +13,21 @@ def compute_regression_metrics(y_true: np.ndarray, y_pred: np.ndarray) -> dict:
 
 
 def compute_classification_metrics(y_true: np.ndarray, y_pred_proba: np.ndarray) -> dict:
-    """AUC, F1, sensitivity, specificity, confusion matrix for 4-class anemia."""
     y_pred = np.argmax(y_pred_proba, axis=1)
     cm = confusion_matrix(y_true, y_pred, labels=[0, 1, 2, 3])
-
     per_class_sens = {}
     per_class_spec = {}
-    for cls in range(4):
-        tp = cm[cls, cls]
-        fn = cm[cls, :].sum() - tp
-        fp = cm[:, cls].sum() - tp
+    for c in range(4):
+        tp = cm[c, c]
+        fn = cm[c, :].sum() - tp
+        fp = cm[:, c].sum() - tp
         tn = cm.sum() - tp - fn - fp
-        per_class_sens[cls] = tp / (tp + fn) if (tp + fn) > 0 else 0.0
-        per_class_spec[cls] = tn / (tn + fp) if (tn + fp) > 0 else 0.0
-
+        per_class_sens[c] = tp / (tp + fn) if (tp + fn) > 0 else 0.0
+        per_class_spec[c] = tn / (tn + fp) if (tn + fp) > 0 else 0.0
     try:
         auc_macro = float(roc_auc_score(y_true, y_pred_proba, multi_class="ovr", average="macro"))
     except ValueError:
         auc_macro = float("nan")
-
     return {
         "auc_macro": auc_macro,
         "f1_macro": float(f1_score(y_true, y_pred, average="macro", zero_division=0)),
@@ -44,7 +38,6 @@ def compute_classification_metrics(y_true: np.ndarray, y_pred_proba: np.ndarray)
 
 
 def bland_altman_stats(y_true: np.ndarray, y_pred: np.ndarray) -> dict:
-    """Bland-Altman agreement statistics."""
     diff = y_true - y_pred
     mean_diff = float(np.mean(diff))
     std_diff = float(np.std(diff, ddof=1))

training/models/convnext_tiny.py

@@ -1,8 +1,6 @@
 # training/models/convnext_tiny.py
 """ConvNeXt-Tiny dual-head model."""
-
 import timm
-import torch
 import torch.nn as nn
 
 
@@ -17,13 +15,10 @@ class AnemiaModel(nn.Module):
             nn.Linear(feature_dim, 256), nn.ReLU(), nn.Dropout(dropout_rate), nn.Linear(256, 1)
         )
         self.classification_head = nn.Sequential(
-            nn.Linear(feature_dim, 256),
-            nn.ReLU(),
-            nn.Dropout(dropout_rate),
-            nn.Linear(256, num_classes),
+            nn.Linear(feature_dim, 256), nn.ReLU(), nn.Dropout(dropout_rate), nn.Linear(256, num_classes)
         )
 
-    def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+    def forward(self, x):
         f = self.backbone(x)
         return self.regression_head(f), self.classification_head(f)
 
@@ -32,7 +27,6 @@ class AnemiaModel(nn.Module):
             p.requires_grad = False
 
     def unfreeze_last_n_blocks(self, n: int = 3):
-        stages = list(self.backbone.stages)
-        for stage in stages[-n:]:
+        for stage in list(self.backbone.stages)[-n:]:
             for p in stage.parameters():
                 p.requires_grad = True

training/models/efficientnet_b4.py

@@ -1,63 +1,36 @@
 # training/models/efficientnet_b4.py
 """EfficientNet-B4 dual-head model for hemoglobin regression + anemia classification."""
-
 import timm
 import torch
 import torch.nn as nn
 
 
 class AnemiaModel(nn.Module):
-    """
-    EfficientNet-B4 backbone with dual prediction heads:
-      - Regression head: predicts Hb (g/dL)
-      - Classification head: predicts 4-class anemia severity
-    """
-
-    def __init__(
-        self,
-        num_classes: int = 4,
-        dropout_rate: float = 0.3,
-        pretrained: bool = True,
-    ):
+    def __init__(self, num_classes: int = 4, dropout_rate: float = 0.3, pretrained: bool = True):
         super().__init__()
         self.backbone = timm.create_model(
-            "efficientnet_b4",
-            pretrained=pretrained,
-            num_classes=0,  # remove classifier head
-            global_pool="avg",
+            "efficientnet_b4", pretrained=pretrained, num_classes=0, global_pool="avg"
         )
         feature_dim = self.backbone.num_features
-
         self.regression_head = nn.Sequential(
-            nn.Linear(feature_dim, 256),
-            nn.ReLU(),
-            nn.Dropout(dropout_rate),
-            nn.Linear(256, 1),
+            nn.Linear(feature_dim, 256), nn.ReLU(), nn.Dropout(dropout_rate), nn.Linear(256, 1)
         )
         self.classification_head = nn.Sequential(
-            nn.Linear(feature_dim, 256),
-            nn.ReLU(),
-            nn.Dropout(dropout_rate),
-            nn.Linear(256, num_classes),
+            nn.Linear(feature_dim, 256), nn.ReLU(), nn.Dropout(dropout_rate), nn.Linear(256, num_classes)
         )
 
     def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
         features = self.backbone(x)
-        hb_pred = self.regression_head(features)
-        class_logits = self.classification_head(features)
-        return hb_pred, class_logits
+        return self.regression_head(features), self.classification_head(features)
 
     def freeze_backbone(self):
         for param in self.backbone.parameters():
             param.requires_grad = False
 
     def unfreeze_last_n_blocks(self, n: int = 3):
-        """Unfreeze last n blocks of the backbone for fine-tuning."""
-        blocks = list(self.backbone.blocks)
-        for block in blocks[-n:]:
+        for block in list(self.backbone.blocks)[-n:]:
             for param in block.parameters():
                 param.requires_grad = True
-        # Always unfreeze the final conv + bn
         for param in self.backbone.conv_head.parameters():
             param.requires_grad = True
         for param in self.backbone.bn2.parameters():

training/models/efficientnetv2_s.py

@@ -1,8 +1,6 @@
 # training/models/efficientnetv2_s.py
 """EfficientNetV2-S dual-head model."""
-
 import timm
-import torch
 import torch.nn as nn
 
 
@@ -17,13 +15,10 @@ class AnemiaModel(nn.Module):
             nn.Linear(feature_dim, 256), nn.ReLU(), nn.Dropout(dropout_rate), nn.Linear(256, 1)
         )
         self.classification_head = nn.Sequential(
-            nn.Linear(feature_dim, 256),
-            nn.ReLU(),
-            nn.Dropout(dropout_rate),
-            nn.Linear(256, num_classes),
+            nn.Linear(feature_dim, 256), nn.ReLU(), nn.Dropout(dropout_rate), nn.Linear(256, num_classes)
         )
 
-    def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+    def forward(self, x):
         f = self.backbone(x)
         return self.regression_head(f), self.classification_head(f)
 
@@ -32,11 +27,9 @@ class AnemiaModel(nn.Module):
             p.requires_grad = False
 
     def unfreeze_last_n_blocks(self, n: int = 3):
-        blocks = list(self.backbone.blocks)
-        for block in blocks[-n:]:
+        for block in list(self.backbone.blocks)[-n:]:
             for p in block.parameters():
                 p.requires_grad = True
-        # Also unfreeze final conv + bn for consistent gradient flow with B4
         if hasattr(self.backbone, "conv_head"):
             for p in self.backbone.conv_head.parameters():
                 p.requires_grad = True

training/models/ensemble.py

@@ -1,12 +1,6 @@
 # training/models/ensemble.py
-"""
-Late-fusion dual-site ensemble.
-
-Loads a conjunctiva model and a nail-bed model.
-Combines predictions with learned weights (optimised on val set).
-Falls back gracefully if only one site image is provided.
-"""
-
+"""Late-fusion dual-site ensemble."""
+import numpy as np
 import torch
 import torch.nn as nn
 from safetensors.torch import load_file
@@ -15,13 +9,7 @@ from training.models.efficientnet_b4 import AnemiaModel
 
 
 class AnemiaEnsemble(nn.Module):
-    def __init__(
-        self,
-        conj_ckpt: str,
-        nail_ckpt: str,
-        w_conj: float = 0.5,
-        w_nail: float = 0.5,
-    ):
+    def __init__(self, conj_ckpt: str, nail_ckpt: str, w_conj: float = 0.5, w_nail: float = 0.5):
         super().__init__()
         self.conj_model = AnemiaModel(pretrained=False)
         self.nail_model = AnemiaModel(pretrained=False)
@@ -30,46 +18,25 @@ class AnemiaEnsemble(nn.Module):
         self.w_conj = w_conj
         self.w_nail = w_nail
 
-    def forward(
-        self,
-        conj_img: torch.Tensor | None = None,
-        nail_img: torch.Tensor | None = None,
-    ) -> tuple[torch.Tensor, torch.Tensor]:
+    def forward(self, conj_img=None, nail_img=None):
         if conj_img is not None and nail_img is not None:
             hb_c, cls_c = self.conj_model(conj_img)
             hb_n, cls_n = self.nail_model(nail_img)
-            hb = self.w_conj * hb_c + self.w_nail * hb_n
-            cls = self.w_conj * cls_c + self.w_nail * cls_n
+            return self.w_conj * hb_c + self.w_nail * hb_n, self.w_conj * cls_c + self.w_nail * cls_n
         elif conj_img is not None:
-            hb, cls = self.conj_model(conj_img)
+            return self.conj_model(conj_img)
         elif nail_img is not None:
-            hb, cls = self.nail_model(nail_img)
-        else:
-            raise ValueError("At least one image (conjunctiva or nail-bed) must be provided")
-        return hb, cls
+            return self.nail_model(nail_img)
+        raise ValueError("At least one image must be provided")
 
     @classmethod
-    def find_best_weights(
-        cls,
-        conj_ckpt: str,
-        nail_ckpt: str,
-        val_rows_conj: list,
-        val_rows_nail: list,
-        config: dict,
-    ) -> tuple[float, float]:
-        """Grid search over ensemble weights on validation set. Returns (w_conj, w_nail).
-
-        IMPORTANT: val_rows_conj and val_rows_nail must be from the same patients
-        in the same order. The ensemble MAE is evaluated against conjunctiva ground-truth
-        (trues_c). Only valid when both sets cover the same patient population.
-        """
+    def find_best_weights(cls, conj_ckpt, nail_ckpt, val_rows_conj, val_rows_nail, config):
+        """Grid search over ensemble weights. Returns (w_conj, w_nail)."""
         if len(val_rows_conj) != len(val_rows_nail):
             raise ValueError(
-                f"val_rows_conj ({len(val_rows_conj)}) and val_rows_nail "
-                f"({len(val_rows_nail)}) must have the same length for ensemble "
-                "weight grid search. Ensure both cover the same patients."
+                f"val_rows_conj ({len(val_rows_conj)}) and val_rows_nail ({len(val_rows_nail)}) "
+                "must have the same length for ensemble weight optimisation."
             )
-        import numpy as np
         from torch.utils.data import DataLoader
 
         from training.evaluation.metrics import compute_regression_metrics
@@ -101,10 +68,7 @@ class AnemiaEnsemble(nn.Module):
 
         best_mae, best_wc = float("inf"), 0.5
         for wc in np.arange(0.0, 1.05, 0.05):
-            wn = 1.0 - wc
-            ensemble_preds = wc * preds_c + wn * preds_n
-            mae = compute_regression_metrics(trues_c, ensemble_preds)["mae"]
+            mae = compute_regression_metrics(trues_c, wc * preds_c + (1.0 - wc) * preds_n)["mae"]
             if mae < best_mae:
-                best_mae, best_wc = mae, wc
-
-        return float(best_wc), float(1.0 - best_wc)
+                best_mae, best_wc = mae, float(wc)
+        return best_wc, 1.0 - best_wc

training/push_model_to_hf.py

@@ -1,6 +1,5 @@
 # training/push_model_to_hf.py
 """Push trained model weights and metrics to HuggingFace Hub."""
-
 import json
 import logging
 import pathlib
@@ -13,16 +12,7 @@ log = logging.getLogger(__name__)
 api = HfApi()
 
 
-def push_model(
-    ckpt_path: str,
-    repo_id: str,
-    metrics: dict,
-    model_name: str,
-    site: str,
-    config: dict,
-    version: str = "v1.0.0",
-):
-    """Push a single model checkpoint + metrics to HF Hub."""
+def push_model(ckpt_path, repo_id, metrics, model_name, site, config, version="v1.0.0"):
     with tempfile.TemporaryDirectory() as tmpdir:
         tmp = pathlib.Path(tmpdir)
         shutil.copy(ckpt_path, tmp / "model.safetensors")
@@ -34,33 +24,16 @@ tags:
   - medical-imaging
   - anemia
   - hemoglobin-estimation
-  - image-classification
 pipeline_tag: image-classification
 ---
 
-# AnemiaScan -- {model_name} ({site})
+# AnemiaScan — {model_name} ({site})
 
 **Task:** Non-invasive hemoglobin estimation + anemia severity classification from {site} images.
 
-**Architecture:** {model_name} (ImageNet pretrained, fine-tuned)
-
-**Input:** 380x380 RGB image of the palpebral {site}
-
-**Outputs:**
-- `hb_estimate` (float, g/dL)
-- `classification` (str: normal / mild / moderate / severe)
-
-## Performance (5-fold CV on public datasets)
-
-| Metric | Mean +/- Std |
-|--------|-----------|
-| MAE (g/dL) | {metrics.get("mae_mean", "TBD")} |
-| Pearson r | {metrics.get("pearson_r_mean", "TBD")} |
-| AUC (macro) | {metrics.get("auc_mean", "TBD")} |
+**Architecture:** {model_name} (ImageNet pretrained, fine-tuned on 380×380 RGB images)
 
-## Disclaimer
-
-**Research tool only. Not a certified diagnostic device. All results require clinical confirmation.**
+**Disclaimer:** Research tool only. Not a certified diagnostic device. Clinical confirmation required.
 """
         (tmp / "README.md").write_text(card)
         api.upload_folder(
@@ -73,13 +46,7 @@ pipeline_tag: image-classification
 
 
 def push_all_models(
-    conj_ckpt: str,
-    nail_ckpt: str,
-    cv_summary_conj: dict,
-    cv_summary_nail: dict,
-    w_conj: float,
-    w_nail: float,
-    config: dict,
+    conj_ckpt, nail_ckpt, cv_summary_conj, cv_summary_nail, w_conj, w_nail, config
 ):
     push_model(
         conj_ckpt,
@@ -89,21 +56,20 @@ def push_all_models(
         "conjunctiva",
         config,
     )
-    push_model(
-        nail_ckpt,
-        "hssling/anemia-efficientnet-b4-nailbed",
-        cv_summary_nail,
-        "efficientnet_b4",
-        "nailbed",
-        config,
-    )
+    if nail_ckpt is not None:
+        push_model(
+            nail_ckpt,
+            "hssling/anemia-efficientnet-b4-nailbed",
+            cv_summary_nail,
+            "efficientnet_b4",
+            "nailbed",
+            config,
+        )
     ensemble_meta = {
         "conj_model": "hssling/anemia-efficientnet-b4-conjunctiva",
-        "nail_model": "hssling/anemia-efficientnet-b4-nailbed",
+        "nail_model": "hssling/anemia-efficientnet-b4-nailbed" if nail_ckpt else None,
         "w_conj": w_conj,
         "w_nail": w_nail,
-        "mae_mean": w_conj * cv_summary_conj.get("mae_mean", 0)
-        + w_nail * cv_summary_nail.get("mae_mean", 0),
     }
     api.upload_file(
         path_or_fileobj=json.dumps(ensemble_meta, indent=2).encode(),

training/train.py

@@ -1,15 +1,5 @@
 # training/train.py
-"""
-Core training loop: two-phase training (head warmup -> backbone fine-tune).
-
-Usage:
-    python training/train.py \
-        --model efficientnet_b4 \
-        --site conjunctiva \
-        --config training/config.yaml \
-        --output-dir outputs/
-"""
-
+"""Two-phase training loop with W&B logging and safetensors checkpointing."""
 import argparse
 import importlib
 import json
@@ -23,10 +13,7 @@ import wandb
 import yaml
 from torch.utils.data import DataLoader
 
-from training.evaluation.metrics import (
-    compute_classification_metrics,
-    compute_regression_metrics,
-)
+from training.evaluation.metrics import compute_classification_metrics, compute_regression_metrics
 from training.utils.dataset import AnemiaDataset
 
 logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
@@ -40,21 +27,11 @@ def load_config(path: pathlib.Path) -> dict:
 
 def get_model(model_name: str, config: dict) -> nn.Module:
     mod = importlib.import_module(f"training.models.{model_name}")
-    return mod.AnemiaModel(
-        dropout_rate=config["model"]["dropout_rate"],
-        pretrained=True,
-    )
+    return mod.AnemiaModel(dropout_rate=config["model"]["dropout_rate"], pretrained=True)
 
 
-def multitask_loss(
-    hb_pred: torch.Tensor,
-    hb_true: torch.Tensor,
-    class_logits: torch.Tensor,
-    class_true: torch.Tensor,
-    w_reg: float = 0.7,
-    w_cls: float = 0.3,
-) -> torch.Tensor:
-    mse = nn.functional.mse_loss(hb_pred.squeeze(), hb_true.float())
+def multitask_loss(hb_pred, hb_true, class_logits, class_true, w_reg=0.7, w_cls=0.3):
+    mse = nn.functional.mse_loss(hb_pred.squeeze(1), hb_true.float())
     ce = nn.functional.cross_entropy(class_logits, class_true.long())
     return w_reg * mse + w_cls * ce
 
@@ -77,10 +54,10 @@ def run_epoch(model, loader, optimizer, device, training: bool, config: dict):
                 loss.backward()
                 optimizer.step()
             total_loss += loss.item()
-            hb_preds.extend(hb_pred.squeeze(1).cpu().numpy().tolist())
-            hb_trues.extend(hb.cpu().numpy().tolist())
-            cls_preds.extend(torch.softmax(cls_logits, dim=1).cpu().numpy().tolist())
-            cls_trues.extend(cls.cpu().numpy().tolist())
+            hb_preds.extend(hb_pred.squeeze(1).detach().cpu().numpy().tolist())
+            hb_trues.extend(hb.detach().cpu().numpy().tolist())
+            cls_preds.extend(torch.softmax(cls_logits, dim=1).detach().cpu().numpy().tolist())
+            cls_trues.extend(cls.detach().cpu().numpy().tolist())
 
     reg_metrics = compute_regression_metrics(np.array(hb_trues), np.array(hb_preds))
     cls_metrics = compute_classification_metrics(np.array(cls_trues), np.array(cls_preds))
@@ -92,31 +69,27 @@ def run_epoch(model, loader, optimizer, device, training: bool, config: dict):
     }
 
 
-def train_model(
-    model_name: str,
-    train_rows: list,
-    val_rows: list,
-    config: dict,
-    output_dir: pathlib.Path,
-    fold: int = 0,
-    run_name: str = "",
-) -> dict:
-    """Full two-phase training. Returns best val metrics dict."""
+def _save_safetensors(model: nn.Module, path: pathlib.Path):
+    from safetensors.torch import save_file
+
+    path.parent.mkdir(parents=True, exist_ok=True)
+    save_file({k: v.contiguous() for k, v in model.state_dict().items()}, str(path))
+
+
+def train_model(model_name, train_rows, val_rows, config, output_dir, fold=0, run_name=""):
     device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
     log.info(f"Training {model_name} fold={fold} on {device}")
 
     img_size = config["data"]["image_size"]
-    train_ds = AnemiaDataset(train_rows, image_size=img_size, augment=True)
-    val_ds = AnemiaDataset(val_rows, image_size=img_size, augment=False)
     train_loader = DataLoader(
-        train_ds,
+        AnemiaDataset(train_rows, image_size=img_size, augment=True),
         batch_size=config["data"]["batch_size"],
         shuffle=True,
         num_workers=config["data"]["num_workers"],
         pin_memory=True,
     )
     val_loader = DataLoader(
-        val_ds,
+        AnemiaDataset(val_rows, image_size=img_size, augment=False),
         batch_size=config["data"]["batch_size"],
         shuffle=False,
         num_workers=config["data"]["num_workers"],
@@ -124,7 +97,6 @@ def train_model(
     )
 
     model = get_model(model_name, config).to(device)
-
     wandb_run = wandb.init(
         project=config["wandb"]["project"],
         name=run_name or f"{model_name}_fold{fold}",
@@ -132,14 +104,13 @@ def train_model(
         reinit=True,
     )
 
-    # Phase 1: freeze backbone, train heads
+    # Phase 1: heads only
     model.freeze_backbone()
     optimizer = torch.optim.AdamW(
         filter(lambda p: p.requires_grad, model.parameters()),
         lr=config["training"]["phase1_lr"],
         weight_decay=config["training"]["weight_decay"],
     )
-    log.info("Phase 1: training heads only")
     for epoch in range(config["training"]["phase1_epochs"]):
         train_m = run_epoch(model, train_loader, optimizer, device, training=True, config=config)
         val_m = run_epoch(model, val_loader, optimizer, device, training=False, config=config)
@@ -150,11 +121,9 @@ def train_model(
                 **{f"val/{k}": v for k, v in val_m.items()},
             }
         )
-        log.info(
-            f"  Phase1 Ep{epoch + 1}: train_mae={train_m['mae']:.3f} val_mae={val_m['mae']:.3f}"
-        )
+        log.info(f"  P1 Ep{epoch+1}: train_mae={train_m['mae']:.3f} val_mae={val_m['mae']:.3f}")
 
-    # Phase 2: unfreeze last 3 blocks
+    # Phase 2: unfreeze last n blocks (look up by model_name)
     arch_cfg = next(
         (a for a in config["model"]["architectures"] if a["name"] == model_name),
         config["model"]["architectures"][0],
@@ -172,9 +141,8 @@ def train_model(
     best_val_mae = float("inf")
     patience_count = 0
     best_metrics = {}
-    best_ckpt_path = output_dir / f"{model_name}_fold{fold}_best.safetensors"
+    best_ckpt_path = pathlib.Path(output_dir) / f"{model_name}_fold{fold}_best.safetensors"
 
-    log.info("Phase 2: fine-tuning last 3 blocks")
     for epoch in range(config["training"]["phase2_epochs"]):
         train_m = run_epoch(model, train_loader, optimizer, device, training=True, config=config)
         val_m = run_epoch(model, val_loader, optimizer, device, training=False, config=config)
@@ -186,7 +154,7 @@ def train_model(
                 **{f"val/{k}": v for k, v in val_m.items()},
             }
         )
-        log.info(f"  Phase2 Ep{epoch + 1}: val_mae={val_m['mae']:.3f} val_auc={val_m['auc']:.3f}")
+        log.info(f"  P2 Ep{epoch+1}: val_mae={val_m['mae']:.3f} val_auc={val_m['auc']:.3f}")
 
         if val_m["mae"] < best_val_mae:
             best_val_mae = val_m["mae"]
@@ -196,34 +164,11 @@ def train_model(
         else:
             patience_count += 1
             if patience_count >= config["training"]["early_stopping_patience"]:
-                log.info(f"  Early stopping at epoch {epoch + 1}")
+                log.info(f"  Early stopping at epoch {epoch+1}")
                 break
 
     wandb_run.finish()
-    metrics_path = output_dir / f"{model_name}_fold{fold}_metrics.json"
+    metrics_path = pathlib.Path(output_dir) / f"{model_name}_fold{fold}_metrics.json"
     with open(metrics_path, "w") as f:
         json.dump(best_metrics, f, indent=2)
-    log.info(f"Best val MAE: {best_val_mae:.3f} -- saved to {best_ckpt_path}")
     return best_metrics
-
-
-def _save_safetensors(model: nn.Module, path: pathlib.Path):
-    from safetensors.torch import save_file
-
-    path.parent.mkdir(parents=True, exist_ok=True)
-    save_file({k: v.contiguous() for k, v in model.state_dict().items()}, str(path))
-
-
-def main():
-    parser = argparse.ArgumentParser()
-    parser.add_argument("--model", default="efficientnet_b4")
-    parser.add_argument("--config", default="training/config.yaml", type=pathlib.Path)
-    parser.add_argument("--output-dir", default="outputs/", type=pathlib.Path)
-    args = parser.parse_args()
-    load_config(args.config)
-    log.info(f"Config loaded: {args.config}")
-    log.info("Pass train_rows and val_rows to train_model() to start training.")
-
-
-if __name__ == "__main__":
-    main()

training/utils/augmentation.py

@@ -1,26 +1,21 @@
 # training/utils/augmentation.py
 """Albumentations pipelines for training and validation."""
-
 import albumentations as A
 
 
 def get_augmentation_pipeline(image_size: int = 380) -> A.Compose:
-    return A.Compose(
-        [
-            A.Resize(image_size, image_size),
-            A.HorizontalFlip(p=0.5),
-            A.Rotate(limit=15, p=0.7),
-            A.RandomBrightnessContrast(brightness_limit=0.2, contrast_limit=0.2, p=0.6),
-            A.HueSaturationValue(hue_shift_limit=10, sat_shift_limit=15, val_shift_limit=10, p=0.4),
-            A.GaussNoise(var_limit=(10, 50), p=0.2),
-            A.CoarseDropout(max_holes=4, max_height=32, max_width=32, p=0.3),
-        ]
-    )
+    return A.Compose([
+        A.Resize(image_size, image_size),
+        A.HorizontalFlip(p=0.5),
+        A.Rotate(limit=15, p=0.7),
+        A.RandomBrightnessContrast(brightness_limit=0.2, contrast_limit=0.2, p=0.6),
+        A.HueSaturationValue(hue_shift_limit=10, sat_shift_limit=15, val_shift_limit=10, p=0.4),
+        A.GaussNoise(std_range=(0.01, 0.05), p=0.2),
+        A.CoarseDropout(num_holes_range=(1, 4), hole_height_range=(1, 32), hole_width_range=(1, 32), p=0.3),
+    ])
 
 
 def get_val_transforms(image_size: int = 380) -> A.Compose:
-    return A.Compose(
-        [
-            A.Resize(image_size, image_size),
-        ]
-    )
+    return A.Compose([
+        A.Resize(image_size, image_size),
+    ])

training/utils/dataset.py

@@ -1,6 +1,5 @@
 # training/utils/dataset.py
 """PyTorch Dataset for anemia screening images."""
-
 from typing import Any
 
 import numpy as np
@@ -28,7 +27,8 @@ class AnemiaDataset(Dataset):
         self.rows = rows
         self.image_size = image_size
         self.transform = (
-            get_augmentation_pipeline(image_size) if augment else get_val_transforms(image_size)
+            get_augmentation_pipeline(image_size) if augment
+            else get_val_transforms(image_size)
         )
 
     def __len__(self) -> int:
@@ -45,7 +45,6 @@ class AnemiaDataset(Dataset):
         transformed = self.transform(image=img_arr)
         img_tensor = torch.from_numpy(transformed["image"]).permute(2, 0, 1).float() / 255.0
 
-        # Normalize with ImageNet stats
         mean = torch.tensor([0.485, 0.456, 0.406]).view(3, 1, 1)
         std = torch.tensor([0.229, 0.224, 0.225]).view(3, 1, 1)
         img_tensor = (img_tensor - mean) / std

training/utils/preprocessing.py

@@ -0,0 +1,5 @@
+# training/utils/preprocessing.py
+"""Image preprocessing utilities."""
+from training.utils.augmentation import get_val_transforms
+
+__all__ = ["get_val_transforms"]