zjc
/
yolov5
1
0
Fork 0
Code Issues Pull Requests Releases Wiki Activity

update test.py

This commit is contained in:
Glenn Jocher 2020-07-18 11:31:22 -07:00
parent 6c01a5b8f5
commit b40852d5a5
1 changed files with 5 additions and 6 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

11
test.py
View File

@ -126,13 +126,13 @@ def test(data,
# Append to pycocotools JSON dictionary # Append to pycocotools JSON dictionary
if save_json: if save_json:
# [{"image_id": 42, "category_id": 18, "bbox": [258.15, 41.29, 348.26, 243.78], "score": 0.236}, ... # [{"image_id": 42, "category_id": 18, "bbox": [258.15, 41.29, 348.26, 243.78], "score": 0.236}, ...
image_id = int(Path(paths[si]).stem.split('_')[-1]) image_id = Path(paths[si]).stem
box = pred[:, :4].clone() # xyxy box = pred[:, :4].clone() # xyxy
scale_coords(img[si].shape[1:], box, shapes[si][0], shapes[si][1]) # to original shape scale_coords(img[si].shape[1:], box, shapes[si][0], shapes[si][1]) # to original shape
box = xyxy2xywh(box) # xywh box = xyxy2xywh(box) # xywh
box[:, :2] -= box[:, 2:] / 2 # xy center to top-left corner box[:, :2] -= box[:, 2:] / 2 # xy center to top-left corner
for p, b in zip(pred.tolist(), box.tolist()): for p, b in zip(pred.tolist(), box.tolist()):
jdict.append({'image_id': image_id, jdict.append({'image_id': int(image_id) if image_id.isnumeric() else image_id,
'category_id': coco91class[int(p[5])], 'category_id': coco91class[int(p[5])],
'bbox': [round(x, 3) for x in b], 'bbox': [round(x, 3) for x in b],
'score': round(p[4], 5)}) 'score': round(p[4], 5)})
@ -200,8 +200,7 @@ def test(data,
print('Speed: %.1f/%.1f/%.1f ms inference/NMS/total per %gx%g image at batch-size %g' % t) print('Speed: %.1f/%.1f/%.1f ms inference/NMS/total per %gx%g image at batch-size %g' % t)
# Save JSON # Save JSON
if save_json and map50 and len(jdict): if save_json and len(jdict):
imgIds = [int(Path(x).stem.split('_')[-1]) for x in dataloader.dataset.img_files]
f = 'detections_val2017_%s_results.json' % \ f = 'detections_val2017_%s_results.json' % \
(weights.split(os.sep)[-1].replace('.pt', '') if isinstance(weights, str) else '') # filename (weights.split(os.sep)[-1].replace('.pt', '') if isinstance(weights, str) else '') # filename
print('\nCOCO mAP with pycocotools... saving %s...' % f) print('\nCOCO mAP with pycocotools... saving %s...' % f)
@ -212,6 +211,7 @@ def test(data,
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]
cocoGt = COCO(glob.glob('../coco/annotations/instances_val*.json')[0]) # initialize COCO ground truth api cocoGt = COCO(glob.glob('../coco/annotations/instances_val*.json')[0]) # initialize COCO ground truth api
cocoDt = cocoGt.loadRes(f) # initialize COCO pred api cocoDt = cocoGt.loadRes(f) # initialize COCO pred api
cocoEval = COCOeval(cocoGt, cocoDt, 'bbox') cocoEval = COCOeval(cocoGt, cocoDt, 'bbox')
@ -221,8 +221,7 @@ def test(data,
cocoEval.summarize() cocoEval.summarize()
map, map50 = cocoEval.stats[:2] # update results (mAP@0.5:0.95, mAP@0.5) map, map50 = cocoEval.stats[:2] # update results (mAP@0.5:0.95, mAP@0.5)
except: except:
print('WARNING: pycocotools must be installed with numpy==1.17 to run correctly. ' print('pycocotools not evaluated')
'See https://github.com/cocodataset/cocoapi/issues/356')
# Return results # Return results
model.float() # for training model.float() # for training