src/dataset.py

from __future__ import annotations

import logging
import re
from pathlib import Path

import pandas as pd

from src.config import ALL_CATEGORIES_LABEL, PROJECT_ROOT


logger = logging.getLogger(__name__)

IMAGE_EXTENSIONS = {".jpg", ".jpeg", ".png", ".bmp", ".webp"}
METADATA_COLUMNS = ["id", "path", "filename", "category"]


class DatasetError(RuntimeError):
    """Raised when local dataset metadata is missing or invalid."""


def _slug(value: str) -> str:
    slug = re.sub(r"[^A-Za-z0-9]+", "-", value.strip()).strip("-").lower()
    return slug or "image"


def _relative_to_project(path: Path) -> str:
    try:
        return path.resolve().relative_to(PROJECT_ROOT).as_posix()
    except ValueError:
        return path.resolve().as_posix()


def scan_images(image_dir: Path) -> list[dict[str, str]]:
    image_dir = Path(image_dir)
    if not image_dir.exists():
        logger.warning("Image directory does not exist: %s", image_dir)
        return []

    rows: list[dict[str, str]] = []
    used_ids: dict[str, int] = {}
    files = sorted(
        path
        for path in image_dir.rglob("*")
        if path.is_file() and path.suffix.lower() in IMAGE_EXTENSIONS
    )

    for image_path in files:
        if image_path.parent == image_dir:
            category = "Uncategorized"
        else:
            category = image_path.parent.name
        base_id = f"{_slug(category)}-{_slug(image_path.stem)}"
        occurrence = used_ids.get(base_id, 0)
        used_ids[base_id] = occurrence + 1
        image_id = base_id if occurrence == 0 else f"{base_id}-{occurrence + 1}"
        rows.append(
            {
                "id": image_id,
                "path": _relative_to_project(image_path),
                "filename": image_path.name,
                "category": category,
            }
        )

    return rows


def build_metadata(image_dir: Path, output_csv: Path) -> pd.DataFrame:
    rows = scan_images(image_dir)
    output_csv = Path(output_csv)
    output_csv.parent.mkdir(parents=True, exist_ok=True)
    metadata = pd.DataFrame(rows, columns=METADATA_COLUMNS)
    metadata.to_csv(output_csv, index=False)
    logger.info("Wrote metadata for %d image(s) to %s", len(metadata), output_csv)
    return metadata


def load_metadata(metadata_csv: Path) -> pd.DataFrame:
    metadata_csv = Path(metadata_csv)
    if not metadata_csv.exists():
        raise DatasetError(
            f"metadata.csv was not found at {metadata_csv}. "
            "Run `python scripts/build_metadata.py` after adding images."
        )
    metadata = pd.read_csv(metadata_csv, dtype=str).fillna("")
    missing = [column for column in METADATA_COLUMNS if column not in metadata.columns]
    if missing:
        raise DatasetError(f"metadata.csv is missing required column(s): {', '.join(missing)}")
    return metadata[METADATA_COLUMNS]


def get_categories(metadata_csv: Path) -> list[str]:
    try:
        metadata = load_metadata(metadata_csv)
    except DatasetError:
        return [ALL_CATEGORIES_LABEL]
    categories = sorted(category for category in metadata["category"].dropna().unique() if category)
    return [ALL_CATEGORIES_LABEL, *categories]