247 lines
11 KiB
Python
247 lines
11 KiB
Python
#!/usr/bin/env python3
|
|
"""Refine official footprint boxes into auditable image-visible SAM roof boxes."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import gc
|
|
import hashlib
|
|
import json
|
|
import math
|
|
import os
|
|
import shutil
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
|
|
def sha256(path: Path) -> str:
|
|
digest = hashlib.sha256()
|
|
with path.open("rb") as stream:
|
|
for chunk in iter(lambda: stream.read(1024 * 1024), b""):
|
|
digest.update(chunk)
|
|
return digest.hexdigest()
|
|
|
|
|
|
def iou(left: tuple[float, float, float, float], right: tuple[float, float, float, float]) -> float:
|
|
ix = max(0.0, min(left[2], right[2]) - max(left[0], right[0]))
|
|
iy = max(0.0, min(left[3], right[3]) - max(left[1], right[1]))
|
|
intersection = ix * iy
|
|
union = (left[2] - left[0]) * (left[3] - left[1]) + (right[2] - right[0]) * (
|
|
right[3] - right[1]
|
|
) - intersection
|
|
return intersection / union if union > 0 else 0.0
|
|
|
|
|
|
def plausible_refinement(
|
|
source: tuple[float, float, float, float],
|
|
refined: tuple[float, float, float, float],
|
|
*,
|
|
min_iou: float,
|
|
min_area_ratio: float,
|
|
max_area_ratio: float,
|
|
max_center_shift_ratio: float,
|
|
min_dimension_ratio: float,
|
|
max_dimension_ratio: float,
|
|
) -> bool:
|
|
source_width, source_height = source[2] - source[0], source[3] - source[1]
|
|
refined_width, refined_height = refined[2] - refined[0], refined[3] - refined[1]
|
|
source_area = source_width * source_height
|
|
refined_area = refined_width * refined_height
|
|
ratio = refined_area / source_area if source_area > 0 else 0.0
|
|
if source_width <= 0 or source_height <= 0:
|
|
return False
|
|
width_ratio = refined_width / source_width
|
|
height_ratio = refined_height / source_height
|
|
source_centre = ((source[0] + source[2]) / 2, (source[1] + source[3]) / 2)
|
|
refined_centre = ((refined[0] + refined[2]) / 2, (refined[1] + refined[3]) / 2)
|
|
centre_shift_ratio = math.dist(source_centre, refined_centre) / math.hypot(source_width, source_height)
|
|
return (
|
|
min_area_ratio <= ratio <= max_area_ratio
|
|
and min_dimension_ratio <= width_ratio <= max_dimension_ratio
|
|
and min_dimension_ratio <= height_ratio <= max_dimension_ratio
|
|
and centre_shift_ratio <= max_center_shift_ratio
|
|
and iou(source, refined) >= min_iou
|
|
)
|
|
|
|
|
|
def read_boxes(path: Path, width: int, height: int) -> list[tuple[float, float, float, float]]:
|
|
boxes = []
|
|
for line in path.read_text(encoding="utf-8").splitlines() if path.is_file() else []:
|
|
class_id, cx, cy, bw, bh = map(float, line.split())
|
|
if class_id != 0:
|
|
raise ValueError(f"Unexpected class in {path}: {class_id}")
|
|
boxes.append(((cx - bw / 2) * width, (cy - bh / 2) * height, (cx + bw / 2) * width, (cy + bh / 2) * height))
|
|
return boxes
|
|
|
|
|
|
def yolo_line(box: tuple[float, float, float, float], width: int, height: int) -> str:
|
|
x1, y1, x2, y2 = box
|
|
return f"0 {(x1+x2)/(2*width):.8f} {(y1+y2)/(2*height):.8f} {(x2-x1)/width:.8f} {(y2-y1)/height:.8f}"
|
|
|
|
|
|
def expanded(box: tuple[float, float, float, float], factor: float, width: int, height: int) -> list[float]:
|
|
x1, y1, x2, y2 = box
|
|
cx, cy = (x1 + x2) / 2, (y1 + y2) / 2
|
|
hw, hh = (x2 - x1) * factor / 2, (y2 - y1) * factor / 2
|
|
return [max(0, cx - hw), max(0, cy - hh), min(width - 1, cx + hw), min(height - 1, cy + hh)]
|
|
|
|
|
|
def link_or_copy(source: Path, target: Path) -> None:
|
|
target.parent.mkdir(parents=True, exist_ok=True)
|
|
try:
|
|
os.link(source, target)
|
|
except OSError:
|
|
shutil.copy2(source, target)
|
|
|
|
|
|
def main() -> int:
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("--summary", type=Path, required=True)
|
|
parser.add_argument("--model", type=Path, required=True)
|
|
parser.add_argument("--output-dir", type=Path, required=True)
|
|
parser.add_argument("--device", default="cuda:0")
|
|
parser.add_argument("--prompt-scale", type=float, default=1.5)
|
|
parser.add_argument("--min-source-iou", type=float, default=0.15)
|
|
parser.add_argument("--min-area-ratio", type=float, default=0.25)
|
|
parser.add_argument("--max-area-ratio", type=float, default=4.0)
|
|
parser.add_argument("--max-center-shift-ratio", type=float, default=0.75)
|
|
parser.add_argument("--min-dimension-ratio", type=float, default=0.5)
|
|
parser.add_argument("--max-dimension-ratio", type=float, default=2.0)
|
|
parser.add_argument("--max-prompts-per-pass", type=int, default=96)
|
|
parser.add_argument("--fallback-policy", choices=("retain", "drop"), default="retain")
|
|
parser.add_argument("--force", action="store_true")
|
|
args = parser.parse_args()
|
|
if args.output_dir.exists():
|
|
if not args.force:
|
|
raise SystemExit(f"Output exists: {args.output_dir}")
|
|
shutil.rmtree(args.output_dir)
|
|
from PIL import Image
|
|
import torch
|
|
from ultralytics import SAM
|
|
|
|
summary = json.loads(args.summary.read_text(encoding="utf-8"))
|
|
model = SAM(str(args.model))
|
|
refined_count = fallback_count = 0
|
|
reason_counts: dict[str, int] = {}
|
|
output_tiles = []
|
|
for index, tile in enumerate(summary["tiles"], start=1):
|
|
source_image = Path(tile["image_path"])
|
|
source_label = Path(tile["label_path"])
|
|
split = tile["split"]
|
|
target_image = args.output_dir / "images" / split / source_image.name
|
|
target_label = args.output_dir / "labels" / split / source_label.name
|
|
link_or_copy(source_image, target_image)
|
|
target_label.parent.mkdir(parents=True, exist_ok=True)
|
|
with Image.open(source_image) as image:
|
|
width, height = image.size
|
|
source_boxes = read_boxes(source_label, width, height)
|
|
output_boxes: list[tuple[float, float, float, float] | None] = list(source_boxes)
|
|
if source_boxes:
|
|
for start in range(0, len(source_boxes), args.max_prompts_per_pass):
|
|
source_chunk = source_boxes[start : start + args.max_prompts_per_pass]
|
|
prompts = [expanded(box, args.prompt_scale, width, height) for box in source_chunk]
|
|
result = model.predict(str(source_image), bboxes=prompts, device=args.device, verbose=False)[0]
|
|
masks = result.masks.data.cpu().numpy() if result.masks is not None else []
|
|
candidates = []
|
|
for mask in masks:
|
|
ys, xs = mask.nonzero()
|
|
if len(xs):
|
|
candidates.append((float(xs.min()), float(ys.min()), float(xs.max() + 1), float(ys.max() + 1)))
|
|
unmatched = set(range(len(candidates)))
|
|
for local_index, source_box in enumerate(source_chunk):
|
|
ranked = sorted(((iou(source_box, candidates[item]), item) for item in unmatched), reverse=True)
|
|
overlap, candidate_index = ranked[0] if ranked else (0.0, -1)
|
|
if candidate_index < 0 or overlap < args.min_source_iou:
|
|
reason_counts["unmatched_mask"] = reason_counts.get("unmatched_mask", 0) + 1
|
|
fallback_count += 1
|
|
if args.fallback_policy == "drop":
|
|
output_boxes[start + local_index] = None
|
|
continue
|
|
candidate = candidates[candidate_index]
|
|
if plausible_refinement(
|
|
source_box,
|
|
candidate,
|
|
min_iou=args.min_source_iou,
|
|
min_area_ratio=args.min_area_ratio,
|
|
max_area_ratio=args.max_area_ratio,
|
|
max_center_shift_ratio=args.max_center_shift_ratio,
|
|
min_dimension_ratio=args.min_dimension_ratio,
|
|
max_dimension_ratio=args.max_dimension_ratio,
|
|
):
|
|
output_boxes[start + local_index] = candidate
|
|
unmatched.remove(candidate_index)
|
|
refined_count += 1
|
|
else:
|
|
reason_counts["geometry_gate"] = reason_counts.get("geometry_gate", 0) + 1
|
|
fallback_count += 1
|
|
if args.fallback_policy == "drop":
|
|
output_boxes[start + local_index] = None
|
|
del result, masks
|
|
if model.predictor is not None:
|
|
model.predictor.reset_image()
|
|
gc.collect()
|
|
torch.cuda.empty_cache()
|
|
retained_boxes = [box for box in output_boxes if box is not None]
|
|
target_label.write_text("\n".join(yolo_line(box, width, height) for box in retained_boxes) + ("\n" if retained_boxes else ""), encoding="utf-8")
|
|
output_tile = dict(tile)
|
|
output_tile.update({
|
|
"image_path": str(target_image),
|
|
"label_path": str(target_label),
|
|
"label_count": len(retained_boxes),
|
|
"is_negative": not retained_boxes,
|
|
})
|
|
output_tiles.append(output_tile)
|
|
print(f"{index}/{len(summary['tiles'])} {tile['sample_slug']}: {len(source_boxes)}", flush=True)
|
|
|
|
output_summary = dict(summary)
|
|
output_summary.update(
|
|
{
|
|
"output_dir": str(args.output_dir),
|
|
"dataset_yaml": str(args.output_dir / "dataset.yaml"),
|
|
"tiles": output_tiles,
|
|
"label_semantics": (
|
|
"sam_visible_roof_only"
|
|
if args.fallback_policy == "drop"
|
|
else "sam_visible_roof_with_official_footprint_fallback"
|
|
),
|
|
"label_count": refined_count if args.fallback_policy == "drop" else summary["label_count"],
|
|
"positive_tile_count": sum(not tile["is_negative"] for tile in output_tiles),
|
|
"negative_tile_count": sum(tile["is_negative"] for tile in output_tiles),
|
|
}
|
|
)
|
|
summary_path = args.output_dir / "yolo_tile_dataset_summary.json"
|
|
summary_path.write_text(json.dumps(output_summary, indent=2), encoding="utf-8")
|
|
(args.output_dir / "dataset.yaml").write_text(
|
|
f"path: {args.output_dir}\ntrain: images/train\nval: images/val\nnames:\n 0: building\n",
|
|
encoding="utf-8",
|
|
)
|
|
evidence = {
|
|
"schema_version": 1,
|
|
"status": "ok",
|
|
"source_summary": str(args.summary),
|
|
"source_summary_sha256": sha256(args.summary),
|
|
"sam_model": str(args.model),
|
|
"sam_model_sha256": sha256(args.model),
|
|
"device": args.device,
|
|
"prompt_scale": args.prompt_scale,
|
|
"min_source_iou": args.min_source_iou,
|
|
"min_area_ratio": args.min_area_ratio,
|
|
"max_area_ratio": args.max_area_ratio,
|
|
"max_center_shift_ratio": args.max_center_shift_ratio,
|
|
"min_dimension_ratio": args.min_dimension_ratio,
|
|
"max_dimension_ratio": args.max_dimension_ratio,
|
|
"max_prompts_per_pass": args.max_prompts_per_pass,
|
|
"fallback_policy": args.fallback_policy,
|
|
"refined_label_count": refined_count,
|
|
"fallback_label_count": fallback_count,
|
|
"dropped_fallback_label_count": fallback_count if args.fallback_policy == "drop" else 0,
|
|
"fallback_reason_counts": reason_counts,
|
|
}
|
|
(args.output_dir / "sam-refinement.json").write_text(json.dumps(evidence, indent=2), encoding="utf-8")
|
|
print(json.dumps(evidence, indent=2))
|
|
return 0
|
|
|
|
|
|
if __name__ == "__main__":
|
|
raise SystemExit(main())
|