hafnia 0.1.27__py3-none-any.whl → 0.2.0__py3-none-any.whl

cli/__main__.py

@@ -1,7 +1,7 @@
 #!/usr/bin/env python
 import click
 
-from cli import consts, data_cmds, experiment_cmds, profile_cmds, recipe_cmds, runc_cmds
+from cli import consts, dataset_cmds, experiment_cmds, profile_cmds, recipe_cmds, runc_cmds
 from cli.config import Config, ConfigSchema
 
 
@@ -46,7 +46,7 @@ def clear(cfg: Config) -> None:
 
 
 main.add_command(profile_cmds.profile)
-main.add_command(data_cmds.data)
+main.add_command(dataset_cmds.dataset)
 main.add_command(runc_cmds.runc)
 main.add_command(experiment_cmds.experiment)
 main.add_command(recipe_cmds.recipe)

cli/dataset_cmds.py

@@ -0,0 +1,60 @@
+from pathlib import Path
+from typing import Optional
+
+import click
+from rich import print as rprint
+
+import cli.consts as consts
+from cli.config import Config
+from hafnia import utils
+from hafnia.platform.datasets import create_rich_table_from_dataset
+
+
+@click.group()
+def dataset():
+    """Manage dataset interaction"""
+    pass
+
+
+@dataset.command("ls")
+@click.pass_obj
+def dataset_list(cfg: Config) -> None:
+    """List available datasets on Hafnia platform"""
+
+    from hafnia.platform.datasets import dataset_list
+
+    try:
+        datasets = dataset_list(cfg=cfg)
+    except Exception:
+        raise click.ClickException(consts.ERROR_GET_RESOURCE)
+
+    table = create_rich_table_from_dataset(datasets)
+    rprint(table)
+
+
+@dataset.command("download")
+@click.argument("dataset_name")
+@click.option(
+    "--destination",
+    "-d",
+    default=None,
+    required=False,
+    help=f"Destination folder to save the dataset. Defaults to '{utils.PATH_DATASETS}/<dataset_name>'",
+)
+@click.option("--force", "-f", is_flag=True, default=False, help="Flag to enable force redownload")
+@click.pass_obj
+def data_download(cfg: Config, dataset_name: str, destination: Optional[click.Path], force: bool) -> Path:
+    """Download dataset from Hafnia platform"""
+
+    from hafnia.platform import datasets
+
+    try:
+        path_dataset = datasets.download_or_get_dataset_path(
+            dataset_name=dataset_name,
+            cfg=cfg,
+            path_datasets_folder=destination,
+            force_redownload=force,
+        )
+    except Exception:
+        raise click.ClickException(consts.ERROR_GET_RESOURCE)
+    return path_dataset

cli/runc_cmds.py

@@ -38,7 +38,7 @@ def runc():
 @click.pass_obj
 def launch_local(cfg: Config, exec_cmd: str, dataset: str, image_name: str) -> None:
     """Launch a job within the image."""
-    from hafnia.data.factory import download_or_get_dataset_path
+    from hafnia.platform.datasets import download_or_get_dataset_path
 
     is_local_dataset = "/" in dataset
     if is_local_dataset:

hafnia/data/__init__.py

@@ -1,3 +1,3 @@
-from hafnia.data.factory import load_dataset
+from hafnia.data.factory import get_dataset_path, load_dataset
 
-__all__ = ["load_dataset"]
+__all__ = ["load_dataset", "get_dataset_path"]

hafnia/data/factory.py

@@ -1,67 +1,20 @@
-import os
-import shutil
 from pathlib import Path
-from typing import Optional, Union
 
-from datasets import Dataset, DatasetDict, load_from_disk
+from hafnia.dataset.hafnia_dataset import HafniaDataset
+from hafnia.platform.datasets import download_or_get_dataset_path
 
-from cli.config import Config
-from hafnia import utils
-from hafnia.log import user_logger
-from hafnia.platform import download_resource, get_dataset_id
 
-
-def load_local(dataset_path: Path) -> Union[Dataset, DatasetDict]:
-    """Load a Hugging Face dataset from a local directory path."""
-    if not dataset_path.exists():
-        raise ValueError(f"Can not load dataset, directory does not exist -- {dataset_path}")
-    user_logger.info(f"Loading data from {dataset_path.as_posix()}")
-    return load_from_disk(dataset_path.as_posix())
-
-
-def download_or_get_dataset_path(
-    dataset_name: str,
-    cfg: Optional[Config] = None,
-    output_dir: Optional[str] = None,
-    force_redownload: bool = False,
-) -> Path:
-    """Download or get the path of the dataset."""
-
-    cfg = cfg or Config()
-    endpoint_dataset = cfg.get_platform_endpoint("datasets")
-    api_key = cfg.api_key
-
-    output_dir = output_dir or str(utils.PATH_DATASET)
-    dataset_path_base = Path(output_dir).absolute() / dataset_name
-    dataset_path_base.mkdir(exist_ok=True, parents=True)
-    dataset_path_sample = dataset_path_base / "sample"
-
-    if dataset_path_sample.exists() and not force_redownload:
-        user_logger.info("Dataset found locally. Set 'force=True' or add `--force` flag with cli to re-download")
-        return dataset_path_sample
-
-    dataset_id = get_dataset_id(dataset_name, endpoint_dataset, api_key)
-    dataset_access_info_url = f"{endpoint_dataset}/{dataset_id}/temporary-credentials"
-
-    if force_redownload and dataset_path_sample.exists():
-        # Remove old files to avoid old files conflicting with new files
-        shutil.rmtree(dataset_path_sample, ignore_errors=True)
-    status = download_resource(dataset_access_info_url, str(dataset_path_base), api_key)
-    if status:
-        return dataset_path_sample
-    raise RuntimeError("Failed to download dataset")
-
-
-def load_dataset(dataset_name: str, force_redownload: bool = False) -> Union[Dataset, DatasetDict]:
+def load_dataset(dataset_name: str, force_redownload: bool = False) -> HafniaDataset:
     """Load a dataset either from a local path or from the Hafnia platform."""
 
-    if utils.is_remote_job():
-        path_dataset = Path(os.getenv("MDI_DATASET_DIR", "/opt/ml/input/data/training"))
-        return load_local(path_dataset)
+    path_dataset = get_dataset_path(dataset_name, force_redownload=force_redownload)
+    dataset = HafniaDataset.read_from_path(path_dataset)
+    return dataset
+
 
+def get_dataset_path(dataset_name: str, force_redownload: bool = False) -> Path:
     path_dataset = download_or_get_dataset_path(
         dataset_name=dataset_name,
         force_redownload=force_redownload,
     )
-    dataset = load_local(path_dataset)
-    return dataset
+    return path_dataset

hafnia/dataset/dataset_helpers.py

@@ -0,0 +1,91 @@
+import io
+import math
+import random
+from pathlib import Path
+from typing import Dict, List
+
+import numpy as np
+import xxhash
+from PIL import Image
+
+
+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)
+
+    split_name_column = []
+    for split_name, n_split_samples in samples_per_split.items():
+        split_name_column.extend([split_name] * n_split_samples)
+    random.Random(seed).shuffle(split_name_column)  # Shuffle the split names
+
+    return split_name_column
+
+
+def hash_file_xxhash(path: Path, chunk_size: int = 262144) -> str:
+    hasher = xxhash.xxh3_64()
+
+    with open(path, "rb") as f:
+        for chunk in iter(lambda: f.read(chunk_size), b""):  # 8192, 16384, 32768, 65536
+            hasher.update(chunk)
+    return hasher.hexdigest()
+
+
+def hash_from_bytes(data: bytes) -> str:
+    hasher = xxhash.xxh3_64()
+    hasher.update(data)
+    return hasher.hexdigest()
+
+
+def save_image_with_hash_name(image: np.ndarray, path_folder: Path) -> Path:
+    pil_image = Image.fromarray(image)
+    buffer = io.BytesIO()
+    pil_image.save(buffer, format="PNG")
+    hash_value = hash_from_bytes(buffer.getvalue())
+    path_image = Path(path_folder) / f"{hash_value}.png"
+    pil_image.save(path_image)
+    return path_image
+
+
+def filename_as_hash_from_path(path_image: Path) -> str:
+    hash = hash_file_xxhash(path_image)
+    return f"{hash}{path_image.suffix}"
+
+
+def split_sizes_from_ratios(n_items: int, split_ratios: Dict[str, float]) -> Dict[str, int]:
+    summed_ratios = sum(split_ratios.values())
+    abs_tols = 0.0011  # Allow some tolerance for floating point errors {"test": 0.333, "val": 0.333, "train": 0.333}
+    if not math.isclose(summed_ratios, 1.0, abs_tol=abs_tols):  # Allow tolerance to allow e.g. (0.333, 0.333, 0.333)
+        raise ValueError(f"Split ratios must sum to 1.0. The summed values of {split_ratios} is {summed_ratios}")
+
+    # recaculate split sizes
+    split_ratios = {split_name: split_ratio / summed_ratios for split_name, split_ratio in split_ratios.items()}
+    split_sizes = {split_name: int(n_items * split_ratio) for split_name, split_ratio in split_ratios.items()}
+
+    remaining_items = n_items - sum(split_sizes.values())
+    if remaining_items > 0:  # Distribute remaining items evenly across splits
+        for _ in range(remaining_items):
+            # Select name by the largest error from the expected distribution
+            total_size = sum(split_sizes.values())
+            distribution_error = {
+                split_name: abs(split_ratios[split_name] - (size / total_size))
+                for split_name, size in split_sizes.items()
+            }
+
+            split_with_largest_error = sorted(distribution_error.items(), key=lambda x: x[1], reverse=True)[0][0]
+            split_sizes[split_with_largest_error] += 1
+
+    if sum(split_sizes.values()) != n_items:
+        raise ValueError("Something is wrong. The split sizes do not match the number of items.")
+
+    return split_sizes
+
+
+def select_evenly_across_list(lst: list, num_samples: int):
+    if num_samples >= len(lst):
+        return lst  # No need to sample
+    step = (len(lst) - 1) / (num_samples - 1)
+    indices = [int(round(step * i)) for i in range(num_samples)]  # noqa: RUF046
+    return [lst[index] for index in indices]
+
+
+def prefix_dict(d: dict, prefix: str) -> dict:
+    return {f"{prefix}.{k}": v for k, v in d.items()}

hafnia/dataset/dataset_names.py

@@ -0,0 +1,71 @@
+from enum import Enum
+from typing import List
+
+FILENAME_DATASET_INFO = "dataset_info.json"
+FILENAME_ANNOTATIONS_JSONL = "annotations.jsonl"
+FILENAME_ANNOTATIONS_PARQUET = "annotations.parquet"
+
+DATASET_FILENAMES = [
+    FILENAME_DATASET_INFO,
+    FILENAME_ANNOTATIONS_JSONL,
+    FILENAME_ANNOTATIONS_PARQUET,
+]
+
+
+class DeploymentStage(Enum):
+    STAGING = "staging"
+    PRODUCTION = "production"
+
+
+class FieldName:
+    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
+    )
+    OBJECT_ID: str = "object_id"  # Unique identifier for the object, e.g. "12345123"
+    CONFIDENCE: str = "confidence"  # Confidence score (0-1.0) for the primitive, e.g. 0.95 for Bbox
+
+    META: str = "meta"  # Contains metadata about each primitive, e.g. attributes color, occluded, iscrowd, etc.
+    TASK_NAME: str = "task_name"  # Name of the task this primitive is associated with, e.g. "bboxes" for Bbox
+
+    @staticmethod
+    def fields() -> List[str]:
+        """
+        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,
+        ]
+
+
+class ColumnName:
+    SAMPLE_INDEX: str = "sample_index"
+    FILE_NAME: str = "file_name"
+    HEIGHT: str = "height"
+    WIDTH: str = "width"
+    SPLIT: str = "split"
+    IS_SAMPLE: str = "is_sample"
+    REMOTE_PATH: str = "remote_path"  # Path to the file in remote storage, e.g. S3
+    META: str = "meta"
+
+
+class SplitName:
+    TRAIN = "train"
+    VAL = "validation"
+    TEST = "test"
+    UNDEFINED = "UNDEFINED"
+
+    @staticmethod
+    def valid_splits() -> List[str]:
+        return [SplitName.TRAIN, SplitName.VAL, SplitName.TEST]
+
+
+class DatasetVariant(Enum):
+    DUMP = "dump"
+    SAMPLE = "sample"
+    HIDDEN = "hidden"

hafnia/dataset/dataset_transformation.py

@@ -0,0 +1,187 @@
+import hashlib
+import shutil
+from pathlib import Path
+from random import Random
+from typing import TYPE_CHECKING, Callable, Dict
+
+import cv2
+import numpy as np
+import polars as pl
+from PIL import Image
+from tqdm import tqdm
+
+from hafnia.dataset import dataset_helpers
+from hafnia.dataset.dataset_names import ColumnName
+from hafnia.log import user_logger
+
+if TYPE_CHECKING:
+    from hafnia.dataset.hafnia_dataset import HafniaDataset
+
+
+### Image transformations ###
+class AnonymizeByPixelation:
+    def __init__(self, resize_factor: float = 0.10):
+        self.resize_factor = resize_factor
+
+    def __call__(self, frame: np.ndarray) -> np.ndarray:
+        org_size = frame.shape[:2]
+        frame = cv2.resize(frame, (0, 0), fx=self.resize_factor, fy=self.resize_factor)
+        frame = cv2.resize(frame, org_size[::-1], interpolation=cv2.INTER_NEAREST)
+        return frame
+
+
+def splits_by_ratios(dataset: "HafniaDataset", split_ratios: Dict[str, float], seed: int = 42) -> "HafniaDataset":
+    """
+    Divides the dataset into splits based on the provided ratios.
+
+    Example: Defining split ratios and applying the transformation
+
+    >>> dataset = HafniaDataset.read_from_path(Path("path/to/dataset"))
+    >>> split_ratios = {SplitName.TRAIN: 0.8, SplitName.VAL: 0.1, SplitName.TEST: 0.1}
+    >>> dataset_with_splits = splits_by_ratios(dataset, split_ratios, seed=42)
+    Or use the function as a
+    >>> dataset_with_splits = dataset.splits_by_ratios(split_ratios, seed=42)
+    """
+    n_items = len(dataset)
+    split_name_column = dataset_helpers.create_split_name_list_from_ratios(
+        split_ratios=split_ratios, n_items=n_items, seed=seed
+    )
+    table = dataset.samples.with_columns(pl.Series(split_name_column).alias("split"))
+    return dataset.update_table(table)
+
+
+def divide_split_into_multiple_splits(
+    dataset: "HafniaDataset",
+    divide_split_name: str,
+    split_ratios: Dict[str, float],
+) -> "HafniaDataset":
+    """
+    Divides a dataset split ('divide_split_name') into multiple splits based on the provided split
+    ratios ('split_ratios'). This is especially useful for some open datasets where they have only provide
+    two splits or only provide annotations for two splits. This function allows you to create additional
+    splits based on the provided ratios.
+
+    Example: Defining split ratios and applying the transformation
+    >>> dataset = HafniaDataset.read_from_path(Path("path/to/dataset"))
+    >>> divide_split_name = SplitName.TEST
+    >>> split_ratios = {SplitName.TEST: 0.8, SplitName.VAL: 0.2}
+    >>> dataset_with_splits = divide_split_into_multiple_splits(dataset, divide_split_name, split_ratios)
+    """
+    dataset_split_to_be_divided = dataset.create_split_dataset(split_name=divide_split_name)
+    if len(dataset_split_to_be_divided) == 0:
+        split_counts = dict(dataset.samples.select(pl.col(ColumnName.SPLIT).value_counts()).iter_rows())
+        raise ValueError(
+            f"No samples in the '{divide_split_name}' split to divide into multiple splits. {split_counts=}"
+        )
+    assert len(dataset_split_to_be_divided) > 0, f"No samples in the '{divide_split_name}' split!"
+    dataset_split_to_be_divided = dataset_split_to_be_divided.split_by_ratios(split_ratios=split_ratios, seed=42)
+
+    remaining_data = dataset.samples.filter(pl.col(ColumnName.SPLIT).is_in([divide_split_name]).not_())
+    new_table = pl.concat([remaining_data, dataset_split_to_be_divided.samples], how="vertical")
+    dataset_new = dataset.update_table(new_table)
+    return dataset_new
+
+
+def shuffle_dataset(dataset: "HafniaDataset", seed: int = 42) -> "HafniaDataset":
+    table = dataset.samples.sample(n=len(dataset), with_replacement=False, seed=seed, shuffle=True)
+    return dataset.update_table(table)
+
+
+def sample(dataset: "HafniaDataset", n_samples: int, shuffle: bool = True, seed: int = 42) -> "HafniaDataset":
+    table = dataset.samples.sample(n=n_samples, with_replacement=False, seed=seed, shuffle=shuffle)
+    return dataset.update_table(table)
+
+
+def define_sample_set_by_size(dataset: "HafniaDataset", n_samples: int, seed: int = 42) -> "HafniaDataset":
+    is_sample_indices = Random(seed).sample(range(len(dataset)), n_samples)
+    is_sample_column = [False for _ in range(len(dataset))]
+    for idx in is_sample_indices:
+        is_sample_column[idx] = True
+
+    table = dataset.samples.with_columns(pl.Series(is_sample_column).alias("is_sample"))
+    return dataset.update_table(table)
+
+
+def transform_images(
+    dataset: "HafniaDataset",
+    transform: Callable[[np.ndarray], np.ndarray],
+    path_output: Path,
+) -> "HafniaDataset":
+    new_paths = []
+    path_image_folder = path_output / "data"
+    path_image_folder.mkdir(parents=True, exist_ok=True)
+
+    for org_path in tqdm(dataset.samples["file_name"].to_list(), desc="Transform images"):
+        org_path = Path(org_path)
+        if not org_path.exists():
+            raise FileNotFoundError(f"File {org_path} does not exist in the dataset.")
+
+        image = np.array(Image.open(org_path))
+        image_transformed = transform(image)
+        new_path = dataset_helpers.save_image_with_hash_name(image_transformed, path_image_folder)
+
+        if not new_path.exists():
+            raise FileNotFoundError(f"Transformed file {new_path} does not exist in the dataset.")
+        new_paths.append(str(new_path))
+
+    table = dataset.samples.with_columns(pl.Series(new_paths).alias("file_name"))
+    return dataset.update_table(table)
+
+
+def rename_to_unique_image_names(dataset: "HafniaDataset", path_output: Path) -> "HafniaDataset":
+    user_logger.info(f"Copy images to have unique filenames. New path is '{path_output}'")
+    shutil.rmtree(path_output, ignore_errors=True)  # Remove the output folder if it exists
+    new_paths = []
+    for org_path in tqdm(dataset.samples["file_name"].to_list(), desc="- Rename/copy images"):
+        org_path = Path(org_path)
+        if not org_path.exists():
+            raise FileNotFoundError(f"File {org_path} does not exist in the dataset.")
+
+        hash_name = hashlib.md5(str(org_path).encode()).hexdigest()[
+            :6
+        ]  # Generate a unique name based on the original file name
+        new_path = path_output / "data" / f"{hash_name}_{org_path.name}"
+        if not new_path.parent.exists():
+            new_path.parent.mkdir(parents=True, exist_ok=True)
+
+        shutil.copyfile(org_path, new_path)  # Copy the original file to the new path
+        new_paths.append(str(new_path))
+
+    table = dataset.samples.with_columns(pl.Series(new_paths).alias("file_name"))
+    return dataset.update_table(table)
+
+
+### Hafnia Dataset Transformations ###
+class SplitsByRatios:
+    def __init__(self, split_ratios: dict, seed: int = 42):
+        self.split_ratios = split_ratios
+        self.seed = seed
+
+    def __call__(self, dataset: "HafniaDataset") -> "HafniaDataset":
+        return splits_by_ratios(dataset, self.split_ratios, self.seed)
+
+
+class ShuffleDataset:
+    def __init__(self, seed: int = 42):
+        self.seed = seed
+
+    def __call__(self, dataset: "HafniaDataset") -> "HafniaDataset":
+        return shuffle_dataset(dataset, self.seed)
+
+
+class SampleSetBySize:
+    def __init__(self, n_samples: int, seed: int = 42):
+        self.n_samples = n_samples
+        self.seed = seed
+
+    def __call__(self, dataset: "HafniaDataset") -> "HafniaDataset":
+        return define_sample_set_by_size(dataset, self.n_samples, self.seed)
+
+
+class TransformImages:
+    def __init__(self, transform: Callable[[np.ndarray], np.ndarray], path_output: Path):
+        self.transform = transform
+        self.path_output = path_output
+
+    def __call__(self, dataset: "HafniaDataset") -> "HafniaDataset":
+        return transform_images(dataset, self.transform, self.path_output)