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

hafnia/dataset/hafnia_dataset.py

@@ -8,14 +8,16 @@ from dataclasses import dataclass
 from datetime import datetime
 from pathlib import Path
 from random import Random
-from typing import Any, Dict, List, Optional, Type, Union
+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 tqdm import tqdm
+from rich.progress import track
 
 import hafnia
 from hafnia.dataset import dataset_helpers
@@ -26,13 +28,20 @@ from hafnia.dataset.dataset_names import (
     FILENAME_DATASET_INFO,
     FILENAME_RECIPE_JSON,
     TAG_IS_SAMPLE,
-    ColumnName,
+    AwsCredentials,
+    PrimitiveField,
+    SampleField,
     SplitName,
+    StorageFormat,
 )
-from hafnia.dataset.operations import dataset_stats, dataset_transformations, table_transformations
-from hafnia.dataset.operations.table_transformations import (
-    check_image_paths,
-    read_table_from_path,
+from hafnia.dataset.format_conversions import (
+    format_image_classification_folder,
+    format_yolo,
+)
+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
@@ -44,14 +53,30 @@ from hafnia.log import user_logger
 
 
 class TaskInfo(BaseModel):
-    primitive: Type[Primitive]  # Primitive class or string name of the primitive, e.g. "Bbox" or "bitmask"
-    class_names: Optional[List[str]]  # Class names for the tasks. To get consistent class indices specify class_names.
-    name: Optional[str] = None  # Use 'None' to use default name Bbox ->"bboxes", Bitmask -> "bitmasks" etc.
+    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
@@ -87,6 +112,10 @@ class TaskInfo(BaseModel):
             )
         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 []
@@ -99,17 +128,36 @@ class TaskInfo(BaseModel):
 
 
 class DatasetInfo(BaseModel):
-    dataset_name: str
-    version: str  # Dataset version. This is not the same as the Hafnia dataset format version.
-    tasks: List[TaskInfo]
-    distributions: Optional[List[TaskInfo]] = None  # Distributions. TODO: FIX/REMOVE/CHANGE this
-    meta: Optional[Dict[str, Any]] = None  # Metadata about the dataset, e.g. description, etc.
-    format_version: str = hafnia.__dataset_format_version__  # Version of the Hafnia dataset format
-    updated_at: datetime = datetime.now()
+    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: List[TaskInfo]) -> List[TaskInfo]:
+    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:
@@ -118,6 +166,35 @@ class DatasetInfo(BaseModel):
             )
         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)
 
@@ -182,14 +259,12 @@ class DatasetInfo(BaseModel):
                 f"Hafnia format version '{hafnia.__dataset_format_version__}'."
             )
         unique_tasks = set(info0.tasks + info1.tasks)
-        distributions = set((info0.distributions or []) + (info1.distributions or []))
         meta = (info0.meta or {}).copy()
         meta.update(info1.meta or {})
         return DatasetInfo(
             dataset_name=info0.dataset_name + "+" + info1.dataset_name,
-            version="merged",
+            version=None,
             tasks=list(unique_tasks),
-            distributions=list(distributions),
             meta=meta,
             format_version=dataset_format_version,
         )
@@ -205,16 +280,24 @@ class DatasetInfo(BaseModel):
             raise ValueError(f"Multiple tasks found with name '{task_name}'. This should not happen!")
         return tasks_with_name[0]
 
-    def get_task_by_primitive(self, primitive: Union[Type[Primitive], str]) -> TaskInfo:
+    def get_tasks_by_primitive(self, primitive: Union[Type[Primitive], str]) -> List[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.
+        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]
-        if not tasks_with_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(
@@ -258,22 +341,44 @@ class DatasetInfo(BaseModel):
 
 
 class Sample(BaseModel):
-    file_name: str
-    height: int
-    width: int
-    split: str  # Split name, e.g., "train", "val", "test"
-    tags: List[str] = []  # tags for a given sample. Used for creating subsets of the dataset.
-    collection_index: Optional[int] = None  # Optional e.g. frame number for video datasets
-    collection_id: Optional[str] = None  # Optional e.g. video name for video datasets
-    remote_path: Optional[str] = None  # Optional remote path for the image, if applicable
-    sample_index: Optional[int] = None  # Don't manually set this, it is used for indexing samples in the dataset.
-    classifications: Optional[List[Classification]] = None  # Optional classification primitive
-    objects: Optional[List[Bbox]] = None  # List of coordinate primitives, e.g., Bbox, Bitmask, etc.
-    bitmasks: Optional[List[Bitmask]] = None  # List of bitmasks, if applicable
-    polygons: Optional[List[Polygon]] = None  # List of polygons, if applicable
-
-    attribution: Optional[Attribution] = None  # Attribution information for the image
-    meta: Optional[Dict] = None  # Additional metadata, e.g., camera settings, GPS data, etc.
+    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]:
         """
@@ -294,7 +399,9 @@ class Sample(BaseModel):
         Reads the image from the file path and returns it as a PIL Image.
         Raises FileNotFoundError if the image file does not exist.
         """
-        path_image = Path(self.file_name)
+        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.")
 
@@ -302,8 +409,22 @@ class Sample(BaseModel):
         return image
 
     def read_image(self) -> np.ndarray:
-        image_pil = self.read_image_pillow()
-        image = np.array(image_pil)
+        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:
@@ -386,9 +507,11 @@ class HafniaDataset:
     samples: pl.DataFrame
 
     # Function mapping: Dataset stats
-    split_counts = dataset_stats.split_counts
-    class_counts_for_task = dataset_stats.class_counts_for_task
-    class_counts_all = dataset_stats.class_counts_all
+    calculate_split_counts = dataset_stats.calculate_split_counts
+    calculate_split_counts_extended = dataset_stats.calculate_split_counts_extended
+    calculate_task_class_counts = dataset_stats.calculate_task_class_counts
+    calculate_class_counts = dataset_stats.calculate_class_counts
+    calculate_primitive_counts = dataset_stats.calculate_primitive_counts
 
     # Function mapping: Print stats
     print_stats = dataset_stats.print_stats
@@ -401,6 +524,13 @@ class HafniaDataset:
 
     # Function mapping: Dataset transformations
     transform_images = dataset_transformations.transform_images
+    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
+    to_image_classification_folder = format_image_classification_folder.to_image_classification_folder
+    from_image_classification_folder = format_image_classification_folder.from_image_classification_folder
 
     def __getitem__(self, item: int) -> Dict[str, Any]:
         return self.samples.row(index=item, named=True)
@@ -413,30 +543,23 @@ class HafniaDataset:
             yield row
 
     def __post_init__(self):
-        samples = self.samples
-        if ColumnName.SAMPLE_INDEX not in samples.columns:
-            samples = samples.with_row_index(name=ColumnName.SAMPLE_INDEX)
-
-        # Backwards compatibility: If tags-column doesn't exist, create it with empty lists
-        if ColumnName.TAGS not in samples.columns:
-            tags_column: List[List[str]] = [[] for _ in range(len(self))]  # type: ignore[annotation-unchecked]
-            samples = samples.with_columns(pl.Series(tags_column, dtype=pl.List(pl.String)).alias(ColumnName.TAGS))
-
-        self.samples = samples
+        self.samples, self.info = _dataset_corrections(self.samples, self.info)
 
     @staticmethod
     def from_path(path_folder: Path, check_for_images: bool = True) -> "HafniaDataset":
+        path_folder = Path(path_folder)
         HafniaDataset.check_dataset_path(path_folder, raise_error=True)
 
         dataset_info = DatasetInfo.from_json_file(path_folder / FILENAME_DATASET_INFO)
-        table = read_table_from_path(path_folder)
+        samples = table_transformations.read_samples_from_path(path_folder)
+        samples, dataset_info = _dataset_corrections(samples, dataset_info)
 
         # Convert from relative paths to absolute paths
         dataset_root = path_folder.absolute().as_posix() + "/"
-        table = table.with_columns((dataset_root + pl.col("file_name")).alias("file_name"))
+        samples = samples.with_columns((dataset_root + pl.col(SampleField.FILE_PATH)).alias(SampleField.FILE_PATH))
         if check_for_images:
-            check_image_paths(table)
-        return HafniaDataset(samples=table, info=dataset_info)
+            table_transformations.check_image_paths(samples)
+        return HafniaDataset(samples=samples, info=dataset_info)
 
     @staticmethod
     def from_name(name: str, force_redownload: bool = False, download_files: bool = True) -> "HafniaDataset":
@@ -462,8 +585,12 @@ class HafniaDataset:
         else:
             raise TypeError(f"Unsupported sample type: {type(sample)}. Expected Sample or dict.")
 
-        table = pl.from_records(json_samples)
-        table = table.drop(ColumnName.SAMPLE_INDEX).with_row_index(name=ColumnName.SAMPLE_INDEX)
+        # To ensure that the 'file_path' column is of type string even if all samples have 'None' as file_path
+        schema_override = {SampleField.FILE_PATH: pl.String}
+        table = pl.from_records(json_samples, schema_overrides=schema_override)
+        table = table.drop(pl.selectors.by_dtype(pl.Null))
+        table = table_transformations.add_sample_index(table)
+        table = table_transformations.add_dataset_name_if_missing(table, dataset_name=info.dataset_name)
         return HafniaDataset(info=info, samples=table)
 
     @staticmethod
@@ -518,6 +645,28 @@ class HafniaDataset:
             merged_dataset = HafniaDataset.merge(merged_dataset, dataset)
         return merged_dataset
 
+    @staticmethod
+    def from_name_public_dataset(
+        name: str,
+        force_redownload: bool = False,
+        n_samples: Optional[int] = None,
+    ) -> HafniaDataset:
+        from hafnia.dataset.format_conversions.torchvision_datasets import (
+            torchvision_to_hafnia_converters,
+        )
+
+        name_to_torchvision_function = torchvision_to_hafnia_converters()
+
+        if name not in name_to_torchvision_function:
+            raise ValueError(
+                f"Unknown torchvision dataset name: {name}. Supported: {list(name_to_torchvision_function.keys())}"
+            )
+        vision_dataset = name_to_torchvision_function[name]
+        return vision_dataset(
+            force_redownload=force_redownload,
+            n_samples=n_samples,
+        )
+
     def shuffle(dataset: HafniaDataset, seed: int = 42) -> HafniaDataset:
         table = dataset.samples.sample(n=len(dataset), with_replacement=False, seed=seed, shuffle=True)
         return dataset.update_samples(table)
@@ -575,12 +724,12 @@ class HafniaDataset:
         """
         dataset_split_to_be_divided = dataset.create_split_dataset(split_name=split_name)
         if len(dataset_split_to_be_divided) == 0:
-            split_counts = dict(dataset.samples.select(pl.col(ColumnName.SPLIT).value_counts()).iter_rows())
+            split_counts = dict(dataset.samples.select(pl.col(SampleField.SPLIT).value_counts()).iter_rows())
             raise ValueError(f"No samples in the '{split_name}' split to divide into multiple splits. {split_counts=}")
         assert len(dataset_split_to_be_divided) > 0, f"No samples in the '{split_name}' split!"
         dataset_split_to_be_divided = dataset_split_to_be_divided.splits_by_ratios(split_ratios=split_ratios, seed=42)
 
-        remaining_data = dataset.samples.filter(pl.col(ColumnName.SPLIT).is_in([split_name]).not_())
+        remaining_data = dataset.samples.filter(pl.col(SampleField.SPLIT).is_in([split_name]).not_())
         new_table = pl.concat([remaining_data, dataset_split_to_be_divided.samples], how="vertical")
         dataset_new = dataset.update_samples(new_table)
         return dataset_new
@@ -593,21 +742,23 @@ class HafniaDataset:
 
         # Remove any pre-existing "sample"-tags
         samples = samples.with_columns(
-            pl.col(ColumnName.TAGS).list.eval(pl.element().filter(pl.element() != TAG_IS_SAMPLE)).alias(ColumnName.TAGS)
+            pl.col(SampleField.TAGS)
+            .list.eval(pl.element().filter(pl.element() != TAG_IS_SAMPLE))
+            .alias(SampleField.TAGS)
         )
 
         # Add "sample" to tags column for the selected samples
         is_sample_indices = Random(seed).sample(range(len(dataset)), n_samples)
         samples = samples.with_columns(
             pl.when(pl.int_range(len(samples)).is_in(is_sample_indices))
-            .then(pl.col(ColumnName.TAGS).list.concat(pl.lit([TAG_IS_SAMPLE])))
-            .otherwise(pl.col(ColumnName.TAGS))
+            .then(pl.col(SampleField.TAGS).list.concat(pl.lit([TAG_IS_SAMPLE])))
+            .otherwise(pl.col(SampleField.TAGS))
         )
         return dataset.update_samples(samples)
 
     def class_mapper(
         dataset: "HafniaDataset",
-        class_mapping: Dict[str, str],
+        class_mapping: Union[Dict[str, str], List[Tuple[str, str]]],
         method: str = "strict",
         primitive: Optional[Type[Primitive]] = None,
         task_name: Optional[str] = None,
@@ -659,6 +810,47 @@ class HafniaDataset:
             dataset=dataset, old_task_name=old_task_name, new_task_name=new_task_name
         )
 
+    def drop_task(
+        dataset: "HafniaDataset",
+        task_name: str,
+    ) -> "HafniaDataset":
+        """
+        Drop a task from the dataset.
+        If 'task_name' and 'primitive' are not provided, the function will attempt to infer the task.
+        """
+        dataset = copy.copy(dataset)  # To avoid mutating the original dataset. Shallow copy is sufficient
+        drop_task = dataset.info.get_task_by_name(task_name=task_name)
+        tasks_with_same_primitive = dataset.info.get_tasks_by_primitive(drop_task.primitive)
+
+        no_other_tasks_with_same_primitive = len(tasks_with_same_primitive) == 1
+        if no_other_tasks_with_same_primitive:
+            return dataset.drop_primitive(primitive=drop_task.primitive)
+
+        dataset.info = dataset.info.replace_task(old_task=drop_task, new_task=None)
+        dataset.samples = dataset.samples.with_columns(
+            pl.col(drop_task.primitive.column_name())
+            .list.filter(pl.element().struct.field(PrimitiveField.TASK_NAME) != drop_task.name)
+            .alias(drop_task.primitive.column_name())
+        )
+
+        return dataset
+
+    def drop_primitive(
+        dataset: "HafniaDataset",
+        primitive: Type[Primitive],
+    ) -> "HafniaDataset":
+        """
+        Drop a primitive from the dataset.
+        """
+        dataset = copy.copy(dataset)  # To avoid mutating the original dataset. Shallow copy is sufficient
+        tasks_to_drop = dataset.info.get_tasks_by_primitive(primitive=primitive)
+        for task in tasks_to_drop:
+            dataset.info = dataset.info.replace_task(old_task=task, new_task=None)
+
+        # Drop the primitive column from the samples table
+        dataset.samples = dataset.samples.drop(primitive.column_name())
+        return dataset
+
     def select_samples_by_class_name(
         dataset: HafniaDataset,
         name: Union[List[str], str],
@@ -695,13 +887,63 @@ class HafniaDataset:
 
         return HafniaDataset(info=merged_info, samples=merged_samples)
 
-    def as_dict_dataset_splits(self) -> Dict[str, "HafniaDataset"]:
+    def download_files_aws(
+        dataset: HafniaDataset,
+        path_output_folder: Path,
+        aws_credentials: AwsCredentials,
+        force_redownload: bool = False,
+    ) -> HafniaDataset:
+        from hafnia.platform.datasets import fast_copy_files_s3
+
+        remote_src_paths = dataset.samples[SampleField.REMOTE_PATH].unique().to_list()
+        update_rows = []
+        local_dst_paths = []
+        for remote_src_path in remote_src_paths:
+            local_path_str = (path_output_folder / "data" / Path(remote_src_path).name).absolute().as_posix()
+            local_dst_paths.append(local_path_str)
+            update_rows.append(
+                {
+                    SampleField.REMOTE_PATH: remote_src_path,
+                    SampleField.FILE_PATH: local_path_str,
+                }
+            )
+        update_df = pl.DataFrame(update_rows)
+        samples = dataset.samples.update(update_df, on=[SampleField.REMOTE_PATH])
+        dataset = dataset.update_samples(samples)
+
+        if not force_redownload:
+            download_indices = [idx for idx, local_path in enumerate(local_dst_paths) if not Path(local_path).exists()]
+            n_files = len(local_dst_paths)
+            skip_files = n_files - len(download_indices)
+            if skip_files > 0:
+                user_logger.info(
+                    f"Found {skip_files}/{n_files} files already exists. Downloading {len(download_indices)} files."
+                )
+            remote_src_paths = [remote_src_paths[idx] for idx in download_indices]
+            local_dst_paths = [local_dst_paths[idx] for idx in download_indices]
+
+        if len(remote_src_paths) == 0:
+            user_logger.info(
+                "All files already exist locally. Skipping download. Set 'force_redownload=True' to re-download."
+            )
+            return dataset
+
+        environment_vars = aws_credentials.aws_credentials()
+        fast_copy_files_s3(
+            src_paths=remote_src_paths,
+            dst_paths=local_dst_paths,
+            append_envs=environment_vars,
+            description="Downloading images",
+        )
+        return dataset
+
+    def to_dict_dataset_splits(self) -> Dict[str, "HafniaDataset"]:
         """
         Splits the dataset into multiple datasets based on the 'split' column.
         Returns a dictionary with split names as keys and HafniaDataset objects as values.
         """
-        if ColumnName.SPLIT not in self.samples.columns:
-            raise ValueError(f"Dataset must contain a '{ColumnName.SPLIT}' column.")
+        if SampleField.SPLIT not in self.samples.columns:
+            raise ValueError(f"Dataset must contain a '{SampleField.SPLIT}' column.")
 
         splits = {}
         for split_name in SplitName.valid_splits():
@@ -710,20 +952,11 @@ class HafniaDataset:
         return splits
 
     def create_sample_dataset(self) -> "HafniaDataset":
-        # Backwards compatibility. Remove in future versions when dataset have been updated
-        if "is_sample" in self.samples.columns:
-            user_logger.warning(
-                "'is_sample' column found in the dataset. This column is deprecated and will be removed in future versions. "
-                "Please use the 'tags' column with the tag 'sample' instead."
-            )
-            table = self.samples.filter(pl.col("is_sample") == True)  # noqa: E712
-            return self.update_samples(table)
-
-        if ColumnName.TAGS not in self.samples.columns:
-            raise ValueError(f"Dataset must contain an '{ColumnName.TAGS}' column.")
+        if SampleField.TAGS not in self.samples.columns:
+            raise ValueError(f"Dataset must contain an '{SampleField.TAGS}' column.")
 
         table = self.samples.filter(
-            pl.col(ColumnName.TAGS).list.eval(pl.element().filter(pl.element() == TAG_IS_SAMPLE)).list.len() > 0
+            pl.col(SampleField.TAGS).list.eval(pl.element().filter(pl.element() == TAG_IS_SAMPLE)).list.len() > 0
         )
         return self.update_samples(table)
 
@@ -734,10 +967,10 @@ class HafniaDataset:
             split_names = split_name
 
         for name in split_names:
-            if name not in SplitName.valid_splits():
+            if name not in SplitName.all_split_names():
                 raise ValueError(f"Invalid split name: {split_name}. Valid splits are: {SplitName.valid_splits()}")
 
-        filtered_dataset = self.samples.filter(pl.col(ColumnName.SPLIT).is_in(split_names))
+        filtered_dataset = self.samples.filter(pl.col(SampleField.SPLIT).is_in(split_names))
         return self.update_samples(filtered_dataset)
 
     def update_samples(self, table: pl.DataFrame) -> "HafniaDataset":
@@ -772,29 +1005,69 @@ class HafniaDataset:
     def copy(self) -> "HafniaDataset":
         return HafniaDataset(info=self.info.model_copy(deep=True), samples=self.samples.clone())
 
+    def create_primitive_table(
+        self,
+        primitive: Type[Primitive],
+        task_name: Optional[str] = None,
+        keep_sample_data: bool = False,
+    ) -> pl.DataFrame:
+        return table_transformations.create_primitive_table(
+            samples_table=self.samples,
+            PrimitiveType=primitive,
+            task_name=task_name,
+            keep_sample_data=keep_sample_data,
+        )
+
     def write(self, path_folder: Path, add_version: bool = False, drop_null_cols: bool = True) -> None:
         user_logger.info(f"Writing dataset to {path_folder}...")
+        path_folder = path_folder.absolute()
         if not path_folder.exists():
             path_folder.mkdir(parents=True)
-
-        new_relative_paths = []
-        for org_path in tqdm(self.samples["file_name"].to_list(), desc="- Copy images"):
+        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"):
             new_path = dataset_helpers.copy_and_rename_file_to_hash_value(
                 path_source=Path(org_path),
                 path_dataset_root=path_folder,
             )
-            new_relative_paths.append(str(new_path.relative_to(path_folder)))
-        table = self.samples.with_columns(pl.Series(new_relative_paths).alias("file_name"))
+            new_paths.append(str(new_path))
+        hafnia_dataset.samples = hafnia_dataset.samples.with_columns(pl.Series(new_paths).alias(SampleField.FILE_PATH))
+        hafnia_dataset.write_annotations(
+            path_folder=path_folder,
+            drop_null_cols=drop_null_cols,
+            add_version=add_version,
+        )
 
+    def write_annotations(
+        dataset: HafniaDataset,
+        path_folder: Path,
+        drop_null_cols: bool = True,
+        add_version: bool = False,
+    ) -> None:
+        """
+        Writes only the annotations files (JSONL and Parquet) to the specified folder.
+        """
+        user_logger.info(f"Writing dataset annotations to {path_folder}...")
+        path_folder = path_folder.absolute()
+        if not path_folder.exists():
+            path_folder.mkdir(parents=True)
+        dataset.info.write_json(path_folder / FILENAME_DATASET_INFO)
+
+        samples = dataset.samples
         if drop_null_cols:  # Drops all unused/Null columns
-            table = table.drop(pl.selectors.by_dtype(pl.Null))
+            samples = samples.drop(pl.selectors.by_dtype(pl.Null))
+
+        # Store only relative paths in the annotations files
+        absolute_paths = samples[SampleField.FILE_PATH].to_list()
+        relative_paths = [str(Path(path).relative_to(path_folder)) for path in absolute_paths]
+        samples = samples.with_columns(pl.Series(relative_paths).alias(SampleField.FILE_PATH))
 
-        table.write_ndjson(path_folder / FILENAME_ANNOTATIONS_JSONL)  # Json for readability
-        table.write_parquet(path_folder / FILENAME_ANNOTATIONS_PARQUET)  # Parquet for speed
-        self.info.write_json(path_folder / FILENAME_DATASET_INFO)
+        samples.write_ndjson(path_folder / FILENAME_ANNOTATIONS_JSONL)  # Json for readability
+        samples.write_parquet(path_folder / FILENAME_ANNOTATIONS_PARQUET)  # Parquet for speed
 
         if add_version:
-            path_version = path_folder / "versions" / f"{self.info.version}"
+            path_version = path_folder / "versions" / f"{dataset.info.version}"
             path_version.mkdir(parents=True, exist_ok=True)
             for filename in DATASET_FILENAMES_REQUIRED:
                 shutil.copy2(path_folder / filename, path_version / filename)
@@ -846,3 +1119,30 @@ def get_or_create_dataset_path_from_recipe(
     dataset.write(path_dataset)
 
     return path_dataset
+
+
+def _dataset_corrections(samples: pl.DataFrame, dataset_info: DatasetInfo) -> Tuple[pl.DataFrame, DatasetInfo]:
+    format_version_of_dataset = Version(dataset_info.format_version)
+
+    ## Backwards compatibility fixes for older dataset versions
+    if format_version_of_dataset < Version("0.2.0"):
+        samples = table_transformations.add_dataset_name_if_missing(samples, dataset_info.dataset_name)
+
+        if "file_name" in samples.columns:
+            samples = samples.rename({"file_name": SampleField.FILE_PATH})
+
+        if SampleField.SAMPLE_INDEX not in samples.columns:
+            samples = table_transformations.add_sample_index(samples)
+
+        # Backwards compatibility: If tags-column doesn't exist, create it with empty lists
+        if SampleField.TAGS not in samples.columns:
+            tags_column: List[List[str]] = [[] for _ in range(len(samples))]  # type: ignore[annotation-unchecked]
+            samples = samples.with_columns(pl.Series(tags_column, dtype=pl.List(pl.String)).alias(SampleField.TAGS))
+
+        if SampleField.STORAGE_FORMAT not in samples.columns:
+            samples = samples.with_columns(pl.lit(StorageFormat.IMAGE).alias(SampleField.STORAGE_FORMAT))
+
+        if SampleField.SAMPLE_INDEX in samples.columns and samples[SampleField.SAMPLE_INDEX].dtype != pl.UInt64:
+            samples = samples.cast({SampleField.SAMPLE_INDEX: pl.UInt64})
+
+    return samples, dataset_info