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

hafnia/dataset/primitives/polygon.py

@@ -0,0 +1,100 @@
+from typing import Any, Dict, List, Optional, Sequence, Tuple
+
+import cv2
+import numpy as np
+
+from hafnia.dataset.primitives.bitmask import Bitmask
+from hafnia.dataset.primitives.point import Point
+from hafnia.dataset.primitives.primitive import Primitive
+from hafnia.dataset.primitives.utils import class_color_by_name, get_class_name
+
+
+class Polygon(Primitive):
+    # Names should match names in FieldName
+    points: List[Point]
+    class_name: Optional[str] = None  # This should match the string in 'FieldName.CLASS_NAME'
+    class_idx: Optional[int] = None  # This should match the string in 'FieldName.CLASS_IDX'
+    object_id: Optional[str] = None  # This should match the string in 'FieldName.OBJECT_ID'
+    confidence: Optional[float] = None  # Confidence score (0-1.0) for the primitive, e.g. 0.95 for Bbox
+    ground_truth: bool = True  # Whether this is ground truth or a prediction
+
+    task_name: str = ""  # Task name to support multiple Polygon tasks in the same dataset. "" defaults to "polygon"
+    meta: Optional[Dict[str, Any]] = None  # This can be used to store additional information about the bitmask
+
+    @staticmethod
+    def from_list_of_points(
+        points: Sequence[Sequence[float]],
+        class_name: Optional[str] = None,
+        class_idx: Optional[int] = None,
+        object_id: Optional[str] = None,
+    ) -> "Polygon":
+        list_points = [Point(x=point[0], y=point[1]) for point in points]
+        return Polygon(points=list_points, class_name=class_name, class_idx=class_idx, object_id=object_id)
+
+    @staticmethod
+    def default_task_name() -> str:
+        return "polygon"
+
+    @staticmethod
+    def column_name() -> str:
+        return "polygons"
+
+    def calculate_area(self) -> float:
+        raise NotImplementedError()
+
+    def to_pixel_coordinates(
+        self, image_shape: Tuple[int, int], as_int: bool = True, clip_values: bool = True
+    ) -> List[Tuple]:
+        points = [
+            point.to_pixel_coordinates(image_shape=image_shape, as_int=as_int, clip_values=clip_values)
+            for point in self.points
+        ]
+        return points
+
+    def draw(self, image: np.ndarray, inplace: bool = False) -> np.ndarray:
+        if not inplace:
+            image = image.copy()
+        points = np.array(self.to_pixel_coordinates(image_shape=image.shape[:2]))
+
+        bottom_left_idx = np.lexsort((-points[:, 1], points[:, 0]))[0]
+        bottom_left_np = points[bottom_left_idx, :]
+        margin = 5
+        bottom_left = (bottom_left_np[0] + margin, bottom_left_np[1] - margin)
+
+        class_name = self.get_class_name()
+        color = class_color_by_name(class_name)
+        font = cv2.FONT_HERSHEY_SIMPLEX
+        cv2.polylines(image, [points], isClosed=True, color=(0, 255, 0), thickness=2)
+        cv2.putText(
+            img=image, text=class_name, org=bottom_left, fontFace=font, fontScale=0.75, color=color, thickness=2
+        )
+        return image
+
+    def anonymize_by_blurring(self, image: np.ndarray, inplace: bool = False, max_resolution: int = 20) -> np.ndarray:
+        if not inplace:
+            image = image.copy()
+        points = np.array(self.to_pixel_coordinates(image_shape=image.shape[:2]))
+        mask = np.zeros(image.shape[:2], dtype=np.uint8)
+        mask = cv2.fillPoly(mask, [points], color=255).astype(bool)
+        bitmask = Bitmask.from_mask(mask=mask, top=0, left=0).squeeze_mask()
+        image = bitmask.anonymize_by_blurring(image=image, inplace=inplace, max_resolution=max_resolution)
+
+        return image
+
+    def mask(
+        self, image: np.ndarray, inplace: bool = False, color: Optional[Tuple[np.uint8, np.uint8, np.uint8]] = None
+    ) -> np.ndarray:
+        if not inplace:
+            image = image.copy()
+        points = self.to_pixel_coordinates(image_shape=image.shape[:2])
+
+        if color is None:
+            mask = np.zeros_like(image[:, :, 0])
+            bitmask = cv2.fillPoly(mask, pts=[np.array(points)], color=255).astype(bool)  # type: ignore[assignment]
+            color = tuple(int(value) for value in np.mean(image[bitmask], axis=0))  # type: ignore[assignment]
+
+        cv2.fillPoly(image, [np.array(points)], color=color)
+        return image
+
+    def get_class_name(self) -> str:
+        return get_class_name(self.class_name, self.class_idx)

hafnia/dataset/primitives/primitive.py

@@ -0,0 +1,44 @@
+from __future__ import annotations
+
+from abc import ABCMeta, abstractmethod
+
+import numpy as np
+from pydantic import BaseModel
+
+
+class Primitive(BaseModel, metaclass=ABCMeta):
+    def model_post_init(self, context) -> None:
+        if self.task_name == "":  # type: ignore[has-type] # Hack because 'task_name' doesn't exist in base-class yet.
+            self.task_name = self.default_task_name()
+
+    @staticmethod
+    @abstractmethod
+    def default_task_name() -> str:
+        # E.g. "return bboxes" for Bbox
+        raise NotImplementedError
+
+    @staticmethod
+    @abstractmethod
+    def column_name() -> str:
+        """
+        Name of field used in hugging face datasets for storing annotations
+        E.g. "objects" for Bbox.
+        """
+        pass
+
+    @abstractmethod
+    def calculate_area(self) -> float:
+        # Calculate the area of the primitive
+        pass
+
+    @abstractmethod
+    def draw(self, image: np.ndarray, inplace: bool = False) -> np.ndarray:
+        pass
+
+    @abstractmethod
+    def mask(self, image: np.ndarray, inplace: bool = False) -> np.ndarray:
+        pass
+
+    @abstractmethod
+    def anonymize_by_blurring(self, image: np.ndarray, inplace: bool = False, max_resolution: int = 20) -> np.ndarray:
+        pass

hafnia/dataset/primitives/segmentation.py

@@ -0,0 +1,51 @@
+from typing import Any, Dict, List, Optional, Tuple
+
+import cv2
+import numpy as np
+
+from hafnia.dataset.primitives.primitive import Primitive
+from hafnia.dataset.primitives.utils import get_class_name
+from hafnia.visualizations.colors import get_n_colors
+
+
+class Segmentation(Primitive):
+    # mask: np.ndarray
+    class_names: Optional[List[str]] = None  # This should match the string in 'FieldName.CLASS_NAME'
+    ground_truth: bool = True  # Whether this is ground truth or a prediction
+
+    # confidence: Optional[float] = None  # Confidence score (0-1.0) for the primitive, e.g. 0.95 for Classification
+    task_name: str = (
+        ""  # Task name to support multiple Segmentation tasks in the same dataset. "" defaults to "segmentation"
+    )
+    meta: Optional[Dict[str, Any]] = None  # This can be used to store additional information about the bitmask
+
+    @staticmethod
+    def default_task_name() -> str:
+        return "segmentation"
+
+    @staticmethod
+    def column_name() -> str:
+        return "segmentation"
+
+    def calculate_area(self) -> float:
+        raise NotImplementedError()
+
+    def draw(self, image: np.ndarray, inplace: bool = False) -> np.ndarray:
+        if not inplace:
+            image = image.copy()
+
+        color_mapping = np.asarray(get_n_colors(len(self.class_names)), dtype=np.uint8)  # type: ignore[arg-type]
+        label_image = color_mapping[self.mask]
+        blended = cv2.addWeighted(image, 0.5, label_image, 0.5, 0)
+        return blended
+
+    def mask(
+        self, image: np.ndarray, inplace: bool = False, color: Optional[Tuple[np.uint8, np.uint8, np.uint8]] = None
+    ) -> np.ndarray:
+        return image
+
+    def anonymize_by_blurring(self, image: np.ndarray, inplace: bool = False, max_resolution: int = 20) -> np.ndarray:
+        return image
+
+    def get_class_name(self) -> str:
+        return get_class_name(self.class_name, self.class_idx)

hafnia/dataset/primitives/utils.py

@@ -0,0 +1,51 @@
+import hashlib
+from typing import Optional, Tuple, Union
+
+import cv2
+import numpy as np
+
+
+def text_org_from_left_bottom_to_centered(xy_org: tuple, text: str, font, font_scale: float, thickness: int) -> tuple:
+    xy_text_size = cv2.getTextSize(text, fontFace=font, fontScale=font_scale, thickness=thickness)[0]
+    xy_text_size_half = np.array(xy_text_size) / 2
+    xy_centered_np = xy_org + xy_text_size_half * np.array([-1, 1])
+    xy_centered = tuple(int(value) for value in xy_centered_np)
+    return xy_centered
+
+
+def round_int_clip_value(value: Union[int, float], max_value: int) -> int:
+    return clip(value=int(round(value)), v_min=0, v_max=max_value)  # noqa: RUF046
+
+
+def class_color_by_name(name: str) -> Tuple[int, int, int]:
+    # Create a hash of the class name
+    hash_object = hashlib.md5(name.encode())
+    # Use the hash to generate a color
+    hash_digest = hash_object.hexdigest()
+    color = (int(hash_digest[0:2], 16), int(hash_digest[2:4], 16), int(hash_digest[4:6], 16))
+    return color
+
+
+# Define an abstract base class
+def clip(value, v_min, v_max):
+    return min(max(v_min, value), v_max)
+
+
+def get_class_name(class_name: Optional[str], class_idx: Optional[int]) -> str:
+    if class_name is not None:
+        return class_name
+    if class_idx is not None:
+        return f"IDX:{class_idx}"
+    return "NoName"
+
+
+def anonymize_by_resizing(blur_region: np.ndarray, max_resolution: int = 20) -> np.ndarray:
+    """
+    Removes high-frequency details from a region of an image by resizing it down and then back up.
+    """
+    original_shape = blur_region.shape[:2]
+    resize_factor = max(original_shape) / max_resolution
+    new_size = (int(original_shape[0] / resize_factor), int(original_shape[1] / resize_factor))
+    blur_region_downsized = cv2.resize(blur_region, new_size[::-1], interpolation=cv2.INTER_LINEAR)
+    blur_region_upsized = cv2.resize(blur_region_downsized, original_shape[::-1], interpolation=cv2.INTER_LINEAR)
+    return blur_region_upsized

hafnia/experiment/hafnia_logger.py

@@ -9,12 +9,12 @@ from typing import Dict, Optional, Union
 
 import pyarrow as pa
 import pyarrow.parquet as pq
-from datasets import DatasetDict
 from pydantic import BaseModel, field_validator
 
 from hafnia.data.factory import load_dataset
+from hafnia.dataset.hafnia_dataset import HafniaDataset
 from hafnia.log import sys_logger, user_logger
-from hafnia.utils import is_remote_job, now_as_str
+from hafnia.utils import is_hafnia_cloud_job, now_as_str
 
 
 class EntityType(Enum):
@@ -92,7 +92,7 @@ class HafniaLogger:
         self.schema = Entity.create_schema()
         self.log_environment()
 
-    def load_dataset(self, dataset_name: str) -> DatasetDict:
+    def load_dataset(self, dataset_name: str) -> HafniaDataset:
         """
         Load a dataset from the specified path.
         """
@@ -101,7 +101,7 @@ class HafniaLogger:
 
     def path_local_experiment(self) -> Path:
         """Get the path for local experiment."""
-        if is_remote_job():
+        if is_hafnia_cloud_job():
             raise RuntimeError("Cannot access local experiment path in remote job.")
         return self._local_experiment_path
 
@@ -110,7 +110,7 @@ class HafniaLogger:
         if "MDI_CHECKPOINT_DIR" in os.environ:
             return Path(os.environ["MDI_CHECKPOINT_DIR"])
 
-        if is_remote_job():
+        if is_hafnia_cloud_job():
             return Path("/opt/ml/checkpoints")
         return self.path_local_experiment() / "checkpoints"
 
@@ -119,7 +119,7 @@ class HafniaLogger:
         if "MDI_ARTIFACT_DIR" in os.environ:
             return Path(os.environ["MDI_ARTIFACT_DIR"])
 
-        if is_remote_job():
+        if is_hafnia_cloud_job():
             return Path("/opt/ml/output/data")
 
         return self.path_local_experiment() / "data"
@@ -129,7 +129,7 @@ class HafniaLogger:
         if "MDI_MODEL_DIR" in os.environ:
             return Path(os.environ["MDI_MODEL_DIR"])
 
-        if is_remote_job():
+        if is_hafnia_cloud_job():
             return Path("/opt/ml/model")
 
         return self.path_local_experiment() / "model"

hafnia/helper_testing.py

@@ -0,0 +1,108 @@
+from inspect import getmembers, isfunction, signature
+from pathlib import Path
+from types import FunctionType
+from typing import Any, Callable, Dict, Union, get_origin
+
+from hafnia import utils
+from hafnia.dataset.dataset_names import FILENAME_ANNOTATIONS_JSONL, DatasetVariant
+from hafnia.dataset.hafnia_dataset import HafniaDataset, Sample
+
+MICRO_DATASETS = {
+    "tiny-dataset": utils.PATH_DATASETS / "tiny-dataset",
+    "coco-2017": utils.PATH_DATASETS / "coco-2017",
+}
+
+
+def get_path_workspace() -> Path:
+    return Path(__file__).parents[2]
+
+
+def get_path_expected_images() -> Path:
+    return get_path_workspace() / "tests" / "data" / "expected_images"
+
+
+def get_path_test_data() -> Path:
+    return get_path_workspace() / "tests" / "data"
+
+
+def get_path_micro_hafnia_dataset_no_check() -> Path:
+    return get_path_test_data() / "micro_test_datasets"
+
+
+def get_path_micro_hafnia_dataset(dataset_name: str, force_update=False) -> Path:
+    import pytest
+
+    if dataset_name not in MICRO_DATASETS:
+        raise ValueError(f"Dataset name '{dataset_name}' is not recognized. Available options: {list(MICRO_DATASETS)}")
+    path_dataset = MICRO_DATASETS[dataset_name]
+
+    path_test_dataset = get_path_micro_hafnia_dataset_no_check() / dataset_name
+    path_test_dataset_annotations = path_test_dataset / FILENAME_ANNOTATIONS_JSONL
+    if path_test_dataset_annotations.exists() and not force_update:
+        return path_test_dataset
+
+    hafnia_dataset = HafniaDataset.from_path(path_dataset / DatasetVariant.SAMPLE.value)
+    hafnia_dataset = hafnia_dataset.select_samples(n_samples=3, seed=42)
+    hafnia_dataset.write(path_test_dataset)
+
+    if force_update:
+        pytest.fail(
+            "Sample image and metadata have been updated using 'force_update=True'. Set 'force_update=False' and rerun the test."
+        )
+    pytest.fail("Missing test sample image. Please rerun the test.")
+    return path_test_dataset
+
+
+def get_sample_micro_hafnia_dataset(dataset_name: str, force_update=False) -> Sample:
+    micro_dataset = get_micro_hafnia_dataset(dataset_name=dataset_name, force_update=force_update)
+    sample_dict = micro_dataset[0]
+    sample = Sample(**sample_dict)
+    return sample
+
+
+def get_micro_hafnia_dataset(dataset_name: str, force_update: bool = False) -> HafniaDataset:
+    path_dataset = get_path_micro_hafnia_dataset(dataset_name=dataset_name, force_update=force_update)
+    hafnia_dataset = HafniaDataset.from_path(path_dataset)
+    return hafnia_dataset
+
+
+def is_hafnia_configured() -> bool:
+    """
+    Check if Hafnia is configured by verifying if the API key is set.
+    """
+    from cli.config import Config
+
+    return Config().is_configured()
+
+
+def is_typing_type(annotation: Any) -> bool:
+    return get_origin(annotation) is not None
+
+
+def annotation_as_string(annotation: Union[type, str]) -> str:
+    """Convert type annotation to string."""
+    if isinstance(annotation, str):
+        return annotation.replace("'", "")
+    if is_typing_type(annotation):  # Is using typing types like List, Dict, etc.
+        return str(annotation).replace("typing.", "")
+    if hasattr(annotation, "__name__"):
+        return annotation.__name__
+    return str(annotation)
+
+
+def get_hafnia_functions_from_module(python_module) -> Dict[str, FunctionType]:
+    def dataset_is_first_arg(func: Callable) -> bool:
+        """
+        Check if the function has 'HafniaDataset' as the first parameter.
+        """
+        func_signature = signature(func)
+        params = func_signature.parameters
+        if len(params) == 0:
+            return False
+        first_argument_type = list(params.values())[0]
+
+        annotation_as_str = annotation_as_string(first_argument_type.annotation)
+        return annotation_as_str == "HafniaDataset"
+
+    functions = {func[0]: func[1] for func in getmembers(python_module, isfunction) if dataset_is_first_arg(func[1])}
+    return functions

hafnia/http.py

@@ -31,7 +31,7 @@ def fetch(endpoint: str, headers: Dict, params: Optional[Dict] = None) -> Dict:
         http.clear()
 
 
-def post(endpoint: str, headers: Dict, data: Union[Path, Dict, bytes], multipart: bool = False) -> Dict:
+def post(endpoint: str, headers: Dict, data: Union[Path, Dict, bytes, str], multipart: bool = False) -> Dict:
     """Posts data to backend endpoint.
 
     Args:
@@ -64,9 +64,11 @@ def post(endpoint: str, headers: Dict, data: Union[Path, Dict, bytes], multipart
                 with open(data, "rb") as f:
                     body = f.read()
                 response = http.request("POST", endpoint, body=body, headers=headers)
-            elif isinstance(data, dict):
+            elif isinstance(data, (str, dict)):
+                if isinstance(data, dict):
+                    data = json.dumps(data)
                 headers["Content-Type"] = "application/json"
-                response = http.request("POST", endpoint, body=json.dumps(data), headers=headers)
+                response = http.request("POST", endpoint, body=data, headers=headers)
             elif isinstance(data, bytes):
                 response = http.request("POST", endpoint, body=data, headers=headers)
             else:

hafnia/platform/__init__.py

@@ -1,7 +1,7 @@
 from hafnia.platform.download import (
     download_resource,
     download_single_object,
-    get_resource_creds,
+    get_resource_credentials,
 )
 from hafnia.platform.experiment import (
     create_experiment,
@@ -17,5 +17,5 @@ __all__ = [
     "create_experiment",
     "download_resource",
     "download_single_object",
-    "get_resource_creds",
+    "get_resource_credentials",
 ]

hafnia/platform/datasets.py

@@ -0,0 +1,197 @@
+import os
+import shutil
+import subprocess
+import tempfile
+from pathlib import Path
+from typing import Any, Dict, List, Optional
+
+import rich
+from tqdm import tqdm
+
+from cli.config import Config
+from hafnia import utils
+from hafnia.dataset.dataset_names import DATASET_FILENAMES_REQUIRED, ColumnName
+from hafnia.dataset.dataset_recipe.dataset_recipe import (
+    DatasetRecipe,
+    get_dataset_path_from_recipe,
+)
+from hafnia.dataset.hafnia_dataset import HafniaDataset
+from hafnia.http import fetch
+from hafnia.log import sys_logger, user_logger
+from hafnia.platform import get_dataset_id
+from hafnia.platform.download import get_resource_credentials
+from hafnia.utils import timed
+
+
+@timed("Fetching dataset list.")
+def dataset_list(cfg: Optional[Config] = None) -> List[Dict[str, str]]:
+    """List available datasets on the Hafnia platform."""
+    cfg = cfg or Config()
+    endpoint_dataset = cfg.get_platform_endpoint("datasets")
+    header = {"Authorization": cfg.api_key}
+    datasets: List[Dict[str, str]] = fetch(endpoint_dataset, headers=header)  # type: ignore
+    if not datasets:
+        raise ValueError("No datasets found on the Hafnia platform.")
+
+    return datasets
+
+
+def download_or_get_dataset_path(
+    dataset_name: 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."""
+    recipe_explicit = DatasetRecipe.from_implicit_form(dataset_name)
+    path_dataset = get_dataset_path_from_recipe(recipe_explicit, path_datasets=path_datasets_folder)
+
+    is_dataset_valid = HafniaDataset.check_dataset_path(path_dataset, raise_error=False)
+    if is_dataset_valid 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()
+    api_key = cfg.api_key
+
+    shutil.rmtree(path_dataset, ignore_errors=True)
+
+    endpoint_dataset = cfg.get_platform_endpoint("datasets")
+    dataset_id = get_dataset_id(dataset_name=dataset_name, endpoint=endpoint_dataset, api_key=api_key)
+    if dataset_id is None:
+        sys_logger.error(f"Dataset '{dataset_name}' not found on the Hafnia platform.")
+    access_dataset_endpoint = f"{endpoint_dataset}/{dataset_id}/temporary-credentials"
+
+    download_dataset_from_access_endpoint(
+        endpoint=access_dataset_endpoint,
+        api_key=api_key,
+        path_dataset=path_dataset,
+        download_files=download_files,
+    )
+    return path_dataset
+
+
+def download_dataset_from_access_endpoint(
+    endpoint: str,
+    api_key: str,
+    path_dataset: Path,
+    download_files: bool = True,
+) -> None:
+    resource_credentials = get_resource_credentials(endpoint, api_key)
+
+    local_dataset_paths = [str(path_dataset / filename) for filename in DATASET_FILENAMES_REQUIRED]
+    s3_uri = resource_credentials.s3_uri()
+    s3_dataset_files = [f"{s3_uri}/{filename}" for filename in DATASET_FILENAMES_REQUIRED]
+
+    envs = resource_credentials.aws_credentials()
+    fast_copy_files_s3(
+        src_paths=s3_dataset_files,
+        dst_paths=local_dataset_paths,
+        append_envs=envs,
+        description="Downloading annotations",
+    )
+
+    if not download_files:
+        return
+
+    dataset = HafniaDataset.from_path(path_dataset, check_for_images=False)
+    fast_copy_files_s3(
+        src_paths=dataset.samples[ColumnName.REMOTE_PATH].to_list(),
+        dst_paths=dataset.samples[ColumnName.FILE_NAME].to_list(),
+        append_envs=envs,
+        description="Downloading images",
+    )
+
+
+def fast_copy_files_s3(
+    src_paths: List[str],
+    dst_paths: List[str],
+    append_envs: Optional[Dict[str, str]] = None,
+    description: str = "Copying files",
+) -> List[str]:
+    if len(src_paths) != len(dst_paths):
+        raise ValueError("Source and destination paths must have the same length.")
+
+    cmds = [f"cp {src} {dst}" for src, dst in zip(src_paths, dst_paths)]
+    lines = execute_s5cmd_commands(cmds, append_envs=append_envs, description=description)
+    return lines
+
+
+def execute_s5cmd_commands(
+    commands: List[str],
+    append_envs: Optional[Dict[str, str]] = None,
+    description: str = "Executing s5cmd commands",
+) -> List[str]:
+    append_envs = append_envs or {}
+    with tempfile.NamedTemporaryFile(suffix=".txt") as tmp_file:
+        tmp_file_path = Path(tmp_file.name)
+        tmp_file_path.write_text("\n".join(commands))
+        run_cmds = [
+            "s5cmd",
+            "run",
+            str(tmp_file_path),
+        ]
+        envs = os.environ.copy()
+        envs.update(append_envs)
+
+        process = subprocess.Popen(
+            run_cmds,
+            stdout=subprocess.PIPE,
+            stderr=subprocess.STDOUT,
+            universal_newlines=True,
+            env=envs,
+        )
+
+        error_lines = []
+        lines = []
+        for line in tqdm(process.stdout, total=len(commands), desc=description):
+            if "ERROR" in line or "error" in line:
+                error_lines.append(line.strip())
+            lines.append(line.strip())
+
+        if len(error_lines) > 0:
+            show_n_lines = min(5, len(error_lines))
+            str_error_lines = "\n".join(error_lines[:show_n_lines])
+            user_logger.error(
+                f"Detected {len(error_lines)} errors occurred while executing a total of {len(commands)} "
+                f" commands with s5cmd. The first {show_n_lines} is printed below:\n{str_error_lines}"
+            )
+            raise RuntimeError("Errors occurred during s5cmd execution.")
+    return lines
+
+
+TABLE_FIELDS = {
+    "ID": "id",
+    "Hidden\nSamples": "hidden.samples",
+    "Hidden\nSize": "hidden.size",
+    "Sample\nSamples": "sample.samples",
+    "Sample\nSize": "sample.size",
+    "Name": "name",
+    "Title": "title",
+}
+
+
+def create_rich_table_from_dataset(datasets: List[Dict[str, str]]) -> rich.table.Table:
+    datasets = extend_dataset_details(datasets)
+    datasets = sorted(datasets, key=lambda x: x["name"].lower())
+
+    table = rich.table.Table(title="Available Datasets")
+    for i_dataset, dataset in enumerate(datasets):
+        if i_dataset == 0:
+            for column_name, _ in TABLE_FIELDS.items():
+                table.add_column(column_name, justify="left", style="cyan", no_wrap=True)
+        row = [str(dataset.get(field, "")) for field in TABLE_FIELDS.values()]
+        table.add_row(*row)
+
+    return table
+
+
+def extend_dataset_details(datasets: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
+    """Extends dataset details with number of samples and size"""
+    for dataset in datasets:
+        for variant in dataset["dataset_variants"]:
+            variant_type = variant["variant_type"]
+            dataset[f"{variant_type}.samples"] = variant["number_of_data_items"]
+            dataset[f"{variant_type}.size"] = utils.size_human_readable(variant["size_bytes"])
+    return datasets