Sync training package

training/config.yaml

@@ -0,0 +1,49 @@
+# 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
+  batch_size: 32
+  num_workers: 4
+
+model:
+  architectures:
+    - name: efficientnet_b4
+      pretrained: true
+      unfreeze_last_n_blocks: 3
+    - name: efficientnetv2_s
+      pretrained: true
+      unfreeze_last_n_blocks: 3
+    - name: convnext_tiny
+      pretrained: true
+      unfreeze_last_n_blocks: 3
+  dropout_rate: 0.3
+  mc_dropout_samples: 30   # for uncertainty (CI95) estimation
+
+training:
+  phase1_epochs: 10
+  phase2_epochs: 30
+  phase1_lr: 1.0e-3
+  phase2_lr: 1.0e-5
+  weight_decay: 1.0e-4
+  early_stopping_patience: 5
+  loss_regression_weight: 0.7
+  loss_classification_weight: 0.3
+  random_seed: 42
+  n_folds: 5
+
+classes:
+  - normal
+  - mild
+  - moderate
+  - severe
+
+output:
+  model_dir: "outputs/models"
+  metrics_dir: "outputs/metrics"
+  figures_dir: "outputs/figures"
+
+wandb:
+  project: "anemiascan"
+  entity: null    # set to your W&B username if needed

training/cross_validation.py

@@ -0,0 +1,95 @@
+# 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
+import json
+import logging
+import pathlib
+
+import numpy as np
+from sklearn.model_selection import StratifiedKFold
+
+from training.train import load_config, train_model
+
+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.
+    """
+    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"]
+    )
+
+    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.mkdir(parents=True, exist_ok=True)
+        metrics = train_model(
+            model_name=model_name,
+            train_rows=train_rows,
+            val_rows=val_rows,
+            config=config,
+            output_dir=fold_out,
+            fold=fold,
+            run_name=f"{model_name}_cv_fold{fold}",
+        )
+        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)]
+        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"
+    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}"
+    )
+    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

@@ -0,0 +1,56 @@
+# 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)
+    return {"mae": mae, "rmse": rmse, "pearson_r": float(r), "pearson_p": float(p_val)}
+
+
+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
+        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
+
+    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)),
+        "sensitivity_per_class": per_class_sens,
+        "specificity_per_class": per_class_spec,
+        "confusion_matrix": cm,
+    }
+
+
+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))
+    return {
+        "mean_diff": mean_diff,
+        "std_diff": std_diff,
+        "loa_upper": mean_diff + 1.96 * std_diff,
+        "loa_lower": mean_diff - 1.96 * std_diff,
+    }

training/models/convnext_tiny.py

@@ -0,0 +1,38 @@
+# training/models/convnext_tiny.py
+"""ConvNeXt-Tiny dual-head model."""
+
+import timm
+import torch
+import torch.nn as nn
+
+
+class AnemiaModel(nn.Module):
+    def __init__(self, num_classes: int = 4, dropout_rate: float = 0.3, pretrained: bool = True):
+        super().__init__()
+        self.backbone = timm.create_model(
+            "convnext_tiny", 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)
+        )
+        self.classification_head = nn.Sequential(
+            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]:
+        f = self.backbone(x)
+        return self.regression_head(f), self.classification_head(f)
+
+    def freeze_backbone(self):
+        for p in self.backbone.parameters():
+            p.requires_grad = False
+
+    def unfreeze_last_n_blocks(self, n: int = 3):
+        stages = list(self.backbone.stages)
+        for stage in stages[-n:]:
+            for p in stage.parameters():
+                p.requires_grad = True

training/models/efficientnet_b4.py

@@ -0,0 +1,64 @@
+# 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,
+    ):
+        super().__init__()
+        self.backbone = timm.create_model(
+            "efficientnet_b4",
+            pretrained=pretrained,
+            num_classes=0,  # remove classifier head
+            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),
+        )
+        self.classification_head = nn.Sequential(
+            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
+
+    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 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():
+            param.requires_grad = True

training/models/efficientnetv2_s.py

@@ -0,0 +1,45 @@
+# training/models/efficientnetv2_s.py
+"""EfficientNetV2-S dual-head model."""
+
+import timm
+import torch
+import torch.nn as nn
+
+
+class AnemiaModel(nn.Module):
+    def __init__(self, num_classes: int = 4, dropout_rate: float = 0.3, pretrained: bool = True):
+        super().__init__()
+        self.backbone = timm.create_model(
+            "tf_efficientnetv2_s", 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)
+        )
+        self.classification_head = nn.Sequential(
+            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]:
+        f = self.backbone(x)
+        return self.regression_head(f), self.classification_head(f)
+
+    def freeze_backbone(self):
+        for p in self.backbone.parameters():
+            p.requires_grad = False
+
+    def unfreeze_last_n_blocks(self, n: int = 3):
+        blocks = list(self.backbone.blocks)
+        for block in 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
+        if hasattr(self.backbone, "bn2"):
+            for p in self.backbone.bn2.parameters():
+                p.requires_grad = True

training/models/ensemble.py

@@ -0,0 +1,110 @@
+# 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.
+"""
+
+import torch
+import torch.nn as nn
+from safetensors.torch import load_file
+
+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,
+    ):
+        super().__init__()
+        self.conj_model = AnemiaModel(pretrained=False)
+        self.nail_model = AnemiaModel(pretrained=False)
+        self.conj_model.load_state_dict(load_file(conj_ckpt))
+        self.nail_model.load_state_dict(load_file(nail_ckpt))
+        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]:
+        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
+        elif conj_img is not None:
+            hb, cls = 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
+
+    @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.
+        """
+        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."
+            )
+        import numpy as np
+        from torch.utils.data import DataLoader
+
+        from training.evaluation.metrics import compute_regression_metrics
+        from training.utils.dataset import AnemiaDataset
+
+        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        img_size = config["data"]["image_size"]
+
+        conj_model = AnemiaModel(pretrained=False).to(device)
+        nail_model = AnemiaModel(pretrained=False).to(device)
+        conj_model.load_state_dict(load_file(conj_ckpt))
+        nail_model.load_state_dict(load_file(nail_ckpt))
+        conj_model.eval()
+        nail_model.eval()
+
+        def get_preds(model, rows):
+            ds = AnemiaDataset(rows, image_size=img_size, augment=False)
+            loader = DataLoader(ds, batch_size=32)
+            preds, trues = [], []
+            with torch.no_grad():
+                for imgs, hb, _ in loader:
+                    hb_pred, _ = model(imgs.to(device))
+                    preds.extend(hb_pred.squeeze(1).cpu().numpy())
+                    trues.extend(hb.numpy())
+            return np.array(preds), np.array(trues)
+
+        preds_c, trues_c = get_preds(conj_model, val_rows_conj)
+        preds_n, _ = get_preds(nail_model, val_rows_nail)
+
+        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"]
+            if mae < best_mae:
+                best_mae, best_wc = mae, wc
+
+        return float(best_wc), float(1.0 - best_wc)

training/push_model_to_hf.py

@@ -0,0 +1,115 @@
+# training/push_model_to_hf.py
+"""Push trained model weights and metrics to HuggingFace Hub."""
+
+import json
+import logging
+import pathlib
+import shutil
+import tempfile
+
+from huggingface_hub import HfApi
+
+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."""
+    with tempfile.TemporaryDirectory() as tmpdir:
+        tmp = pathlib.Path(tmpdir)
+        shutil.copy(ckpt_path, tmp / "model.safetensors")
+        (tmp / "metrics.json").write_text(json.dumps(metrics, indent=2))
+        card = f"""---
+language: en
+license: cc-by-nc-4.0
+tags:
+  - medical-imaging
+  - anemia
+  - hemoglobin-estimation
+  - image-classification
+pipeline_tag: image-classification
+---
+
+# 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")} |
+
+## Disclaimer
+
+**Research tool only. Not a certified diagnostic device. All results require clinical confirmation.**
+"""
+        (tmp / "README.md").write_text(card)
+        api.upload_folder(
+            folder_path=str(tmp),
+            repo_id=repo_id,
+            repo_type="model",
+            commit_message=f"Add {model_name} {site} weights {version}",
+        )
+    log.info(f"Pushed to https://huggingface.co/{repo_id}")
+
+
+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,
+):
+    push_model(
+        conj_ckpt,
+        "hssling/anemia-efficientnet-b4-conjunctiva",
+        cv_summary_conj,
+        "efficientnet_b4",
+        "conjunctiva",
+        config,
+    )
+    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",
+        "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(),
+        path_in_repo="ensemble_config.json",
+        repo_id="hssling/anemia-ensemble",
+        repo_type="model",
+        commit_message="Add ensemble configuration",
+    )
+    log.info("Ensemble config pushed")

training/train.py

@@ -0,0 +1,229 @@
+# 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/
+"""
+
+import argparse
+import importlib
+import json
+import logging
+import pathlib
+
+import numpy as np
+import torch
+import torch.nn as nn
+import wandb
+import yaml
+from torch.utils.data import DataLoader
+
+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")
+log = logging.getLogger(__name__)
+
+
+def load_config(path: pathlib.Path) -> dict:
+    with open(path) as f:
+        return yaml.safe_load(f)
+
+
+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,
+    )
+
+
+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())
+    ce = nn.functional.cross_entropy(class_logits, class_true.long())
+    return w_reg * mse + w_cls * ce
+
+
+def run_epoch(model, loader, optimizer, device, training: bool, config: dict):
+    model.train() if training else model.eval()
+    total_loss, hb_preds, hb_trues, cls_preds, cls_trues = 0.0, [], [], [], []
+    w_reg = config["training"]["loss_regression_weight"]
+    w_cls = config["training"]["loss_classification_weight"]
+
+    ctx = torch.enable_grad() if training else torch.no_grad()
+    with ctx:
+        for imgs, hb, cls in loader:
+            imgs, hb, cls = imgs.to(device), hb.to(device), cls.to(device)
+            if training:
+                optimizer.zero_grad()
+            hb_pred, cls_logits = model(imgs)
+            loss = multitask_loss(hb_pred, hb, cls_logits, cls, w_reg, w_cls)
+            if training:
+                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())
+
+    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))
+    return {
+        "loss": total_loss / len(loader),
+        **reg_metrics,
+        "auc": cls_metrics["auc_macro"],
+        "f1": cls_metrics["f1_macro"],
+    }
+
+
+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."""
+    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,
+        batch_size=config["data"]["batch_size"],
+        shuffle=True,
+        num_workers=config["data"]["num_workers"],
+        pin_memory=True,
+    )
+    val_loader = DataLoader(
+        val_ds,
+        batch_size=config["data"]["batch_size"],
+        shuffle=False,
+        num_workers=config["data"]["num_workers"],
+        pin_memory=True,
+    )
+
+    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}",
+        config=config,
+        reinit=True,
+    )
+
+    # Phase 1: freeze backbone, train heads
+    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)
+        wandb.log(
+            {
+                "epoch": epoch,
+                **{f"train/{k}": v for k, v in train_m.items()},
+                **{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}"
+        )
+
+    # Phase 2: unfreeze last 3 blocks
+    arch_cfg = next(
+        (a for a in config["model"]["architectures"] if a["name"] == model_name),
+        config["model"]["architectures"][0],
+    )
+    model.unfreeze_last_n_blocks(arch_cfg["unfreeze_last_n_blocks"])
+    optimizer = torch.optim.AdamW(
+        filter(lambda p: p.requires_grad, model.parameters()),
+        lr=config["training"]["phase2_lr"],
+        weight_decay=config["training"]["weight_decay"],
+    )
+    scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(
+        optimizer, T_max=config["training"]["phase2_epochs"]
+    )
+
+    best_val_mae = float("inf")
+    patience_count = 0
+    best_metrics = {}
+    best_ckpt_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)
+        scheduler.step()
+        wandb.log(
+            {
+                "epoch": epoch + config["training"]["phase1_epochs"],
+                **{f"train/{k}": v for k, v in train_m.items()},
+                **{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}")
+
+        if val_m["mae"] < best_val_mae:
+            best_val_mae = val_m["mae"]
+            best_metrics = val_m
+            patience_count = 0
+            _save_safetensors(model, best_ckpt_path)
+        else:
+            patience_count += 1
+            if patience_count >= config["training"]["early_stopping_patience"]:
+                log.info(f"  Early stopping at epoch {epoch + 1}")
+                break
+
+    wandb_run.finish()
+    metrics_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

@@ -0,0 +1,26 @@
+# 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),
+        ]
+    )
+
+
+def get_val_transforms(image_size: int = 380) -> A.Compose:
+    return A.Compose(
+        [
+            A.Resize(image_size, image_size),
+        ]
+    )

training/utils/dataset.py

@@ -0,0 +1,55 @@
+# training/utils/dataset.py
+"""PyTorch Dataset for anemia screening images."""
+
+from typing import Any
+
+import numpy as np
+import torch
+from PIL import Image
+from torch.utils.data import Dataset
+
+from training.utils.augmentation import get_augmentation_pipeline, get_val_transforms
+
+CLASS_TO_IDX = {"normal": 0, "mild": 1, "moderate": 2, "severe": 3}
+IDX_TO_CLASS = {v: k for k, v in CLASS_TO_IDX.items()}
+
+
+class AnemiaDataset(Dataset):
+    """
+    Dataset wrapping a list of HuggingFace-style row dicts.
+
+    Each row must have:
+        image       : PIL Image
+        hb_value    : float
+        anemia_class: str (normal | mild | moderate | severe)
+    """
+
+    def __init__(self, rows: list[dict[str, Any]], image_size: int = 380, augment: bool = False):
+        self.rows = rows
+        self.image_size = image_size
+        self.transform = (
+            get_augmentation_pipeline(image_size) if augment else get_val_transforms(image_size)
+        )
+
+    def __len__(self) -> int:
+        return len(self.rows)
+
+    def __getitem__(self, idx: int) -> tuple[torch.Tensor, float, int]:
+        row = self.rows[idx]
+        img = row["image"]
+        if not isinstance(img, Image.Image):
+            img = Image.fromarray(np.array(img))
+        img = img.convert("RGB")
+        img_arr = np.array(img)
+
+        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
+
+        hb_val = float(row["hb_value"]) if row["hb_value"] is not None else 0.0
+        cls_idx = CLASS_TO_IDX.get(row.get("anemia_class", "normal"), 0)
+        return img_tensor, hb_val, cls_idx