minor bug fixes + support for slide semantic segmentation + minor improvements (#132)

* bug fix, default argument for segment_tissue throws an exception

* add num_workers arg to segment_tissue

* bug fix, display mpp for patcher vis if magnification is None

* don t install heavy libraries by default

* add ability to create WSIPatcher from CLAM type coordinates

* fix from_legacy_coords

* add coords_only arg to from_legacy

* add pil argument to from_legacy_coords

* pass slide_path as positional argument in CuCIMWSI, ImageWSI + fix src_pixel_size exception WSIPatcher

* add semantic segmentation inference

* add semantic segmentation inference

* test

* add forward_kwargs and create save_coords_h5

* cleanup before PR

* fix parameter propagation from ImageWSI to WSI

* cleanup before PR

* add README newline

* remove forward_kwargs

* remove get_channels

* fix: set correct wsi path in processor unit test

---------

Co-authored-by: Richard Chen <[email protected]>
Co-authored-by: guillaumejaume <[email protected]>

Author: 3 people

pyproject.toml

@@ -14,20 +14,22 @@ repository = "https://github.com/mahmoodlab/TRIDENT"
 python = "^3.10"  # Specify the Python version compatibility.
 ipywidgets = "*"
 torch = "*"
-transformers = "*"
 tqdm = "*"
 h5py = "*"
 matplotlib = "*"
-segmentation-models-pytorch = "*"
 opencv-python = "*"
 openslide-python = "*"
 Pillow = "*"
-timm = "0.9.16"
-einops_exts = "*"
 geopandas = "*"
 huggingface_hub = "*"
 openslide-bin = "*"
+
+# Optional dependencies (can be marked as optional in later versions)
+transformers = "*"
+timm = "0.9.16"
+einops_exts = "*"
 scipy = "*"
+segmentation-models-pytorch = "*"
 
 [tool.poetry.dev-dependencies]
 # Optional development dependencies 

tests/test_processor.py

@@ -70,7 +70,7 @@ def test_processor_with_wsis(self):
 
         self.processor = Processor(
             job_dir=self.TEST_OUTPUT_DIR,
-            wsi_source=os.path.join(TestProcessor.TEST_OUTPUT_DIR),
+            wsi_source=os.path.join(TestProcessor.TEST_OUTPUT_DIR, 'wsis'),
             wsi_ext=self.TEST_WSI_EXT,
             custom_list_of_wsis=os.path.join(self.custom_list_of_wsis, 'valid_list_of_wsis.csv')
         )

trident/IO.py

@@ -517,6 +517,40 @@ def read_coords_legacy(coords_path):
     return patch_size, patch_level, custom_downsample, coords
 
 
+def coords_to_h5(
+    coords: List[List[int]], 
+    save_path, 
+    patch_size, 
+    src_mag, 
+    target_mag, 
+    save_coords, 
+    width, 
+    height, 
+    name,
+    overlap
+):
+    """ Save tissue coordinates to .h5 """
+    # Prepare assets for saving
+    assets = {'coords' : np.array(coords)}
+    attributes = {
+        'patch_size': patch_size, # Reference frame: patch_level
+        'patch_size_level0': patch_size * src_mag // target_mag, # Reference frame: level0
+        'level0_magnification': src_mag,
+        'target_magnification': target_mag,
+        'overlap': overlap,
+        'name': name,
+        'savetodir': save_coords,
+        'level0_width': width,
+        'level0_height': height
+    }
+
+    # Save the assets and attributes to an hdf5 file
+    save_h5(save_path,
+            assets = assets,
+            attributes = {'coords': attributes},
+            mode='w')
+
+
 def mask_to_gdf(
     mask: np.ndarray,
     keep_ids: List[int] = [],

trident/Visualization.py

@@ -1,10 +1,10 @@
 import numpy as np
 import cv2
 import matplotlib.pyplot as plt
-from scipy.stats import rankdata
 from PIL import Image
-from typing import Optional, Tuple
+from typing import Optional, Tuple, Union
 import os 
+from shapely import Polygon, MultiPolygon
 
 
 def create_overlay(
@@ -73,6 +73,9 @@ def visualize_heatmap(
     normalize: bool = True,
     num_top_patches_to_save: int = -1,
     output_dir: Optional[str] = "output",
+    vis_mag: Optional[int] = None,
+    overlay_only = False,
+    filename = 'heatmap.png'
 ) -> str:
     """
     Generate a heatmap visualization overlayed on a whole slide image (WSI).
@@ -87,30 +90,45 @@ def visualize_heatmap(
         normalize (bool): Whether to normalize the scores.
         num_top_patches_to_save (int): Number of high-score patches to save. If set to -1, do not save any. Defaults to -1.
         output_dir (Optional[str]): Directory to save heatmap and top-k patches.
-    
+        vis_mag (Optional[int]): Visualization Magnification. This will overwrite vis_level
+        overlay_only bool: Whenever to save the overlay only. If set to True, save the overlay on top of downscaled version of the WSI. Defaults to False.
+        filename (str): file will be saved in `output_dir`/`filename`
+
     Returns:
         str: Path to the saved heatmap image.
     """
 
     if normalize:
+        from scipy.stats import rankdata
         scores = rankdata(scores, 'average') / len(scores) * 100 / 100
 
-    downsample = wsi.level_downsamples[vis_level]
+    if vis_mag is None:
+        downsample = wsi.level_downsamples[vis_level]
+    else:
+        src_mag = wsi.mag
+        downsample = src_mag / vis_mag
+        if not overlay_only:
+            vis_level, _ = wsi.get_best_level_and_custom_downsample(downsample)
+    
     scale = np.array([1 / downsample, 1 / downsample])
     region_size = tuple((np.array(wsi.level_dimensions[0]) * scale).astype(int))
-    
     overlay = create_overlay(scores, coords, patch_size_level0, scale, region_size)
+
+    overlay_colored = apply_colormap(overlay, cmap)
 
-    img = wsi.read_region((0, 0), vis_level, wsi.level_dimensions[vis_level]).convert("RGB")
-    img = img.resize(region_size, resample=Image.Resampling.BICUBIC)
-    img = np.array(img)
+    if overlay_only:
+        blended_img = overlay_colored
+    else:
+        img = wsi.read_region((0, 0), vis_level, wsi.level_dimensions[vis_level]).convert("RGB")
+        img = img.resize(region_size, resample=Image.Resampling.BICUBIC)
+        img = np.array(img)
+        
+        blended_img = cv2.addWeighted(img, 0.6, overlay_colored, 0.4, 0)
 
-    overlay_colored = apply_colormap(overlay, cmap)
-    blended_img = cv2.addWeighted(img, 0.6, overlay_colored, 0.4, 0)
     blended_img = Image.fromarray(blended_img)
 
     os.makedirs(output_dir, exist_ok=True)
-    heatmap_path = os.path.join(output_dir, "heatmap.png")
+    heatmap_path = os.path.join(output_dir, filename)
     blended_img.save(heatmap_path)
 
     if num_top_patches_to_save > 0:

trident/wsi_objects/CuCIMWSI.py

@@ -8,9 +8,28 @@
 
 class CuCIMWSI(WSI):
 
-    def __init__(self, **kwargs) -> None:
+    def __init__(self, slide_path, **kwargs) -> None:
+        """
+        Initialize a WSI instance using CuCIM as a backend.
+
+        Parameters
+        ----------
+        slide_path : str
+            Path to the WSI file.
+        **kwargs : dict
+            Keyword arguments forwarded to the base `WSI` class. Most important key is:
+            - lazy_init (bool, default=True): Whether to defer loading WSI and metadata.
+
+        Please refer to WSI constructor for all parameters. 
+
+        Example
+        -------
+        >>> wsi = CuCIMWSI(slide_path="path/to/wsi.svs", lazy_init=False)
+        >>> print(wsi)
+        <width=100000, height=80000, backend=CuCIMWSI, mpp=0.25, mag=40>
+        """
         self.img = None
-        super().__init__(**kwargs)
+        super().__init__(slide_path, **kwargs)
 
     def _lazy_initialize(self) -> None:
         """

trident/wsi_objects/ImageWSI.py

@@ -1,14 +1,14 @@
 from __future__ import annotations
 import numpy as np
 from PIL import Image
-from typing import Tuple, Union
+from typing import List, Tuple, Union
 
 from trident.wsi_objects.WSI import WSI, ReadMode
 
 
 class ImageWSI(WSI):
 
-    def __init__(self, **kwargs) -> None:
+    def __init__(self, slide_path, **kwargs) -> None:
         """
         Initialize a WSI object from a standard image file (e.g., PNG, JPEG, etc.).
 
@@ -30,7 +30,7 @@ def __init__(self, **kwargs) -> None:
 
         Example
         -------
-        >>> wsi = ImageWSI(slide_path="path/to/image.png", lazy_init=False, mpp=0.51)
+        >>> wsi = ImageWSI("path/to/image.png", lazy_init=False, mpp=0.51)
         >>> print(wsi)
         <width=5120, height=3840, backend=ImageWSI, mpp=0.51, mag=20>
         """
@@ -48,7 +48,7 @@ def __init__(self, **kwargs) -> None:
         PngImagePlugin.MAX_IMAGE_PIXELS = None  # Optional: disables large image warning
 
         self.img = None
-        super().__init__(**kwargs)
+        super().__init__(slide_path, **kwargs)
 
     def _lazy_initialize(self) -> None:
         """
@@ -177,28 +177,28 @@ def read_region(
         else:
             raise ValueError(f"Invalid `read_as` value: {read_as}. Must be 'pil' or 'numpy'.")
 
-    def segment_tissue(self, **kwargs):
-        out = super().segment_tissue(**kwargs)
+    def segment_tissue(self, *args, **kwargs):
+        out = super().segment_tissue(*args, **kwargs)
         self.close()
         return out
 
-    def extract_tissue_coords(self, **kwargs):
-        out = super().extract_tissue_coords(**kwargs)
+    def extract_tissue_coords(self, *args, **kwargs):
+        out = super().extract_tissue_coords(*args, **kwargs)
         self.close()
         return out
 
-    def visualize_coords(self, **kwargs):
-        out = super().visualize_coords(**kwargs)
+    def visualize_coords(self, *args, **kwargs):
+        out = super().visualize_coords(*args, **kwargs)
         self.close()
         return out
 
-    def extract_patch_features(self, **kwargs):
-        out = super().extract_patch_features(**kwargs)
+    def extract_patch_features(self, *args, **kwargs):
+        out = super().extract_patch_features(*args, **kwargs)
         self.close()
         return out
 
-    def extract_slide_features(self, **kwargs):
-        out = super().extract_slide_features(**kwargs)
+    def extract_slide_features(self, *args, **kwargs):
+        out = super().extract_slide_features(*args, **kwargs)
         self.close()
         return out
 

trident/wsi_objects/OpenSlideWSI.py

@@ -9,15 +9,16 @@
 
 class OpenSlideWSI(WSI):
 
-    def __init__(self, **kwargs) -> None:
+    def __init__(self, slide_path, **kwargs) -> None:
         """
         Initialize an OpenSlideWSI instance.
 
         Parameters
         ----------
+        slide_path : str
+            Path to the WSI file.
         **kwargs : dict
             Keyword arguments forwarded to the base `WSI` class. Most important key is:
-            - slide_path (str): Path to the WSI.
             - lazy_init (bool, default=True): Whether to defer loading WSI and metadata.
 
         Please refer to WSI constructor for all parameters. 
@@ -28,7 +29,7 @@ def __init__(self, **kwargs) -> None:
         >>> print(wsi)
         <width=100000, height=80000, backend=OpenSlideWSI, mpp=0.25, mag=40>
         """
-        super().__init__(**kwargs)
+        super().__init__(slide_path, **kwargs)
 
     def _lazy_initialize(self) -> None:
         """
@@ -244,4 +245,4 @@ def get_thumbnail(self, size: tuple[int, int]) -> Image.Image:
         PIL.Image.Image
             RGB thumbnail as a PIL Image.
         """
-        return self.img.get_thumbnail(size).convert('RGB')
+        return self.img.get_thumbnail(size).convert('RGB')