Add all classes' confidence scores to the detection output

[yarinHayun]

I am trying to add all classes' confidence scores to the detection output.

I change the following lines in the NMS function:

output = [torch.zeros((0, nc + 6), device=prediction.device)] * bs
(The original line was: output = [torch.zeros((0, 6), device=prediction.device)] * bs)

x = torch.cat((box, conf, j.float(), x[:, 5:]), 1)[conf.view(-1) > conf_thres]
(The original line was: x = torch.cat((box, conf, j.float()), 1)[conf.view(-1) > conf_thres])

After these changes, the returned pred object from the NMS function is list of detections, on (n,nc+6) tensor per image [xyxy, conf, predicted class, all classes conf].

Then, I change the code in detect.py:
for *xyxy, conf, cls in reversed(det[:, :6]):
if save_txt:
line = (cls, *xywh, conf, *det[:, 6:][0]) if save_conf else (cls, *xywh)

But, when I use my trained model to detect some PIL image ( model(PILimg) ), I still get the original detection object [xyxy, confidence, predicted class, class name] without the all classes confidence score.
What do I need to do to get a detection object that includes all classes' confidence scores?

[faizch-ai]

Hello, were you able to solve this?

[kalfasyan]

Also interested in this

[AkkuTUD]

Also interested in this