197 lines
6.8 KiB
Python
197 lines
6.8 KiB
Python
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
|
|
|