mirrors
/
mmyolo
mirror of https://github.com/open-mmlab/mmyolo.git
1
0
Fork 0
Code Issues Projects Releases Wiki Activity

fix image-demo.py empty class_name fatal (#615)

pull/612/head
Xin Li 2023-03-02 15:14:51 +08:00 committed by GitHub
parent 146cd930c5
commit 9e8f764399
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 7 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

9
projects/easydeploy/tools/image-demo.py
View File

@ -74,6 +74,7 @@ def main():
model.to(args.device)
cfg = Config.fromfile(args.config)
class_names = cfg.get('class_name')
test_pipeline = get_test_pipeline_cfg(cfg)
test_pipeline[0] = ConfigDict({'type': 'mmdet.LoadImageFromNDArray'})
@ -125,8 +126,12 @@ def main():
for (bbox, score, label) in zip(bboxes, scores, labels):
bbox = bbox.tolist()
color = colors[label]
label_name = cfg.get('class_name', {})[label]
name = f'cls:{label_name}_score:{score:0.4f}'
if class_names is not None:
label_name = class_names[label]
name = f'cls:{label_name}_score:{score:0.4f}'
else:
name = f'cls:{label}_score:{score:0.4f}'
cv2.rectangle(bgr, bbox[:2], bbox[2:], color, 2)
cv2.putText(