deploy changes

Author: github-actions[bot]

asparagus/functional/loading.py

@@ -0,0 +1,15 @@
+import nibabel as nib
+import numpy as np
+import torch
+
+
+def load_image_file(file: str) -> torch.Tensor:
+    if file.endswith(".pt"):
+        return torch.load(file)
+    elif file.endswith(".nii.gz") or file.endswith(".nii"):
+        nii = nib.load(file)
+        data = nii.get_fdata(dtype=np.float32)
+        tensor = torch.from_numpy(data)
+        return tensor.unsqueeze(0)  # (H,W,D) -> (1,H,W,D) to match .pt channel convention
+    else:
+        raise ValueError(f"Unsupported file format: {file}. Expected .pt, .nii, or .nii.gz")

asparagus/modules/datasets/PretrainDataset.py

@@ -1,5 +1,6 @@
 import torch
 import torchvision
+from asparagus.functional.loading import load_image_file
 from torch.utils.data import Dataset
 from typing import Optional
 
@@ -20,7 +21,7 @@ def __len__(self):
 
     def __getitem__(self, idx):
         file = self.files[idx]
-        data = torch.load(file)
+        data = load_image_file(file)
         data_dict = {"file_path": file, "image": data, "transforms_applied": {}}
         data_dict = self._transform(data_dict)  # CPU transforms only here
 

asparagus/modules/networks/unet.py

@@ -84,6 +84,7 @@ def unet_tiny(
     output_channels: int = 1,
     dimensions: str = "3D",
     deep_supervision: bool = False,
+    use_skip_connections: bool = True,
 ):
     return UNet(
         input_channels=input_channels,
@@ -93,6 +94,7 @@ def unet_tiny(
         encoder_basic_block=MultiLayerConvDropoutNormNonlin.get_block_constructor(1),
         decoder_basic_block=MultiLayerConvDropoutNormNonlin.get_block_constructor(1),
         deep_supervision=deep_supervision,
+        use_skip_connections=use_skip_connections,
     )