Merge pull request #76 from Hendrik-code/new_features_robert

New features robert

Author: robert-graf

TPTBox/core/nii_wrapper.py

@@ -62,7 +62,7 @@
 )
 
 if TYPE_CHECKING:
-    pass
+    from torch import device
 MODES = Literal["constant", "nearest", "reflect", "wrap"]
 _unpacked_nii = tuple[np.ndarray, AFFINE, nib.nifti1.Nifti1Header]
 _formatwarning = warnings.formatwarning
@@ -1042,7 +1042,7 @@ def to_simpleITK(self):
         from TPTBox.core.sitk_utils import nii_to_sitk
         return nii_to_sitk(self)
 
-    def to_deepali(self,align_corners: bool = True,dtype=None,device = "cuda"):
+    def to_deepali(self,align_corners: bool = True,dtype=None,device:device|str = "cpu"):
         import torch
         try:
             from deepali.data import Image as deepaliImage  # type: ignore

TPTBox/core/poi_fun/poi_abstract.py

@@ -19,7 +19,7 @@
 POI_ID = Union[
     tuple[int, int],
     slice,
-    tuple[Abstract_lvl | int, Abstract_lvl | int],
+    tuple[Union[Abstract_lvl, int], Union[Abstract_lvl, int]],
     tuple[Abstract_lvl, Abstract_lvl],
     tuple[Abstract_lvl, int],
     tuple[int, Abstract_lvl],

TPTBox/core/vert_constants.py

@@ -107,6 +107,9 @@ def _get_id(cls, s: str | int, no_raise=True) -> int:
         try:
             return cls[s].value
         except KeyError:
+            for c in cls:
+                if c.name.lower() == s.lower():
+                    return c.value
             if not no_raise:
                 raise
             return int(s)

TPTBox/registration/deformable/deformable_reg.py

@@ -231,9 +231,9 @@ def save(self, path: str | Path):
             pickle.dump(self.get_dump(), w)
 
     @classmethod
-    def load(cls, path):
+    def load(cls, path, gpu=0, ddevice: DEVICES = "cuda"):
         with open(path, "rb") as w:
-            return cls.load_(pickle.load(w))
+            return cls.load_(pickle.load(w), gpu, ddevice)
 
     @classmethod
     def load_(cls, w, gpu=0, ddevice: DEVICES = "cuda"):

TPTBox/segmentation/TotalVibeSeg/auto_download.py

@@ -67,10 +67,11 @@ def user_guard(func: Any) -> Any:
 
 
 @user_guard
-def _download_weights(idx=85) -> None:
+def _download_weights(idx=85, addendum="") -> None:
     weights_dir = get_weights_dir(idx)
-    weights_url = WEIGHTS_URL_ + f"{idx:03}.zip"
+    weights_url = WEIGHTS_URL_ + f"{idx:03}{addendum}.zip"
     _download(weights_url, weights_dir, text="pretrained weights")
+    addendum_download(idx)
 
 
 def _download(weights_url, weights_dir, text="") -> None:
@@ -100,10 +101,21 @@ def update_progress(block_num: int, block_size: int, total_size: int) -> None:
     os.remove(zip_path)  # noqa: PTH107
 
 
+def addendum_download(idx):
+    weights_dir = get_weights_dir(idx)
+    next_zip = weights_dir / "other_downloads.json"
+    if next_zip.exists():
+        with open(next_zip) as f:
+            add = json.load(f)
+        [_download_weights(idx, addendum=a) for a in add]
+        next_zip.unlink()
+
+
 def download_weights(idx, model_path: Path | None = None) -> Path:
     weights_dir = get_weights_dir(idx, model_path)
 
     # Check if weights are downloaded
     if not weights_dir.exists():
         _download_weights(idx)
+    addendum_download(idx)
     return weights_dir

TPTBox/segmentation/TotalVibeSeg/inference_nnunet.py

@@ -22,9 +22,12 @@ def get_ds_info(idx) -> dict:
     try:
         nnunet_path = next(next(iter(model_path.glob(f"*{idx}*"))).glob("*__nnUNetPlans*"))
     except StopIteration:
-        Print_Logger().print(f"Please add Dataset {idx} to {model_path}", Log_Type.FAIL)
-        model_path.mkdir(exist_ok=True, parents=True)
-        sys.exit()
+        try:
+            nnunet_path = next(next(iter(model_path.glob(f"*{idx}*"))).glob("*__nnUNet*ResEnc*"))
+        except StopIteration:
+            Print_Logger().print(f"Please add Dataset {idx} to {model_path}", Log_Type.FAIL)
+            model_path.mkdir(exist_ok=True, parents=True)
+            sys.exit()
     with open(Path(nnunet_path, "dataset.json")) as f:
         ds_info = json.load(f)
     return ds_info

TPTBox/segmentation/nnUnet_utils/predictor.py

@@ -166,22 +166,22 @@ def mapp(d: dict):
         if "ignore" in dataset_json["labels"]:
             num_output_channels -= 1
 
-            network = get_network_from_plans(
-                plans_manager,  # type: ignore
-                dataset_json,
-                configuration_manager,  # type: ignore
-                num_input_channels,
-                num_output_channels=num_output_channels,
-                deep_supervision=False,
-            )
-            self.network = network
-
-            self.plans_manager = plans_manager
-            self.configuration_manager = configuration_manager
-            self.list_of_parameters = parameters  # Lists of model folds
-            self.dataset_json = dataset_json
-            self.trainer_name = trainer_name
-            self.allowed_mirroring_axes = inference_allowed_mirroring_axes
+        network = get_network_from_plans(
+            plans_manager,  # type: ignore
+            dataset_json,
+            configuration_manager,  # type: ignore
+            num_input_channels,
+            num_output_channels=num_output_channels,
+            deep_supervision=False,
+        )
+        self.network = network
+
+        self.plans_manager = plans_manager
+        self.configuration_manager = configuration_manager
+        self.list_of_parameters = parameters  # Lists of model folds
+        self.dataset_json = dataset_json
+        self.trainer_name = trainer_name
+        self.allowed_mirroring_axes = inference_allowed_mirroring_axes
         self.label_manager = plans_manager.get_label_manager(dataset_json)
         if (
             ("nnUNet_compile" in os.environ.keys())

unit_tests/test_testsamples.py

@@ -214,10 +214,10 @@ def _test_deformable(self, dim=0, save=True):  # type: ignore
         poi[123, 44] = (random.randint(0, mri.shape[0] - 1), random.randint(0, mri.shape[1] - 1), random.randint(0, mri.shape[2] - 1))
         poi[123, 45] = (random.randint(0, mri.shape[0] - 1), random.randint(0, mri.shape[1] - 1), random.randint(0, mri.shape[2] - 1))
 
-        deform = Deformable_Registration(mov, mov, reference_image=mov)
+        deform = Deformable_Registration(mov, mov, reference_image=mov, ddevice="cpu")
         if save:
             deform.save(test_save)
-            deform = Deformable_Registration.load(test_save)
+            deform = Deformable_Registration.load(test_save, ddevice="cpu")
         mov2 = mov.copy()
         mov2.seg = True
         mov2[mov > -10000] = 0