hafnia 0.1.26__py3-none-any.whl → 0.2.0__py3-none-any.whl

hafnia/visualizations/colors.py

@@ -0,0 +1,267 @@
+from typing import List, Tuple
+
+
+def get_n_colors(index: int) -> List[Tuple[int, int, int]]:
+    n_colors = len(COLORS)
+    colors = [COLORS[index % n_colors] for index in range(index)]
+    return colors
+
+
+COLORS = [
+    (210, 24, 32),
+    (24, 105, 255),
+    (0, 138, 0),
+    (243, 109, 255),
+    (113, 0, 121),
+    (170, 251, 0),
+    (0, 190, 194),
+    (255, 162, 53),
+    (93, 61, 4),
+    (8, 0, 138),
+    (0, 93, 93),
+    (154, 125, 130),
+    (162, 174, 255),
+    (150, 182, 117),
+    (158, 40, 255),
+    (77, 0, 20),
+    (255, 174, 190),
+    (206, 0, 146),
+    (0, 255, 182),
+    (0, 45, 0),
+    (158, 117, 0),
+    (61, 53, 65),
+    (243, 235, 146),
+    (101, 97, 138),
+    (138, 61, 77),
+    (89, 4, 186),
+    (85, 138, 113),
+    (178, 190, 194),
+    (255, 93, 130),
+    (28, 198, 0),
+    (146, 247, 255),
+    (45, 134, 166),
+    (57, 93, 40),
+    (235, 206, 255),
+    (255, 93, 0),
+    (166, 97, 170),
+    (134, 0, 0),
+    (53, 0, 89),
+    (0, 81, 142),
+    (158, 73, 16),
+    (206, 190, 0),
+    (0, 40, 40),
+    (0, 178, 255),
+    (202, 166, 134),
+    (190, 154, 194),
+    (45, 32, 12),
+    (117, 101, 69),
+    (130, 121, 223),
+    (0, 194, 138),
+    (186, 231, 194),
+    (134, 142, 166),
+    (202, 113, 89),
+    (130, 154, 0),
+    (45, 0, 255),
+    (210, 4, 247),
+    (255, 215, 190),
+    (146, 206, 247),
+    (186, 93, 125),
+    (255, 65, 194),
+    (190, 134, 255),
+    (146, 142, 101),
+    (166, 4, 170),
+    (134, 227, 117),
+    (73, 0, 61),
+    (251, 239, 12),
+    (105, 85, 93),
+    (89, 49, 45),
+    (105, 53, 255),
+    (182, 4, 77),
+    (93, 109, 113),
+    (65, 69, 53),
+    (101, 113, 0),
+    (121, 0, 73),
+    (28, 49, 81),
+    (121, 65, 158),
+    (255, 146, 113),
+    (255, 166, 243),
+    (186, 158, 65),
+    (130, 170, 154),
+    (215, 121, 0),
+    (73, 61, 113),
+    (81, 162, 85),
+    (231, 130, 182),
+    (210, 227, 251),
+    (0, 73, 49),
+    (109, 219, 194),
+    (61, 77, 93),
+    (97, 53, 85),
+    (0, 113, 81),
+    (93, 24, 0),
+    (154, 93, 81),
+    (85, 142, 219),
+    (202, 202, 154),
+    (53, 24, 32),
+    (57, 61, 0),
+    (0, 154, 150),
+    (235, 16, 109),
+    (138, 69, 121),
+    (117, 170, 194),
+    (202, 146, 154),
+    (210, 186, 198),
+    (154, 206, 0),
+    (69, 109, 170),
+    (117, 89, 0),
+    (206, 77, 12),
+    (0, 223, 251),
+    (255, 61, 65),
+    (255, 202, 73),
+    (45, 49, 146),
+    (134, 105, 134),
+    (158, 130, 190),
+    (206, 174, 255),
+    (121, 69, 45),
+    (198, 251, 130),
+    (93, 117, 73),
+    (182, 69, 73),
+    (255, 223, 239),
+    (162, 0, 113),
+    (77, 77, 166),
+    (166, 170, 202),
+    (113, 28, 40),
+    (40, 121, 121),
+    (8, 73, 0),
+    (0, 105, 134),
+    (166, 117, 73),
+    (251, 182, 130),
+    (85, 24, 125),
+    (0, 255, 89),
+    (0, 65, 77),
+    (109, 142, 146),
+    (170, 36, 0),
+    (190, 210, 109),
+    (138, 97, 186),
+    (210, 65, 190),
+    (73, 97, 81),
+    (206, 243, 239),
+    (97, 194, 97),
+    (20, 138, 77),
+    (0, 255, 231),
+    (0, 105, 0),
+    (178, 121, 158),
+    (170, 178, 158),
+    (186, 85, 255),
+    (198, 121, 206),
+    (32, 49, 32),
+    (125, 4, 219),
+    (194, 198, 247),
+    (138, 198, 206),
+    (231, 235, 206),
+    (40, 28, 57),
+    (158, 255, 174),
+    (130, 206, 154),
+    (49, 166, 12),
+    (0, 162, 117),
+    (219, 146, 85),
+    (61, 20, 4),
+    (255, 138, 154),
+    (130, 134, 53),
+    (105, 77, 113),
+    (182, 97, 0),
+    (125, 45, 0),
+    (162, 178, 57),
+    (49, 4, 125),
+    (166, 61, 202),
+    (154, 32, 45),
+    (4, 223, 134),
+    (117, 125, 109),
+    (138, 150, 210),
+    (8, 162, 202),
+    (247, 109, 93),
+    (16, 85, 202),
+    (219, 182, 101),
+    (146, 89, 109),
+    (162, 255, 227),
+    (89, 85, 40),
+    (113, 121, 170),
+    (215, 89, 101),
+    (73, 32, 81),
+    (223, 77, 146),
+    (0, 0, 202),
+    (93, 101, 210),
+    (223, 166, 0),
+    (178, 73, 146),
+    (182, 138, 117),
+    (97, 77, 61),
+    (166, 150, 162),
+    (85, 28, 53),
+    (49, 65, 65),
+    (117, 117, 134),
+    (146, 158, 162),
+    (117, 154, 113),
+    (255, 130, 32),
+    (134, 85, 255),
+    (154, 198, 182),
+    (223, 150, 243),
+    (202, 223, 49),
+    (142, 93, 40),
+    (53, 190, 227),
+    (113, 166, 255),
+    (89, 138, 49),
+    (255, 194, 235),
+    (170, 61, 105),
+    (73, 97, 125),
+    (73, 53, 28),
+    (69, 178, 158),
+    (28, 36, 49),
+    (247, 49, 239),
+    (117, 0, 166),
+    (231, 182, 170),
+    (130, 105, 101),
+    (227, 162, 202),
+    (32, 36, 0),
+    (121, 182, 16),
+    (158, 142, 255),
+    (210, 117, 138),
+    (202, 182, 219),
+    (174, 154, 223),
+    (255, 113, 219),
+    (210, 247, 178),
+    (198, 215, 206),
+    (255, 210, 138),
+    (93, 223, 53),
+    (93, 121, 146),
+    (162, 142, 0),
+    (174, 223, 239),
+    (113, 77, 194),
+    (125, 69, 0),
+    (101, 146, 182),
+    (93, 121, 255),
+    (81, 73, 89),
+    (150, 158, 81),
+    (206, 105, 174),
+    (101, 53, 117),
+    (219, 210, 227),
+    (182, 174, 117),
+    (81, 89, 0),
+    (182, 89, 57),
+    (85, 4, 235),
+    (61, 117, 45),
+    (146, 130, 154),
+    (130, 36, 105),
+    (186, 134, 57),
+    (138, 178, 227),
+    (109, 178, 130),
+    (150, 65, 53),
+    (109, 65, 73),
+    (138, 117, 61),
+    (178, 113, 117),
+    (146, 28, 73),
+    (223, 109, 49),
+    (0, 227, 223),
+    (146, 4, 202),
+    (49, 40, 89),
+    (0, 125, 210),
+    (162, 109, 255),
+    (130, 89, 146),
+]

hafnia/visualizations/image_visualizations.py

@@ -0,0 +1,202 @@
+import shutil
+from pathlib import Path
+from typing import Dict, List, Optional, Tuple, Type, Union
+
+import cv2
+import numpy as np
+import numpy.typing as npt
+from PIL import Image
+
+from hafnia.dataset.hafnia_dataset import HafniaDataset, Sample
+from hafnia.dataset.primitives import (
+    Bbox,
+    Bitmask,
+    Classification,
+    Polygon,
+    Segmentation,
+)
+from hafnia.dataset.primitives.primitive import Primitive
+
+
+def draw_anonymize_by_blurring(
+    image: np.ndarray,
+    primitives: List[Primitive],
+    inplace: bool = False,
+    class_names: Union[List[str], str] = "all",
+    anonymization_settings: Optional[Dict[Type[Primitive], Dict]] = None,
+) -> np.ndarray:
+    if not inplace:
+        image = image.copy()
+
+    anonymization_settings = anonymization_settings or {}
+    if isinstance(class_names, str) and class_names == "all":
+        primitives = primitives
+    elif isinstance(class_names, list):
+        primitives = [primitive for primitive in primitives if primitive.class_name in class_names]
+    else:
+        raise ValueError(f"Invalid class_names type: {type(class_names)}. Expected 'all' or a list of class names.")
+    for primitive in primitives:
+        settings = anonymization_settings.get(type(primitive), {})
+        image = primitive.anonymize_by_blurring(image, inplace=True, **settings)
+    return image
+
+
+def draw_masks(image: np.ndarray, primitives: List[Primitive], inplace: bool = False) -> np.ndarray:
+    if not inplace:
+        image = image.copy()
+
+    for primitive in primitives:
+        primitive.mask(image, inplace=True)
+    return image
+
+
+def draw_annotations(
+    image: np.ndarray,
+    primitives: List[Primitive],
+    inplace: bool = False,
+    draw_settings: Optional[Dict[Type[Primitive], Dict]] = None,
+) -> np.ndarray:
+    if not inplace:
+        image = image.copy()
+    draw_settings = draw_settings or {}
+    primitives_order = [Segmentation, Bitmask, Bbox, Polygon, Classification]
+    primitives = sorted(primitives, key=lambda x: primitives_order.index(type(x)))
+    for primitive in primitives:
+        draw_settings_for_primitive = draw_settings.get(type(primitive), {})
+        image = primitive.draw(image, **draw_settings_for_primitive)
+    return image
+
+
+def concatenate_below(img0: np.ndarray, below_img: np.ndarray) -> np.ndarray:
+    scale_factor = img0.shape[1] / below_img.shape[1]
+    new_height = int(below_img.shape[0] * scale_factor)
+    text_region_resized = cv2.resize(below_img, (img0.shape[1], new_height))
+    if len(img0.shape) == 2:
+        img0 = cv2.cvtColor(img0, cv2.COLOR_GRAY2BGR)
+    if len(text_region_resized.shape) == 2:
+        text_region_resized = cv2.cvtColor(text_region_resized, cv2.COLOR_GRAY2BGR)
+    frame_visualized = np.concatenate([img0, text_region_resized], axis=0)
+    return frame_visualized
+
+
+def concatenate_below_resize_by_padding(img0: np.ndarray, below_img: np.ndarray):
+    max_width = max(img0.shape[1], below_img.shape[1])
+
+    if len(img0.shape) == 2:
+        img0 = cv2.cvtColor(img0, cv2.COLOR_GRAY2RGB)
+    if len(below_img.shape) == 2:
+        below_img = cv2.cvtColor(below_img, cv2.COLOR_GRAY2RGB)
+    img0_padded = resize_width_by_padding(img0, new_width=max_width)
+    below_img_padded = resize_width_by_padding(below_img, new_width=max_width)
+
+    return np.concatenate([img0_padded, below_img_padded], axis=0)
+
+
+def resize_width_by_padding(img0: np.ndarray, new_width: int) -> np.ndarray:
+    img0_new_shape = list(img0.shape)
+    img0_new_shape[1] = new_width
+    img0_padded = np.zeros(img0_new_shape, dtype=img0.dtype)
+    extra_width = new_width - img0.shape[1]
+    left_margin = extra_width // 2
+    img0_padded[:, left_margin : left_margin + img0.shape[1]] = img0
+    return img0_padded
+
+
+def append_text_below_frame(frame: np.ndarray, text: str) -> np.ndarray:
+    font_size_px = int(frame.shape[0] * 0.1)  # 10% of the frame height
+    font_size_px = max(font_size_px, 7)  # Ensure a minimum font size
+    font_size_px = min(font_size_px, 50)  # Ensure a maximum font size
+
+    text_region = create_text_img(text, font_size_px=font_size_px)
+    frame_with_text = concatenate_below_resize_by_padding(frame, text_region)
+    return frame_with_text
+
+
+def create_text_img(
+    text_strings: Union[List[str], str],
+    font_size_px: int,
+    text_width: Optional[int] = None,
+    color: Tuple[int, int, int] = (255, 255, 255),
+    bg_color: Tuple[int, int, int] = (0, 0, 0),
+) -> npt.NDArray[np.uint8]:
+    """Private code borrowed by Peter."""
+    font_face = cv2.FONT_HERSHEY_SIMPLEX
+    thickness = 2
+    if font_size_px < 15:
+        thickness = 1
+    line_type = cv2.LINE_AA
+    if isinstance(text_strings, str):
+        text_strings = [text_strings]
+    font_scale = cv2.getFontScaleFromHeight(fontFace=font_face, pixelHeight=font_size_px, thickness=thickness)
+
+    text_w_max = 0
+    for text in text_strings:
+        (text_w, text_h), baseline = cv2.getTextSize(
+            text=text, fontFace=font_face, fontScale=font_scale, thickness=thickness
+        )
+        text_w_max = max(text_w_max, text_w)
+
+    text_width = text_width or text_w_max + baseline * 2
+    text_height = text_h + baseline * 2
+
+    y_pos = text_h + baseline
+    text_imgs = []
+    for text in text_strings:
+        shape_color_image = (text_height, text_width, 3)
+        img = np.full(shape_color_image, bg_color, dtype=np.uint8)
+        text_img = cv2.putText(
+            img=img,
+            text=text,
+            org=(baseline, y_pos),
+            fontFace=font_face,
+            fontScale=font_scale,
+            color=color,
+            thickness=thickness,
+            lineType=line_type,
+        )
+        text_imgs.append(text_img)
+    img_text = np.vstack(text_imgs)
+    return img_text
+
+
+def concatenate_right(img0: np.ndarray, below_img: np.ndarray) -> np.ndarray:
+    scale_factor = img0.shape[0] / below_img.shape[0]
+    new_width = int(below_img.shape[1] * scale_factor)
+    text_region_resized = cv2.resize(below_img, (new_width, img0.shape[0]))
+
+    frame_visualized = np.concatenate([img0, text_region_resized], axis=1)
+    return frame_visualized
+
+
+def save_dataset_sample_set_visualizations(
+    path_dataset: Path,
+    path_output_folder: Path,
+    max_samples: int = 10,
+    draw_settings: Optional[Dict[Type[Primitive], Dict]] = None,
+    anonymize_settings: Optional[Dict[Type[Primitive], Dict]] = None,
+) -> List[Path]:
+    dataset = HafniaDataset.read_from_path(path_dataset)
+    shutil.rmtree(path_output_folder, ignore_errors=True)
+    path_output_folder.mkdir(parents=True)
+
+    draw_settings = draw_settings or {}
+
+    paths = []
+    dataset_shuffled = dataset.shuffle(seed=42)
+    for sample_dict in dataset_shuffled:
+        sample = Sample(**sample_dict)
+        image = sample.read_image()
+        annotations = sample.get_annotations()
+
+        if anonymize_settings:
+            image = draw_anonymize_by_blurring(image, annotations, anonymization_settings=anonymize_settings)
+        image = draw_annotations(image, annotations, draw_settings=draw_settings)
+
+        pil_image = Image.fromarray(image)
+        path_image = path_output_folder / Path(sample.file_name).name
+        pil_image.save(path_image)
+        paths.append(path_image)
+
+        if len(paths) >= max_samples:
+            return paths  # Return early if we have enough samples
+    return paths