Dense-UAV-prepare/uavdense_prepare/processing.py

from __future__ import annotations

import logging
from dataclasses import dataclass
from pathlib import Path

from PIL import Image
from tqdm import tqdm

from .indexing import DenseUavLayout, DRONE_NAMES, SAT_NAMES


log = logging.getLogger("uavdense_prepare")


@dataclass(frozen=True)
class ProcessingOptions:
    target_size: int = 256
    exclude_old: bool = False
    jpeg_quality: int = 95
    limit_ids: int | None = None


def _ensure_dir(p: Path) -> None:
    p.mkdir(parents=True, exist_ok=True)


def _resize_to_square(img: Image.Image, target_size: int) -> Image.Image:
    if img.mode not in ("RGB", "L"):
        img = img.convert("RGB")
    return img.resize((target_size, target_size), Image.LANCZOS)


def _copy_resize_jpg(src: Path, dst: Path, target_size: int, quality: int) -> None:
    if dst.exists():
        return
    with Image.open(src) as im:
        out = _resize_to_square(im, target_size)
        _ensure_dir(dst.parent)
        out.save(dst, "JPEG", quality=quality)


def _copy_resize_tif_to_png(src: Path, dst: Path, target_size: int) -> None:
    if dst.exists():
        return
    with Image.open(src) as im:
        out = _resize_to_square(im, target_size)
        _ensure_dir(dst.parent)
        out.save(dst, "PNG")


def _iter_ids(dir_path: Path) -> list[str]:
    if not dir_path.exists():
        return []
    ids = [p.name for p in dir_path.iterdir() if p.is_dir()]
    return sorted(ids)


def process_denseuav_to_new_root(
    *,
    src_root: Path,
    dst_root: Path,
    opts: ProcessingOptions,
    strict: bool = True,
) -> dict:
    """
    Создаёт новую структуру датасета в dst_root:
      - drone: JPEG resized до target_size
      - satellite: PNG resized до target_size (из .tif)

    Важно: DenseUAV уже “разрезан” по ID, поэтому satellite tiling не требуется.
    """
    src = DenseUavLayout(root=src_root)
    dst = DenseUavLayout(root=dst_root)

    # Validate IDs
    train_ids = _iter_ids(src.train_drone_dir)
    train_sat_ids = _iter_ids(src.train_sat_dir)
    if strict and train_ids != train_sat_ids:
        raise ValueError("train/drone IDs differ from train/satellite IDs")

    test_query_ids = _iter_ids(src.test_query_drone_dir)
    test_gallery_ids = _iter_ids(src.test_gallery_sat_dir)

    if opts.limit_ids is not None:
        train_ids = train_ids[: opts.limit_ids]
        # For test, keep query IDs (first N) and ensure gallery contains them,
        # otherwise indexing will not find positives.
        test_query_ids = test_query_ids[: opts.limit_ids]
        gallery_set = set(test_gallery_ids)
        test_gallery_ids = [i for i in test_query_ids if i in gallery_set]

    issues: list[str] = []
    copied = {"train_drone": 0, "train_sat": 0, "test_query_drone": 0, "test_gallery_sat": 0}

    # Process train drone
    log.info("Processing train drone (%d IDs)", len(train_ids))
    for id_ in tqdm(train_ids, desc="train/drone IDs"):
        for name in DRONE_NAMES:
            s = src.train_drone_dir / id_ / name
            d = dst.train_drone_dir / id_ / name
            if not s.exists():
                if strict:
                    raise FileNotFoundError(s)
                issues.append(f"missing: {s}")
                continue
            _copy_resize_jpg(s, d, opts.target_size, opts.jpeg_quality)
            copied["train_drone"] += 1

    # Process train satellite (.tif -> .png)
    sat_names = list(SAT_NAMES)
    if opts.exclude_old:
        sat_names = [n for n in sat_names if not n.endswith("_old.tif")]

    log.info("Processing train satellite (%d IDs)", len(train_ids))
    for id_ in tqdm(train_ids, desc="train/satellite IDs"):
        for name in sat_names:
            s = src.train_sat_dir / id_ / name
            d = dst.train_sat_dir / id_ / name.replace(".tif", ".png")
            if not s.exists():
                if strict:
                    raise FileNotFoundError(s)
                issues.append(f"missing: {s}")
                continue
            _copy_resize_tif_to_png(s, d, opts.target_size)
            copied["train_sat"] += 1

    # Process test query drone
    log.info("Processing test query drone (%d IDs)", len(test_query_ids))
    for id_ in tqdm(test_query_ids, desc="test/query_drone IDs"):
        for name in DRONE_NAMES:
            s = src.test_query_drone_dir / id_ / name
            d = dst.test_query_drone_dir / id_ / name
            if not s.exists():
                if strict:
                    raise FileNotFoundError(s)
                issues.append(f"missing: {s}")
                continue
            _copy_resize_jpg(s, d, opts.target_size, opts.jpeg_quality)
            copied["test_query_drone"] += 1

    # Process test gallery satellite
    log.info("Processing test gallery satellite (%d IDs)", len(test_gallery_ids))
    for id_ in tqdm(test_gallery_ids, desc="test/gallery_satellite IDs"):
        for name in sat_names:
            s = src.test_gallery_sat_dir / id_ / name
            d = dst.test_gallery_sat_dir / id_ / name.replace(".tif", ".png")
            if not s.exists():
                if strict:
                    raise FileNotFoundError(s)
                issues.append(f"missing: {s}")
                continue
            _copy_resize_tif_to_png(s, d, opts.target_size)
            copied["test_gallery_sat"] += 1

    # Also copy GPS txt as-is (if present)
    for name in ("Dense_GPS_ALL.txt", "Dense_GPS_train.txt", "Dense_GPS_test.txt"):
        s = src_root / name
        if s.exists():
            d = dst_root / name
            _ensure_dir(d.parent)
            if not d.exists():
                d.write_text(s.read_text(encoding="utf-8", errors="replace"), encoding="utf-8")

    # Minimal marker
    (dst_root / ".prepared_by_uavdense_prepare.txt").write_text(
        f"target_size={opts.target_size}\nexclude_old={opts.exclude_old}\n",
        encoding="utf-8",
    )

    log.info("Processing completed. Copied: %s", copied)
    return {"copied": copied, "issues": issues[:200], "dst_root": str(dst_root)}