slide2vec 1.1.0__py3-none-any.whl

slide2vec/__init__.py

@@ -0,0 +1 @@
+__version__ = "1.1.0"

slide2vec/configs/__init__.py

@@ -0,0 +1,17 @@
+import pathlib
+
+from omegaconf import OmegaConf
+
+
+def load_config(config_name: str):
+    config_filename = config_name + ".yaml"
+    return OmegaConf.load(pathlib.Path(__file__).parent.resolve() / config_filename)
+
+
+default_config = load_config("default")
+
+
+def load_and_merge_config(config_name: str):
+    default_config = OmegaConf.create(default_config)
+    loaded_config = load_config(config_name)
+    return OmegaConf.merge(default_config, loaded_config)

slide2vec/data/__init__.py

@@ -0,0 +1,2 @@
+from .dataset import TileDataset
+from .augmentations import RegionUnfolding

slide2vec/data/augmentations.py

@@ -0,0 +1,57 @@
+import torch
+import torchvision.transforms.functional as F
+
+from typing import Sequence
+from einops import rearrange
+from torchvision import transforms
+
+# Use timm's names
+IMAGENET_DEFAULT_MEAN = (0.485, 0.456, 0.406)
+IMAGENET_DEFAULT_STD = (0.229, 0.224, 0.225)
+
+
+def make_normalize_transform(
+    mean: Sequence[float] = IMAGENET_DEFAULT_MEAN,
+    std: Sequence[float] = IMAGENET_DEFAULT_STD,
+) -> transforms.Normalize:
+    return transforms.Normalize(mean=mean, std=std)
+
+
+class MaybeToTensor(transforms.ToTensor):
+    """
+    Convert a PIL Image or ndarray to tensor if it's not already one.
+    """
+
+    def __init__(self):
+        super().__init__()
+
+    def __call__(self, pic):
+        """
+        Args:
+            pic (PIL Image or numpy.ndarray): Image to be converted to tensor.
+
+        Returns:
+            Tensor: Converted image.
+        """
+        if isinstance(pic, torch.Tensor):
+            return pic
+        return F.to_tensor(pic)
+
+    def __repr__(self):
+        return f"{self.__class__.__name__}()"
+
+
+class RegionUnfolding:
+    def __init__(self, tile_size):
+        self.tile_size = tile_size
+
+    def __call__(self, x):
+        # x = [3, region_size, region_size]
+        # unfold into tilees and rearrange
+        x = x.unfold(1, self.tile_size, self.tile_size).unfold(
+            2, self.tile_size, self.tile_size
+        )  # [3, ntile, region_size, tile_size] -> [3, ntile, ntile, tile_size, tile_size]
+        x = rearrange(
+            x, "c p1 p2 w h -> (p1 p2) c w h"
+        )  # [num_tilees, 3, tile_size, tile_size]
+        return x

slide2vec/data/dataset.py

@@ -0,0 +1,62 @@
+import torch
+import numpy as np
+import wholeslidedata as wsd
+
+from PIL import Image
+from pathlib import Path
+
+
+class TileDataset(torch.utils.data.Dataset):
+    def __init__(self, wsi_path, tile_dir, target_spacing, backend, transforms=None):
+        self.path = wsi_path
+        self.target_spacing = target_spacing
+        self.backend = backend
+        self.name = wsi_path.stem.replace(" ", "_")
+        self.load_coordinates(tile_dir)
+        self.transforms = transforms
+
+    def load_coordinates(self, tile_dir):
+        coordinates = np.load(Path(tile_dir, f"{self.name}.npy"), allow_pickle=True)
+        self.x = coordinates["x"]
+        self.y = coordinates["y"]
+        self.coordinates = (np.array([self.x, self.y]).T).astype(int)
+        self.scaled_coordinates = self.scale_coordinates()
+        self.tile_level = coordinates["tile_level"]
+        self.tile_size_resized = coordinates["tile_size_resized"]
+        self.resize_factor = coordinates["resize_factor"]
+        self.tile_size = (self.tile_size_resized / self.resize_factor).astype(int)
+        self.tile_size_lv0 = coordinates["tile_size_lv0"][0]
+
+    def scale_coordinates(self):
+        # coordinates are defined w.r.t. level 0
+        # i need to scale them to target_spacing
+        wsi = wsd.WholeSlideImage(self.path, backend=self.backend)
+        min_spacing = wsi.spacings[0]
+        scale = min_spacing / self.target_spacing
+        # create a [N, 2] array with x and y coordinates
+        scaled_coordinates = (self.coordinates * scale).astype(int)
+        return scaled_coordinates
+
+    def __len__(self):
+        return len(self.x)
+
+    def __getitem__(self, idx):
+        wsi = wsd.WholeSlideImage(
+            self.path, backend=self.backend
+        )  # cannot be defined in __init__ because of multiprocessing
+        tile_level = self.tile_level[idx]
+        tile_spacing = wsi.spacings[tile_level]
+        tile_arr = wsi.get_patch(
+            self.x[idx],
+            self.y[idx],
+            self.tile_size_resized[idx],
+            self.tile_size_resized[idx],
+            spacing=tile_spacing,
+            center=False,
+        )
+        tile = Image.fromarray(tile_arr).convert("RGB")
+        if self.tile_size[idx] != self.tile_size_resized[idx]:
+            tile = tile.resize((self.tile_size[idx], self.tile_size[idx]))
+        if self.transforms:
+            tile = self.transforms(tile)
+        return idx, tile

slide2vec/distributed/__init__.py

@@ -0,0 +1,344 @@
+import os
+import random
+import re
+import socket
+import datetime
+from typing import Dict, List
+
+import torch
+import torch.distributed as dist
+
+_LOCAL_RANK = -1
+_LOCAL_WORLD_SIZE = -1
+
+
+def is_enabled() -> bool:
+    """
+    Returns:
+        True if distributed training is enabled
+    """
+    return dist.is_available() and dist.is_initialized()
+
+
+def is_enabled_and_multiple_gpus() -> bool:
+    """
+    Returns:
+        True if distributed training is enabled and there are multiple GPUs available.
+    """
+    return (
+        dist.is_available() and dist.is_initialized() and torch.cuda.device_count() > 1
+    )
+
+
+def get_global_size() -> int:
+    """
+    Returns:
+        The number of processes in the process group
+    """
+    return dist.get_world_size() if is_enabled() else 1
+
+
+def get_global_rank() -> int:
+    """
+    Returns:
+        The rank of the current process within the global process group.
+    """
+    return dist.get_rank() if is_enabled() else 0
+
+
+def get_local_rank() -> int:
+    """
+    Returns:
+        The rank of the current process within the local (per-machine) process group.
+    """
+    if not is_enabled():
+        return 0
+    assert 0 <= _LOCAL_RANK < _LOCAL_WORLD_SIZE
+    return _LOCAL_RANK
+
+
+def get_local_size() -> int:
+    """
+    Returns:
+        The size of the per-machine process group,
+        i.e. the number of processes per machine.
+    """
+    if not is_enabled():
+        return 1
+    assert 0 <= _LOCAL_RANK < _LOCAL_WORLD_SIZE
+    return _LOCAL_WORLD_SIZE
+
+
+def is_main_process() -> bool:
+    """
+    Returns:
+        True if the current process is the main one.
+    """
+    return get_global_rank() == 0
+
+
+def _restrict_print_to_main_process() -> None:
+    """
+    This function disables printing when not in the main process
+    """
+    import builtins as __builtin__
+
+    builtin_print = __builtin__.print
+
+    def print(*args, **kwargs):
+        force = kwargs.pop("force", False)
+        if is_main_process() or force:
+            builtin_print(*args, **kwargs)
+
+    __builtin__.print = print
+
+
+def _get_master_port(seed: int = 0) -> int:
+    MIN_MASTER_PORT, MAX_MASTER_PORT = (20_000, 60_000)
+
+    master_port_str = os.environ.get("MASTER_PORT")
+    if master_port_str is None:
+        rng = random.Random(seed)
+        return rng.randint(MIN_MASTER_PORT, MAX_MASTER_PORT)
+
+    return int(master_port_str)
+
+
+def _get_available_port() -> int:
+    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
+        # A "" host address means INADDR_ANY i.e. binding to all interfaces.
+        # Note this is not compatible with IPv6.
+        s.bind(("", 0))
+        port = s.getsockname()[1]
+        return port
+
+
+_TORCH_DISTRIBUTED_ENV_VARS = (
+    "MASTER_ADDR",
+    "MASTER_PORT",
+    "RANK",
+    "WORLD_SIZE",
+    "LOCAL_RANK",
+    "LOCAL_WORLD_SIZE",
+)
+
+
+def _collect_env_vars() -> Dict[str, str]:
+    return {
+        env_var: os.environ[env_var]
+        for env_var in _TORCH_DISTRIBUTED_ENV_VARS
+        if env_var in os.environ
+    }
+
+
+def _is_slurm_job_process() -> bool:
+    return "SLURM_JOB_ID" in os.environ
+
+
+def _parse_slurm_node_list(s: str) -> List[str]:
+    nodes = []
+    # Extract "hostname", "hostname[1-2,3,4-5]," substrings
+    p = re.compile(r"(([^\[]+)(?:\[([^\]]+)\])?),?")
+    for m in p.finditer(s):
+        prefix, suffixes = s[m.start(2) : m.end(2)], s[m.start(3) : m.end(3)]
+        for suffix in suffixes.split(","):
+            span = suffix.split("-")
+            if len(span) == 1:
+                nodes.append(prefix + suffix)
+            else:
+                width = len(span[0])
+                start, end = int(span[0]), int(span[1]) + 1
+                nodes.extend([prefix + f"{i:0{width}}" for i in range(start, end)])
+    return nodes
+
+
+def _check_env_variable(key: str, new_value: str):
+    # Only check for difference with preset environment variables
+    if key in os.environ and os.environ[key] != new_value:
+        raise RuntimeError(
+            f"Cannot export environment variables as {key} is already set"
+        )
+
+
+class _TorchDistributedEnvironment:
+    def __init__(self):
+        self.master_addr = "127.0.0.1"
+        self.master_port = 0
+        self.rank = -1
+        self.world_size = -1
+        self.local_rank = -1
+        self.local_world_size = -1
+
+        # if _is_slurm_job_process():
+        #     return self._set_from_slurm_env()
+
+        env_vars = _collect_env_vars()
+        if not env_vars:
+            # Environment is not set
+            pass
+        elif len(env_vars) == len(_TORCH_DISTRIBUTED_ENV_VARS):
+            # Environment is fully set
+            return self._set_from_preset_env()
+        else:
+            # Environment is partially set
+            collected_env_vars = ", ".join(env_vars.keys())
+            raise RuntimeError(f"Partially set environment: {collected_env_vars}")
+
+        if torch.cuda.device_count() > 0:
+            return self._set_from_local()
+
+        raise RuntimeError("Can't initialize PyTorch distributed environment")
+
+    # Slurm job created with sbatch, submitit, etc...
+    def _set_from_slurm_env(self):
+        # logger.info("Initialization from Slurm environment")
+        job_id = int(os.environ["SLURM_JOB_ID"])
+        node_count = int(os.environ["SLURM_JOB_NUM_NODES"])
+        nodes = _parse_slurm_node_list(os.environ["SLURM_JOB_NODELIST"])
+        assert len(nodes) == node_count
+
+        self.master_addr = nodes[0]
+        self.master_port = _get_master_port(seed=job_id)
+        self.rank = int(os.environ["SLURM_PROCID"])
+        self.world_size = int(os.environ["SLURM_NTASKS"])
+        assert self.rank < self.world_size
+        self.local_rank = int(os.environ["SLURM_LOCALID"])
+        self.local_world_size = self.world_size // node_count
+        assert self.local_rank < self.local_world_size
+
+    # Single node job with preset environment (i.e. torchrun)
+    def _set_from_preset_env(self):
+        # logger.info("Initialization from preset environment")
+        self.master_addr = os.environ["MASTER_ADDR"]
+        self.master_port = os.environ["MASTER_PORT"]
+        self.rank = int(os.environ["RANK"])
+        self.world_size = int(os.environ["WORLD_SIZE"])
+        assert self.rank < self.world_size
+        self.local_rank = int(os.environ["LOCAL_RANK"])
+        self.local_world_size = int(os.environ["LOCAL_WORLD_SIZE"])
+        assert self.local_rank < self.local_world_size
+
+    # Single node and GPU job (i.e. local script run)
+    def _set_from_local(self):
+        # logger.info("Initialization from local")
+        self.master_addr = "127.0.0.1"
+        self.master_port = _get_available_port()
+        self.rank = 0
+        self.world_size = 1
+        self.local_rank = 0
+        self.local_world_size = 1
+
+    def export(self, *, overwrite: bool) -> "_TorchDistributedEnvironment":
+        # See the "Environment variable initialization" section from
+        # https://pytorch.org/docs/stable/distributed.html for the complete list of
+        # environment variables required for the env:// initialization method.
+        env_vars = {
+            "MASTER_ADDR": self.master_addr,
+            "MASTER_PORT": str(self.master_port),
+            "RANK": str(self.rank),
+            "WORLD_SIZE": str(self.world_size),
+            "LOCAL_RANK": str(self.local_rank),
+            "LOCAL_WORLD_SIZE": str(self.local_world_size),
+        }
+        if not overwrite:
+            for k, v in env_vars.items():
+                _check_env_variable(k, v)
+
+        os.environ.update(env_vars)
+        return self
+
+
+def enable(
+    *,
+    set_cuda_current_device: bool = True,
+    overwrite: bool = False,
+    allow_nccl_timeout: bool = False,
+):
+    """Enable distributed mode
+
+    Args:
+        set_cuda_current_device: If True, call torch.cuda.set_device() to set the
+            current PyTorch CUDA device to the one matching the local rank.
+        overwrite: If True, overwrites already set variables. Else fails.
+    """
+
+    global _LOCAL_RANK, _LOCAL_WORLD_SIZE
+    if _LOCAL_RANK >= 0 or _LOCAL_WORLD_SIZE >= 0:
+        raise RuntimeError("Distributed mode has already been enabled")
+    torch_env = _TorchDistributedEnvironment()
+    torch_env.export(overwrite=overwrite)
+
+    if set_cuda_current_device:
+        torch.cuda.set_device(torch_env.local_rank)
+
+    if allow_nccl_timeout:
+        # This allows to use torch distributed timeout in a NCCL backend
+        key, value = "NCCL_ASYNC_ERROR_HANDLING", "1"
+        if not overwrite:
+            _check_env_variable(key, value)
+        os.environ[key] = value
+
+    dist.init_process_group(backend="nccl", timeout=datetime.timedelta(days=14))
+    dist.barrier()
+
+    # Finalize setup
+    _LOCAL_RANK = torch_env.local_rank
+    _LOCAL_WORLD_SIZE = torch_env.local_world_size
+    _restrict_print_to_main_process()
+
+
+def gather_tensor(t: torch.Tensor, dst_rank: int = 0):
+    """
+    Gather a tensor t from all ranks to dst_rank.
+    Returns a list of size world_size on dst_rank, otherwise an empty list on other ranks.
+    """
+    gather_list = []
+    if dist.get_rank() == dst_rank:
+        gather_list = [torch.empty_like(t) for _ in range(get_global_size())]
+    dist.gather(t, gather_list, dst=dst_rank)
+    return gather_list
+
+
+def gather_features(
+    features, indices, device, features_dim, level, scaled_coordinates=None
+):
+    gathered_feature = [
+        torch.zeros_like(features, device=device) for _ in range(get_global_size())
+    ]
+    gathered_indices = [
+        torch.zeros_like(indices, device=device) for _ in range(get_global_size())
+    ]
+    torch.distributed.all_gather(gathered_feature, features)
+    torch.distributed.all_gather(gathered_indices, indices)
+    # concatenate the gathered features and indices
+    wsi_feature = torch.cat(gathered_feature, dim=0)
+    tile_indices = torch.cat(gathered_indices, dim=0)
+    # handle duplicates
+    unique_indices, inverse_indices = torch.unique(
+        tile_indices, sorted=False, return_inverse=True
+    )
+    # create a final tensor to store the features in the correct order
+    coordinates = None
+    if level == "tile" or level == "slide":
+        wsi_feature_ordered = torch.zeros(
+            (unique_indices.size(0), features_dim), device=device
+        )
+        if level == "slide":
+            coordinates = torch.zeros(
+                (unique_indices.size(0), 2), dtype=torch.int64, device=device
+            )
+    elif level == "region":
+        num_tiles = features.shape[1]
+        wsi_feature_ordered = torch.zeros(
+            (unique_indices.size(0), num_tiles, features_dim), device=device
+        )
+    # insert each feature into its correct position based on tile_indices
+    for idx in inverse_indices:
+        unique_idx = unique_indices[idx]
+        wsi_feature_ordered[unique_idx] = wsi_feature[idx]
+        if level == "slide":
+            assert scaled_coordinates is not None
+            coordinates[unique_idx] = torch.tensor(
+                scaled_coordinates[idx], device=device
+            )
+    return wsi_feature_ordered, coordinates