hafnia 0.4.2__py3-none-any.whl → 0.4.3__py3-none-any.whl

hafnia/dataset/hafnia_dataset.py

@@ -1,25 +1,15 @@
 from __future__ import annotations
 
-import collections
 import copy
-import json
 import shutil
 from dataclasses import dataclass
-from datetime import datetime
 from pathlib import Path
 from random import Random
 from typing import Any, Dict, List, Optional, Tuple, Type, Union
 
-import cv2
-import more_itertools
-import numpy as np
 import polars as pl
 from packaging.version import Version
-from PIL import Image
-from pydantic import BaseModel, Field, field_serializer, field_validator
-from rich.progress import track
 
-import hafnia
 from hafnia.dataset import dataset_helpers
 from hafnia.dataset.dataset_names import (
     DATASET_FILENAMES_REQUIRED,
@@ -35,470 +25,19 @@ from hafnia.dataset.dataset_names import (
     StorageFormat,
 )
 from hafnia.dataset.format_conversions import (
+    format_coco,
     format_image_classification_folder,
     format_yolo,
 )
+from hafnia.dataset.hafnia_dataset_types import DatasetInfo, Sample
 from hafnia.dataset.operations import (
     dataset_stats,
     dataset_transformations,
     table_transformations,
 )
-from hafnia.dataset.primitives import PRIMITIVE_TYPES, get_primitive_type_from_string
-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.polygon import Polygon
 from hafnia.dataset.primitives.primitive import Primitive
 from hafnia.log import user_logger
-
-
-class TaskInfo(BaseModel):
-    primitive: Type[Primitive] = Field(
-        description="Primitive class or string name of the primitive, e.g. 'Bbox' or 'bitmask'"
-    )
-    class_names: Optional[List[str]] = Field(default=None, description="Optional list of class names for the primitive")
-    name: Optional[str] = Field(
-        default=None,
-        description=(
-            "Optional name for the task. 'None' will use default name of the provided primitive. "
-            "e.g. Bbox ->'bboxes', Bitmask -> 'bitmasks' etc."
-        ),
-    )
-
-    def model_post_init(self, __context: Any) -> None:
-        if self.name is None:
-            self.name = self.primitive.default_task_name()
-
-    def get_class_index(self, class_name: str) -> int:
-        """Get class index for a given class name"""
-        if self.class_names is None:
-            raise ValueError(f"Task '{self.name}' has no class names defined.")
-        if class_name not in self.class_names:
-            raise ValueError(f"Class name '{class_name}' not found in task '{self.name}'.")
-        return self.class_names.index(class_name)
-
-    # The 'primitive'-field of type 'Type[Primitive]' is not supported by pydantic out-of-the-box as
-    # the 'Primitive' class is an abstract base class and for the actual primtives such as Bbox, Bitmask, Classification.
-    # Below magic functions ('ensure_primitive' and 'serialize_primitive') ensures that the 'primitive' field can
-    # correctly validate and serialize sub-classes (Bbox, Classification, ...).
-    @field_validator("primitive", mode="plain")
-    @classmethod
-    def ensure_primitive(cls, primitive: Any) -> Any:
-        if isinstance(primitive, str):
-            return get_primitive_type_from_string(primitive)
-
-        if issubclass(primitive, Primitive):
-            return primitive
-
-        raise ValueError(f"Primitive must be a string or a Primitive subclass, got {type(primitive)} instead.")
-
-    @field_serializer("primitive")
-    @classmethod
-    def serialize_primitive(cls, primitive: Type[Primitive]) -> str:
-        if not issubclass(primitive, Primitive):
-            raise ValueError(f"Primitive must be a subclass of Primitive, got {type(primitive)} instead.")
-        return primitive.__name__
-
-    @field_validator("class_names", mode="after")
-    @classmethod
-    def validate_unique_class_names(cls, class_names: Optional[List[str]]) -> Optional[List[str]]:
-        """Validate that class names are unique"""
-        if class_names is None:
-            return None
-        duplicate_class_names = set([name for name in class_names if class_names.count(name) > 1])
-        if duplicate_class_names:
-            raise ValueError(
-                f"Class names must be unique. The following class names appear multiple times: {duplicate_class_names}."
-            )
-        return class_names
-
-    def full_name(self) -> str:
-        """Get qualified name for the task: <primitive_name>:<task_name>"""
-        return f"{self.primitive.__name__}:{self.name}"
-
-    # To get unique hash value for TaskInfo objects
-    def __hash__(self) -> int:
-        class_names = self.class_names or []
-        return hash((self.name, self.primitive.__name__, tuple(class_names)))
-
-    def __eq__(self, other: Any) -> bool:
-        if not isinstance(other, TaskInfo):
-            return False
-        return self.name == other.name and self.primitive == other.primitive and self.class_names == other.class_names
-
-
-class DatasetInfo(BaseModel):
-    dataset_name: str = Field(description="Name of the dataset, e.g. 'coco'")
-    version: Optional[str] = Field(default=None, description="Version of the dataset")
-    tasks: List[TaskInfo] = Field(default=None, description="List of tasks in the dataset")
-    reference_bibtex: Optional[str] = Field(
-        default=None,
-        description="Optional, BibTeX reference to dataset publication",
-    )
-    reference_paper_url: Optional[str] = Field(
-        default=None,
-        description="Optional, URL to dataset publication",
-    )
-    reference_dataset_page: Optional[str] = Field(
-        default=None,
-        description="Optional, URL to the dataset page",
-    )
-    meta: Optional[Dict[str, Any]] = Field(default=None, description="Optional metadata about the dataset")
-    format_version: str = Field(
-        default=hafnia.__dataset_format_version__,
-        description="Version of the Hafnia dataset format. You should not set this manually.",
-    )
-    updated_at: datetime = Field(
-        default_factory=datetime.now,
-        description="Timestamp of the last update to the dataset info. You should not set this manually.",
-    )
-
-    @field_validator("tasks", mode="after")
-    @classmethod
-    def _validate_check_for_duplicate_tasks(cls, tasks: Optional[List[TaskInfo]]) -> List[TaskInfo]:
-        if tasks is None:
-            return []
-        task_name_counts = collections.Counter(task.name for task in tasks)
-        duplicate_task_names = [name for name, count in task_name_counts.items() if count > 1]
-        if duplicate_task_names:
-            raise ValueError(
-                f"Tasks must be unique. The following tasks appear multiple times: {duplicate_task_names}."
-            )
-        return tasks
-
-    @field_validator("format_version")
-    @classmethod
-    def _validate_format_version(cls, format_version: str) -> str:
-        try:
-            Version(format_version)
-        except Exception as e:
-            raise ValueError(f"Invalid format_version '{format_version}'. Must be a valid version string.") from e
-
-        if Version(format_version) > Version(hafnia.__dataset_format_version__):
-            user_logger.warning(
-                f"The loaded dataset format version '{format_version}' is newer than the format version "
-                f"'{hafnia.__dataset_format_version__}' used in your version of Hafnia. Please consider "
-                f"updating Hafnia package."
-            )
-        return format_version
-
-    @field_validator("version")
-    @classmethod
-    def _validate_version(cls, dataset_version: Optional[str]) -> Optional[str]:
-        if dataset_version is None:
-            return None
-
-        try:
-            Version(dataset_version)
-        except Exception as e:
-            raise ValueError(f"Invalid dataset_version '{dataset_version}'. Must be a valid version string.") from e
-
-        return dataset_version
-
-    def check_for_duplicate_task_names(self) -> List[TaskInfo]:
-        return self._validate_check_for_duplicate_tasks(self.tasks)
-
-    def write_json(self, path: Path, indent: Optional[int] = 4) -> None:
-        json_str = self.model_dump_json(indent=indent)
-        path.write_text(json_str)
-
-    @staticmethod
-    def from_json_file(path: Path) -> DatasetInfo:
-        json_str = path.read_text()
-
-        # TODO: Deprecated support for old dataset info without format_version
-        # Below 4 lines can be replaced by 'dataset_info = DatasetInfo.model_validate_json(json_str)'
-        # when all datasets include a 'format_version' field
-        json_dict = json.loads(json_str)
-        if "format_version" not in json_dict:
-            json_dict["format_version"] = "0.0.0"
-
-        if "updated_at" not in json_dict:
-            json_dict["updated_at"] = datetime.min.isoformat()
-        dataset_info = DatasetInfo.model_validate(json_dict)
-
-        return dataset_info
-
-    @staticmethod
-    def merge(info0: DatasetInfo, info1: DatasetInfo) -> DatasetInfo:
-        """
-        Merges two DatasetInfo objects into one and validates if they are compatible.
-        """
-        for task_ds0 in info0.tasks:
-            for task_ds1 in info1.tasks:
-                same_name = task_ds0.name == task_ds1.name
-                same_primitive = task_ds0.primitive == task_ds1.primitive
-                same_name_different_primitive = same_name and not same_primitive
-                if same_name_different_primitive:
-                    raise ValueError(
-                        f"Cannot merge datasets with different primitives for the same task name: "
-                        f"'{task_ds0.name}' has primitive '{task_ds0.primitive}' in dataset0 and "
-                        f"'{task_ds1.primitive}' in dataset1."
-                    )
-
-                is_same_name_and_primitive = same_name and same_primitive
-                if is_same_name_and_primitive:
-                    task_ds0_class_names = task_ds0.class_names or []
-                    task_ds1_class_names = task_ds1.class_names or []
-                    if task_ds0_class_names != task_ds1_class_names:
-                        raise ValueError(
-                            f"Cannot merge datasets with different class names for the same task name and primitive: "
-                            f"'{task_ds0.name}' with primitive '{task_ds0.primitive}' has class names "
-                            f"{task_ds0_class_names} in dataset0 and {task_ds1_class_names} in dataset1."
-                        )
-
-        if info1.format_version != info0.format_version:
-            user_logger.warning(
-                "Dataset format version of the two datasets do not match. "
-                f"'{info1.format_version}' vs '{info0.format_version}'."
-            )
-        dataset_format_version = info0.format_version
-        if hafnia.__dataset_format_version__ != dataset_format_version:
-            user_logger.warning(
-                f"Dataset format version '{dataset_format_version}' does not match the current "
-                f"Hafnia format version '{hafnia.__dataset_format_version__}'."
-            )
-        unique_tasks = set(info0.tasks + info1.tasks)
-        meta = (info0.meta or {}).copy()
-        meta.update(info1.meta or {})
-        return DatasetInfo(
-            dataset_name=info0.dataset_name + "+" + info1.dataset_name,
-            version=None,
-            tasks=list(unique_tasks),
-            meta=meta,
-            format_version=dataset_format_version,
-        )
-
-    def get_task_by_name(self, task_name: str) -> TaskInfo:
-        """
-        Get task by its name. Raises an error if the task name is not found or if multiple tasks have the same name.
-        """
-        tasks_with_name = [task for task in self.tasks if task.name == task_name]
-        if not tasks_with_name:
-            raise ValueError(f"Task with name '{task_name}' not found in dataset info.")
-        if len(tasks_with_name) > 1:
-            raise ValueError(f"Multiple tasks found with name '{task_name}'. This should not happen!")
-        return tasks_with_name[0]
-
-    def get_tasks_by_primitive(self, primitive: Union[Type[Primitive], str]) -> List[TaskInfo]:
-        """
-        Get all tasks by their primitive type.
-        """
-        if isinstance(primitive, str):
-            primitive = get_primitive_type_from_string(primitive)
-
-        tasks_with_primitive = [task for task in self.tasks if task.primitive == primitive]
-        return tasks_with_primitive
-
-    def get_task_by_primitive(self, primitive: Union[Type[Primitive], str]) -> TaskInfo:
-        """
-        Get task by its primitive type. Raises an error if the primitive type is not found or if multiple tasks
-        have the same primitive type.
-        """
-
-        tasks_with_primitive = self.get_tasks_by_primitive(primitive)
-        if len(tasks_with_primitive) == 0:
-            raise ValueError(f"Task with primitive {primitive} not found in dataset info.")
-        if len(tasks_with_primitive) > 1:
-            raise ValueError(
-                f"Multiple tasks found with primitive {primitive}. Use '{self.get_task_by_name.__name__}' instead."
-            )
-        return tasks_with_primitive[0]
-
-    def get_task_by_task_name_and_primitive(
-        self,
-        task_name: Optional[str],
-        primitive: Optional[Union[Type[Primitive], str]],
-    ) -> TaskInfo:
-        """
-        Logic to get a unique task based on the provided 'task_name' and/or 'primitive'.
-        If both 'task_name' and 'primitive' are None, the dataset must have only one task.
-        """
-        task = dataset_transformations.get_task_info_from_task_name_and_primitive(
-            tasks=self.tasks,
-            primitive=primitive,
-            task_name=task_name,
-        )
-        return task
-
-    def replace_task(self, old_task: TaskInfo, new_task: Optional[TaskInfo]) -> DatasetInfo:
-        dataset_info = self.model_copy(deep=True)
-        has_task = any(t for t in dataset_info.tasks if t.name == old_task.name and t.primitive == old_task.primitive)
-        if not has_task:
-            raise ValueError(f"Task '{old_task.__repr__()}' not found in dataset info.")
-
-        new_tasks = []
-        for task in dataset_info.tasks:
-            if task.name == old_task.name and task.primitive == old_task.primitive:
-                if new_task is None:
-                    continue  # Remove the task
-                new_tasks.append(new_task)
-            else:
-                new_tasks.append(task)
-
-        dataset_info.tasks = new_tasks
-        return dataset_info
-
-
-class Sample(BaseModel):
-    file_path: Optional[str] = Field(description="Path to the image/video file.")
-    height: int = Field(description="Height of the image")
-    width: int = Field(description="Width of the image")
-    split: str = Field(description="Split name, e.g., 'train', 'val', 'test'")
-    tags: List[str] = Field(
-        default_factory=list,
-        description="Tags for a given sample. Used for creating subsets of the dataset.",
-    )
-    storage_format: str = Field(
-        default=StorageFormat.IMAGE,
-        description="Storage format. Sample data is stored as image or inside a video or zip file.",
-    )
-    collection_index: Optional[int] = Field(default=None, description="Optional e.g. frame number for video datasets")
-    collection_id: Optional[str] = Field(default=None, description="Optional e.g. video name for video datasets")
-    remote_path: Optional[str] = Field(default=None, description="Optional remote path for the image, if applicable")
-    sample_index: Optional[int] = Field(
-        default=None,
-        description="Don't manually set this, it is used for indexing samples in the dataset.",
-    )
-    classifications: Optional[List[Classification]] = Field(
-        default=None, description="Optional list of classifications"
-    )
-    bboxes: Optional[List[Bbox]] = Field(default=None, description="Optional list of bounding boxes")
-    bitmasks: Optional[List[Bitmask]] = Field(default=None, description="Optional list of bitmasks")
-    polygons: Optional[List[Polygon]] = Field(default=None, description="Optional list of polygons")
-
-    attribution: Optional[Attribution] = Field(default=None, description="Attribution information for the image")
-    dataset_name: Optional[str] = Field(
-        default=None,
-        description=(
-            "Don't manually set this, it will be automatically defined during initialization. "
-            "Name of the dataset the sample belongs to. E.g. 'coco-2017' or 'midwest-vehicle-detection'."
-        ),
-    )
-    meta: Optional[Dict] = Field(
-        default=None,
-        description="Additional metadata, e.g., camera settings, GPS data, etc.",
-    )
-
-    def get_annotations(self, primitive_types: Optional[List[Type[Primitive]]] = None) -> List[Primitive]:
-        """
-        Returns a list of all annotations (classifications, objects, bitmasks, polygons) for the sample.
-        """
-        primitive_types = primitive_types or PRIMITIVE_TYPES
-        annotations_primitives = [
-            getattr(self, primitive_type.column_name(), None) for primitive_type in primitive_types
-        ]
-        annotations = more_itertools.flatten(
-            [primitives for primitives in annotations_primitives if primitives is not None]
-        )
-
-        return list(annotations)
-
-    def read_image_pillow(self) -> Image.Image:
-        """
-        Reads the image from the file path and returns it as a PIL Image.
-        Raises FileNotFoundError if the image file does not exist.
-        """
-        if self.file_path is None:
-            raise ValueError(f"Sample has no '{SampleField.FILE_PATH}' defined.")
-        path_image = Path(self.file_path)
-        if not path_image.exists():
-            raise FileNotFoundError(f"Image file {path_image} does not exist. Please check the file path.")
-
-        image = Image.open(str(path_image))
-        return image
-
-    def read_image(self) -> np.ndarray:
-        if self.storage_format == StorageFormat.VIDEO:
-            video = cv2.VideoCapture(str(self.file_path))
-            if self.collection_index is None:
-                raise ValueError("collection_index must be set for video storage format to read the correct frame.")
-            video.set(cv2.CAP_PROP_POS_FRAMES, self.collection_index)
-            success, image = video.read()
-            video.release()
-            if not success:
-                raise ValueError(f"Could not read frame {self.collection_index} from video file {self.file_path}.")
-            return image
-
-        elif self.storage_format == StorageFormat.IMAGE:
-            image_pil = self.read_image_pillow()
-            image = np.array(image_pil)
-        else:
-            raise ValueError(f"Unsupported storage format: {self.storage_format}")
-        return image
-
-    def draw_annotations(self, image: Optional[np.ndarray] = None) -> np.ndarray:
-        from hafnia.visualizations import image_visualizations
-
-        image = image or self.read_image()
-        annotations = self.get_annotations()
-        annotations_visualized = image_visualizations.draw_annotations(image=image, primitives=annotations)
-        return annotations_visualized
-
-
-class License(BaseModel):
-    """License information"""
-
-    name: Optional[str] = Field(
-        default=None,
-        description="License name. E.g. 'Creative Commons: Attribution 2.0 Generic'",
-        max_length=100,
-    )
-    name_short: Optional[str] = Field(
-        default=None,
-        description="License short name or abbreviation. E.g. 'CC BY 4.0'",
-        max_length=100,
-    )
-    url: Optional[str] = Field(
-        default=None,
-        description="License URL e.g. https://creativecommons.org/licenses/by/4.0/",
-    )
-    description: Optional[str] = Field(
-        default=None,
-        description=(
-            "License description e.g. 'You must give appropriate credit, provide a "
-            "link to the license, and indicate if changes were made.'"
-        ),
-    )
-
-    valid_date: Optional[datetime] = Field(
-        default=None,
-        description="License valid date. E.g. '2023-01-01T00:00:00Z'",
-    )
-
-    permissions: Optional[List[str]] = Field(
-        default=None,
-        description="License permissions. Allowed to Access, Label, Distribute, Represent and Modify data.",
-    )
-    liability: Optional[str] = Field(
-        default=None,
-        description="License liability. Optional and not always applicable.",
-    )
-    location: Optional[str] = Field(
-        default=None,
-        description=(
-            "License Location. E.g. Iowa state. This is essential to understand the industry and "
-            "privacy location specific rules that applies to the data. Optional and not always applicable."
-        ),
-    )
-    notes: Optional[str] = Field(
-        default=None,
-        description="Additional license notes. Optional and not always applicable.",
-    )
-
-
-class Attribution(BaseModel):
-    """Attribution information for the image: Giving source and credit to the original creator"""
-
-    title: Optional[str] = Field(default=None, description="Title of the image", max_length=255)
-    creator: Optional[str] = Field(default=None, description="Creator of the image", max_length=255)
-    creator_url: Optional[str] = Field(default=None, description="URL of the creator", max_length=255)
-    date_captured: Optional[datetime] = Field(default=None, description="Date when the image was captured")
-    copyright_notice: Optional[str] = Field(default=None, description="Copyright notice for the image", max_length=255)
-    licenses: Optional[List[License]] = Field(default=None, description="List of licenses for the image")
-    disclaimer: Optional[str] = Field(default=None, description="Disclaimer for the image", max_length=255)
-    changes: Optional[str] = Field(default=None, description="Changes made to the image", max_length=255)
-    source_url: Optional[str] = Field(default=None, description="Source URL for the image", max_length=255)
+from hafnia.utils import progress_bar
 
 
 @dataclass
@@ -527,8 +66,10 @@ class HafniaDataset:
     convert_to_image_storage_format = dataset_transformations.convert_to_image_storage_format
 
     # Import / export functions
-    from_yolo_format = format_yolo.from_yolo_format
     to_yolo_format = format_yolo.to_yolo_format
+    from_yolo_format = format_yolo.from_yolo_format
+    to_coco_format = format_coco.to_coco_format
+    from_coco_format = format_coco.from_coco_format
     to_image_classification_folder = format_image_classification_folder.to_image_classification_folder
     from_image_classification_folder = format_image_classification_folder.from_image_classification_folder
 
@@ -978,6 +519,10 @@ class HafniaDataset:
         dataset.check_dataset_tasks()
         return dataset
 
+    def has_primitive(dataset: HafniaDataset, PrimitiveType: Type[Primitive]) -> bool:
+        table = dataset.samples if isinstance(dataset, HafniaDataset) else dataset
+        return table_transformations.has_primitive(table, PrimitiveType)
+
     @staticmethod
     def check_dataset_path(path_dataset: Path, raise_error: bool = True) -> bool:
         """
@@ -1026,7 +571,7 @@ class HafniaDataset:
         hafnia_dataset = self.copy()  # To avoid inplace modifications
         new_paths = []
         org_paths = hafnia_dataset.samples[SampleField.FILE_PATH].to_list()
-        for org_path in track(org_paths, description="- Copy images"):
+        for org_path in progress_bar(org_paths, description="- Copy images"):
             new_path = dataset_helpers.copy_and_rename_file_to_hash_value(
                 path_source=Path(org_path),
                 path_dataset_root=path_folder,
@@ -1145,4 +690,13 @@ def _dataset_corrections(samples: pl.DataFrame, dataset_info: DatasetInfo) -> Tu
         if SampleField.SAMPLE_INDEX in samples.columns and samples[SampleField.SAMPLE_INDEX].dtype != pl.UInt64:
             samples = samples.cast({SampleField.SAMPLE_INDEX: pl.UInt64})
 
+    if format_version_of_dataset <= Version("0.2.0"):
+        if SampleField.BITMASKS in samples.columns and samples[SampleField.BITMASKS].dtype == pl.List(pl.Struct):
+            struct_schema = samples.schema[SampleField.BITMASKS].inner
+            struct_names = [f.name for f in struct_schema.fields]
+            if "rleString" in struct_names:
+                struct_names[struct_names.index("rleString")] = "rle_string"
+                samples = samples.with_columns(
+                    pl.col(SampleField.BITMASKS).list.eval(pl.element().struct.rename_fields(struct_names))
+                )
     return samples, dataset_info