CSV Data Loading API¶

class CSVImageDataset(Dataset):
    """Dataset for single-label classification from CSV.

    CSV format::

        image_path,label
        img001.jpg,cat
        img002.jpg,dog

    Parameters:
        csv_path: Path to CSV file.
        image_dir: Root directory for resolving image paths.
        image_column: Name of the column containing image paths.
            If ``None``, uses the first column.
        label_column: Name of the column containing class labels.
            If ``None``, uses the second column.
        transform: Transform to apply to images. Supports both
            torchvision transforms (PIL input) and albumentations
            transforms (numpy input with ``image`` key).
        use_albumentations: If ``True``, load images with OpenCV and
            pass as numpy arrays to an albumentations transform.
            Default is ``False`` (PIL + torchvision).
    """

    def __init__(
        self,
        csv_path: str | Path,
        image_dir: str | Path = ".",
        image_column: str | None = None,
        label_column: str | None = None,
        transform: Callable | None = None,
        use_albumentations: bool = False,
    ):
        self.csv_path = Path(csv_path)
        self.image_dir = Path(image_dir)
        self.transform = transform
        self.use_albumentations = use_albumentations

        # Parse CSV
        with open(self.csv_path, newline="") as f:
            reader = csv.DictReader(f)
            fieldnames = list(reader.fieldnames or [])
            if not fieldnames:
                raise ValueError(f"CSV file has no columns: {self.csv_path}")

            # Determine image column
            self._image_column = image_column or fieldnames[0]
            if self._image_column not in fieldnames:
                raise ValueError(
                    f"Image column '{self._image_column}' not found in CSV. "
                    f"Available columns: {fieldnames}"
                )

            # Determine label column
            if label_column is not None:
                if label_column not in fieldnames:
                    raise ValueError(
                        f"Label column '{label_column}' not found in CSV. "
                        f"Available columns: {fieldnames}"
                    )
                self._label_column = label_column
            else:
                non_image = [c for c in fieldnames if c != self._image_column]
                if not non_image:
                    raise ValueError(
                        "No label column found. Provide label_column explicitly "
                        "or ensure the CSV has a column beyond the image column."
                    )
                self._label_column = non_image[0]

            # Read rows
            self._image_paths: list[str] = []
            self._labels_raw: list[str] = []
            for row in reader:
                self._image_paths.append(row[self._image_column])
                self._labels_raw.append(row[self._label_column])

        if len(self._image_paths) == 0:
            raise ValueError(f"CSV file has no data rows: {self.csv_path}")

        # Build class mapping from unique label values
        unique_labels = sorted(set(self._labels_raw))
        self.classes: list[str] = unique_labels
        self.class_to_idx: dict[str, int] = {
            cls: idx for idx, cls in enumerate(self.classes)
        }

        # Convert string labels to indices
        self.samples: list[tuple[str, int]] = [
            (img, self.class_to_idx[lbl])
            for img, lbl in zip(self._image_paths, self._labels_raw)
        ]

    @property
    def num_classes(self) -> int:
        return len(self.classes)

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

    def __getitem__(self, index: int) -> tuple[Any, int]:
        rel_path, target = self.samples[index]
        img_path = self.image_dir / rel_path

        if self.use_albumentations:
            import cv2

            image = cv2.imread(str(img_path), cv2.IMREAD_COLOR)
            if image is None:
                raise RuntimeError(f"Failed to load image: {img_path}")
            image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)

            if self.transform is not None:
                transformed = self.transform(image=image)
                image = transformed["image"]
        else:
            from PIL import Image

            image = Image.open(img_path).convert("RGB")
            if self.transform is not None:
                image = self.transform(image)

        return image, target

class MultiLabelImageDataset(Dataset):
    """Dataset for multi-label classification from CSV.

    CSV format::

        image_path,label_0,label_1,...,label_N
        img1.jpg,1,0,1,...,0
        img2.jpg,0,1,0,...,1

    Parameters:
        csv_path: Path to CSV file.
        image_dir: Root directory for resolving image paths.
        label_columns: List of column names to use as labels.
            If ``None``, uses all columns except the image column.
        image_column: Name of the column containing image paths.
            If ``None``, uses the first column.
        transform: Transform to apply to images. Supports both
            torchvision transforms (PIL input) and albumentations
            transforms (numpy input with ``image`` key).
        use_albumentations: If ``True``, load images with OpenCV and
            pass as numpy arrays to an albumentations transform.
            Default is ``False`` (PIL + torchvision).
    """

    def __init__(
        self,
        csv_path: str | Path,
        image_dir: str | Path = ".",
        label_columns: list[str] | None = None,
        image_column: str | None = None,
        transform: Callable | None = None,
        use_albumentations: bool = False,
    ):
        self.csv_path = Path(csv_path)
        self.image_dir = Path(image_dir)
        self.transform = transform
        self.use_albumentations = use_albumentations

        # Parse CSV
        with open(self.csv_path, newline="") as f:
            reader = csv.DictReader(f)
            fieldnames = list(reader.fieldnames or [])
            if not fieldnames:
                raise ValueError(f"CSV file has no columns: {self.csv_path}")

            # Determine image column
            self._image_column = image_column or fieldnames[0]
            if self._image_column not in fieldnames:
                raise ValueError(
                    f"Image column '{self._image_column}' not found in CSV. "
                    f"Available columns: {fieldnames}"
                )

            # Determine label columns
            if label_columns is not None:
                for col in label_columns:
                    if col not in fieldnames:
                        raise ValueError(
                            f"Label column '{col}' not found in CSV. "
                            f"Available columns: {fieldnames}"
                        )
                self._label_columns = label_columns
            else:
                self._label_columns = [c for c in fieldnames if c != self._image_column]

            if not self._label_columns:
                raise ValueError(
                    "No label columns found. Provide label_columns explicitly "
                    "or ensure the CSV has columns beyond the image column."
                )

            # Read rows
            self._image_paths: list[str] = []
            self._labels: list[list[float]] = []
            for row in reader:
                self._image_paths.append(row[self._image_column])
                self._labels.append([float(row[col]) for col in self._label_columns])

        if len(self._image_paths) == 0:
            raise ValueError(f"CSV file has no data rows: {self.csv_path}")

    @property
    def num_labels(self) -> int:
        """Number of label columns."""
        return len(self._label_columns)

    @property
    def label_names(self) -> list[str]:
        """Names of the label columns."""
        return list(self._label_columns)

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

    def __getitem__(self, index: int) -> tuple[Any, torch.Tensor]:
        img_path = self.image_dir / self._image_paths[index]
        label_tensor = torch.tensor(self._labels[index], dtype=torch.float32)

        if self.use_albumentations:
            import cv2

            image = cv2.imread(str(img_path), cv2.IMREAD_COLOR)
            if image is None:
                raise RuntimeError(f"Failed to load image: {img_path}")
            image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)

            if self.transform is not None:
                transformed = self.transform(image=image)
                image = transformed["image"]
        else:
            from PIL import Image

            image = Image.open(img_path).convert("RGB")
            if self.transform is not None:
                image = self.transform(image)

        return image, label_tensor

class CSVDetectionDataset(Dataset):
    """Dataset for object detection from CSV.

    CSV format (one row per bounding box)::

        image_path,x_min,y_min,x_max,y_max,label
        img001.jpg,10,20,100,200,cat
        img001.jpg,50,60,150,250,dog
        img002.jpg,30,40,120,220,cat

    Multiple rows per image are grouped automatically.

    Parameters:
        csv_path: Path to CSV file.
        image_dir: Directory containing image files.
        image_column: Column name for image paths. Default ``"image_path"``.
        bbox_columns: Column names for bounding box coordinates.
            Default ``["x_min", "y_min", "x_max", "y_max"]``.
        label_column: Column name for class labels. Default ``"label"``.
        transform: Albumentations transform with bbox_params.
        min_bbox_area: Minimum bbox area to include. Default 0.

    Attributes:
        class_names: List of class names.
        num_classes: Number of classes.
    """

    def __init__(
        self,
        csv_path: str | Path,
        image_dir: str | Path = ".",
        image_column: str = "image_path",
        bbox_columns: list[str] | None = None,
        label_column: str = "label",
        transform: Callable | None = None,
        min_bbox_area: float = 0.0,
    ):
        self.csv_path = Path(csv_path)
        self.image_dir = Path(image_dir)
        self.transform = transform
        self.min_bbox_area = min_bbox_area

        if bbox_columns is None:
            bbox_columns = ["x_min", "y_min", "x_max", "y_max"]

        # Parse CSV and group by image
        image_anns: dict[str, list[dict]] = defaultdict(list)
        with open(self.csv_path, newline="") as f:
            reader = csv.DictReader(f)
            fieldnames = list(reader.fieldnames or [])

            # Validate columns exist
            for col in [image_column, label_column] + bbox_columns:
                if col not in fieldnames:
                    raise ValueError(
                        f"Column '{col}' not found in CSV. "
                        f"Available columns: {fieldnames}"
                    )

            for row in reader:
                img_path = row[image_column]
                x1 = float(row[bbox_columns[0]])
                y1 = float(row[bbox_columns[1]])
                x2 = float(row[bbox_columns[2]])
                y2 = float(row[bbox_columns[3]])

                # Filter by area
                area = (x2 - x1) * (y2 - y1)
                if area < min_bbox_area:
                    continue

                image_anns[img_path].append(
                    {"bbox": [x1, y1, x2, y2], "label": row[label_column]}
                )

        if not image_anns:
            raise ValueError(
                f"No images with valid annotations found in {self.csv_path}."
            )

        # Build class mapping
        all_labels = sorted(
            {ann["label"] for anns in image_anns.values() for ann in anns}
        )
        self.class_names: list[str] = all_labels
        self._class_to_idx: dict[str, int] = {
            name: idx for idx, name in enumerate(all_labels)
        }
        self.num_classes: int = len(all_labels)

        # Store as ordered list for indexing
        self._image_paths: list[str] = sorted(image_anns.keys())
        self._image_anns = image_anns

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

    def __getitem__(self, index: int) -> dict[str, Any]:
        import cv2

        img_rel = self._image_paths[index]
        img_path = self.image_dir / img_rel
        anns = self._image_anns[img_rel]

        image = cv2.imread(str(img_path))
        if image is None:
            raise RuntimeError(f"Failed to load image: {img_path}")
        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)

        orig_h, orig_w = image.shape[:2]

        # Boxes are already in xyxy format
        bboxes = [ann["bbox"] for ann in anns]
        labels = [self._class_to_idx[ann["label"]] for ann in anns]

        if self.transform is not None:
            transformed = self.transform(image=image, bboxes=bboxes, labels=labels)
            image = transformed["image"]
            bboxes = transformed["bboxes"]
            labels = transformed["labels"]

        if len(bboxes) > 0:
            boxes = torch.tensor(bboxes, dtype=torch.float32)
            labels = torch.tensor(labels, dtype=torch.int64)
        else:
            boxes = torch.zeros((0, 4), dtype=torch.float32)
            labels = torch.zeros((0,), dtype=torch.int64)

        return {
            "image": image,
            "boxes": boxes,
            "labels": labels,
            "image_id": index,
            "orig_size": torch.tensor([orig_h, orig_w]),
        }