datamint 1.2.4__py3-none-any.whl

datamintapi/utils/io_utils.py

@@ -0,0 +1,149 @@
+import numpy as np
+import nibabel as nib
+from PIL import Image
+from .dicom_utils import load_image_normalized, is_dicom
+import pydicom
+import os
+from typing import Any
+import logging
+from PIL import ImageFile
+import cv2
+
+ImageFile.LOAD_TRUNCATED_IMAGES = True
+
+_LOGGER = logging.getLogger(__name__)
+
+IMAGE_EXTS = ('.png', '.jpg', '.jpeg')
+NII_EXTS = ('.nii', '.nii.gz')
+VIDEO_EXTS = ('.mp4', '.avi', '.mov', '.mkv')
+
+
+def read_video(file_path: str, index: int = None) -> np.ndarray:
+    cap = cv2.VideoCapture(file_path)
+    if not cap.isOpened():
+        raise ValueError(f"Failed to open video file: {file_path}")
+    try:
+        if index is None:
+            frames = []
+            while True:
+                ret, frame = cap.read()
+                if not ret:
+                    break
+                # Convert BGR to RGB and transpose to (C, H, W) format
+                frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
+                frame = frame.transpose(2, 0, 1)
+                frames.append(frame)
+            imgs = np.array(frames)  # shape: (#frames, C, H, W)
+        else:
+            while index > 0:
+                cap.grab()
+                index -= 1
+            ret, frame = cap.read()
+            if not ret:
+                raise ValueError(f"Failed to read frame {index} from video file: {file_path}")
+            # Convert BGR to RGB and transpose to (C, H, W) format
+            frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
+            imgs = frame.transpose(2, 0, 1)
+    finally:
+        cap.release()
+
+    if imgs is None or len(imgs) == 0:
+        raise ValueError(f"No frames found in video file: {file_path}")
+
+    return imgs
+
+
+def read_nifti(file_path: str) -> np.ndarray:
+    imgs = nib.load(file_path).get_fdata()  # shape: (W, H, #frame) or (W, H)
+    if imgs.ndim == 2:
+        imgs = imgs.transpose(1, 0)
+        imgs = imgs[np.newaxis, np.newaxis]
+    elif imgs.ndim == 3:
+        imgs = imgs.transpose(2, 1, 0)
+        imgs = imgs[:, np.newaxis]
+    else:
+        raise ValueError(f"Unsupported number of dimensions in '{file_path}': {imgs.ndim}")
+
+    return imgs
+
+
+def read_image(file_path: str) -> np.ndarray:
+    with Image.open(file_path) as pilimg:
+        imgs = np.array(pilimg)
+    if imgs.ndim == 2:  # (H, W)
+        imgs = imgs[np.newaxis, np.newaxis]
+    elif imgs.ndim == 3:  # (H, W, C)
+        imgs = imgs.transpose(2, 0, 1)[np.newaxis]  # (H, W, C) -> (1, C, H, W)
+
+    return imgs
+
+
+def read_array_normalized(file_path: str,
+                          index: int = None,
+                          return_metainfo: bool = False,
+                          use_magic=False) -> np.ndarray | tuple[np.ndarray, Any]:
+    """
+    Read an array from a file.
+
+    Args:
+        file_path: The path to the file.
+        Supported file formats are NIfTI (.nii, .nii.gz), PNG (.png), JPEG (.jpg, .jpeg) and npy (.npy).
+
+    Returns:
+        The array read from the file in shape (#frames, C, H, W), if `index=None`,
+            or (C, H, W) if `index` is specified.
+    """
+    if not os.path.exists(file_path):
+        raise FileNotFoundError(f"File not found: {file_path}")
+
+    metainfo = None
+
+    try:
+        if is_dicom(file_path):
+            ds = pydicom.dcmread(file_path)
+            if index is not None:
+                imgs = load_image_normalized(ds, index=index)[0]
+            else:
+                imgs = load_image_normalized(ds)
+            # Free up memory
+            if hasattr(ds, '_pixel_array'):
+                ds._pixel_array = None
+            if hasattr(ds, 'PixelData'):
+                ds.PixelData = None
+            metainfo = ds
+        else:
+            if use_magic:
+                import magic  # it is important to import here because magic has an OS lib dependency.
+                mime_type = magic.from_file(file_path, mime=True)
+            else:
+                mime_type = ""
+
+            if mime_type.startswith('video/') or file_path.endswith(VIDEO_EXTS):
+                imgs = read_video(file_path, index)
+            else:
+                if mime_type == 'image/x.nifti' or file_path.endswith(NII_EXTS):
+                    imgs = read_nifti(file_path)
+                elif mime_type.startswith('image/') or file_path.endswith(IMAGE_EXTS):
+                    imgs = read_image(file_path)
+                elif file_path.endswith('.npy') or mime_type == 'application/x-numpy-data':
+                    imgs = np.load(file_path)
+                    if imgs.ndim != 4:
+                        raise ValueError(f"Unsupported number of dimensions in '{file_path}': {imgs.ndim}")
+                else:
+                    raise ValueError(f"Unsupported file format '{mime_type}' of '{file_path}'")
+
+                if index is not None:
+                    if len(imgs) > 1:
+                        _LOGGER.warning(f"It is inefficient to load all frames from '{file_path}' to access a single frame." +
+                                        " Consider converting the file to a format that supports random access (DICOM), or" +
+                                        " convert to png/jpeg files or" +
+                                        " manually handle all frames at once instead of loading a specific frame.")
+                    imgs = imgs[index]
+
+        if return_metainfo:
+            return imgs, metainfo
+        return imgs
+
+    except Exception as e:
+        _LOGGER.error(f"Failed to read array from '{file_path}': {e}")
+        raise e

datamintapi/utils/logging_utils.py

@@ -0,0 +1,55 @@
+import logging
+import logging.config
+from rich.console import ConsoleRenderable
+from rich.logging import RichHandler
+from rich.traceback import Traceback
+import yaml
+import importlib
+
+_LOGGER = logging.getLogger(__name__)
+
+
+class ConditionalRichHandler(RichHandler):
+    """
+    Class that uses 'show_level=True' only if the message level is WARNING or higher.
+    """
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+
+    def handle(self, record):
+        if record.levelno >= logging.WARNING:
+            self.show_level = True
+        else:
+            self.show_level = False
+        super().handle(record)
+
+    def render(self, *, record: logging.LogRecord,
+               traceback: Traceback | None,
+               message_renderable: ConsoleRenderable) -> ConsoleRenderable:
+        # if level is WARNING or higher, add the level column
+        try:
+            self._log_render.show_level = record.levelno >= logging.WARNING
+            ret = super().render(record=record, traceback=traceback, message_renderable=message_renderable)
+            self._log_render.show_level = False
+        except Exception as e:
+            _LOGGER.error(f"Error rendering log. {e}")
+        return ret
+
+
+def load_cmdline_logging_config():
+    # Load the logging configuration file
+    try:
+        try:
+            # try loading the developer's logging config
+            with open('logging_dev.yaml', 'r') as f:
+                config = yaml.safe_load(f)
+        except:
+            with importlib.resources.open_text('datamintapi', 'logging.yaml') as f:
+                config = yaml.safe_load(f.read())
+
+        logging.config.dictConfig(config)
+    except Exception as e:
+        print(f"Warning: Error loading logging configuration file: {e}")
+        _LOGGER.exception(e)
+        logging.basicConfig(level=logging.INFO)

datamintapi/utils/torchmetrics.py

@@ -0,0 +1,70 @@
+from torchmetrics.classification import Recall, Precision, F1Score, Specificity
+from torchmetrics.wrappers.abstract import WrapperMetric
+import torchmetrics
+from torch import Tensor
+
+
+class SegmentationToClassificationWrapper(WrapperMetric):
+    """
+    Enables applying classification metrics to segmentation masks.
+    The segmentation masks are converted to binary one-hot enconding vectors using a IoU threshold.
+
+    Args:
+        metric_cls: Segmentation metric class.
+        iou_threshold: IoU threshold to convert segmentation to binary classification.
+    """
+
+    def __init__(self,
+                 metric_cls: torchmetrics.Metric,
+                 iou_threshold=0.5):
+        super().__init__()
+        self.metric = metric_cls
+        self.iou_threshold = iou_threshold
+
+    def update(self, preds: Tensor, target: Tensor):
+        cls_pred, cls_target = self.transform_mask_to_binary(preds, target, self.iou_threshold)
+        self.metric.update(cls_pred, cls_target)
+
+    def compute(self):
+        return self.metric.compute()
+
+    def reset(self):
+        super().reset()
+        self.metric.reset()
+
+    def forward(self, preds: Tensor, target: Tensor):
+        cls_pred, cls_target = self.transform_mask_to_binary(preds, target, self.iou_threshold)
+        return self.metric.forward(cls_pred, cls_target)
+
+    @staticmethod
+    def transform_mask_to_binary(pred: Tensor, target: Tensor, iou_threshold: float = 0.5) -> tuple[Tensor, Tensor]:
+        """
+        Convert both the segmentation masks with shape (B,L,H,W) to shape (B,L), so that classification metrics can be used.
+        The conversion is done by applying using a IoU threshold that must be satisfied to consider a prediction as true positive.
+        Args:
+            pred: Segmentation prediction of shape (B, L, H, W).
+            target: Segmentation target of shape (B, L, H, W).
+            iou_threshold: IoU threshold to convert segmentation to binary classification.
+        Returns:
+            Tuple of binary classification predictions (first) and targets (second) of shape (B, L).
+        """
+
+        if pred.ndim != 4 or target.ndim != 4:
+            raise ValueError("Input tensors must have 4 dimensions (B, L, H, W).")
+
+        # Calculate IoU for each sample
+        intersection = (pred & target).float().sum(dim=(2, 3))
+        union = (pred | target).float().sum(dim=(2, 3))
+        iou = intersection / union
+
+        cls_target = target.amax(dim=(2, 3))
+        cls_pred = (iou >= iou_threshold).float()
+
+        return cls_pred, cls_target
+
+# # test
+# cls_metric = Recall(task="multilabel", average="macro", num_labels=3)
+# metric = SegmentationToClassificationWrapper(cls_metric, iou_threshold=0.33)
+# metric(preds=torch.randint(0, 2, size=(4, 3, 32, 32)),
+#        target=torch.randint(0, 2, size=(4, 3, 32, 32))
+#        )

datamintapi/utils/visualization.py

@@ -0,0 +1,129 @@
+import matplotlib.pyplot as plt
+import numpy as np
+from torchvision.transforms import functional as F
+from torch import Tensor
+import torchvision.utils
+import torch
+import colorsys
+
+
+def show(imgs: list[Tensor] | Tensor,
+         figsize: tuple[int, int] = None,
+         normalize: bool = False):
+    """
+    Show a list of images in a grid.
+    Args:
+        imgs (list[Tensor] | Tensor): List of images to show.
+            Each image should be a tensor of shape (C, H, W) and dtype uint8 or float.
+        figsize (tuple[int, int], optional): Size of the figure. Defaults to None.
+        normalize (bool, optional): Whether to normalize the images to [0, 1] range by min-max scaling.
+    """
+
+    if not isinstance(imgs, list):
+        imgs = [imgs]
+
+    if normalize:
+        for i, img in enumerate(imgs):
+            img = img.float()
+            img = img - img.min()
+            img = img / img.max()
+            imgs[i] = img
+
+    if figsize is not None:
+        fig, axs = plt.subplots(ncols=len(imgs), squeeze=False, figsize=figsize)
+    else:
+        fig, axs = plt.subplots(ncols=len(imgs), squeeze=False)
+    for i, img in enumerate(imgs):
+        img = img.detach()
+        img = F.to_pil_image(img)
+        axs[0, i].imshow(np.asarray(img))
+        axs[0, i].set(xticklabels=[], yticklabels=[], xticks=[], yticks=[])
+
+
+_COLORS = [(255, 0, 0), (0, 255, 0), (0, 0, 255),
+           (200, 128, 0), (128, 200, 0), (0, 128, 200), (200, 0, 128), (128, 0, 200), (0, 200, 128)]
+
+
+def generate_color_palette(num_objects: int) -> list[tuple[int, int, int]]:
+    """
+    Generate a list of colors for segmentation masks.
+
+    Args:
+        num_objects (int): Number of objects to generate colors for.
+
+    Returns:
+        List of RGB colors.
+    """
+    if num_objects <= 0:
+        raise ValueError("Number of objects must be greater than 0.")
+
+    colors = _COLORS[:num_objects]
+    if len(colors) == num_objects:
+        return colors
+
+    num_objects -= len(colors)
+
+    # generate random colors
+    for _ in range(num_objects):
+        hue = np.random.rand()
+        rgb = colorsys.hsv_to_rgb(hue, 0.9, 0.9)
+        colors.append(tuple(int(c * 255) for c in rgb))
+
+    return colors
+
+
+@torch.no_grad()
+def draw_masks(
+    image: Tensor,
+    masks: Tensor,
+    alpha: float = 0.5,
+    colors: list[str | tuple[int, int, int]] | str | tuple[int, int, int] | None = None,
+) -> Tensor:
+    """
+    Draws segmentation masks on given RGB image.
+    This is different from `torchvision.utils.draw_segmentation_masks` as overlapping masks are blended together correctly.
+    The image values should be uint8 or float.
+
+    Args:
+        image (Tensor): Tensor of shape (3, H, W) and dtype uint8 or float.
+        masks (Tensor): Tensor of shape (num_masks, H, W) or (H, W) and dtype bool.
+        alpha (float): Float number between 0 and 1 denoting the transparency of the masks.
+            0 means full transparency, 1 means no transparency.
+        colors (color or list of colors, optional): List containing the colors
+            of the masks or single color for all masks. The color can be represented as
+            PIL strings e.g. "red" or "#FF00FF", or as RGB tuples e.g. ``(240, 10, 157)``.
+            By default, random colors are generated for each mask.
+
+    Returns:
+        img (Tensor[C, H, W]): Image Tensor, with segmentation masks drawn on top.
+    """
+    
+
+    if image.ndim == 3 and image.shape[0] == 1:
+        # convert to RGB
+        image = image.expand(3, -1, -1)
+
+    if masks.dtype != torch.bool:
+        masks = masks.bool()
+
+    # if image has negative values, scale to [0, 1]
+    if image.min() < 0:
+        image = image.float()
+        image = image - image.min()
+        image = image / image.max()
+
+    if masks.ndim == 2:
+        return torchvision.utils.draw_segmentation_masks(image=image,
+                                                         masks=masks,
+                                                         alpha=alpha,
+                                                         colors=colors)
+
+    if colors is None:
+        colors = generate_color_palette(len(masks))
+
+    for color, mask in zip(colors, masks):
+        image = torchvision.utils.draw_segmentation_masks(image=image,
+                                                          masks=mask,
+                                                          alpha=alpha,
+                                                          colors=color)
+    return image