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

hafnia/dataset/hafnia_dataset.py

@@ -1,7 +1,11 @@
 from __future__ import annotations
 
+import collections
+import copy
+import json
 import shutil
 from dataclasses import dataclass
+from datetime import datetime
 from pathlib import Path
 from random import Random
 from typing import Any, Dict, List, Optional, Type, Union
@@ -9,13 +13,11 @@ from typing import Any, Dict, List, Optional, Type, Union
 import more_itertools
 import numpy as np
 import polars as pl
-import rich
 from PIL import Image
-from pydantic import BaseModel, field_serializer, field_validator
-from rich import print as rprint
-from rich.table import Table
+from pydantic import BaseModel, Field, field_serializer, field_validator
 from tqdm import tqdm
 
+import hafnia
 from hafnia.dataset import dataset_helpers
 from hafnia.dataset.dataset_names import (
     DATASET_FILENAMES_REQUIRED,
@@ -23,20 +25,16 @@ from hafnia.dataset.dataset_names import (
     FILENAME_ANNOTATIONS_PARQUET,
     FILENAME_DATASET_INFO,
     FILENAME_RECIPE_JSON,
+    TAG_IS_SAMPLE,
     ColumnName,
-    FieldName,
     SplitName,
 )
-from hafnia.dataset.operations import dataset_stats, dataset_transformations
+from hafnia.dataset.operations import dataset_stats, dataset_transformations, table_transformations
 from hafnia.dataset.operations.table_transformations import (
     check_image_paths,
-    create_primitive_table,
     read_table_from_path,
 )
-from hafnia.dataset.primitives import (
-    PRIMITIVE_NAME_TO_TYPE,
-    PRIMITIVE_TYPES,
-)
+from hafnia.dataset.primitives import PRIMITIVE_TYPES, get_primitive_type_from_string
 from hafnia.dataset.primitives.bbox import Bbox
 from hafnia.dataset.primitives.bitmask import Bitmask
 from hafnia.dataset.primitives.classification import Classification
@@ -48,9 +46,7 @@ from hafnia.log import user_logger
 class TaskInfo(BaseModel):
     primitive: Type[Primitive]  # Primitive class or string name of the primitive, e.g. "Bbox" or "bitmask"
     class_names: Optional[List[str]]  # Class names for the tasks. To get consistent class indices specify class_names.
-    name: Optional[str] = (
-        None  # None to use the default primitive task name Bbox ->"bboxes", Bitmask -> "bitmasks" etc.
-    )
+    name: Optional[str] = None  # Use 'None' to use default name Bbox ->"bboxes", Bitmask -> "bitmasks" etc.
 
     def model_post_init(self, __context: Any) -> None:
         if self.name is None:
@@ -64,12 +60,7 @@ class TaskInfo(BaseModel):
     @classmethod
     def ensure_primitive(cls, primitive: Any) -> Any:
         if isinstance(primitive, str):
-            if primitive not in PRIMITIVE_NAME_TO_TYPE:
-                raise ValueError(
-                    f"Primitive '{primitive}' is not recognized. Available primitives: {list(PRIMITIVE_NAME_TO_TYPE.keys())}"
-                )
-
-            return PRIMITIVE_NAME_TO_TYPE[primitive]
+            return get_primitive_type_from_string(primitive)
 
         if issubclass(primitive, Primitive):
             return primitive
@@ -83,22 +74,187 @@ class TaskInfo(BaseModel):
             raise ValueError(f"Primitive must be a subclass of Primitive, got {type(primitive)} instead.")
         return primitive.__name__
 
+    @field_validator("class_names", mode="after")
+    @classmethod
+    def validate_unique_class_names(cls, class_names: Optional[List[str]]) -> Optional[List[str]]:
+        """Validate that class names are unique"""
+        if class_names is None:
+            return None
+        duplicate_class_names = set([name for name in class_names if class_names.count(name) > 1])
+        if duplicate_class_names:
+            raise ValueError(
+                f"Class names must be unique. The following class names appear multiple times: {duplicate_class_names}."
+            )
+        return class_names
+
+    # To get unique hash value for TaskInfo objects
+    def __hash__(self) -> int:
+        class_names = self.class_names or []
+        return hash((self.name, self.primitive.__name__, tuple(class_names)))
+
+    def __eq__(self, other: Any) -> bool:
+        if not isinstance(other, TaskInfo):
+            return False
+        return self.name == other.name and self.primitive == other.primitive and self.class_names == other.class_names
+
 
 class DatasetInfo(BaseModel):
     dataset_name: str
-    version: str
-    tasks: list[TaskInfo]
+    version: str  # Dataset version. This is not the same as the Hafnia dataset format version.
+    tasks: List[TaskInfo]
     distributions: Optional[List[TaskInfo]] = None  # Distributions. TODO: FIX/REMOVE/CHANGE this
     meta: Optional[Dict[str, Any]] = None  # Metadata about the dataset, e.g. description, etc.
+    format_version: str = hafnia.__dataset_format_version__  # Version of the Hafnia dataset format
+    updated_at: datetime = datetime.now()
+
+    @field_validator("tasks", mode="after")
+    @classmethod
+    def _validate_check_for_duplicate_tasks(cls, tasks: List[TaskInfo]) -> List[TaskInfo]:
+        task_name_counts = collections.Counter(task.name for task in tasks)
+        duplicate_task_names = [name for name, count in task_name_counts.items() if count > 1]
+        if duplicate_task_names:
+            raise ValueError(
+                f"Tasks must be unique. The following tasks appear multiple times: {duplicate_task_names}."
+            )
+        return tasks
+
+    def check_for_duplicate_task_names(self) -> List[TaskInfo]:
+        return self._validate_check_for_duplicate_tasks(self.tasks)
 
     def write_json(self, path: Path, indent: Optional[int] = 4) -> None:
         json_str = self.model_dump_json(indent=indent)
         path.write_text(json_str)
 
     @staticmethod
-    def from_json_file(path: Path) -> "DatasetInfo":
+    def from_json_file(path: Path) -> DatasetInfo:
         json_str = path.read_text()
-        return DatasetInfo.model_validate_json(json_str)
+
+        # TODO: Deprecated support for old dataset info without format_version
+        # Below 4 lines can be replaced by 'dataset_info = DatasetInfo.model_validate_json(json_str)'
+        # when all datasets include a 'format_version' field
+        json_dict = json.loads(json_str)
+        if "format_version" not in json_dict:
+            json_dict["format_version"] = "0.0.0"
+
+        if "updated_at" not in json_dict:
+            json_dict["updated_at"] = datetime.min.isoformat()
+        dataset_info = DatasetInfo.model_validate(json_dict)
+
+        return dataset_info
+
+    @staticmethod
+    def merge(info0: DatasetInfo, info1: DatasetInfo) -> DatasetInfo:
+        """
+        Merges two DatasetInfo objects into one and validates if they are compatible.
+        """
+        for task_ds0 in info0.tasks:
+            for task_ds1 in info1.tasks:
+                same_name = task_ds0.name == task_ds1.name
+                same_primitive = task_ds0.primitive == task_ds1.primitive
+                same_name_different_primitive = same_name and not same_primitive
+                if same_name_different_primitive:
+                    raise ValueError(
+                        f"Cannot merge datasets with different primitives for the same task name: "
+                        f"'{task_ds0.name}' has primitive '{task_ds0.primitive}' in dataset0 and "
+                        f"'{task_ds1.primitive}' in dataset1."
+                    )
+
+                is_same_name_and_primitive = same_name and same_primitive
+                if is_same_name_and_primitive:
+                    task_ds0_class_names = task_ds0.class_names or []
+                    task_ds1_class_names = task_ds1.class_names or []
+                    if task_ds0_class_names != task_ds1_class_names:
+                        raise ValueError(
+                            f"Cannot merge datasets with different class names for the same task name and primitive: "
+                            f"'{task_ds0.name}' with primitive '{task_ds0.primitive}' has class names "
+                            f"{task_ds0_class_names} in dataset0 and {task_ds1_class_names} in dataset1."
+                        )
+
+        if info1.format_version != info0.format_version:
+            user_logger.warning(
+                "Dataset format version of the two datasets do not match. "
+                f"'{info1.format_version}' vs '{info0.format_version}'."
+            )
+        dataset_format_version = info0.format_version
+        if hafnia.__dataset_format_version__ != dataset_format_version:
+            user_logger.warning(
+                f"Dataset format version '{dataset_format_version}' does not match the current "
+                f"Hafnia format version '{hafnia.__dataset_format_version__}'."
+            )
+        unique_tasks = set(info0.tasks + info1.tasks)
+        distributions = set((info0.distributions or []) + (info1.distributions or []))
+        meta = (info0.meta or {}).copy()
+        meta.update(info1.meta or {})
+        return DatasetInfo(
+            dataset_name=info0.dataset_name + "+" + info1.dataset_name,
+            version="merged",
+            tasks=list(unique_tasks),
+            distributions=list(distributions),
+            meta=meta,
+            format_version=dataset_format_version,
+        )
+
+    def get_task_by_name(self, task_name: str) -> TaskInfo:
+        """
+        Get task by its name. Raises an error if the task name is not found or if multiple tasks have the same name.
+        """
+        tasks_with_name = [task for task in self.tasks if task.name == task_name]
+        if not tasks_with_name:
+            raise ValueError(f"Task with name '{task_name}' not found in dataset info.")
+        if len(tasks_with_name) > 1:
+            raise ValueError(f"Multiple tasks found with name '{task_name}'. This should not happen!")
+        return tasks_with_name[0]
+
+    def get_task_by_primitive(self, primitive: Union[Type[Primitive], str]) -> TaskInfo:
+        """
+        Get task by its primitive type. Raises an error if the primitive type is not found or if multiple tasks
+        have the same primitive type.
+        """
+        if isinstance(primitive, str):
+            primitive = get_primitive_type_from_string(primitive)
+
+        tasks_with_primitive = [task for task in self.tasks if task.primitive == primitive]
+        if not tasks_with_primitive:
+            raise ValueError(f"Task with primitive {primitive} not found in dataset info.")
+        if len(tasks_with_primitive) > 1:
+            raise ValueError(
+                f"Multiple tasks found with primitive {primitive}. Use '{self.get_task_by_name.__name__}' instead."
+            )
+        return tasks_with_primitive[0]
+
+    def get_task_by_task_name_and_primitive(
+        self,
+        task_name: Optional[str],
+        primitive: Optional[Union[Type[Primitive], str]],
+    ) -> TaskInfo:
+        """
+        Logic to get a unique task based on the provided 'task_name' and/or 'primitive'.
+        If both 'task_name' and 'primitive' are None, the dataset must have only one task.
+        """
+        task = dataset_transformations.get_task_info_from_task_name_and_primitive(
+            tasks=self.tasks,
+            primitive=primitive,
+            task_name=task_name,
+        )
+        return task
+
+    def replace_task(self, old_task: TaskInfo, new_task: Optional[TaskInfo]) -> DatasetInfo:
+        dataset_info = self.model_copy(deep=True)
+        has_task = any(t for t in dataset_info.tasks if t.name == old_task.name and t.primitive == old_task.primitive)
+        if not has_task:
+            raise ValueError(f"Task '{old_task.__repr__()}' not found in dataset info.")
+
+        new_tasks = []
+        for task in dataset_info.tasks:
+            if task.name == old_task.name and task.primitive == old_task.primitive:
+                if new_task is None:
+                    continue  # Remove the task
+                new_tasks.append(new_task)
+            else:
+                new_tasks.append(task)
+
+        dataset_info.tasks = new_tasks
+        return dataset_info
 
 
 class Sample(BaseModel):
@@ -106,7 +262,7 @@ class Sample(BaseModel):
     height: int
     width: int
     split: str  # Split name, e.g., "train", "val", "test"
-    is_sample: bool  # Indicates if this is a sample (True) or a metadata entry (False)
+    tags: List[str] = []  # tags for a given sample. Used for creating subsets of the dataset.
     collection_index: Optional[int] = None  # Optional e.g. frame number for video datasets
     collection_id: Optional[str] = None  # Optional e.g. video name for video datasets
     remote_path: Optional[str] = None  # Optional remote path for the image, if applicable
@@ -116,6 +272,7 @@ class Sample(BaseModel):
     bitmasks: Optional[List[Bitmask]] = None  # List of bitmasks, if applicable
     polygons: Optional[List[Polygon]] = None  # List of polygons, if applicable
 
+    attribution: Optional[Attribution] = None  # Attribution information for the image
     meta: Optional[Dict] = None  # Additional metadata, e.g., camera settings, GPS data, etc.
 
     def get_annotations(self, primitive_types: Optional[List[Type[Primitive]]] = None) -> List[Primitive]:
@@ -158,11 +315,93 @@ class Sample(BaseModel):
         return annotations_visualized
 
 
+class License(BaseModel):
+    """License information"""
+
+    name: Optional[str] = Field(
+        default=None,
+        description="License name. E.g. 'Creative Commons: Attribution 2.0 Generic'",
+        max_length=100,
+    )
+    name_short: Optional[str] = Field(
+        default=None,
+        description="License short name or abbreviation. E.g. 'CC BY 4.0'",
+        max_length=100,
+    )
+    url: Optional[str] = Field(
+        default=None,
+        description="License URL e.g. https://creativecommons.org/licenses/by/4.0/",
+    )
+    description: Optional[str] = Field(
+        default=None,
+        description=(
+            "License description e.g. 'You must give appropriate credit, provide a "
+            "link to the license, and indicate if changes were made.'"
+        ),
+    )
+
+    valid_date: Optional[datetime] = Field(
+        default=None,
+        description="License valid date. E.g. '2023-01-01T00:00:00Z'",
+    )
+
+    permissions: Optional[List[str]] = Field(
+        default=None,
+        description="License permissions. Allowed to Access, Label, Distribute, Represent and Modify data.",
+    )
+    liability: Optional[str] = Field(
+        default=None,
+        description="License liability. Optional and not always applicable.",
+    )
+    location: Optional[str] = Field(
+        default=None,
+        description=(
+            "License Location. E.g. Iowa state. This is essential to understand the industry and "
+            "privacy location specific rules that applies to the data. Optional and not always applicable."
+        ),
+    )
+    notes: Optional[str] = Field(
+        default=None,
+        description="Additional license notes. Optional and not always applicable.",
+    )
+
+
+class Attribution(BaseModel):
+    """Attribution information for the image: Giving source and credit to the original creator"""
+
+    title: Optional[str] = Field(default=None, description="Title of the image", max_length=255)
+    creator: Optional[str] = Field(default=None, description="Creator of the image", max_length=255)
+    creator_url: Optional[str] = Field(default=None, description="URL of the creator", max_length=255)
+    date_captured: Optional[datetime] = Field(default=None, description="Date when the image was captured")
+    copyright_notice: Optional[str] = Field(default=None, description="Copyright notice for the image", max_length=255)
+    licenses: Optional[List[License]] = Field(default=None, description="List of licenses for the image")
+    disclaimer: Optional[str] = Field(default=None, description="Disclaimer for the image", max_length=255)
+    changes: Optional[str] = Field(default=None, description="Changes made to the image", max_length=255)
+    source_url: Optional[str] = Field(default=None, description="Source URL for the image", max_length=255)
+
+
 @dataclass
 class HafniaDataset:
     info: DatasetInfo
     samples: pl.DataFrame
 
+    # Function mapping: Dataset stats
+    split_counts = dataset_stats.split_counts
+    class_counts_for_task = dataset_stats.class_counts_for_task
+    class_counts_all = dataset_stats.class_counts_all
+
+    # Function mapping: Print stats
+    print_stats = dataset_stats.print_stats
+    print_sample_and_task_counts = dataset_stats.print_sample_and_task_counts
+    print_class_distribution = dataset_stats.print_class_distribution
+
+    # Function mapping: Dataset checks
+    check_dataset = dataset_stats.check_dataset
+    check_dataset_tasks = dataset_stats.check_dataset_tasks
+
+    # Function mapping: Dataset transformations
+    transform_images = dataset_transformations.transform_images
+
     def __getitem__(self, item: int) -> Dict[str, Any]:
         return self.samples.row(index=item, named=True)
 
@@ -173,6 +412,18 @@ class HafniaDataset:
         for row in self.samples.iter_rows(named=True):
             yield row
 
+    def __post_init__(self):
+        samples = self.samples
+        if ColumnName.SAMPLE_INDEX not in samples.columns:
+            samples = samples.with_row_index(name=ColumnName.SAMPLE_INDEX)
+
+        # Backwards compatibility: If tags-column doesn't exist, create it with empty lists
+        if ColumnName.TAGS not in samples.columns:
+            tags_column: List[List[str]] = [[] for _ in range(len(self))]  # type: ignore[annotation-unchecked]
+            samples = samples.with_columns(pl.Series(tags_column, dtype=pl.List(pl.String)).alias(ColumnName.TAGS))
+
+        self.samples = samples
+
     @staticmethod
     def from_path(path_folder: Path, check_for_images: bool = True) -> "HafniaDataset":
         HafniaDataset.check_dataset_path(path_folder, raise_error=True)
@@ -192,11 +443,12 @@ class HafniaDataset:
         """
         Load a dataset by its name. The dataset must be registered in the Hafnia platform.
         """
-        from hafnia.dataset.hafnia_dataset import HafniaDataset
         from hafnia.platform.datasets import download_or_get_dataset_path
 
         dataset_path = download_or_get_dataset_path(
-            dataset_name=name, force_redownload=force_redownload, download_files=download_files
+            dataset_name=name,
+            force_redownload=force_redownload,
+            download_files=download_files,
         )
         return HafniaDataset.from_path(dataset_path, check_for_images=download_files)
 
@@ -210,9 +462,8 @@ class HafniaDataset:
         else:
             raise TypeError(f"Unsupported sample type: {type(sample)}. Expected Sample or dict.")
 
-        table = pl.from_records(json_samples).drop(ColumnName.SAMPLE_INDEX)
-        table = table.with_row_index(name=ColumnName.SAMPLE_INDEX)  # Add sample index column
-
+        table = pl.from_records(json_samples)
+        table = table.drop(ColumnName.SAMPLE_INDEX).with_row_index(name=ColumnName.SAMPLE_INDEX)
         return HafniaDataset(info=info, samples=table)
 
     @staticmethod
@@ -241,7 +492,11 @@ class HafniaDataset:
         If the dataset is already cached, it will be loaded from the cache.
         """
 
-        path_dataset = get_or_create_dataset_path_from_recipe(dataset_recipe, path_datasets=path_datasets)
+        path_dataset = get_or_create_dataset_path_from_recipe(
+            dataset_recipe,
+            path_datasets=path_datasets,
+            force_redownload=force_redownload,
+        )
         return HafniaDataset.from_path(path_dataset, check_for_images=False)
 
     @staticmethod
@@ -263,20 +518,24 @@ class HafniaDataset:
             merged_dataset = HafniaDataset.merge(merged_dataset, dataset)
         return merged_dataset
 
-    # Dataset transformations
-    transform_images = dataset_transformations.transform_images
-
     def shuffle(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)
+        return dataset.update_samples(table)
 
     def select_samples(
-        dataset: "HafniaDataset", n_samples: int, shuffle: bool = True, seed: int = 42, with_replacement: bool = False
+        dataset: "HafniaDataset",
+        n_samples: int,
+        shuffle: bool = True,
+        seed: int = 42,
+        with_replacement: bool = False,
     ) -> "HafniaDataset":
+        """
+        Create a new dataset with a subset of samples.
+        """
         if not with_replacement:
             n_samples = min(n_samples, len(dataset))
         table = dataset.samples.sample(n=n_samples, with_replacement=with_replacement, seed=seed, shuffle=shuffle)
-        return dataset.update_table(table)
+        return dataset.update_samples(table)
 
     def splits_by_ratios(dataset: "HafniaDataset", split_ratios: Dict[str, float], seed: int = 42) -> "HafniaDataset":
         """
@@ -295,7 +554,7 @@ class HafniaDataset:
             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)
+        return dataset.update_samples(table)
 
     def split_into_multiple_splits(
         dataset: "HafniaDataset",
@@ -323,33 +582,124 @@ class HafniaDataset:
 
         remaining_data = dataset.samples.filter(pl.col(ColumnName.SPLIT).is_in([split_name]).not_())
         new_table = pl.concat([remaining_data, dataset_split_to_be_divided.samples], how="vertical")
-        dataset_new = dataset.update_table(new_table)
+        dataset_new = dataset.update_samples(new_table)
         return dataset_new
 
     def define_sample_set_by_size(dataset: "HafniaDataset", n_samples: int, seed: int = 42) -> "HafniaDataset":
+        """
+        Defines a sample set randomly by selecting 'n_samples' samples from the dataset.
+        """
+        samples = dataset.samples
+
+        # Remove any pre-existing "sample"-tags
+        samples = samples.with_columns(
+            pl.col(ColumnName.TAGS).list.eval(pl.element().filter(pl.element() != TAG_IS_SAMPLE)).alias(ColumnName.TAGS)
+        )
+
+        # Add "sample" to tags column for the selected samples
         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
+        samples = samples.with_columns(
+            pl.when(pl.int_range(len(samples)).is_in(is_sample_indices))
+            .then(pl.col(ColumnName.TAGS).list.concat(pl.lit([TAG_IS_SAMPLE])))
+            .otherwise(pl.col(ColumnName.TAGS))
+        )
+        return dataset.update_samples(samples)
+
+    def class_mapper(
+        dataset: "HafniaDataset",
+        class_mapping: Dict[str, str],
+        method: str = "strict",
+        primitive: Optional[Type[Primitive]] = None,
+        task_name: Optional[str] = None,
+    ) -> "HafniaDataset":
+        """
+        Map class names to new class names using a strict mapping.
+        A strict mapping means that all class names in the dataset must be mapped to a new class name.
+        If a class name is not mapped, an error is raised.
+
+        The class indices are determined by the order of appearance of the new class names in the mapping.
+        Duplicates in the new class names are removed, preserving the order of first appearance.
+
+        E.g.
+
+        mnist = HafniaDataset.from_name("mnist")
+        strict_class_mapping = {
+            "1 - one": "odd",   # 'odd' appears first and becomes class index 0
+            "3 - three": "odd",
+            "5 - five": "odd",
+            "7 - seven": "odd",
+            "9 - nine": "odd",
+            "0 - zero": "even",  # 'even' appears second and becomes class index 1
+            "2 - two": "even",
+            "4 - four": "even",
+            "6 - six": "even",
+            "8 - eight": "even",
+        }
+
+        dataset_new = class_mapper(dataset=mnist, class_mapping=strict_class_mapping)
 
-        table = dataset.samples.with_columns(pl.Series(is_sample_column).alias("is_sample"))
-        return dataset.update_table(table)
+        """
+        return dataset_transformations.class_mapper(
+            dataset=dataset,
+            class_mapping=class_mapping,
+            method=method,
+            primitive=primitive,
+            task_name=task_name,
+        )
+
+    def rename_task(
+        dataset: "HafniaDataset",
+        old_task_name: str,
+        new_task_name: str,
+    ) -> "HafniaDataset":
+        """
+        Rename a task in the dataset.
+        """
+        return dataset_transformations.rename_task(
+            dataset=dataset, old_task_name=old_task_name, new_task_name=new_task_name
+        )
+
+    def select_samples_by_class_name(
+        dataset: HafniaDataset,
+        name: Union[List[str], str],
+        task_name: Optional[str] = None,
+        primitive: Optional[Type[Primitive]] = None,
+    ) -> HafniaDataset:
+        """
+        Select samples that contain at least one annotation with the specified class name(s).
+        If 'task_name' and 'primitive' are not provided, the function will attempt to infer the task.
+        """
+        return dataset_transformations.select_samples_by_class_name(
+            dataset=dataset, name=name, task_name=task_name, primitive=primitive
+        )
 
     def merge(dataset0: "HafniaDataset", dataset1: "HafniaDataset") -> "HafniaDataset":
         """
         Merges two Hafnia datasets by concatenating their samples and updating the split names.
         """
-        ## Currently, only a very naive merging is implemented.
-        # In the future we need to verify that the class and tasks are compatible.
-        # Do they have similar classes and tasks? What to do if they don't?
-        # For now, we just concatenate the samples and keep the split names as they are.
-        merged_samples = pl.concat([dataset0.samples, dataset1.samples], how="vertical")
-        return dataset0.update_table(merged_samples)
 
-    # Dataset stats
-    split_counts = dataset_stats.split_counts
+        # Merges dataset info and checks for compatibility
+        merged_info = DatasetInfo.merge(dataset0.info, dataset1.info)
+
+        # Merges samples tables (removes incompatible columns)
+        merged_samples = table_transformations.merge_samples(samples0=dataset0.samples, samples1=dataset1.samples)
+
+        # Check if primitives have been removed during the merge_samples
+        for task in copy.deepcopy(merged_info.tasks):
+            if task.primitive.column_name() not in merged_samples.columns:
+                user_logger.warning(
+                    f"Task '{task.name}' with primitive '{task.primitive.__name__}' has been removed during the merge. "
+                    "This happens if the two datasets do not have the same primitives."
+                )
+                merged_info = merged_info.replace_task(old_task=task, new_task=None)
+
+        return HafniaDataset(info=merged_info, samples=merged_samples)
 
     def as_dict_dataset_splits(self) -> Dict[str, "HafniaDataset"]:
+        """
+        Splits the dataset into multiple datasets based on the 'split' column.
+        Returns a dictionary with split names as keys and HafniaDataset objects as values.
+        """
         if ColumnName.SPLIT not in self.samples.columns:
             raise ValueError(f"Dataset must contain a '{ColumnName.SPLIT}' column.")
 
@@ -360,10 +710,22 @@ class HafniaDataset:
         return splits
 
     def create_sample_dataset(self) -> "HafniaDataset":
-        if ColumnName.IS_SAMPLE not in self.samples.columns:
-            raise ValueError(f"Dataset must contain an '{ColumnName.IS_SAMPLE}' column.")
-        table = self.samples.filter(pl.col(ColumnName.IS_SAMPLE))
-        return self.update_table(table)
+        # Backwards compatibility. Remove in future versions when dataset have been updated
+        if "is_sample" in self.samples.columns:
+            user_logger.warning(
+                "'is_sample' column found in the dataset. This column is deprecated and will be removed in future versions. "
+                "Please use the 'tags' column with the tag 'sample' instead."
+            )
+            table = self.samples.filter(pl.col("is_sample") == True)  # noqa: E712
+            return self.update_samples(table)
+
+        if ColumnName.TAGS not in self.samples.columns:
+            raise ValueError(f"Dataset must contain an '{ColumnName.TAGS}' column.")
+
+        table = self.samples.filter(
+            pl.col(ColumnName.TAGS).list.eval(pl.element().filter(pl.element() == TAG_IS_SAMPLE)).list.len() > 0
+        )
+        return self.update_samples(table)
 
     def create_split_dataset(self, split_name: Union[str | List[str]]) -> "HafniaDataset":
         if isinstance(split_name, str):
@@ -376,16 +738,12 @@ class HafniaDataset:
                 raise ValueError(f"Invalid split name: {split_name}. Valid splits are: {SplitName.valid_splits()}")
 
         filtered_dataset = self.samples.filter(pl.col(ColumnName.SPLIT).is_in(split_names))
-        return self.update_table(filtered_dataset)
+        return self.update_samples(filtered_dataset)
 
-    def get_task_by_name(self, task_name: str) -> TaskInfo:
-        for task in self.info.tasks:
-            if task.name == task_name:
-                return task
-        raise ValueError(f"Task with name {task_name} not found in dataset info.")
-
-    def update_table(self, table: pl.DataFrame) -> "HafniaDataset":
-        return HafniaDataset(info=self.info.model_copy(), samples=table)
+    def update_samples(self, table: pl.DataFrame) -> "HafniaDataset":
+        dataset = HafniaDataset(info=self.info.model_copy(deep=True), samples=table)
+        dataset.check_dataset_tasks()
+        return dataset
 
     @staticmethod
     def check_dataset_path(path_dataset: Path, raise_error: bool = True) -> bool:
@@ -411,7 +769,10 @@ class HafniaDataset:
 
         return True
 
-    def write(self, path_folder: Path, add_version: bool = False) -> None:
+    def copy(self) -> "HafniaDataset":
+        return HafniaDataset(info=self.info.model_copy(deep=True), samples=self.samples.clone())
+
+    def write(self, path_folder: Path, add_version: bool = False, drop_null_cols: bool = True) -> None:
         user_logger.info(f"Writing dataset to {path_folder}...")
         if not path_folder.exists():
             path_folder.mkdir(parents=True)
@@ -424,6 +785,10 @@ class HafniaDataset:
             )
             new_relative_paths.append(str(new_path.relative_to(path_folder)))
         table = self.samples.with_columns(pl.Series(new_relative_paths).alias("file_name"))
+
+        if drop_null_cols:  # Drops all unused/Null columns
+            table = table.drop(pl.selectors.by_dtype(pl.Null))
+
         table.write_ndjson(path_folder / FILENAME_ANNOTATIONS_JSONL)  # Json for readability
         table.write_parquet(path_folder / FILENAME_ANNOTATIONS_PARQUET)  # Parquet for speed
         self.info.write_json(path_folder / FILENAME_DATASET_INFO)
@@ -448,51 +813,10 @@ class HafniaDataset:
             return False
         return True
 
-    def print_stats(self) -> None:
-        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", self.info.dataset_name)
-        table_base.add_row("Version", self.info.version)
-        table_base.add_row("Number of samples", str(len(self.samples)))
-        rprint(table_base)
-        rprint(self.info.tasks)
-
-        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 = self.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_hafnia_dataset_from_path(path_dataset: Path) -> None:
     dataset = HafniaDataset.from_path(path_dataset, check_for_images=True)
-    check_hafnia_dataset(dataset)
+    dataset.check_dataset()
 
 
 def get_or_create_dataset_path_from_recipe(
@@ -522,89 +846,3 @@ def get_or_create_dataset_path_from_recipe(
     dataset.write(path_dataset)
 
     return path_dataset
-
-
-def check_hafnia_dataset(dataset: HafniaDataset):
-    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
-
-    is_sample_list = set(dataset.samples.select(pl.col(ColumnName.IS_SAMPLE)).unique().to_series().to_list())
-    if True not in is_sample_list:
-        raise ValueError(f"The dataset should contain '{ColumnName.IS_SAMPLE}=True' samples")
-
-    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}'. ")
-
-    expected_tasks = dataset.info.tasks
-    for task in expected_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 '{primtive_name}' task '{task.name}' in dataset '{dataset.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.")
-
-        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
-
-    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:
-                continue
-            if 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)  # Checks format of all samples with pydantic validation  # noqa: F841