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

hafnia/platform/download.py

@@ -1,15 +1,87 @@
 from pathlib import Path
-from typing import Any, Dict
+from typing import Dict
 
 import boto3
 from botocore.exceptions import ClientError
+from pydantic import BaseModel, field_validator
 from tqdm import tqdm
 
 from hafnia.http import fetch
 from hafnia.log import sys_logger, user_logger
 
-
-def get_resource_creds(endpoint: str, api_key: str) -> Dict[str, Any]:
+ARN_PREFIX = "arn:aws:s3:::"
+
+
+class ResourceCredentials(BaseModel):
+    access_key: str
+    secret_key: str
+    session_token: str
+    s3_arn: str
+    region: str
+
+    @staticmethod
+    def fix_naming(payload: Dict[str, str]) -> "ResourceCredentials":
+        """
+        The endpoint returns a payload with a key called 's3_path', but it
+        is actually an ARN path (starts with arn:aws:s3::). This method renames it to 's3_arn' for consistency.
+        """
+        if "s3_path" in payload and payload["s3_path"].startswith(ARN_PREFIX):
+            payload["s3_arn"] = payload.pop("s3_path")
+
+        if "region" not in payload:
+            payload["region"] = "eu-west-1"
+        return ResourceCredentials(**payload)
+
+    @field_validator("s3_arn")
+    @classmethod
+    def validate_s3_arn(cls, value: str) -> str:
+        """Validate s3_arn to ensure it starts with 'arn:aws:s3:::'"""
+        if not value.startswith("arn:aws:s3:::"):
+            raise ValueError(f"Invalid S3 ARN: {value}. It should start with 'arn:aws:s3:::'")
+        return value
+
+    def s3_path(self) -> str:
+        """
+        Extracts the S3 path from the ARN.
+        Example: arn:aws:s3:::my-bucket/my-prefix -> my-bucket/my-prefix
+        """
+        return self.s3_arn[len(ARN_PREFIX) :]
+
+    def s3_uri(self) -> str:
+        """
+        Converts the S3 ARN to a URI format.
+        Example: arn:aws:s3:::my-bucket/my-prefix -> s3://my-bucket/my-prefix
+        """
+        return f"s3://{self.s3_path()}"
+
+    def bucket_name(self) -> str:
+        """
+        Extracts the bucket name from the S3 ARN.
+        Example: arn:aws:s3:::my-bucket/my-prefix -> my-bucket
+        """
+        return self.s3_path().split("/")[0]
+
+    def object_key(self) -> str:
+        """
+        Extracts the object key from the S3 ARN.
+        Example: arn:aws:s3:::my-bucket/my-prefix -> my-prefix
+        """
+        return "/".join(self.s3_path().split("/")[1:])
+
+    def aws_credentials(self) -> Dict[str, str]:
+        """
+        Returns the AWS credentials as a dictionary.
+        """
+        environment_vars = {
+            "AWS_ACCESS_KEY_ID": self.access_key,
+            "AWS_SECRET_ACCESS_KEY": self.secret_key,
+            "AWS_SESSION_TOKEN": self.session_token,
+            "AWS_REGION": self.region,
+        }
+        return environment_vars
+
+
+def get_resource_credentials(endpoint: str, api_key: str) -> ResourceCredentials:
     """
     Retrieve credentials for accessing the recipe stored in S3 (or another resource)
     by calling a DIP endpoint with the API key.
@@ -18,21 +90,16 @@ def get_resource_creds(endpoint: str, api_key: str) -> Dict[str, Any]:
         endpoint (str): The endpoint URL to fetch credentials from.
 
     Returns:
-        Dict[str, Any]: Dictionary containing the credentials, for example:
-            {
-                "access_key": str,
-                "secret_key": str,
-                "session_token": str,
-                "s3_path": str
-            }
+        ResourceCredentials
 
     Raises:
         RuntimeError: If the call to fetch the credentials fails for any reason.
     """
     try:
-        creds = fetch(endpoint, headers={"Authorization": api_key, "accept": "application/json"})
+        credentials_dict = fetch(endpoint, headers={"Authorization": api_key, "accept": "application/json"})
+        credentials = ResourceCredentials.fix_naming(credentials_dict)
         sys_logger.debug("Successfully retrieved credentials from DIP endpoint.")
-        return creds
+        return credentials
     except Exception as e:
         sys_logger.error(f"Failed to fetch credentials from endpoint: {e}")
         raise RuntimeError(f"Failed to retrieve credentials: {e}") from e
@@ -76,23 +143,18 @@ def download_resource(resource_url: str, destination: str, api_key: str) -> Dict
         ValueError: If the S3 ARN is invalid or no objects found under prefix.
         RuntimeError: If S3 calls fail with an unexpected error.
     """
-    res_creds = get_resource_creds(resource_url, api_key)
-    s3_arn = res_creds["s3_path"]
-    arn_prefix = "arn:aws:s3:::"
-    if not s3_arn.startswith(arn_prefix):
-        raise ValueError(f"Invalid S3 ARN: {s3_arn}")
+    res_credentials = get_resource_credentials(resource_url, api_key)
 
-    s3_path = s3_arn[len(arn_prefix) :]
-    bucket_name, *key_parts = s3_path.split("/")
-    key = "/".join(key_parts)
+    bucket_name = res_credentials.bucket_name()
+    key = res_credentials.object_key()
 
     output_path = Path(destination)
     output_path.mkdir(parents=True, exist_ok=True)
     s3_client = boto3.client(
         "s3",
-        aws_access_key_id=res_creds["access_key"],
-        aws_secret_access_key=res_creds["secret_key"],
-        aws_session_token=res_creds["session_token"],
+        aws_access_key_id=res_credentials.access_key,
+        aws_secret_access_key=res_credentials.secret_key,
+        aws_session_token=res_credentials.session_token,
     )
     downloaded_files = []
     try:

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

@@ -1,3 +1,4 @@
+import hashlib
 import os
 import time
 import zipfile
@@ -14,7 +15,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 = [
@@ -132,6 +133,24 @@ def show_recipe_content(recipe_path: Path, style: str = "emoji", depth_limit: in
     user_logger.info(f"Recipe size: {size_human_readable(os.path.getsize(recipe_path))}. Max size 800 MiB")
 
 
-def is_remote_job() -> bool:
+def is_hafnia_cloud_job() -> bool:
     """Check if the current job is running in HAFNIA cloud environment."""
     return os.getenv("HAFNIA_CLOUD", "false").lower() == "true"
+
+
+def pascal_to_snake_case(name: str) -> str:
+    """
+    Convert PascalCase to snake_case.
+    """
+    return "".join(["_" + char.lower() if char.isupper() else char for char in name]).lstrip("_")
+
+
+def snake_to_pascal_case(name: str) -> str:
+    """
+    Convert snake_case to PascalCase.
+    """
+    return "".join(word.capitalize() for word in name.split("_"))
+
+
+def hash_from_string(s: str) -> str:
+    return hashlib.md5(s.encode("utf-8")).hexdigest()