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

hafnia/dataset/license_types.py

@@ -0,0 +1,63 @@
+from typing import List, Optional
+
+from hafnia.dataset.hafnia_dataset import License
+
+LICENSE_TYPES: List[License] = [
+    License(
+        name="Creative Commons: Attribution-NonCommercial-ShareAlike 2.0 Generic",
+        name_short="CC BY-NC-SA 2.0",
+        url="https://creativecommons.org/licenses/by-nc-sa/2.0/",
+    ),
+    License(
+        name="Creative Commons: Attribution-NonCommercial 2.0 Generic",
+        name_short="CC BY-NC 2.0",
+        url="https://creativecommons.org/licenses/by-nc/2.0/",
+    ),
+    License(
+        name="Creative Commons: Attribution-NonCommercial-NoDerivs 2.0 Generic",
+        name_short="CC BY-NC-ND 2.0",
+        url="https://creativecommons.org/licenses/by-nc-nd/2.0/",
+    ),
+    License(
+        name="Creative Commons: Attribution 2.0 Generic",
+        name_short="CC BY 2.0",
+        url="https://creativecommons.org/licenses/by/2.0/",
+    ),
+    License(
+        name="Creative Commons: Attribution-ShareAlike 2.0 Generic",
+        name_short="CC BY-SA 2.0",
+        url="https://creativecommons.org/licenses/by-sa/2.0/",
+    ),
+    License(
+        name="Creative Commons: Attribution-NoDerivs 2.0 Generic",
+        name_short="CC BY-ND 2.0",
+        url="https://creativecommons.org/licenses/by-nd/2.0/",
+    ),
+    License(
+        name="Flickr: No known copyright restrictions",
+        name_short="Flickr",
+        url="https://flickr.com/commons/usage/",
+    ),
+    License(
+        name="United States Government Work",
+        name_short="US Gov",
+        url="http://www.usa.gov/copyright.shtml",
+    ),
+]
+
+
+def get_license_by_url(url: str) -> Optional[License]:
+    for license in LICENSE_TYPES:
+        # To handle http urls
+        license_url = (license.url or "").replace("http://", "https://")
+        url_https = url.replace("http://", "https://")
+        if license_url == url_https:
+            return license
+    raise ValueError(f"License with URL '{url}' not found.")
+
+
+def get_license_by_short_name(short_name: str) -> Optional[License]:
+    for license in LICENSE_TYPES:
+        if license.name_short == short_name:
+            return license
+    raise ValueError(f"License with short name '{short_name}' not found.")

hafnia/dataset/operations/dataset_stats.py

@@ -1,11 +1,21 @@
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, Dict
+from typing import TYPE_CHECKING, Dict, Optional, Type
 
-from hafnia.dataset.dataset_names import ColumnName
+import polars as pl
+import rich
+from rich import print as rprint
+from rich.table import Table
+from tqdm import tqdm
 
-if TYPE_CHECKING:
+from hafnia.dataset.dataset_names import ColumnName, FieldName, SplitName
+from hafnia.dataset.operations.table_transformations import create_primitive_table
+from hafnia.dataset.primitives import PRIMITIVE_TYPES
+from hafnia.log import user_logger
+
+if TYPE_CHECKING:  # Using 'TYPE_CHECKING' to avoid circular imports during type checking
     from hafnia.dataset.hafnia_dataset import HafniaDataset
+    from hafnia.dataset.primitives.primitive import Primitive
 
 
 def split_counts(dataset: HafniaDataset) -> Dict[str, int]:
@@ -13,3 +23,250 @@ def split_counts(dataset: HafniaDataset) -> Dict[str, int]:
     Returns a dictionary with the counts of samples in each split of the dataset.
     """
     return dict(dataset.samples[ColumnName.SPLIT].value_counts().iter_rows())
+
+
+def class_counts_for_task(
+    dataset: HafniaDataset,
+    primitive: Optional[Type[Primitive]] = None,
+    task_name: Optional[str] = None,
+) -> Dict[str, int]:
+    """
+    Determines class name counts for a specific task in the dataset.
+
+    The counts are returned as a dictionary where keys are class names and values are their respective counts.
+    Note that class names with zero counts are included in the dictionary and that
+    the order of the dictionary matches the class idx.
+
+    >>> counts = dataset.class_counts_for_task(primitive=Bbox)
+    >>> print(counts)
+    {
+        'person': 0,  # Note: Zero count classes are included to maintain order with class idx
+        'car': 500,
+        'bicycle': 0,
+        'bus': 100,
+        'truck': 150,
+        'motorcycle': 50,
+    }
+    """
+    task = dataset.info.get_task_by_task_name_and_primitive(task_name=task_name, primitive=primitive)
+    class_counts_df = (
+        dataset.samples[task.primitive.column_name()]
+        .explode()
+        .struct.unnest()
+        .filter(pl.col(FieldName.TASK_NAME) == task.name)[FieldName.CLASS_NAME]
+        .value_counts()
+    )
+
+    # Initialize counts with zero for all classes to ensure zero-count classes are included
+    # and to have class names in the order of class idx
+    class_counts = {name: 0 for name in task.class_names or []}
+    class_counts.update(dict(class_counts_df.iter_rows()))
+
+    return class_counts
+
+
+def class_counts_all(dataset: HafniaDataset) -> Dict[str, int]:
+    """
+    Get class counts for all tasks in the dataset.
+    The counts are returned as a dictionary where keys are in the format
+    '{primitive_name}/{task_name}/{class_name}' and values are their respective counts.
+
+    Example:
+    >>> counts = dataset.class_counts_all()
+    >>> print(counts)
+    {
+        objects/bboxes/car: 500
+        objects/bboxes/person: 0
+        classifications/weather/sunny: 300
+        classifications/weather/rainy: 0
+        ...
+    }
+    """
+    class_counts = {}
+    for task in dataset.info.tasks:
+        if task.class_names is None:
+            raise ValueError(f"Task '{task.name}' does not have class names defined.")
+        class_counts_task = dataset.class_counts_for_task(primitive=task.primitive, task_name=task.name)
+
+        for class_idx, (class_name, count) in enumerate(class_counts_task.items()):
+            count_name = f"{task.primitive.__name__}/{task.name}/{class_name}"
+            class_counts[count_name] = count
+
+    return class_counts
+
+
+def print_stats(dataset: HafniaDataset) -> None:
+    """
+    Prints verbose statistics about the dataset, including dataset name, version,
+    number of samples, and detailed counts of samples and tasks.
+    """
+    table_base = Table(title="Dataset Statistics", show_lines=True, box=rich.box.SIMPLE)
+    table_base.add_column("Property", style="cyan")
+    table_base.add_column("Value")
+    table_base.add_row("Dataset Name", dataset.info.dataset_name)
+    table_base.add_row("Version", dataset.info.version)
+    table_base.add_row("Number of samples", str(len(dataset.samples)))
+    rprint(table_base)
+    print_sample_and_task_counts(dataset)
+    print_class_distribution(dataset)
+
+
+def print_class_distribution(dataset: HafniaDataset) -> None:
+    """
+    Prints the class distribution for each task in the dataset.
+    """
+    for task in dataset.info.tasks:
+        if task.class_names is None:
+            raise ValueError(f"Task '{task.name}' does not have class names defined.")
+        class_counts = dataset.class_counts_for_task(primitive=task.primitive, task_name=task.name)
+
+        # Print class distribution
+        rich_table = Table(title=f"Class Count for '{task.primitive.__name__}/{task.name}'", show_lines=False)
+        rich_table.add_column("Class Name", style="cyan")
+        rich_table.add_column("Class Idx", style="cyan")
+        rich_table.add_column("Count", justify="right")
+        for class_name, count in class_counts.items():
+            class_idx = task.class_names.index(class_name)  # Get class idx from task info
+            rich_table.add_row(class_name, str(class_idx), str(count))
+        rprint(rich_table)
+
+
+def print_sample_and_task_counts(dataset: HafniaDataset) -> None:
+    """
+    Prints a table with sample counts and task counts for each primitive type
+    in total and for each split (train, val, test).
+    """
+    from hafnia.dataset.operations.table_transformations import create_primitive_table
+    from hafnia.dataset.primitives import PRIMITIVE_TYPES
+
+    splits_sets = {
+        "All": SplitName.valid_splits(),
+        "Train": [SplitName.TRAIN],
+        "Validation": [SplitName.VAL],
+        "Test": [SplitName.TEST],
+    }
+    rows = []
+    for split_name, splits in splits_sets.items():
+        dataset_split = dataset.create_split_dataset(splits)
+        table = dataset_split.samples
+        row = {}
+        row["Split"] = split_name
+        row["Sample "] = str(len(table))
+        for PrimitiveType in PRIMITIVE_TYPES:
+            column_name = PrimitiveType.column_name()
+            objects_df = create_primitive_table(table, PrimitiveType=PrimitiveType, keep_sample_data=False)
+            if objects_df is None:
+                continue
+            for (task_name,), object_group in objects_df.group_by(FieldName.TASK_NAME):
+                count = len(object_group[FieldName.CLASS_NAME])
+                row[f"{PrimitiveType.__name__}\n{task_name}"] = str(count)
+        rows.append(row)
+
+    rich_table = Table(title="Dataset Statistics", show_lines=True, box=rich.box.SIMPLE)
+    for i_row, row in enumerate(rows):
+        if i_row == 0:
+            for column_name in row.keys():
+                rich_table.add_column(column_name, justify="left", style="cyan")
+        rich_table.add_row(*[str(value) for value in row.values()])
+    rprint(rich_table)
+
+
+def check_dataset(dataset: HafniaDataset):
+    """
+    Performs various checks on the dataset to ensure its integrity and consistency.
+    Raises errors if any issues are found.
+    """
+    from hafnia.dataset.hafnia_dataset import Sample
+
+    user_logger.info("Checking Hafnia dataset...")
+    assert isinstance(dataset.info.version, str) and len(dataset.info.version) > 0
+    assert isinstance(dataset.info.dataset_name, str) and len(dataset.info.dataset_name) > 0
+
+    sample_dataset = dataset.create_sample_dataset()
+    if len(sample_dataset) == 0:
+        raise ValueError("The dataset does not include a sample dataset")
+
+    actual_splits = dataset.samples.select(pl.col(ColumnName.SPLIT)).unique().to_series().to_list()
+    expected_splits = SplitName.valid_splits()
+    if set(actual_splits) != set(expected_splits):
+        raise ValueError(f"Expected all splits '{expected_splits}' in dataset, but got '{actual_splits}'. ")
+
+    dataset.check_dataset_tasks()
+
+    expected_tasks = dataset.info.tasks
+    distribution = dataset.info.distributions or []
+    distribution_names = [task.name for task in distribution]
+    # Check that tasks found in the 'dataset.table' matches the tasks defined in 'dataset.info.tasks'
+    for PrimitiveType in PRIMITIVE_TYPES:
+        column_name = PrimitiveType.column_name()
+        if column_name not in dataset.samples.columns:
+            continue
+        objects_df = create_primitive_table(dataset.samples, PrimitiveType=PrimitiveType, keep_sample_data=False)
+        if objects_df is None:
+            continue
+        for (task_name,), object_group in objects_df.group_by(FieldName.TASK_NAME):
+            has_task = any([t for t in expected_tasks if t.name == task_name and t.primitive == PrimitiveType])
+            if has_task or (task_name in distribution_names):
+                continue
+            class_names = object_group[FieldName.CLASS_NAME].unique().to_list()
+            raise ValueError(
+                f"Task name '{task_name}' for the '{PrimitiveType.__name__}' primitive is missing in "
+                f"'dataset.info.tasks' for dataset '{task_name}'. Missing task has the following "
+                f"classes: {class_names}. "
+            )
+
+    for sample_dict in tqdm(dataset, desc="Checking samples in dataset"):
+        sample = Sample(**sample_dict)  # noqa: F841
+
+
+def check_dataset_tasks(dataset: HafniaDataset):
+    """
+    Checks that the tasks defined in 'dataset.info.tasks' are consistent with the data in 'dataset.samples'.
+    """
+    dataset.info.check_for_duplicate_task_names()
+
+    for task in dataset.info.tasks:
+        primitive = task.primitive.__name__
+        column_name = task.primitive.column_name()
+        primitive_column = dataset.samples[column_name]
+        msg_something_wrong = (
+            f"Something is wrong with the defined tasks ('info.tasks') in dataset '{dataset.info.dataset_name}'. \n"
+            f"For '{primitive=}' and '{task.name=}' "
+        )
+        if primitive_column.dtype == pl.Null:
+            raise ValueError(msg_something_wrong + "the column is 'Null'. Please check the dataset.")
+
+        if len(dataset) > 0:  # Check only performed for non-empty datasets
+            primitive_table = (
+                primitive_column.explode().struct.unnest().filter(pl.col(FieldName.TASK_NAME) == task.name)
+            )
+            if primitive_table.is_empty():
+                raise ValueError(
+                    msg_something_wrong
+                    + f"the column '{column_name}' has no {task.name=} objects. Please check the dataset."
+                )
+
+            actual_classes = set(primitive_table[FieldName.CLASS_NAME].unique().to_list())
+            if task.class_names is None:
+                raise ValueError(
+                    msg_something_wrong
+                    + f"the column '{column_name}' with {task.name=} has no defined classes. Please check the dataset."
+                )
+            defined_classes = set(task.class_names)
+
+            if not actual_classes.issubset(defined_classes):
+                raise ValueError(
+                    msg_something_wrong
+                    + f"the column '{column_name}' with {task.name=} we expected the actual classes in the dataset to \n"
+                    f"to be a subset of the defined classes\n\t{actual_classes=} \n\t{defined_classes=}."
+                )
+            # Check class_indices
+            mapped_indices = primitive_table[FieldName.CLASS_NAME].map_elements(
+                lambda x: task.class_names.index(x), return_dtype=pl.Int64
+            )
+            table_indices = primitive_table[FieldName.CLASS_IDX]
+
+            error_msg = msg_something_wrong + (
+                f"class indices in '{FieldName.CLASS_IDX}' column does not match classes ordering in 'task.class_names'"
+            )
+            assert mapped_indices.equals(table_indices), error_msg

hafnia/dataset/operations/dataset_transformations.py

@@ -29,19 +29,27 @@ HafniaDataset class and a RecipeTransform class in the `data_recipe/recipe_trans
 that the signatures match.
 """
 
+import json
+import re
+import textwrap
 from pathlib import Path
-from typing import TYPE_CHECKING, Callable
+from typing import TYPE_CHECKING, Callable, Dict, List, Optional, Set, Tuple, Type, Union
 
 import cv2
+import more_itertools
 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 OPS_REMOVE_CLASS, FieldName
+from hafnia.dataset.primitives import get_primitive_type_from_string
+from hafnia.dataset.primitives.primitive import Primitive
+from hafnia.utils import remove_duplicates_preserve_order
 
-if TYPE_CHECKING:
-    from hafnia.dataset.hafnia_dataset import HafniaDataset
+if TYPE_CHECKING:  # Using 'TYPE_CHECKING' to avoid circular imports during type checking
+    from hafnia.dataset.hafnia_dataset import HafniaDataset, TaskInfo
 
 
 ### Image transformations ###
@@ -79,4 +87,317 @@ def transform_images(
         new_paths.append(str(new_path))
 
     table = dataset.samples.with_columns(pl.Series(new_paths).alias("file_name"))
-    return dataset.update_table(table)
+    return dataset.update_samples(table)
+
+
+def get_task_info_from_task_name_and_primitive(
+    tasks: List["TaskInfo"],
+    task_name: Optional[str] = None,
+    primitive: Union[None, str, Type[Primitive]] = None,
+) -> "TaskInfo":
+    if len(tasks) == 0:
+        raise ValueError("Dataset has no tasks defined.")
+
+    tasks_str = "\n".join([f"\t{task.__repr__()}" for task in tasks])
+    if task_name is None and primitive is None:
+        if len(tasks) == 1:
+            return tasks[0]
+        else:
+            raise ValueError(
+                "For multiple tasks, you will need to specify 'task_name' or 'type_primitive' "
+                "to return a unique task. The dataset contains the following tasks: \n" + tasks_str
+            )
+
+    if isinstance(primitive, str):
+        primitive = get_primitive_type_from_string(primitive)
+
+    tasks_filtered = tasks
+    if primitive is None:
+        tasks_filtered = [task for task in tasks if task.name == task_name]
+
+        if len(tasks_filtered) == 0:
+            raise ValueError(f"No task found with {task_name=}. Available tasks: \n {tasks_str}")
+
+        unique_primitives = set(task.primitive for task in tasks_filtered)
+        if len(unique_primitives) > 1:
+            raise ValueError(
+                f"Found multiple tasks with {task_name=} using different primitives {unique_primitives=}. "
+                "Please specify the primitive type to make it unique. "
+                f"The dataset contains the following tasks: \n {tasks_str}"
+            )
+        primitive = list(unique_primitives)[0]
+
+    if task_name is None:
+        tasks_filtered = [task for task in tasks if task.primitive == primitive]
+        if len(tasks_filtered) == 0:
+            raise ValueError(f"No task found with {primitive=}. Available tasks: \n {tasks_str}")
+
+        unique_task_names = set(task.name for task in tasks_filtered)
+        if len(unique_task_names) > 1:
+            raise ValueError(
+                f"Found multiple tasks with {primitive=} using different task names {unique_task_names=}. "
+                "Please specify the 'task_name' to make it unique."
+                f"The dataset contains the following tasks: \n {tasks_str}"
+            )
+        task_name = list(unique_task_names)[0]
+
+    tasks_filtered = [task for task in tasks_filtered if task.primitive == primitive and task.name == task_name]
+    if len(tasks_filtered) == 0:
+        raise ValueError(f"No task found with {task_name=} and {primitive=}. Available tasks: \n {tasks_str}")
+
+    if len(tasks_filtered) > 1:
+        raise ValueError(
+            f"Multiple tasks found with {task_name=} and {primitive=}. "
+            f"This should never happen. The dataset contains the following tasks: \n {tasks_str}"
+        )
+    task = tasks_filtered[0]
+    return task
+
+
+def class_mapper(
+    dataset: "HafniaDataset",
+    class_mapping: Dict[str, str],
+    method: str = "strict",
+    primitive: Optional[Type[Primitive]] = None,
+    task_name: Optional[str] = None,
+) -> "HafniaDataset":
+    from hafnia.dataset.hafnia_dataset import HafniaDataset
+
+    allowed_methods = ("strict", "remove_undefined", "keep_undefined")
+    if method not in allowed_methods:
+        raise ValueError(f"Method '{method}' is not recognized. Allowed methods are: {allowed_methods}")
+
+    task = dataset.info.get_task_by_task_name_and_primitive(task_name=task_name, primitive=primitive)
+    current_names = task.class_names or []
+
+    # Expand wildcard mappings
+    class_mapping = expand_class_mapping(class_mapping, current_names)
+
+    non_existing_mapping_names = set(class_mapping) - set(current_names)
+    if len(non_existing_mapping_names) > 0:
+        raise ValueError(
+            f"The specified class mapping contains class names {list(non_existing_mapping_names)} "
+            f"that do not exist in the dataset task '{task.name}'. "
+            f"Available class names: {current_names}"
+        )
+
+    missing_class_names = [c for c in current_names if c not in class_mapping]  # List-comprehension to preserve order
+    class_mapping = class_mapping.copy()
+    if method == "strict":
+        pass  # Continue to strict mapping below
+    elif method == "remove_undefined":
+        for missing_class_name in missing_class_names:
+            class_mapping[missing_class_name] = OPS_REMOVE_CLASS
+    elif method == "keep_undefined":
+        for missing_class_name in missing_class_names:
+            class_mapping[missing_class_name] = missing_class_name
+    else:
+        raise ValueError(f"Method '{method}' is not recognized. Allowed methods are: {allowed_methods}")
+
+    missing_class_names = [c for c in current_names if c not in class_mapping]
+    if len(missing_class_names) > 0:
+        error_msg = f"""\
+        The specified class mapping is not a strict mapping - meaning that all class names have not 
+        been mapped to a new class name.
+        In the current mapping, the following classes {list(missing_class_names)} have not been mapped.
+        The currently specified mapping is:
+        {json.dumps(class_mapping, indent=2)}
+        A strict mapping will replace all old class names (dictionary keys) to new class names (dictionary values).
+        Please update the mapping to include all class names from the dataset task '{task.name}'.
+        To keep class map to the same name e.g. 'person' = 'person' 
+        or remove class by using the '__REMOVE__' key, e.g. 'person': '__REMOVE__'."""
+        raise ValueError(textwrap.dedent(error_msg))
+
+    new_class_names = remove_duplicates_preserve_order(class_mapping.values())
+
+    if OPS_REMOVE_CLASS in new_class_names:
+        # Move __REMOVE__ to the end of the list if it exists
+        new_class_names.append(new_class_names.pop(new_class_names.index(OPS_REMOVE_CLASS)))
+    name_2_idx_mapping: Dict[str, int] = {name: idx for idx, name in enumerate(new_class_names)}
+
+    samples = dataset.samples
+    samples_updated = samples.with_columns(
+        pl.col(task.primitive.column_name())
+        .list.eval(
+            pl.element().struct.with_fields(
+                pl.when(pl.field(FieldName.TASK_NAME) == task.name)
+                .then(pl.field(FieldName.CLASS_NAME).replace_strict(class_mapping))
+                .otherwise(pl.field(FieldName.CLASS_NAME))
+                .alias(FieldName.CLASS_NAME)
+            )
+        )
+        .alias(task.primitive.column_name())
+    )
+
+    # Update class indices too
+    samples_updated = samples_updated.with_columns(
+        pl.col(task.primitive.column_name())
+        .list.eval(
+            pl.element().struct.with_fields(
+                pl.when(pl.field(FieldName.TASK_NAME) == task.name)
+                .then(pl.field(FieldName.CLASS_NAME).replace_strict(name_2_idx_mapping))
+                .otherwise(pl.field(FieldName.CLASS_IDX))
+                .alias(FieldName.CLASS_IDX)
+            )
+        )
+        .alias(task.primitive.column_name())
+    )
+
+    if OPS_REMOVE_CLASS in new_class_names:  # Remove class_names that are mapped to REMOVE_CLASS
+        samples_updated = samples_updated.with_columns(
+            pl.col(task.primitive.column_name())
+            .list.filter(pl.element().struct.field(FieldName.CLASS_NAME) != OPS_REMOVE_CLASS)
+            .alias(task.primitive.column_name())
+        )
+
+        new_class_names = [c for c in new_class_names if c != OPS_REMOVE_CLASS]
+
+    new_task = task.model_copy(deep=True)
+    new_task.class_names = new_class_names
+    dataset_info = dataset.info.replace_task(old_task=task, new_task=new_task)
+    return HafniaDataset(info=dataset_info, samples=samples_updated)
+
+
+def expand_class_mapping(wildcard_mapping: Dict[str, str], class_names: List[str]) -> Dict[str, str]:
+    """
+    Expand a wildcard class mapping to a full explicit mapping.
+
+    This function takes a mapping that may contain wildcard patterns (using '*')
+    and expands them to match actual class names from a dataset. Exact matches
+    take precedence over wildcard patterns.
+
+    Examples:
+        >>> from hafnia.dataset.dataset_names import OPS_REMOVE_CLASS
+        >>> wildcard_mapping = {
+        ...     "Person": "Person",
+        ...     "Vehicle.*": "Vehicle",
+        ...     "Vehicle.Trailer": OPS_REMOVE_CLASS
+        ... }
+        >>> class_names = [
+        ...     "Person", "Vehicle.Car", "Vehicle.Trailer", "Vehicle.Bus", "Animal.Dog"
+        ... ]
+        >>> result = expand_wildcard_mapping(wildcard_mapping, class_names)
+        >>> print(result)
+        {
+            "Person": "Person",
+            "Vehicle.Car": "Vehicle",
+            "Vehicle.Trailer": OPS_REMOVE_CLASS,  # Exact match overrides wildcard
+            "Vehicle.Bus": "Vehicle",
+            # Note: "Animal.Dog" is not included as it doesn't match any pattern
+        }
+    """
+    expanded_mapping = {}
+    for match_pattern, mapping_value in wildcard_mapping.items():
+        if "*" in match_pattern:
+            # Convert wildcard pattern to regex: Escape special regex characters except *, then replace * with .*
+            regex_pattern = re.escape(match_pattern).replace("\\*", ".*")
+            class_names_matched = [cn for cn in class_names if re.fullmatch(regex_pattern, cn)]
+            expanded_mapping.update({cn: mapping_value for cn in class_names_matched})
+        else:
+            expanded_mapping.pop(match_pattern, None)
+            expanded_mapping[match_pattern] = mapping_value
+    return expanded_mapping
+
+
+def rename_task(
+    dataset: "HafniaDataset",
+    old_task_name: str,
+    new_task_name: str,
+) -> "HafniaDataset":
+    from hafnia.dataset.hafnia_dataset import HafniaDataset
+
+    old_task = dataset.info.get_task_by_name(task_name=old_task_name)
+    new_task = old_task.model_copy(deep=True)
+    new_task.name = new_task_name
+    samples = dataset.samples.with_columns(
+        pl.col(old_task.primitive.column_name())
+        .list.eval(
+            pl.element().struct.with_fields(
+                pl.field(FieldName.TASK_NAME).replace(old_task.name, new_task.name).alias(FieldName.TASK_NAME)
+            )
+        )
+        .alias(new_task.primitive.column_name())
+    )
+
+    dataset_info = dataset.info.replace_task(old_task=old_task, new_task=new_task)
+    return HafniaDataset(info=dataset_info, samples=samples)
+
+
+def select_samples_by_class_name(
+    dataset: "HafniaDataset",
+    name: Union[List[str], str],
+    task_name: Optional[str] = None,
+    primitive: Optional[Type[Primitive]] = None,
+) -> "HafniaDataset":
+    task, class_names = _validate_inputs_select_samples_by_class_name(
+        dataset=dataset,
+        name=name,
+        task_name=task_name,
+        primitive=primitive,
+    )
+
+    samples = dataset.samples.filter(
+        pl.col(task.primitive.column_name())
+        .list.eval(
+            pl.element().struct.field(FieldName.CLASS_NAME).is_in(class_names)
+            & (pl.element().struct.field(FieldName.TASK_NAME) == task.name)
+        )
+        .list.any()
+    )
+
+    dataset_updated = dataset.update_samples(samples)
+    return dataset_updated
+
+
+def _validate_inputs_select_samples_by_class_name(
+    dataset: "HafniaDataset",
+    name: Union[List[str], str],
+    task_name: Optional[str] = None,
+    primitive: Optional[Type[Primitive]] = None,
+) -> Tuple["TaskInfo", Set[str]]:
+    if isinstance(name, str):
+        name = [name]
+    names = set(name)
+
+    # Check that specified names are available in at least one of the tasks
+    available_names_across_tasks = set(more_itertools.flatten([t.class_names for t in dataset.info.tasks]))
+    missing_class_names_across_tasks = names - available_names_across_tasks
+    if len(missing_class_names_across_tasks) > 0:
+        raise ValueError(
+            f"The specified names {list(names)} have not been found in any of the tasks. "
+            f"Available class names: {available_names_across_tasks}"
+        )
+
+    # Auto infer task if task_name and primitive are not provided
+    if task_name is None and primitive is None:
+        tasks_with_names = [t for t in dataset.info.tasks if names.issubset(t.class_names or [])]
+        if len(tasks_with_names) == 0:
+            raise ValueError(
+                f"The specified names {list(names)} have not been found in any of the tasks. "
+                f"Available class names: {available_names_across_tasks}"
+            )
+        if len(tasks_with_names) > 1:
+            raise ValueError(
+                f"Found multiple tasks containing the specified names {list(names)}. "
+                f"Specify either 'task_name' or 'primitive' to only select from one task. "
+                f"Tasks containing all provided names: {[t.name for t in tasks_with_names]}"
+            )
+
+        task = tasks_with_names[0]
+
+    else:
+        task = get_task_info_from_task_name_and_primitive(
+            tasks=dataset.info.tasks,
+            task_name=task_name,
+            primitive=primitive,
+        )
+
+    task_class_names = set(task.class_names or [])
+    missing_class_names = names - task_class_names
+    if len(missing_class_names) > 0:
+        raise ValueError(
+            f"The specified names {list(missing_class_names)} have not been found for the '{task.name}' task. "
+            f"Available class names: {task_class_names}"
+        )
+
+    return task, names