207 lines
7.4 KiB
Python
207 lines
7.4 KiB
Python
#!/usr/bin/env python3
|
|
"""Compare local detection checkpoints on one non-protected YOLO validation split."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import hashlib
|
|
import json
|
|
import re
|
|
from datetime import UTC, datetime
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
|
|
def sha256_file(path: Path) -> str:
|
|
digest = hashlib.sha256()
|
|
with path.open("rb") as handle:
|
|
for chunk in iter(lambda: handle.read(1024 * 1024), b""):
|
|
digest.update(chunk)
|
|
return digest.hexdigest()
|
|
|
|
|
|
def safe_name(path: Path) -> str:
|
|
value = re.sub(r"[^a-z0-9]+", "-", path.stem.casefold()).strip("-")
|
|
return (value or path.stem.casefold())[:80]
|
|
|
|
|
|
def validation_images(dataset_yaml: Path) -> list[Path]:
|
|
import yaml
|
|
|
|
payload = yaml.safe_load(dataset_yaml.read_text(encoding="utf-8"))
|
|
if not isinstance(payload, dict):
|
|
raise ValueError("dataset YAML must contain a mapping")
|
|
source = payload.get("val")
|
|
root = Path(str(payload.get("path") or dataset_yaml.parent)).expanduser()
|
|
if not root.is_absolute():
|
|
root = (dataset_yaml.parent / root).resolve(strict=False)
|
|
if not isinstance(source, str) or not source.strip():
|
|
raise ValueError("dataset YAML requires one explicit val image directory")
|
|
directory = Path(source).expanduser()
|
|
if not directory.is_absolute():
|
|
directory = root / directory
|
|
if not directory.is_dir():
|
|
raise ValueError(f"validation image directory is unavailable: {directory}")
|
|
images = sorted(
|
|
path
|
|
for path in directory.iterdir()
|
|
if path.is_file() and path.suffix.casefold() in {".jpg", ".jpeg", ".png", ".tif", ".tiff"}
|
|
)
|
|
if not images:
|
|
raise ValueError("validation image directory is empty")
|
|
return images
|
|
|
|
|
|
def metric_value(metrics: Any, attribute: str) -> float:
|
|
value = getattr(metrics.box, attribute)
|
|
return float(value)
|
|
|
|
|
|
def background_detection_count(
|
|
model: Any,
|
|
images: list[Path],
|
|
*,
|
|
prefixes: tuple[str, ...],
|
|
confidence: float,
|
|
image_size: int,
|
|
device: str,
|
|
) -> tuple[int, int]:
|
|
selected = [path for path in images if path.stem.casefold().startswith(prefixes)]
|
|
if not selected:
|
|
raise ValueError("no validation images match the declared pure-background prefixes")
|
|
count = 0
|
|
for start in range(0, len(selected), 16):
|
|
results = model.predict(
|
|
[str(path) for path in selected[start : start + 16]],
|
|
conf=confidence,
|
|
iou=0.7,
|
|
max_det=1000,
|
|
imgsz=image_size,
|
|
device=device,
|
|
verbose=False,
|
|
)
|
|
count += sum(len(result.boxes) for result in results)
|
|
return len(selected), count
|
|
|
|
|
|
def main() -> int:
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("--dataset-yaml", type=Path, required=True)
|
|
parser.add_argument("--model", type=Path, action="append", required=True)
|
|
parser.add_argument("--output", type=Path, required=True)
|
|
parser.add_argument("--background-prefix", action="append", required=True)
|
|
parser.add_argument("--background-confidence", type=float, default=0.15)
|
|
parser.add_argument("--device", default="cuda:0")
|
|
parser.add_argument("--imgsz", type=int, default=640)
|
|
parser.add_argument("--batch", type=int, default=8)
|
|
args = parser.parse_args()
|
|
|
|
if args.output.exists():
|
|
parser.error(f"output already exists: {args.output}")
|
|
if not 0.0 < args.background_confidence < 1.0:
|
|
parser.error("--background-confidence must be between zero and one")
|
|
dataset_yaml = args.dataset_yaml.expanduser().resolve(strict=True)
|
|
images = validation_images(dataset_yaml)
|
|
prefixes = tuple(value.casefold() for value in args.background_prefix)
|
|
|
|
import torch
|
|
from ultralytics import YOLO
|
|
|
|
if not torch.cuda.is_available() or not args.device.casefold().startswith("cuda"):
|
|
raise SystemExit("checkpoint matrix requires the configured CUDA device")
|
|
|
|
output_root = args.output.parent / f"{args.output.stem}-runs"
|
|
rows: list[dict[str, Any]] = []
|
|
seen_hashes: set[str] = set()
|
|
for raw_model in args.model:
|
|
model_path = raw_model.expanduser().resolve(strict=True)
|
|
model_sha256 = sha256_file(model_path)
|
|
if model_sha256 in seen_hashes:
|
|
continue
|
|
seen_hashes.add(model_sha256)
|
|
row: dict[str, Any] = {
|
|
"model_path": str(model_path),
|
|
"model_sha256": model_sha256,
|
|
"size_bytes": model_path.stat().st_size,
|
|
}
|
|
model = None
|
|
try:
|
|
model = YOLO(str(model_path))
|
|
metrics = model.val(
|
|
data=str(dataset_yaml),
|
|
split="val",
|
|
imgsz=args.imgsz,
|
|
batch=args.batch,
|
|
workers=0,
|
|
device=args.device,
|
|
project=str(output_root),
|
|
name=safe_name(model_path),
|
|
exist_ok=False,
|
|
plots=False,
|
|
save_json=False,
|
|
verbose=False,
|
|
)
|
|
background_images, background_detections = background_detection_count(
|
|
model,
|
|
images,
|
|
prefixes=prefixes,
|
|
confidence=args.background_confidence,
|
|
image_size=args.imgsz,
|
|
device=args.device,
|
|
)
|
|
row.update(
|
|
{
|
|
"status": "ok",
|
|
"precision": metric_value(metrics, "mp"),
|
|
"recall": metric_value(metrics, "mr"),
|
|
"map50": metric_value(metrics, "map50"),
|
|
"map50_95": metric_value(metrics, "map"),
|
|
"pure_background_image_count": background_images,
|
|
"pure_background_detection_count": background_detections,
|
|
}
|
|
)
|
|
except Exception as exc: # preserve the complete attempted matrix
|
|
row.update({"status": "error", "error_type": type(exc).__name__, "error": str(exc)[:1000]})
|
|
finally:
|
|
del model
|
|
if torch.cuda.is_available():
|
|
torch.cuda.empty_cache()
|
|
rows.append(row)
|
|
print(json.dumps(row, sort_keys=True), flush=True)
|
|
|
|
successful = [row for row in rows if row["status"] == "ok"]
|
|
successful.sort(
|
|
key=lambda row: (
|
|
int(row["pure_background_detection_count"] == 0),
|
|
row["map50_95"],
|
|
row["map50"],
|
|
row["precision"],
|
|
row["recall"],
|
|
),
|
|
reverse=True,
|
|
)
|
|
payload = {
|
|
"schema_version": 1,
|
|
"generated_at": datetime.now(UTC).isoformat(),
|
|
"status": "ok" if successful else "failed",
|
|
"claim_boundary": "Non-protected validation ranking only; no test, challenge or promotion claim.",
|
|
"dataset_yaml": str(dataset_yaml),
|
|
"dataset_yaml_sha256": sha256_file(dataset_yaml),
|
|
"validation_image_count": len(images),
|
|
"pure_background_prefixes": list(prefixes),
|
|
"pure_background_confidence": args.background_confidence,
|
|
"device": args.device,
|
|
"torch_version": torch.__version__,
|
|
"candidate_count": len(rows),
|
|
"successful_candidate_count": len(successful),
|
|
"ranking": successful,
|
|
"attempts": rows,
|
|
}
|
|
args.output.parent.mkdir(parents=True, exist_ok=True)
|
|
args.output.write_text(json.dumps(payload, indent=2, sort_keys=True) + "\n", encoding="utf-8")
|
|
return 0 if successful else 2
|
|
|
|
|
|
if __name__ == "__main__":
|
|
raise SystemExit(main())
|