heavy unet tests - v0.3.5 (#327)

multiple unitteststs + fixes

dacapo/__init__.py

@@ -1,4 +1,4 @@
-__version__ = "0.3.2"
+__version__ = "0.3.5"
 __version_info__ = tuple(int(i) for i in __version__.split("."))
 
 from .options import Options  # noqa

dacapo/compute_context/local_torch.py

@@ -56,9 +56,10 @@ def device(self):
         if self._device is None:
             if torch.cuda.is_available():
                 # TODO: make this more sophisticated, for multiple GPUs for instance
-                free = torch.cuda.mem_get_info()[0] / 1024**3
-                if free < self.oom_limit:  # less than 1 GB free, decrease chance of OOM
-                    return torch.device("cpu")
+                # commented out code below is for checking free memory and falling back on CPU, whhen model in GPU and memory is low model get moved to CPU
+                # free = torch.cuda.mem_get_info()[0] / 1024**3
+                # if free < self.oom_limit:  # less than 1 GB free, decrease chance of OOM
+                #     return torch.device("cpu")
                 return torch.device("cuda")
             # Multiple MPS ops are not available yet : https://github.com/pytorch/pytorch/issues/77764
             # got error aten::max_pool3d_with_indices

dacapo/experiments/architectures/cnnectome_unet.py

@@ -3,6 +3,8 @@
 import torch
 import torch.nn as nn
 
+from funlib.geometry import Coordinate
+
 import math
 
 
@@ -176,7 +178,7 @@ def __init__(self, architecture_config):
         self.unet = self.module()
 
     @property
-    def eval_shape_increase(self):
+    def eval_shape_increase(self) -> Coordinate:
         """
         The increase in shape due to the U-Net.
 
@@ -192,7 +194,7 @@ def eval_shape_increase(self):
         """
         if self._eval_shape_increase is None:
             return super().eval_shape_increase
-        return self._eval_shape_increase
+        return Coordinate(self._eval_shape_increase)
 
     def module(self):
         """
@@ -235,16 +237,15 @@ def module(self):
         """
         fmaps_in = self.fmaps_in
         levels = len(self.downsample_factors) + 1
-        dims = len(self.downsample_factors[0])
 
-        if hasattr(self, "kernel_size_down"):
+        if self.kernel_size_down is not None:
             kernel_size_down = self.kernel_size_down
         else:
-            kernel_size_down = [[(3,) * dims, (3,) * dims]] * levels
-        if hasattr(self, "kernel_size_up"):
+            kernel_size_down = [[(3,) * self.dims, (3,) * self.dims]] * levels
+        if self.kernel_size_up is not None:
             kernel_size_up = self.kernel_size_up
         else:
-            kernel_size_up = [[(3,) * dims, (3,) * dims]] * (levels - 1)
+            kernel_size_up = [[(3,) * self.dims, (3,) * self.dims]] * (levels - 1)
 
         # downsample factors has to be a list of tuples
         downsample_factors = [tuple(x) for x in self.downsample_factors]
@@ -280,7 +281,7 @@ def module(self):
                 conv = ConvPass(
                     self.fmaps_out,
                     self.fmaps_out,
-                    [(3,) * len(upsample_factor)] * 2,
+                    kernel_size_up[-1],
                     activation="ReLU",
                     batch_norm=self.batch_norm,
                 )
@@ -306,11 +307,11 @@ def scale(self, voxel_size):
             The voxel size should be given as a tuple ``(z, y, x)``.
         """
         for upsample_factor in self.upsample_factors:
-            voxel_size = voxel_size / upsample_factor
+            voxel_size = voxel_size / Coordinate(upsample_factor)
         return voxel_size
 
     @property
-    def input_shape(self):
+    def input_shape(self) -> Coordinate:
         """
         Return the input shape of the U-Net.
 
@@ -324,7 +325,7 @@ def input_shape(self):
         Note:
             The input shape should be given as a tuple ``(batch, channels, [length,] depth, height, width)``.
         """
-        return self._input_shape
+        return Coordinate(self._input_shape)
 
     @property
     def num_in_channels(self) -> int:

dacapo/experiments/architectures/cnnectome_unet_config.py

@@ -128,6 +128,6 @@ class CNNectomeUNetConfig(ArchitectureConfig):
         },
     )
     batch_norm: bool = attr.ib(
-        default=True,
+        default=False,
         metadata={"help_text": "Whether to use batch normalization."},
     )

dacapo/experiments/datasplits/datasets/dataset_config.py

@@ -63,3 +63,6 @@ def verify(self) -> Tuple[bool, str]:
             This method is used to validate the configuration of the dataset.
         """
         return True, "No validation for this DataSet"
+
+    def __hash__(self) -> int:
+        return hash(self.name)

dacapo/experiments/datasplits/simple_config.py

@@ -44,7 +44,7 @@ def get_paths(self, group_name: str) -> list[Path]:
                 len(level_2_matches) == 0
             ), f"Found raw data at {level_1} and {level_2}"
             return [Path(x).parent for x in level_1_matches]
-        elif len(level_2_matches).parent > 0:
+        elif len(level_2_matches) > 0:
             return [Path(x) for x in level_2_matches]
 
         raise Exception(f"No raw data found at {level_0} or {level_1} or {level_2}")

dacapo/experiments/tasks/hot_distance_task.py

@@ -4,6 +4,8 @@
 from .predictors import HotDistancePredictor
 from .task import Task
 
+import warnings
+
 
 class HotDistanceTask(Task):
     """
@@ -34,10 +36,19 @@ def __init__(self, task_config):
             >>> task = HotDistanceTask(task_config)
 
         """
+
+        if task_config.kernel_size is None:
+            warnings.warn(
+                "The default kernel size of 3 will be changing to 1. "
+                "Please specify the kernel size explicitly.",
+                DeprecationWarning,
+            )
+            task_config.kernel_size = 3
         self.predictor = HotDistancePredictor(
             channels=task_config.channels,
             scale_factor=task_config.scale_factor,
             mask_distances=task_config.mask_distances,
+            kernel_size=task_config.kernel_size,
         )
         self.loss = HotDistanceLoss()
         self.post_processor = ThresholdPostProcessor()

dacapo/experiments/tasks/hot_distance_task_config.py

@@ -56,3 +56,7 @@ class HotDistanceTaskConfig(TaskConfig):
             "is less than the distance to object boundary."
         },
     )
+
+    kernel_size: int | None = attr.ib(
+        default=None,
+    )

dacapo/experiments/tasks/one_hot_task.py

@@ -4,6 +4,8 @@
 from .predictors import OneHotPredictor
 from .task import Task
 
+import warnings
+
 
 class OneHotTask(Task):
     """
@@ -30,7 +32,17 @@ def __init__(self, task_config):
         Examples:
             >>> task = OneHotTask(task_config)
         """
-        self.predictor = OneHotPredictor(classes=task_config.classes)
+
+        if task_config.kernel_size is None:
+            warnings.warn(
+                "The default kernel size of 3 will be changing to 1. "
+                "Please specify the kernel size explicitly.",
+                DeprecationWarning,
+            )
+            task_config.kernel_size = 3
+        self.predictor = OneHotPredictor(
+            classes=task_config.classes, kernel_size=task_config.kernel_size
+        )
         self.loss = DummyLoss()
         self.post_processor = ArgmaxPostProcessor()
         self.evaluator = DummyEvaluator()

dacapo/experiments/tasks/one_hot_task_config.py

@@ -28,3 +28,6 @@ class OneHotTaskConfig(TaskConfig):
     classes: List[str] = attr.ib(
         metadata={"help_text": "The classes corresponding with each id starting from 0"}
     )
+    kernel_size: int | None = attr.ib(
+        default=None,
+    )

dacapo/experiments/tasks/post_processors/argmax_post_processor.py

@@ -133,7 +133,7 @@ def process(
             overwrite=True,
         )
 
-        read_roi = Roi((0, 0, 0), block_size[-self.prediction_array.dims :])
+        read_roi = Roi((0,) * block_size.dims, block_size)
         input_array = open_ds(
             f"{self.prediction_array_identifier.container.path}/{self.prediction_array_identifier.dataset}"
         )

dacapo/experiments/tasks/post_processors/threshold_post_processor.py

@@ -18,6 +18,9 @@
 from funlib.persistence import Array
 
 from typing import Iterable
+import logging
+
+logger = logging.getLogger(__name__)
 
 
 class ThresholdPostProcessor(PostProcessor):
@@ -108,13 +111,15 @@ def process(
         if self.prediction_array._source_data.chunks is not None:
             block_size = self.prediction_array._source_data.chunks
 
-        write_size = [
-            b * v
-            for b, v in zip(
-                block_size[-self.prediction_array.dims :],
-                self.prediction_array.voxel_size,
-            )
-        ]
+        write_size = Coordinate(
+            [
+                b * v
+                for b, v in zip(
+                    block_size[-self.prediction_array.dims :],
+                    self.prediction_array.voxel_size,
+                )
+            ]
+        )
         output_array = create_from_identifier(
             output_array_identifier,
             self.prediction_array.axis_names,
@@ -125,7 +130,7 @@ def process(
             overwrite=True,
         )
 
-        read_roi = Roi((0, 0, 0), write_size[-self.prediction_array.dims :])
+        read_roi = Roi(write_size * 0, write_size)
         input_array = open_ds(
             f"{self.prediction_array_identifier.container.path}/{self.prediction_array_identifier.dataset}"
         )
@@ -135,7 +140,7 @@ def process_block(block):
             data = input_array[write_roi] > parameters.threshold
             data = data.astype(np.uint8)
             if int(data.max()) == 0:
-                print("No data in block", write_roi)
+                logger.debug("No data in block", write_roi)
                 return
             output_array[write_roi] = data
 

dacapo/experiments/tasks/predictors/distance_predictor.py

@@ -223,6 +223,10 @@ def create_distance_mask(
             >>> predictor.create_distance_mask(distances, mask, voxel_size, normalize, normalize_args)
 
         """
+        no_channel_dim = len(mask.shape) == len(distances.shape) - 1
+        if no_channel_dim:
+            mask = mask[np.newaxis]
+
         mask_output = mask.copy()
         for i, (channel_distance, channel_mask) in enumerate(zip(distances, mask)):
             tmp = np.zeros(
@@ -231,9 +235,11 @@ def create_distance_mask(
             )
             slices = tmp.ndim * (slice(1, -1),)
             tmp[slices] = channel_mask
+            sampling = tuple(float(v) / 2 for v in voxel_size)
+            sampling = sampling[-len(tmp.shape) :]
             boundary_distance = distance_transform_edt(
                 tmp,
-                sampling=voxel_size,
+                sampling=sampling,
             )
             if self.epsilon is None:
                 add = 0
@@ -273,6 +279,8 @@ def create_distance_mask(
                     np.sum(channel_mask_output)
                 )
             )
+        if no_channel_dim:
+            mask_output = mask_output[0]
         return mask_output
 
     def process(
@@ -298,7 +306,20 @@ def process(
             >>> predictor.process(labels, voxel_size, normalize, normalize_args)
 
         """
+
+        num_dims = len(labels.shape)
+        if num_dims == voxel_size.dims:
+            channel_dim = False
+        elif num_dims == voxel_size.dims + 1:
+            channel_dim = True
+        else:
+            raise ValueError("Cannot handle multiple channel dims")
+
+        if not channel_dim:
+            labels = labels[np.newaxis]
+
         all_distances = np.zeros(labels.shape, dtype=np.float32) - 1
+
         for ii, channel in enumerate(labels):
             boundaries = self.__find_boundaries(channel)
 
@@ -315,13 +336,15 @@ def process(
                     distances = np.ones(channel.shape, dtype=np.float32) * max_distance
             else:
                 # get distances (voxel_size/2 because image is doubled)
-                distances = distance_transform_edt(
-                    boundaries, sampling=tuple(float(v) / 2 for v in voxel_size)
-                )
+                sampling = tuple(float(v) / 2 for v in voxel_size)
+                # fixing the sampling for 2D images
+                if len(boundaries.shape) < len(sampling):
+                    sampling = sampling[-len(boundaries.shape) :]
+                distances = distance_transform_edt(boundaries, sampling=sampling)
                 distances = distances.astype(np.float32)
 
                 # restore original shape
-                downsample = (slice(None, None, 2),) * len(voxel_size)
+                downsample = (slice(None, None, 2),) * distances.ndim
                 distances = distances[downsample]
 
                 # todo: inverted distance
@@ -354,7 +377,7 @@ def __find_boundaries(self, labels: np.ndarray):
         # bound.: 00000001000100000001000      2n - 1
 
         if labels.dtype == bool:
-            raise ValueError("Labels should not be bools")
+            # raise ValueError("Labels should not be bools")
             labels = labels.astype(np.uint8)
 
         logger.debug(f"computing boundaries for {labels.shape}")

dacapo/experiments/tasks/predictors/dummy_predictor.py

@@ -50,7 +50,7 @@ def create_model(self, architecture):
             >>> model = predictor.create_model(architecture)
         """
         head = torch.nn.Conv3d(
-            architecture.num_out_channels, self.embedding_dims, kernel_size=3
+            architecture.num_out_channels, self.embedding_dims, kernel_size=1
         )
 
         return Model(architecture, head)
@@ -71,9 +71,8 @@ def create_target(self, gt):
         # zeros
         return np_to_funlib_array(
             np.zeros((self.embedding_dims,) + gt.data.shape[-gt.dims :]),
-            gt.roi,
+            gt.roi.offset,
             gt.voxel_size,
-            ["c^"] + gt.axis_names,
         )
 
     def create_weight(self, gt, target, mask, moving_class_counts=None):
@@ -96,9 +95,8 @@ def create_weight(self, gt, target, mask, moving_class_counts=None):
         return (
             np_to_funlib_array(
                 np.ones(target.data.shape),
-                target.roi,
+                target.roi.offset,
                 target.voxel_size,
-                target.axis_names,
             ),
             None,
         )