merge baseline

Author: ChristianHinge

.gitignore

@@ -2,10 +2,7 @@
 conversion/
 CLAUDE.md
 outputs*/
-outputs2/
-
 *.nii.gz
-v3/
 # Created by https://www.toptal.com/developers/gitignore/api/visualstudiocode,python
 # Edit at https://www.toptal.com/developers/gitignore?templates=visualstudiocode,python
 

src/baseline/config.yaml

@@ -1,9 +1,10 @@
-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
+data_dir: /data/bic-mac-data/train #CHANGE to your dataset path
 num_workers: 2
+output_dir: outputs
 
 epochs: 250
 learning_rate: 0.0003
 
 patch_size: [192,192,192]
+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,6 +18,7 @@
 
 
 MODEL_PATH = Path(__file__).parent / "weights/best_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
@@ -28,7 +29,7 @@ def predict(features_dir, out_path):
     transforms = Compose([
         LoadImaged(keys=["nacpet"]),
         EnsureChannelFirstd(keys=["nacpet"]),
-        NormalizeIntensityd(keys=["nacpet"], nonzero=True, 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

@@ -2,30 +2,23 @@
 import torch
 import yaml
 import matplotlib.pyplot as plt
-import torch.nn.functional as F
-
-from monai.data import Dataset, DataLoader, CacheDataset, PersistentDataset
+from monai.data import DataLoader, CacheDataset
 from tqdm import tqdm
 
 from dataset import get_dataset
-from transforms import get_train_transforms
+from transforms import get_transforms
 from unet import build_model
 
 
 torch.backends.cudnn.benchmark = True
 
-# -----------------------------
-# CONFIG
-# -----------------------------
+
 
 def load_config():
     with open("config.yaml") as f:
         return yaml.safe_load(f)
 
 
-# -----------------------------
-# TRAIN
-# -----------------------------
 
 def main():
 
@@ -35,30 +28,39 @@ def main():
 
     print("Using device:", device)
 
-    data = get_dataset(cfg["data_dir"])
+    all_data = get_dataset(cfg["data_dir"])
+    val_data, train_data = all_data[:2], all_data[2:]
 
-    transforms = get_train_transforms(
-        cfg["patch_size"],
-    )
+    train_transforms = get_transforms(cfg["patch_size"], cfg["train_num_samples"])
+    val_transforms = get_transforms(cfg["patch_size"], cfg["val_num_samples"])
 
-    print("Preparing dataset ...")
-    dataset = PersistentDataset(
-        data=data,
-        transform=transforms,
-        cache_dir=cfg["cache_dir"]
+    print("Caching train dataset...")
+    train_dataset = CacheDataset(
+        data=train_data,
+        transform=train_transforms,
+        cache_rate=1.0, # Change this to reduce memory footprint
+        num_workers=cfg["num_workers"],
+    )
+    loader = DataLoader(
+        train_dataset,
+        batch_size=cfg["batch_size"],
+        shuffle=True,
+        num_workers=cfg["num_workers"],
+        pin_memory=True,
+        persistent_workers=True
     )
 
-    print("Caching dataset...")
-    dataset = CacheDataset(
-        data=dataset,
+    print("Caching val dataset...")
+    val_dataset = CacheDataset(
+        data=val_data,
+        transform=val_transforms,
         cache_rate=1.0,
-        num_workers=8,
+        num_workers=cfg["num_workers"],
     )
-
-    loader = DataLoader(
-        dataset,
+    val_loader = DataLoader(
+        val_dataset,
         batch_size=cfg["batch_size"],
-        shuffle=True,
+        shuffle=False,
         num_workers=cfg["num_workers"],
         pin_memory=True,
         persistent_workers=True
@@ -77,17 +79,18 @@ def main():
         T_max=cfg["epochs"]
     )
 
+    scaler  = torch.amp.GradScaler("cuda")
     l1_loss = torch.nn.L1Loss()
 
-    scaler = torch.amp.GradScaler("cuda")
-
-    os.makedirs("outputs/checkpoints", exist_ok=True)
-    os.makedirs("outputs/logs", exist_ok=True)
-    os.makedirs("outputs/plots", exist_ok=True)
+    out = cfg["output_dir"]
+    os.makedirs(f"{out}/checkpoints", exist_ok=True)
+    os.makedirs(f"{out}/logs", exist_ok=True)
+    os.makedirs(f"{out}/plots", exist_ok=True)
 
-    best_loss = float("inf")
+    best_val_loss = float("inf")
 
-    loss_history = []
+    train_loss_history = []
+    val_loss_history = []
 
     print("Starting training...")
 
@@ -101,9 +104,10 @@ def main():
 
         for batch in pbar:
 
-            x = batch["input"].to(device)
-            y = batch["ct"].to(device)
-
+            x    = batch["input"].to(device)
+            y    = batch["ct"].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"):
@@ -120,41 +124,60 @@ def main():
 
             pbar.set_description(f"loss {loss.item():.4f}")
 
-        avg_loss = epoch_loss / len(loader)
+        avg_train_loss = epoch_loss / len(loader)
 
-        print("Epoch", epoch, "Loss", avg_loss)
+        scheduler.step()
 
-        loss_history.append(avg_loss)
+        # validation
+        model.eval()
+        val_loss = 0
+        with torch.no_grad():
+            for batch in val_loader:
+                x    = batch["input"].to(device)
+                y    = batch["ct"].to(device)
+                mask = batch["prediction_mask"].bool().to(device)
+                y[~mask] = 0 # don't bother trying to predict the bed 
 
-        scheduler.step()
+                with torch.amp.autocast("cuda"):
+                    pred = model(x)
+                    loss = l1_loss(pred, y)
+                val_loss += loss.item()
+        avg_val_loss = val_loss / len(val_loader)
+
+        print(f"Epoch {epoch}  train={avg_train_loss:.4f}  val={avg_val_loss:.4f}")
+
+        train_loss_history.append(avg_train_loss)
+        val_loss_history.append(avg_val_loss)
 
-        # best checkpoint
-        if avg_loss < best_loss:
+        # best checkpoint (by val)
+        if avg_val_loss < best_val_loss:
 
-            best_loss = avg_loss
+            best_val_loss = avg_val_loss
 
             torch.save(
                 model.state_dict(),
-                "outputs/checkpoints/best_model.pth"
+                f"{out}/checkpoints/best_model.pth"
             )
 
         # last checkpoint
         torch.save(
             model.state_dict(),
-            "outputs/checkpoints/last_model.pth"
+            f"{out}/checkpoints/last_model.pth"
         )
 
         # log
-        with open("outputs/logs/train_log.txt","a") as f:
-            f.write(f"{epoch},{avg_loss}\n")
+        with open(f"{out}/logs/train_log.txt", "a") as f:
+            f.write(f"{epoch},{avg_train_loss},{avg_val_loss}\n")
 
         # plot loss
         plt.figure()
-        plt.plot(loss_history)
+        plt.plot(train_loss_history, label="train")
+        plt.plot(val_loss_history, label="val")
         plt.xlabel("Epoch")
         plt.ylabel("Loss")
-        plt.title("Training Loss")
-        plt.savefig("outputs/plots/loss_curve.png")
+        plt.title("Train / Val Loss")
+        plt.legend()
+        plt.savefig(f"{out}/plots/loss_curve.png")
         plt.close()
 
 

src/baseline/transforms.py

@@ -5,19 +5,21 @@
 # however, your model may use all images and metadata available
 # under the /features folder
 
-def get_train_transforms(patch_size):
+def get_transforms(patch_size, num_samples=2):
 
     transforms = Compose(
         [
 
-            LoadImaged(keys=["nacpet", "ct"]),
+            LoadImaged(keys=["nacpet", "ct", "prediction_mask"]),
 
-            EnsureChannelFirstd(keys=["nacpet", "ct"]),
+            EnsureChannelFirstd(keys=["nacpet", "ct", "prediction_mask"]),
 
             NormalizeIntensityd(
                 keys=["nacpet"],
-                nonzero=True,
-                channel_wise=True
+                nonzero=False,
+                channel_wise=True,
+                subtrahend=[0]
+                
             ),
 
             ScaleIntensityRanged(
@@ -35,26 +37,33 @@ def get_train_transforms(patch_size):
                 name="input"
             ),
 
+            # Crop first so all random augmentations run on small patches
             RandSpatialCropSamplesd(
-                keys=["input","ct"],
+                keys=["input", "ct", "prediction_mask"],
                 roi_size=patch_size,
                 random_size=False,
-                num_samples=2
+                num_samples=num_samples
             ),
 
-            RandFlipd(
-                keys=["input","ct"],
-                spatial_axis=0,
-                prob=0.5
-            ),
+            #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,
+            #),
 
-            RandFlipd(
-                keys=["input","ct"],
-                spatial_axis=1,
-                prob=0.5
-            ),
+            # 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"]),
+            EnsureTyped(keys=["input", "ct", "prediction_mask"]),
 
         ]
     )