yixionghuo
commited on
Commit
•
5c271c7
1
Parent(s):
4b68df4
Upload onnx_inference.py
Browse filesFix the bug about no object detected.
- 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
|
|
|
|
|
|
|
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 |
-
|
154 |
-
|
|
|
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)
|