Load features for classification

ahcore/backends.py

@@ -1,10 +1,14 @@
-from typing import Any
+from enum import Enum
+from typing import Any, Callable
 
 import pyvips
 from dlup.backends.common import AbstractSlideBackend
+from dlup.backends.openslide_backend import OpenSlideSlide
+from dlup.backends.pyvips_backend import PyVipsSlide
+from dlup.backends.tifffile_backend import TifffileSlide
 from dlup.types import PathLike  # type: ignore
 
-from ahcore.readers import StitchingMode, ZarrFileImageReader
+from ahcore.readers import H5FileImageReader, StitchingMode, ZarrFileImageReader
 
 
 class ZarrSlide(AbstractSlideBackend):
@@ -42,3 +46,53 @@ def read_region(self, coordinates: tuple[int, int], level: int, size: tuple[int,
 
     def close(self) -> None:
         self._reader.close()
+
+
+class H5Slide(AbstractSlideBackend):
+    def __init__(self, filename: PathLike, stitching_mode: StitchingMode | str = StitchingMode.CROP) -> None:
+        super().__init__(filename)
+        self._reader: H5FileImageReader = H5FileImageReader(filename, stitching_mode=stitching_mode)
+        self._spacings = [(self._reader.mpp, self._reader.mpp)]
+
+    @property
+    def size(self) -> tuple[int, int]:
+        return self._reader.size
+
+    @property
+    def level_dimensions(self) -> tuple[tuple[int, int], ...]:
+        return (self._reader.size,)
+
+    @property
+    def level_downsamples(self) -> tuple[float, ...]:
+        return (1.0,)
+
+    @property
+    def vendor(self) -> str:
+        return "H5FileImageReader"
+
+    @property
+    def properties(self) -> dict[str, Any]:
+        return self._reader.metadata
+
+    @property
+    def magnification(self) -> None:
+        return None
+
+    def read_region(self, coordinates: tuple[int, int], level: int, size: tuple[int, int]) -> pyvips.Image:
+        return self._reader.read_region(coordinates, level, size)
+
+    def close(self) -> None:
+        self._reader.close()
+
+
+class ImageBackend(Enum):
+    """Available image backends."""
+
+    OPENSLIDE: Callable[[PathLike], OpenSlideSlide] = OpenSlideSlide
+    PYVIPS: Callable[[PathLike], PyVipsSlide] = PyVipsSlide
+    TIFFFILE: Callable[[PathLike], TifffileSlide] = TifffileSlide
+    H5: Callable[[PathLike], H5Slide] = H5Slide
+    ZARR: Callable[[PathLike], ZarrSlide] = ZarrSlide
+
+    def __call__(self, *args: Any) -> OpenSlideSlide | PyVipsSlide | TifffileSlide | H5Slide | ZarrSlide:
+        return self.value(*args)  # type: ignore

ahcore/callbacks/file_writer_callback.py

@@ -4,14 +4,15 @@
 from typing import Type
 
 from dlup.data.dataset import TiledWsiDataset
+from dlup.tiling import Grid
 
 from ahcore.callbacks.abstract_writer_callback import AbstractWriterCallback
 from ahcore.callbacks.converters.common import ConvertCallbacks
 from ahcore.lit_module import AhCoreLightningModule
 from ahcore.utils.callbacks import get_output_filename as get_output_filename_
 from ahcore.utils.data import DataDescription, GridDescription
 from ahcore.utils.io import get_logger
-from ahcore.utils.types import InferencePrecision, NormalizationType
+from ahcore.utils.types import DataFormat, InferencePrecision, NormalizationType
 from ahcore.writers import Writer
 
 logger = get_logger(__name__)
@@ -27,7 +28,8 @@ def __init__(
         normalization_type: str = NormalizationType.LOGITS,
         precision: str = InferencePrecision.FP32,
         callbacks: list[ConvertCallbacks] | None = None,
-    ):
+        data_format: str = DataFormat.IMAGE,
+    ) -> None:
         """
         Callback to write predictions to H5 files. This callback is used to write whole-slide predictions to single H5
         files in a separate thread.
@@ -54,6 +56,7 @@ def __init__(
         self._suffix = ".cache"
         self._normalization_type: NormalizationType = NormalizationType(normalization_type)
         self._precision: InferencePrecision = InferencePrecision(precision)
+        self._data_format = DataFormat(data_format)
 
         super().__init__(
             writer_class=writer_class,
@@ -97,41 +100,64 @@ def build_writer_class(self, pl_module: AhCoreLightningModule, stage: str, filen
         with open(link_fn, "a" if link_fn.is_file() else "w") as file:
             file.write(f"{filename},{output_filename}\n")
 
-        current_dataset: TiledWsiDataset
-        current_dataset, _ = self._total_dataset.index_to_dataset(self._dataset_index)  # type: ignore
-        slide_image = current_dataset.slide_image
-        num_samples = len(current_dataset)
-
-        data_description: DataDescription = pl_module.data_description
-        inference_grid: GridDescription = data_description.inference_grid
-
-        mpp = inference_grid.mpp
-        if mpp is None:
-            mpp = slide_image.mpp
-
-        _, size = slide_image.get_scaled_slide_bounds(slide_image.get_scaling(mpp))
-
-        # Let's get the data_description, so we can figure out the tile size and things like that
-        tile_size = inference_grid.tile_size
-        tile_overlap = inference_grid.tile_overlap
-
-        if stage == "validate":
-            grid = current_dataset._grids[0][0]  # pylint: disable=protected-access
-        else:
-            grid = None  # During inference we don't have a grid around ROI
+        size, mpp, tile_size, tile_overlap, num_samples, grid = self._get_writer_data_args(
+            pl_module, data_format=self._data_format, stage=stage
+        )
 
         writer = self._writer_class(
             output_filename,
-            size=size,
+            size=size,  # --> (num_samples,1)
             mpp=mpp,
-            tile_size=tile_size,
+            tile_size=tile_size,  # --> (1,1)
             tile_overlap=tile_overlap,
             num_samples=num_samples,
             color_profile=None,
-            is_compressed_image=False,
+            data_format=self._data_format,
             progress=None,
             precision=InferencePrecision(self._precision),
             grid=grid,
         )
 
         return writer
+
+    def _get_writer_data_args(
+        self, pl_module: AhCoreLightningModule, data_format: DataFormat, stage: str
+    ) -> tuple[tuple[int, int], float, tuple[int, int], tuple[int, int], int, Grid | None]:
+        current_dataset: TiledWsiDataset
+        current_dataset, _ = self._total_dataset.index_to_dataset(self._dataset_index)  # type: ignore
+        slide_image = current_dataset.slide_image
+        num_samples = len(current_dataset)
+
+        if data_format == DataFormat.IMAGE or data_format == DataFormat.COMPRESSED_IMAGE:
+            data_description: DataDescription = pl_module.data_description
+            if data_description.inference_grid is None:
+                raise ValueError("Inference grid is not defined in the data description.")
+            inference_grid: GridDescription = data_description.inference_grid
+
+            mpp = inference_grid.mpp
+            if mpp is None:
+                mpp = slide_image.mpp
+
+            _, size = slide_image.get_scaled_slide_bounds(slide_image.get_scaling(mpp))
+
+            # Let's get the data_description, so we can figure out the tile size and things like that
+            tile_size = inference_grid.tile_size
+            tile_overlap = inference_grid.tile_overlap
+
+            if stage == "validate":
+                grid = current_dataset._grids[0][0]  # pylint: disable=protected-access
+            else:
+                grid = None  # During inference we don't have a grid around ROI
+
+        elif data_format == DataFormat.FEATURE:
+            size = (num_samples, 1)
+            mpp = 1.0
+            tile_size = (1, 1)
+            tile_overlap = (0, 0)
+            num_samples = num_samples
+            grid = current_dataset._grids[0][0]  # give grid, bc doesn't work otherwise
+
+        else:
+            raise NotImplementedError(f"Data format {data_format} is not yet supported.")
+
+        return size, mpp, tile_size, tile_overlap, num_samples, grid

ahcore/cli/tiling.py

@@ -33,6 +33,7 @@
 from rich.progress import Progress
 
 from ahcore.cli import dir_path, file_path
+from ahcore.utils.types import DataFormat
 from ahcore.writers import H5FileImageWriter, Writer, ZarrFileImageWriter
 
 _WriterClass = Type[Writer]
@@ -363,7 +364,7 @@ def _tiling_pipeline(
             tile_size=dataset_cfg.tile_size,
             tile_overlap=dataset_cfg.tile_overlap,
             num_samples=len(dataset),
-            is_compressed_image=compression != "none",
+            data_format=DataFormat.IMAGE,
             color_profile=color_profile,
             extra_metadata=extra_metadata,
             grid=dataset.grids[0][0],

ahcore/data/dataset.py

@@ -119,6 +119,7 @@ def __init__(
         num_workers: int = 16,
         persistent_workers: bool = False,
         pin_memory: bool = False,
+        use_cache: bool = True,
     ) -> None:
         """
         Construct a DataModule based on a manifest.
@@ -178,6 +179,7 @@ def __init__(
         self._num_workers = num_workers
         self._persistent_workers = persistent_workers
         self._pin_memory = pin_memory
+        self._use_cache = use_cache
 
         self._fit_data_iterator: Iterator[_DlupDataset] | None = None
         self._validate_data_iterator: Iterator[_DlupDataset] | None = None
@@ -245,7 +247,7 @@ def construct_dataset() -> ConcatDataset:
             return ConcatDataset(datasets=datasets)
 
         self._logger.info("Constructing dataset for stage %s (this can take a while)", stage)
-        dataset = self._load_from_cache(construct_dataset, stage=stage)
+        dataset = self._load_from_cache(construct_dataset, stage=stage) if self._use_cache else construct_dataset()
         setattr(self, f"{stage}_dataset", dataset)
 
         lengths = np.asarray([len(ds) for ds in dataset.datasets])
@@ -365,4 +367,6 @@ def uuid(self) -> uuid_module.UUID:
         str
             A unique identifier for this datamodule.
         """
+
+        # todo: It doesn't take into account different types of pretransforms, which can be important.
         return basemodel_to_uuid(self.data_description)

ahcore/lit_module.py

@@ -36,9 +36,11 @@ class AhCoreLightningModule(pl.LightningModule):
         "grid_index",
     ]
 
+    _model: nn.Module | BaseAhcoreJitModel
+
     def __init__(
         self,
-        model: nn.Module | BaseAhcoreJitModel,
+        model: nn.Module | BaseAhcoreJitModel | functools.partial[nn.Module],
         optimizer: torch.optim.optimizer.Optimizer,  # noqa
         data_description: DataDescription,
         loss: nn.Module | None = None,
@@ -66,6 +68,8 @@ def __init__(
             except AttributeError:
                 raise AttributeError("num_classes must be specified in data_description")
             self._model = model(out_channels=self._num_classes)
+        elif isinstance(model, nn.Module):
+            self._model = model
         else:
             raise TypeError(f"The class of models: {model.__class__} is not supported on ahcore")
         self._augmentations = augmentations

ahcore/losses.py

@@ -67,7 +67,7 @@ def __init__(
         if class_proportions is not None:
             _class_weights = 1 / class_proportions
             _class_weights[_class_weights.isnan()] = 0.0
-            _class_weights = _class_weights / _class_weights.max()
+            _class_weights = _class_weights / _class_weights.max()  # todo: check, shouldn't this be .sum?
             self._class_weights = _class_weights
         else:
             self._class_weights = None
@@ -141,26 +141,34 @@ def cross_entropy(
     else:
         roi_sum = torch.tensor([np.prod(tuple(input.shape)[2:])]).to(input.device)
 
+    if input.dim() != target.dim():
+        raise ValueError(f"Dimension do not match for input and target. Got {input.dim()} and {target.dim()}")
+
+    if input.dim() == 2 and target.dim() == 2:
+        # handle cls task as an image of size 1x1
+        input = input.unsqueeze(-1).unsqueeze(-1)
+        target = target.unsqueeze(-1).unsqueeze(-1)
+
     if ignore_index is None:
         ignore_index = -100
 
     # compute cross_entropy pixel by pixel
-    if not multiclass:
-        _cross_entropy = F.cross_entropy(
+    if multiclass:
+        _cross_entropy = F.binary_cross_entropy_with_logits(
             input,
-            target.argmax(dim=1),
-            ignore_index=ignore_index,
+            target,
             weight=None if weight is None else weight.to(input.device),
             reduction="none",
-            label_smoothing=label_smoothing,
+            pos_weight=None,
         )
     else:
-        _cross_entropy = F.binary_cross_entropy_with_logits(
+        _cross_entropy = F.cross_entropy(
             input,
-            target,
+            target.argmax(dim=1),
+            ignore_index=ignore_index,
             weight=None if weight is None else weight.to(input.device),
             reduction="none",
-            pos_weight=None,
+            label_smoothing=label_smoothing,
         )
 
     if limit is not None: