Merge pull request #64 from mprires/main

Changed vertebrae model. Made custom trainers available.

moosez/__init__.py

@@ -8,3 +8,4 @@
 from . import predict
 from . import resources
 from .moosez import moose
+from .nnUNet_custom_trainer import utility

moosez/moosez.py

@@ -39,6 +39,7 @@
 from moosez import resources
 from moosez.image_processing import ImageResampler
 from moosez.resources import MODELS, AVAILABLE_MODELS
+from moosez.nnUNet_custom_trainer.utility import add_custom_trainers_to_local_nnunetv2
 
 logging.basicConfig(format='%(asctime)s %(levelname)-8s [%(filename)s:%(lineno)d] %(message)s', level=logging.INFO,
                     filename=datetime.now().strftime('moosez-v.2.0.0.%H-%M-%d-%m-%Y.log'),
@@ -108,6 +109,8 @@ def main():
     print(f'{constants.ANSI_VIOLET} {emoji.emojize(":memo:")} NOTE:{constants.ANSI_RESET}')
     print(' ')
     modalities = display.expectations(model_name)
+    custom_trainer_status = add_custom_trainers_to_local_nnunetv2()
+    logging.info('- Custom trainer: ' + custom_trainer_status)
     accelerator = resources.check_cuda()
 
     # ----------------------------------
@@ -279,6 +282,8 @@ def moose(model_name: str, input_dir: str, output_dir: str, accelerator: str) ->
     model_path = constants.NNUNET_RESULTS_FOLDER
     file_utilities.create_directory(model_path)
     download.model(model_name, model_path)
+    custom_trainer_status = add_custom_trainers_to_local_nnunetv2()
+    logging.info('- Custom trainer: ' + custom_trainer_status)
     input_validation.make_nnunet_compatible(input_dir)
     predict.predict(model_name, input_dir, output_dir, accelerator)
 

moosez/nnUNet_custom_trainer/MOOSE_custom_trainers.py

@@ -0,0 +1,23 @@
+from nnunetv2.training.nnUNetTrainer.variants.data_augmentation.nnUNetTrainerDA5 import nnUNetTrainerDA5
+import torch
+
+
+class nnUNetTrainer_2000_epochs_DA5NoMirroring(nnUNetTrainerDA5):
+    def __init__(self, plans: dict, configuration: str, fold: int, dataset_json: dict, unpack_dataset: bool = True,
+                 device: torch.device = torch.device('cuda')):
+        super().__init__(plans, configuration, fold, dataset_json, unpack_dataset, device)
+        self.num_epochs = 2000
+
+        print(f"Epochs: {self.num_epochs}")
+
+
+    def configure_rotation_dummyDA_mirroring_and_inital_patch_size(self):
+        rotation_for_DA, do_dummy_2d_data_aug, initial_patch_size, mirror_axes = \
+            super().configure_rotation_dummyDA_mirroring_and_inital_patch_size()
+        print(f"Mirroring: {mirror_axes} | {self.inference_allowed_mirroring_axes}")
+
+        mirror_axes = None
+        self.inference_allowed_mirroring_axes = None
+        print(f"Mirroring: {mirror_axes} | {self.inference_allowed_mirroring_axes}")
+
+        return rotation_for_DA, do_dummy_2d_data_aug, initial_patch_size, mirror_axes

moosez/nnUNet_custom_trainer/utility.py

@@ -0,0 +1,28 @@
+import site
+import os
+import platform
+from moosez import file_utilities
+
+
+def add_custom_trainers_to_local_nnunetv2():
+    # Locate the site-packages directory
+    site_packages = site.getsitepackages()[0]
+
+    if platform.system() == 'Windows':
+        site_packages = os.path.join(site_packages, 'Lib', 'site-packages')
+
+    source_file_path = os.path.join(site_packages, 'moosez', 'nnUNet_custom_trainer', 'MOOSE_custom_trainers.py')
+    target_file_path = os.path.join(site_packages, 'nnunetv2', 'training', 'nnUNetTrainer', 'variants',
+                                    'MOOSE_custom_trainers.py')
+
+    # Check if the file exists
+    if not os.path.exists(source_file_path):
+        return f'Could not find custom trainer file: {source_file_path}'
+
+    # Check if the file exists
+    if os.path.exists(target_file_path):
+        return f'Custom trainer already installed: {target_file_path}.'
+
+    file_utilities.copy_file(source_file_path, target_file_path)
+
+    return f"Custom trainer added to {target_file_path}"

moosez/resources.py

@@ -128,10 +128,10 @@
         "multilabel_prefix": "CT_Fat_"
     },
     "clin_ct_vertebrae":{
-        "url": "https://moose-files.s3.eu-de.cloud-object-storage.appdomain.cloud/clin_ct_vertebrae_04092023.zip",
+        "url": "https://moose-files.s3.eu-de.cloud-object-storage.appdomain.cloud/clin_ct_vertebrae_da5_03102023.zip",
         "filename": "Dataset111_Vertebrae.zip",
         "directory": "Dataset111_Vertebrae",
-        "trainer": "nnUNetTrainer_2000epochs_NoMirroring",
+        "trainer": "nnUNetTrainer_2000_epochs_DA5NoMirroring",
         "voxel_spacing": [1.5, 1.5, 1.5],
         "multilabel_prefix": "CT_Vertebrae_"
     },

setup.py

@@ -2,7 +2,7 @@
 
 setup(
     name='moosez',
-    version='2.2.32',
+    version='2.2.33',
     author='Lalith Kumar Shiyam Sundar | Sebastian Gutschmayer',
     author_email='[email protected]',
     description='An AI-inference engine for 3D clinical and preclinical whole-body segmentation tasks',