For consistency, introduce types Shape2d and Shape3d (#148)

* Added `Shape2d`, `Shape3d`, and `Shape` type definitions. (#137)

* Refactored type annotations to use `Shape`, `Shape2d`, and `Shape3d` across modules. (#137)

* Expanded type annotations to include `AmbiguousShape` and updated affected modules accordingly. (#137)

Author: ATATC

mipcandy/__init__.py

@@ -14,4 +14,5 @@
 from mipcandy.run import config
 from mipcandy.sanity_check import num_trainable_params, model_complexity_info, SanityCheckResult, sanity_check
 from mipcandy.training import TrainerToolbox, Trainer, SWMetadata, SlidingTrainer
-from mipcandy.types import Setting, Settings, Params, Transform, SupportedPredictant, Colormap, Device
+from mipcandy.types import Setting, Settings, Params, Transform, SupportedPredictant, Colormap, Device, Shape2d, \
+    Shape3d, Shape, AmbiguousShape

mipcandy/common/module/preprocess.py

@@ -4,7 +4,7 @@
 import torch
 from torch import nn
 
-from mipcandy.types import Colormap
+from mipcandy.types import Colormap, Shape2d, Shape3d
 
 
 class Pad(nn.Module):
@@ -34,10 +34,10 @@ def _c_p(size: int, min_factor: int) -> tuple[int, int]:
 
 
 class Pad2d(Pad):
-    def __init__(self, min_factor: int | tuple[int, int], *, value: int = 0, mode: str = "constant",
+    def __init__(self, min_factor: int | Shape2d, *, value: int = 0, mode: str = "constant",
                  batch: bool = True) -> None:
         super().__init__(value=value, mode=mode, batch=batch)
-        self._min_factor: tuple[int, int] = (min_factor,) * 2 if isinstance(min_factor, int) else min_factor
+        self._min_factor: Shape2d = (min_factor,) * 2 if isinstance(min_factor, int) else min_factor
 
     def paddings(self) -> tuple[int, int, int, int] | None:
         return self._paddings
@@ -58,10 +58,10 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
 
 
 class Pad3d(Pad):
-    def __init__(self, min_factor: int | tuple[int, int, int], *, value: int = 0, mode: str = "constant",
+    def __init__(self, min_factor: int | Shape3d, *, value: int = 0, mode: str = "constant",
                  batch: bool = True) -> None:
         super().__init__(value=value, mode=mode, batch=batch)
-        self._min_factor: tuple[int, int, int] = (min_factor,) * 3 if isinstance(min_factor, int) else min_factor
+        self._min_factor: Shape3d = (min_factor,) * 3 if isinstance(min_factor, int) else min_factor
 
     def paddings(self) -> tuple[int, int, int, int, int, int] | None:
         return self._paddings

mipcandy/data/inspection.py

@@ -12,7 +12,7 @@
 from mipcandy.data.dataset import SupervisedDataset
 from mipcandy.data.geometric import crop
 from mipcandy.layer import HasDevice
-from mipcandy.types import Device
+from mipcandy.types import Device, Shape, AmbiguousShape
 
 
 def format_bbox(bbox: Sequence[int]) -> tuple[int, int, int, int] | tuple[int, int, int, int, int, int]:
@@ -26,11 +26,11 @@ def format_bbox(bbox: Sequence[int]) -> tuple[int, int, int, int] | tuple[int, i
 
 @dataclass
 class InspectionAnnotation(object):
-    shape: tuple[int, ...]
+    shape: AmbiguousShape
     foreground_bbox: tuple[int, int, int, int] | tuple[int, int, int, int, int, int]
     ids: tuple[int, ...]
 
-    def foreground_shape(self) -> tuple[int, int] | tuple[int, int, int]:
+    def foreground_shape(self) -> Shape:
         r = (self.foreground_bbox[1] - self.foreground_bbox[0], self.foreground_bbox[3] - self.foreground_bbox[2])
         return r if len(self.foreground_bbox) == 4 else r + (self.foreground_bbox[5] - self.foreground_bbox[4],)
 
@@ -50,13 +50,13 @@ def __init__(self, dataset: SupervisedDataset, background: int, *annotations: In
         self._dataset: SupervisedDataset = dataset
         self._background: int = background
         self._annotations: tuple[InspectionAnnotation, ...] = annotations
-        self._shapes: tuple[tuple[int, ...] | None, tuple[int, ...], tuple[int, ...]] | None = None
-        self._foreground_shapes: tuple[tuple[int, ...] | None, tuple[int, ...], tuple[int, ...]] | None = None
-        self._statistical_foreground_shape: tuple[int, int] | tuple[int, int, int] | None = None
+        self._shapes: tuple[AmbiguousShape | None, AmbiguousShape, AmbiguousShape] | None = None
+        self._foreground_shapes: tuple[AmbiguousShape | None, AmbiguousShape, AmbiguousShape] | None = None
+        self._statistical_foreground_shape: Shape | None = None
         self._foreground_heatmap: torch.Tensor | None = None
         self._center_of_foregrounds: tuple[int, int] | tuple[int, int, int] | None = None
         self._foreground_offsets: tuple[int, int] | tuple[int, int, int] | None = None
-        self._roi_shape: tuple[int, int] | tuple[int, int, int] | None = None
+        self._roi_shape: Shape | None = None
 
     def dataset(self) -> SupervisedDataset:
         return self._dataset
@@ -79,8 +79,8 @@ def save(self, path: str | PathLike[str]) -> None:
         with open(path, "w") as f:
             dump({"background": self._background, "annotations": [a.to_dict() for a in self._annotations]}, f)
 
-    def _get_shapes(self, get_shape: Callable[[InspectionAnnotation], tuple[int, ...]]) -> tuple[
-        tuple[int, ...] | None, tuple[int, ...], tuple[int, ...]]:
+    def _get_shapes(self, get_shape: Callable[[InspectionAnnotation], AmbiguousShape]) -> tuple[
+        AmbiguousShape | None, AmbiguousShape, AmbiguousShape]:
         depths = []
         widths = []
         heights = []
@@ -95,19 +95,19 @@ def _get_shapes(self, get_shape: Callable[[InspectionAnnotation], tuple[int, ...
                 widths.append(shape[2])
         return tuple(depths) if depths else None, tuple(heights), tuple(widths)
 
-    def shapes(self) -> tuple[tuple[int, ...] | None, tuple[int, ...], tuple[int, ...]]:
+    def shapes(self) -> tuple[AmbiguousShape | None, AmbiguousShape, AmbiguousShape]:
         if self._shapes:
             return self._shapes
         self._shapes = self._get_shapes(lambda annotation: annotation.shape)
         return self._shapes
 
-    def foreground_shapes(self) -> tuple[tuple[int, ...] | None, tuple[int, ...], tuple[int, ...]]:
+    def foreground_shapes(self) -> tuple[AmbiguousShape | None, AmbiguousShape, AmbiguousShape]:
         if self._foreground_shapes:
             return self._foreground_shapes
         self._foreground_shapes = self._get_shapes(lambda annotation: annotation.foreground_shape())
         return self._foreground_shapes
 
-    def statistical_foreground_shape(self, *, percentile: float = .95) -> tuple[int, int] | tuple[int, int, int]:
+    def statistical_foreground_shape(self, *, percentile: float = .95) -> Shape:
         if self._statistical_foreground_shape:
             return self._statistical_foreground_shape
         depths, heights, widths = self.foreground_shapes()
@@ -172,7 +172,7 @@ def center_of_foregrounds_offsets(self) -> tuple[int, int] | tuple[int, int, int
         self._foreground_offsets = offsets + (round(center[2] - max_shape[2] * .5),) if depths else offsets
         return self._foreground_offsets
 
-    def set_roi_shape(self, roi_shape: tuple[int, int] | tuple[int, int, int] | None) -> None:
+    def set_roi_shape(self, roi_shape: Shape | None) -> None:
         if roi_shape is not None:
             depths, heights, widths = self.shapes()
             if depths:
@@ -183,7 +183,7 @@ def set_roi_shape(self, roi_shape: tuple[int, int] | tuple[int, int, int] | None
                     raise ValueError(f"ROI shape {roi_shape} exceeds minimum image shape ({min(heights)}, {min(widths)})")
         self._roi_shape = roi_shape
 
-    def roi_shape(self, *, percentile: float = .95) -> tuple[int, int] | tuple[int, int, int]:
+    def roi_shape(self, *, percentile: float = .95) -> Shape:
         if self._roi_shape:
             return self._roi_shape
         sfs = self.statistical_foreground_shape(percentile=percentile)

mipcandy/inference.py

@@ -11,7 +11,7 @@
 from mipcandy.data import save_image, Loader, UnsupervisedDataset, PathBasedUnsupervisedDataset
 from mipcandy.layer import WithPaddingModule, WithNetwork
 from mipcandy.sliding_window import SlidingWindow
-from mipcandy.types import SupportedPredictant, Device
+from mipcandy.types import SupportedPredictant, Device, AmbiguousShape
 
 
 def parse_predictant(x: SupportedPredictant, loader: type[Loader], *, as_label: bool = False) -> tuple[list[
@@ -40,20 +40,21 @@ def parse_predictant(x: SupportedPredictant, loader: type[Loader], *, as_label:
 
 
 class Predictor(WithPaddingModule, WithNetwork, metaclass=ABCMeta):
-    def __init__(self, experiment_folder: str | PathLike[str], example_shape: tuple[int, ...], *,
+    def __init__(self, experiment_folder: str | PathLike[str], example_shape: AmbiguousShape, *,
                  checkpoint: str = "checkpoint_best.pth", device: Device = "cpu") -> None:
         WithPaddingModule.__init__(self, device)
         WithNetwork.__init__(self, device)
         self._experiment_folder: str = experiment_folder
-        self._example_shape: tuple[int, ...] = example_shape
+        self._example_shape: AmbiguousShape = example_shape
         self._checkpoint: str = checkpoint
         self._model: nn.Module | None = None
 
     def lazy_load_model(self) -> None:
         if self._model:
             return
-        self._model = self.load_model(self._example_shape,
-                                      checkpoint=torch.load(f"{self._experiment_folder}/{self._checkpoint}"))
+        self._model = self.load_model(self._example_shape, checkpoint=torch.load(
+            f"{self._experiment_folder}/{self._checkpoint}"
+        ))
         self._model.eval()
 
     def predict_image(self, image: torch.Tensor, *, batch: bool = False) -> torch.Tensor:

mipcandy/layer.py

@@ -4,7 +4,7 @@
 import torch
 from torch import nn
 
-from mipcandy.types import Device
+from mipcandy.types import Device, AmbiguousShape
 
 
 def batch_int_multiply(f: float, *n: int) -> Generator[int, None, None]:
@@ -101,18 +101,18 @@ def __init__(self, device: Device) -> None:
         super().__init__(device)
 
     @abstractmethod
-    def build_network(self, example_shape: tuple[int, ...]) -> nn.Module:
+    def build_network(self, example_shape: AmbiguousShape) -> nn.Module:
         raise NotImplementedError
 
-    def build_network_from_checkpoint(self, example_shape: tuple[int, ...], checkpoint: Mapping[str, Any]) -> nn.Module:
+    def build_network_from_checkpoint(self, example_shape: AmbiguousShape, checkpoint: Mapping[str, Any]) -> nn.Module:
         """
         Internally exposed interface for overriding. Use `load_model()` instead.
         """
         network = self.build_network(example_shape)
         network.load_state_dict(checkpoint)
         return network
 
-    def load_model(self, example_shape: tuple[int, ...], *, checkpoint: Mapping[str, Any] | None = None) -> nn.Module:
+    def load_model(self, example_shape: AmbiguousShape, *, checkpoint: Mapping[str, Any] | None = None) -> nn.Module:
         if checkpoint:
             return self.build_network_from_checkpoint(example_shape, checkpoint).to(self._device)
         return self.build_network(example_shape).to(self._device)

mipcandy/presets/segmentation.py

@@ -8,7 +8,7 @@
 from mipcandy.data import visualize2d, visualize3d, overlay, auto_convert
 from mipcandy.sliding_window import SWMetadata
 from mipcandy.training import Trainer, TrainerToolbox, SlidingTrainer
-from mipcandy.types import Params
+from mipcandy.types import Params, Shape
 
 
 class SegmentationTrainer(Trainer, metaclass=ABCMeta):
@@ -65,7 +65,7 @@ def validate_case(self, image: torch.Tensor, label: torch.Tensor, toolbox: Train
 
 
 class SlidingSegmentationTrainer(SlidingTrainer, SegmentationTrainer, metaclass=ABCMeta):
-    sliding_window_shape: tuple[int, int] | tuple[int, int, int] = (128, 128)
+    sliding_window_shape: Shape = (128, 128)
 
     @override
     def backward_windowed(self, images: torch.Tensor, labels: torch.Tensor, toolbox: TrainerToolbox,
@@ -93,7 +93,7 @@ def validate_case_windowed(self, images: torch.Tensor, label: torch.Tensor, tool
         return -loss.item(), metrics, outputs.squeeze(0)
 
     @override
-    def get_window_shape(self) -> tuple[int, int] | tuple[int, int, int]:
+    def get_window_shape(self) -> Shape:
         return self.sliding_window_shape
 
 
@@ -118,5 +118,5 @@ def validate_case_windowed(self, images: torch.Tensor, label: torch.Tensor, tool
         return super().validate_case_windowed(images, label, toolbox, metadata)
 
     @override
-    def get_window_shape(self) -> tuple[int, int] | tuple[int, int, int]:
+    def get_window_shape(self) -> Shape:
         return self.sliding_window_shape

mipcandy/sliding_window.py

@@ -6,23 +6,24 @@
 from torch import nn
 
 from mipcandy.layer import HasDevice
+from mipcandy.types import Shape
 
 
 @dataclass
 class SWMetadata(object):
-    kernel: tuple[int, int] | tuple[int, int, int]
+    kernel: Shape
     stride: tuple[int, int] | tuple[int, int, int]
     ndim: Literal[2, 3]
     batch_size: int
-    out_size: tuple[int, int] | tuple[int, int, int]
+    out_size: Shape
     n: int
 
 
 class SlidingWindow(HasDevice, metaclass=ABCMeta):
     sliding_window_batch_size: int | None = None
 
     @abstractmethod
-    def get_window_shape(self) -> tuple[int, int] | tuple[int, int, int]:
+    def get_window_shape(self) -> Shape:
         raise NotImplementedError
 
     def get_batch_size(self) -> int | None:

mipcandy/training.py

@@ -27,7 +27,7 @@
 from mipcandy.layer import WithPaddingModule, WithNetwork
 from mipcandy.sanity_check import sanity_check
 from mipcandy.sliding_window import SWMetadata, SlidingWindow
-from mipcandy.types import Params, Setting
+from mipcandy.types import Params, Setting, AmbiguousShape
 
 
 def try_append(new: float, to: dict[str, list[float]], key: str) -> None:
@@ -101,7 +101,7 @@ def load_metrics(self) -> dict[str, list[float]]:
         df = read_csv(f"{self.experiment_folder()}/metrics.csv", index_col="epoch")
         return {column: df[column].astype(float).tolist() for column in df.columns}
 
-    def load_toolbox(self, num_epochs: int, example_shape: tuple[int, ...]) -> TrainerToolbox:
+    def load_toolbox(self, num_epochs: int, example_shape: AmbiguousShape) -> TrainerToolbox:
         toolbox = self._build_toolbox(num_epochs, example_shape, model=self.load_model(
             example_shape, checkpoint=torch.load(f"{self.experiment_folder()}/checkpoint_latest.pth")
         ))
@@ -330,7 +330,7 @@ def build_scheduler(self, optimizer: optim.Optimizer, num_epochs: int) -> optim.
     def build_criterion(self) -> nn.Module:
         raise NotImplementedError
 
-    def _build_toolbox(self, num_epochs: int, example_shape: tuple[int, ...], *,
+    def _build_toolbox(self, num_epochs: int, example_shape: AmbiguousShape, *,
                        model: nn.Module | None = None) -> TrainerToolbox:
         if not model:
             model = self.load_model(example_shape)
@@ -339,7 +339,7 @@ def _build_toolbox(self, num_epochs: int, example_shape: tuple[int, ...], *,
         criterion = self.build_criterion().to(self._device)
         return TrainerToolbox(model, optimizer, scheduler, criterion)
 
-    def build_toolbox(self, num_epochs: int, example_shape: tuple[int, ...]) -> TrainerToolbox:
+    def build_toolbox(self, num_epochs: int, example_shape: AmbiguousShape) -> TrainerToolbox:
         return self._build_toolbox(num_epochs, example_shape)
 
     # Training methods

mipcandy/types.py

@@ -12,3 +12,7 @@
 type SupportedPredictant = Sequence[torch.Tensor] | str | PathLike[str] | Sequence[str] | torch.Tensor
 type Colormap = Sequence[int | tuple[int, int, int]]
 type Device = torch.device | str
+type Shape2d = tuple[int, int]
+type Shape3d = tuple[int, int, int]
+type Shape = Shape2d | Shape3d
+type AmbiguousShape = tuple[int, ...]