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

hafnia/dataset/hafnia_dataset.py

@@ -10,6 +10,7 @@ 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
@@ -27,18 +28,21 @@ from hafnia.dataset.dataset_names import (
     FILENAME_DATASET_INFO,
     FILENAME_RECIPE_JSON,
     TAG_IS_SAMPLE,
-    ColumnName,
+    AwsCredentials,
+    PrimitiveField,
+    SampleField,
     SplitName,
+    StorageFormat,
+)
+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.operations.table_transformations import (
-    check_image_paths,
-    read_samples_from_path,
-)
 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
@@ -65,6 +69,14 @@ class TaskInfo(BaseModel):
         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
@@ -100,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 []
@@ -115,7 +131,6 @@ 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")
-    distributions: Optional[List[TaskInfo]] = Field(default=None, description="Optional list of task distributions")
     reference_bibtex: Optional[str] = Field(
         default=None,
         description="Optional, BibTeX reference to dataset publication",
@@ -244,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=None,
             tasks=list(unique_tasks),
-            distributions=list(distributions),
             meta=meta,
             format_version=dataset_format_version,
         )
@@ -267,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(
@@ -320,7 +341,7 @@ class DatasetInfo(BaseModel):
 
 
 class Sample(BaseModel):
-    file_path: str = Field(description="Path to the image file")
+    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'")
@@ -328,6 +349,10 @@ class Sample(BaseModel):
         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")
@@ -338,7 +363,7 @@ class Sample(BaseModel):
     classifications: Optional[List[Classification]] = Field(
         default=None, description="Optional list of classifications"
     )
-    objects: Optional[List[Bbox]] = Field(default=None, description="Optional list of objects (bounding boxes)")
+    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")
 
@@ -374,6 +399,8 @@ 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.
         """
+        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.")
@@ -382,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:
@@ -466,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
@@ -481,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)
@@ -501,14 +551,14 @@ class HafniaDataset:
         HafniaDataset.check_dataset_path(path_folder, raise_error=True)
 
         dataset_info = DatasetInfo.from_json_file(path_folder / FILENAME_DATASET_INFO)
-        samples = read_samples_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() + "/"
-        samples = samples.with_columns((dataset_root + pl.col(ColumnName.FILE_PATH)).alias(ColumnName.FILE_PATH))
+        samples = samples.with_columns((dataset_root + pl.col(SampleField.FILE_PATH)).alias(SampleField.FILE_PATH))
         if check_for_images:
-            check_image_paths(samples)
+            table_transformations.check_image_paths(samples)
         return HafniaDataset(samples=samples, info=dataset_info)
 
     @staticmethod
@@ -535,16 +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)
-
-        # Add 'dataset_name' to samples
-        table = table.with_columns(
-            pl.when(pl.col(ColumnName.DATASET_NAME).is_null())
-            .then(pl.lit(info.dataset_name))
-            .otherwise(pl.col(ColumnName.DATASET_NAME))
-            .alias(ColumnName.DATASET_NAME)
-        )
+        # 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
@@ -678,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
@@ -696,15 +742,17 @@ 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)
 
@@ -762,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],
@@ -798,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():
@@ -813,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)
 
@@ -837,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":
@@ -875,30 +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 = []
-        org_paths = self.samples[ColumnName.FILE_PATH].to_list()
+        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(ColumnName.FILE_PATH))
+            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)
@@ -956,19 +1125,24 @@ def _dataset_corrections(samples: pl.DataFrame, dataset_info: DatasetInfo) -> Tu
     format_version_of_dataset = Version(dataset_info.format_version)
 
     ## Backwards compatibility fixes for older dataset versions
-    if format_version_of_dataset <= Version("0.3.0"):
-        if ColumnName.DATASET_NAME not in samples.columns:
-            samples = samples.with_columns(pl.lit(dataset_info.dataset_name).alias(ColumnName.DATASET_NAME))
+    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": ColumnName.FILE_PATH})
+            samples = samples.rename({"file_name": SampleField.FILE_PATH})
 
-        if ColumnName.SAMPLE_INDEX not in samples.columns:
-            samples = samples.with_row_index(name=ColumnName.SAMPLE_INDEX)
+        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 ColumnName.TAGS not in samples.columns:
+        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(ColumnName.TAGS))
+            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