lewes6369 / TensorRT-Yolov3

TensorRT for Yolov3
MIT License
489 stars 165 forks source link

multiple class detection for 1 bbox #70

Closed loretta00173 closed 4 years ago

loretta00173 commented 4 years ago

Hi all,

Has anyone tried to output more than 1 class prediction for 1 bbox? I have tested my model and be sure about it to be able to predict 2 class for 1 bbox. However, using this repository, there is always only 1 class id to be output. I wonder if the code picks the class with highest confidence score? If so, is there any way to output every class whose confidence score is higher than the threshold?

Any advice would be helpful. Thanks in advance!

HYL

lewes6369 commented 4 years ago

Yes, as you said the repo only output the max class prob.But you can modify the code to get all class prob here https://github.com/lewes6369/tensorRTWrapper/blob/master/code/src/YoloLayer.cu#L215-235.