Sync all source code, docs, and configs

scripts/verify_data.py

@@ -0,0 +1,94 @@
+"""
+Data verification script: loads all images and annotations,
+validates counts, and saves visual overlays.
+
+Usage:
+    python scripts/verify_data.py --config config/config.yaml
+"""
+
+import argparse
+import sys
+from pathlib import Path
+
+import numpy as np
+import yaml
+
+sys.path.insert(0, str(Path(__file__).parent.parent))
+
+from src.preprocessing import discover_synapse_data, load_synapse
+from src.visualize import overlay_annotations
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--config", default="config/config.yaml")
+    args = parser.parse_args()
+
+    with open(args.config) as f:
+        cfg = yaml.safe_load(f)
+
+    print("=" * 60)
+    print("Immunogold Data Verification")
+    print("=" * 60)
+
+    records = discover_synapse_data(cfg["data"]["root"], cfg["data"]["synapse_ids"])
+
+    total_6nm = 0
+    total_12nm = 0
+    output_dir = Path("results/verification")
+
+    for record in records:
+        print(f"\n--- {record.synapse_id} ---")
+        print(f"  Image: {record.image_path.name}")
+        print(f"  Mask:  {record.mask_path.name if record.mask_path else 'NONE'}")
+        print(f"  6nm CSVs:  {[p.name for p in record.csv_6nm_paths]}")
+        print(f"  12nm CSVs: {[p.name for p in record.csv_12nm_paths]}")
+
+        data = load_synapse(record)
+        img = data["image"]
+        annots = data["annotations"]
+
+        n6 = len(annots["6nm"])
+        n12 = len(annots["12nm"])
+        total_6nm += n6
+        total_12nm += n12
+
+        print(f"  Image shape: {img.shape}")
+        print(f"  6nm particles: {n6}")
+        print(f"  12nm particles: {n12}")
+
+        if data["mask"] is not None:
+            # Check how many particles fall within mask
+            mask = data["mask"]
+            for cls, coords in annots.items():
+                if len(coords) == 0:
+                    continue
+                inside = sum(
+                    1 for x, y in coords
+                    if 0 <= int(y) < mask.shape[0] and
+                       0 <= int(x) < mask.shape[1] and
+                       mask[int(y), int(x)]
+                )
+                print(f"  {cls} in mask: {inside}/{len(coords)}")
+
+        # Save overlay
+        overlay_annotations(
+            img, annots,
+            title=f"{record.synapse_id}: {n6} 6nm, {n12} 12nm",
+            save_path=output_dir / f"{record.synapse_id}_annotations.png",
+        )
+
+    print(f"\n{'=' * 60}")
+    print(f"TOTAL: {total_6nm} 6nm + {total_12nm} 12nm = {total_6nm + total_12nm}")
+    print(f"Expected: 403 6nm + 50 12nm = 453")
+
+    if total_6nm + total_12nm >= 400:
+        print("PASS: Particle counts look reasonable")
+    else:
+        print("WARNING: Total count is lower than expected")
+
+    print(f"\nOverlays saved to: {output_dir}")
+
+
+if __name__ == "__main__":
+    main()

slurm/02_train_single_fold.sh

@@ -0,0 +1,40 @@
+#!/bin/bash
+#SBATCH --job-name=immunogold_train
+#SBATCH --partition=gpu          # TODO: adjust to your GPU partition
+#SBATCH --gres=gpu:1             # single A100 or V100
+#SBATCH --nodes=1
+#SBATCH --ntasks=1
+#SBATCH --cpus-per-task=8
+#SBATCH --mem=80G
+#SBATCH --time=12:00:00
+#SBATCH --array=0-49             # 10 folds x 5 seeds = 50 jobs
+#SBATCH --output=logs/train_%A_%a.out
+#SBATCH --error=logs/train_%A_%a.err
+
+set -euo pipefail
+mkdir -p logs
+
+eval "$(conda shell.bash hook)"
+conda activate immunogold
+
+# Map array task ID to fold and seed
+SYNAPSE_IDS=("S1" "S4" "S7" "S8" "S13" "S15" "S22" "S25" "S27" "S29")
+FOLD_IDX=$((SLURM_ARRAY_TASK_ID / 5))
+SEED_IDX=$((SLURM_ARRAY_TASK_ID % 5))
+SEED=$((SEED_IDX + 42))
+FOLD_NAME=${SYNAPSE_IDS[$FOLD_IDX]}
+
+echo "=== Training immunogold CenterNet ==="
+echo "Date: $(date)"
+echo "Array task: ${SLURM_ARRAY_TASK_ID}"
+echo "Fold: ${FOLD_NAME} (idx ${FOLD_IDX})"
+echo "Seed: ${SEED} (idx ${SEED_IDX})"
+echo "GPU: $(nvidia-smi --query-gpu=name --format=csv,noheader 2>/dev/null || echo 'N/A')"
+
+python train.py \
+    --fold "${FOLD_NAME}" \
+    --seed "${SEED}" \
+    --config config/config.yaml \
+    --device cuda:0
+
+echo "=== Training complete for fold=${FOLD_NAME} seed=${SEED} ==="

slurm/03_evaluate_ensemble.sh

@@ -0,0 +1,39 @@
+#!/bin/bash
+#SBATCH --job-name=immunogold_eval
+#SBATCH --partition=gpu          # TODO: adjust
+#SBATCH --gres=gpu:1             # single GPU per fold (model loaded sequentially)
+#SBATCH --nodes=1
+#SBATCH --ntasks=1
+#SBATCH --cpus-per-task=8
+#SBATCH --mem=64G
+#SBATCH --time=04:00:00
+#SBATCH --array=0-9              # one job per test fold
+#SBATCH --output=logs/eval_%A_%a.out
+#SBATCH --error=logs/eval_%A_%a.err
+
+set -euo pipefail
+mkdir -p logs
+
+eval "$(conda shell.bash hook)"
+conda activate immunogold
+
+SYNAPSE_IDS=("S1" "S4" "S7" "S8" "S13" "S15" "S22" "S25" "S27" "S29")
+FOLD_NAME=${SYNAPSE_IDS[$SLURM_ARRAY_TASK_ID]}
+
+echo "=== Ensemble evaluation ==="
+echo "Date: $(date)"
+echo "Test fold: ${FOLD_NAME}"
+echo "Array task: ${SLURM_ARRAY_TASK_ID}"
+echo "GPU: $(nvidia-smi --query-gpu=name --format=csv,noheader 2>/dev/null || echo 'N/A')"
+
+# Evaluate this single fold using the same script as local evaluation
+# evaluate_loocv.py handles loading all ensemble members per fold
+python evaluate_loocv.py \
+    --config config/config.yaml \
+    --ensemble-dir checkpoints \
+    --device cuda:0 \
+    --use-tta \
+    --fold "${FOLD_NAME}" \
+    --output "results/per_fold_predictions/${FOLD_NAME}_metrics.csv"
+
+echo "=== Evaluation complete for fold ${FOLD_NAME} ==="

slurm/04_full_pipeline.sh

@@ -0,0 +1,40 @@
+#!/bin/bash
+# Master submission script for the full immunogold detection pipeline.
+# Submits all stages with dependency chaining so they run sequentially.
+#
+# Usage:
+#   bash slurm/04_full_pipeline.sh
+
+set -euo pipefail
+
+echo "=== Immunogold Detection Pipeline ==="
+echo "Date: $(date)"
+echo "Submitting all pipeline stages..."
+
+# Create log directory
+mkdir -p logs
+
+# Stage 0: Environment setup (single job)
+JOB0=$(sbatch --parsable slurm/00_setup_env.sh)
+echo "Stage 0 (setup):     Job ${JOB0}"
+
+# Stage 1: Training — 50 parallel jobs (depends on setup)
+JOB1=$(sbatch --parsable --dependency=afterok:${JOB0} slurm/02_train_single_fold.sh)
+echo "Stage 1 (training):  Job ${JOB1} (array 0-49, depends on ${JOB0})"
+
+# Stage 2: Ensemble evaluation — 10 parallel jobs (depends on ALL training tasks)
+# afterok on an array job waits for ALL array tasks to complete successfully
+JOB2=$(sbatch --parsable --dependency=afterok:${JOB1} slurm/03_evaluate_ensemble.sh)
+echo "Stage 2 (evaluate):  Job ${JOB2} (array 0-9, depends on ${JOB1})"
+
+echo ""
+echo "Pipeline submitted successfully!"
+echo "Final results job: ${JOB2}"
+echo ""
+echo "Monitor with:"
+echo "  squeue -u \${USER}"
+echo "  sacct -j ${JOB0},${JOB1},${JOB2}"
+echo ""
+echo "Results will be in:"
+echo "  results/loocv_metrics.csv"
+echo "  results/per_fold_predictions/"

slurm/05_train_final.sh

@@ -0,0 +1,27 @@
+#!/bin/bash
+#SBATCH --job-name=immunogold_final
+#SBATCH --partition=shortq7-gpu
+#SBATCH --gres=gpu:1
+#SBATCH --nodes=1
+#SBATCH --ntasks=1
+#SBATCH --cpus-per-task=8
+#SBATCH --mem=32G
+#SBATCH --time=06:00:00
+#SBATCH --output=logs/train_final_%j.out
+#SBATCH --error=logs/train_final_%j.err
+
+set -euo pipefail
+mkdir -p logs
+
+module load miniconda3/24.3.0-gcc-13.2.0-rslr3to
+module load cuda/12.4.0-gcc-13.2.0-bxjolrw
+eval "$(conda shell.bash hook)"
+conda activate immunogold
+
+echo "=== Training final deployable model ==="
+echo "Date: $(date)"
+echo "GPU: $(nvidia-smi --query-gpu=name --format=csv,noheader 2>/dev/null)"
+
+python train_final.py --config config/config.yaml --device cuda:0
+
+echo "=== Final model training complete ==="

src/heatmap.py

@@ -0,0 +1,187 @@
+"""
+Ground truth heatmap generation and peak extraction for CenterNet.
+
+Generates Gaussian-splat heatmaps at stride-2 resolution with
+class-specific sigma values calibrated to bead size.
+"""
+
+import numpy as np
+import torch
+import torch.nn.functional as F
+from typing import Dict, List, Tuple, Optional
+
+# Class index mapping
+CLASS_IDX = {"6nm": 0, "12nm": 1}
+CLASS_NAMES = ["6nm", "12nm"]
+STRIDE = 2
+
+
+def generate_heatmap_gt(
+    coords_6nm: np.ndarray,
+    coords_12nm: np.ndarray,
+    image_h: int,
+    image_w: int,
+    sigmas: Optional[Dict[str, float]] = None,
+    stride: int = STRIDE,
+    confidence_6nm: Optional[np.ndarray] = None,
+    confidence_12nm: Optional[np.ndarray] = None,
+) -> Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
+    """
+    Generate CenterNet ground truth heatmaps and offset maps.
+
+    Args:
+        coords_6nm: (N, 2) array of (x, y) in ORIGINAL pixel space
+        coords_12nm: (M, 2) array of (x, y) in ORIGINAL pixel space
+        image_h: original image height
+        image_w: original image width
+        sigmas: per-class Gaussian sigma in feature space
+        stride: output stride (default 2)
+        confidence_6nm: optional per-particle confidence weights
+        confidence_12nm: optional per-particle confidence weights
+
+    Returns:
+        heatmap: (2, H//stride, W//stride) float32 in [0, 1]
+        offsets: (2, H//stride, W//stride) float32 sub-pixel offsets
+        offset_mask: (H//stride, W//stride) bool — True at particle centers
+        conf_map: (2, H//stride, W//stride) float32 confidence weights
+    """
+    if sigmas is None:
+        sigmas = {"6nm": 1.0, "12nm": 1.5}
+
+    h_feat = image_h // stride
+    w_feat = image_w // stride
+
+    heatmap = np.zeros((2, h_feat, w_feat), dtype=np.float32)
+    offsets = np.zeros((2, h_feat, w_feat), dtype=np.float32)
+    offset_mask = np.zeros((h_feat, w_feat), dtype=bool)
+    conf_map = np.ones((2, h_feat, w_feat), dtype=np.float32)
+
+    # Prepare coordinate lists with class labels and confidences
+    all_entries = []
+    if len(coords_6nm) > 0:
+        confs = confidence_6nm if confidence_6nm is not None else np.ones(len(coords_6nm))
+        for i, (x, y) in enumerate(coords_6nm):
+            all_entries.append((x, y, "6nm", confs[i]))
+    if len(coords_12nm) > 0:
+        confs = confidence_12nm if confidence_12nm is not None else np.ones(len(coords_12nm))
+        for i, (x, y) in enumerate(coords_12nm):
+            all_entries.append((x, y, "12nm", confs[i]))
+
+    for x, y, cls, conf in all_entries:
+        cidx = CLASS_IDX[cls]
+        sigma = sigmas[cls]
+
+        # Feature-space center (float)
+        cx_f = x / stride
+        cy_f = y / stride
+
+        # Integer grid center
+        cx_i = int(round(cx_f))
+        cy_i = int(round(cy_f))
+
+        # Sub-pixel offset
+        off_x = cx_f - cx_i
+        off_y = cy_f - cy_i
+
+        # Gaussian radius: truncate at 3 sigma
+        r = max(int(3 * sigma + 1), 2)
+
+        # Bounds-clipped grid
+        y0 = max(0, cy_i - r)
+        y1 = min(h_feat, cy_i + r + 1)
+        x0 = max(0, cx_i - r)
+        x1 = min(w_feat, cx_i + r + 1)
+
+        if y0 >= y1 or x0 >= x1:
+            continue
+
+        yy, xx = np.meshgrid(
+            np.arange(y0, y1),
+            np.arange(x0, x1),
+            indexing="ij",
+        )
+
+        # Gaussian centered at INTEGER center (not float)
+        # The integer center MUST be exactly 1.0 — the CornerNet focal loss
+        # uses pos_mask = (gt == 1.0) and treats everything else as negative.
+        # Centering the Gaussian at the float position produces peaks of 0.78-0.93
+        # which the loss sees as negatives → zero positive training signal.
+        gaussian = np.exp(
+            -((xx - cx_i) ** 2 + (yy - cy_i) ** 2) / (2 * sigma ** 2)
+        )
+
+        # Scale by confidence (for pseudo-label weighting)
+        gaussian = gaussian * conf
+
+        # Element-wise max (handles overlapping particles correctly)
+        heatmap[cidx, y0:y1, x0:x1] = np.maximum(
+            heatmap[cidx, y0:y1, x0:x1], gaussian
+        )
+
+        # Offset and confidence only at the integer center pixel
+        if 0 <= cy_i < h_feat and 0 <= cx_i < w_feat:
+            offsets[0, cy_i, cx_i] = off_x
+            offsets[1, cy_i, cx_i] = off_y
+            offset_mask[cy_i, cx_i] = True
+            conf_map[cidx, cy_i, cx_i] = conf
+
+    return heatmap, offsets, offset_mask, conf_map
+
+
+def extract_peaks(
+    heatmap: torch.Tensor,
+    offset_map: torch.Tensor,
+    stride: int = STRIDE,
+    conf_threshold: float = 0.3,
+    nms_kernel_sizes: Optional[Dict[str, int]] = None,
+) -> List[dict]:
+    """
+    Extract detections from predicted heatmap via max-pool NMS.
+
+    Args:
+        heatmap: (2, H/stride, W/stride) sigmoid-activated
+        offset_map: (2, H/stride, W/stride) raw offset predictions
+        stride: output stride
+        conf_threshold: minimum confidence to keep
+        nms_kernel_sizes: per-class NMS kernel sizes
+
+    Returns:
+        List of {'x': float, 'y': float, 'class': str, 'conf': float}
+    """
+    if nms_kernel_sizes is None:
+        nms_kernel_sizes = {"6nm": 3, "12nm": 5}
+
+    detections = []
+
+    for cls_idx, cls_name in enumerate(CLASS_NAMES):
+        hm_cls = heatmap[cls_idx].unsqueeze(0).unsqueeze(0)  # (1,1,H,W)
+        kernel = nms_kernel_sizes[cls_name]
+
+        # Max-pool NMS
+        hmax = F.max_pool2d(
+            hm_cls, kernel_size=kernel, stride=1, padding=kernel // 2
+        )
+        peaks = (hmax.squeeze() == heatmap[cls_idx]) & (
+            heatmap[cls_idx] > conf_threshold
+        )
+
+        ys, xs = torch.where(peaks)
+        for y_idx, x_idx in zip(ys, xs):
+            y_i = y_idx.item()
+            x_i = x_idx.item()
+            conf = heatmap[cls_idx, y_i, x_i].item()
+            dx = offset_map[0, y_i, x_i].item()
+            dy = offset_map[1, y_i, x_i].item()
+
+            # Back to input space with sub-pixel offset
+            det_x = (x_i + dx) * stride
+            det_y = (y_i + dy) * stride
+
+            detections.append({
+                "x": det_x,
+                "y": det_y,
+                "class": cls_name,
+                "conf": conf,
+            })
+
+    return detections

tests/test_evaluate.py

@@ -0,0 +1,110 @@
+"""Unit tests for evaluation matching and metrics."""
+
+import numpy as np
+import pytest
+
+from src.evaluate import match_detections_to_gt, compute_f1, compute_average_precision
+
+
+class TestComputeF1:
+    def test_perfect_score(self):
+        f1, p, r = compute_f1(10, 0, 0)
+        assert f1 == pytest.approx(1.0, abs=0.001)
+        assert p == pytest.approx(1.0, abs=0.001)
+        assert r == pytest.approx(1.0, abs=0.001)
+
+    def test_zero_detections(self):
+        f1, p, r = compute_f1(0, 0, 10)
+        assert f1 == pytest.approx(0.0, abs=0.01)
+        assert r == pytest.approx(0.0, abs=0.01)
+
+    def test_all_false_positives(self):
+        f1, p, r = compute_f1(0, 10, 0)
+        assert p == pytest.approx(0.0, abs=0.01)
+
+
+class TestMatchDetections:
+    def test_perfect_matching(self):
+        """Detections at exact GT locations should all match."""
+        gt_6nm = np.array([[100.0, 100.0], [200.0, 200.0]])
+        gt_12nm = np.array([[300.0, 300.0]])
+
+        dets = [
+            {"x": 100.0, "y": 100.0, "class": "6nm", "conf": 0.9},
+            {"x": 200.0, "y": 200.0, "class": "6nm", "conf": 0.8},
+            {"x": 300.0, "y": 300.0, "class": "12nm", "conf": 0.7},
+        ]
+
+        results = match_detections_to_gt(dets, gt_6nm, gt_12nm)
+        assert results["6nm"]["tp"] == 2
+        assert results["6nm"]["fp"] == 0
+        assert results["6nm"]["fn"] == 0
+        assert results["12nm"]["tp"] == 1
+
+    def test_wrong_class_no_match(self):
+        """Detection near GT but wrong class should not match."""
+        gt_6nm = np.array([[100.0, 100.0]])
+        gt_12nm = np.empty((0, 2))
+
+        dets = [
+            {"x": 100.0, "y": 100.0, "class": "12nm", "conf": 0.9},
+        ]
+
+        results = match_detections_to_gt(dets, gt_6nm, gt_12nm)
+        assert results["6nm"]["fn"] == 1  # missed
+        assert results["12nm"]["fp"] == 1  # false positive
+
+    def test_beyond_radius_no_match(self):
+        """Detection beyond match radius should not match."""
+        gt_6nm = np.array([[100.0, 100.0]])
+        gt_12nm = np.empty((0, 2))
+
+        dets = [
+            {"x": 120.0, "y": 100.0, "class": "6nm", "conf": 0.9},  # 20px away > 9px radius
+        ]
+
+        results = match_detections_to_gt(
+            dets, gt_6nm, gt_12nm, match_radii={"6nm": 9.0, "12nm": 15.0}
+        )
+        assert results["6nm"]["tp"] == 0
+        assert results["6nm"]["fp"] == 1
+        assert results["6nm"]["fn"] == 1
+
+    def test_within_radius_matches(self):
+        """Detection within match radius should match."""
+        gt_6nm = np.array([[100.0, 100.0]])
+        gt_12nm = np.empty((0, 2))
+
+        dets = [
+            {"x": 105.0, "y": 100.0, "class": "6nm", "conf": 0.9},  # 5px away < 9px
+        ]
+
+        results = match_detections_to_gt(
+            dets, gt_6nm, gt_12nm, match_radii={"6nm": 9.0, "12nm": 15.0}
+        )
+        assert results["6nm"]["tp"] == 1
+
+    def test_no_detections(self):
+        """No detections: all GT are false negatives."""
+        gt_6nm = np.array([[100.0, 100.0], [200.0, 200.0]])
+        results = match_detections_to_gt([], gt_6nm, np.empty((0, 2)))
+        assert results["6nm"]["fn"] == 2
+        assert results["6nm"]["f1"] == pytest.approx(0.0, abs=0.01)
+
+    def test_no_ground_truth(self):
+        """No GT: all detections are false positives."""
+        dets = [{"x": 100.0, "y": 100.0, "class": "6nm", "conf": 0.9}]
+        results = match_detections_to_gt(dets, np.empty((0, 2)), np.empty((0, 2)))
+        assert results["6nm"]["fp"] == 1
+
+
+class TestAveragePrecision:
+    def test_perfect_ap(self):
+        """All detections match in rank order → AP = 1.0."""
+        gt = np.array([[100.0, 100.0], [200.0, 200.0]])
+        dets = [
+            {"x": 100.0, "y": 100.0, "class": "6nm", "conf": 0.9},
+            {"x": 200.0, "y": 200.0, "class": "6nm", "conf": 0.8},
+        ]
+        ap = compute_average_precision(dets, gt, match_radius=9.0)
+        assert ap == pytest.approx(1.0, abs=0.01)

tests/test_heatmap.py

@@ -0,0 +1,132 @@
+"""Unit tests for heatmap GT generation and peak extraction."""
+
+import numpy as np
+import pytest
+import torch
+
+from src.heatmap import generate_heatmap_gt, extract_peaks
+
+
+class TestHeatmapGeneration:
+    def test_single_particle_peak(self):
+        """A single particle should produce a Gaussian peak at the correct location."""
+        coords_6nm = np.array([[100.0, 200.0]])
+        coords_12nm = np.empty((0, 2))
+
+        hm, off, mask, conf = generate_heatmap_gt(
+            coords_6nm, coords_12nm, 512, 512, stride=2,
+        )
+
+        assert hm.shape == (2, 256, 256)
+        assert off.shape == (2, 256, 256)
+        assert mask.shape == (256, 256)
+
+        # Peak should be at (50, 100) in stride-2 space
+        peak_y, peak_x = np.unravel_index(hm[0].argmax(), hm[0].shape)
+        assert abs(peak_x - 50) <= 1
+        assert abs(peak_y - 100) <= 1
+
+        # Peak value should be 1.0 (confidence=1.0 default)
+        assert hm[0].max() == pytest.approx(1.0, abs=0.01)
+
+        # 12nm channel should be empty
+        assert hm[1].max() == 0.0
+
+    def test_two_classes(self):
+        """Both classes should produce peaks in their respective channels."""
+        coords_6nm = np.array([[100.0, 100.0]])
+        coords_12nm = np.array([[200.0, 200.0]])
+
+        hm, _, _, _ = generate_heatmap_gt(
+            coords_6nm, coords_12nm, 512, 512, stride=2,
+        )
+
+        assert hm[0].max() > 0.9  # 6nm channel has peak
+        assert hm[1].max() > 0.9  # 12nm channel has peak
+
+    def test_offset_values(self):
+        """Offsets should encode sub-pixel correction."""
+        # Place particle at (101.5, 200.5) → stride-2 center at (50.75, 100.25)
+        # Integer center: (51, 100) → offset: (-0.25, 0.25)
+        coords_6nm = np.array([[101.5, 200.5]])
+        coords_12nm = np.empty((0, 2))
+
+        _, off, mask, _ = generate_heatmap_gt(
+            coords_6nm, coords_12nm, 512, 512, stride=2,
+        )
+
+        # Mask should have exactly one True pixel
+        assert mask.sum() == 1
+
+    def test_empty_annotations(self):
+        """Empty annotations should produce zero heatmap."""
+        hm, off, mask, conf = generate_heatmap_gt(
+            np.empty((0, 2)), np.empty((0, 2)), 512, 512,
+        )
+        assert hm.max() == 0.0
+        assert mask.sum() == 0
+
+    def test_confidence_weighting(self):
+        """Confidence < 1 should scale peak value."""
+        coords = np.array([[100.0, 100.0]])
+        confidences = np.array([0.5])
+
+        hm, _, _, _ = generate_heatmap_gt(
+            coords, np.empty((0, 2)), 512, 512,
+            confidence_6nm=confidences,
+        )
+
+        assert hm[0].max() == pytest.approx(0.5, abs=0.05)
+
+    def test_overlapping_particles_use_max(self):
+        """Overlapping Gaussians should use element-wise max, not sum."""
+        coords = np.array([[100.0, 100.0], [104.0, 100.0]])  # close together
+        hm, _, _, _ = generate_heatmap_gt(
+            coords, np.empty((0, 2)), 512, 512, stride=2,
+        )
+        # Max should be 1.0, not >1.0
+        assert hm[0].max() <= 1.0
+
+
+class TestPeakExtraction:
+    def test_single_peak(self):
+        """Extract a single peak from synthetic heatmap."""
+        heatmap = torch.zeros(2, 256, 256)
+        heatmap[0, 100, 50] = 0.9  # 6nm peak
+
+        offset_map = torch.zeros(2, 256, 256)
+        offset_map[0, 100, 50] = 0.3  # dx
+        offset_map[1, 100, 50] = 0.1  # dy
+
+        dets = extract_peaks(heatmap, offset_map, stride=2, conf_threshold=0.5)
+
+        assert len(dets) == 1
+        assert dets[0]["class"] == "6nm"
+        assert dets[0]["conf"] == pytest.approx(0.9, abs=0.01)
+        # x = (50 + 0.3) * 2 = 100.6
+        assert dets[0]["x"] == pytest.approx(100.6, abs=0.1)
+        # y = (100 + 0.1) * 2 = 200.2
+        assert dets[0]["y"] == pytest.approx(200.2, abs=0.1)
+
+    def test_nms_suppresses_neighbors(self):
+        """NMS should suppress weaker neighboring peaks."""
+        heatmap = torch.zeros(2, 256, 256)
+        heatmap[0, 100, 50] = 0.9  # strong
+        heatmap[0, 101, 50] = 0.7  # weaker neighbor (within NMS kernel)
+
+        dets = extract_peaks(
+            heatmap, torch.zeros(2, 256, 256),
+            stride=2, conf_threshold=0.5,
+            nms_kernel_sizes={"6nm": 5, "12nm": 5},
+        )
+
+        # Only the stronger peak should survive
+        assert len([d for d in dets if d["class"] == "6nm"]) == 1
+
+    def test_below_threshold_filtered(self):
+        """Peaks below threshold should not be extracted."""
+        heatmap = torch.zeros(2, 256, 256)
+        heatmap[0, 100, 50] = 0.2  # below 0.3 threshold
+
+        dets = extract_peaks(heatmap, torch.zeros(2, 256, 256), conf_threshold=0.3)
+        assert len(dets) == 0

tests/test_loss.py

@@ -0,0 +1,109 @@
+"""Unit tests for loss functions."""
+
+import pytest
+import torch
+
+from src.loss import cornernet_focal_loss, offset_loss, total_loss
+
+
+class TestCornerNetFocalLoss:
+    def test_perfect_prediction_zero_loss(self):
+        """Perfect predictions should produce near-zero loss."""
+        gt = torch.zeros(1, 2, 64, 64)
+        gt[0, 0, 32, 32] = 1.0  # one particle
+
+        # Near-perfect prediction
+        pred = torch.zeros(1, 2, 64, 64) + 1e-6
+        pred[0, 0, 32, 32] = 1.0 - 1e-6
+
+        loss = cornernet_focal_loss(pred, gt)
+        assert loss.item() < 0.1
+
+    def test_all_zeros_prediction_nonzero_loss(self):
+        """Predicting all zeros when particles exist should give positive loss."""
+        gt = torch.zeros(1, 2, 64, 64)
+        gt[0, 0, 32, 32] = 1.0
+
+        pred = torch.zeros(1, 2, 64, 64) + 1e-6
+        loss = cornernet_focal_loss(pred, gt)
+        assert loss.item() > 0
+
+    def test_high_false_positive_penalized(self):
+        """Predicting high confidence where GT is zero should be penalized."""
+        gt = torch.zeros(1, 2, 64, 64)
+        pred_low_fp = torch.zeros(1, 2, 64, 64) + 0.01
+        pred_high_fp = torch.zeros(1, 2, 64, 64) + 0.9
+
+        loss_low = cornernet_focal_loss(pred_low_fp, gt)
+        loss_high = cornernet_focal_loss(pred_high_fp, gt)
+
+        assert loss_high.item() > loss_low.item()
+
+    def test_near_peak_reduced_penalty(self):
+        """Pixels near GT peaks should have reduced negative penalty via beta term."""
+        gt = torch.zeros(1, 2, 64, 64)
+        gt[0, 0, 32, 32] = 1.0
+        gt[0, 0, 31, 32] = 0.8  # nearby pixel with Gaussian falloff
+
+        # Moderate prediction near peak should have low loss
+        pred = torch.zeros(1, 2, 64, 64) + 0.01
+        pred[0, 0, 31, 32] = 0.5
+
+        loss = cornernet_focal_loss(pred, gt)
+        # Should be a reasonable value, not extremely high
+        assert loss.item() < 10
+
+    def test_confidence_weighting(self):
+        """Confidence weights should scale the loss."""
+        gt = torch.zeros(1, 2, 64, 64)
+        gt[0, 0, 32, 32] = 1.0
+        pred = torch.zeros(1, 2, 64, 64) + 0.5
+
+        weights_full = torch.ones(1, 2, 64, 64)
+        weights_half = torch.ones(1, 2, 64, 64) * 0.5
+
+        loss_full = cornernet_focal_loss(pred, gt, conf_weights=weights_full)
+        loss_half = cornernet_focal_loss(pred, gt, conf_weights=weights_half)
+
+        # Half weights should produce lower loss
+        assert loss_half.item() < loss_full.item()
+
+
+class TestOffsetLoss:
+    def test_zero_when_no_particles(self):
+        """Offset loss should be zero when mask is empty."""
+        pred = torch.randn(1, 2, 64, 64)
+        gt = torch.zeros(1, 2, 64, 64)
+        mask = torch.zeros(1, 64, 64, dtype=torch.bool)
+
+        loss = offset_loss(pred, gt, mask)
+        assert loss.item() == 0.0
+
+    def test_nonzero_with_particles(self):
+        """Offset loss should be nonzero when predictions differ from GT."""
+        pred = torch.randn(1, 2, 64, 64)
+        gt = torch.zeros(1, 2, 64, 64)
+        mask = torch.zeros(1, 64, 64, dtype=torch.bool)
+        mask[0, 32, 32] = True
+
+        loss = offset_loss(pred, gt, mask)
+        assert loss.item() > 0
+
+
+class TestTotalLoss:
+    def test_returns_three_values(self):
+        """total_loss should return (total, hm_loss, off_loss)."""
+        hm_pred = torch.sigmoid(torch.randn(1, 2, 64, 64))
+        hm_gt = torch.zeros(1, 2, 64, 64)
+        off_pred = torch.randn(1, 2, 64, 64)
+        off_gt = torch.zeros(1, 2, 64, 64)
+        mask = torch.zeros(1, 64, 64, dtype=torch.bool)
+
+        total, hm_val, off_val = total_loss(
+            hm_pred, hm_gt, off_pred, off_gt, mask,
+        )
+
+        assert isinstance(total, torch.Tensor)
+        assert isinstance(hm_val, float)
+        assert isinstance(off_val, float)
+        assert total.requires_grad

tests/test_model.py

@@ -0,0 +1,112 @@
+"""Unit tests for model architecture."""
+
+import pytest
+import torch
+
+from src.model import ImmunogoldCenterNet, BiFPN
+
+
+class TestModelForwardPass:
+    def test_output_shapes(self):
+        """Verify output shapes match stride-2 specification."""
+        model = ImmunogoldCenterNet(pretrained_path=None)
+        x = torch.randn(1, 1, 512, 512)
+        hm, off = model(x)
+
+        assert hm.shape == (1, 2, 256, 256), f"Expected (1,2,256,256), got {hm.shape}"
+        assert off.shape == (1, 2, 256, 256), f"Expected (1,2,256,256), got {off.shape}"
+
+    def test_heatmap_sigmoid_range(self):
+        """Heatmap outputs should be in [0, 1] from sigmoid."""
+        model = ImmunogoldCenterNet(pretrained_path=None)
+        x = torch.randn(1, 1, 512, 512)
+        hm, _ = model(x)
+
+        assert hm.min() >= 0.0
+        assert hm.max() <= 1.0
+
+    def test_batch_dimension(self):
+        """Model should handle batch size > 1."""
+        model = ImmunogoldCenterNet(pretrained_path=None)
+        x = torch.randn(4, 1, 512, 512)
+        hm, off = model(x)
+
+        assert hm.shape[0] == 4
+        assert off.shape[0] == 4
+
+    def test_variable_input_size(self):
+        """Model should handle different input sizes (multiples of 32)."""
+        model = ImmunogoldCenterNet(pretrained_path=None)
+
+        for size in [256, 384, 512]:
+            x = torch.randn(1, 1, size, size)
+            hm, off = model(x)
+            assert hm.shape == (1, 2, size // 2, size // 2)
+
+    def test_parameter_count(self):
+        """Model should have approximately 25M parameters."""
+        model = ImmunogoldCenterNet(pretrained_path=None)
+        n_params = sum(p.numel() for p in model.parameters())
+        # ResNet-50 is ~25M, plus BiFPN and heads
+        assert 20_000_000 < n_params < 40_000_000
+
+
+class TestFreezeUnfreeze:
+    def test_freeze_encoder(self):
+        """Frozen encoder should have no gradients."""
+        model = ImmunogoldCenterNet(pretrained_path=None)
+        model.freeze_encoder()
+
+        for name, param in model.named_parameters():
+            if any(x in name for x in ["stem", "layer1", "layer2", "layer3", "layer4"]):
+                assert not param.requires_grad, f"{name} should be frozen"
+
+        # BiFPN and heads should still be trainable
+        for name, param in model.bifpn.named_parameters():
+            assert param.requires_grad, f"bifpn.{name} should be trainable"
+
+    def test_unfreeze_deep(self):
+        """Unfreezing deep layers should enable gradients for layer3/4."""
+        model = ImmunogoldCenterNet(pretrained_path=None)
+        model.freeze_encoder()
+        model.unfreeze_deep_layers()
+
+        for param in model.layer3.parameters():
+            assert param.requires_grad
+        for param in model.layer4.parameters():
+            assert param.requires_grad
+        # Stem and layer1/2 still frozen
+        for param in model.stem.parameters():
+            assert not param.requires_grad
+
+    def test_unfreeze_all(self):
+        """Unfreeze all should enable all gradients."""
+        model = ImmunogoldCenterNet(pretrained_path=None)
+        model.freeze_encoder()
+        model.unfreeze_all()
+
+        for param in model.parameters():
+            assert param.requires_grad
+
+
+class TestBiFPN:
+    def test_bifpn_output_shapes(self):
+        """BiFPN should output 4 feature maps at 128 channels."""
+        bifpn = BiFPN(
+            in_channels=[256, 512, 1024, 2048],
+            out_channels=128,
+            num_rounds=2,
+        )
+        features = [
+            torch.randn(1, 256, 128, 128),   # P2: stride 4
+            torch.randn(1, 512, 64, 64),      # P3: stride 8
+            torch.randn(1, 1024, 32, 32),     # P4: stride 16
+            torch.randn(1, 2048, 16, 16),     # P5: stride 32
+        ]
+
+        outputs = bifpn(features)
+        assert len(outputs) == 4
+        for i, out in enumerate(outputs):
+            assert out.shape[1] == 128, f"P{i+2} channels should be 128"
+            assert out.shape[2:] == features[i].shape[2:], \
+                f"P{i+2} spatial dims should match input"