fixed runtime issues in isolated inference mode

Author: drgmo

src/stamp/__main__.py

@@ -86,6 +86,7 @@ def _run_cli(args: argparse.Namespace) -> None:
                 tile_size_um=config.preprocessing.tile_size_um,
                 tile_size_px=config.preprocessing.tile_size_px,
                 extractor=config.preprocessing.extractor,
+                tile_extractor=config.preprocessing.tile_extractor,
                 max_workers=config.preprocessing.max_workers,
                 device=config.preprocessing.device,
                 default_slide_mpp=config.preprocessing.default_slide_mpp,

src/stamp/preprocessing/__init__.py

@@ -122,6 +122,7 @@ def extract_(
     cache_dir: Path | None,
     cache_tiles_ext: ImageExtension,
     extractor: ExtractorName | Extractor,
+    tile_extractor: ExtractorName,
     tile_size_px: TilePixels,
     tile_size_um: Microns,
     max_workers: int,
@@ -222,6 +223,11 @@ def extract_(
 
             extractor = plip()
 
+        case ExtractorName.TICON:
+            from stamp.preprocessing.extractor.ticon import ticon
+
+            extractor = ticon(tile_extractor=tile_extractor)
+
         case ExtractorName.EMPTY:
             from stamp.preprocessing.extractor.empty import empty
 
@@ -238,7 +244,8 @@ def extract_(
     code_hash = get_processing_code_hash(Path(__file__))[:8]
 
     extractor_id = extractor.identifier
-
+    if extractor_id == ExtractorName.TICON and tile_extractor is not None:
+        extractor_id = f"{extractor_id}-{tile_extractor}"
     _logger.info(f"Using extractor {extractor.identifier}")
 
     if cache_dir:
@@ -330,6 +337,8 @@ def extract_(
 
                 h5_fp.attrs["stamp_version"] = stamp.__version__
                 h5_fp.attrs["extractor"] = str(extractor.identifier)
+                if tile_extractor is not None:
+                    h5_fp.attrs["tile_extractor"] = str(tile_extractor)
                 h5_fp.attrs["unit"] = "um"
                 h5_fp.attrs["tile_size_um"] = tile_size_um  # changed in v2.1.0
                 h5_fp.attrs["tile_size_px"] = tile_size_px

src/stamp/preprocessing/config.py

@@ -45,7 +45,7 @@ class PreprocessingConfig(BaseModel, arbitrary_types_allowed=True):
     tile_size_um: Microns = Microns(256.0)
     tile_size_px: TilePixels = TilePixels(224)
     extractor: ExtractorName
-    tile_extractor: ExtractorName | None = None
+    tile_extractor: ExtractorName
     max_workers: int = 8
     device: str = "cuda" if torch.cuda.is_available() else "cpu"
     generate_hash: bool = True

…amp/preprocessing/extractor/ticon_iso.py …c/stamp/preprocessing/extractor/ticon.pysrc/stamp/preprocessing/extractor/ticon_iso.py renamed to src/stamp/preprocessing/extractor/ticon.py src/stamp/preprocessing/extractor/ticon_iso.py renamed to src/stamp/preprocessing/extractor/ticon.py

@@ -146,7 +146,7 @@ class TICON(nn.Module):
 
     def __init__(
         self,
-        tile_extractor: ExtractorName = ExtractorName.H_OPTIMUS_1,
+        tile_extractor: ExtractorName,
         device: str = "cuda",
     ):
         super().__init__()
@@ -221,13 +221,9 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
         return out.squeeze(1)
 
 
-def ticon_iso(
-    tile_extractor: ExtractorName = ExtractorName.H_OPTIMUS_1,
-    device: str = "cuda",
-) -> Extractor[TICON]:
+def ticon(tile_extractor: ExtractorName) -> Extractor[TICON]:
     """Create TICON Isolated Mode extractor."""
-    model = TICON(tile_extractor=tile_extractor, device=device)
-
+    model = TICON(tile_extractor=tile_extractor)
     return Extractor(
         model=model,
         transform=model.get_transform(),