Initial commit: add docs, requirements, and prep package

Co-authored-by: Cursor <cursoragent@cursor.com>

uavdense_prepare/processing.py

@@ -0,0 +1,173 @@
+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)}
+