SupervisedDataset.fold() does not pass the desired types of arguments to construct_new()

mipcandy/data/dataset.py

@@ -168,12 +168,8 @@ def transform(self) -> JointTransform | None:
     def set_transform(self, transform: JointTransform | None) -> None:
         self._transform = transform.to(self._device) if transform else None
 
-    def _construct_new(self, images: D, labels: D) -> Self:
-        new = self.construct_new(images, labels)
-        return new
-
     @abstractmethod
-    def construct_new(self, images: D, labels: D) -> Self:
+    def construct_new(self, images: list[Any], labels: list[Any]) -> Self:
         raise NotImplementedError
 
     def preload(self, output_folder: str | PathLike[str], *, do_transform: bool = False) -> None:
@@ -207,7 +203,7 @@ def fold(self, *, fold: Literal[0, 1, 2, 3, 4, "all"] = "all", picker: type[KFPi
             else:
                 images_train.append(self._images[i])
                 labels_train.append(self._labels[i])
-        return self._construct_new(images_train, labels_train), self._construct_new(images_val, labels_val)
+        return self.construct_new(images_train, labels_train), self.construct_new(images_val, labels_val)
 
 
 class DatasetFromMemory(UnsupervisedDataset[Sequence[torch.Tensor]]):
@@ -234,7 +230,7 @@ def load_label(self, idx: int) -> torch.Tensor:
         return self._labels[idx]
 
     @override
-    def construct_new(self, images: UnsupervisedDataset, labels: UnsupervisedDataset) -> Self:
+    def construct_new(self, images: list[Any], labels: list[Any]) -> Self:
         return MergedDataset(DatasetFromMemory(images), DatasetFromMemory(labels), transform=self._transform,
                              device=self._device)
 
@@ -378,7 +374,7 @@ def save(self, split: str | Literal["Tr", "Ts"], *, target_folder: str | PathLik
         self._folded = False
 
     @override
-    def construct_new(self, images: list[str], labels: list[str]) -> Self:
+    def construct_new(self, images: list[Any], labels: list[Any]) -> Self:
         if self._folded:
             raise ValueError("Cannot construct a new dataset from a fold")
         new = self.__class__(self._folder, split=self._split, prefix=self._prefix, align_spacing=self._align_spacing,
@@ -401,7 +397,7 @@ def __len__(self) -> int:
         return len(self._base)
 
     @override
-    def construct_new(self, images: tuple[None], labels: tuple[None]) -> Self:
+    def construct_new(self, images: list[Any], labels: list[Any]) -> Self:
         raise NotImplementedError
 
     @override

mipcandy/data/inspection.py

@@ -322,7 +322,7 @@ def __init__(self, annotations: InspectionAnnotations, *, clamp: bool = True, pe
         self._percentile: float = percentile
 
     @override
-    def construct_new(self, images: list[int], labels: list[int]) -> Self:
+    def construct_new(self, images: list[Any], labels: list[Any]) -> Self:
         new = self.__class__(self._annotations, percentile=self._percentile)
         new._images = images
         new._labels = labels
@@ -388,7 +388,7 @@ def roi_shape(self, *, roi_shape: Shape | None = None) -> None | Shape:
         self._roi_shape = roi_shape
 
     @override
-    def construct_new(self, images: list[int], labels: list[int]) -> Self:
+    def construct_new(self, images: list[Any], labels: list[Any]) -> Self:
         new = self.__class__(self._annotations, self._batch_size, oversample_rate=self._oversample_rate,
                              clamp=self._clamp, percentile=self._percentile)
         new._images = images