| from __future__ import annotations |
|
|
| import argparse |
| import json |
| from pathlib import Path |
|
|
| import torch |
| from torch.utils.data import DataLoader |
|
|
| from .boxes import box_cxcywh_to_xyxy |
| from .config import apply_overrides, load_config |
| from .data import build_dataset, detection_collate |
| from .model import build_model |
|
|
|
|
| @torch.inference_mode() |
| def evaluate_coco( |
| model, |
| data_loader, |
| device: torch.device, |
| output_file: str | Path, |
| confidence: float = 0.001, |
| max_detections: int = 300, |
| ) -> dict[str, float]: |
| try: |
| from pycocotools.coco import COCO |
| from pycocotools.cocoeval import COCOeval |
| except ImportError as error: |
| raise RuntimeError("COCO evaluation requires: pip install -e '.[coco]'") from error |
|
|
| model.eval() |
| input_size = model.spec.input_size |
| predictions = [] |
| label_to_category = data_loader.dataset.label_to_category |
| for images, targets in data_loader: |
| images = images.to(device, non_blocking=True) |
| outputs = model(images) |
| probabilities = outputs["pred_logits"].sigmoid() |
| normalized_boxes = box_cxcywh_to_xyxy(outputs["pred_boxes"]).clamp(0.0, 1.0) |
| for batch_index, target in enumerate(targets): |
| scores, labels = probabilities[batch_index].max(dim=-1) |
| count = min(max_detections, scores.numel()) |
| scores, indices = scores.topk(count) |
| keep = scores >= confidence |
| scores, indices = scores[keep], indices[keep] |
| labels = labels[indices] |
| boxes = normalized_boxes[batch_index, indices] * input_size |
| ratio, offset_x, offset_y = target["transform"].tolist() |
| original_height, original_width = target["original_size"].tolist() |
| boxes[:, [0, 2]] = (boxes[:, [0, 2]] - offset_x) / ratio |
| boxes[:, [1, 3]] = (boxes[:, [1, 3]] - offset_y) / ratio |
| boxes[:, [0, 2]].clamp_(0, original_width) |
| boxes[:, [1, 3]].clamp_(0, original_height) |
| boxes[:, 2:] -= boxes[:, :2] |
| for box, score, label in zip(boxes.cpu(), scores.cpu(), labels.cpu(), strict=True): |
| predictions.append( |
| { |
| "image_id": int(target["image_id"]), |
| "category_id": int(label_to_category[int(label)]), |
| "bbox": [round(float(value), 3) for value in box], |
| "score": float(score), |
| } |
| ) |
|
|
| output_file = Path(output_file) |
| output_file.parent.mkdir(parents=True, exist_ok=True) |
| with output_file.open("w", encoding="utf-8") as handle: |
| json.dump(predictions, handle) |
| ground_truth = COCO(str(data_loader.dataset.annotation_file)) |
| if not predictions: |
| return {"AP": 0.0, "AP50": 0.0, "AP75": 0.0, "APS": 0.0, "APM": 0.0, "APL": 0.0} |
| detections = ground_truth.loadRes(str(output_file)) |
| evaluator = COCOeval(ground_truth, detections, "bbox") |
| evaluator.params.imgIds = [int(item["id"]) for item in data_loader.dataset.images] |
| evaluator.evaluate() |
| evaluator.accumulate() |
| evaluator.summarize() |
| names = ("AP", "AP50", "AP75", "APS", "APM", "APL") |
| return {name: float(evaluator.stats[index]) for index, name in enumerate(names)} |
|
|
|
|
| def main() -> None: |
| parser = argparse.ArgumentParser(description="Evaluate ObjectModel-v1 on COCO") |
| parser.add_argument("--config", default="configs/objectmodel_v1.yaml") |
| parser.add_argument("--checkpoint", required=True) |
| parser.add_argument("--data-root", required=True) |
| parser.add_argument("--output", default="outputs/predictions.json") |
| parser.add_argument("--device", default="cuda" if torch.cuda.is_available() else "cpu") |
| parser.add_argument("--batch-size", type=int, default=8) |
| parser.add_argument("--workers", type=int, default=4) |
| parser.add_argument("--set", action="append", default=[]) |
| args = parser.parse_args() |
| config = apply_overrides(load_config(args.config), args.set) |
| model = build_model(config) |
| checkpoint = torch.load(args.checkpoint, map_location="cpu", weights_only=False) |
| model.load_state_dict(checkpoint.get("ema", checkpoint.get("model", checkpoint))) |
| device = torch.device(args.device) |
| model.to(device) |
| dataset = build_dataset(config, args.data_root, "val") |
| loader = DataLoader( |
| dataset, |
| batch_size=args.batch_size, |
| shuffle=False, |
| num_workers=args.workers, |
| pin_memory=device.type == "cuda", |
| collate_fn=detection_collate, |
| ) |
| print(json.dumps(evaluate_coco(model, loader, device, args.output), indent=2)) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
