"""Run the first GeoAI object-detection baseline on drone images.
|
|
This baseline uses the opengeos/geoai environment plus a general-purpose
|
Ultralytics model. It intentionally keeps georeferenced GeoTIFF processing for
|
the next stage, because the current JPEGs do not carry usable GPS metadata.
|
"""
|
|
from __future__ import annotations
|
|
import argparse
|
import json
|
import os
|
import sys
|
import tempfile
|
import time
|
from datetime import UTC, datetime
|
from pathlib import Path
|
from typing import Any
|
|
from PIL import Image
|
|
|
IMAGE_SUFFIXES = {".jpg", ".jpeg", ".png"}
|
COCO_TARGET_CLASS_IDS = [0, 1, 2, 3, 5, 7]
|
|
|
def parse_args() -> argparse.Namespace:
|
root = Path(__file__).resolve().parents[2]
|
default_input = root / "shared" / "data" / "raw" / "01-object-detection"
|
default_output = root / "shared" / "outputs" / "01-object-detection"
|
parser = argparse.ArgumentParser(description="Detect objects in drone images.")
|
parser.add_argument("--input", type=Path, default=default_input)
|
parser.add_argument("--output", type=Path, default=default_output)
|
parser.add_argument("--model", default="yolo11n.pt", help="Ultralytics model name or path.")
|
parser.add_argument("--confidence", type=float, default=0.20)
|
parser.add_argument("--image-size", type=int, default=1024, help="Model input size for each tile.")
|
parser.add_argument("--tile-size", type=int, default=1024, help="Pixel size of the sliding detection window.")
|
parser.add_argument("--tile-overlap", type=float, default=0.20, help="Overlap ratio between adjacent windows.")
|
return parser.parse_args()
|
|
|
def load_image_paths(input_dir: Path) -> tuple[list[Path], list[dict[str, str]]]:
|
paths: list[Path] = []
|
skipped: list[dict[str, str]] = []
|
candidates = [input_dir] if input_dir.is_file() else sorted(input_dir.iterdir())
|
for path in candidates:
|
if not path.is_file() or path.suffix.lower() not in IMAGE_SUFFIXES:
|
continue
|
try:
|
with Image.open(path) as image:
|
image.verify()
|
if image.format == "MPO":
|
skipped.append({"file": path.name, "reason": "MPO is not supported in baseline"})
|
continue
|
except Exception as exc: # pragma: no cover - depends on source files
|
skipped.append({"file": path.name, "reason": f"unreadable: {exc}"})
|
continue
|
paths.append(path)
|
return paths, skipped
|
|
|
def tile_starts(length: int, tile_size: int, overlap: float) -> list[int]:
|
if not 0 <= overlap < 1:
|
raise ValueError("tile-overlap must be between 0 and 1")
|
if length <= tile_size:
|
return [0]
|
stride = max(1, int(tile_size * (1 - overlap)))
|
starts = list(range(0, length - tile_size + 1, stride))
|
last = length - tile_size
|
if starts[-1] != last:
|
starts.append(last)
|
return starts
|
|
|
def intersection_over_union(box_a: list[float], box_b: list[float]) -> float:
|
left = max(box_a[0], box_b[0])
|
top = max(box_a[1], box_b[1])
|
right = min(box_a[2], box_b[2])
|
bottom = min(box_a[3], box_b[3])
|
intersection = max(0.0, right - left) * max(0.0, bottom - top)
|
area_a = max(0.0, box_a[2] - box_a[0]) * max(0.0, box_a[3] - box_a[1])
|
area_b = max(0.0, box_b[2] - box_b[0]) * max(0.0, box_b[3] - box_b[1])
|
union = area_a + area_b - intersection
|
return intersection / union if union else 0.0
|
|
|
def intersection_over_smaller(box_a: list[float], box_b: list[float]) -> float:
|
left = max(box_a[0], box_b[0])
|
top = max(box_a[1], box_b[1])
|
right = min(box_a[2], box_b[2])
|
bottom = min(box_a[3], box_b[3])
|
intersection = max(0.0, right - left) * max(0.0, bottom - top)
|
area_a = max(0.0, box_a[2] - box_a[0]) * max(0.0, box_a[3] - box_a[1])
|
area_b = max(0.0, box_b[2] - box_b[0]) * max(0.0, box_b[3] - box_b[1])
|
smaller_area = min(area_a, area_b)
|
return intersection / smaller_area if smaller_area else 0.0
|
|
|
def class_aware_nms(
|
detections: list[dict[str, Any]],
|
iou_threshold: float = 0.50,
|
containment_threshold: float = 0.80,
|
) -> list[dict[str, Any]]:
|
kept: list[dict[str, Any]] = []
|
for candidate in sorted(detections, key=lambda item: item["confidence"], reverse=True):
|
if all(
|
candidate["class_id"] != existing["class_id"]
|
or (
|
intersection_over_union(candidate["bbox_xyxy"], existing["bbox_xyxy"]) < iou_threshold
|
and intersection_over_smaller(candidate["bbox_xyxy"], existing["bbox_xyxy"])
|
< containment_threshold
|
)
|
for existing in kept
|
):
|
kept.append(candidate)
|
return kept
|
|
|
def predict_tiled(model: Any, image_path: Path, args: argparse.Namespace, np: Any) -> tuple[list[dict[str, Any]], int, int]:
|
with Image.open(image_path) as source:
|
image = source.convert("RGB")
|
width, height = image.size
|
detections: list[dict[str, Any]] = []
|
names: dict[int, str] = {}
|
for y0 in tile_starts(height, args.tile_size, args.tile_overlap):
|
for x0 in tile_starts(width, args.tile_size, args.tile_overlap):
|
tile = image.crop((x0, y0, min(x0 + args.tile_size, width), min(y0 + args.tile_size, height)))
|
result = model.predict(
|
source=np.asarray(tile),
|
device="cpu",
|
imgsz=args.image_size,
|
conf=args.confidence,
|
classes=COCO_TARGET_CLASS_IDS,
|
max_det=100,
|
verbose=False,
|
)[0]
|
names = result.names
|
if result.boxes is None:
|
continue
|
for box in result.boxes:
|
class_id = int(box.cls.item())
|
local_box = [float(value) for value in box.xyxy[0].tolist()]
|
detections.append(
|
{
|
"class_id": class_id,
|
"class_name": names[class_id],
|
"confidence": round(float(box.conf.item()), 4),
|
"bbox_xyxy": [
|
round(local_box[0] + x0, 2),
|
round(local_box[1] + y0, 2),
|
round(local_box[2] + x0, 2),
|
round(local_box[3] + y0, 2),
|
],
|
}
|
)
|
return class_aware_nms(detections), width, height
|
|
|
def draw_detections(image_path: Path, detections: list[dict[str, Any]], output_path: Path) -> None:
|
import cv2
|
import numpy as np
|
|
canvas = cv2.cvtColor(np.asarray(Image.open(image_path).convert("RGB")), cv2.COLOR_RGB2BGR)
|
for detection in detections:
|
x1, y1, x2, y2 = [int(round(value)) for value in detection["bbox_xyxy"]]
|
label = f"{detection['class_name']} {detection['confidence']:.2f}"
|
cv2.rectangle(canvas, (x1, y1), (x2, y2), (255, 80, 0), 4)
|
cv2.putText(canvas, label, (x1, max(30, y1 - 8)), cv2.FONT_HERSHEY_SIMPLEX, 1.0, (255, 80, 0), 2)
|
cv2.imwrite(str(output_path), canvas)
|
|
|
def main() -> int:
|
args = parse_args()
|
input_dir = args.input.resolve()
|
output_dir = args.output.resolve()
|
annotated_dir = output_dir / "annotated"
|
annotated_dir.mkdir(parents=True, exist_ok=True)
|
cache_root = Path(tempfile.gettempdir()) / "geoai-object-detection"
|
os.environ.setdefault("YOLO_CONFIG_DIR", str(cache_root / "ultralytics"))
|
os.environ.setdefault("MPLCONFIGDIR", str(cache_root / "matplotlib"))
|
if not input_dir.exists():
|
print(f"Input directory does not exist: {input_dir}", file=sys.stderr)
|
return 2
|
|
image_paths, skipped = load_image_paths(input_dir)
|
if not image_paths:
|
print("No supported images found.", file=sys.stderr)
|
return 2
|
|
# Imports happen after argument validation so the CLI can explain path
|
# mistakes without requiring the heavyweight ML stack.
|
import geoai
|
import numpy as np
|
import torch
|
import ultralytics
|
from ultralytics import YOLO
|
|
started = time.perf_counter()
|
model = YOLO(args.model)
|
detections: list[dict[str, Any]] = []
|
for image_path in image_paths:
|
image_detections, width, height = predict_tiled(model, image_path, args, np)
|
annotated_path = annotated_dir / image_path.name
|
draw_detections(image_path, image_detections, annotated_path)
|
detections.append(
|
{
|
"file": image_path.name,
|
"annotated_file": str(annotated_path.relative_to(output_dir)),
|
"width": width,
|
"height": height,
|
"detections": image_detections,
|
}
|
)
|
|
elapsed = round(time.perf_counter() - started, 3)
|
(output_dir / "detections.json").write_text(
|
json.dumps({"images": detections}, ensure_ascii=False, indent=2), encoding="utf-8"
|
)
|
metadata = {
|
"created_at": datetime.now(UTC).isoformat(),
|
"geoai_package": getattr(geoai, "__version__", "unknown"),
|
"ultralytics": ultralytics.__version__,
|
"torch": torch.__version__,
|
"device": "cpu",
|
"cuda_available": bool(torch.cuda.is_available()),
|
"model": args.model,
|
"confidence": args.confidence,
|
"image_size": args.image_size,
|
"tile_size": args.tile_size,
|
"tile_overlap": args.tile_overlap,
|
"merge_iou_threshold": 0.50,
|
"merge_containment_threshold": 0.80,
|
"target_class_ids": COCO_TARGET_CLASS_IDS,
|
"target_class_note": "COCO person, bicycle, car, motorcycle, bus and truck; tree is not a COCO class.",
|
"input_dir": str(input_dir),
|
"processed_images": len(image_paths),
|
"skipped_images": skipped,
|
"detection_count": sum(len(item["detections"]) for item in detections),
|
"elapsed_seconds": elapsed,
|
"notes": [
|
"This is a baseline for people and common vehicle classes.",
|
"The default COCO model does not provide a tree class.",
|
"JPEG inputs have no usable GPS metadata; GeoJSON is deferred to GeoTIFF stage.",
|
],
|
}
|
(output_dir / "run_metadata.json").write_text(
|
json.dumps(metadata, ensure_ascii=False, indent=2), encoding="utf-8"
|
)
|
print(f"Processed {len(image_paths)} images; detections={metadata['detection_count']}; elapsed={elapsed}s")
|
print(f"Annotated images: {annotated_dir}")
|
print(f"JSON: {output_dir / 'detections.json'}")
|
return 0
|
|
|
if __name__ == "__main__":
|
raise SystemExit(main())
|
