686db62c25
4-modality annotation pipeline (depth, edges, segmentation, chmv2) for 973K drone/satellite images. SegEarth-OV3 open-vocabulary segmentation with 11 classes optimized for cross-view geo-localization. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
115 lines
3.9 KiB
Python
115 lines
3.9 KiB
Python
"""Quick test: run 12-class segmentation on 10 diverse images and save results."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import sys
|
|
import shutil
|
|
from pathlib import Path
|
|
|
|
import torch
|
|
import numpy as np
|
|
from PIL import Image
|
|
|
|
# Ensure project root is on path.
|
|
ROOT = Path(__file__).resolve().parent
|
|
sys.path.insert(0, str(ROOT))
|
|
sys.path.insert(0, str(ROOT / "src" / "nn" / "segearth_ov3"))
|
|
|
|
from src.conf.config_loader import load_all_configs
|
|
from src.augmentor.models import load_segmentation_model
|
|
from src.augmentor.inference import infer_segmentation_batch
|
|
from src.augmentor.io_utils import save_segmentation
|
|
|
|
INPUT_ROOT = Path("/mnt/data1tb/cvgl_datasets/UAV-GeoLoc")
|
|
|
|
# 10 diverse test images: satellite (DB) — urban + terrain
|
|
TEST_IMAGES = [
|
|
# Urban (satellite)
|
|
"Country/German/Munich/Ludwigs/DB/img/crop_12_4.png",
|
|
"Country/French/Paris/LeMarais/DB/img/crop_12_4.png",
|
|
"Country/Italy/Venice/SanMarco/DB/img/crop_12_4.png",
|
|
"Country/USA/NewYork/Manhattan/DB/img/crop_12_4.png",
|
|
"Country/Australia/Sydney/SydneyCBD/DB/img/crop_12_4.png",
|
|
"Country/Korea/Busan/Gwangalli/DB/img/crop_12_4.png",
|
|
# Terrain (satellite)
|
|
"Terrain/Desert/GobiDesert/DB/img/crop_50_33.png",
|
|
"Terrain/Glacier/AthabascaGlacier/DB/img/crop_12_4.png",
|
|
"Terrain/Volcano/KilaueaVolcano/DB/img/crop_12_4.png",
|
|
"Terrain/Danxia/GrandCanyon/DB/img/crop_50_33.png",
|
|
]
|
|
|
|
OUTPUT_DIR = ROOT / "test_seg_output"
|
|
|
|
|
|
def main():
|
|
# Load configs (reads all .gin files)
|
|
configs = load_all_configs(str(ROOT / "in" / "config_files") + "/")
|
|
pipeline_conf = configs["pipeline"]
|
|
hw_conf = configs["hardware"]
|
|
models_conf = configs["models"]
|
|
input_conf = configs["input"]
|
|
seg_conf = configs["seg"]
|
|
|
|
print(f"Prompts ({len(seg_conf.prompts)}): {seg_conf.prompts}")
|
|
print(f"Threshold: {seg_conf.threshold}")
|
|
print(f"Resolution: {seg_conf.default_resolution}")
|
|
|
|
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
|
print(f"Device: {device}")
|
|
|
|
# Load model
|
|
print("\nLoading SegEarth-OV3...")
|
|
model, seg_config = load_segmentation_model(models_conf, hw_conf, seg_conf, device)
|
|
num_classes = seg_config.get("num_classes", 150)
|
|
print(f"Model loaded. Type: {seg_config['type']}, num_classes: {num_classes}")
|
|
|
|
# Prepare output dir
|
|
OUTPUT_DIR.mkdir(exist_ok=True)
|
|
|
|
# Process each image
|
|
for i, rel_path in enumerate(TEST_IMAGES):
|
|
img_path = INPUT_ROOT / rel_path
|
|
if not img_path.exists():
|
|
print(f"[{i+1}/10] SKIP (not found): {rel_path}")
|
|
continue
|
|
|
|
print(f"\n[{i+1}/10] Processing: {rel_path}")
|
|
|
|
# Load and preprocess
|
|
pil_img = Image.open(img_path).convert("RGB")
|
|
img_np = np.array(pil_img).astype(np.float32) / 255.0
|
|
img_tensor = torch.from_numpy(img_np).permute(2, 0, 1).unsqueeze(0) # [1, 3, H, W]
|
|
|
|
# Infer
|
|
seg_ids = infer_segmentation_batch(model, seg_config, img_tensor, device)
|
|
# seg_ids: [1, 1, H, W] uint8
|
|
|
|
# Save original + segmentation side by side
|
|
stem = img_path.stem
|
|
tag = rel_path.split("/")[0] # Country or Terrain
|
|
if tag == "Country":
|
|
tag = rel_path.split("/")[2] # city name
|
|
elif tag == "Terrain":
|
|
tag = rel_path.split("/")[1] # terrain type
|
|
|
|
out_stem = f"{i:02d}_{tag}_{stem}"
|
|
|
|
# Save original for reference
|
|
pil_img.save(OUTPUT_DIR / f"{out_stem}_orig.png")
|
|
|
|
# Save segmentation vis
|
|
save_segmentation(
|
|
seg_ids[0], OUTPUT_DIR, out_stem,
|
|
save_npy=False, save_vis=True, num_classes=num_classes,
|
|
)
|
|
|
|
unique_classes = torch.unique(seg_ids[0]).tolist()
|
|
class_names = [seg_conf.prompts[c] for c in unique_classes if c < len(seg_conf.prompts)]
|
|
print(f" -> Classes found: {class_names}")
|
|
|
|
print(f"\nDone! Results saved to: {OUTPUT_DIR}")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|