| |
| """Write deterministic per-image predictions from a CartoLegend point detector.""" |
|
|
| from __future__ import annotations |
|
|
| import argparse |
| import json |
| from datetime import datetime, timezone |
| from pathlib import Path |
|
|
| import torch |
| from PIL import Image |
| from torch.utils.data import DataLoader |
|
|
| from train_cartolegend_point_detector import ( |
| PointSymbolDataset, |
| build_model, |
| collate, |
| load_unique_records, |
| sha256, |
| ) |
| from cartolegend_point_detector_artifact import load_point_detector_artifact |
|
|
|
|
| ROOT = Path(__file__).resolve().parents[1] |
|
|
|
|
| def portable_path(path: Path) -> str: |
| resolved = path.expanduser().resolve() |
| try: |
| return "project://" + resolved.relative_to(ROOT).as_posix() |
| except ValueError: |
| return resolved.name |
|
|
|
|
| def load_prediction_records(path: Path, image_field: str) -> list[dict]: |
| if image_field == "images": |
| return load_unique_records(path) |
| records = [] |
| seen: set[Path] = set() |
| for line_number, line in enumerate(path.read_text().splitlines(), start=1): |
| if not line.strip(): |
| continue |
| row = json.loads(line) |
| value = str(row.get(image_field) or "") |
| image = Path(value) |
| image = (image if image.is_absolute() else ROOT / image).resolve() |
| if not image.is_file(): |
| raise FileNotFoundError(f"line {line_number}: {image}") |
| if image in seen: |
| continue |
| seen.add(image) |
| with Image.open(image) as source: |
| width, height = source.size |
| records.append( |
| { |
| "image": str(image), |
| "width": width, |
| "height": height, |
| "boxes": [], |
| "text_labels": [], |
| "group": str(row.get("source_id") or image.stem), |
| } |
| ) |
| return records |
|
|
|
|
| def parse_args() -> argparse.Namespace: |
| parser = argparse.ArgumentParser(description=__doc__) |
| parser.add_argument("--checkpoint", type=Path, required=True) |
| parser.add_argument("--config", type=Path, default=None) |
| parser.add_argument( |
| "--allow-legacy-pt", |
| action="store_true", |
| help="Explicitly allow a legacy .pt checkpoint through weights_only=True.", |
| ) |
| parser.add_argument("--input", type=Path, required=True) |
| parser.add_argument("--output", type=Path, required=True) |
| parser.add_argument("--image-field", default="images") |
| parser.add_argument("--workers", type=int, default=2) |
| return parser.parse_args() |
|
|
|
|
| def main() -> int: |
| args = parse_args() |
| checkpoint_path = args.checkpoint.expanduser().resolve() |
| input_path = args.input.expanduser().resolve() |
| output_path = args.output.expanduser().resolve() |
| artifact = load_point_detector_artifact( |
| checkpoint_path, |
| config_path=args.config, |
| device="cpu", |
| allow_legacy_pt=args.allow_legacy_pt, |
| ) |
| config = artifact.model_config |
| model = build_model( |
| False, |
| int(config["min_size"]), |
| int(config["max_size"]), |
| int(config.get("trainable_backbone_layers", 6)), |
| str(config.get("architecture", "mobilenet")), |
| ) |
| model.load_state_dict(artifact.state_dict, strict=True) |
| device = torch.device("cuda" if torch.cuda.is_available() else "cpu") |
| model.to(device).eval() |
|
|
| records = load_prediction_records(input_path, args.image_field) |
| loader = DataLoader( |
| PointSymbolDataset(records, augment=False), |
| batch_size=1, |
| shuffle=False, |
| num_workers=args.workers, |
| collate_fn=collate, |
| pin_memory=device.type == "cuda", |
| persistent_workers=args.workers > 0, |
| ) |
| output_path.parent.mkdir(parents=True, exist_ok=True) |
| with output_path.open("w") as output_file, torch.inference_mode(): |
| for record, (images, _targets) in zip(records, loader, strict=True): |
| prediction = model([images[0].to(device)])[0] |
| boxes = prediction["boxes"].detach().cpu().tolist() |
| scores = prediction["scores"].detach().cpu().tolist() |
| row = { |
| "schema": "cartolegend_point_detector_predictions_v1", |
| "image": portable_path(Path(record["image"])), |
| "width": record["width"], |
| "height": record["height"], |
| "detections": [ |
| { |
| "symbol_bbox": [round(float(value), 4) for value in box], |
| "score": round(float(score), 8), |
| } |
| for box, score in zip(boxes, scores, strict=True) |
| ], |
| } |
| output_file.write(json.dumps(row, sort_keys=True, separators=(",", ":")) + "\n") |
|
|
| manifest = { |
| "schema": "cartolegend_point_detector_prediction_manifest_v1", |
| "generated_utc": datetime.now(timezone.utc).isoformat(timespec="seconds"), |
| "checkpoint": portable_path(checkpoint_path), |
| "checkpoint_sha256": artifact.sha256, |
| "checkpoint_epoch": artifact.metadata.get("completed_epochs"), |
| "input": portable_path(input_path), |
| "input_sha256": sha256(input_path), |
| "output": portable_path(output_path), |
| "output_sha256": sha256(output_path), |
| "images": len(records), |
| "image_field": args.image_field, |
| "device": str(device), |
| "config": config, |
| } |
| manifest_path = output_path.with_suffix(output_path.suffix + ".manifest.json") |
| manifest_path.write_text(json.dumps(manifest, indent=2, sort_keys=True) + "\n") |
| print(json.dumps(manifest, indent=2, sort_keys=True)) |
| return 0 |
|
|
|
|
| if __name__ == "__main__": |
| raise SystemExit(main()) |
|
|
