Increase pycocotools robustness (#1396)

This commit is contained in:
Glenn Jocher 2020-11-14 13:09:04 +01:00 committed by GitHub
parent cf581305db
commit 9b0f6e33ee
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23

26
test.py
View File

@ -66,6 +66,7 @@ def test(data,
# Configure # Configure
model.eval() model.eval()
is_coco = data.endswith('coco.yaml') # is COCO dataset
with open(data) as f: with open(data) as f:
data = yaml.load(f, Loader=yaml.FullLoader) # model dict data = yaml.load(f, Loader=yaml.FullLoader) # model dict
check_dataset(data) # check check_dataset(data) # check
@ -240,24 +241,25 @@ def test(data,
# Save JSON # Save JSON
if save_json and len(jdict): if save_json and len(jdict):
w = Path(weights[0] if isinstance(weights, list) else weights).stem if weights is not None else '' # weights w = Path(weights[0] if isinstance(weights, list) else weights).stem if weights is not None else '' # weights
file = save_dir / f"detections_val2017_{w}_results.json" # predicted annotations file anno_json = glob.glob('../coco/annotations/instances_val*.json')[0] # annotations json
print('\nCOCO mAP with pycocotools... saving %s...' % file) pred_json = str(save_dir / f"{w}_predictions.json") # predictions json
with open(file, 'w') as f: print('\nEvaluating pycocotools mAP... saving %s...' % pred_json)
with open(pred_json, 'w') as f:
json.dump(jdict, f) json.dump(jdict, f)
try: # https://github.com/cocodataset/cocoapi/blob/master/PythonAPI/pycocoEvalDemo.ipynb try: # https://github.com/cocodataset/cocoapi/blob/master/PythonAPI/pycocoEvalDemo.ipynb
from pycocotools.coco import COCO from pycocotools.coco import COCO
from pycocotools.cocoeval import COCOeval from pycocotools.cocoeval import COCOeval
imgIds = [int(Path(x).stem) for x in dataloader.dataset.img_files] anno = COCO(anno_json) # init annotations api
cocoAnno = COCO(glob.glob('../coco/annotations/instances_val*.json')[0]) # initialize COCO annotations api pred = anno.loadRes(pred_json) # init predictions api
cocoPred = cocoAnno.loadRes(str(file)) # initialize COCO pred api eval = COCOeval(anno, pred, 'bbox')
cocoEval = COCOeval(cocoAnno, cocoPred, 'bbox') if is_coco:
cocoEval.params.imgIds = imgIds # image IDs to evaluate eval.params.imgIds = [int(Path(x).stem) for x in dataloader.dataset.img_files] # image IDs to evaluate
cocoEval.evaluate() eval.evaluate()
cocoEval.accumulate() eval.accumulate()
cocoEval.summarize() eval.summarize()
map, map50 = cocoEval.stats[:2] # update results (mAP@0.5:0.95, mAP@0.5) map, map50 = eval.stats[:2] # update results (mAP@0.5:0.95, mAP@0.5)
except Exception as e: except Exception as e:
print('ERROR: pycocotools unable to run: %s' % e) print('ERROR: pycocotools unable to run: %s' % e)