Initial commit: add docs, requirements, and prep package

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

uavdense_prepare/__init__.py

@@ -0,0 +1,4 @@
+__all__ = ["__version__"]
+
+__version__ = "0.1.0"
+

uavdense_prepare/__main__.py

@@ -0,0 +1,83 @@
+from __future__ import annotations
+
+import argparse
+import logging
+from pathlib import Path
+
+import coloredlogs
+
+from .gps import parse_dense_gps_txt, write_gps_csv
+from .indexing import build_indices
+from .processing import ProcessingOptions, process_denseuav_to_new_root
+
+
+def _p(s: str) -> Path:
+    return Path(s).expanduser().resolve()
+
+
+def main() -> None:
+    ap = argparse.ArgumentParser(
+        prog="uavdense_prepare",
+        description="Prepare DenseUAV (UAVDense): optionally build processed dataset + indices + stats",
+    )
+    ap.add_argument("--root", required=True, help="Path to raw DenseUAV root directory")
+    ap.add_argument("--out", required=True, help="Output directory for indices/stats (always created)")
+    ap.add_argument(
+        "--dst",
+        default=None,
+        help="If set, creates a new processed dataset root (resized images, satellite tif->png).",
+    )
+    ap.add_argument("--target-size", type=int, default=256, help="Resize to NxN in --dst mode")
+    ap.add_argument(
+        "--limit-ids",
+        type=int,
+        default=None,
+        help="Debug option: in --dst mode process only first N IDs from each split",
+    )
+    ap.add_argument("--exclude-old", action="store_true", help="Exclude *_old.tif from DB/positives")
+    ap.add_argument("--no-strict", action="store_true", help="Do not fail on missing files/ids (best-effort)")
+    ap.add_argument("--log-level", default="INFO", help="Logging level (DEBUG, INFO, WARNING, ERROR)")
+    args = ap.parse_args()
+
+    coloredlogs.install(level=args.log_level.upper(), fmt="%(asctime)s %(levelname)s %(message)s")
+    log = logging.getLogger("uavdense_prepare")
+
+    root = _p(args.root)
+    out = _p(args.out)
+    strict = not args.no_strict
+    dst = _p(args.dst) if args.dst else None
+
+    dataset_root_for_index = root
+    satellite_ext = ".tif"
+
+    if dst is not None:
+        log.info("Building processed dataset at %s", dst)
+        opts = ProcessingOptions(target_size=args.target_size, exclude_old=args.exclude_old, limit_ids=args.limit_ids)
+        process_denseuav_to_new_root(src_root=root, dst_root=dst, opts=opts, strict=strict)
+        dataset_root_for_index = dst
+        satellite_ext = ".png"
+
+    # GPS exports (optional files) from the dataset root we index
+    gps_dir = out / "gps"
+    for name in ("Dense_GPS_ALL.txt", "Dense_GPS_train.txt", "Dense_GPS_test.txt"):
+        src = dataset_root_for_index / name
+        if src.exists():
+            log.info("Parsing GPS file %s", src)
+            recs = parse_dense_gps_txt(src)
+            out_csv = gps_dir / (name.replace(".txt", ".csv").lower())
+            write_gps_csv(recs, out_csv)
+
+    log.info("Building indices from %s into %s", dataset_root_for_index, out)
+    build_indices(
+        dataset_root_for_index,
+        out,
+        exclude_old=args.exclude_old,
+        strict=strict,
+        satellite_ext=satellite_ext,
+    )
+    log.info("Done")
+
+
+if __name__ == "__main__":
+    main()
+

uavdense_prepare/gps.py

@@ -0,0 +1,47 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from pathlib import Path
+
+
+@dataclass(frozen=True)
+class GpsRecord:
+    rel_path: str
+    lon: float
+    lat: float
+    height: float
+
+
+def parse_dense_gps_txt(path: Path) -> list[GpsRecord]:
+    """
+    DenseUAV формат строки (эмпирически):
+        <rel_path> E<lon> N<lat> <height>
+    пример:
+        train/satellite/000001/H80.tif E120.387... N30.324... 94.761
+    """
+    records: list[GpsRecord] = []
+    with path.open("r", encoding="utf-8", errors="replace") as f:
+        for i, line in enumerate(f, start=1):
+            line = line.strip()
+            if not line:
+                continue
+            parts = line.split()
+            if len(parts) != 4:
+                raise ValueError(f"Bad GPS line {path}:{i}: expected 4 fields, got {len(parts)}: {line!r}")
+            rel_path, e_lon, n_lat, height_s = parts
+            if not e_lon.startswith("E") or not n_lat.startswith("N"):
+                raise ValueError(f"Bad GPS line {path}:{i}: expected E.. N.. fields: {line!r}")
+            lon = float(e_lon[1:])
+            lat = float(n_lat[1:])
+            height = float(height_s)
+            records.append(GpsRecord(rel_path=rel_path, lon=lon, lat=lat, height=height))
+    return records
+
+
+def write_gps_csv(records: list[GpsRecord], out_csv: Path) -> None:
+    out_csv.parent.mkdir(parents=True, exist_ok=True)
+    with out_csv.open("w", encoding="utf-8") as f:
+        f.write("rel_path,lon,lat,height\n")
+        for r in records:
+            f.write(f"{r.rel_path},{r.lon:.12f},{r.lat:.12f},{r.height:.6f}\n")
+

uavdense_prepare/indexing.py

@@ -0,0 +1,196 @@
+from __future__ import annotations
+
+import json
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Iterable
+
+
+DRONE_NAMES = ("H80.JPG", "H90.JPG", "H100.JPG")
+SAT_NAMES = ("H80.tif", "H90.tif", "H100.tif", "H80_old.tif", "H90_old.tif", "H100_old.tif")
+
+
+@dataclass(frozen=True)
+class DenseUavLayout:
+    root: Path
+
+    @property
+    def train_drone_dir(self) -> Path:
+        return self.root / "train" / "drone"
+
+    @property
+    def train_sat_dir(self) -> Path:
+        return self.root / "train" / "satellite"
+
+    @property
+    def test_query_drone_dir(self) -> Path:
+        return self.root / "test" / "query_drone"
+
+    @property
+    def test_gallery_sat_dir(self) -> Path:
+        return self.root / "test" / "gallery_satellite"
+
+
+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 _collect_files_for_id(id_dir: Path, expected_names: Iterable[str]) -> dict[str, str]:
+    """
+    Returns mapping name -> relative path (posix) for files that exist.
+    """
+    out: dict[str, str] = {}
+    for name in expected_names:
+        p = id_dir / name
+        if p.exists():
+            out[name] = p.as_posix()
+    return out
+
+
+def _write_lines(path: Path, lines: Iterable[str]) -> None:
+    path.parent.mkdir(parents=True, exist_ok=True)
+    with path.open("w", encoding="utf-8") as f:
+        for line in lines:
+            f.write(line.rstrip("\n") + "\n")
+
+
+def build_indices(
+    root: Path,
+    out_dir: Path,
+    *,
+    exclude_old: bool = False,
+    strict: bool = True,
+    satellite_ext: str = ".tif",
+) -> dict:
+    """
+    Создаёт train/test индексы в стиле:
+        query_path label pos1 pos2 ...
+    где label = целочисленный id класса (по порядку).
+
+    Positive list:
+      - для train query: все доступные satellite варианты того же ID (3 или 6 файлов)
+      - для test query: satellite варианты из test/gallery_satellite/<ID> (если есть)
+    """
+    layout = DenseUavLayout(root=root)
+
+    # IDs
+    train_ids = _iter_ids(layout.train_drone_dir)
+    train_sat_ids = _iter_ids(layout.train_sat_dir)
+    if strict and train_ids != train_sat_ids:
+        raise ValueError("train/drone IDs differ from train/satellite IDs")
+
+    gallery_ids = _iter_ids(layout.test_gallery_sat_dir)
+    query_ids = _iter_ids(layout.test_query_drone_dir)
+
+    # label mapping: use gallery id universe for stable evaluation labels
+    # (train ids are subset of gallery ids in typical setting)
+    all_ids = sorted(set(gallery_ids) | set(train_ids) | set(query_ids))
+    id_to_label = {id_: i for i, id_ in enumerate(all_ids)}
+
+    # Collect DB lists
+    def sat_paths_for(id_: str, sat_root: Path) -> list[str]:
+        id_dir = sat_root / id_
+        if not id_dir.exists():
+            return []
+        names = list(SAT_NAMES)
+        if exclude_old:
+            names = [n for n in names if not n.endswith("_old.tif")]
+        # allow processed datasets where tif were converted to png
+        names_fs = [n.replace(".tif", satellite_ext) for n in names]
+        got = _collect_files_for_id(id_dir, names_fs)
+        # stable order: H80,H90,H100,(old...)
+        ordered = [got[n] for n in names_fs if n in got]
+        # make relative to dataset root
+        rel = [str(Path(p).relative_to(root).as_posix()) for p in ordered]
+        return rel
+
+    def drone_paths_for(id_: str, drone_root: Path) -> list[str]:
+        id_dir = drone_root / id_
+        if not id_dir.exists():
+            return []
+        got = _collect_files_for_id(id_dir, DRONE_NAMES)
+        ordered = [got[n] for n in DRONE_NAMES if n in got]
+        rel = [str(Path(p).relative_to(root).as_posix()) for p in ordered]
+        return rel
+
+    # train_db: all train satellite images (optionally exclude old)
+    train_db: list[str] = []
+    for id_ in train_ids:
+        train_db.extend(sat_paths_for(id_, layout.train_sat_dir))
+
+    # test_db: all test gallery satellite images
+    test_db: list[str] = []
+    for id_ in gallery_ids:
+        test_db.extend(sat_paths_for(id_, layout.test_gallery_sat_dir))
+
+    # Queries:
+    train_query_lines: list[str] = []
+    train_missing: list[str] = []
+    for id_ in train_ids:
+        q_paths = drone_paths_for(id_, layout.train_drone_dir)
+        pos = sat_paths_for(id_, layout.train_sat_dir)
+        if strict:
+            if len(q_paths) != 3:
+                train_missing.append(f"{id_}: drone files {len(q_paths)}/3")
+            if (exclude_old and len(pos) != 3) or ((not exclude_old) and len(pos) != 6):
+                train_missing.append(f"{id_}: satellite files {len(pos)}/expected")
+        for q in q_paths:
+            label = id_to_label[id_]
+            if not pos:
+                if strict:
+                    raise ValueError(f"No positives for train query id={id_}")
+                continue
+            train_query_lines.append(" ".join([q, str(label), *pos]))
+
+    test_query_lines: list[str] = []
+    test_missing: list[str] = []
+    for id_ in query_ids:
+        q_paths = drone_paths_for(id_, layout.test_query_drone_dir)
+        pos = sat_paths_for(id_, layout.test_gallery_sat_dir)
+        if strict:
+            if len(q_paths) != 3:
+                test_missing.append(f"{id_}: query files {len(q_paths)}/3")
+            if not pos:
+                test_missing.append(f"{id_}: no gallery satellite folder/files")
+        for q in q_paths:
+            label = id_to_label[id_]
+            if not pos:
+                if strict:
+                    raise ValueError(f"No positives for test query id={id_}")
+                continue
+            test_query_lines.append(" ".join([q, str(label), *pos]))
+
+    # Write outputs
+    index_dir = out_dir / "index"
+    _write_lines(index_dir / "train_db.txt", train_db)
+    _write_lines(index_dir / "test_db.txt", test_db)
+    _write_lines(index_dir / "train_query.txt", train_query_lines)
+    _write_lines(index_dir / "test_query.txt", test_query_lines)
+
+    stats = {
+        "root": str(root),
+        "out_dir": str(out_dir),
+        "exclude_old": exclude_old,
+        "strict": strict,
+        "counts": {
+            "n_train_ids": len(train_ids),
+            "n_gallery_ids": len(gallery_ids),
+            "n_query_ids": len(query_ids),
+            "n_all_ids_universe": len(all_ids),
+            "n_train_db_images": len(train_db),
+            "n_test_db_images": len(test_db),
+            "n_train_query_images": len(train_query_lines),
+            "n_test_query_images": len(test_query_lines),
+        },
+        "integrity": {
+            "train_issues": train_missing[:200],
+            "test_issues": test_missing[:200],
+        },
+    }
+    (out_dir / "stats").mkdir(parents=True, exist_ok=True)
+    (out_dir / "stats" / "stats.json").write_text(json.dumps(stats, ensure_ascii=False, indent=2), encoding="utf-8")
+    return stats
+

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)}
+