How to draw bounding box with xmin,xmax,ymin,ymax

Viewed 1588

I used Yolov5 as my detection modules and the input is based on a screen grab. The size of the screen is bounding_box = {'top': 340, 'left': 650, 'width': 350, 'height': 400}. I successfully get the xmin,ymin,xmax and ymax value from an object detection. But how can I used these value to draw a bounding box over my detection object?

Sample xmin,ymin,xmax and ymax value

         xmin       ymin        xmax        ymax  confidence  class    name
0  205.366241  68.419243  279.225586  272.266388    0.808757      0  person
1  134.863235  45.008553  208.445160  268.448029    0.777079      0  person

Current screen

Expected output: enter image description here

My full code:

# PyTorch Hub
import torch
import cv2 as cv
import numpy as np
import threading
import time

import numpy as np
import cv2
from PIL.ImageDraw import ImageDraw
from mss import mss
from PIL import Image
from tensorboard.summary.v1 import image
from torch import nn
from torch.utils.checkpoint import checkpoint

bounding_box = {'top': 340, 'left': 650, 'width': 350, 'height': 400}
# Model

#model = torch.hub.load('ultralytics/yolov5', 'custom', path='runs/train/exp/weights/last.pt')
#model = model.fuse().autoshape()
model = torch.hub.load('ultralytics/yolov5', 'yolov5s', pretrained=True).autoshape()

# Images
dir = 'https://ultralytics.com/images/'
#imgs = [dir + f for f in ('zidane.jpg', 'bus.jpg')]  # batch of images

sct = mss()

class ScreenGrab(threading.Thread):
    def __init__(self, id_manager):
        threading.Thread.__init__(self)
        self.id_manager = int(id_manager)

    def run(self):
        while True:
            global scr_img

            sct_img = sct.grab(bounding_box)
            scr_img = np.array(sct_img)
            cv.imshow('Testing', scr_img)

            if (cv2.waitKey(1) & 0xFF) == ord('q'):
                cv2.destroyAllWindows()
                break


class DetectionModule(threading.Thread):
    def __init__(self, id_manager):
        threading.Thread.__init__(self)
        self.id_manager = int(id_manager)

    def run(self):
        while True:
            global x, y, w, h
        # Inference
            results = model(scr_img)
            results.print()  # or .show(), .save()
            boxes = results.pandas().xyxy[0]

        #GET ALL COLUMN WITH ONLY CLASS == PERSON
            person = boxes.loc[boxes['class'] == 0]
            print(person)


def ScreenGrabMain():
    thread_id = ("0")
    led_index = 0
    thread_list = list()
    for objs in thread_id:
        thread = ScreenGrab(led_index)
        thread_list.append(thread)
        led_index += 1
    for thread in thread_list:
        thread.start()
        time.sleep(1)

def DetectionMain2():
    thread_id = ("0")
    led_index = 0
    thread_list = list()
    for objs in thread_id:
        thread = DetectionModule(led_index)
        thread_list.append(thread)
        led_index += 1
    for thread in thread_list:
        thread.start()
        time.sleep(1)

if __name__ == "__main__":
    ScreenGrabMain()
    DetectionMain2()
2 Answers

You can use cv2.rectangle in to draw a bounding box. An example is given below:


# cv2

image = cv2.rectangle(image, (xmin,ymin), (xmax,ymax), (0,255,0), 2)


#just convert each to an int type and crop giving the original image.

for index, row in results.pandas().xyxy[0].iterrows():
    print(row['xmin'], row['ymin'], row['xmax'], row['ymax'], row['confidence'])
    x1 = int(row['xmin'])
    y1 = int(row['ymin'])
    x2 = int(row['xmax'])
    y2 = int(row['ymax'])
    cropped_image = imgOrg[y1:y2, x1:x2]

    cv2.imwrite(imageName+'_'+str(counter)+'_cropped.png', cropped_image)
    counter = counter + 1 # initialize this to 1 before the for loop.
Related