hafnia 0.4.0__py3-none-any.whl → 0.4.1__py3-none-any.whl

hafnia/__init__.py

@@ -3,4 +3,4 @@ from importlib.metadata import version
 __package_name__ = "hafnia"
 __version__ = version(__package_name__)
 
-__dataset_format_version__ = "0.1.0"  # Hafnia dataset format version
+__dataset_format_version__ = "0.2.0"  # Hafnia dataset format version

hafnia/dataset/dataset_names.py

@@ -1,5 +1,8 @@
 from enum import Enum
-from typing import List
+from typing import Dict, List, Optional
+
+import boto3
+from pydantic import BaseModel, field_validator
 
 FILENAME_RECIPE_JSON = "recipe.json"
 FILENAME_DATASET_INFO = "dataset_info.json"
@@ -23,7 +26,7 @@ TAG_IS_SAMPLE = "sample"
 OPS_REMOVE_CLASS = "__REMOVE__"
 
 
-class FieldName:
+class PrimitiveField:
     CLASS_NAME: str = "class_name"  # Name of the class this primitive is associated with, e.g. "car" for Bbox
     CLASS_IDX: str = "class_idx"  # Index of the class this primitive is associated with, e.g. 0 for "car" if it is the first class  # noqa: E501
     OBJECT_ID: str = "object_id"  # Unique identifier for the object, e.g. "12345123"
@@ -38,40 +41,150 @@ class FieldName:
         Returns a list of expected field names for primitives.
         """
         return [
-            FieldName.CLASS_NAME,
-            FieldName.CLASS_IDX,
-            FieldName.OBJECT_ID,
-            FieldName.CONFIDENCE,
-            FieldName.META,
-            FieldName.TASK_NAME,
+            PrimitiveField.CLASS_NAME,
+            PrimitiveField.CLASS_IDX,
+            PrimitiveField.OBJECT_ID,
+            PrimitiveField.CONFIDENCE,
+            PrimitiveField.META,
+            PrimitiveField.TASK_NAME,
         ]
 
 
-class ColumnName:
-    SAMPLE_INDEX: str = "sample_index"
+class SampleField:
     FILE_PATH: str = "file_path"
     HEIGHT: str = "height"
     WIDTH: str = "width"
     SPLIT: str = "split"
+    TAGS: str = "tags"
+
+    CLASSIFICATIONS: str = "classifications"
+    BBOXES: str = "bboxes"
+    BITMASKS: str = "bitmasks"
+    POLYGONS: str = "polygons"
+
+    STORAGE_FORMAT: str = "storage_format"  # E.g. "image", "video", "zip"
+    COLLECTION_INDEX: str = "collection_index"
+    COLLECTION_ID: str = "collection_id"
     REMOTE_PATH: str = "remote_path"  # Path to the file in remote storage, e.g. S3
+    SAMPLE_INDEX: str = "sample_index"
+
     ATTRIBUTION: str = "attribution"  # Attribution for the sample (image/video), e.g. creator, license, source, etc.
-    TAGS: str = "tags"
     META: str = "meta"
     DATASET_NAME: str = "dataset_name"
 
 
+class StorageFormat:
+    IMAGE: str = "image"
+    VIDEO: str = "video"
+    ZIP: str = "zip"
+
+
 class SplitName:
-    TRAIN = "train"
-    VAL = "validation"
-    TEST = "test"
-    UNDEFINED = "UNDEFINED"
+    TRAIN: str = "train"
+    VAL: str = "validation"
+    TEST: str = "test"
+    UNDEFINED: str = "UNDEFINED"
 
     @staticmethod
     def valid_splits() -> List[str]:
         return [SplitName.TRAIN, SplitName.VAL, SplitName.TEST]
 
+    @staticmethod
+    def all_split_names() -> List[str]:
+        return [*SplitName.valid_splits(), SplitName.UNDEFINED]
+
 
 class DatasetVariant(Enum):
     DUMP = "dump"
     SAMPLE = "sample"
     HIDDEN = "hidden"
+
+
+class AwsCredentials(BaseModel):
+    access_key: str
+    secret_key: str
+    session_token: str
+    region: Optional[str]
+
+    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,
+        }
+        if self.region:
+            environment_vars["AWS_REGION"] = self.region
+
+        return environment_vars
+
+    @staticmethod
+    def from_session(session: boto3.Session) -> "AwsCredentials":
+        """
+        Creates AwsCredentials from a Boto3 session.
+        """
+        frozen_credentials = session.get_credentials().get_frozen_credentials()
+        return AwsCredentials(
+            access_key=frozen_credentials.access_key,
+            secret_key=frozen_credentials.secret_key,
+            session_token=frozen_credentials.token,
+            region=session.region_name,
+        )
+
+
+ARN_PREFIX = "arn:aws:s3:::"
+
+
+class ResourceCredentials(AwsCredentials):
+    s3_arn: 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:])

hafnia/dataset/dataset_upload_helper.py

@@ -14,10 +14,10 @@ from pydantic import BaseModel, ConfigDict, field_validator
 from cli.config import Config
 from hafnia.dataset import primitives
 from hafnia.dataset.dataset_names import (
-    ColumnName,
     DatasetVariant,
     DeploymentStage,
-    FieldName,
+    PrimitiveField,
+    SampleField,
     SplitName,
 )
 from hafnia.dataset.hafnia_dataset import Attribution, HafniaDataset, Sample, TaskInfo
@@ -193,7 +193,7 @@ class Annotations(BaseModel):
     in gallery images on the dataset detail page.
     """
 
-    objects: Optional[List[Bbox]] = None
+    bboxes: Optional[List[Bbox]] = None
     classifications: Optional[List[Classification]] = None
     polygons: Optional[List[Polygon]] = None
     bitmasks: Optional[List[Bitmask]] = None
@@ -210,13 +210,15 @@ class DatasetImageMetadata(BaseModel):
     @classmethod
     def from_sample(cls, sample: Sample) -> "DatasetImageMetadata":
         sample = sample.model_copy(deep=True)
+        if sample.file_path is None:
+            raise ValueError("Sample has no file_path defined.")
         sample.file_path = "/".join(Path(sample.file_path).parts[-3:])
         metadata = {}
         metadata_field_names = [
-            ColumnName.FILE_PATH,
-            ColumnName.HEIGHT,
-            ColumnName.WIDTH,
-            ColumnName.SPLIT,
+            SampleField.FILE_PATH,
+            SampleField.HEIGHT,
+            SampleField.WIDTH,
+            SampleField.SPLIT,
         ]
         for field_name in metadata_field_names:
             if hasattr(sample, field_name) and getattr(sample, field_name) is not None:
@@ -224,7 +226,7 @@ class DatasetImageMetadata(BaseModel):
 
         obj = DatasetImageMetadata(
             annotations=Annotations(
-                objects=sample.objects,
+                bboxes=sample.bboxes,
                 classifications=sample.classifications,
                 polygons=sample.polygons,
                 bitmasks=sample.bitmasks,
@@ -343,13 +345,13 @@ def calculate_distribution_values(
     classifications = dataset_split.select(pl.col(classification_column).explode())
     classifications = classifications.filter(pl.col(classification_column).is_not_null()).unnest(classification_column)
     classifications = classifications.filter(
-        pl.col(FieldName.TASK_NAME).is_in([task.name for task in distribution_tasks])
+        pl.col(PrimitiveField.TASK_NAME).is_in([task.name for task in distribution_tasks])
     )
     dist_values = []
-    for (task_name,), task_group in classifications.group_by(FieldName.TASK_NAME):
+    for (task_name,), task_group in classifications.group_by(PrimitiveField.TASK_NAME):
         distribution_type = DbDistributionType(name=task_name)
         n_annotated_total = len(task_group)
-        for (class_name,), class_group in task_group.group_by(FieldName.CLASS_NAME):
+        for (class_name,), class_group in task_group.group_by(PrimitiveField.CLASS_NAME):
             class_count = len(class_group)
 
             dist_values.append(
@@ -383,6 +385,7 @@ def dataset_info_from_dataset(
     path_hidden: Optional[Path],
     path_gallery_images: Optional[Path] = None,
     gallery_image_names: Optional[List[str]] = None,
+    distribution_task_names: Optional[List[TaskInfo]] = None,
 ) -> DbDataset:
     dataset_variants = []
     dataset_reports = []
@@ -427,13 +430,15 @@ def dataset_info_from_dataset(
             )
         )
 
+        distribution_task_names = distribution_task_names or []
+        distribution_tasks = [t for t in dataset.info.tasks if t.name in distribution_task_names]
         for split_name in SplitChoices:
             split_names = SPLIT_CHOICE_MAPPING[split_name]
-            dataset_split = dataset_variant.samples.filter(pl.col(ColumnName.SPLIT).is_in(split_names))
+            dataset_split = dataset_variant.samples.filter(pl.col(SampleField.SPLIT).is_in(split_names))
 
             distribution_values = calculate_distribution_values(
                 dataset_split=dataset_split,
-                distribution_tasks=dataset.info.distributions,
+                distribution_tasks=distribution_tasks,
             )
             report = DbSplitAnnotationsReport(
                 variant_type=VARIANT_TYPE_MAPPING[variant_type],  # type: ignore[index]
@@ -461,7 +466,7 @@ def dataset_info_from_dataset(
 
                 annotation_type = DbAnnotationType(name=AnnotationType.ObjectDetection.value)
                 for (class_name, task_name), class_group in df_per_instance.group_by(
-                    FieldName.CLASS_NAME, FieldName.TASK_NAME
+                    PrimitiveField.CLASS_NAME, PrimitiveField.TASK_NAME
                 ):
                     if class_name is None:
                         continue
@@ -473,10 +478,10 @@ def dataset_info_from_dataset(
                                 annotation_type=annotation_type,
                                 task_name=task_name,
                             ),
-                            unique_obj_ids=class_group[FieldName.OBJECT_ID].n_unique(),
+                            unique_obj_ids=class_group[PrimitiveField.OBJECT_ID].n_unique(),
                             obj_instances=len(class_group),
                             annotation_type=[annotation_type],
-                            images_with_obj=class_group[ColumnName.SAMPLE_INDEX].n_unique(),
+                            images_with_obj=class_group[SampleField.SAMPLE_INDEX].n_unique(),
                             area_avg_ratio=class_group["area"].mean(),
                             area_min_ratio=class_group["area"].min(),
                             area_max_ratio=class_group["area"].max(),
@@ -495,7 +500,7 @@ def dataset_info_from_dataset(
                             width_avg_px=class_group["width_px"].mean(),
                             width_min_px=int(class_group["width_px"].min()),
                             width_max_px=int(class_group["width_px"].max()),
-                            average_count_per_image=len(class_group) / class_group[ColumnName.SAMPLE_INDEX].n_unique(),
+                            average_count_per_image=len(class_group) / class_group[SampleField.SAMPLE_INDEX].n_unique(),
                         )
                     )
 
@@ -509,13 +514,13 @@ def dataset_info_from_dataset(
 
                     # Include only classification tasks that are defined in the dataset info
                     classification_df = classification_df.filter(
-                        pl.col(FieldName.TASK_NAME).is_in(classification_tasks)
+                        pl.col(PrimitiveField.TASK_NAME).is_in(classification_tasks)
                     )
 
                     for (
                         task_name,
                         class_name,
-                    ), class_group in classification_df.group_by(FieldName.TASK_NAME, FieldName.CLASS_NAME):
+                    ), class_group in classification_df.group_by(PrimitiveField.TASK_NAME, PrimitiveField.CLASS_NAME):
                         if class_name is None:
                             continue
                         if task_name == Classification.default_task_name():
@@ -544,7 +549,7 @@ def dataset_info_from_dataset(
             if has_primitive(dataset_split, PrimitiveType=Bitmask):
                 col_name = Bitmask.column_name()
                 drop_columns = [col for col in primitive_columns if col != col_name]
-                drop_columns.append(FieldName.META)
+                drop_columns.append(PrimitiveField.META)
 
                 df_per_instance = table_transformations.create_primitive_table(
                     dataset_split, PrimitiveType=Bitmask, keep_sample_data=True
@@ -562,7 +567,7 @@ def dataset_info_from_dataset(
 
                 annotation_type = DbAnnotationType(name=AnnotationType.InstanceSegmentation)
                 for (class_name, task_name), class_group in df_per_instance.group_by(
-                    FieldName.CLASS_NAME, FieldName.TASK_NAME
+                    PrimitiveField.CLASS_NAME, PrimitiveField.TASK_NAME
                 ):
                     if class_name is None:
                         continue
@@ -574,11 +579,11 @@ def dataset_info_from_dataset(
                                 annotation_type=annotation_type,
                                 task_name=task_name,
                             ),
-                            unique_obj_ids=class_group[FieldName.OBJECT_ID].n_unique(),
+                            unique_obj_ids=class_group[PrimitiveField.OBJECT_ID].n_unique(),
                             obj_instances=len(class_group),
                             annotation_type=[annotation_type],
-                            average_count_per_image=len(class_group) / class_group[ColumnName.SAMPLE_INDEX].n_unique(),
-                            images_with_obj=class_group[ColumnName.SAMPLE_INDEX].n_unique(),
+                            average_count_per_image=len(class_group) / class_group[SampleField.SAMPLE_INDEX].n_unique(),
+                            images_with_obj=class_group[SampleField.SAMPLE_INDEX].n_unique(),
                             area_avg_ratio=class_group["area"].mean(),
                             area_min_ratio=class_group["area"].min(),
                             area_max_ratio=class_group["area"].max(),
@@ -646,7 +651,7 @@ def create_gallery_images(
         path_gallery_images.mkdir(parents=True, exist_ok=True)
         COL_IMAGE_NAME = "image_name"
         samples = dataset.samples.with_columns(
-            dataset.samples[ColumnName.FILE_PATH].str.split("/").list.last().alias(COL_IMAGE_NAME)
+            dataset.samples[SampleField.FILE_PATH].str.split("/").list.last().alias(COL_IMAGE_NAME)
         )
         gallery_samples = samples.filter(pl.col(COL_IMAGE_NAME).is_in(gallery_image_names))
 

hafnia/dataset/format_conversions/{image_classification_from_directory.py → format_image_classification_folder.py}

@@ -1,23 +1,27 @@
 import shutil
 from pathlib import Path
-from typing import List, Optional
+from typing import TYPE_CHECKING, List, Optional
 
 import more_itertools
 import polars as pl
 from PIL import Image
 from rich.progress import track
 
-from hafnia.dataset.dataset_names import ColumnName, FieldName
-from hafnia.dataset.hafnia_dataset import DatasetInfo, HafniaDataset, Sample, TaskInfo
+from hafnia.dataset.dataset_names import PrimitiveField, SampleField
 from hafnia.dataset.primitives import Classification
 from hafnia.utils import is_image_file
 
+if TYPE_CHECKING:
+    from hafnia.dataset.hafnia_dataset import HafniaDataset
 
-def import_image_classification_directory_tree(
+
+def from_image_classification_folder(
     path_folder: Path,
     split: str,
     n_samples: Optional[int] = None,
-) -> HafniaDataset:
+) -> "HafniaDataset":
+    from hafnia.dataset.hafnia_dataset import DatasetInfo, HafniaDataset, Sample, TaskInfo
+
     class_folder_paths = [path for path in path_folder.iterdir() if path.is_dir()]
     class_names = sorted([folder.name for folder in class_folder_paths])  # Sort for determinism
 
@@ -62,8 +66,8 @@ def import_image_classification_directory_tree(
     return hafnia_dataset
 
 
-def export_image_classification_directory_tree(
-    dataset: HafniaDataset,
+def to_image_classification_folder(
+    dataset: "HafniaDataset",
     path_output: Path,
     task_name: Optional[str] = None,
     clean_folder: bool = False,
@@ -72,7 +76,7 @@ def export_image_classification_directory_tree(
 
     samples = dataset.samples.with_columns(
         pl.col(task.primitive.column_name())
-        .list.filter(pl.element().struct.field(FieldName.TASK_NAME) == task.name)
+        .list.filter(pl.element().struct.field(PrimitiveField.TASK_NAME) == task.name)
         .alias(task.primitive.column_name())
     )
 
@@ -95,11 +99,11 @@ def export_image_classification_directory_tree(
         if len(classifications) != 1:
             raise ValueError("Each sample should have exactly one classification.")
         classification = classifications[0]
-        class_name = classification[FieldName.CLASS_NAME].replace("/", "_")  # Avoid issues with subfolders
+        class_name = classification[PrimitiveField.CLASS_NAME].replace("/", "_")  # Avoid issues with subfolders
         path_class_folder = path_output / class_name
         path_class_folder.mkdir(parents=True, exist_ok=True)
 
-        path_image_org = Path(sample_dict[ColumnName.FILE_PATH])
+        path_image_org = Path(sample_dict[SampleField.FILE_PATH])
         path_image_new = path_class_folder / path_image_org.name
         shutil.copy2(path_image_org, path_image_new)
 

hafnia/dataset/format_conversions/format_yolo.py

@@ -0,0 +1,164 @@
+import shutil
+from pathlib import Path
+from typing import TYPE_CHECKING, List, Optional
+
+from PIL import Image
+from rich.progress import track
+
+from hafnia.dataset import primitives
+from hafnia.dataset.dataset_names import SplitName
+
+if TYPE_CHECKING:
+    from hafnia.dataset.hafnia_dataset import HafniaDataset
+
+FILENAME_YOLO_CLASS_NAMES = "obj.names"
+FILENAME_YOLO_IMAGES_TXT = "images.txt"
+
+
+def get_image_size(path: Path) -> tuple[int, int]:
+    with Image.open(path) as img:
+        return img.size  # (width, height)
+
+
+def from_yolo_format(
+    path_yolo_dataset: Path,
+    split_name: str = SplitName.UNDEFINED,
+    dataset_name: str = "yolo-dataset",
+    filename_class_names: str = FILENAME_YOLO_CLASS_NAMES,
+    filename_images_txt: str = FILENAME_YOLO_IMAGES_TXT,
+) -> "HafniaDataset":
+    """
+    Imports a YOLO (Darknet) formatted dataset as a HafniaDataset.
+    """
+    from hafnia.dataset.hafnia_dataset import DatasetInfo, HafniaDataset, Sample, TaskInfo
+
+    path_class_names = path_yolo_dataset / filename_class_names
+
+    if split_name not in SplitName.all_split_names():
+        raise ValueError(f"Invalid split name: {split_name}. Must be one of {SplitName.all_split_names()}")
+
+    if not path_class_names.exists():
+        raise FileNotFoundError(f"File with class names not found at '{path_class_names.resolve()}'.")
+
+    class_names_text = path_class_names.read_text()
+    if class_names_text.strip() == "":
+        raise ValueError(f"File with class names not found at '{path_class_names.resolve()}' is empty")
+
+    class_names = [class_name for class_name in class_names_text.splitlines() if class_name.strip() != ""]
+
+    if len(class_names) == 0:
+        raise ValueError(f"File with class names not found at '{path_class_names.resolve()}' has no class names")
+
+    path_images_txt = path_yolo_dataset / filename_images_txt
+
+    if not path_images_txt.exists():
+        raise FileNotFoundError(f"File with images not found at '{path_images_txt.resolve()}'")
+
+    images_txt_text = path_images_txt.read_text()
+    if len(images_txt_text.strip()) == 0:
+        raise ValueError(f"File is empty at '{path_images_txt.resolve()}'")
+
+    image_paths_raw = [line.strip() for line in images_txt_text.splitlines()]
+
+    samples: List[Sample] = []
+    for image_path_raw in track(image_paths_raw):
+        path_image = path_yolo_dataset / image_path_raw
+        if not path_image.exists():
+            raise FileNotFoundError(f"File with image not found at '{path_image.resolve()}'")
+        width, height = get_image_size(path_image)
+
+        path_label = path_image.with_suffix(".txt")
+        if not path_label.exists():
+            raise FileNotFoundError(f"File with labels not found at '{path_label.resolve()}'")
+
+        boxes: List[primitives.Bbox] = []
+        bbox_strings = path_label.read_text().splitlines()
+        for bbox_string in bbox_strings:
+            parts = bbox_string.strip().split()
+            if len(parts) != 5:
+                raise ValueError(f"Invalid bbox format in file {path_label.resolve()}: {bbox_string}")
+
+            class_idx = int(parts[0])
+            x_center, y_center, bbox_width, bbox_height = (float(value) for value in parts[1:5])
+
+            top_left_x = x_center - bbox_width / 2
+            top_left_y = y_center - bbox_height / 2
+
+            bbox = primitives.Bbox(
+                top_left_x=top_left_x,
+                top_left_y=top_left_y,
+                width=bbox_width,
+                height=bbox_height,
+                class_idx=class_idx,
+                class_name=class_names[class_idx] if 0 <= class_idx < len(class_names) else None,
+            )
+            boxes.append(bbox)
+
+        sample = Sample(
+            file_path=path_image.absolute().as_posix(),
+            height=height,
+            width=width,
+            split=split_name,
+            bboxes=boxes,
+        )
+        samples.append(sample)
+
+    tasks = [TaskInfo(primitive=primitives.Bbox, class_names=class_names)]
+    info = DatasetInfo(dataset_name=dataset_name, tasks=tasks)
+    hafnia_dataset = HafniaDataset.from_samples_list(samples, info=info)
+    return hafnia_dataset
+
+
+def to_yolo_format(
+    dataset: "HafniaDataset",
+    path_export_yolo_dataset: Path,
+    task_name: Optional[str] = None,
+):
+    """Exports a HafniaDataset as YOLO (Darknet) format."""
+    from hafnia.dataset.hafnia_dataset import Sample
+
+    bbox_task = dataset.info.get_task_by_task_name_and_primitive(task_name=task_name, primitive=primitives.Bbox)
+
+    class_names = bbox_task.class_names or []
+    if len(class_names) == 0:
+        raise ValueError(
+            f"Hafnia dataset task '{bbox_task.name}' has no class names defined. This is required for YOLO export."
+        )
+    path_export_yolo_dataset.mkdir(parents=True, exist_ok=True)
+    path_class_names = path_export_yolo_dataset / FILENAME_YOLO_CLASS_NAMES
+    path_class_names.write_text("\n".join(class_names))
+
+    path_data_folder = path_export_yolo_dataset / "data"
+    path_data_folder.mkdir(parents=True, exist_ok=True)
+    image_paths: List[str] = []
+    for sample_dict in dataset:
+        sample = Sample(**sample_dict)
+        if sample.file_path is None:
+            raise ValueError("Sample has no file_path defined.")
+        path_image_src = Path(sample.file_path)
+        path_image_dst = path_data_folder / path_image_src.name
+        shutil.copy2(path_image_src, path_image_dst)
+        image_paths.append(path_image_dst.relative_to(path_export_yolo_dataset).as_posix())
+        path_label = path_image_dst.with_suffix(".txt")
+        bboxes = sample.bboxes or []
+        bbox_strings = [bbox_to_yolo_format(bbox) for bbox in bboxes]
+        path_label.write_text("\n".join(bbox_strings))
+
+    path_images_txt = path_export_yolo_dataset / FILENAME_YOLO_IMAGES_TXT
+    path_images_txt.write_text("\n".join(image_paths))
+
+
+def bbox_to_yolo_format(bbox: primitives.Bbox) -> str:
+    """
+    From hafnia bbox to yolo bbox string conversion
+    Both yolo and hafnia use normalized coordinates [0, 1]
+        Hafnia: top_left_x, top_left_y, width, height
+        Yolo (darknet): "<object-class> <x_center> <y_center> <width> <height>"
+    Example (3 bounding boxes):
+        1 0.716797 0.395833 0.216406 0.147222
+        0 0.687109 0.379167 0.255469 0.158333
+        1 0.420312 0.395833 0.140625 0.166667
+    """
+    x_center = bbox.top_left_x + bbox.width / 2
+    y_center = bbox.top_left_y + bbox.height / 2
+    return f"{bbox.class_idx} {x_center} {y_center} {bbox.width} {bbox.height}"

hafnia/dataset/format_conversions/torchvision_datasets.py

@@ -14,8 +14,8 @@ from torchvision.datasets.utils import download_and_extract_archive, extract_arc
 from hafnia import utils
 from hafnia.dataset.dataset_helpers import save_pil_image_with_hash_name
 from hafnia.dataset.dataset_names import SplitName
-from hafnia.dataset.format_conversions.image_classification_from_directory import (
-    import_image_classification_directory_tree,
+from hafnia.dataset.format_conversions.format_image_classification_folder import (
+    from_image_classification_folder,
 )
 from hafnia.dataset.hafnia_dataset import DatasetInfo, HafniaDataset, Sample, TaskInfo
 from hafnia.dataset.primitives import Classification
@@ -72,7 +72,7 @@ def caltech_101_as_hafnia_dataset(
     path_image_classification_folder = _download_and_extract_caltech_dataset(
         dataset_name, force_redownload=force_redownload
     )
-    hafnia_dataset = import_image_classification_directory_tree(
+    hafnia_dataset = from_image_classification_folder(
         path_image_classification_folder,
         split=SplitName.TRAIN,
         n_samples=n_samples,
@@ -102,7 +102,7 @@ def caltech_256_as_hafnia_dataset(
     path_image_classification_folder = _download_and_extract_caltech_dataset(
         dataset_name, force_redownload=force_redownload
     )
-    hafnia_dataset = import_image_classification_directory_tree(
+    hafnia_dataset = from_image_classification_folder(
         path_image_classification_folder,
         split=SplitName.TRAIN,
         n_samples=n_samples,
@@ -122,6 +122,12 @@ def caltech_256_as_hafnia_dataset(
         }""")
     hafnia_dataset.info.reference_dataset_page = "https://data.caltech.edu/records/nyy15-4j048"
 
+    task = hafnia_dataset.info.get_task_by_primitive(Classification)
+
+    # Class Mapping: To remove numeric prefixes from class names
+    # E.g. "001.ak47 --> ak47", "002.american-flag --> american-flag", ...
+    class_mapping = {name: name.split(".")[-1] for name in task.class_names or []}
+    hafnia_dataset = hafnia_dataset.class_mapper(class_mapping=class_mapping, task_name=task.name)
     return hafnia_dataset