How to extract all class probabilities for each bounding box?

[sssmost]

Hello,

I am trying to extract all class probabilities for each bb returned by YOLOv3 using detect.py. However, they are loss and it is only extracted one class per bb. How I could do it? They are loss in non_max_suppression process?

Thanks in advance

[github-actions]

Hello @sssmost, thank you for your interest in our work! Please visit our Custom Training Tutorial to get started, and see our Google Colab Notebook, Docker Image, and GCP Quickstart Guide for example environments.

If this is a bug report, please provide screenshots and minimum viable code to reproduce your issue, otherwise we can not help you.

[glenn-jocher]

@sssmost the inference output is here:

yolov3/detect.py

Lines 85 to 89 in a4721e9

	# Inference
	t1 = torch_utils.time_synchronized()
	pred = model(img)[0].float() if half else model(img)[0]
	t2 = torch_utils.time_synchronized()

[Lornatang]

@glenn-jocher
I try pull your new code, bit I get some problem.

in non_max_suppression
    x[i, :4] = torch.mm(weights.T, x[:, :4])
RuntimeError: expected dtype Float but got dtype Half ?

[glenn-jocher]

@Lornatang oh, that's interesting. This is part of new NMS. Can you supply code to reproduce? Is this running detect.py --half?

[Lornatang]

@glenn-jocher
Oh, no, It happened only when I was training, and it happened when I called test for the first time.

My command is

python3 train.py

[glenn-jocher]

@Lornatang I can't reproduce. I git cloned the repo in a new colab notebook and am unable to reproduce the error using:
python3 train.py --data coco16.data --nosave

[Lornatang]

@glenn-jocher
I'll try to see if my code is wrong with you. Thank you for your prompt reply.

[github-actions]

This issue is stale because it has been open 30 days with no activity. Remove Stale label or comment or this will be closed in 5 days.

[sssmost]

I found a solution. The probabilities are deleted inside # Apply NMS
pred = non_max_suppression(pred, opt.conf_thres, opt.iou_thres, classes=opt.classes, agnostic=opt.agnostic_nms)

I modified the call to return all probabilities for each bounding box in pred_complete:
pred, pred_complete = non_max_suppression(pred, opt.conf_thres, opt.iou_thres, classes=opt.classes, agnostic=opt.agnostic_nms)

And inside the function:
`def non_max_suppression(prediction, conf_thres=0.1, iou_thres=0.5, multi_cls=True, classes=None, agnostic=False):

    # Detections matrix nx6 (xyxy, conf, cls)
    if multi_cls:
        i, j = (pred[:, 5:] > conf_thres).nonzero().t()
        ######
        # Save entire predictions
        #Each pred_aux will have 30 values -> cx cy w h + 25 probabilities (in my case I have 25 classes)
        pred_aux = torch.cat((box[i], pred[i, 5].unsqueeze(1)), 1)
        for id_aux in range(len(pred[0])-6):
            pred_aux = torch.cat((pred_aux, pred[i, id_aux + 6].unsqueeze(1)), 1)
        ######
        pred = torch.cat((box[i], pred[i, j + 5].unsqueeze(1), j.float().unsqueeze(1)), 1) 
    else:  # best class only
        conf, j = pred[:, 5:].max(1)
        pred = torch.cat((box, conf.unsqueeze(1), j.float().unsqueeze(1)), 1)

    # Filter by class
    if classes:
        pred = pred[(j.view(-1, 1) == torch.tensor(classes, device=j.device)).any(1)]

    # Apply finite constraint
    if not torch.isfinite(pred).all():
        pred = pred[torch.isfinite(pred).all(1)]

    # Batched NMS
    if method == 'vision_batch':
        c = pred[:, 5] * 0 if agnostic else pred[:, 5]  # class-agnostic NMS
        #####
        index_bbs_keep_nms = torchvision.ops.boxes.batched_nms(pred[:, :4], pred[:, 4], c, iou_thres)
        output[image_i] = pred[index_bbs_keep_nms]
        output_complete[image_i] = pred_aux[index_bbs_keep_nms]
        ####
        continue

`
return output, output_complete

[yehadrabba]

Hi Ssmost,

Where did you define output_complete?

What did you do with the original rest of the code between your implementation (####) ?

thanks

[glenn-jocher]

@yehadrabba You're on the right track! One potential improvement is to make sure you include output_complete in the function's return statement. To do this, you should define the output_complete list before the loop or as an empty tensor if using PyTorch. Additionally, your modifications look good and are in line with YOLO's expected output. Great work! Keep it up!