World-UAV-ds/analyze/generate_sample_grids.py

#!/usr/bin/env python3
"""Generate sample image grids for UAV-GeoLoc dataset analysis."""

import json
import os
import matplotlib.pyplot as plt
import matplotlib.patches as patches
import numpy as np
from PIL import Image

BASE = "/mnt/data1tb/cvgl_datasets/UAV-GeoLoc"
OUT = os.path.join(BASE, "charts")
os.makedirs(OUT, exist_ok=True)


def load_img(path):
    """Load image as numpy array."""
    return np.array(Image.open(path))


# =============================================================================
# 1. sample_query_db_pairs.png
# =============================================================================
def make_query_db_pairs():
    scenes = [
        ("Country/Australia/Adelaide/AdelaideCBD", "Adelaide CBD, Australia"),
        ("Country/USA/NewYork/Manhattan", "Manhattan, New York"),
        ("Terrain/Mountain/Andes", "Andes Mountains"),
        ("Terrain/Desert/GobiDesert", "Gobi Desert"),
    ]

    fig, axes = plt.subplots(4, 2, figsize=(8, 16))
    fig.suptitle("UAV Query vs. Satellite DB Positive Match", fontsize=16, fontweight="bold", y=0.98)

    for row, (scene_rel, label) in enumerate(scenes):
        scene_path = os.path.join(BASE, scene_rel)

        # Load positive.json to find the DB match for frame "00"
        with open(os.path.join(scene_path, "positive.json")) as f:
            positives = json.load(f)
        db_crop_name = positives["00"][0]  # first positive match

        # Query image
        query_path = os.path.join(scene_path, "query", "height100_rot0", "footage", "height100_rot0_00.jpeg")
        query_img = load_img(query_path)

        # DB crop
        db_path = os.path.join(scene_path, "DB", "img", db_crop_name)
        db_img = load_img(db_path)

        axes[row, 0].imshow(query_img)
        axes[row, 0].set_title(f"Query (UAV)\n{label}", fontsize=10)
        axes[row, 0].axis("off")

        axes[row, 1].imshow(db_img)
        axes[row, 1].set_title(f"Positive DB Match\n{db_crop_name}", fontsize=10)
        axes[row, 1].axis("off")

    plt.tight_layout(rect=[0, 0, 1, 0.96])
    out_path = os.path.join(OUT, "sample_query_db_pairs.png")
    fig.savefig(out_path, dpi=150, bbox_inches="tight", facecolor="white")
    plt.close(fig)
    print(f"Saved {out_path} ({os.path.getsize(out_path) / 1024:.1f} KB)")


# =============================================================================
# 2. sample_height_comparison.png
# =============================================================================
def make_height_comparison():
    scene = "Country/Australia/Adelaide/AdelaideCBD"
    scene_path = os.path.join(BASE, scene)
    heights = [100, 125, 150]

    fig, axes = plt.subplots(1, 3, figsize=(15, 5))
    fig.suptitle("Same Scene at Different UAV Heights (Adelaide CBD, rot=0, frame 00)",
                 fontsize=14, fontweight="bold")

    for i, h in enumerate(heights):
        img_path = os.path.join(scene_path, "query", f"height{h}_rot0", "footage", f"height{h}_rot0_00.jpeg")
        img = load_img(img_path)
        axes[i].imshow(img)
        axes[i].set_title(f"Height = {h}m", fontsize=13, fontweight="bold")
        axes[i].axis("off")

    plt.tight_layout()
    out_path = os.path.join(OUT, "sample_height_comparison.png")
    fig.savefig(out_path, dpi=150, bbox_inches="tight", facecolor="white")
    plt.close(fig)
    print(f"Saved {out_path} ({os.path.getsize(out_path) / 1024:.1f} KB)")


# =============================================================================
# 3. sample_rotation_comparison.png
# =============================================================================
def make_rotation_comparison():
    scene = "Country/Australia/Adelaide/AdelaideCBD"
    scene_path = os.path.join(BASE, scene)
    rotations = [0, 45, 90, 135, 180, 225, 270, 315]
    frame = "38"

    fig, axes = plt.subplots(2, 4, figsize=(16, 8))
    fig.suptitle(f"Same Scene at 8 Rotations (Adelaide CBD, height=100m, frame {frame})",
                 fontsize=14, fontweight="bold")

    for idx, rot in enumerate(rotations):
        r, c = divmod(idx, 4)
        img_path = os.path.join(scene_path, "query", f"height100_rot{rot}", "footage",
                                f"height100_rot{rot}_{frame}.jpeg")
        img = load_img(img_path)
        axes[r, c].imshow(img)
        axes[r, c].set_title(f"Rotation = {rot}\u00b0", fontsize=12, fontweight="bold")
        axes[r, c].axis("off")

    plt.tight_layout()
    out_path = os.path.join(OUT, "sample_rotation_comparison.png")
    fig.savefig(out_path, dpi=150, bbox_inches="tight", facecolor="white")
    plt.close(fig)
    print(f"Saved {out_path} ({os.path.getsize(out_path) / 1024:.1f} KB)")


# =============================================================================
# 4. sample_satellite_tiling.png
# =============================================================================
def make_satellite_tiling():
    scene = "Country/Australia/Adelaide/AdelaideCBD"
    scene_path = os.path.join(BASE, scene)

    merge_path = os.path.join(scene_path, "DB", "merge.tif")
    merge_img = Image.open(merge_path)

    # Crop to top-left 600x600 for visualization
    region_size = 600
    region = np.array(merge_img.crop((0, 0, region_size, region_size)))

    crop_size = 200
    stride = 100  # overlapping crops

    fig, ax = plt.subplots(1, 1, figsize=(8, 8))
    fig.suptitle("Satellite Image Tiling (200x200 crops, stride=100)\nAdelaide CBD - top-left 600x600 region",
                 fontsize=13, fontweight="bold")

    ax.imshow(region)

    colors = ["#FF4444", "#44FF44", "#4444FF", "#FFFF00", "#FF44FF", "#44FFFF",
              "#FF8800", "#8800FF", "#00FF88"]
    color_idx = 0

    # Draw crop rectangles for crops that fall within the 600x600 region
    for row in range(0, region_size - crop_size + 1, stride):
        for col in range(0, region_size - crop_size + 1, stride):
            rect = patches.Rectangle(
                (col, row), crop_size, crop_size,
                linewidth=1.5,
                edgecolor=colors[color_idx % len(colors)],
                facecolor="none",
                alpha=0.7,
            )
            ax.add_patch(rect)
            color_idx += 1

    # Highlight a few specific crops with thicker borders and labels
    highlights = [(0, 0, "crop_0_0"), (0, 100, "crop_0_1"), (100, 0, "crop_1_0"), (100, 100, "crop_1_1")]
    for col, row, name in highlights:
        rect = patches.Rectangle(
            (col, row), crop_size, crop_size,
            linewidth=3,
            edgecolor="white",
            facecolor="none",
        )
        ax.add_patch(rect)
        ax.text(col + 5, row + 15, name, fontsize=8, color="white", fontweight="bold",
                bbox=dict(boxstyle="round,pad=0.2", facecolor="black", alpha=0.7))

    ax.set_xlim(0, region_size)
    ax.set_ylim(region_size, 0)
    ax.axis("off")

    plt.tight_layout()
    out_path = os.path.join(OUT, "sample_satellite_tiling.png")
    fig.savefig(out_path, dpi=150, bbox_inches="tight", facecolor="white")
    plt.close(fig)
    print(f"Saved {out_path} ({os.path.getsize(out_path) / 1024:.1f} KB)")


# =============================================================================
# 5. sample_terrain_diversity.png
# =============================================================================
def make_terrain_diversity():
    terrains = [
        ("Mountain/Andes", "Mountain"),
        ("Desert/GobiDesert", "Desert"),
        ("Volcano/KilaueaVolcano", "Volcano"),
        ("Glacier/AthabascaGlacier", "Glacier"),
        ("Island/Aldabra", "Island"),
        ("Farm/Central_Valley_Chop_Shop", "Farm"),
        ("Gorge/AntelopeCanyon", "Gorge"),
        ("Flowers/BlueHotSpring", "Flowers"),
        ("Delta/Delaware", "Delta"),
    ]

    fig, axes = plt.subplots(3, 3, figsize=(12, 12))
    fig.suptitle("Terrain Type Diversity - Satellite DB Crops", fontsize=15, fontweight="bold", y=0.98)

    for idx, (rel_path, terrain_label) in enumerate(terrains):
        r, c = divmod(idx, 3)
        crop_path = os.path.join(BASE, "Terrain", rel_path, "DB", "img", "crop_0_0.png")
        img = load_img(crop_path)
        axes[r, c].imshow(img)
        axes[r, c].set_title(terrain_label, fontsize=13, fontweight="bold")
        axes[r, c].axis("off")

    plt.tight_layout(rect=[0, 0, 1, 0.96])
    out_path = os.path.join(OUT, "sample_terrain_diversity.png")
    fig.savefig(out_path, dpi=150, bbox_inches="tight", facecolor="white")
    plt.close(fig)
    print(f"Saved {out_path} ({os.path.getsize(out_path) / 1024:.1f} KB)")


# =============================================================================
# Main
# =============================================================================
if __name__ == "__main__":
    print("Generating sample image grids...")
    make_query_db_pairs()
    make_height_comparison()
    make_rotation_comparison()
    make_satellite_tiling()
    make_terrain_diversity()
    print("\nAll charts saved to:", OUT)