Pikaliov
4ff36ce188
Add dataloaders (v1/v2), analysis scripts, and documentation for working with UAV-GeoLoc (World-UAV). Co-authored-by: Cursor <cursoragent@cursor.com>
118 lines
4.0 KiB
Python
118 lines
4.0 KiB
Python
"""
|
||
Анализ схемы нарезки спутниковых снимков в датасете UAV-GeoLoc.
|
||
|
||
Скрипт определяет crop_size, stride и overlap для каждой сцены,
|
||
сопоставляя кропы с исходным merge.tif через попиксельное сравнение.
|
||
|
||
Результат: stride = crop_size // 2 (50% overlap) для всех сцен.
|
||
Naming: crop_X_Y.png — X по ширине (col), Y по высоте (row).
|
||
Позиция в merge.tif: merge[Y*stride : Y*stride+crop_size, X*stride : X*stride+crop_size]
|
||
"""
|
||
|
||
import glob
|
||
import os
|
||
|
||
import numpy as np
|
||
from PIL import Image
|
||
|
||
Image.MAX_IMAGE_PIXELS = None # некоторые merge.tif очень большие
|
||
|
||
|
||
def analyze_scene(scene_db_dir: str) -> dict:
|
||
"""Определяет параметры нарезки для одной сцены.
|
||
|
||
Args:
|
||
scene_db_dir: путь к папке DB сцены (содержит merge.tif и img/).
|
||
|
||
Returns:
|
||
dict с ключами: merge_size, crop_size, grid, stride, overlap.
|
||
"""
|
||
merge_path = os.path.join(scene_db_dir, "merge.tif")
|
||
img_dir = os.path.join(scene_db_dir, "img")
|
||
|
||
merge = np.array(Image.open(merge_path))
|
||
mh, mw = merge.shape[:2]
|
||
|
||
# Размер кропа
|
||
c00 = np.array(Image.open(os.path.join(img_dir, "crop_0_0.png")))
|
||
ch, cw = c00.shape[:2]
|
||
|
||
# Размер сетки
|
||
crops = os.listdir(img_dir)
|
||
xs, ys = [], []
|
||
for name in crops:
|
||
parts = name.replace("crop_", "").replace(".png", "").split("_")
|
||
xs.append(int(parts[0]))
|
||
ys.append(int(parts[1]))
|
||
grid_x, grid_y = max(xs) + 1, max(ys) + 1
|
||
|
||
# Проверяем что crop_0_0 начинается с (0, 0)
|
||
assert np.array_equal(c00, merge[0:ch, 0:cw, :3]), "crop_0_0 не совпадает с merge[0:ch, 0:cw]"
|
||
|
||
# Ищем stride по X: сдвигаем crop_1_0 вдоль ширины merge
|
||
c10 = np.array(Image.open(os.path.join(img_dir, "crop_1_0.png")))
|
||
stride_x = None
|
||
for s in range(1, cw + 1):
|
||
if s + cw <= mw and np.array_equal(c10, merge[0:ch, s:s + cw, :3]):
|
||
stride_x = s
|
||
break
|
||
assert stride_x is not None, "Не удалось найти stride по X"
|
||
|
||
# Ищем stride по Y: сдвигаем crop_0_1 вдоль высоты merge
|
||
c01 = np.array(Image.open(os.path.join(img_dir, "crop_0_1.png")))
|
||
stride_y = None
|
||
for s in range(1, ch + 1):
|
||
if s + ch <= mh and np.array_equal(c01, merge[s:s + ch, 0:cw, :3]):
|
||
stride_y = s
|
||
break
|
||
assert stride_y is not None, "Не удалось найти stride по Y"
|
||
|
||
return {
|
||
"merge_size": (mw, mh),
|
||
"crop_size": (cw, ch),
|
||
"grid": (grid_x, grid_y),
|
||
"stride": (stride_x, stride_y),
|
||
"overlap": (cw - stride_x, ch - stride_y),
|
||
}
|
||
|
||
|
||
def main():
|
||
base = os.path.dirname(os.path.abspath(__file__))
|
||
|
||
patterns = [
|
||
os.path.join(base, "Country", "*", "*", "*", "DB"),
|
||
os.path.join(base, "Terrain", "*", "*", "DB"),
|
||
]
|
||
|
||
scene_dirs = []
|
||
for pat in patterns:
|
||
scene_dirs.extend(sorted(glob.glob(pat)))
|
||
|
||
print(f"{'Scene':<50} {'merge WxH':>14} {'crop':>8} {'grid':>8} {'stride':>8} {'overlap':>8}")
|
||
print("-" * 100)
|
||
|
||
for scene_db in scene_dirs:
|
||
if not os.path.isfile(os.path.join(scene_db, "merge.tif")):
|
||
continue
|
||
|
||
# Короткое имя сцены
|
||
rel = os.path.relpath(scene_db, base)
|
||
name = rel.replace("/DB", "")
|
||
|
||
try:
|
||
info = analyze_scene(scene_db)
|
||
mw, mh = info["merge_size"]
|
||
cw, ch = info["crop_size"]
|
||
gx, gy = info["grid"]
|
||
sx, sy = info["stride"]
|
||
ox, oy = info["overlap"]
|
||
print(
|
||
f"{name:<50} {mw:>6}x{mh:<6} {cw:>3}x{ch:<4} {gx:>3}x{gy:<4} {sx:>3}x{sy:<4} {ox:>3}x{oy:<4}"
|
||
)
|
||
except Exception as e:
|
||
print(f"{name:<50} ERROR: {e}")
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main()
|