Beat nnU-Net

benchmark/__main__.py

@@ -0,0 +1,37 @@
+from argparse import ArgumentParser
+from os.path import exists
+
+from benchmark.data import DataTest, SlidingWindowTest, RandomROIDatasetTest
+from benchmark.training import TrainingTest, ResizeTrainingTest, SlidingTrainingTest
+from mipcandy import auto_device, download_dataset, Frontend, NotionFrontend, WandBFrontend
+
+BENCHMARK_DATASET: str = "AbdomenCT-1K-ss1"
+
+if __name__ == "__main__":
+    tests = {
+        "SlidingWindow": SlidingWindowTest,
+        "RandomROI": RandomROIDatasetTest,
+        "Training": TrainingTest,
+        "ResizeTraining": ResizeTrainingTest,
+        "SlidingTraining": SlidingTrainingTest
+    }
+    parser = ArgumentParser(prog="MIP Candy Benchmark", description="MIP Candy Benchmark",
+                            epilog="GitHub: https://github.com/ProjectNeura/MIPCandy")
+    parser.add_argument("test", choices=tests.keys())
+    parser.add_argument("-i", "--input-folder")
+    parser.add_argument("-o", "--output-folder")
+    parser.add_argument("--num-epochs", type=int, default=100)
+    parser.add_argument("--device", default=None)
+    parser.add_argument("--front-end", choices=(None, "n", "w"), default=None)
+    args = parser.parse_args()
+    DataTest.dataset = BENCHMARK_DATASET
+    test = tests[args.test](
+        args.input_folder, args.output_folder, args.num_epochs, args.device if args.device else auto_device(), {
+            None: Frontend, "n": NotionFrontend, "w": WandBFrontend
+        }[args.front_end]
+    )
+    if not exists(f"{args.input_folder}/{BENCHMARK_DATASET}"):
+        download_dataset(f"nnunet_datasets/{BENCHMARK_DATASET}", f"{args.input_folder}/{BENCHMARK_DATASET}")
+    stat, err = test.run()
+    if not stat:
+        raise err

benchmark/data.py

@@ -0,0 +1,57 @@
+from time import time
+from typing import override, Literal
+
+from benchmark.prototype import UnitTest
+from mipcandy import NNUNetDataset, do_sliding_window, visualize3d, revert_sliding_window, JointTransform, inspect, \
+    RandomROIDataset
+
+
+class DataTest(UnitTest):
+    dataset: str = "AbdomenCT-1K-ss1"
+    transform: JointTransform | None = None
+
+    @override
+    def set_up(self) -> None:
+        self["dataset"] = NNUNetDataset(f"{self.input_folder}/{DataTest.dataset}", transform=self.transform,
+                                        device=self.device)
+        self["dataset"].preload(f"{self.input_folder}/{DataTest.dataset}/preloaded")
+
+
+class FoldedDataTest(DataTest):
+    fold: Literal[0, 1, 2, 3, 4, "all"] = 0
+
+    @override
+    def set_up(self) -> None:
+        super().set_up()
+        self["train_dataset"], self["val_dataset"] = self["dataset"].fold(fold=self.fold)
+
+
+class SlidingWindowTest(DataTest):
+    @override
+    def execute(self) -> None:
+        image, _ = self["dataset"][0]
+        print(image.shape)
+        visualize3d(image, title="raw")
+        t0 = time()
+        windows, layout, pad = do_sliding_window(image, (128, 128, 128))
+        print(f"took {time() - t0:.2f}s")
+        print(windows[0].shape, layout)
+        t0 = time()
+        recon = revert_sliding_window(windows, layout, pad)
+        print(f"took {time() - t0:.2f}s")
+        print(recon.shape)
+        visualize3d(recon, title="reconstructed")
+
+
+class RandomROIDatasetTest(DataTest):
+    @override
+    def execute(self) -> None:
+        annotations = inspect(self["dataset"])
+        dataset = RandomROIDataset(annotations)
+        print(len(dataset))
+        image, label = self["dataset"][0]
+        image_roi, label_roi = dataset[0]
+        visualize3d(image, title="image raw")
+        visualize3d(label, title="label raw", is_label=True)
+        visualize3d(image_roi, title="image roi")
+        visualize3d(label_roi, title="label roi", is_label=True)

benchmark/prototype.py

@@ -0,0 +1,41 @@
+from os import PathLike
+from typing import Any
+
+from mipcandy import Device, Frontend
+
+
+class UnitTest(object):
+    def __init__(self, input_folder: str | PathLike[str], output_folder: str | PathLike[str], num_epochs: int,
+                 device: Device, frontend: type[Frontend]) -> None:
+        self.input_folder: str = input_folder
+        self.output_folder: str = output_folder
+        self.num_epochs: int = num_epochs
+        self.device: Device = device
+        self.frontend: type[Frontend] = frontend
+
+    def set_up(self) -> None:
+        pass
+
+    def execute(self) -> None:
+        pass
+
+    def clean_up(self) -> None:
+        pass
+
+    def run(self) -> tuple[bool, Exception | None]:
+        try:
+            self.set_up()
+            self.execute()
+        except Exception as e:
+            try:
+                self.clean_up()
+            except Exception as e2:
+                print(f"Failed to clean up after exception: {e2}")
+            return False, e
+        return True, None
+
+    def __setitem__(self, key: str, value: Any) -> None:
+        setattr(self, "_x_" + key, value)
+
+    def __getitem__(self, item: str) -> Any:
+        return getattr(self, "_x_" + item)

benchmark/training.py

@@ -0,0 +1,123 @@
+from os import removedirs
+from os.path import exists
+from typing import override
+
+from monai.transforms import Compose, Resized
+from torch.utils.data import DataLoader
+
+from benchmark.data import DataTest, FoldedDataTest
+from benchmark.transforms import training_transforms, validation_transforms
+from benchmark.unet import UNetTrainer, UNetSlidingTrainer
+from mipcandy import SegmentationTrainer, slide_dataset, Shape, SupervisedSWDataset, JointTransform, inspect, \
+    ROIDataset, PadTo, MONAITransform, load_inspection_annotations
+
+
+class TrainingTest(DataTest):
+    trainer: type[SegmentationTrainer] = UNetTrainer
+    resize: Shape = (128, 128, 128)
+    num_classes: int = 5
+    _continue: str | None = None  # internal flag for continued training
+
+    def set_up_datasets(self) -> None:
+        super().set_up()
+        self["dataset"].device(device="cpu")
+        self["dataset"].set_transform(
+            JointTransform(transform=MONAITransform(PadTo(self.resize, batch=False)))
+        )
+        path = f"{self.input_folder}/training_test.json"
+        if exists(path):
+            annotations = load_inspection_annotations(path, self["dataset"])
+        else:
+            annotations = inspect(self["dataset"])
+            annotations.save(path)
+        annotations.set_roi_shape(self.resize)
+        dataset = ROIDataset(annotations)
+        self["train_dataset"], self["val_dataset"] = dataset.fold(fold=0)
+
+    @override
+    def set_up(self) -> None:
+        self.set_up_datasets()
+        train, val = self["train_dataset"], self["val_dataset"]
+        train.set_transform(JointTransform(transform=Compose([
+            train.transform().transform, training_transforms()
+        ])))
+        val.set_transform(JointTransform(transform=Compose([
+            val.transform().transform, validation_transforms()
+        ])))
+        train_dataloader = DataLoader(train, batch_size=2, shuffle=True, pin_memory=True)
+        val_dataloader = DataLoader(val, batch_size=1, shuffle=False, pin_memory=True)
+        trainer = self.trainer(self.output_folder, train_dataloader, val_dataloader, device=self.device)
+        trainer.num_classes = self.num_classes
+        trainer.set_frontend(self.frontend)
+        self["trainer"] = trainer
+
+    @override
+    def execute(self) -> None:
+        if not self._continue:
+            return self["trainer"].train(self.num_epochs, note=f"Training test {self.resize}")
+        self["trainer"].recover_from(self._continue)
+        return self["trainer"].continue_training(self.num_epochs)
+
+    @override
+    def clean_up(self) -> None:
+        removedirs(self["trainer"].experiment_folder())
+
+
+class ResizeTrainingTest(FoldedDataTest):
+    trainer: type[SegmentationTrainer] = UNetTrainer
+    resize: Shape = (256, 256, 256)
+    num_classes: int = 5
+
+    @override
+    def set_up(self) -> None:
+        self.transform = JointTransform(transform=Resized(("image", "label"), self.resize))
+        super().set_up()
+        train_dataloader = DataLoader(self["train_dataset"], batch_size=2, shuffle=True)
+        val_dataloader = DataLoader(self["val_dataset"], batch_size=1, shuffle=False)
+        trainer = self.trainer(self.output_folder, train_dataloader, val_dataloader, recoverable=False,
+                                       profiler=True, device=self.device)
+        trainer.num_classes = self.num_classes
+        trainer.set_frontend(self.frontend)
+        self["trainer"] = trainer
+
+    @override
+    def execute(self) -> None:
+        self["trainer"].train(self.num_epochs, note=f"Resize Training test {self.resize}")
+
+    @override
+    def clean_up(self) -> None:
+        removedirs(self["trainer"].experiment_folder())
+
+
+class SlidingTrainingTest(TrainingTest, FoldedDataTest):
+    trainer: type[SegmentationTrainer] = UNetSlidingTrainer
+    window_shape: Shape = (128, 128, 128)
+    overlap: float = .5
+
+    @override
+    def set_up(self) -> None:
+        self.set_up_datasets()
+        train, val = self["train_dataset"], self["val_dataset"]
+        FoldedDataTest.set_up(self)
+        full_val = self["val_dataset"]
+        path = f"{self.output_folder}/val_slided"
+        if not exists(path):
+            slide_dataset(full_val, path, self.window_shape, overlap=self.overlap)
+        slided_val = SupervisedSWDataset(path)
+        train_dataloader = DataLoader(train, batch_size=2, shuffle=True)
+        val_dataloader = DataLoader(val, batch_size=1, shuffle=False)
+        trainer = self.trainer(self.output_folder, train_dataloader, val_dataloader, recoverable=False,
+                                       profiler=True, device=self.device)
+        trainer.set_datasets(full_val, slided_val)
+        trainer.num_classes = self.num_classes
+        trainer.overlap = self.overlap
+        trainer.set_frontend(self.frontend)
+        self["trainer"] = trainer
+
+    @override
+    def execute(self) -> None:
+        self["trainer"].train(self.num_epochs, note="Training test with sliding window")
+
+    @override
+    def clean_up(self) -> None:
+        removedirs(self["trainer"].experiment_folder())