yixionghuo commited on
Commit
5c271c7
1 Parent(s): 4b68df4

Upload onnx_inference.py

Browse files

Fix the bug about no object detected.

Files changed (1) hide show
  1. onnx_inference.py +6 -3
onnx_inference.py CHANGED
@@ -141,7 +141,10 @@ if __name__ == '__main__':
141
  det = pred[0]
142
  im0 = img0.copy()
143
 
144
- if len(det):
 
 
 
145
  # Rescale boxes from imgsz to im0 size
146
  det[:, :4] = scale_coords(img.shape[2:], det[:, :4], im0.shape).round()
147
 
@@ -150,5 +153,5 @@ if __name__ == '__main__':
150
  label = '%s %.2f' % (names[int(cls)], conf)
151
  plot_one_box(xyxy, im0, label=label, color=colors[int(cls)])
152
 
153
- # Stream results
154
- cv2.imwrite(new_path, im0)
 
141
  det = pred[0]
142
  im0 = img0.copy()
143
 
144
+ if det is None:
145
+ print('No objects detected!')
146
+
147
+ elif len(det):
148
  # Rescale boxes from imgsz to im0 size
149
  det[:, :4] = scale_coords(img.shape[2:], det[:, :4], im0.shape).round()
150
 
 
153
  label = '%s %.2f' % (names[int(cls)], conf)
154
  plot_one_box(xyxy, im0, label=label, color=colors[int(cls)])
155
 
156
+ # Stream results
157
+ cv2.imwrite(new_path, im0)