hafnia 0.5.0__py3-none-any.whl → 0.5.2__py3-none-any.whl

hafnia/dataset/dataset_helpers.py

@@ -3,12 +3,70 @@ import math
 import random
 import shutil
 from pathlib import Path
-from typing import Dict, List
+from typing import Dict, List, Optional, Tuple
 
 import numpy as np
 import xxhash
+from packaging.version import InvalidVersion, Version
 from PIL import Image
 
+from hafnia.log import user_logger
+
+
+def is_valid_version_string(version: Optional[str], allow_none: bool = False, allow_latest: bool = False) -> bool:
+    if allow_none and version is None:
+        return True
+    if allow_latest and version == "latest":
+        return True
+    return version_from_string(version, raise_error=False) is not None
+
+
+def version_from_string(version: Optional[str], raise_error: bool = True) -> Optional[Version]:
+    if version is None:
+        if raise_error:
+            raise ValueError("Version is 'None'. A valid version string is required e.g '1.0.0'")
+        return None
+
+    try:
+        version_casted = Version(version)
+    except (InvalidVersion, TypeError) as e:
+        if raise_error:
+            raise ValueError(f"Invalid version string/type: {version}") from e
+        return None
+
+    # Check if version is semantic versioning (MAJOR.MINOR.PATCH)
+    if len(version_casted.release) < 3:
+        if raise_error:
+            raise ValueError(f"Version string '{version}' is not semantic versioning (MAJOR.MINOR.PATCH)")
+        return None
+    return version_casted
+
+
+def dataset_name_and_version_from_string(
+    string: str,
+    resolve_missing_version: bool = True,
+) -> Tuple[str, Optional[str]]:
+    if not isinstance(string, str):
+        raise TypeError(f"'{type(string)}' for '{string}' is an unsupported type. Expected 'str' e.g 'mnist:1.0.0'")
+
+    parts = string.split(":")
+    if len(parts) == 1:
+        dataset_name = parts[0]
+        if resolve_missing_version:
+            version = "latest"  # Default to 'latest' if version is missing. This will be resolved to a specific version later.
+            user_logger.info(f"Version is missing in dataset name: {string}. Defaulting to version='latest'.")
+        else:
+            raise ValueError(f"Version is missing in dataset name: {string}. Use 'name:version'")
+    elif len(parts) == 2:
+        dataset_name, version = parts
+    else:
+        raise ValueError(f"Invalid dataset name format: {string}. Use 'name' or 'name:version' ")
+
+    if not is_valid_version_string(version, allow_none=True, allow_latest=True):
+        raise ValueError(f"Invalid version string: {version}. Use semantic versioning e.g. '1.0.0' or 'latest'")
+
+    return dataset_name, version
+
 
 def create_split_name_list_from_ratios(split_ratios: Dict[str, float], n_items: int, seed: int = 42) -> List[str]:
     samples_per_split = split_sizes_from_ratios(split_ratios=split_ratios, n_items=n_items)

hafnia/dataset/dataset_names.py

@@ -1,9 +1,5 @@
 from enum import Enum
-from typing import Dict, List, Optional
-
-import boto3
-from botocore.exceptions import UnauthorizedSSOTokenError
-from pydantic import BaseModel, field_validator
+from typing import List
 
 FILENAME_RECIPE_JSON = "recipe.json"
 FILENAME_DATASET_INFO = "dataset_info.json"
@@ -22,7 +18,6 @@ class DeploymentStage(Enum):
     PRODUCTION = "production"
 
 
-ARN_PREFIX = "arn:aws:s3:::"
 TAG_IS_SAMPLE = "sample"
 
 OPS_REMOVE_CLASS = "__REMOVE__"
@@ -126,105 +121,3 @@ class DatasetVariant(Enum):
     DUMP = "dump"
     SAMPLE = "sample"
     HIDDEN = "hidden"
-
-
-class AwsCredentials(BaseModel):
-    access_key: str
-    secret_key: str
-    session_token: str
-    region: Optional[str]
-
-    def aws_credentials(self) -> Dict[str, str]:
-        """
-        Returns the AWS credentials as a dictionary.
-        """
-        environment_vars = {
-            "AWS_ACCESS_KEY_ID": self.access_key,
-            "AWS_SECRET_ACCESS_KEY": self.secret_key,
-            "AWS_SESSION_TOKEN": self.session_token,
-        }
-        if self.region:
-            environment_vars["AWS_REGION"] = self.region
-
-        return environment_vars
-
-    @staticmethod
-    def from_session(session: boto3.Session) -> "AwsCredentials":
-        """
-        Creates AwsCredentials from a Boto3 session.
-        """
-        try:
-            frozen_credentials = session.get_credentials().get_frozen_credentials()
-        except UnauthorizedSSOTokenError as e:
-            raise RuntimeError(
-                f"Failed to get AWS credentials from the session for profile '{session.profile_name}'.\n"
-                f"Ensure the profile exists in your AWS config in '~/.aws/config' and that you are logged in via AWS SSO.\n"
-                f"\tUse 'aws sso login --profile {session.profile_name}' to log in."
-            ) from e
-        return AwsCredentials(
-            access_key=frozen_credentials.access_key,
-            secret_key=frozen_credentials.secret_key,
-            session_token=frozen_credentials.token,
-            region=session.region_name,
-        )
-
-    def to_resource_credentials(self, bucket_name: str) -> "ResourceCredentials":
-        """
-        Converts AwsCredentials to ResourceCredentials by adding the S3 ARN.
-        """
-        payload = self.model_dump()
-        payload["s3_arn"] = f"{ARN_PREFIX}{bucket_name}"
-        return ResourceCredentials(**payload)
-
-
-class ResourceCredentials(AwsCredentials):
-    s3_arn: str
-
-    @staticmethod
-    def fix_naming(payload: Dict[str, str]) -> "ResourceCredentials":
-        """
-        The endpoint returns a payload with a key called 's3_path', but it
-        is actually an ARN path (starts with arn:aws:s3::). This method renames it to 's3_arn' for consistency.
-        """
-        if "s3_path" in payload and payload["s3_path"].startswith(ARN_PREFIX):
-            payload["s3_arn"] = payload.pop("s3_path")
-
-        if "region" not in payload:
-            payload["region"] = "eu-west-1"
-        return ResourceCredentials(**payload)
-
-    @field_validator("s3_arn")
-    @classmethod
-    def validate_s3_arn(cls, value: str) -> str:
-        """Validate s3_arn to ensure it starts with 'arn:aws:s3:::'"""
-        if not value.startswith("arn:aws:s3:::"):
-            raise ValueError(f"Invalid S3 ARN: {value}. It should start with 'arn:aws:s3:::'")
-        return value
-
-    def s3_path(self) -> str:
-        """
-        Extracts the S3 path from the ARN.
-        Example: arn:aws:s3:::my-bucket/my-prefix -> my-bucket/my-prefix
-        """
-        return self.s3_arn[len(ARN_PREFIX) :]
-
-    def s3_uri(self) -> str:
-        """
-        Converts the S3 ARN to a URI format.
-        Example: arn:aws:s3:::my-bucket/my-prefix -> s3://my-bucket/my-prefix
-        """
-        return f"s3://{self.s3_path()}"
-
-    def bucket_name(self) -> str:
-        """
-        Extracts the bucket name from the S3 ARN.
-        Example: arn:aws:s3:::my-bucket/my-prefix -> my-bucket
-        """
-        return self.s3_path().split("/")[0]
-
-    def object_key(self) -> str:
-        """
-        Extracts the object key from the S3 ARN.
-        Example: arn:aws:s3:::my-bucket/my-prefix -> my-prefix
-        """
-        return "/".join(self.s3_path().split("/")[1:])

hafnia/dataset/dataset_recipe/dataset_recipe.py

@@ -11,14 +11,19 @@ from pydantic import (
 )
 
 from hafnia import utils
+from hafnia.dataset.dataset_helpers import dataset_name_and_version_from_string
 from hafnia.dataset.dataset_recipe import recipe_transforms
 from hafnia.dataset.dataset_recipe.recipe_types import (
     RecipeCreation,
     RecipeTransform,
     Serializable,
 )
-from hafnia.dataset.hafnia_dataset import HafniaDataset
+from hafnia.dataset.hafnia_dataset import (
+    HafniaDataset,
+    available_dataset_versions_from_name,
+)
 from hafnia.dataset.primitives.primitive import Primitive
+from hafnia.log import user_logger
 
 
 class DatasetRecipe(Serializable):
@@ -41,8 +46,31 @@ class DatasetRecipe(Serializable):
 
     ### Creation Methods (using the 'from_X' )###
     @staticmethod
-    def from_name(name: str, force_redownload: bool = False, download_files: bool = True) -> DatasetRecipe:
-        creation = FromName(name=name, force_redownload=force_redownload, download_files=download_files)
+    def from_name(
+        name: str,
+        version: Optional[str] = None,
+        force_redownload: bool = False,
+        download_files: bool = True,
+    ) -> DatasetRecipe:
+        if version == "latest":
+            user_logger.info(
+                f"The dataset '{name}' in a dataset recipe uses 'latest' as version. For dataset recipes the "
+                "version is pinned to a specific version. Consider specifying a specific version to ensure "
+                "reproducibility of your experiments. "
+            )
+            available_versions = available_dataset_versions_from_name(name)
+            version = str(max(available_versions))
+        if version is None:
+            available_versions = available_dataset_versions_from_name(name)
+            str_versions = ", ".join([str(v) for v in available_versions])
+            raise ValueError(
+                f"Version must be specified when creating a DatasetRecipe from name. "
+                f"Available versions are: {str_versions}"
+            )
+
+        creation = FromName(
+            name=name, version=version, force_redownload=force_redownload, download_files=download_files
+        )
         return DatasetRecipe(creation=creation)
 
     @staticmethod
@@ -125,6 +153,21 @@ class DatasetRecipe(Serializable):
         recipe_id = recipe["id"]
         return DatasetRecipe.from_recipe_id(recipe_id)
 
+    @staticmethod
+    def from_name_and_version_string(string: str, resolve_missing_version: bool = False) -> "DatasetRecipe":
+        """
+        Validates and converts a dataset name and version string (name:version) to a DatasetRecipe.from_name recipe.
+        If version is missing and 'resolve_missing_version' is True, it will default to 'latest'.
+        If resolve_missing_version is False, it will raise an error if version is missing.
+        """
+
+        dataset_name, version = dataset_name_and_version_from_string(
+            string=string,
+            resolve_missing_version=resolve_missing_version,
+        )
+
+        return DatasetRecipe.from_name(name=dataset_name, version=version)
+
     @staticmethod
     def from_implicit_form(recipe: Any) -> DatasetRecipe:
         """
@@ -180,7 +223,7 @@ class DatasetRecipe(Serializable):
             return recipe
 
         if isinstance(recipe, str):  # str-type is convert to DatasetFromName
-            return DatasetRecipe.from_name(name=recipe)
+            return DatasetRecipe.from_name_and_version_string(string=recipe, resolve_missing_version=True)
 
         if isinstance(recipe, Path):  # Path-type is convert to DatasetFromPath
             return DatasetRecipe.from_path(path_folder=recipe)
@@ -409,6 +452,7 @@ class FromPath(RecipeCreation):
 
 class FromName(RecipeCreation):
     name: str
+    version: Optional[str] = None
     force_redownload: bool = False
     download_files: bool = True
 

hafnia/dataset/format_conversions/torchvision_datasets.py

@@ -78,7 +78,7 @@ def caltech_101_as_hafnia_dataset(
         n_samples=n_samples,
         dataset_name=dataset_name,
     )
-    hafnia_dataset.info.version = "1.1.0"
+    hafnia_dataset.info.version = "1.0.0"
     hafnia_dataset.info.reference_bibtex = textwrap.dedent("""\
         @article{FeiFei2004LearningGV,
             title={Learning Generative Visual Models from Few Training Examples: An Incremental Bayesian
@@ -108,7 +108,7 @@ def caltech_256_as_hafnia_dataset(
         n_samples=n_samples,
         dataset_name=dataset_name,
     )
-    hafnia_dataset.info.version = "1.1.0"
+    hafnia_dataset.info.version = "1.0.0"
     hafnia_dataset.info.reference_bibtex = textwrap.dedent("""\
         @misc{griffin_2023_5sv1j-ytw97,
             author       = {Griffin, Gregory and

hafnia/dataset/hafnia_dataset.py

@@ -10,14 +10,12 @@ from typing import Any, Dict, List, Optional, Tuple, Type, Union
 import polars as pl
 from packaging.version import Version
 
+from hafnia import utils
 from hafnia.dataset import dataset_helpers
+from hafnia.dataset.dataset_helpers import is_valid_version_string, version_from_string
 from hafnia.dataset.dataset_names import (
-    FILENAME_ANNOTATIONS_JSONL,
-    FILENAME_ANNOTATIONS_PARQUET,
-    FILENAME_DATASET_INFO,
     FILENAME_RECIPE_JSON,
     TAG_IS_SAMPLE,
-    AwsCredentials,
     PrimitiveField,
     SampleField,
     SplitName,
@@ -28,7 +26,7 @@ from hafnia.dataset.format_conversions import (
     format_image_classification_folder,
     format_yolo,
 )
-from hafnia.dataset.hafnia_dataset_types import DatasetInfo, Sample
+from hafnia.dataset.hafnia_dataset_types import DatasetInfo, DatasetMetadataFilePaths, Sample
 from hafnia.dataset.operations import (
     dataset_stats,
     dataset_transformations,
@@ -36,6 +34,9 @@ from hafnia.dataset.operations import (
 )
 from hafnia.dataset.primitives.primitive import Primitive
 from hafnia.log import user_logger
+from hafnia.platform import s5cmd_utils
+from hafnia.platform.datasets import get_read_credentials_by_name
+from hafnia.platform.s5cmd_utils import AwsCredentials, ResourceCredentials
 from hafnia.utils import progress_bar
 from hafnia_cli.config import Config
 
@@ -89,10 +90,11 @@ class HafniaDataset:
     @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)
+        metadata_file_paths = DatasetMetadataFilePaths.from_path(path_folder)
+        metadata_file_paths.exists(raise_error=True)
 
-        dataset_info = DatasetInfo.from_json_file(path_folder / FILENAME_DATASET_INFO)
-        samples = table_transformations.read_samples_from_path(path_folder)
+        dataset_info = DatasetInfo.from_json_file(Path(metadata_file_paths.dataset_info))
+        samples = metadata_file_paths.read_samples()
         samples, dataset_info = _dataset_corrections(samples, dataset_info)
 
         # Convert from relative paths to absolute paths
@@ -103,14 +105,24 @@ class HafniaDataset:
         return HafniaDataset(samples=samples, info=dataset_info)
 
     @staticmethod
-    def from_name(name: str, force_redownload: bool = False, download_files: bool = True) -> "HafniaDataset":
+    def from_name(
+        name: str,
+        version: Optional[str] = None,
+        force_redownload: bool = False,
+        download_files: bool = True,
+    ) -> "HafniaDataset":
         """
         Load a dataset by its name. The dataset must be registered in the Hafnia platform.
         """
-        from hafnia.platform.datasets import download_or_get_dataset_path
-
+        if ":" in name:
+            name, version = dataset_helpers.dataset_name_and_version_from_string(name)
+            raise ValueError(
+                "The 'from_name' does not support the 'name:version' format. Please provide the version separately.\n"
+                f"E.g., HafniaDataset.from_name(name='{name}', version='{version}')"
+            )
         dataset_path = download_or_get_dataset_path(
             dataset_name=name,
+            version=version,
             force_redownload=force_redownload,
             download_files=download_files,
         )
@@ -523,30 +535,6 @@ class HafniaDataset:
         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:
-        """
-        Checks if the dataset path exists and contains the required files.
-        Returns True if the dataset is valid, otherwise raises an error or returns False.
-        """
-        if not path_dataset.exists():
-            if raise_error:
-                raise FileNotFoundError(f"Dataset path {path_dataset} does not exist.")
-            return False
-
-        required_files = [
-            FILENAME_DATASET_INFO,
-            FILENAME_ANNOTATIONS_JSONL,
-            FILENAME_ANNOTATIONS_PARQUET,
-        ]
-        for filename in required_files:
-            if not (path_dataset / filename).exists():
-                if raise_error:
-                    raise FileNotFoundError(f"Required file {filename} not found in {path_dataset}.")
-                return False
-
-        return True
-
     def copy(self) -> "HafniaDataset":
         return HafniaDataset(info=self.info.model_copy(deep=True), samples=self.samples.clone())
 
@@ -584,16 +572,18 @@ class HafniaDataset:
         """
         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)
+        metadata_file_paths = DatasetMetadataFilePaths.from_path(path_folder)
+        path_dataset_info = Path(metadata_file_paths.dataset_info)
+        path_dataset_info.parent.mkdir(parents=True, exist_ok=True)
+        dataset.info.write_json(path_dataset_info)
 
         samples = dataset.samples
         if drop_null_cols:  # Drops all unused/Null columns
             samples = samples.drop(pl.selectors.by_dtype(pl.Null))
 
+        path_folder = path_folder.absolute()
         # Store only relative paths in the annotations files
         if SampleField.FILE_PATH in samples.columns:  # We drop column for remote datasets
             absolute_paths = samples[SampleField.FILE_PATH].to_list()
@@ -601,8 +591,11 @@ class HafniaDataset:
             samples = samples.with_columns(pl.Series(relative_paths).alias(SampleField.FILE_PATH))
         else:
             samples = samples.with_columns(pl.lit("").alias(SampleField.FILE_PATH))
-        samples.write_ndjson(path_folder / FILENAME_ANNOTATIONS_JSONL)  # Json for readability
-        samples.write_parquet(path_folder / FILENAME_ANNOTATIONS_PARQUET)  # Parquet for speed
+
+        if metadata_file_paths.annotations_jsonl:
+            samples.write_ndjson(Path(metadata_file_paths.annotations_jsonl))  # Json for readability
+        if metadata_file_paths.annotations_parquet:
+            samples.write_parquet(Path(metadata_file_paths.annotations_parquet))  # Parquet for speed
 
     def delete_on_platform(dataset: HafniaDataset, interactive: bool = True) -> None:
         """
@@ -707,6 +700,42 @@ class HafniaDataset:
         return True
 
 
+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})
+
+    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
+
+
 def check_hafnia_dataset_from_path(path_dataset: Path) -> None:
     dataset = HafniaDataset.from_path(path_dataset, check_for_images=True)
     dataset.check_dataset()
@@ -728,7 +757,8 @@ def get_or_create_dataset_path_from_recipe(
     if force_redownload:
         shutil.rmtree(path_dataset, ignore_errors=True)
 
-    if HafniaDataset.check_dataset_path(path_dataset, raise_error=False):
+    dataset_metadata_files = DatasetMetadataFilePaths.from_path(path_dataset)
+    if dataset_metadata_files.exists(raise_error=False):
         return path_dataset
 
     path_dataset.mkdir(parents=True, exist_ok=True)
@@ -741,37 +771,101 @@ def get_or_create_dataset_path_from_recipe(
     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)
+def available_dataset_versions_from_name(dataset_name: str) -> Dict[Version, "DatasetMetadataFilePaths"]:
+    credentials: ResourceCredentials = get_read_credentials_by_name(dataset_name=dataset_name)
+    return available_dataset_versions(credentials=credentials)
 
-    ## 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})
+def available_dataset_versions(
+    credentials: ResourceCredentials,
+) -> Dict[Version, "DatasetMetadataFilePaths"]:
+    envs = credentials.aws_credentials()
+    bucket_prefix_sample_versions = f"{credentials.s3_uri()}/versions"
+    all_s3_annotation_files = s5cmd_utils.list_bucket(bucket_prefix=bucket_prefix_sample_versions, append_envs=envs)
+    available_versions = DatasetMetadataFilePaths.available_versions_from_files_list(all_s3_annotation_files)
+    return available_versions
 
-        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))
+def select_version_from_available_versions(
+    available_versions: Dict[Version, "DatasetMetadataFilePaths"],
+    version: Optional[str],
+) -> "DatasetMetadataFilePaths":
+    if len(available_versions) == 0:
+        raise ValueError("No versions were found in the dataset.")
 
-        if SampleField.STORAGE_FORMAT not in samples.columns:
-            samples = samples.with_columns(pl.lit(StorageFormat.IMAGE).alias(SampleField.STORAGE_FORMAT))
+    if version is None:
+        str_versions = [str(v) for v in available_versions]
+        raise ValueError(f"Version must be specified. Available versions: {str_versions}")
+    elif version == "latest":
+        version_casted = max(available_versions)
+        user_logger.info(f"'latest' version '{version_casted}' has been selected")
+    else:
+        version_casted = version_from_string(version)
 
-        if SampleField.SAMPLE_INDEX in samples.columns and samples[SampleField.SAMPLE_INDEX].dtype != pl.UInt64:
-            samples = samples.cast({SampleField.SAMPLE_INDEX: pl.UInt64})
+    if version_casted not in available_versions:
+        raise ValueError(f"Selected version '{version}' not found in available versions: {available_versions}")
 
-    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
+    return available_versions[version_casted]
+
+
+def download_meta_dataset_files_from_version(
+    resource_credentials: ResourceCredentials, version: Optional[str], path_dataset: Path
+) -> list[str]:
+    envs = resource_credentials.aws_credentials()
+    available_versions = available_dataset_versions(credentials=resource_credentials)
+    metadata_files = select_version_from_available_versions(available_versions=available_versions, version=version)
+
+    s3_files = metadata_files.as_list()
+    path_dataset.mkdir(parents=True, exist_ok=True)
+    local_paths = [(path_dataset / filename.split("/")[-1]).as_posix() for filename in s3_files]
+    s5cmd_utils.fast_copy_files(
+        src_paths=s3_files,
+        dst_paths=local_paths,
+        append_envs=envs,
+        description="Downloading meta dataset files",
+    )
+
+    return local_paths
+
+
+def download_or_get_dataset_path(
+    dataset_name: str,
+    version: Optional[str],
+    cfg: Optional[Config] = None,
+    path_datasets_folder: Optional[str] = None,
+    force_redownload: bool = False,
+    download_files: bool = True,
+) -> Path:
+    """Download or get the path of the dataset."""
+
+    path_datasets = path_datasets_folder or utils.PATH_DATASETS
+    path_dataset = Path(path_datasets) / dataset_name
+    if not is_valid_version_string(version, allow_none=True, allow_latest=True):
+        raise ValueError(
+            f"Invalid version string: {version}. Should be a valid version (e.g. '0.1.0'), 'latest' or None."
+        )
+
+    # Only valid versions (e.g. '0.1.0', '1.0.0') can use local cache. Using either "latest"/None will always redownload
+    if is_valid_version_string(version, allow_none=False, allow_latest=False):
+        dataset_metadata_files = DatasetMetadataFilePaths.from_path(path_dataset)
+        dataset_exists = dataset_metadata_files.exists(version=version, raise_error=False)
+        if dataset_exists and not force_redownload:
+            user_logger.info("Dataset found locally. Set 'force=True' or add `--force` flag with cli to re-download")
+            return path_dataset
+
+    cfg = cfg or Config()
+    resource_credentials = get_read_credentials_by_name(dataset_name=dataset_name, cfg=cfg)
+    if resource_credentials is None:
+        raise ValueError(f"Failed to get read credentials for dataset '{dataset_name}' from the platform.")
+
+    download_meta_dataset_files_from_version(
+        resource_credentials=resource_credentials, version=version, path_dataset=path_dataset
+    )
+
+    if not download_files:
+        return path_dataset
+
+    dataset = HafniaDataset.from_path(path_dataset, check_for_images=False)
+    dataset = dataset.download_files_aws(path_dataset, aws_credentials=resource_credentials, force_redownload=True)
+    dataset.write_annotations(path_folder=path_dataset)  # Overwrite annotations as files have been re-downloaded
+    return path_dataset