refactoring

Author: ChristianHinge

pyproject.toml

@@ -8,3 +8,6 @@ dependencies = [
     "nibabel>=5.3.3",
     "tqdm>=4.67.3",
 ]
+
+[tool.setuptools.packages.find]
+where = ["src"]

src/baseline/README.md

@@ -14,22 +14,19 @@ Despite MRI and topogram being available, this baseline uses only NAC-PET to kee
 
 ## Inference
 
-Build and run with Docker:
+The baseline model has already been pretrained
+
+Run with Docker:
 
 ```bash
-docker build -t bic-mac-baseline .
-timeout 5m docker run --rm --gpus 1 \
-  --memory 120g \
+docker run ghcr.io/bic-mac-challenge/baseline \
+ --memory 120g \
   -v /path/to/sub-XXX/features:/data/features:ro \
   -v /path/to/output:/data/output \
   bic-mac-baseline
 ```
 
-The predicted CT is written to `/data/output/ct.nii.gz`.
-
-Inference uses sliding window with 192³ patches, 50% overlap, and Gaussian blending.
-
-Without Docker:
+Or without Docker:
 
 ```bash
 pip install -r requirements.txt
@@ -46,3 +43,8 @@ python train.py
 
 Weights are saved to `outputs/checkpoints/best_model.pth` (lowest validation loss).
 
+Move new weights to `weights/best_model.pth` and build the image:
+
+```bash
+docker build -t my-baseline .
+```

src/baseline/config.yaml

@@ -1,7 +1,7 @@
-data_dir: /home/hinge_local/bic-mac-data/train
-
+data_dir: /data/bic-mac-data/train
+cache_dir: /data/bic-mac-cache 
 batch_size: 1 #batch_size is actually 2, since monai samples two patches per volume
-num_workers: 8
+num_workers: 2
 
 epochs: 250
 learning_rate: 0.0003

src/baseline/requirements.txt

@@ -2,4 +2,5 @@ monai
 nibabel
 numpy
 tqdm
-pyyaml
+pyyaml
+matplotlib

src/baseline/train.py

@@ -4,7 +4,7 @@
 import matplotlib.pyplot as plt
 import torch.nn.functional as F
 
-from monai.data import Dataset, DataLoader
+from monai.data import Dataset, DataLoader, CacheDataset, PersistentDataset
 from tqdm import tqdm
 
 from dataset import get_dataset
@@ -41,11 +41,18 @@ def main():
         cfg["patch_size"],
     )
 
-    print("Preparing dataset cache...")
-
-    dataset = Dataset(
+    print("Preparing dataset ...")
+    dataset = PersistentDataset(
         data=data,
         transform=transforms,
+        cache_dir=cfg["cache_dir"]
+    )
+
+    print("Caching dataset...")
+    dataset = CacheDataset(
+        data=dataset,
+        cache_rate=1.0,
+        num_workers=8,
     )
 
     loader = DataLoader(
@@ -82,6 +89,8 @@ def main():
 
     loss_history = []
 
+    print("Starting training...")
+
     for epoch in range(cfg["epochs"]):
 
         model.train()

src/evaluation/readme.md src/evaluation/README.mdsrc/evaluation/readme.md renamed to src/evaluation/README.md

@@ -9,4 +9,14 @@
 - Brain outlier robustness score
 - Organ bias
 - TAC bias
-"""
+"""
+
+from .eval_dataset import eval_dataset
+from .eval_case import evaluate_case
+from .metrics import (
+    compute_whole_body_suv_mae,
+    compute_brain_outlier_score,
+    compute_organ_bias_from_totalseg,
+    compute_tac_bias,
+    compute_whole_body_mu_mae,
+)

src/evaluation/__init__.py

@@ -15,7 +15,7 @@
 import os
 import numpy as np
 
-from metrics import (
+from .metrics import (
     compute_whole_body_suv_mae,
     compute_organ_bias_from_totalseg,
     compute_whole_body_mu_mae,

src/evaluation/eval_case.py

@@ -26,8 +26,8 @@
 
 import numpy as np
 
-from eval_case import evaluate_case
-from metrics import compute_brain_outlier_score
+from .eval_case import evaluate_case
+from .metrics import compute_brain_outlier_score
 
 
 def eval_dataset(dataset_path, pred_dir, subjects=None):

src/evaluation/eval_dataset.py

@@ -203,17 +203,23 @@ def reconstruction_pipeline(
     recon_template       = os.path.join(_recon_dir, "recon_OSEM_template.par")
     acf_forwardprojector = os.path.join(_recon_dir, "acf_forwardprojector.par")
 
-    reconstruction_pipeline(
-        output_dir=output_dir,
-        ct_path=args.ct,
-        ct_face_and_bed_path=ct_face_and_bed_path,
-        face_and_bed_mask_path=face_and_bed_mask_path,
-        add_sino_path=add_sino_path,
-        mult_sino_path=mult_sino_path,
-        prompts_sino_path=prompts_sino_path,
-        offset_json_path=offset_json_path,
-        recon_template=recon_template,
-        acf_forwardprojector=acf_forwardprojector,
-        overwrite=args.overwrite,
-        verbose=args.verbose,
-    )
+    log_path = os.path.join(output_dir, "intermediates", "recon.log")
+    try:
+        reconstruction_pipeline(
+            output_dir=output_dir,
+            ct_path=args.ct,
+            ct_face_and_bed_path=ct_face_and_bed_path,
+            face_and_bed_mask_path=face_and_bed_mask_path,
+            add_sino_path=add_sino_path,
+            mult_sino_path=mult_sino_path,
+            prompts_sino_path=prompts_sino_path,
+            offset_json_path=offset_json_path,
+            recon_template=recon_template,
+            acf_forwardprojector=acf_forwardprojector,
+            overwrite=args.overwrite,
+            verbose=args.verbose,
+        )
+    except Exception:
+        if not args.verbose:
+            print(f"\nReconstruction failed. Check {log_path} for details, or rerun with -v/--verbose for more output.", flush=True)
+        raise