Skip to content

Instantly share code, notes, and snippets.

@AlessandroMondin

AlessandroMondin/prediction.py

Last active December 20, 2022 19:23
Show Gist options
  • Save AlessandroMondin/52f83f41580bdd0217330c4851d3ac4b to your computer and use it in GitHub Desktop.
Save AlessandroMondin/52f83f41580bdd0217330c4851d3ac4b to your computer and use it in GitHub Desktop.
Download ZIP
Raw
prediction.py
S = [8, 16, 32]
with torch.no_grad():
out = model(img)
boxes = cells_to_bboxes(out, model.head.anchors, S, list_output=False, is_pred=True)
boxes = non_max_suppression(boxes, iou_threshold=0.6, threshold=.25, max_detections=300)
plot_image(img[0].permute(1, 2, 0).to("cpu"), boxes[0])
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment