How to filter the results from pytorch with NMS

Question

I have loaded a yolov7 model with pyTorch and I also get the result out of the model.
Now I wonder how to filter these results to prevent duplicate boxes. With OpenCV and onnx I know that this is possible with NMS, but how can I do this with pytorch?

Here is the code I have so far

import torch
from matplotlib import pyplot as plt
import numpy as np
import cv2

model = torch.hub.load("WongKinYiu/yolov7", 'custom', r"best.pt")
cap = cv2.VideoCapture(r"Test.mp4")

while (True):
    ret, img = cap.read()
    if ret ==True:

        
        results=model(img)

        boxes = results.xyxy[0].cpu().numpy()

        for i in boxes:
            x1,y1,x2,y2,a,c=i

            x1=int(x1)
            x2=int(x2)
            y1=int(y1)
            y2=int(y2)
            
            if a>0.4:
                BLUE   = (255,178,50)
                cv2.rectangle(img, (x1, y1), (x2 , y2), BLUE, 2)

            
        cv2.imshow("Output",img)
        if cv2.waitKey(1) & 0xFF == ord('q'):
            break

Asked By: Peter T.

||

Source

Answer 1

As suggested is possible to use nms from torchvision.ops :

from torchvision.ops import nms

res = results.xyxy[0].cpu()
boxes = res[:,:4] # x1,y1,x2,y2
scores = res[:,5] # confidence
filtered_boxes = nms(boxes, scores, iou_threshold = 0.2)

Buy yolov7 must have nms inside model, like yolov5. So you can control it behavior through model.iou = ... See sample: https://github.com/ultralytics/yolov5/issues/36

Answered By: Anton Ganichev

Answer 2

I didn’t know about the NMS. Thank you very much for your help.

Merry Christmas and a Happy New Year

Answered By: Peter T.

How to filter the results from pytorch with NMS

Question:

Answers: