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

hafnia/dataset/hafnia_dataset_types.py

@@ -1,5 +1,6 @@
 import collections
 import json
+from dataclasses import dataclass
 from datetime import datetime
 from pathlib import Path
 from typing import Any, Dict, List, Optional, Type, Union
@@ -7,12 +8,21 @@ from typing import Any, Dict, List, Optional, 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
 
 import hafnia
-from hafnia.dataset.dataset_names import SampleField, StorageFormat
+from hafnia.dataset import dataset_helpers
+from hafnia.dataset.dataset_helpers import version_from_string
+from hafnia.dataset.dataset_names import (
+    FILENAME_ANNOTATIONS_JSONL,
+    FILENAME_ANNOTATIONS_PARQUET,
+    FILENAME_DATASET_INFO,
+    SampleField,
+    StorageFormat,
+)
 from hafnia.dataset.primitives import (
     PRIMITIVE_TYPES,
     Bbox,
@@ -51,7 +61,7 @@ class TaskInfo(BaseModel):
         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.
+    # the 'Primitive' class is an abstract base class and for the actual primitives 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")
@@ -102,7 +112,9 @@ class TaskInfo(BaseModel):
 
 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")
+    version: str = Field(default="0.0.0", description="Version of the dataset")
+    dataset_title: Optional[str] = Field(default=None, description="Optional, human-readable title of the dataset")
+    description: Optional[str] = Field(default=None, description="Optional, description of the dataset")
     tasks: List[TaskInfo] = Field(default=None, description="List of tasks in the dataset")
     reference_bibtex: Optional[str] = Field(
         default=None,
@@ -142,31 +154,21 @@ class DatasetInfo(BaseModel):
     @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
+        version_casted: Version = dataset_helpers.version_from_string(format_version, raise_error=True)
 
-        if Version(format_version) > Version(hafnia.__dataset_format_version__):
+        if version_casted > 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
+        return str(version_casted)
 
     @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
+        version_casted: Version = dataset_helpers.version_from_string(dataset_version, raise_error=True)
+        return str(version_casted)
 
     def check_for_duplicate_task_names(self) -> List[TaskInfo]:
         return self._validate_check_for_duplicate_tasks(self.tasks)
@@ -236,7 +238,7 @@ class DatasetInfo(BaseModel):
         meta.update(info1.meta or {})
         return DatasetInfo(
             dataset_name=info0.dataset_name + "+" + info1.dataset_name,
-            version=None,
+            version="0.0.0",
             tasks=list(unique_tasks),
             meta=meta,
             format_version=dataset_format_version,
@@ -475,3 +477,127 @@ class Sample(BaseModel):
         annotations = self.get_annotations()
         annotations_visualized = image_visualizations.draw_annotations(image=image, primitives=annotations)
         return annotations_visualized
+
+
+@dataclass
+class DatasetMetadataFilePaths:
+    dataset_info: str  # Use 'str' to also support s3 paths
+    annotations_jsonl: Optional[str]
+    annotations_parquet: Optional[str]
+
+    def as_list(self) -> List[str]:
+        files = [self.dataset_info]
+        if self.annotations_jsonl is not None:
+            files.append(self.annotations_jsonl)
+        if self.annotations_parquet is not None:
+            files.append(self.annotations_parquet)
+        return files
+
+    def read_samples(self) -> pl.DataFrame:
+        if self.annotations_parquet is not None:
+            if not Path(self.annotations_parquet).exists():
+                raise FileNotFoundError(f"Parquet annotations file '{self.annotations_parquet}' does not exist.")
+            user_logger.info(f"Reading dataset annotations from Parquet file: {self.annotations_parquet}")
+            return pl.read_parquet(self.annotations_parquet)
+
+        if self.annotations_jsonl is not None:
+            if not Path(self.annotations_jsonl).exists():
+                raise FileNotFoundError(f"JSONL annotations file '{self.annotations_jsonl}' does not exist.")
+            user_logger.info(f"Reading dataset annotations from JSONL file: {self.annotations_jsonl}")
+            return pl.read_ndjson(self.annotations_jsonl)
+
+        raise ValueError(
+            "No annotations file available to read samples from. Dataset is missing both JSONL and Parquet files."
+        )
+
+    @staticmethod
+    def from_path(path_dataset: Path) -> "DatasetMetadataFilePaths":
+        path_dataset = path_dataset.absolute()
+        metadata_files = DatasetMetadataFilePaths(
+            dataset_info=str(path_dataset / FILENAME_DATASET_INFO),
+            annotations_jsonl=str(path_dataset / FILENAME_ANNOTATIONS_JSONL),
+            annotations_parquet=str(path_dataset / FILENAME_ANNOTATIONS_PARQUET),
+        )
+
+        return metadata_files
+
+    @staticmethod
+    def available_versions_from_files_list(files: list[str]) -> Dict[Version, "DatasetMetadataFilePaths"]:
+        versions_and_files: Dict[Version, Dict[str, str]] = collections.defaultdict(dict)
+        for metadata_file in files:
+            version_str, filename = metadata_file.split("/")[-2:]
+            versions_and_files[version_str][filename] = metadata_file
+
+        available_versions: Dict[Version, DatasetMetadataFilePaths] = {}
+        for version_str, version_files in versions_and_files.items():
+            version_casted: Version = dataset_helpers.version_from_string(version_str, raise_error=False)
+            if version_casted is None:
+                continue
+
+            if FILENAME_DATASET_INFO not in version_files:
+                continue
+            dataset_metadata_file = DatasetMetadataFilePaths(
+                dataset_info=version_files[FILENAME_DATASET_INFO],
+                annotations_jsonl=version_files.get(FILENAME_ANNOTATIONS_JSONL, None),
+                annotations_parquet=version_files.get(FILENAME_ANNOTATIONS_PARQUET, None),
+            )
+
+            available_versions[version_casted] = dataset_metadata_file
+
+        return available_versions
+
+    def check_version(self, version: str, raise_error: bool = True) -> bool:
+        """
+        Check if the dataset metadata files match the given version.
+        If raise_error is True, raises ValueError if the version does not match.
+        """
+        valid_version = version_from_string(version, raise_error=raise_error)
+        if valid_version is None:
+            return False
+
+        path_dataset_info = Path(self.dataset_info)
+        if not path_dataset_info.exists():
+            raise FileNotFoundError(f"Dataset info file missing '{self.dataset_info}' in dataset folder.")
+
+        dataset_info = json.loads(path_dataset_info.read_text())
+        dataset_version = dataset_info.get("version", None)
+        if dataset_version != version:
+            if raise_error:
+                raise ValueError(
+                    f"Dataset version mismatch. Expected version '{version}' but found "
+                    f"version '{dataset_version}' in dataset info."
+                )
+            return False
+
+        return True
+
+    def exists(self, version: Optional[str] = None, raise_error: bool = True) -> bool:
+        """
+        Check if all metadata files exist.
+        Add version to check if it matches the version in dataset info.
+        If raise_error is True, raises FileNotFoundError if any file is missing.
+        """
+        path_dataset_info = Path(self.dataset_info)
+        if not path_dataset_info.exists():
+            if raise_error:
+                raise FileNotFoundError(f"Dataset info file missing '{self.dataset_info}' in dataset folder.")
+            return False
+
+        if version is not None and self.check_version(version, raise_error=raise_error) is False:
+            return False
+
+        has_jsonl_file = self.annotations_jsonl is not None and Path(self.annotations_jsonl).exists()
+        if has_jsonl_file:
+            return True
+
+        has_parquet_file = self.annotations_parquet is not None and Path(self.annotations_parquet).exists()
+        if has_parquet_file:
+            return True
+
+        if raise_error:
+            raise FileNotFoundError(
+                f"Missing annotation file. Expected either '{FILENAME_ANNOTATIONS_JSONL}' or "
+                f"'{FILENAME_ANNOTATIONS_PARQUET}' in dataset folder."
+            )
+
+        return False

hafnia/dataset/operations/dataset_s3_storage.py

@@ -0,0 +1,216 @@
+import tempfile
+import time
+from pathlib import Path
+from typing import Dict, Optional
+
+import polars as pl
+
+from hafnia.dataset.dataset_helpers import hash_file_xxhash
+from hafnia.dataset.dataset_names import (
+    DatasetVariant,
+    SampleField,
+)
+from hafnia.dataset.hafnia_dataset import HafniaDataset
+from hafnia.log import user_logger
+from hafnia.platform import s5cmd_utils
+from hafnia.platform.datasets import get_upload_credentials
+from hafnia.platform.s5cmd_utils import ResourceCredentials
+from hafnia.utils import progress_bar
+from hafnia_cli.config import Config
+
+
+def delete_hafnia_dataset_files_on_platform(
+    dataset_name: str,
+    interactive: bool = True,
+    cfg: Optional[Config] = None,
+) -> bool:
+    cfg = cfg or Config()
+    resource_credentials = get_upload_credentials(dataset_name, cfg=cfg)
+
+    if resource_credentials is None:
+        raise RuntimeError("Failed to get upload credentials from the platform.")
+
+    return delete_hafnia_dataset_files_from_resource_credentials(
+        interactive=interactive,
+        resource_credentials=resource_credentials,
+    )
+
+
+def delete_hafnia_dataset_files_from_resource_credentials(
+    resource_credentials: ResourceCredentials,
+    interactive: bool = True,
+    remove_bucket: bool = True,
+) -> bool:
+    envs = resource_credentials.aws_credentials()
+    bucket_name = resource_credentials.bucket_name()
+    if interactive:
+        confirmation = (
+            input(
+                f"WARNING THIS WILL delete all files stored in 's3://{bucket_name}'.\n"
+                "Meaning that all previous versions of the dataset will be deleted. \n"
+                "Normally this is not needed, but if you have changed the dataset structure or want to start from fresh, "
+                "you can delete all files in the S3 bucket. "
+                "\nDo you really want to delete all files? (yes/NO): "
+            )
+            .strip()
+            .lower()
+        )
+        if confirmation != "yes":
+            user_logger.info("Delete operation cancelled by the user.")
+            return False
+    user_logger.info(f"Deleting all files in S3 bucket '{bucket_name}'...")
+    s5cmd_utils.delete_bucket_content(
+        bucket_prefix=f"s3://{bucket_name}",
+        remove_bucket=remove_bucket,
+        append_envs=envs,
+    )
+    return True
+
+
+def sync_hafnia_dataset_to_s3(
+    dataset: HafniaDataset,
+    bucket_prefix: str,
+    allow_version_overwrite: bool = False,
+    interactive: bool = True,
+    envs: Optional[Dict[str, str]] = None,
+) -> None:
+    t0 = time.time()
+    # bucket_prefix e.g. 's3://bucket-name/sample'
+    remote_paths = []
+    for file_str in progress_bar(dataset.samples[SampleField.FILE_PATH], description="Hashing data files"):
+        path_file = Path(file_str)
+        file_hash = hash_file_xxhash(path_file)
+
+        # Relative path in S3 bucket e.g. 'data/e2/b0/e2b000ac47b19a999bee5456a6addb88.png'
+        relative_path = s3_prefix_from_hash(hash=file_hash, suffix=path_file.suffix)
+
+        # Remote path in S3 bucket e.g. 's3://bucket-name/sample/data/e2/b0/e2b000ac47b19a999bee5456a6addb88.png'
+        remote_path = f"{bucket_prefix}/{relative_path}"
+        remote_paths.append(remote_path)
+
+    dataset.samples = dataset.samples.with_columns(pl.Series(remote_paths).alias(SampleField.REMOTE_PATH))
+
+    user_logger.info(f"Syncing dataset to S3 bucket '{bucket_prefix}'")
+    files_in_s3 = set(s5cmd_utils.list_bucket(bucket_prefix=bucket_prefix, append_envs=envs))
+
+    # Discover data files (images, videos, etc.) missing in s3
+    data_files_missing = dataset.samples.filter(~pl.col(SampleField.REMOTE_PATH).is_in(files_in_s3))
+    files_already_in_s3 = dataset.samples.filter(pl.col(SampleField.REMOTE_PATH).is_in(files_in_s3))
+
+    with tempfile.TemporaryDirectory() as temp_dir:  # Temp folder to store metadata files
+        path_temp = Path(temp_dir)
+        # File paths are dropped when uploading to S3
+        dataset = dataset.update_samples(dataset.samples.drop(SampleField.FILE_PATH))
+        dataset.write_annotations(path_temp)
+
+        # Discover versioned metadata files (e.g. "annotations.jsonl", "dataset_info.json") missing in s3
+        metadata_files_local = []
+        metadata_files_s3 = []
+        for filename in path_temp.iterdir():
+            metadata_files_s3.append(f"{bucket_prefix}/versions/{dataset.info.version}/{filename.name}")
+            metadata_files_local.append(filename.as_posix())
+
+        overwrite_metadata_files = files_in_s3.intersection(set(metadata_files_s3))
+        will_overwrite_metadata_files = len(overwrite_metadata_files) > 0
+
+        n_files_already_in_s3 = len(files_already_in_s3)
+        user_logger.info(f"Sync dataset to {bucket_prefix}")
+        user_logger.info(
+            f"- Found that {n_files_already_in_s3} / {len(dataset.samples)} data files already exist. "
+            f"Meaning {len(data_files_missing)} data files will be uploaded. \n"
+            f"- Will upload {len(metadata_files_local)} metadata files. \n"
+            f"- Total files to upload: {len(data_files_missing) + len(metadata_files_local)}"
+        )
+        if will_overwrite_metadata_files:
+            msg = f"Metadata files for dataset version '{dataset.info.version}' already exist"
+            if allow_version_overwrite:
+                user_logger.warning(
+                    f"- WARNING: {msg}. Version will be overwritten as 'allow_version_overwrite=True' is set."
+                )
+            else:
+                raise ValueError(
+                    f"Upload cancelled. {msg}. \nTo overwrite existing metadata files, "
+                    "you will need to set 'allow_version_overwrite=True' explicitly."
+                )
+
+        has_missing_files = len(data_files_missing) > 0
+        if interactive and (has_missing_files or will_overwrite_metadata_files):
+            print("Please type 'yes' to upload files.")
+            confirmation = input("Do you want to continue? (yes/NO): ").strip().lower()
+
+            if confirmation != "yes":
+                raise RuntimeError("Upload cancelled by user.")
+
+        local_paths = metadata_files_local + data_files_missing[SampleField.FILE_PATH].to_list()
+        s3_paths = metadata_files_s3 + data_files_missing[SampleField.REMOTE_PATH].to_list()
+        s5cmd_utils.fast_copy_files(local_paths, s3_paths, append_envs=envs, description="Uploading files")
+    user_logger.info(f"- Synced dataset in {time.time() - t0:.2f} seconds.")
+
+
+def sync_dataset_files_to_platform(
+    dataset: HafniaDataset,
+    sample_dataset: Optional[HafniaDataset] = None,
+    interactive: bool = True,
+    allow_version_overwrite: bool = False,
+    cfg: Optional[Config] = None,
+) -> None:
+    cfg = cfg or Config()
+    resource_credentials = get_upload_credentials(dataset.info.dataset_name, cfg=cfg)
+
+    if resource_credentials is None:
+        raise RuntimeError("Failed to get upload credentials from the platform.")
+
+    sync_dataset_files_to_platform_from_resource_credentials(
+        dataset=dataset,
+        sample_dataset=sample_dataset,
+        interactive=interactive,
+        allow_version_overwrite=allow_version_overwrite,
+        resource_credentials=resource_credentials,
+    )
+
+
+def sync_dataset_files_to_platform_from_resource_credentials(
+    dataset: HafniaDataset,
+    sample_dataset: Optional[HafniaDataset],
+    interactive: bool,
+    allow_version_overwrite: bool,
+    resource_credentials: ResourceCredentials,
+):
+    envs = resource_credentials.aws_credentials()
+    bucket_name = resource_credentials.bucket_name()
+
+    for dataset_variant_type in [DatasetVariant.SAMPLE, DatasetVariant.HIDDEN]:
+        if dataset_variant_type == DatasetVariant.SAMPLE:
+            if sample_dataset is None:
+                dataset_variant = dataset.create_sample_dataset()
+            else:
+                dataset_variant = sample_dataset
+        else:
+            dataset_variant = dataset
+
+        sync_hafnia_dataset_to_s3(
+            dataset=dataset_variant,
+            bucket_prefix=f"s3://{bucket_name}/{dataset_variant_type.value}",
+            interactive=interactive,
+            allow_version_overwrite=allow_version_overwrite,
+            envs=envs,
+        )
+
+
+def s3_prefix_from_hash(hash: str, suffix: str) -> str:
+    """
+    Generate a relative S3 path from a hash value for objects stored in S3.
+
+    This function deliberately uses a hierarchical directory layout based on the
+    hash prefix to avoid putting too many objects in a single S3 prefix, which
+    can run into AWS S3 rate limits and performance issues. For example, for
+    hash "dfe8f3b1c2a4f5b6c7d8e9f0a1b2c3d4" and suffix ".png", the returned
+    path will be:
+
+        "data/df/e8/dfe8f3b1c2a4f5b6c7d8e9f0a1b2c3d4.png"
+
+    Note: This intentionally differs from when images are stored to disk locally, where
+    a flat path of the form ``data/<hash><suffix>`` is used.
+    """
+    s3_prefix = f"data/{hash[:2]}/{hash[2:4]}/{hash}{suffix}"
+    return s3_prefix

hafnia/dataset/operations/table_transformations.py

@@ -4,8 +4,6 @@ from typing import List, Optional, Tuple, Type
 import polars as pl
 
 from hafnia.dataset.dataset_names import (
-    FILENAME_ANNOTATIONS_JSONL,
-    FILENAME_ANNOTATIONS_PARQUET,
     PrimitiveField,
     SampleField,
 )
@@ -45,7 +43,8 @@ def create_primitive_table(
         remove_no_object_frames = remove_no_object_frames.drop(drop_columns_names)
         # Rename columns "height", "width" and "meta" for sample to avoid conflicts with object fields names
         remove_no_object_frames = remove_no_object_frames.rename(
-            {"height": "image.height", "width": "image.width", "meta": "image.meta"}
+            {"height": "image.height", "width": "image.width", "meta": "image.meta"},
+            strict=False,
         )
         objects_df = remove_no_object_frames.explode(column_name).unnest(column_name)
     else:
@@ -203,22 +202,6 @@ def split_primitive_columns_by_task_name(
     return samples_table
 
 
-def read_samples_from_path(path: Path) -> pl.DataFrame:
-    path_annotations = path / FILENAME_ANNOTATIONS_PARQUET
-    if path_annotations.exists():
-        user_logger.info(f"Reading dataset annotations from Parquet file: {path_annotations}")
-        return pl.read_parquet(path_annotations)
-
-    path_annotations_jsonl = path / FILENAME_ANNOTATIONS_JSONL
-    if path_annotations_jsonl.exists():
-        user_logger.info(f"Reading dataset annotations from JSONL file: {path_annotations_jsonl}")
-        return pl.read_ndjson(path_annotations_jsonl)
-
-    raise FileNotFoundError(
-        f"Unable to read annotations. No json file '{path_annotations.name}' or Parquet file '{{path_annotations.name}} in in '{path}'."
-    )
-
-
 def check_image_paths(table: pl.DataFrame) -> bool:
     missing_files = []
     org_paths = table[SampleField.FILE_PATH].to_list()

hafnia/http.py

@@ -24,7 +24,8 @@ def fetch(endpoint: str, headers: Dict, params: Optional[Dict] = None) -> Union[
     try:
         response = http.request("GET", endpoint, fields=params, headers=headers)
         if response.status != 200:
-            raise urllib3.exceptions.HTTPError(f"Request failed with status {response.status}")
+            error_details = response.data.decode("utf-8")
+            raise urllib3.exceptions.HTTPError(f"Request failed with status {response.status}: {error_details}")
 
         return json.loads(response.data.decode("utf-8"))
     finally: