Merge branch 'main' of github.com:NKI-AI/ahcore

ahcore/cli/tiling.py

@@ -185,12 +185,12 @@ def _generator(
     dataset: TiledWsiDataset, quality: int | None = 80, compression: str = "JPEG"
 ) -> Generator[Any, Any, Any]:
     for idx in range(len(dataset)):
-        # TODO: To use:
-        # for idx, sample in enumerate(dataset):
-        # The following needs to be added to TiledWsiDataset:
-        # def __iter__(self) -> Iterator[RegionFromWsiDatasetSample]:
-        #     for i in range(len(self)):
-        #         yield self[i]
+    # TODO: To use:
+    # for idx, sample in enumerate(dataset):
+    # The following needs to be added to TiledWsiDataset:
+    # def __iter__(self) -> Iterator[RegionFromWsiDatasetSample]:
+    #     for i in range(len(self)):
+    #         yield self[i]
         sample = dataset[idx]
         buffered = io.BytesIO()
         if quality is not None:

ahcore/lit_module.py

@@ -5,7 +5,7 @@
 """
 from __future__ import annotations
 
-from typing import Any, Literal, MutableMapping, Optional, Type, TypedDict
+from typing import Any, Callable
 
 import pytorch_lightning as pl
 import torch.optim.optimizer
@@ -18,9 +18,6 @@
 from ahcore.utils.data import DataDescription
 from ahcore.utils.io import get_logger
 from ahcore.utils.types import DlupDatasetSample
-
-LitModuleSample = dict[str, Any]
-
 logger = get_logger(__name__)
 
 
@@ -88,7 +85,7 @@ def wsi_metrics(self) -> WSIMetricFactory | None:
     def name(self) -> str:
         return str(self._model.__class__.__name__)
 
-    def forward(self, sample: torch.Tensor) -> Any:
+    def forward(self, sample: DlupDatasetSample) -> DlupDatasetSample:
         """This function is only used during inference"""
         self._model.eval()
         return self._model.forward(sample)
@@ -206,9 +203,9 @@ def predict_step(self, batch: Any, batch_idx: int, dataloader_idx: int = 0) -> A
             raise ValueError("Filenames are not constant across the batch.")
         return output
 
-    def configure_optimizers(self) -> Any:
+    def configure_optimizers(self):
         optimizer = self.hparams.optimizer(params=self.parameters())  # type: ignore
-        if self.hparams.scheduler is not None:  # type: ignore
+        if self.hparams.scheduler is not None: # type: ignore
             scheduler = self.hparams.scheduler(optimizer=optimizer)  # type: ignore
             return {
                 "optimizer": optimizer,

ahcore/transforms/pre_transforms.py

@@ -3,7 +3,7 @@
 dataset."""
 from __future__ import annotations
 
-from typing import Callable
+from typing import Callable, Any
 
 import numpy as np
 import numpy.typing as npt
@@ -15,8 +15,9 @@
 from ahcore.utils.data import DataDescription
 from ahcore.utils.io import get_logger
 from ahcore.utils.types import DlupDatasetSample
+from dlup.data.dataset import TileSample, TileSampleWithAnnotationData
 
-PreTransformCallable = Callable[[DlupDatasetSample], DlupDatasetSample]
+PreTransformCallable = Callable[[TileSample], TileSampleWithAnnotationData]
 
 logger = get_logger(__name__)
 
@@ -189,7 +190,7 @@ class AllowCollate:
     This transform converts the path to a string. Same holds for the annotations and labels
     """
 
-    def __call__(self, sample: DlupDatasetSample) -> DlupDatasetSample:
+    def __call__(self, sample: TileSample) -> dict[str, Any]:
         # Path objects cannot be collated
         sample["path"] = str(sample["path"])
 

ahcore/utils/manifest.py

@@ -11,8 +11,9 @@
 from typing import Any, Callable, Generator, Literal, Optional, Type, TypedDict, cast
 
 from dlup import SlideImage
+from dlup.data.dataset import TileSample, RegionFromWsiDatasetSample
 from dlup.annotations import WsiAnnotations
-from dlup.data.dataset import RegionFromWsiDatasetSample, TiledWsiDataset, TileSample
+from dlup.data.dataset import TiledWsiDataset
 from dlup.experimental_backends import ImageBackend  # type: ignore
 from dlup.tiling import GridOrder, TilingMode
 from pydantic import BaseModel