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

cli/__main__.py

@@ -37,7 +37,9 @@ def configure(cfg: Config) -> None:
 
     platform_url = click.prompt("Hafnia Platform URL", type=str, default=consts.DEFAULT_API_URL)
 
-    cfg_profile = ConfigSchema(api_key=api_key, platform_url=platform_url)
+    use_keychain = click.confirm("Store API key in system keychain?", default=False)
+
+    cfg_profile = ConfigSchema(platform_url=platform_url, api_key=api_key, use_keychain=use_keychain)
     cfg.add_profile(profile_name, cfg_profile, set_active=True)
     cfg.save_config()
     profile_cmds.profile_show(cfg)

cli/config.py

@@ -6,6 +6,7 @@ from typing import Dict, List, Optional
 from pydantic import BaseModel, field_validator
 
 import cli.consts as consts
+import cli.keychain as keychain
 from hafnia.log import sys_logger, user_logger
 
 PLATFORM_API_MAPPING = {
@@ -19,9 +20,18 @@ PLATFORM_API_MAPPING = {
 }
 
 
+class SecretStr(str):
+    def __repr__(self):
+        return "********"
+
+    def __str__(self):
+        return "********"
+
+
 class ConfigSchema(BaseModel):
     platform_url: str = ""
     api_key: Optional[str] = None
+    use_keychain: bool = False
 
     @field_validator("api_key")
     def validate_api_key(cls, value: Optional[str]) -> Optional[str]:
@@ -35,7 +45,7 @@ class ConfigSchema(BaseModel):
             sys_logger.warning("API key is missing the 'ApiKey ' prefix. Prefix is being added automatically.")
             value = f"ApiKey {value}"
 
-        return value
+        return SecretStr(value)  # Keeps the API key masked in logs and repr
 
 
 class ConfigFileSchema(BaseModel):
@@ -70,13 +80,32 @@ class Config:
 
     @property
     def api_key(self) -> str:
+        # Check keychain first if enabled
+        if self.config.use_keychain:
+            keychain_key = keychain.get_api_key(self.active_profile)
+            if keychain_key is not None:
+                return keychain_key
+
+        # Fall back to config file
         if self.config.api_key is not None:
             return self.config.api_key
+
         raise ValueError(consts.ERROR_API_KEY_NOT_SET)
 
     @api_key.setter
     def api_key(self, value: str) -> None:
-        self.config.api_key = value
+        # Store in keychain if enabled
+        if self.config.use_keychain:
+            if keychain.store_api_key(self.active_profile, value):
+                # Successfully stored in keychain, don't store in config
+                self.config.api_key = None
+            else:
+                # Keychain storage failed, fall back to config file
+                sys_logger.warning("Failed to store in keychain, falling back to config file")
+                self.config.api_key = value
+        else:
+            # Not using keychain, store in config file
+            self.config.api_key = value
 
     @property
     def platform_url(self) -> str:
@@ -152,8 +181,19 @@ class Config:
             raise ValueError("Failed to parse configuration file")
 
     def save_config(self) -> None:
+        # Create a copy to avoid modifying the original data
+        config_to_save = self.config_data.model_dump()
+
+        # Store API key in keychain if enabled, and don't write to file
+        for profile_name, profile_data in config_to_save["profiles"].items():
+            if profile_data.get("use_keychain", False):
+                api_key = profile_data.get("api_key")
+                if api_key:
+                    keychain.store_api_key(profile_name, api_key)
+                profile_data["api_key"] = None
+
         with open(self.config_path, "w") as f:
-            json.dump(self.config_data.model_dump(), f, indent=4)
+            json.dump(config_to_save, f, indent=4)
 
     def remove_profile(self, profile_name: str) -> None:
         if profile_name not in self.config_data.profiles:

cli/keychain.py

@@ -0,0 +1,88 @@
+"""Keychain storage for API keys using the system keychain."""
+
+from typing import Optional
+
+from hafnia.log import sys_logger
+
+# Keyring is optional - gracefully degrade if not available
+try:
+    import keyring
+
+    KEYRING_AVAILABLE = True
+except ImportError:
+    KEYRING_AVAILABLE = False
+    sys_logger.debug("keyring library not available, keychain storage disabled")
+
+KEYRING_SERVICE_NAME = "hafnia-cli"
+
+
+def store_api_key(profile_name: str, api_key: str) -> bool:
+    """
+    Store an API key in the system keychain.
+
+    Args:
+        profile_name: The profile name to associate with the key
+        api_key: The API key to store
+
+    Returns:
+        True if successfully stored, False otherwise
+    """
+    if not KEYRING_AVAILABLE:
+        sys_logger.warning("Keyring library not available, cannot store API key in keychain")
+        return False
+
+    try:
+        keyring.set_password(KEYRING_SERVICE_NAME, profile_name, api_key)
+        sys_logger.debug(f"Stored API key for profile '{profile_name}' in keychain")
+        return True
+    except Exception as e:
+        sys_logger.warning(f"Failed to store API key in keychain: {e}")
+        return False
+
+
+def get_api_key(profile_name: str) -> Optional[str]:
+    """
+    Retrieve an API key from the system keychain.
+
+    Args:
+        profile_name: The profile name to retrieve the key for
+
+    Returns:
+        The API key if found, None otherwise
+    """
+    if not KEYRING_AVAILABLE:
+        return None
+
+    try:
+        api_key = keyring.get_password(KEYRING_SERVICE_NAME, profile_name)
+        if api_key:
+            sys_logger.debug(f"Retrieved API key for profile '{profile_name}' from keychain")
+        return api_key
+    except Exception as e:
+        sys_logger.warning(f"Failed to retrieve API key from keychain: {e}")
+        return None
+
+
+def delete_api_key(profile_name: str) -> bool:
+    """
+    Delete an API key from the system keychain.
+
+    Args:
+        profile_name: The profile name to delete the key for
+
+    Returns:
+        True if successfully deleted or didn't exist, False on error
+    """
+    if not KEYRING_AVAILABLE:
+        return False
+
+    try:
+        keyring.delete_password(KEYRING_SERVICE_NAME, profile_name)
+        sys_logger.debug(f"Deleted API key for profile '{profile_name}' from keychain")
+        return True
+    except keyring.errors.PasswordDeleteError:
+        # Key didn't exist, which is fine
+        return True
+    except Exception as e:
+        sys_logger.warning(f"Failed to delete API key from keychain: {e}")
+        return False

cli/profile_cmds.py

@@ -50,10 +50,13 @@ def cmd_profile_use(cfg: Config, profile_name: str) -> None:
 @click.option(
     "--activate/--no-activate", help="Activate the created profile after creation", default=True, show_default=True
 )
+@click.option(
+    "--use-keychain", is_flag=True, help="Store API key in system keychain instead of config file", default=False
+)
 @click.pass_obj
-def cmd_profile_create(cfg: Config, name: str, api_url: str, api_key: str, activate: bool) -> None:
+def cmd_profile_create(cfg: Config, name: str, api_url: str, api_key: str, activate: bool, use_keychain: bool) -> None:
     """Create a new profile."""
-    cfg_profile = ConfigSchema(platform_url=api_url, api_key=api_key)
+    cfg_profile = ConfigSchema(platform_url=api_url, api_key=api_key, use_keychain=use_keychain)
 
     cfg.add_profile(profile_name=name, profile=cfg_profile, set_active=activate)
     profile_show(cfg)

hafnia/__init__.py

@@ -3,4 +3,4 @@ from importlib.metadata import version
 __package_name__ = "hafnia"
 __version__ = version(__package_name__)
 
-__dataset_format_version__ = "0.0.2"  # Hafnia dataset format version
+__dataset_format_version__ = "0.2.0"  # Hafnia dataset format version

hafnia/dataset/dataset_helpers.py

@@ -38,12 +38,19 @@ def hash_from_bytes(data: bytes) -> str:
 
 def save_image_with_hash_name(image: np.ndarray, path_folder: Path) -> Path:
     pil_image = Image.fromarray(image)
+    path_image = save_pil_image_with_hash_name(pil_image, path_folder)
+    return path_image
+
+
+def save_pil_image_with_hash_name(image: Image.Image, path_folder: Path, allow_skip: bool = True) -> Path:
     buffer = io.BytesIO()
-    pil_image.save(buffer, format="PNG")
+    image.save(buffer, format="PNG")
     hash_value = hash_from_bytes(buffer.getvalue())
     path_image = Path(path_folder) / relative_path_from_hash(hash=hash_value, suffix=".png")
+    if allow_skip and path_image.exists():
+        return path_image
     path_image.parent.mkdir(parents=True, exist_ok=True)
-    pil_image.save(path_image)
+    image.save(path_image)
     return path_image
 
 

hafnia/dataset/dataset_names.py

@@ -1,5 +1,8 @@
 from enum import Enum
-from typing import List
+from typing import Dict, List, Optional
+
+import boto3
+from pydantic import BaseModel, field_validator
 
 FILENAME_RECIPE_JSON = "recipe.json"
 FILENAME_DATASET_INFO = "dataset_info.json"
@@ -23,7 +26,7 @@ TAG_IS_SAMPLE = "sample"
 OPS_REMOVE_CLASS = "__REMOVE__"
 
 
-class FieldName:
+class PrimitiveField:
     CLASS_NAME: str = "class_name"  # Name of the class this primitive is associated with, e.g. "car" for Bbox
     CLASS_IDX: str = "class_idx"  # Index of the class this primitive is associated with, e.g. 0 for "car" if it is the first class  # noqa: E501
     OBJECT_ID: str = "object_id"  # Unique identifier for the object, e.g. "12345123"
@@ -38,39 +41,150 @@ class FieldName:
         Returns a list of expected field names for primitives.
         """
         return [
-            FieldName.CLASS_NAME,
-            FieldName.CLASS_IDX,
-            FieldName.OBJECT_ID,
-            FieldName.CONFIDENCE,
-            FieldName.META,
-            FieldName.TASK_NAME,
+            PrimitiveField.CLASS_NAME,
+            PrimitiveField.CLASS_IDX,
+            PrimitiveField.OBJECT_ID,
+            PrimitiveField.CONFIDENCE,
+            PrimitiveField.META,
+            PrimitiveField.TASK_NAME,
         ]
 
 
-class ColumnName:
-    SAMPLE_INDEX: str = "sample_index"
-    FILE_NAME: str = "file_name"
+class SampleField:
+    FILE_PATH: str = "file_path"
     HEIGHT: str = "height"
     WIDTH: str = "width"
     SPLIT: str = "split"
+    TAGS: str = "tags"
+
+    CLASSIFICATIONS: str = "classifications"
+    BBOXES: str = "bboxes"
+    BITMASKS: str = "bitmasks"
+    POLYGONS: str = "polygons"
+
+    STORAGE_FORMAT: str = "storage_format"  # E.g. "image", "video", "zip"
+    COLLECTION_INDEX: str = "collection_index"
+    COLLECTION_ID: str = "collection_id"
     REMOTE_PATH: str = "remote_path"  # Path to the file in remote storage, e.g. S3
+    SAMPLE_INDEX: str = "sample_index"
+
     ATTRIBUTION: str = "attribution"  # Attribution for the sample (image/video), e.g. creator, license, source, etc.
-    TAGS: str = "tags"
     META: str = "meta"
+    DATASET_NAME: str = "dataset_name"
+
+
+class StorageFormat:
+    IMAGE: str = "image"
+    VIDEO: str = "video"
+    ZIP: str = "zip"
 
 
 class SplitName:
-    TRAIN = "train"
-    VAL = "validation"
-    TEST = "test"
-    UNDEFINED = "UNDEFINED"
+    TRAIN: str = "train"
+    VAL: str = "validation"
+    TEST: str = "test"
+    UNDEFINED: str = "UNDEFINED"
 
     @staticmethod
     def valid_splits() -> List[str]:
         return [SplitName.TRAIN, SplitName.VAL, SplitName.TEST]
 
+    @staticmethod
+    def all_split_names() -> List[str]:
+        return [*SplitName.valid_splits(), SplitName.UNDEFINED]
+
 
 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.
+        """
+        frozen_credentials = session.get_credentials().get_frozen_credentials()
+        return AwsCredentials(
+            access_key=frozen_credentials.access_key,
+            secret_key=frozen_credentials.secret_key,
+            session_token=frozen_credentials.token,
+            region=session.region_name,
+        )
+
+
+ARN_PREFIX = "arn:aws:s3:::"
+
+
+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

@@ -3,7 +3,7 @@ from __future__ import annotations
 import json
 import os
 from pathlib import Path
-from typing import Any, Callable, Dict, List, Optional, Type, Union
+from typing import Any, Callable, Dict, List, Optional, Tuple, Type, Union
 
 from pydantic import (
     field_serializer,
@@ -12,7 +12,11 @@ from pydantic import (
 
 from hafnia import utils
 from hafnia.dataset.dataset_recipe import recipe_transforms
-from hafnia.dataset.dataset_recipe.recipe_types import RecipeCreation, RecipeTransform, Serializable
+from hafnia.dataset.dataset_recipe.recipe_types import (
+    RecipeCreation,
+    RecipeTransform,
+    Serializable,
+)
 from hafnia.dataset.hafnia_dataset import HafniaDataset
 from hafnia.dataset.primitives.primitive import Primitive
 
@@ -41,6 +45,17 @@ class DatasetRecipe(Serializable):
         creation = FromName(name=name, force_redownload=force_redownload, download_files=download_files)
         return DatasetRecipe(creation=creation)
 
+    @staticmethod
+    def from_name_public_dataset(
+        name: str, force_redownload: bool = False, n_samples: Optional[int] = None
+    ) -> DatasetRecipe:
+        creation = FromNamePublicDataset(
+            name=name,
+            force_redownload=force_redownload,
+            n_samples=n_samples,
+        )
+        return DatasetRecipe(creation=creation)
+
     @staticmethod
     def from_path(path_folder: Path, check_for_images: bool = True) -> DatasetRecipe:
         creation = FromPath(path_folder=path_folder, check_for_images=check_for_images)
@@ -222,7 +237,7 @@ class DatasetRecipe(Serializable):
         """Serialize the dataset recipe to a dictionary."""
         return self.model_dump(mode="json")
 
-    def as_platform_recipe(self, recipe_name: Optional[str]) -> Dict:
+    def as_platform_recipe(self, recipe_name: Optional[str], overwrite: bool = False) -> Dict:
         """Uploads dataset recipe to the hafnia platform."""
         from cli.config import Config
         from hafnia.platform.dataset_recipe import get_or_create_dataset_recipe
@@ -235,6 +250,7 @@ class DatasetRecipe(Serializable):
             endpoint=endpoint_dataset,
             api_key=cfg.api_key,
             name=recipe_name,
+            overwrite=overwrite,
         )
 
         return recipe_dict
@@ -246,10 +262,17 @@ class DatasetRecipe(Serializable):
         return recipe
 
     def select_samples(
-        recipe: DatasetRecipe, n_samples: int, shuffle: bool = True, seed: int = 42, with_replacement: bool = False
+        recipe: DatasetRecipe,
+        n_samples: int,
+        shuffle: bool = True,
+        seed: int = 42,
+        with_replacement: bool = False,
     ) -> DatasetRecipe:
         operation = recipe_transforms.SelectSamples(
-            n_samples=n_samples, shuffle=shuffle, seed=seed, with_replacement=with_replacement
+            n_samples=n_samples,
+            shuffle=shuffle,
+            seed=seed,
+            with_replacement=with_replacement,
         )
         recipe.append_operation(operation)
         return recipe
@@ -273,7 +296,7 @@ class DatasetRecipe(Serializable):
 
     def class_mapper(
         recipe: DatasetRecipe,
-        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,
@@ -400,6 +423,22 @@ class FromName(RecipeCreation):
         return [self.name]
 
 
+class FromNamePublicDataset(RecipeCreation):
+    name: str
+    force_redownload: bool = False
+    n_samples: Optional[int] = None
+
+    @staticmethod
+    def get_function() -> Callable[..., "HafniaDataset"]:
+        return HafniaDataset.from_name_public_dataset
+
+    def as_short_name(self) -> str:
+        return f"Torchvision('{self.name}')"
+
+    def get_dataset_names(self) -> List[str]:
+        return []
+
+
 class FromMerge(RecipeCreation):
     recipe0: DatasetRecipe
     recipe1: DatasetRecipe
@@ -414,7 +453,10 @@ class FromMerge(RecipeCreation):
 
     def get_dataset_names(self) -> List[str]:
         """Get the dataset names from the merged recipes."""
-        names = [*self.recipe0.creation.get_dataset_names(), *self.recipe1.creation.get_dataset_names()]
+        names = [
+            *self.recipe0.creation.get_dataset_names(),
+            *self.recipe1.creation.get_dataset_names(),
+        ]
         return names
 
 
@@ -439,33 +481,3 @@ class FromMerger(RecipeCreation):
         for recipe in self.recipes:
             names.extend(recipe.creation.get_dataset_names())
         return names
-
-
-def extract_dataset_names_from_json_dict(data: dict) -> list[str]:
-    """
-    Extract dataset names recursively from a JSON dictionary added with 'from_name'.
-
-    Even if the same functionality is achieved with `DatasetRecipe.get_dataset_names()`,
-    we want to keep this function in 'dipdatalib' to extract dataset names from json dictionaries
-    directly.
-    """
-    creation_field = data.get("creation")
-    if creation_field is None:
-        return []
-    if creation_field.get("__type__") == "FromName":
-        return [creation_field["name"]]
-    elif creation_field.get("__type__") == "FromMerge":
-        recipe_names = ["recipe0", "recipe1"]
-        dataset_name = []
-        for recipe_name in recipe_names:
-            recipe = creation_field.get(recipe_name)
-            if recipe is None:
-                continue
-            dataset_name.extend(extract_dataset_names_from_json_dict(recipe))
-        return dataset_name
-    elif creation_field.get("__type__") == "FromMerger":
-        dataset_name = []
-        for recipe in creation_field.get("recipes", []):
-            dataset_name.extend(extract_dataset_names_from_json_dict(recipe))
-        return dataset_name
-    return []

hafnia/dataset/dataset_recipe/recipe_transforms.py

@@ -1,4 +1,6 @@
-from typing import Callable, Dict, List, Optional, Type, Union
+from typing import Callable, Dict, List, Optional, Tuple, Type, Union
+
+from pydantic import field_validator
 
 from hafnia.dataset.dataset_recipe.recipe_types import RecipeTransform
 from hafnia.dataset.hafnia_dataset import HafniaDataset
@@ -52,11 +54,25 @@ class DefineSampleSetBySize(RecipeTransform):
 
 
 class ClassMapper(RecipeTransform):
-    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
 
+    @field_validator("class_mapping", mode="after")
+    @classmethod
+    def serialize_class_mapping(cls, value: Union[Dict[str, str], List[Tuple[str, str]]]) -> List[Tuple[str, str]]:
+        # Converts the dictionary class mapping to a list of tuples
+        #  e.g. {"old_class": "new_class", } --> [("old_class", "new_class")]
+        # The reason is that storing class mappings as a dictionary does not preserve order of json fields
+        # when stored in a database as a jsonb field (postgres).
+        # Preserving order of class mapping fields is important as it defines the indices of the classes.
+        # So to ensure that class indices are maintained, we preserve order of json fields, by converting the
+        # dictionary to a list of tuples.
+        if isinstance(value, dict):
+            value = list(value.items())
+        return value
+
     @staticmethod
     def get_function() -> Callable[..., "HafniaDataset"]:
         return HafniaDataset.class_mapper