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

hafnia/torch_helpers.py

@@ -1,80 +1,126 @@
-from typing import List, Optional
+from typing import Dict, List, Optional, Tuple, Type, Union
 
-import datasets
+import numpy as np
 import torch
 import torchvision
-from flatten_dict import flatten
+from flatten_dict import flatten, unflatten
 from PIL import Image, ImageDraw, ImageFont
 from torchvision import tv_tensors
 from torchvision import utils as tv_utils
 from torchvision.transforms import v2
 
+from hafnia.dataset.dataset_names import FieldName
+from hafnia.dataset.hafnia_dataset import HafniaDataset, Sample
+from hafnia.dataset.primitives import (
+    PRIMITIVE_COLUMN_NAMES,
+    class_color_by_name,
+)
+from hafnia.dataset.primitives.bbox import Bbox
+from hafnia.dataset.primitives.bitmask import Bitmask
+from hafnia.dataset.primitives.classification import Classification
+from hafnia.dataset.primitives.primitive import Primitive
+from hafnia.dataset.primitives.segmentation import Segmentation
+from hafnia.log import user_logger
+
+
+def get_primitives_per_task_name_for_primitive(
+    sample: Sample, PrimitiveType: Type[Primitive], split_by_task_name: bool = True
+) -> Dict[str, List[Primitive]]:
+    if not hasattr(sample, PrimitiveType.column_name()):
+        return {}
+
+    primitives = getattr(sample, PrimitiveType.column_name())
+    if primitives is None:
+        return {}
+
+    primitives_by_task_name: Dict[str, List[Primitive]] = {}
+    for primitive in primitives:
+        if primitive.task_name not in primitives_by_task_name:
+            primitives_by_task_name[primitive.task_name] = []
+        primitives_by_task_name[primitive.task_name].append(primitive)
+    return primitives_by_task_name
+
 
 class TorchvisionDataset(torch.utils.data.Dataset):
     def __init__(
         self,
-        hf_dataset: datasets.Dataset,
+        dataset: HafniaDataset,
         transforms=None,
-        classification_tasks: Optional[List] = None,
-        object_tasks: Optional[List] = None,
-        segmentation_tasks: Optional[List] = None,
         keep_metadata: bool = False,
     ):
-        self.dataset = hf_dataset
+        self.dataset = dataset
+
         self.transforms = transforms
-        self.object_tasks = object_tasks or ["objects"]
-        self.segmentation_tasks = segmentation_tasks or ["segmentation"]
-        self.classification_tasks = classification_tasks or ["classification"]
         self.keep_metadata = keep_metadata
 
-    def __getitem__(self, idx):
-        sample = self.dataset[idx]
-
-        # For now, we expect a dataset to always have an image field
-        image = tv_tensors.Image(sample.pop("image"))
-
-        img_shape = image.shape[-2:]
-        target = {}
+    def __getitem__(self, idx: int) -> Tuple[torch.Tensor, Dict]:
+        sample_dict = self.dataset[idx]
+        sample = Sample(**sample_dict)
+        image = tv_tensors.Image(sample.read_image_pillow())
+        h, w = image.shape[-2:]
+        target_flat = {}
+        mask_tasks: Dict[str, List[Segmentation]] = get_primitives_per_task_name_for_primitive(sample, Segmentation)
+        for task_name, masks in mask_tasks.items():
+            raise NotImplementedError("Segmentation tasks are not yet implemented")
+            # target[f"{mask.task_name}.mask"] = tv_tensors.Mask(mask.mask)
+
+        class_tasks: Dict[str, List] = get_primitives_per_task_name_for_primitive(sample, Classification)
+        for task_name, classifications in class_tasks.items():
+            assert len(classifications) == 1, "Expected exactly one classification task per sample"
+            target_flat[f"{Classification.column_name()}.{task_name}"] = {
+                FieldName.CLASS_IDX: classifications[0].class_idx,
+                FieldName.CLASS_NAME: classifications[0].class_name,
+            }
+
+        bbox_tasks: Dict[str, List[Bbox]] = get_primitives_per_task_name_for_primitive(sample, Bbox)
+        for task_name, bboxes in bbox_tasks.items():
+            bboxes_list = [bbox.to_coco(image_height=h, image_width=w) for bbox in bboxes]
+            bboxes_tensor = torch.as_tensor(bboxes_list).reshape(-1, 4)
+            target_flat[f"{Bbox.column_name()}.{task_name}"] = {
+                FieldName.CLASS_IDX: [bbox.class_idx for bbox in bboxes],
+                FieldName.CLASS_NAME: [bbox.class_name for bbox in bboxes],
+                "bbox": tv_tensors.BoundingBoxes(bboxes_tensor, format="XYWH", canvas_size=(h, w)),
+            }
+
+        bitmask_tasks: Dict[str, List[Bitmask]] = get_primitives_per_task_name_for_primitive(sample, Bitmask)
+        for task_name, bitmasks in bitmask_tasks.items():
+            bitmasks_np = np.array([bitmask.to_mask(img_height=h, img_width=w) for bitmask in bitmasks])
+            target_flat[f"{Bitmask.column_name()}.{task_name}"] = {
+                FieldName.CLASS_IDX: [bitmask.class_idx for bitmask in bitmasks],
+                FieldName.CLASS_NAME: [bitmask.class_name for bitmask in bitmasks],
+                "mask": tv_tensors.Mask(bitmasks_np),
+            }
 
-        for segmentation_task in self.segmentation_tasks:
-            if segmentation_task in sample:
-                target[f"{segmentation_task}.mask"] = tv_tensors.Mask(sample[segmentation_task].pop("mask"))
-
-        for classification_task in self.classification_tasks:
-            if classification_task in sample:
-                target[f"{classification_task}.class_idx"] = sample[classification_task].pop("class_idx")
-
-        for object_task in self.object_tasks:
-            if object_task in sample:
-                bboxes_list = sample[object_task].pop("bbox")
-                bboxes = tv_tensors.BoundingBoxes(bboxes_list, format="XYWH", canvas_size=img_shape)
-                if bboxes.numel() == 0:
-                    bboxes = bboxes.reshape(-1, 4)
-                target[f"{object_task}.bbox"] = bboxes
-                target[f"{object_task}.class_idx"] = torch.tensor(sample[object_task].pop("class_idx"))
+        if self.transforms:
+            image, target_flat = self.transforms(image, target_flat)
 
         if self.keep_metadata:
-            target.update(flatten(sample, reducer="dot"))
-
-        if self.transforms:
-            image, target = self.transforms(image, target)
+            sample_dict = sample_dict.copy()
+            drop_columns = PRIMITIVE_COLUMN_NAMES
+            for column in drop_columns:
+                if column in sample_dict:
+                    sample_dict.pop(column)
 
+        target = flatten(target_flat, reducer="dot")
         return image, target
 
     def __len__(self):
         return len(self.dataset)
 
 
-def draw_image_classification(visualize_image: torch.Tensor, text_label: str) -> torch.Tensor:
+def draw_image_classification(visualize_image: torch.Tensor, text_labels: Union[str, List[str]]) -> torch.Tensor:
+    if isinstance(text_labels, str):
+        text_labels = [text_labels]
+    text = "\n".join(text_labels)
     max_dim = max(visualize_image.shape[-2:])
-    font_size = max(int(max_dim * 0.1), 10)  # Minimum font size of 10
+    font_size = max(int(max_dim * 0.06), 10)  # Minimum font size of 10
     txt_font = ImageFont.load_default(font_size)
     dummie_draw = ImageDraw.Draw(Image.new("RGB", (10, 10)))
-    _, _, w, h = dummie_draw.textbbox((0, 0), text=text_label, font=txt_font)  # type: ignore[arg-type]
+    _, _, w, h = dummie_draw.textbbox((0, 0), text=text, font=txt_font)  # type: ignore[arg-type]
 
     text_image = Image.new("RGB", (int(w), int(h)))
     draw = ImageDraw.Draw(text_image)
-    draw.text((0, 0), text=text_label, font=txt_font)  # type: ignore[arg-type]
+    draw.text((0, 0), text=text, font=txt_font)  # type: ignore[arg-type]
     text_tensor = v2.functional.to_image(text_image)
 
     height = text_tensor.shape[-2] + visualize_image.shape[-2]
@@ -94,77 +140,116 @@ def draw_image_classification(visualize_image: torch.Tensor, text_label: str) ->
 def draw_image_and_targets(
     image: torch.Tensor,
     targets,
-    detection_tasks: Optional[List[str]] = None,
-    segmentation_tasks: Optional[List[str]] = None,
-    classification_tasks: Optional[List[str]] = None,
 ) -> torch.Tensor:
-    detection_tasks = detection_tasks or ["objects"]
-    segmentation_tasks = segmentation_tasks or ["segmentation"]
-    classification_tasks = classification_tasks or ["classification"]
-
     visualize_image = image.clone()
     if visualize_image.is_floating_point():
         visualize_image = image - torch.min(image)
         visualize_image = visualize_image / visualize_image.max()
 
     visualize_image = v2.functional.to_dtype(visualize_image, torch.uint8, scale=True)
-
-    for object_task in detection_tasks:
-        bbox_field = f"{object_task}.bbox"
-        if bbox_field in targets:
-            hugging_face_format = "xywh"
-            bbox = torchvision.ops.box_convert(targets[bbox_field], in_fmt=hugging_face_format, out_fmt="xyxy")
-            class_names_field = f"{object_task}.class_name"
-            class_names = targets.get(class_names_field, None)
-            visualize_image = tv_utils.draw_bounding_boxes(visualize_image, bbox, labels=class_names, width=2)
-
-    for segmentation_task in segmentation_tasks:
-        mask_field = f"{segmentation_task}.mask"
-        if mask_field in targets:
-            mask = targets[mask_field].squeeze(0)
-            masks_list = [mask == value for value in mask.unique()]
-            masks = torch.stack(masks_list, dim=0).to(torch.bool)
-            visualize_image = tv_utils.draw_segmentation_masks(visualize_image, masks=masks, alpha=0.5)
-
-    for classification_task in classification_tasks:
-        classification_field = f"{classification_task}.class_idx"
-        if classification_field in targets:
-            text_label = f"[{targets[classification_field]}]"
-            classification_name_field = f"{classification_task}.class_name"
-            if classification_name_field in targets:
-                text_label = text_label + f" {targets[classification_name_field]}"
-            visualize_image = draw_image_classification(visualize_image, text_label)
+    targets = unflatten(targets, splitter="dot")  # Nested dictionary format
+    # NOTE: Order of drawing is important so visualizations are not overlapping in an undesired way
+    if Segmentation.column_name() in targets:
+        primitive_annotations = targets[Segmentation.column_name()]
+        for task_name, task_annotations in primitive_annotations.items():
+            raise NotImplementedError("Segmentation tasks are not yet implemented")
+            # mask = targets[mask_field].squeeze(0)
+            # masks_list = [mask == value for value in mask.unique()]
+            # masks = torch.stack(masks_list, dim=0).to(torch.bool)
+            # visualize_image = tv_utils.draw_segmentation_masks(visualize_image, masks=masks, alpha=0.5)
+
+    if Bitmask.column_name() in targets:
+        primitive_annotations = targets[Bitmask.column_name()]
+        for task_name, task_annotations in primitive_annotations.items():
+            colors = [class_color_by_name(class_name) for class_name in task_annotations[FieldName.CLASS_NAME]]
+            visualize_image = tv_utils.draw_segmentation_masks(
+                image=visualize_image,
+                masks=task_annotations["mask"],
+                colors=colors,
+            )
+
+    if Bbox.column_name() in targets:
+        primitive_annotations = targets[Bbox.column_name()]
+        for task_name, task_annotations in primitive_annotations.items():
+            bboxes = torchvision.ops.box_convert(task_annotations["bbox"], in_fmt="xywh", out_fmt="xyxy")
+            colors = [class_color_by_name(class_name) for class_name in task_annotations[FieldName.CLASS_NAME]]
+            visualize_image = tv_utils.draw_bounding_boxes(
+                image=visualize_image,
+                boxes=bboxes,
+                labels=task_annotations[FieldName.CLASS_NAME],
+                width=2,
+                colors=colors,
+            )
+
+    # Important that classification is drawn last as it will change image dimensions
+    if Classification.column_name() in targets:
+        primitive_annotations = targets[Classification.column_name()]
+        text_labels = []
+        for task_name, task_annotations in primitive_annotations.items():
+            if task_name == Classification.default_task_name():
+                text_label = task_annotations[FieldName.CLASS_NAME]
+            else:
+                text_label = f"{task_name}: {task_annotations[FieldName.CLASS_NAME]}"
+            text_labels.append(text_label)
+        visualize_image = draw_image_classification(visualize_image, text_labels)
     return visualize_image
 
 
 class TorchVisionCollateFn:
     def __init__(self, skip_stacking: Optional[List] = None):
         if skip_stacking is None:
-            skip_stacking = []
-        self.skip_stacking_list = skip_stacking
+            skip_stacking = [f"{Bbox.column_name()}.*", f"{Bitmask.column_name()}.*"]
+
+        self.wild_card_skip_stacking = []
+        self.skip_stacking_list = []
+        for skip_name in skip_stacking:
+            if skip_name.endswith("*"):
+                self.wild_card_skip_stacking.append(skip_name[:-1])  # Remove the trailing '*'
+            else:
+                self.skip_stacking_list.append(skip_name)
+
+    def skip_key_name(self, key_name: str) -> bool:
+        if key_name in self.skip_stacking_list:
+            return True
+        if any(key_name.startswith(wild_card) for wild_card in self.wild_card_skip_stacking):
+            return True
+        return False
 
     def __call__(self, batch):
         images, targets = tuple(zip(*batch, strict=False))
         if "image" not in self.skip_stacking_list:
             images = torch.stack(images)
 
-        targets_modified = {k: [d[k] for d in targets] for k in targets[0]}
+        keys_min = set(targets[0])
+        keys_max = set(targets[0])
+        for target in targets:
+            keys_min = keys_min.intersection(target)
+            keys_max = keys_max.union(target)
+
+        if keys_min != keys_max:
+            user_logger.warning(
+                "Not all images in the batch contain the same targets. To solve for missing targets "
+                f"the following keys {keys_max - keys_min} are dropped from the batch "
+            )
+
+        targets_modified = {k: [d[k] for d in targets] for k in keys_min}
         for key_name, item_values in targets_modified.items():
-            if key_name not in self.skip_stacking_list:
-                first_element = item_values[0]
-                if isinstance(first_element, torch.Tensor):
-                    item_values = torch.stack(item_values)
-                elif isinstance(first_element, (int, float)):
-                    item_values = torch.tensor(item_values)
-                elif isinstance(first_element, (str, list)):
-                    # Skip stacking for certain types such as strings and lists
-                    pass
-                if isinstance(first_element, tv_tensors.Mask):
-                    item_values = tv_tensors.Mask(item_values)
-                elif isinstance(first_element, tv_tensors.Image):
-                    item_values = tv_tensors.Image(item_values)
-                elif isinstance(first_element, tv_tensors.BoundingBoxes):
-                    item_values = tv_tensors.BoundingBoxes(item_values)
-                targets_modified[key_name] = item_values
+            if self.skip_key_name(key_name):
+                continue
+            first_element = item_values[0]
+            if isinstance(first_element, torch.Tensor):
+                item_values = torch.stack(item_values)
+            elif isinstance(first_element, (int, float)):
+                item_values = torch.tensor(item_values)
+            elif isinstance(first_element, (str, list)):
+                # Skip stacking for certain types such as strings and lists
+                pass
+            if isinstance(first_element, tv_tensors.Mask):
+                item_values = tv_tensors.Mask(item_values)
+            elif isinstance(first_element, tv_tensors.Image):
+                item_values = tv_tensors.Image(item_values)
+            elif isinstance(first_element, tv_tensors.BoundingBoxes):
+                item_values = tv_tensors.BoundingBoxes(item_values)
+            targets_modified[key_name] = item_values
 
         return images, targets_modified

hafnia/utils.py

@@ -14,7 +14,7 @@ from rich import print as rprint
 from hafnia.log import sys_logger, user_logger
 
 PATH_DATA = Path("./.data")
-PATH_DATASET = PATH_DATA / "datasets"
+PATH_DATASETS = PATH_DATA / "datasets"
 PATH_RECIPES = PATH_DATA / "recipes"
 FILENAME_HAFNIAIGNORE = ".hafniaignore"
 DEFAULT_IGNORE_SPECIFICATION = [

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),
+]