File size: 4,676 Bytes
9b92c75 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 | 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()
|