quadra 0.0.1__py3-none-any.whl → 2.1.13__py3-none-any.whl

quadra/utils/patch/visualization.py

@@ -0,0 +1,217 @@
+from __future__ import annotations
+
+import os
+
+import cv2
+import matplotlib.pyplot as plt
+import numpy as np
+from matplotlib.cm import get_cmap
+from matplotlib.colors import Colormap
+from matplotlib.lines import Line2D
+from matplotlib.pyplot import Figure
+
+from quadra.utils import utils
+
+log = utils.get_logger(__name__)
+
+
+def plot_patch_reconstruction(
+    reconstruction: dict,
+    idx_to_class: dict[int, str],
+    class_to_idx: dict[str, int],
+    ignore_classes: list[int] | None = None,
+    is_polygon: bool = True,
+) -> Figure:
+    """Helper function for plotting the patch reconstruction.
+
+    Args:
+        reconstruction: Dict following this structure
+            {
+                "file_path": str,
+                "mask_path": str,
+                "prediction": {
+                    "label": str,
+                    "points": [{"x": int, "y": int}]
+                }
+            } if is_polygon else
+            {
+                "file_path": str,
+                "mask_path": str,
+                "prediction": np.ndarray
+            }
+        idx_to_class: Dictionary mapping indices to label names
+        class_to_idx: Dictionary mapping class names to indices
+        ignore_classes: Eventually the classes to not plot
+        is_polygon: Boolean indicating if the prediction is a polygon or a mask.
+
+    Returns:
+        Matplotlib plot showing predicted patch regions and eventually gt
+
+    """
+    cmap_name = "tab10"
+
+    # 10 classes + good
+    if len(idx_to_class.values()) > 11:
+        cmap_name = "tab20"
+
+    cmap = get_cmap(cmap_name)
+    test_img = cv2.imread(reconstruction["image_path"])
+    test_img = cv2.cvtColor(test_img, cv2.COLOR_BGR2RGB)
+    gt_img = None
+
+    if reconstruction["mask_path"] is not None and os.path.isfile(reconstruction["mask_path"]):
+        gt_img = cv2.imread(reconstruction["mask_path"], 0)
+
+    out = np.zeros((test_img.shape[0], test_img.shape[1]), dtype=np.uint8)
+
+    if is_polygon:
+        for _, region in enumerate(reconstruction["prediction"]):
+            points = [[item["x"], item["y"]] for item in region["points"]]
+            c_label = region["label"]
+
+            out = cv2.drawContours(
+                out,
+                np.array([points], np.int32),
+                -1,
+                class_to_idx[c_label],
+                thickness=cv2.FILLED,
+            )  # type: ignore[call-overload]
+    else:
+        out = reconstruction["prediction"]
+
+    fig = plot_patch_results(
+        image=test_img,
+        prediction_image=out,
+        ground_truth_image=gt_img,
+        plot_original=True,
+        ignore_classes=ignore_classes,
+        save_path=None,
+        class_to_idx=class_to_idx,
+        cmap=cmap,
+    )
+
+    return fig
+
+
+def show_mask_on_image(image: np.ndarray, mask: np.ndarray):
+    """Plot mask on top of the original image."""
+    image = image.astype(np.float32) / 255
+    mask = mask.astype(np.float32) / 255
+    out = mask + image.astype(np.float32)
+    out = out / np.max(out)
+    return np.uint8(255 * out)
+
+
+def create_rgb_mask(
+    mask: np.ndarray,
+    color_map: dict,
+    ignore_classes: list[int] | None = None,
+    ground_truth_mask: np.ndarray | None = None,
+):
+    """Convert index mask to RGB mask."""
+    output_mask = np.zeros([mask.shape[0], mask.shape[1], 3])
+    for c in np.unique(mask):
+        if ignore_classes is not None and c in ignore_classes:
+            continue
+
+        output_mask[mask == c] = color_map[str(c)]
+    if ignore_classes is not None and ground_truth_mask is not None:
+        output_mask[np.isin(ground_truth_mask, ignore_classes)] = [0, 0, 0]
+
+    return output_mask
+
+
+def plot_patch_results(
+    image: np.ndarray,
+    prediction_image: np.ndarray,
+    ground_truth_image: np.ndarray | None,
+    class_to_idx: dict[str, int],
+    plot_original: bool = True,
+    ignore_classes: list[int] | None = None,
+    image_height: int = 10,
+    save_path: str | None = None,
+    cmap: Colormap | None = None,
+) -> Figure:
+    """Function used to plot the image predicted.
+
+    Args:
+        prediction_image: The prediction image
+        image: The original image to plot
+        ground_truth_image: The ground truth image
+        class_to_idx: Dictionary mapping class names to indices
+        plot_original: Boolean to plot the original image
+        ignore_classes: The classes to ignore, default is 0
+        image_height: The height of the output figure
+        save_path: The path to save the figure
+        cmap: The colormap to use. If None, tab20 is used
+
+    Returns:
+        The matplotlib figure
+    """
+    if ignore_classes is None:
+        ignore_classes = [0]
+
+    if cmap is None:
+        cmap = get_cmap("tab20")
+
+    image = image[0 : prediction_image.shape[0], 0 : prediction_image.shape[1], :]
+    idx_to_class = {v: k for k, v in class_to_idx.items()}
+
+    if ignore_classes is not None:
+        class_to_idx = {k: v for k, v in class_to_idx.items() if v not in ignore_classes}
+
+    class_idxs = list(class_to_idx.values())
+
+    cmap = {str(c): tuple(int(i * 255) for i in cmap(c / len(class_idxs))[:-1]) for c in class_idxs}
+    output_images = []
+    titles = []
+
+    if plot_original:
+        output_images.append(image)
+        titles.append("Original Image")
+
+    if ground_truth_image is not None:
+        ground_truth_image = ground_truth_image[0 : prediction_image.shape[0], 0 : prediction_image.shape[1]]
+        ground_truth_mask = create_rgb_mask(ground_truth_image, cmap, ignore_classes=ignore_classes)
+        output_images.append(ground_truth_mask)
+        titles.append("Ground Truth Mask")
+
+    prediction_mask = create_rgb_mask(
+        prediction_image,
+        cmap,
+        ignore_classes=ignore_classes,
+    )
+
+    output_images.append(prediction_mask)
+    titles.append("Prediction Mask")
+    if ignore_classes is not None and ground_truth_image is not None:
+        prediction_mask = create_rgb_mask(
+            prediction_image, cmap, ignore_classes=ignore_classes, ground_truth_mask=ground_truth_image
+        )
+
+        ignored_classes_str = [idx_to_class[c] for c in ignore_classes]
+        prediction_title = f"Prediction Mask \n (Ignoring Ground Truth Class: {ignored_classes_str})"
+        output_images.append(prediction_mask)
+        titles.append(prediction_title)
+
+    fig, axs = plt.subplots(
+        ncols=len(output_images),
+        nrows=1,
+        figsize=(len(output_images) * image_height, image_height),
+        squeeze=False,
+        facecolor="white",
+    )
+
+    for i, output_image in enumerate(output_images):
+        axs[0, i].imshow(show_mask_on_image(image, output_image))
+        axs[0, i].set_title(titles[i])
+        axs[0, i].axis("off")
+
+    custom_lines = [Line2D([0], [0], color=tuple(i / 255.0 for i in cmap[str(c)]), lw=4) for c in class_idxs]
+    custom_labels = list(class_to_idx.keys())
+    axs[0, -1].legend(custom_lines, custom_labels, loc="center left", bbox_to_anchor=(1.01, 0.81), borderaxespad=0)
+    if save_path is not None:
+        plt.savefig(save_path, bbox_inches="tight")
+        plt.close()
+
+    return fig

quadra/utils/resolver.py

@@ -0,0 +1,42 @@
+from __future__ import annotations
+
+from typing import Any
+
+from hydra.core.hydra_config import HydraConfig
+from omegaconf import OmegaConf
+
+
+def multirun_subdir_beautify(subdir: str) -> str:
+    """Change the subdir name to be more readable and usable, this function will replace / with | to avoid creating
+    undesired subdirectories and remove the left part of the equals sign to avoid having too long names.
+
+    Args:
+        subdir: The subdir name.
+
+    Returns:
+        The beautified subdir name.
+
+    Examples:
+        >>> multirun_subdir_beautify("experiment=pippo/anomaly/padim,trainer.batch_size=32")
+        "pippo|anomaly|padim,32"
+    """
+    hydra_cfg = HydraConfig.get()
+    if hydra_cfg.mode is None or hydra_cfg.mode.name == "RUN":
+        return subdir
+    # Remove slashes to avoid creating multiple subdirs
+    # TODO: if right side of the equals sign has `,` this will not work.
+    subdir_list = subdir.replace("/", "|").split(",")
+    subdir = ",".join([x.split("=")[1].replace(" ", "") for x in subdir_list])
+
+    return subdir
+
+
+def as_tuple(*args: Any) -> tuple[Any, ...]:
+    """Resolves a list of arguments to a tuple."""
+    return tuple(args)
+
+
+def register_resolvers() -> None:
+    """Register custom resolver."""
+    OmegaConf.register_new_resolver("multirun_subdir_beautify", multirun_subdir_beautify)
+    OmegaConf.register_new_resolver("as_tuple", as_tuple)

quadra/utils/segmentation.py

@@ -0,0 +1,31 @@
+import numpy as np
+import skimage
+from skimage.morphology import medial_axis
+
+from quadra.utils import utils
+
+log = utils.get_logger(__name__)
+
+
+def smooth_mask(mask: np.ndarray) -> np.ndarray:
+    """Smooths for segmentation.
+
+    Args:
+        mask: Input mask
+
+    Returns:
+        Smoothed mask
+    """
+    labeled_mask = skimage.measure.label(mask)
+    labels = np.arange(0, np.max(labeled_mask) + 1)
+    output_mask = np.zeros_like(mask).astype(np.float32)
+    for l in labels:
+        component_mask = labeled_mask == l
+        _, distance = medial_axis(component_mask, return_distance=True)
+        component_mask_norm = distance ** (1 / 2.2)
+        component_mask_norm = (component_mask_norm - np.min(component_mask_norm)) / (
+            np.max(component_mask_norm) - np.min(component_mask_norm)
+        )
+        output_mask += component_mask_norm
+    output_mask = output_mask * mask
+    return output_mask

quadra/utils/tests/fixtures/__init__.py

@@ -0,0 +1 @@
+from .dataset import *  # noqa: F403

quadra/utils/tests/fixtures/dataset/__init__.py

@@ -0,0 +1,39 @@
+from .anomaly import AnomalyDatasetArguments, anomaly_dataset, base_anomaly_dataset
+from .classification import (
+    ClassificationDatasetArguments,
+    ClassificationMultilabelDatasetArguments,
+    ClassificationPatchDatasetArguments,
+    base_classification_dataset,
+    base_multilabel_classification_dataset,
+    base_patch_classification_dataset,
+    classification_dataset,
+    classification_patch_dataset,
+    multilabel_classification_dataset,
+)
+from .imagenette import imagenette_dataset
+from .segmentation import (
+    SegmentationDatasetArguments,
+    base_binary_segmentation_dataset,
+    base_multiclass_segmentation_dataset,
+    segmentation_dataset,
+)
+
+__all__ = [
+    "anomaly_dataset",
+    "classification_dataset",
+    "AnomalyDatasetArguments",
+    "ClassificationDatasetArguments",
+    "ClassificationPatchDatasetArguments",
+    "classification_patch_dataset",
+    "segmentation_dataset",
+    "SegmentationDatasetArguments",
+    "multilabel_classification_dataset",
+    "ClassificationMultilabelDatasetArguments",
+    "base_anomaly_dataset",
+    "imagenette_dataset",
+    "base_classification_dataset",
+    "base_patch_classification_dataset",
+    "base_binary_segmentation_dataset",
+    "base_multiclass_segmentation_dataset",
+    "base_multilabel_classification_dataset",
+]

quadra/utils/tests/fixtures/dataset/anomaly.py

@@ -0,0 +1,124 @@
+from __future__ import annotations
+
+import shutil
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any
+
+import cv2
+import pytest
+
+from quadra.utils.tests.helpers import _random_image
+
+
+@dataclass
+class AnomalyDatasetArguments:
+    """Anomaly dataset arguments.
+
+    Args:
+        train_samples: number of train samples
+        val_samples: number of validation samples (good, bad)
+        test_samples: number of test samples (good, bad)
+    """
+
+    train_samples: int
+    val_samples: tuple[int, int]
+    test_samples: tuple[int, int]
+
+
+def _build_anomaly_dataset(
+    tmp_path: Path, dataset_arguments: AnomalyDatasetArguments
+) -> tuple[str, AnomalyDatasetArguments]:
+    """Generate anomaly dataset in the standard mvtec format.
+
+    Args:
+        tmp_path: path to temporary directory
+        dataset_arguments: dataset arguments
+
+    Returns:
+        path to anomaly dataset
+    """
+    train_samples = dataset_arguments.train_samples
+    val_samples = dataset_arguments.val_samples
+    test_samples = dataset_arguments.test_samples
+
+    anomaly_dataset_path = tmp_path / "anomaly_dataset"
+    anomaly_dataset_path.mkdir()
+    train_good_path = anomaly_dataset_path / "train" / "good"
+    val_good_path = anomaly_dataset_path / "val" / "good"
+    val_bad_path = anomaly_dataset_path / "val" / "bad"
+    test_good_path = anomaly_dataset_path / "test" / "good"
+    test_bad_path = anomaly_dataset_path / "test" / "bad"
+
+    train_good_path.mkdir(parents=True)
+    val_good_path.mkdir(parents=True)
+    val_bad_path.mkdir(parents=True)
+    test_good_path.mkdir(parents=True)
+    test_bad_path.mkdir(parents=True)
+
+    # Generate train good images
+    for i in range(train_samples):
+        image = _random_image()
+        image_path = train_good_path / f"train_{i}.png"
+        cv2.imwrite(str(image_path), image)
+
+    # Generate val good images
+    for i in range(val_samples[0]):
+        image = _random_image()
+        image_path = val_good_path / f"val_{i}.png"
+        cv2.imwrite(str(image_path), image)
+    # Generate val bad images
+    for i in range(val_samples[1]):
+        image = _random_image()
+        image_path = val_bad_path / f"val_{i}.png"
+        cv2.imwrite(str(image_path), image)
+
+    # Generate test good images
+    for i in range(test_samples[0]):
+        image = _random_image()
+        image_path = test_good_path / f"test_{i}.png"
+        cv2.imwrite(str(image_path), image)
+    # Generate test bad images
+    for i in range(test_samples[1]):
+        image = _random_image()
+        image_path = test_bad_path / f"test_{i}.png"
+        cv2.imwrite(str(image_path), image)
+
+    return str(anomaly_dataset_path), dataset_arguments
+
+
+@pytest.fixture
+def anomaly_dataset(tmp_path: Path, dataset_arguments: AnomalyDatasetArguments) -> tuple[str, AnomalyDatasetArguments]:
+    """Fixture used to dinamically generate anomaly dataset. By default images are random grayscales with size 10x10.
+
+    Args:
+        tmp_path: path to temporary directory
+        dataset_arguments: dataset arguments
+
+    Returns:
+        path to anomaly dataset
+    """
+    yield _build_anomaly_dataset(tmp_path, dataset_arguments)
+    if tmp_path.exists():
+        shutil.rmtree(tmp_path)
+
+
+@pytest.fixture(
+    params=[AnomalyDatasetArguments(**{"train_samples": 10, "val_samples": (1, 1), "test_samples": (1, 1)})]
+)
+def base_anomaly_dataset(tmp_path: Path, request: Any) -> tuple[str, AnomalyDatasetArguments]:
+    """Generate base anomaly dataset with the following parameters:
+        - train_samples: 10
+        - val_samples: (10, 10)
+        - test_samples: (10, 10).
+
+    Args:
+        tmp_path: Path to temporary directory
+        request: Pytest SubRequest object
+
+    Yields:
+        Path to anomaly dataset and dataset arguments
+    """
+    yield _build_anomaly_dataset(tmp_path, request.param)
+    if tmp_path.exists():
+        shutil.rmtree(tmp_path)