diff --git a/detect.py b/detect.py index fb4c867286bda49036473dce483d290c053b0f43..67721d8db916e871893a79be97d27e53fc6abe95 100755 --- a/detect.py +++ b/detect.py @@ -60,7 +60,7 @@ def detect( img = torch.from_numpy(img).unsqueeze(0).to(device) if ONNX_EXPORT: torch.onnx.export(model, img, 'weights/model.onnx', verbose=True) - return # ONNX export + return pred = model(img) pred = pred[pred[:, :, 4] > conf_thres] # remove boxes < threshold