updates
This commit is contained in:
parent
e23b1a3d73
commit
742908257a
|
@ -60,7 +60,7 @@ def detect(
|
||||||
img = torch.from_numpy(img).unsqueeze(0).to(device)
|
img = torch.from_numpy(img).unsqueeze(0).to(device)
|
||||||
if ONNX_EXPORT:
|
if ONNX_EXPORT:
|
||||||
torch.onnx.export(model, img, 'weights/model.onnx', verbose=True)
|
torch.onnx.export(model, img, 'weights/model.onnx', verbose=True)
|
||||||
return # ONNX export
|
return
|
||||||
pred = model(img)
|
pred = model(img)
|
||||||
pred = pred[pred[:, :, 4] > conf_thres] # remove boxes < threshold
|
pred = pred[pred[:, :, 4] > conf_thres] # remove boxes < threshold
|
||||||
|
|
||||||
|
|
Loading…
Reference in New Issue