hafnia 0.5.0__py3-none-any.whl → 0.5.2__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,
@@ -102,7 +112,7 @@ 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")
@@ -144,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)
@@ -238,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,
@@ -477,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

@@ -8,13 +8,13 @@ import polars as pl
 from hafnia.dataset.dataset_helpers import hash_file_xxhash
 from hafnia.dataset.dataset_names import (
     DatasetVariant,
-    ResourceCredentials,
     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
 
@@ -39,6 +39,7 @@ def delete_hafnia_dataset_files_on_platform(
 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()
@@ -58,7 +59,11 @@ def delete_hafnia_dataset_files_from_resource_credentials(
             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=True, append_envs=envs)
+    s5cmd_utils.delete_bucket_content(
+        bucket_prefix=f"s3://{bucket_name}",
+        remove_bucket=remove_bucket,
+        append_envs=envs,
+    )
     return True
 
 

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,
 )
@@ -204,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()