File size: 3,605 Bytes
f4804d7 993959a f4804d7 | 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 | 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)
# id->(h,w)
id2size = {int(ex['image_id']): (ex['height'], ex['width']) for ex in val_ds}
results = [] # COCO result list
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()
# ev.stats: [AP50..AP75, APsmall, APmedium, APlarge, AR1, AR10, AR100, ARsmall, ARmedium, ARlarge]
ap50_95 = ev.stats[0]
ap50 = ev.stats[1]
ap75 = ev.stats[2]
ar100 = ev.stats[8]
prec = ev.eval['precision'] # T,R,K,A,M
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
|