How to wrap a DataFrame-backed object detection dataset¶

import tempfile
from dataclasses import dataclass
from pathlib import Path

import numpy as np
import pandas as pd
from PIL import Image

from dataeval import Metadata
from dataeval.protocols import AnnotatedDataset, DatasetMetadata, DatumMetadata, ObjectDetectionTarget

data_dir = Path(tempfile.mkdtemp())
rng = np.random.default_rng(0)

index2label = {0: "person", 1: "car", 2: "bicycle"}
weather_options = ["clear", "rainy", "foggy"]
occlusion_levels = ["none", "partial", "heavy"]

rows = []
# Creating 30 128x128 images
for i in range(30):
    # Stand-in for a real image file - replace with your own images on disk
    pixels = rng.integers(0, 256, size=(128, 128, 3), dtype=np.uint8)
    filepath = data_dir / f"img_{i:03d}.png"
    Image.fromarray(pixels).save(filepath)

    # A variable number of boxes per image is exactly why OD catalogs use one row
    # per box rather than one row per image.
    n_boxes = int(rng.integers(1, 5))
    for _ in range(n_boxes):
        x, y = rng.integers(0, 80, size=2)
        w, h = rng.integers(16, 40, size=2)
        rows.append({
            "filepath": str(filepath),
            "x": int(x),
            "y": int(y),
            "w": int(w),
            "h": int(h),
            "label": int(rng.integers(0, 3)),
            "weather": weather_options[i % 3],  # image-level factor (same for every box)
            "altitude_m": float(50 + i),  # image-level factor (same for every box)
            "occlusion": occlusion_levels[int(rng.integers(0, 3))],  # box-level factor (per box)
        })

catalog = pd.DataFrame(rows)
catalog.head()

@dataclass
class BoxTarget:
    """A minimal object detection target.

    Implements the ObjectDetectionTarget protocol by exposing the three
    attributes DataEval reads.
    """

    boxes: np.ndarray  # (N, 4) bounding boxes in x0, y0, x1, y1 pixel format
    labels: np.ndarray  # (N,) integer class index per box
    scores: np.ndarray  # (N,) confidence per box; 1.0 for ground truth

class DataFrameODDataset:
    """An object detection dataset backed by a long-format pandas DataFrame.

    Each row describes one bounding box. Rows that share the image-path column
    belong to the same image and are grouped into a single target. Metadata is
    surfaced at two levels: ``metadata_cols`` are image-level (one value per image)
    and ``box_metadata_cols`` are box-level (one value per box).
    """

    def __init__(
        self,
        dataframe: pd.DataFrame,
        index2label: dict[int, str],
        image_col: str = "filepath",
        box_cols: tuple[str, str, str, str] = ("x", "y", "w", "h"),
        label_col: str = "label",
        metadata_cols: list[str] | None = None,
        box_metadata_cols: list[str] | None = None,
        dataset_id: str = "dataframe-od-catalog",
    ) -> None:
        self._df = dataframe.reset_index(drop=True)
        self.index2label = index2label
        self._box_cols = box_cols
        self._label_col = label_col
        self._metadata_cols = metadata_cols or []
        self._box_metadata_cols = box_metadata_cols or []
        # Group rows by image, preserving first-seen order, so __getitem__(i)
        # always refers to the same image. Each group is a (key, sub-frame) pair.
        self._groups = [rows for _, rows in self._df.groupby(image_col, sort=False)]
        self._image_col = image_col
        # DatasetMetadata advertises the class-name mapping to DataEval
        self.metadata: DatasetMetadata = DatasetMetadata(id=dataset_id, index2label=index2label)

    def __len__(self) -> int:
        return len(self._groups)

    def __getitem__(self, index: int) -> tuple[np.ndarray, BoxTarget, DatumMetadata]:
        rows = self._groups[index]
        first = rows.iloc[0]

        # Decode the image and convert to channels-first (C, H, W)
        image = np.asarray(Image.open(first[self._image_col]).convert("RGB"), dtype=np.uint8).transpose(2, 0, 1)

        # Convert COCO-style (x, y, w, h) boxes to the (x0, y0, x1, y1) format
        # ObjectDetectionTarget expects.
        x, y, w, h = (rows[col].to_numpy(dtype=np.float32) for col in self._box_cols)
        boxes = np.stack([x, y, x + w, y + h], axis=1)

        labels = rows[self._label_col].to_numpy(dtype=np.intp)
        # Ground-truth boxes are certain, so every score is 1.0
        scores = np.ones(len(labels), dtype=np.float32)
        target = BoxTarget(boxes=boxes, labels=labels, scores=scores)

        # Image-level metadata is the same for every box, so take it from the first
        # row as a scalar. Box-level metadata is one value per box, passed as a list.
        # DataEval broadcasts the scalars across the image's detections and expands
        # the lists to one value per detection. Per-item metadata must include an "id".
        datum_metadata = DatumMetadata(
            id=index,
            **{col: first[col] for col in self._metadata_cols},
            **{col: rows[col].tolist() for col in self._box_metadata_cols},
        )

        return image, target, datum_metadata

dataset = DataFrameODDataset(
    catalog,
    index2label,
    metadata_cols=["weather", "altitude_m"],  # image-level
    box_metadata_cols=["occlusion"],  # box-level
)
print(f"Catalog rows (boxes): {len(catalog)}")
print(f"Dataset length (images): {len(dataset)}")

image, target, datum_metadata = dataset[0]
print(f"image shape:    {image.shape} ({image.dtype})")
print(f"boxes shape:    {target.boxes.shape} (x0, y0, x1, y1)")
print(f"labels:         {target.labels}")
print(f"scores:         {target.scores}")
print(f"datum metadata: {datum_metadata}")

print(f"Is an AnnotatedDataset:          {isinstance(dataset, AnnotatedDataset)}")
print(f"Target is ObjectDetectionTarget: {isinstance(target, ObjectDetectionTarget)}")

metadata = Metadata(dataset)

# "id" is a per-item identifier, not a meaningful factor for bias analysis
metadata.exclude = ["id"]
print(f"Factor names: {metadata.factor_names}")

print(metadata.target_data.select(["item_index", "target_index", "class_label", "weather", "occlusion"]).head(8))