Initial commit: add docs, requirements, and prep package

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

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
+