final baseline

Author: ChristianHinge

src/baseline/config.yaml

@@ -1,12 +1,10 @@
-data_dir: /data/bic-mac-data/train
-cache_dir: /data/bic-mac-cache 
-batch_size: 1
+data_dir: /data/bic-mac-data/train #CHANGE to your dataset path
 num_workers: 2
-output_dir: outputs4
+output_dir: outputs
 
 epochs: 250
 learning_rate: 0.0003
 
 patch_size: [192,192,192]
-train_num_samples: 2
-val_num_samples: 8
+train_num_samples: 2 #Number of patches to sample per train subject
+val_num_samples: 8 #Number of patches to sample per val subject

src/baseline/predict.py

@@ -18,7 +18,7 @@
 
 
 MODEL_PATH = Path(__file__).parent / "weights/best_model.pth"
-MODEL_PATH = Path("/sonne/hinge/Projects/challenge-codebase/src/baseline/outputs2/checkpoints/last_model.pth")
+MODEL_PATH = Path("/sonne/hinge/Projects/challenge-codebase/src/baseline/outputs4/checkpoints/best_model.pth")
 PATCH_SIZE = (192, 192, 192)
 SW_BATCH = 2
 OVERLAP = 0.5
@@ -29,7 +29,7 @@ def predict(features_dir, out_path):
     transforms = Compose([
         LoadImaged(keys=["nacpet"]),
         EnsureChannelFirstd(keys=["nacpet"]),
-        NormalizeIntensityd(keys=["nacpet"], nonzero=False, subtrahend=0, channel_wise=True),
+        NormalizeIntensityd(keys=["nacpet"], nonzero=False, subtrahend=[0], channel_wise=True),
         ConcatItemsd(keys=["nacpet"], name="input"),
         EnsureTyped(keys=["input"]),
     ])

src/baseline/train.py

@@ -38,8 +38,8 @@ def main():
     train_dataset = CacheDataset(
         data=train_data,
         transform=train_transforms,
-        cache_rate=1.0,
-        num_workers=8,
+        cache_rate=1.0, # Change this to reduce memory footprint
+        num_workers=cfg["num_workers"],
     )
     loader = DataLoader(
         train_dataset,
@@ -55,7 +55,7 @@ def main():
         data=val_data,
         transform=val_transforms,
         cache_rate=1.0,
-        num_workers=4,
+        num_workers=cfg["num_workers"],
     )
     val_loader = DataLoader(
         val_dataset,
@@ -106,15 +106,15 @@ def main():
 
             x    = batch["input"].to(device)
             y    = batch["ct"].to(device)
-            mask = batch["prediction_mask"].to(device)
-
+            mask = batch["prediction_mask"].bool().to(device)
+            y[~mask] = 0 # don't bother trying to predict the bed 
             optimizer.zero_grad()
 
             with torch.amp.autocast("cuda"):
 
                 pred = model(x)
 
-                loss = l1_loss(pred[mask.bool()], y[mask.bool()])
+                loss = l1_loss(pred, y)
 
             scaler.scale(loss).backward()
             scaler.step(optimizer)
@@ -135,10 +135,12 @@ def main():
             for batch in val_loader:
                 x    = batch["input"].to(device)
                 y    = batch["ct"].to(device)
-                mask = batch["prediction_mask"].to(device)
+                mask = batch["prediction_mask"].bool().to(device)
+                y[~mask] = 0 # don't bother trying to predict the bed 
+
                 with torch.amp.autocast("cuda"):
                     pred = model(x)
-                    loss = l1_loss(pred[mask.bool()], y[mask.bool()])
+                    loss = l1_loss(pred, y)
                 val_loss += loss.item()
         avg_val_loss = val_loss / len(val_loader)
 

src/baseline/transforms.py

@@ -45,23 +45,23 @@ def get_transforms(patch_size, num_samples=2):
                 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,
-            ),
+            #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",
-            ),
+            # 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"]),