maite-datasets 0.0.5__py3-none-any.whl → 0.0.7__py3-none-any.whl

maite_datasets/{_reader → object_detection}/_coco.py

@@ -3,20 +3,19 @@
 from __future__ import annotations
 
 import json
-import logging
 from pathlib import Path
 from typing import Any
 
+import maite.protocols.object_detection as od
 import numpy as np
+from maite.protocols import DatasetMetadata, DatumMetadata
 from PIL import Image
 
-from maite_datasets._protocols import DatasetMetadata, DatumMetadata, ObjectDetectionDataset, ObjectDetectionDatum
-from maite_datasets._reader._base import _ObjectDetectionTarget, BaseDatasetReader
+from maite_datasets._base import BaseDataset, ObjectDetectionTarget
+from maite_datasets._reader import BaseDatasetReader
 
-_logger = logging.getLogger(__name__)
 
-
-class COCODatasetReader(BaseDatasetReader):
+class COCODatasetReader(BaseDatasetReader[od.Dataset]):
     """
     COCO format dataset reader conforming to MAITE protocols.
 
@@ -132,9 +131,9 @@ class COCODatasetReader(BaseDatasetReader):
         """Mapping from class index to class name."""
         return self._index2label
 
-    def _create_dataset_implementation(self) -> ObjectDetectionDataset:
+    def create_dataset(self) -> od.Dataset:
         """Create COCO dataset implementation."""
-        return _COCODataset(self)
+        return COCODataset(self)
 
     def _validate_format_specific(self) -> tuple[list[str], dict[str, Any]]:
         """Validate COCO format specific files and structure."""
@@ -198,37 +197,40 @@ class COCODatasetReader(BaseDatasetReader):
         else:
             class_names = [cat["name"] for cat in self._coco_data["categories"]]
 
-        self._index2label = {idx: name for idx, name in enumerate(class_names)}
+        self._index2label = dict(enumerate(class_names))
 
 
-class _COCODataset:
+class COCODataset(BaseDataset):
     """Internal COCO dataset implementation."""
 
     def __init__(self, reader: COCODatasetReader) -> None:
-        self.reader = reader
-        self.image_ids = list(reader._image_id_to_info.keys())
-
-    @property
-    def metadata(self) -> DatasetMetadata:
-        return DatasetMetadata(
-            id=self.reader.dataset_id,
-            index2label=self.reader.index2label,
+        self._reader = reader
+        self._image_ids = list(reader._image_id_to_info.keys())
+
+        self.root = reader.dataset_path
+        self.images_path = reader._images_path
+        self.annotation_path = reader._annotation_path
+        self.size = len(reader._image_id_to_info)
+        self.classes = reader.index2label
+        self.metadata = DatasetMetadata(
+            id=self._reader.dataset_id,
+            index2label=self._reader.index2label,
         )
 
     def __len__(self) -> int:
-        return len(self.image_ids)
+        return len(self._image_ids)
 
-    def __getitem__(self, index: int) -> ObjectDetectionDatum:
-        image_id = self.image_ids[index]
-        image_info = self.reader._image_id_to_info[image_id]
+    def __getitem__(self, index: int) -> tuple[od.InputType, od.ObjectDetectionTarget, DatumMetadata]:
+        image_id = self._image_ids[index]
+        image_info = self._reader._image_id_to_info[image_id]
 
         # Load image
-        image_path = self.reader._images_path / image_info["file_name"]
+        image_path = self._reader._images_path / image_info["file_name"]
         image = np.array(Image.open(image_path).convert("RGB"))
         image = np.transpose(image, (2, 0, 1))  # Convert to CHW format
 
         # Get annotations for this image
-        annotations = self.reader.image_id_to_annotations.get(image_id, [])
+        annotations = self._reader.image_id_to_annotations.get(image_id, [])
 
         if annotations:
             boxes = []
@@ -241,7 +243,7 @@ class _COCODataset:
                 boxes.append([x, y, x + w, y + h])
 
                 # Map category_id to class index
-                cat_idx = self.reader._category_id_to_idx[ann["category_id"]]
+                cat_idx = self._reader._category_id_to_idx[ann["category_id"]]
                 labels.append(cat_idx)
 
                 # Collect annotation metadata
@@ -267,12 +269,12 @@ class _COCODataset:
             scores = np.empty(0, dtype=np.float32)
             annotation_metadata = []
 
-        target = _ObjectDetectionTarget(boxes, labels, scores)
+        target = ObjectDetectionTarget(boxes, labels, scores)
 
         # Create comprehensive datum metadata
         datum_metadata = DatumMetadata(
             **{
-                "id": f"{self.reader.dataset_id}_{image_id}",
+                "id": f"{self._reader.dataset_id}_{image_id}",
                 # Image-level metadata
                 "coco_image_id": image_id,
                 "file_name": image_info["file_name"],

maite_datasets/object_detection/_milco.py

@@ -2,18 +2,20 @@ from __future__ import annotations
 
 __all__ = []
 
+from collections.abc import Sequence
 from pathlib import Path
-from typing import Any, Literal, Sequence
+from typing import Any, Literal
 
-import numpy as np
-from numpy.typing import NDArray
+from maite_datasets._base import (
+    BaseDatasetNumpyMixin,
+    BaseODDataset,
+    DataLocation,
+    NumpyArray,
+    NumpyObjectDetectionTransform,
+)
 
-from maite_datasets._base import BaseODDataset, DataLocation
-from maite_datasets._mixin._numpy import BaseDatasetNumpyMixin
-from maite_datasets._protocols import Transform
 
-
-class MILCO(BaseODDataset[NDArray[np.number[Any]], list[str], str], BaseDatasetNumpyMixin):
+class MILCO(BaseODDataset[NumpyArray, list[str], str], BaseDatasetNumpyMixin):
     """
     A side-scan sonar dataset focused on mine-like object detection.
 
@@ -116,7 +118,7 @@ class MILCO(BaseODDataset[NDArray[np.number[Any]], list[str], str], BaseDatasetN
         self,
         root: str | Path,
         image_set: Literal["train", "operational", "base"] = "train",
-        transforms: Transform[NDArray[np.number[Any]]] | Sequence[Transform[NDArray[np.number[Any]]]] | None = None,
+        transforms: NumpyObjectDetectionTransform | Sequence[NumpyObjectDetectionTransform] | None = None,
         download: bool = False,
         verbose: bool = False,
     ) -> None:

maite_datasets/object_detection/_seadrone.py

@@ -3,21 +3,23 @@ from __future__ import annotations
 __all__ = []
 
 import json
+from collections.abc import Sequence
 from pathlib import Path
-from typing import Any, Literal, Sequence
+from typing import Any, Literal
 
-import numpy as np
-from numpy.typing import NDArray
-
-from maite_datasets._base import BaseODDataset, DataLocation
+from maite_datasets._base import (
+    BaseDatasetNumpyMixin,
+    BaseODDataset,
+    DataLocation,
+    NumpyArray,
+    NumpyObjectDetectionTransform,
+)
 from maite_datasets._fileio import _ensure_exists
-from maite_datasets._mixin._numpy import BaseDatasetNumpyMixin
-from maite_datasets._protocols import Transform
 
 
 class SeaDrone(
     BaseODDataset[
-        NDArray[np.number[Any]],
+        NumpyArray,
         list[tuple[list[int], list[list[float]]]],
         tuple[list[int], list[list[float]]],
     ],
@@ -313,7 +315,7 @@ class SeaDrone(
         self,
         root: str | Path,
         image_set: Literal["train", "val", "test", "base"] = "train",
-        transforms: Transform[NDArray[np.number[Any]]] | Sequence[Transform[NDArray[np.number[Any]]]] | None = None,
+        transforms: NumpyObjectDetectionTransform | Sequence[NumpyObjectDetectionTransform] | None = None,
         download: bool = False,
         verbose: bool = False,
     ) -> None:

maite_datasets/object_detection/_voc.py

@@ -4,24 +4,22 @@ __all__ = []
 
 import os
 import shutil
+from collections.abc import Sequence
 from pathlib import Path
-from typing import Any, Literal, Sequence, TypeVar
+from typing import Any, Literal, TypeVar
 
-import numpy as np
 from defusedxml.ElementTree import parse
-from numpy.typing import NDArray
 
 from maite_datasets._base import (
-    BaseDataset,
+    BaseDatasetNumpyMixin,
+    BaseDownloadedDataset,
     BaseODDataset,
     DataLocation,
+    NumpyArray,
+    NumpyObjectDetectionTransform,
+    ObjectDetectionTarget,
     _ensure_exists,
-    _TArray,
-    _TTarget,
 )
-from maite_datasets._mixin._numpy import BaseDatasetNumpyMixin
-from maite_datasets._protocols import Transform
-from maite_datasets._types import ObjectDetectionTarget
 
 VOCClassStringMap = Literal[
     "aeroplane",
@@ -48,7 +46,7 @@ VOCClassStringMap = Literal[
 TVOCClassMap = TypeVar("TVOCClassMap", VOCClassStringMap, int, list[VOCClassStringMap], list[int])
 
 
-class BaseVOCDataset(BaseDataset[_TArray, _TTarget, list[str], str]):
+class BaseVOCDataset(BaseDownloadedDataset[NumpyArray, ObjectDetectionTarget, list[str], str]):
     _resources = [
         DataLocation(
             url="https://data.brainchip.com/dataset-mirror/voc/VOCtrainval_11-May-2012.tar",
@@ -130,7 +128,7 @@ class BaseVOCDataset(BaseDataset[_TArray, _TTarget, list[str], str]):
         root: str | Path,
         image_set: Literal["train", "val", "test", "base"] = "train",
         year: Literal["2007", "2008", "2009", "2010", "2011", "2012"] = "2012",
-        transforms: Transform[_TArray] | Sequence[Transform[_TArray]] | None = None,
+        transforms: NumpyObjectDetectionTransform | Sequence[NumpyObjectDetectionTransform] | None = None,
         download: bool = False,
         verbose: bool = False,
     ) -> None:
@@ -432,8 +430,8 @@ class BaseVOCDataset(BaseDataset[_TArray, _TTarget, list[str], str]):
 
 
 class VOCDetection(
-    BaseVOCDataset[NDArray[np.number[Any]], ObjectDetectionTarget[NDArray[np.number[Any]]]],
-    BaseODDataset[NDArray[np.number[Any]], list[str], str],
+    BaseVOCDataset,
+    BaseODDataset[NumpyArray, list[str], str],
     BaseDatasetNumpyMixin,
 ):
     """

maite_datasets/{_reader → object_detection}/_yolo.py

@@ -7,14 +7,16 @@ __all__ = []
 from pathlib import Path
 from typing import Any
 
+import maite.protocols.object_detection as od
 import numpy as np
+from maite.protocols import DatasetMetadata, DatumMetadata
 from PIL import Image
 
-from maite_datasets._protocols import DatasetMetadata, DatumMetadata, ObjectDetectionDataset, ObjectDetectionDatum
-from maite_datasets._reader._base import _ObjectDetectionTarget, BaseDatasetReader
+from maite_datasets._base import BaseDataset, ObjectDetectionTarget
+from maite_datasets._reader import BaseDatasetReader
 
 
-class YOLODatasetReader(BaseDatasetReader):
+class YOLODatasetReader(BaseDatasetReader[od.Dataset]):
     """
     YOLO format dataset reader conforming to MAITE protocols.
 
@@ -120,9 +122,9 @@ class YOLODatasetReader(BaseDatasetReader):
         """Mapping from class index to class name."""
         return self._index2label
 
-    def _create_dataset_implementation(self) -> ObjectDetectionDataset:
+    def create_dataset(self) -> od.Dataset:
         """Create YOLO dataset implementation."""
-        return _YOLODataset(self)
+        return YOLODataset(self)
 
     def _validate_format_specific(self) -> tuple[list[str], dict[str, Any]]:
         """Validate YOLO format specific files and structure."""
@@ -200,7 +202,7 @@ class YOLODatasetReader(BaseDatasetReader):
         """Load class names from classes file."""
         with open(self._classes_path) as f:
             class_names = [line.strip() for line in f if line.strip()]
-        self._index2label = {idx: name for idx, name in enumerate(class_names)}
+        self._index2label = dict(enumerate(class_names))
 
     def _find_image_files(self) -> None:
         """Find all valid image files."""
@@ -213,32 +215,35 @@ class YOLODatasetReader(BaseDatasetReader):
             raise ValueError(f"No image files found in {self._images_path}")
 
 
-class _YOLODataset:
+class YOLODataset(BaseDataset):
     """Internal YOLO dataset implementation."""
 
     def __init__(self, reader: YOLODatasetReader) -> None:
-        self.reader = reader
-
-    @property
-    def metadata(self) -> DatasetMetadata:
-        return DatasetMetadata(
-            id=self.reader.dataset_id,
-            index2label=self.reader.index2label,
+        self._reader = reader
+
+        self.root = reader.dataset_path
+        self.images_path = reader._images_path
+        self.annotation_path = reader._labels_path
+        self.size = len(reader._image_files)
+        self.classes = reader.index2label
+        self.metadata = DatasetMetadata(
+            id=self._reader.dataset_id,
+            index2label=self._reader.index2label,
         )
 
     def __len__(self) -> int:
-        return len(self.reader._image_files)
+        return len(self._reader._image_files)
 
-    def __getitem__(self, index: int) -> ObjectDetectionDatum:
-        image_path = self.reader._image_files[index]
+    def __getitem__(self, index: int) -> tuple[od.InputType, od.ObjectDetectionTarget, DatumMetadata]:
+        image_path = self._reader._image_files[index]
 
         # Load image
-        image = np.array(Image.open(image_path).convert("RGB"))
+        image = np.asarray(Image.open(image_path).convert("RGB"), dtype=np.uint8)
         img_height, img_width = image.shape[:2]
         image = np.transpose(image, (2, 0, 1))  # Convert to CHW format
 
         # Load corresponding label file
-        label_path = self.reader._labels_path / f"{image_path.stem}.txt"
+        label_path = self._reader._labels_path / f"{image_path.stem}.txt"
 
         annotation_metadata = []
         if label_path.exists():
@@ -292,12 +297,12 @@ class _YOLODataset:
             labels = np.empty(0, dtype=np.int64)
             scores = np.empty(0, dtype=np.float32)
 
-        target = _ObjectDetectionTarget(boxes, labels, scores)
+        target = ObjectDetectionTarget(boxes, labels, scores)
 
         # Create comprehensive datum metadata
         datum_metadata = DatumMetadata(
             **{
-                "id": f"{self.reader.dataset_id}_{image_path.stem}",
+                "id": f"{self._reader.dataset_id}_{image_path.stem}",
                 # Image-level metadata
                 "file_name": image_path.name,
                 "file_path": str(image_path),

maite_datasets/protocols.py

@@ -0,0 +1,94 @@
+"""
+Common type protocols used for interoperability.
+"""
+
+from collections.abc import Iterable, Iterator, Mapping, Sequence
+from typing import Any, Protocol, overload, runtime_checkable
+
+
+@runtime_checkable
+class Array(Protocol):
+    """
+    Protocol for interoperable array objects.
+
+    Supports common array representations with popular libraries like
+    PyTorch, Tensorflow and JAX, as well as NumPy arrays.
+    """
+
+    @property
+    def shape(self) -> tuple[int, ...]: ...
+    def __array__(self) -> Any: ...
+    def __getitem__(self, key: Any, /) -> Any: ...
+    def __iter__(self) -> Iterator[Any]: ...
+    def __len__(self) -> int: ...
+
+
+@runtime_checkable
+class HFDatasetInfo(Protocol):
+    @property
+    def dataset_name(self) -> str: ...
+
+
+@runtime_checkable
+class HFDataset(Protocol):
+    @property
+    def features(self) -> Mapping[str, Any]: ...
+
+    @property
+    def builder_name(self) -> str | None: ...
+
+    @property
+    def info(self) -> HFDatasetInfo: ...
+
+    @overload
+    def __getitem__(self, key: int | slice | Iterable[int]) -> dict[str, Any]: ...
+    @overload
+    def __getitem__(self, key: str) -> Sequence[int]: ...
+    def __getitem__(self, key: str | int | slice | Iterable[int]) -> dict[str, Any] | Sequence[int]: ...
+
+    def __len__(self) -> int: ...
+
+
+@runtime_checkable
+class HFFeature(Protocol):
+    @property
+    def _type(self) -> str: ...
+
+
+@runtime_checkable
+class HFClassLabel(HFFeature, Protocol):
+    @property
+    def names(self) -> list[str]: ...
+
+    @property
+    def num_classes(self) -> int: ...
+
+
+@runtime_checkable
+class HFImage(HFFeature, Protocol):
+    @property
+    def decode(self) -> bool: ...
+
+
+@runtime_checkable
+class HFArray(HFFeature, Protocol):
+    @property
+    def shape(self) -> tuple[int, ...]: ...
+    @property
+    def dtype(self) -> str: ...
+
+
+@runtime_checkable
+class HFList(HFFeature, Protocol):
+    @property
+    def feature(self) -> Any: ...
+    @property
+    def length(self) -> int: ...
+
+
+@runtime_checkable
+class HFValue(HFFeature, Protocol):
+    @property
+    def pa_type(self) -> Any: ...  # pyarrow type ... not documented
+    @property
+    def dtype(self) -> str: ...

maite_datasets/wrappers/__init__.py

@@ -0,0 +1,8 @@
+import importlib.util
+
+__all__ = []
+
+if importlib.util.find_spec("torch") is not None and importlib.util.find_spec("torchvision") is not None:
+    from ._torch import TorchvisionWrapper
+
+    __all__ += ["TorchvisionWrapper"]

maite_datasets/wrappers/_torch.py

@@ -0,0 +1,109 @@
+from __future__ import annotations
+
+from typing import Any, Callable, Generic, TypeAlias, TypeVar, cast, overload
+
+import torch
+from maite.protocols import DatasetMetadata, DatumMetadata
+from maite.protocols.object_detection import ObjectDetectionTarget as _ObjectDetectionTarget
+from torch import Tensor
+from torchvision.tv_tensors import BoundingBoxes, Image
+
+from maite_datasets._base import BaseDataset, ObjectDetectionTarget
+from maite_datasets.protocols import Array
+
+TArray = TypeVar("TArray", bound=Array)
+TTarget = TypeVar("TTarget")
+
+TorchvisionImageClassificationDatum: TypeAlias = tuple[Image, Tensor, DatumMetadata]
+TorchvisionObjectDetectionDatum: TypeAlias = tuple[Image, ObjectDetectionTarget, DatumMetadata]
+
+
+class TorchvisionWrapper(Generic[TArray, TTarget]):
+    """
+    Lightweight wrapper converting numpy-based datasets to Torchvision tensors.
+
+    Converts images to tv_tensor.Image and targets to the appropriate torchvision format.
+
+    Parameters
+    ----------
+    dataset : Dataset
+        Source dataset with numpy arrays
+    transforms : callable, optional
+        Torchvision v2 transform functions for targets
+    """
+
+    def __init__(
+        self,
+        dataset: BaseDataset[TArray, TTarget],
+        transforms: Callable[[Any], Any] | None = None,
+    ) -> None:
+        self._dataset = dataset
+        self.transforms = transforms
+        self.metadata: DatasetMetadata = {
+            "id": f"TorchvisionWrapper({dataset.metadata['id']})",
+            "index2label": dataset.metadata.get("index2label", {}),
+        }
+
+    def __getattr__(self, name: str) -> Any:
+        """Forward unknown attributes to wrapped dataset."""
+        return getattr(self._dataset, name)
+
+    def __dir__(self) -> list[str]:
+        """Include wrapped dataset attributes in dir() for IDE support."""
+        wrapper_attrs = set(super().__dir__())
+        dataset_attrs = set(dir(self._dataset))
+        return sorted(wrapper_attrs | dataset_attrs)
+
+    def _transform(self, datum: Any) -> Any:
+        return self.transforms(datum) if self.transforms else datum
+
+    @overload
+    def __getitem__(self: TorchvisionWrapper[TArray, TArray], index: int) -> tuple[Image, Tensor, DatumMetadata]: ...
+    @overload
+    def __getitem__(
+        self: TorchvisionWrapper[TArray, TTarget], index: int
+    ) -> tuple[Image, ObjectDetectionTarget, DatumMetadata]: ...
+
+    def __getitem__(self, index: int) -> tuple[Image, Tensor | ObjectDetectionTarget, DatumMetadata]:
+        """Get item with torch tensor conversion."""
+        image, target, metadata = self._dataset[index]
+
+        # Convert image to torch tensor
+        torch_image = Image(torch.tensor(image))
+
+        # Handle different target types
+        if isinstance(target, Array):
+            # Image classification case
+            torch_target = torch.tensor(target, dtype=torch.float32)
+            torch_datum = self._transform((torch_image, torch_target, metadata))
+            return cast(TorchvisionImageClassificationDatum, torch_datum)
+
+        if isinstance(target, _ObjectDetectionTarget):
+            # Object detection case
+            torch_boxes = BoundingBoxes(
+                torch.tensor(target.boxes), format="XYXY", canvas_size=(torch_image.shape[-2], torch_image.shape[-1])
+            )  # type: ignore
+            torch_labels = torch.tensor(target.labels, dtype=torch.int64)
+            torch_scores = torch.tensor(target.scores, dtype=torch.float32)
+            torch_target = ObjectDetectionTarget(torch_boxes, torch_labels, torch_scores)
+            torch_datum = self._transform((torch_image, torch_target, metadata))
+            return cast(TorchvisionObjectDetectionDatum, torch_datum)
+
+        raise TypeError(f"Unsupported target type: {type(target)}")
+
+    def __str__(self) -> str:
+        """String representation showing torch version."""
+        nt = "\n    "
+        base_name = f"{self._dataset.__class__.__name__.replace('Dataset', '')} Dataset"
+        title = f"Torchvision Wrapped {base_name}" if not base_name.startswith("Torchvision") else base_name
+        sep = "-" * len(title)
+        attrs = [
+            f"{' '.join(w.capitalize() for w in k.split('_'))}: {v}"
+            for k, v in self.__dict__.items()
+            if not k.startswith("_")
+        ]
+        wrapped = f"{title}\n{sep}{nt}{nt.join(attrs)}"
+        return f"{wrapped}\n\n{self._dataset}"
+
+    def __len__(self) -> int:
+        return self._dataset.__len__()