search

ultralytics / yolov5

YOLOv5 🚀 in PyTorch > ONNX > CoreML > TFLite
https://docs.ultralytics.com
GNU Affero General Public License v3.0
49.47k stars 16.07k forks source link

Exclude NMS from custom model #8583

Closed michelebechini closed 2 years ago

michelebechini commented 2 years ago

Search before asking

Question

Hi everyone!

I trained a yolov5 model for custom single object detection and for my project I would like to skip the NMS part and make the model to output all the predictions, to subsequently filter them with a custom algorithm that I am developing.

Could you please help me to find where I need to modify the source code to exclude NMS and get all the predictions as output?

Thanks in advance!!!

Additional

No response

glenn-jocher commented 2 years ago

@michelebechini YOLOv5 models don't include NMS. It's processed separately, i.e. see detect.py: https://github.com/ultralytics/yolov5/blob/f8722b4429e80f96be04b36e4efd84ce6583bfa1/detect.py#L120-L129

michelebechini commented 2 years ago

@glenn-jocher so basically if I comment out line 127 the model will return directly all the predictions, is it correct?

glenn-jocher commented 2 years ago

@michelebechini well there's nothing really to comment.

pred on L122 returns what you want I think. You can load one of these base models with PyTorch Hub also by setting autoshape=False.

model = torch.hub.load('ultralytics/yolov5', 'yolov5s', autoshape=False)  # or yolov5n - yolov5x6, custom
github-actions[bot] commented 2 years ago

👋 Hello, this issue has been automatically marked as stale because it has not had recent activity. Please note it will be closed if no further activity occurs.

Access additional YOLOv5 🚀 resources:

Access additional Ultralytics ⚡ resources:

Feel free to inform us of any other issues you discover or feature requests that come to mind in the future. Pull Requests (PRs) are also always welcomed!

Thank you for your contributions to YOLOv5 🚀 and Vision AI ⭐!