depth_edges_annotate_worlduav/src/augmentor/dataset.py

"""Dataset discovery, completion filtering, and PyTorch Dataset for augmentation."""

from __future__ import annotations

import logging
from pathlib import Path
from typing import Any, NamedTuple

import torch
from PIL import Image
from torch.utils.data import Dataset
from torchvision import transforms

from src.augmentor.io_utils import npy_path, vis_path, safetensors_path

logger = logging.getLogger(__name__)

EXTENSIONS = {".png", ".jpg", ".jpeg", ".bmp"}

EXCLUDE_NAMES = {"merge.tif"}
EXCLUDE_DIRS = {"Index", "__MACOSX", "charts", "__pycache__"}

# Incomplete World-UAV Country scenes (16 of 171): no DB crops, no positive.json.
INCOMPLETE_SCENES: set[str] = {
    "CastleHill", "Dalry", "Haymarket", "NewTown", "Stockbridge",
    "CamdenTown", "CoventGarden", "Fitzrovia", "Mayfair", "SoHo",
    "Ancoats", "Castlefield", "Deansgate", "NorthernQuarter", "Piccadilly",
    "JewelleryQuarter",
}


class ImageRecord(NamedTuple):
    """Lightweight descriptor for a single dataset image."""

    abs_path: Path
    rel_path: str
    stem: str
    output_root: Path
    rel_parent: str


def is_query_record(record: ImageRecord) -> bool:
    """Return True if the record belongs to a query (drone) image."""
    return "query" in Path(record.rel_path).parts


def split_by_view(
    records: list[ImageRecord],
) -> tuple[list[ImageRecord], list[ImageRecord]]:
    """Split records into (db_records, query_records)."""
    db: list[ImageRecord] = []
    query: list[ImageRecord] = []
    for r in records:
        if is_query_record(r):
            query.append(r)
        else:
            db.append(r)
    return db, query


# ---------------------------------------------------------------------------
# Discovery
# ---------------------------------------------------------------------------

def discover_images(
    root: Path,
    subset: str | None = None,
    source: str | None = None,
) -> list[ImageRecord]:
    """Recursively find images under *root*, preserving relative paths.

    Args:
        root: Dataset root directory.
        subset: Limit to a World-UAV subset (Country, Terrain, Rot).
        source: Filter by source — 'query' (drone) or 'db' (satellite).

    Returns:
        Sorted list of ImageRecord.
    """
    search_root = root / subset if subset else root
    if not search_root.exists():
        logger.warning("Search root does not exist: %s", search_root)
        return []

    records: list[ImageRecord] = []
    n_skipped_incomplete = 0

    for p in sorted(search_root.rglob("*")):
        if not p.is_file():
            continue
        if p.name in EXCLUDE_NAMES:
            continue
        if p.suffix.lower() not in EXTENSIONS:
            continue
        if any(d in p.parts for d in EXCLUDE_DIRS):
            continue
        if any(scene in p.parts for scene in INCOMPLETE_SCENES):
            n_skipped_incomplete += 1
            continue

        if source is not None:
            rel_parts = p.relative_to(root).parts
            if source == "query" and "DB" in rel_parts:
                continue
            if source == "db" and "query" in rel_parts:
                continue

        rel = p.relative_to(root)
        records.append(ImageRecord(
            abs_path=p,
            rel_path=str(rel),
            stem=p.stem,
            output_root=Path(),
            rel_parent=str(rel.parent),
        ))

    if n_skipped_incomplete > 0:
        logger.info(
            "Skipped %d images from %d incomplete scenes.",
            n_skipped_incomplete, len(INCOMPLETE_SCENES),
        )
    return records


def attach_output_dirs(
    records: list[ImageRecord],
    output_root: Path,
) -> list[ImageRecord]:
    """Set output_root for each record."""
    return [r._replace(output_root=output_root) for r in records]


# Modality name for each stage (used for folder names).
STAGE_MODALITY: dict[str, str] = {
    "depth": "depth",
    "edges": "edge",
    "segmentation": "segm",
    "chmv2": "chm",
}


def filter_completed(
    records: list[ImageRecord],
    stage: str,
) -> tuple[list[ImageRecord], int]:
    """Return (pending_records, n_skipped) for a given stage."""
    if stage == "consolidate":
        return filter_consolidated(records)
    modality = STAGE_MODALITY.get(stage)
    if modality is None:
        return records, 0
    pending: list[ImageRecord] = []
    skipped = 0
    for r in records:
        np_p = npy_path(r.output_root, modality, r.rel_parent, r.stem)
        vis_p = vis_path(r.output_root, modality, r.rel_parent, r.stem)
        if np_p.exists() or vis_p.exists():
            skipped += 1
        else:
            pending.append(r)
    return pending, skipped


def filter_consolidated(
    records: list[ImageRecord],
) -> tuple[list[ImageRecord], int]:
    """Return (pending, skipped) for safetensors consolidation stage."""
    pending: list[ImageRecord] = []
    skipped = 0
    for r in records:
        st = safetensors_path(r.output_root, r.rel_parent, r.stem)
        if st.exists():
            skipped += 1
        else:
            pending.append(r)
    return pending, skipped


# ---------------------------------------------------------------------------
# PyTorch Dataset
# ---------------------------------------------------------------------------

class AugmentDataset(Dataset):
    """Loads RGB images at image_size x image_size for the augmentation pipeline.

    Args:
        records: List of ImageRecord to load.
        image_size: Target spatial resolution (default 256).
    """

    def __init__(self, records: list[ImageRecord], image_size: int = 256) -> None:
        self.records = records
        self.resize = transforms.Resize(
            (image_size, image_size),
            interpolation=transforms.InterpolationMode.BILINEAR,
        )
        self.to_tensor = transforms.ToTensor()

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

    def __getitem__(self, idx: int) -> dict[str, Any]:
        r = self.records[idx]
        img = Image.open(r.abs_path).convert("RGB")
        tensor = self.to_tensor(self.resize(img))
        return {
            "image_raw": tensor,
            "rel_path": r.rel_path,
            "stem": r.stem,
            "output_root": str(r.output_root),
            "rel_parent": r.rel_parent,
        }