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

hafnia/dataset/hafnia_dataset_types.py

@@ -0,0 +1,477 @@
+import collections
+import json
+from datetime import datetime
+from pathlib import Path
+from typing import Any, Dict, List, Optional, Type, Union
+
+import cv2
+import more_itertools
+import numpy as np
+from packaging.version import Version
+from PIL import Image
+from pydantic import BaseModel, Field, field_serializer, field_validator
+
+import hafnia
+from hafnia.dataset.dataset_names import SampleField, StorageFormat
+from hafnia.dataset.primitives import (
+    PRIMITIVE_TYPES,
+    Bbox,
+    Bitmask,
+    Classification,
+    Polygon,
+    get_primitive_type_from_string,
+)
+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.
+        """
+        from hafnia.dataset.operations import dataset_transformations
+
+        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 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)
+
+
+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
+
+        if image is None:
+            image = self.read_image()
+        annotations = self.get_annotations()
+        annotations_visualized = image_visualizations.draw_annotations(image=image, primitives=annotations)
+        return annotations_visualized

hafnia/dataset/license_types.py

@@ -1,6 +1,6 @@
-from typing import List, Optional
+from typing import List
 
-from hafnia.dataset.hafnia_dataset import License
+from hafnia.dataset.hafnia_dataset_types import License
 
 LICENSE_TYPES: List[License] = [
     License(
@@ -46,7 +46,7 @@ LICENSE_TYPES: List[License] = [
 ]
 
 
-def get_license_by_url(url: str) -> Optional[License]:
+def get_license_by_url(url: str) -> License:
     for license in LICENSE_TYPES:
         # To handle http urls
         license_url = (license.url or "").replace("http://", "https://")
@@ -56,7 +56,7 @@ def get_license_by_url(url: str) -> Optional[License]:
     raise ValueError(f"License with URL '{url}' not found.")
 
 
-def get_license_by_short_name(short_name: str) -> Optional[License]:
+def get_license_by_short_name(short_name: str) -> License:
     for license in LICENSE_TYPES:
         if license.name_short == short_name:
             return license

hafnia/dataset/operations/dataset_stats.py

@@ -5,13 +5,14 @@ from typing import TYPE_CHECKING, Any, Dict, List, Optional, Type
 import polars as pl
 import rich
 from rich import print as rprint
-from rich.progress import track
 from rich.table import Table
 
 from hafnia.dataset.dataset_names import PrimitiveField, SampleField, SplitName
+from hafnia.dataset.hafnia_dataset_types import Sample
 from hafnia.dataset.operations.table_transformations import create_primitive_table
 from hafnia.dataset.primitives import PRIMITIVE_TYPES
 from hafnia.log import user_logger
+from hafnia.utils import progress_bar
 
 if TYPE_CHECKING:  # Using 'TYPE_CHECKING' to avoid circular imports during type checking
     from hafnia.dataset.hafnia_dataset import HafniaDataset
@@ -188,7 +189,6 @@ def check_dataset(dataset: HafniaDataset, check_splits: bool = True):
     Performs various checks on the dataset to ensure its integrity and consistency.
     Raises errors if any issues are found.
     """
-    from hafnia.dataset.hafnia_dataset import Sample
 
     user_logger.info("Checking Hafnia dataset...")
     assert isinstance(dataset.info.dataset_name, str) and len(dataset.info.dataset_name) > 0
@@ -226,7 +226,7 @@ def check_dataset(dataset: HafniaDataset, check_splits: bool = True):
                 f"classes: {class_names}. "
             )
 
-    for sample_dict in track(dataset, description="Checking samples in dataset"):
+    for sample_dict in progress_bar(dataset, description="Checking samples in dataset"):
         sample = Sample(**sample_dict)  # noqa: F841
 
 

hafnia/dataset/operations/dataset_transformations.py

@@ -31,7 +31,6 @@ that the signatures match.
 
 import json
 import re
-import shutil
 import textwrap
 from pathlib import Path
 from typing import TYPE_CHECKING, Callable, Dict, List, Optional, Tuple, Type, Union
@@ -40,7 +39,6 @@ import cv2
 import more_itertools
 import numpy as np
 import polars as pl
-from rich.progress import track
 
 from hafnia.dataset import dataset_helpers
 from hafnia.dataset.dataset_names import (
@@ -49,14 +47,15 @@ from hafnia.dataset.dataset_names import (
     SampleField,
     StorageFormat,
 )
+from hafnia.dataset.hafnia_dataset_types import Sample, TaskInfo
 from hafnia.dataset.operations.table_transformations import update_class_indices
 from hafnia.dataset.primitives import get_primitive_type_from_string
 from hafnia.dataset.primitives.primitive import Primitive
 from hafnia.log import user_logger
-from hafnia.utils import remove_duplicates_preserve_order
+from hafnia.utils import progress_bar, remove_duplicates_preserve_order
 
 if TYPE_CHECKING:  # Using 'TYPE_CHECKING' to avoid circular imports during type checking
-    from hafnia.dataset.hafnia_dataset import HafniaDataset, Sample, TaskInfo
+    from hafnia.dataset.hafnia_dataset import HafniaDataset
 
 
 ### Image transformations ###
@@ -64,7 +63,7 @@ class AnonymizeByPixelation:
     def __init__(self, resize_factor: float = 0.10):
         self.resize_factor = resize_factor
 
-    def __call__(self, frame: np.ndarray, sample: "Sample") -> np.ndarray:
+    def __call__(self, frame: np.ndarray, sample: Sample) -> np.ndarray:
         org_size = frame.shape[:2]
         frame = cv2.resize(frame, (0, 0), fx=self.resize_factor, fy=self.resize_factor)
         frame = cv2.resize(frame, org_size[::-1], interpolation=cv2.INTER_NEAREST)
@@ -73,17 +72,15 @@ class AnonymizeByPixelation:
 
 def transform_images(
     dataset: "HafniaDataset",
-    transform: Callable[[np.ndarray, "Sample"], np.ndarray],
+    transform: Callable[[np.ndarray, Sample], np.ndarray],
     path_output: Path,
     description: str = "Transform images",
 ) -> "HafniaDataset":
-    from hafnia.dataset.hafnia_dataset import Sample
-
     new_paths = []
     path_image_folder = path_output / "data"
     path_image_folder.mkdir(parents=True, exist_ok=True)
 
-    for sample_dict in track(dataset, description=description):
+    for sample_dict in progress_bar(dataset, description=description):
         sample = Sample(**sample_dict)
         image = sample.read_image()
         image_transformed = transform(image, sample)
@@ -102,15 +99,15 @@ def convert_to_image_storage_format(
     path_output_folder: Path,
     reextract_frames: bool,
     image_format: str = "png",
-    transform: Optional[Callable[[np.ndarray, "Sample"], np.ndarray]] = None,
+    transform: Optional[Callable[[np.ndarray, Sample], np.ndarray]] = None,
 ) -> "HafniaDataset":
     """
     Convert a video-based dataset ("storage_format" == "video", FieldName.STORAGE_FORMAT == StorageFormat.VIDEO)
     to an image-based dataset by extracting frames.
     """
-    from hafnia.dataset.hafnia_dataset import HafniaDataset, Sample
+    from hafnia.dataset.hafnia_dataset import HafniaDataset
 
-    path_images = path_output_folder / "data"
+    path_images = (path_output_folder / "data").absolute()
     path_images.mkdir(parents=True, exist_ok=True)
 
     # Only video format dataset samples are processed
@@ -128,7 +125,7 @@ def convert_to_image_storage_format(
         video = cv2.VideoCapture(str(path_video))
 
         video_samples = video_samples.sort(SampleField.COLLECTION_INDEX)
-        for sample_dict in track(
+        for sample_dict in progress_bar(
             video_samples.iter_rows(named=True),
             total=video_samples.height,
             description=f"Extracting frames from '{Path(path_video).name}'",
@@ -147,7 +144,7 @@ def convert_to_image_storage_format(
                 }
             )
             if reextract_frames:
-                shutil.rmtree(path_image, ignore_errors=True)
+                path_image.unlink(missing_ok=True)
             if path_image.exists():
                 continue
 
@@ -168,10 +165,10 @@ def convert_to_image_storage_format(
 
 
 def get_task_info_from_task_name_and_primitive(
-    tasks: List["TaskInfo"],
+    tasks: List[TaskInfo],
     task_name: Optional[str] = None,
     primitive: Union[None, str, Type[Primitive]] = None,
-) -> "TaskInfo":
+) -> TaskInfo:
     if len(tasks) == 0:
         raise ValueError("Dataset has no tasks defined.")
 
@@ -423,7 +420,7 @@ def _validate_inputs_select_samples_by_class_name(
     name: Union[List[str], str],
     task_name: Optional[str] = None,
     primitive: Optional[Type[Primitive]] = None,
-) -> Tuple["TaskInfo", List[str]]:
+) -> Tuple[TaskInfo, List[str]]:
     if isinstance(name, str):
         name = [name]
     names = list(name)