| import torch, os, json |
| import numpy as np |
| from pycocotools.coco import COCO |
| from pycocotools.cocoeval import COCOeval |
|
|
| CLASS_NAMES = ["Photograph","Illustration","Map","Comics/Cartoon","Editorial Cartoon","Headline","Advertisement"] |
|
|
| def _build_gt(val_ds): |
| cats = [{'id': i+1, 'name': nm} for i, nm in enumerate(CLASS_NAMES)] |
| images = [] |
| anns = [] |
| ann_id = 1 |
| for ex in val_ds: |
| img_id = int(ex['image_id']) |
| images.append({'id': img_id, 'width': ex['width'], 'height': ex['height'], 'file_name': str(img_id)}) |
| objs = ex['objects'] |
| for obj in objs: |
| x, y, w, h = obj['bbox'] |
| anns.append({ |
| 'id': ann_id, |
| 'image_id': img_id, |
| 'category_id': int(obj['category_id']) + 1, |
| 'bbox': [float(x), float(y), float(w), float(h)], |
| 'area': float(obj['area']) if 'area' in obj else float(w*h), |
| 'iscrowd': int(obj['iscrowd']) if 'iscrowd' in obj else 0, |
| }) |
| ann_id += 1 |
| data = {'images': images, 'annotations': anns, 'categories': cats} |
| return data |
|
|
| def evaluate(model, processor, val_ds, val_dl, device): |
| gt_data = _build_gt(val_ds) |
| |
| id2size = {int(ex['image_id']): (ex['height'], ex['width']) for ex in val_ds} |
|
|
| results = [] |
| with torch.no_grad(): |
| for batch in val_dl: |
| pv = batch['pixel_values'].to(device) |
| pm = batch['pixel_mask'].to(device) |
| out = model(pixel_values=pv, pixel_mask=pm) |
| target_sizes = torch.tensor([list(id2size[int(i)]) for i in batch['img_id']], dtype=torch.int64) |
| post = processor.post_process_object_detection(out, target_sizes=target_sizes, threshold=0.0) |
| for bi, img_id in enumerate(batch['img_id']): |
| imid = int(img_id) |
| boxes = post[bi]['boxes'].cpu() |
| scores = post[bi]['scores'].cpu() |
| labels = post[bi]['labels'].cpu() |
| for score, lab, box in zip(scores, labels, boxes): |
| x0, y0, x1, y1 = box.tolist() |
| w = max(0.0, x1-x0); h = max(0.0, y1-y0) |
| results.append({'image_id': imid, |
| 'category_id': int(lab)+1, |
| 'bbox': [x0, y0, w, h], |
| 'score': float(score)}) |
|
|
| coco_gt = COCO() |
| coco_gt.dataset = gt_data |
| coco_gt.createIndex() |
| if len(results) == 0: |
| return {'mAP': 0.0} |
| coco_dt = coco_gt.loadRes(results) |
| imgids = [im['id'] for im in gt_data['images']] |
| ev = COCOeval(coco_gt, coco_dt, iouType='bbox') |
| ev.params.imgIds = imgids |
| ev.evaluate() |
| ev.accumulate() |
| ev.summarize() |
| |
| ap50_95 = ev.stats[0] |
| ap50 = ev.stats[1] |
| ap75 = ev.stats[2] |
| ar100 = ev.stats[8] |
| prec = ev.eval['precision'] |
| class_ap = {} |
| for i, clsnamed in enumerate(CLASS_NAMES): |
| p = prec[:,:,i,0,-1] |
| p = p[p > -1] |
| class_ap[clsnamed] = float(p.mean()) if p.size else 0.0 |
| for k,v in class_ap.items(): |
| print(f" AP {k}: {v:.4f}") |
| metrics = {'mean_AP_0.50_0.95': float(ap50_95), 'per_class_AP': class_ap, |
| 'AP_0.50': float(ap50), |
| 'AP_0.75': float(ap75), |
| 'AR_max_100': float(ar100), |
| 'num_predictions': len(results), |
| 'pred_file': 'predictions.json'} |
| return metrics |
|
|
