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

hafnia/dataset/dataset_upload_helper.py

@@ -4,7 +4,7 @@ import base64
 from datetime import datetime
 from enum import Enum
 from pathlib import Path
-from typing import Dict, List, Optional, Tuple, Type, Union
+from typing import Any, Dict, List, Optional, Tuple, Type, Union
 
 import boto3
 import polars as pl
@@ -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
@@ -52,6 +52,7 @@ class DbDataset(BaseModel, validate_assignment=True):  # type: ignore[call-arg]
     license_citation: Optional[str] = None
     version: Optional[str] = None
     s3_bucket_name: Optional[str] = None
+    dataset_format_version: Optional[str] = None
     annotation_date: Optional[datetime] = None
     annotation_project_id: Optional[str] = None
     annotation_dataset_id: Optional[str] = None
@@ -186,9 +187,60 @@ class EntityTypeChoices(str, Enum):  # Should match `EntityTypeChoices` in `dipd
     EVENT = "EVENT"
 
 
+class Annotations(BaseModel):
+    """
+    Used in 'DatasetImageMetadata' for visualizing image annotations
+    in gallery images on the dataset detail page.
+    """
+
+    bboxes: Optional[List[Bbox]] = None
+    classifications: Optional[List[Classification]] = None
+    polygons: Optional[List[Polygon]] = None
+    bitmasks: Optional[List[Bitmask]] = None
+
+
+class DatasetImageMetadata(BaseModel):
+    """
+    Metadata for gallery images on the dataset detail page on portal.
+    """
+
+    annotations: Optional[Annotations] = None
+    meta: Optional[Dict[str, Any]] = None
+
+    @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 = [
+            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:
+                metadata[field_name] = getattr(sample, field_name)
+
+        obj = DatasetImageMetadata(
+            annotations=Annotations(
+                bboxes=sample.bboxes,
+                classifications=sample.classifications,
+                polygons=sample.polygons,
+                bitmasks=sample.bitmasks,
+            ),
+            meta=metadata,
+        )
+
+        return obj
+
+
 class DatasetImage(Attribution, validate_assignment=True):  # type: ignore[call-arg]
     img: str  # Base64-encoded image string
     order: Optional[int] = None
+    metadata: Optional[DatasetImageMetadata] = None
 
     @field_validator("img", mode="before")
     def validate_image_path(cls, v: Union[str, Path]) -> str:
@@ -254,7 +306,7 @@ def upload_dataset_details(cfg: Config, data: str, dataset_name: str) -> dict:
     import_endpoint = f"{dataset_endpoint}/{dataset_id}/import"
     headers = {"Authorization": cfg.api_key}
 
-    user_logger.info("Importing dataset details. This may take up to 30 seconds...")
+    user_logger.info("Exporting dataset details to platform. This may take up to 30 seconds...")
     response = post(endpoint=import_endpoint, headers=headers, data=data)  # type: ignore[assignment]
     return response  # type: ignore[return-value]
 
@@ -293,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(
@@ -333,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 = []
@@ -377,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]
@@ -411,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
@@ -423,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(),
@@ -445,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(),
                         )
                     )
 
@@ -459,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():
@@ -494,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
@@ -512,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
@@ -524,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(),
@@ -569,7 +624,9 @@ def dataset_info_from_dataset(
         s3_bucket_name=bucket_sample,
         dataset_variants=dataset_variants,
         split_annotations_reports=dataset_reports,
-        license_citation=dataset_meta_info.get("license_citation", None),
+        latest_update=dataset.info.updated_at,
+        dataset_format_version=dataset.info.format_version,
+        license_citation=dataset.info.reference_bibtex,
         data_captured_start=dataset_meta_info.get("data_captured_start", None),
         data_captured_end=dataset_meta_info.get("data_captured_end", None),
         data_received_start=dataset_meta_info.get("data_received_start", None),
@@ -594,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_NAME].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))
 
@@ -604,6 +661,9 @@ def create_gallery_images(
         gallery_images = []
         for gallery_sample in gallery_samples.iter_rows(named=True):
             sample = Sample(**gallery_sample)
+
+            metadata = DatasetImageMetadata.from_sample(sample=sample)
+            sample.classifications = None  # To not draw classifications in gallery images
             image = sample.draw_annotations()
 
             path_gallery_image = path_gallery_images / gallery_sample[COL_IMAGE_NAME]
@@ -611,6 +671,7 @@ def create_gallery_images(
 
             dataset_image_dict = {
                 "img": path_gallery_image,
+                "metadata": metadata,
             }
             if sample.attribution is not None:
                 sample.attribution.changes = "Annotations have been visualized"

hafnia/dataset/format_conversions/format_image_classification_folder.py

@@ -0,0 +1,110 @@
+import shutil
+from pathlib import Path
+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 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 from_image_classification_folder(
+    path_folder: Path,
+    split: str,
+    n_samples: Optional[int] = None,
+) -> "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
+
+    # Gather all image paths per class
+    path_images_per_class: List[List[Path]] = []
+    for path_class_folder in class_folder_paths:
+        per_class_images = []
+        for path_image in list(path_class_folder.rglob("*.*")):
+            if is_image_file(path_image):
+                per_class_images.append(path_image)
+        path_images_per_class.append(sorted(per_class_images))
+
+    # Interleave to ensure classes are balanced in the output dataset for n_samples < total
+    path_images = list(more_itertools.interleave_longest(*path_images_per_class))
+
+    if n_samples is not None:
+        path_images = path_images[:n_samples]
+
+    samples = []
+    for path_image_org in track(path_images, description="Convert 'image classification' dataset to Hafnia Dataset"):
+        class_name = path_image_org.parent.name
+
+        read_image = Image.open(path_image_org)
+        width, height = read_image.size
+
+        classifications = [Classification(class_name=class_name, class_idx=class_names.index(class_name))]
+        sample = Sample(
+            file_path=str(path_image_org.absolute()),
+            width=width,
+            height=height,
+            split=split,
+            classifications=classifications,
+        )
+        samples.append(sample)
+
+    dataset_info = DatasetInfo(
+        dataset_name="ImageClassificationFromDirectoryTree",
+        tasks=[TaskInfo(primitive=Classification, class_names=class_names)],
+    )
+
+    hafnia_dataset = HafniaDataset.from_samples_list(samples_list=samples, info=dataset_info)
+    return hafnia_dataset
+
+
+def to_image_classification_folder(
+    dataset: "HafniaDataset",
+    path_output: Path,
+    task_name: Optional[str] = None,
+    clean_folder: bool = False,
+) -> Path:
+    task = dataset.info.get_task_by_task_name_and_primitive(task_name=task_name, primitive=Classification)
+
+    samples = dataset.samples.with_columns(
+        pl.col(task.primitive.column_name())
+        .list.filter(pl.element().struct.field(PrimitiveField.TASK_NAME) == task.name)
+        .alias(task.primitive.column_name())
+    )
+
+    classification_counts = samples[task.primitive.column_name()].list.len()
+    has_no_classification_samples = (classification_counts == 0).sum()
+    if has_no_classification_samples > 0:
+        raise ValueError(f"Some samples do not have a classification for task '{task.name}'.")
+
+    has_multi_classification_samples = (classification_counts > 1).sum()
+    if has_multi_classification_samples > 0:
+        raise ValueError(f"Some samples have multiple classifications for task '{task.name}'.")
+
+    if clean_folder:
+        shutil.rmtree(path_output, ignore_errors=True)
+    path_output.mkdir(parents=True, exist_ok=True)
+
+    description = "Export Hafnia Dataset to directory tree"
+    for sample_dict in track(samples.iter_rows(named=True), total=len(samples), description=description):
+        classifications = sample_dict[task.primitive.column_name()]
+        if len(classifications) != 1:
+            raise ValueError("Each sample should have exactly one classification.")
+        classification = classifications[0]
+        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[SampleField.FILE_PATH])
+        path_image_new = path_class_folder / path_image_org.name
+        shutil.copy2(path_image_org, path_image_new)
+
+    return path_output

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}"