merge

Author: ChristianHinge

src/baseline/Dockerfile

@@ -13,11 +13,11 @@ RUN pip install --no-cache-dir --break-system-packages -r requirements.txt
 # Copy code
 # -----------------------------
 
-COPY weights weights
+COPY outputs .
 COPY predict.py . 
 COPY unet.py .
 COPY dataset.py .
 # -----------------------------
 # Run inference
 # -----------------------------
-ENTRYPOINT ["python", "predict.py", "/data/features", "/data/output/ct.nii.gz"]
+ENTRYPOINT ["python", "predict.py", "--features_dir", "/data/features", "--output_ct", "/data/output/ct.nii.gz"]

src/baseline/README.md

@@ -20,7 +20,7 @@ Run with Docker:
 
 ```bash
 docker run --rm \
-  --memory 120g \
+  --memory 128g \
   -v /path/to/sub-000/features:/data/features:ro \
   -v /path/to/output:/data/output \
   ghcr.io/bic-mac-challenge/baseline
@@ -33,6 +33,9 @@ pip install -r requirements.txt
 python predict.py --features_dir /path/to/sub-000/features --output_ct pseudo-ct.nii.gz
 ```
 
+Weights are loaded from `outputs/checkpoints/best_model.pth`
+
+
 ## Training
 
 Edit `config.yaml` to set your data path, then:
@@ -43,8 +46,6 @@ 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/dataset.py

@@ -5,7 +5,7 @@
 def get_subject_features(features_dir):
     """Return all available model inputs for a subject: NAC-PET, topogram, combined and chunked
     DIXON MRI (in/out-phase), MRI face mask, and metadata (sex, age, height, weight).
-    These are everything a participant model may use to predict the pseudo-CT."""
+    These are everything a your model may use to predict the pseudo-CT."""
     paths = {
         "nacpet":                  os.path.join(features_dir, "nacpet.nii.gz"),
         "topogram":                os.path.join(features_dir, "topogram.nii.gz"),

src/baseline/predict.py

@@ -17,8 +17,7 @@
 from dataset import get_subject_features
 
 
-MODEL_PATH = Path(__file__).parent / "weights/best_model.pth"
-MODEL_PATH = Path("/sonne/hinge/Projects/challenge-codebase/src/baseline/outputs4/checkpoints/best_model.pth")
+MODEL_PATH = Path(__file__).parent / "outputs/best_model.pth"
 PATCH_SIZE = (192, 192, 192)
 SW_BATCH = 2
 OVERLAP = 0.5
@@ -51,6 +50,7 @@ def predict(features_dir, out_path):
             overlap=OVERLAP, mode="gaussian", progress=True,
         )
 
+    # Apply inverse of normalization to get HU
     pred_hu = pred.cpu().numpy()[0, 0] * 3000 - 1000
     affine = data["nacpet"].meta["affine"].numpy()
 

src/baseline/transforms.py

@@ -2,7 +2,7 @@
 
 
 #NOTE: The baseline uses only the NAC-PET as input
-# however, your model may use all images and metadata available
+# however, your model may use all the images and metadata available
 # under the /features folder
 
 def get_transforms(patch_size, num_samples=2):
@@ -44,25 +44,6 @@ def get_transforms(patch_size, num_samples=2):
                 random_size=False,
                 num_samples=num_samples
             ),
-
-            #RandGaussianNoised(keys=["input"], prob=0.5, mean=0.0, std=0.05),
-            #RandScaleIntensityd(keys=["input"], factors=0.1, prob=0.5),
-            #RandShiftIntensityd(keys=["input"], offsets=0.1, prob=0.5),
-            #RandGaussianSmoothd(
-            #    keys=["input"],
-            #    sigma_x=(0.5, 1.0), sigma_y=(0.5, 1.0), sigma_z=(0.5, 1.0),
-            #    prob=0.3,
-            #),
-
-            # RandAffined(
-            #    keys=["input", "ct", "prediction_mask"],
-            #    prob=0.5,
-            #    rotate_range=(0.087, 0.087, 0.087),  # ±5°
-            #    scale_range=(0.05, 0.05, 0.05),       # ±5%
-            #    mode=("bilinear", "bilinear", "nearest"),
-            #    padding_mode="border",
-            # ),
-
             EnsureTyped(keys=["input", "ct", "prediction_mask"]),
 
         ]