OpenCV DNN模塊推理YOLOv5 ONNX模型方法

這篇具有很好參考價(jià)值的文章主要介紹了OpenCV DNN模塊推理YOLOv5 ONNX模型方法。希望對大家有所幫助。如果存在錯(cuò)誤或未考慮完全的地方，請大家不吝賜教，您也可以點(diǎn)擊"舉報(bào)違法"按鈕提交疑問。

概述

本文檔主要描述python平臺，使用opencv-python深度神經(jīng)網(wǎng)絡(luò)模塊dnn,推理YOLOv5模型的方法。

文檔主要包含以下內(nèi)容：

opencv-python模塊的安裝
YOLOv5模型格式的說明
ONNX格式模型的加載
圖片數(shù)據(jù)的預(yù)處理
模型推理
推理結(jié)果后處理，包括NMS,cxcywh坐標(biāo)轉(zhuǎn)換為xyxy坐標(biāo)等
關(guān)鍵方法的調(diào)用與參數(shù)說明
完整的示例代碼

1. 環(huán)境部署

YOLOv5算法ONNX模型獲取

可通過官方鏈接下載YOLOv5的官方預(yù)訓(xùn)練模型，模型格式為pt.下載鏈接
YOLOv5官方項(xiàng)目提供了pt格式模型轉(zhuǎn)換為ONNX格式模型的腳本，項(xiàng)目鏈接

模型導(dǎo)出指令：

python export --weights yolov5s.pt --include onnx

注：導(dǎo)出文件執(zhí)行指令所需環(huán)境安裝配置參考官方項(xiàng)目README文檔即可，不在贅述。

opencv-python模塊安裝

創(chuàng)建虛擬環(huán)境并激活

conda create -n opencv python=3.8 -y
conda activate opencv

pip安裝opencv-python模塊
```
pip install opencv-python
```
注: 通過pip安裝opencv-python模塊時(shí)，默認(rèn)安裝僅支持CPU推理，如需支持GPU推理，需從源碼編譯安裝，具體安裝方法較復(fù)雜，這里不在贅述。

2.關(guān)鍵代碼

2.1 模型加載

opencv-python模塊提供了readNetFromONNX方法，用于加載ONNX格式模型。

import cv2
cv2.dnn.readNetFromONNX(model_path)

2.2 圖片數(shù)據(jù)預(yù)處理

數(shù)據(jù)預(yù)處理步驟包括resize，歸一化，顏色通道轉(zhuǎn)換，NCWH維度轉(zhuǎn)換等。

resize之前，有一個(gè)非常常用的trick來處理非方形的圖片，即計(jì)算圖形的最長邊，以此最長邊為基礎(chǔ)，創(chuàng)建一個(gè)正方形，并將原圖形放置到左上角，剩余部分用黑色填充，這樣做的好處是，不會(huì)改變原圖形的長寬比，同時(shí)也不會(huì)改變原圖形的內(nèi)容。

 # image preprocessing, the trick is to make the frame to be a square but not twist the image
row, col, _ = frame.shape  # get the row and column of the origin frame array
_max = max(row, col)  # get the max value of row and column
input_image = np.zeros((_max, _max, 3), dtype=np.uint8)  # create a new array with the max value
input_image[:row, :col, :] = frame  # paste the original frame  to make the input_image to be a square

完成圖片的填充后,繼續(xù)執(zhí)行resize，歸一化，顏色通道轉(zhuǎn)換等操作。

blob = cv2.dnn.blobFromImage(image, scalefactor=1 / 255.0, size=(640,640), swapRB=True, crop=False)

image: 輸入圖片數(shù)據(jù),numpy.ndarray格式,shape為(H,W,C),Channel順序?yàn)?code>BGR。
scalefactor: 圖片數(shù)據(jù)歸一化系數(shù)，一般為1/255.0。
size: 圖片resize尺寸，以模型的輸入要求為準(zhǔn)，這里是(640,640)。
swapRB: 是否交換顏色通道，即轉(zhuǎn)換BGR為RGB True表示交換，False表示不交換，由于opencv讀取圖片數(shù)據(jù)的顏色通道順序?yàn)?code>BGR，而YOLOv5模型的輸入要求為RGB，所以這里需要交換顏色通道。
crop: 是否裁剪圖片，False表示不裁剪。

blobFromImage函數(shù)返回四維Mat對象(NCHW dimensions order),數(shù)據(jù)的shape為(1,3,640,640)

2.3 模型推理

設(shè)置推理Backend和Target
```
model.setPreferableBackend(cv2.dnn.DNN_BACKEND_OPENCV)
model.setPreferableTarget(cv2.dnn.DNN_TARGET_CPU)
```
模型加載完成后，需要設(shè)置推理時(shí)的設(shè)備，一般情況下，推理設(shè)備為CPU，設(shè)置方法如下：
```
model.setPreferableBackend(cv2.dnn.DNN_BACKEND_OPENCV)
model.setPreferableTarget(cv2.dnn.DNN_TARGET_CPU)
```
當(dāng)然，若此時(shí)環(huán)境中的opencv-python模塊支持GPU推理，也可以設(shè)置為GPU推理，設(shè)置方法如下：
```
model.setPreferableBackend(cv2.dnn.DNN_BACKEND_CUDA)
model.setPreferableTarget(cv2.dnn.DNN_TARGET_CUDA)
```
注: 判斷opencv-python模塊是否支持GPU推理的方法如下：cv2.cuda.getCudaEnabledDeviceCount(),返回值大于0表示支持GPU推理，否則表示不支持。
設(shè)置模型輸入數(shù)據(jù)
```
model.setInput(blob)
```
blob為上一步數(shù)據(jù)預(yù)處理得到的數(shù)據(jù)。
調(diào)用模型前向傳播forward方法
```
outputs = model.forward()
```
outputs為模型推理的輸出，輸出格式為(1,25200,5+nc),25200為模型輸出的網(wǎng)格數(shù)量，5+nc為每個(gè)網(wǎng)格預(yù)測的5+nc個(gè)值，5為x,y,w,h,conf，nc為類別數(shù)量。

2.4 推理結(jié)果后處理

由于推理結(jié)果存在大量重疊的bbox，需要進(jìn)行NMS處理，后續(xù)根據(jù)每個(gè)bbox的置信度和用戶設(shè)定的置信度閾值進(jìn)行過濾，最終得到最終的bbox，和對應(yīng)的類別、置信度。

2.4.1 NMS

opencv-python模塊提供了NMSBoxes方法，用于進(jìn)行NMS處理。

cv2.dnn.NMSBoxes(bboxes, scores, score_threshold, nms_threshold, eta=None, top_k=None)

bboxes: bbox列表，shape為(N,4)，N為bbox數(shù)量，4為bbox的x,y,w,h。
scores: bbox對應(yīng)的置信度列表，shape為(N,1)，N為bbox數(shù)量。
score_threshold: 置信度閾值，小于該閾值的bbox將被過濾。
nms_threshold: NMS閾值

NMSBoxes函數(shù)返回值為bbox索引列表，shape為(M,)，M為bbox數(shù)量.

2.4.2 score_threshold過濾

根據(jù)NMS處理后的bbox索引列表，過濾置信度小于score_threshold的bbox。

2.4.3 bbox坐標(biāo)轉(zhuǎn)換與還原

YOLOv5模型輸出的bbox坐標(biāo)為cxcywh格式，需要轉(zhuǎn)換為xyxy格式，此外，由于之前對圖片進(jìn)行了resize操作，所以需要將bbox坐標(biāo)還原到原始圖片的尺寸。
轉(zhuǎn)換方法如下：

# 獲取原始圖片的尺寸(填充后)
image_width, image_height, _ = input_image.shape
# 計(jì)算縮放比
x_factor = image_width / INPUT_WIDTH  #  640
y_factor = image_height / INPUT_HEIGHT #  640 

# 將cxcywh坐標(biāo)轉(zhuǎn)換為xyxy坐標(biāo)
x1 = int((x - w / 2) * x_factor)
y1 = int((y - h / 2) * y_factor)
w = int(w * x_factor)
h = int(h * y_factor)
x2 = x1 + w
y2 = y1 + h

x1,y1,x2,y2即為bbox的xyxy坐標(biāo)。

3. 示例代碼(可運(yùn)行)

源代碼一共有兩份，其中一份是函數(shù)的拼接與調(diào)用，比較方便調(diào)試，另一份是封裝成類，方便集成到其他項(xiàng)目中。文章來源地址http://www.zghlxwxcb.cn/news/detail-604797.html

3.1 未封裝

"""
running the onnx model inference with opencv dnn module

"""
from typing import List

import cv2
import numpy as np
import time
from pathlib import Path


def build_model(model_path: str) -> cv2.dnn_Net:
    """
    build the model with opencv dnn module
    Args:
        model_path: the path of the model, the model should be in onnx format

    Returns:
        the model object
    """
    # check if the model file exists
    if not Path(model_path).exists():
        raise FileNotFoundError(f"model file {model_path} not found")
    model = cv2.dnn.readNetFromONNX(model_path)

    # check if the opencv-python in your environment supports cuda
    cuda_available = cv2.cuda.getCudaEnabledDeviceCount() > 0

    if cuda_available:  # if cuda is available, use cuda
        model.setPreferableBackend(cv2.dnn.DNN_BACKEND_CUDA)
        model.setPreferableTarget(cv2.dnn.DNN_TARGET_CUDA)
    else:  # if cuda is not available, use cpu
        model.setPreferableBackend(cv2.dnn.DNN_BACKEND_OPENCV)
        model.setPreferableTarget(cv2.dnn.DNN_TARGET_CPU)
    return model


def inference(image: np.ndarray, model: cv2.dnn_Net) -> np.ndarray:
    """
    inference the model with the input image
    Args:
        image: the input image in numpy array format, the shape should be (height, width, channel),
        the color channels should be in GBR order, like the original opencv image format
        model: the model object

    Returns:
        the output data of the model, the shape should be (1, 25200, nc+5), nc is the number of classes
    """
    # image preprocessing, include resize, normalization, channel swap like BGR to RGB, and convert to blob format
    # get a 4-dimensional Mat with NCHW dimensions order.
    blob = cv2.dnn.blobFromImage(image, 1 / 255.0, (INPUT_WIDTH, INPUT_HEIGHT), swapRB=True, crop=False)

    # the alternative way to get the blob
    # rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
    # input_image = cv2.resize(src=rgb, dsize=(INPUT_WIDTH, INPUT_HEIGHT))
    # blob_img = np.float32(input_image) / 255.0
    # input_x = blob_img.transpose((2, 0, 1))
    # blob = np.expand_dims(input_x, 0)

    if cv2.cuda.getCudaEnabledDeviceCount() > 0:
        model.setPreferableBackend(cv2.dnn.DNN_BACKEND_CUDA)
        model.setPreferableTarget(cv2.dnn.DNN_TARGET_CUDA)
    else:
        model.setPreferableBackend(cv2.dnn.DNN_BACKEND_OPENCV)
        model.setPreferableTarget(cv2.dnn.DNN_TARGET_CPU)

    # set the input data
    model.setInput(blob)

    start = time.perf_counter()
    # inference
    outs = model.forward()

    end = time.perf_counter()

    print("inference time: ", end - start)

    # the shape of the output data is (1, 25200, nc+5), nc is the number of classes
    return outs


def xywh_to_xyxy(bbox_xywh, image_width, image_height):
    """
    Convert bounding box coordinates from (center_x, center_y, width, height) to (x_min, y_min, x_max, y_max) format.

    Parameters:
        bbox_xywh (list or tuple): Bounding box coordinates in (center_x, center_y, width, height) format.
        image_width (int): Width of the image.
        image_height (int): Height of the image.

    Returns:
        tuple: Bounding box coordinates in (x_min, y_min, x_max, y_max) format.
    """
    center_x, center_y, width, height = bbox_xywh
    x_min = max(0, int(center_x - width / 2))
    y_min = max(0, int(center_y - height / 2))
    x_max = min(image_width - 1, int(center_x + width / 2))
    y_max = min(image_height - 1, int(center_y + height / 2))
    return x_min, y_min, x_max, y_max


def wrap_detection(
        input_image: np.ndarray,
        output_data: np.ndarray,
        labels: List[str],
        confidence_threshold: float = 0.6
) -> (List[int], List[float], List[List[int]]):
    # the shape of the output_data is (25200,5+nc),
    # the first 5 elements are [x, y, w, h, confidence], the rest are prediction scores of each class

    image_width, image_height, _ = input_image.shape
    x_factor = image_width / INPUT_WIDTH
    y_factor = image_height / INPUT_HEIGHT

    # transform the output_data[:, 0:4] from (x, y, w, h) to (x_min, y_min, x_max, y_max)

    indices = cv2.dnn.NMSBoxes(output_data[:, 0:4].tolist(), output_data[:, 4].tolist(), 0.6, 0.4)

    raw_boxes = output_data[:, 0:4][indices]
    raw_confidences = output_data[:, 4][indices]
    raw_class_prediction_probabilities = output_data[:, 5:][indices]

    criteria = raw_confidences > confidence_threshold
    raw_class_prediction_probabilities = raw_class_prediction_probabilities[criteria]
    raw_boxes = raw_boxes[criteria]
    raw_confidences = raw_confidences[criteria]

    bounding_boxes, confidences, class_ids = [], [], []
    for class_prediction_probability, box, confidence in zip(raw_class_prediction_probabilities, raw_boxes,
                                                             raw_confidences):
        #
        # find the least and most probable classes' indices and their probabilities
        # min_val, max_val, min_loc, mac_loc = cv2.minMaxLoc(class_prediction_probability)
        most_probable_class_index = np.argmax(class_prediction_probability)
        label = labels[most_probable_class_index]
        confidence = float(confidence)

        # bounding_boxes.append(box)
        # confidences.append(confidence)
        # class_ids.append(most_probable_class_index)

        x, y, w, h = box
        left = int((x - 0.5 * w) * x_factor)
        top = int((y - 0.5 * h) * y_factor)
        width = int(w * x_factor)
        height = int(h * y_factor)
        bounding_box = [left, top, width, height]
        bounding_boxes.append(bounding_box)
        confidences.append(confidence)
        class_ids.append(most_probable_class_index)

    return class_ids, confidences, bounding_boxes

coco_class_names = ["person", "bicycle", "car", "motorcycle", "airplane", "bus", "train", "truck", "boat",
                    "traffic light", "fire hydrant", "stop sign", "parking meter", "bench", "bird", "cat",
                    "dog", "horse", "sheep", "cow", "elephant", "bear", "zebra", "giraffe", "backpack",
                    "umbrella", "handbag", "tie", "suitcase", "frisbee", "skis", "snowboard", "sports ball",
                    "kite", "baseball bat", "baseball glove", "skateboard", "surfboard", "tennis racket",
                    "bottle", "wine glass", "cup", "fork", "knife", "spoon", "bowl", "banana", "apple",
                    "sandwich", "orange", "broccoli", "carrot", "hot dog", "pizza", "donut", "cake", "chair",
                    "couch", "potted plant", "bed", "dining table", "toilet", "tv", "laptop", "mouse",
                    "remote", "keyboard", "cell phone", "microwave", "oven", "toaster", "sink",
                    "refrigerator", "book", "clock", "vase", "scissors", "teddy bear", "hair drier",
                    "toothbrush"]
# generate different colors for coco classes
colors = np.random.uniform(0, 255, size=(len(coco_class_names), 3))

INPUT_WIDTH = 640
INPUT_HEIGHT = 640
CONFIDENCE_THRESHOLD = 0.7
NMS_THRESHOLD = 0.45

def video_detector(video_src):
    cap = cv2.VideoCapture(video_src)

    # 3. inference and show the result in a loop
    while cap.isOpened():
        success, frame = cap.read()
        start = time.perf_counter()
        if not success:
            break
        # image preprocessing, the trick is to make the frame to be a square but not twist the image
        row, col, _ = frame.shape  # get the row and column of the origin frame array
        _max = max(row, col)  # get the max value of row and column
        input_image = np.zeros((_max, _max, 3), dtype=np.uint8)  # create a new array with the max value
        input_image[:row, :col, :] = frame  # paste the original frame  to make the input_image to be a square
        # inference
        output_data = inference(input_image, net)  # the shape of output_data is (1, 25200, 85)

        # 4. wrap the detection result
        class_ids, confidences, boxes = wrap_detection(input_image, output_data[0], coco_class_names)

        # 5. draw the detection result on the frame
        for (class_id, confidence, box) in zip(class_ids, confidences, boxes):
            color = colors[int(class_id) % len(colors)]
            label = coco_class_names[int(class_id)]

            xmin, ymin, width, height = box
            cv2.rectangle(frame, (xmin, ymin), (xmin + width, ymin + height), color, 2)
            # cv2.rectangle(frame, box, color, 2)
            # cv2.rectangle(frame, [box[0], box[1], box[2], box[3]], color, thickness=2)

            # cv2.rectangle(frame, (box[0], box[1] - 20), (box[0] + 100, box[1]), color, -1)
            cv2.putText(frame, str(label), (box[0], box[1] - 5), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 0), 2)
        finish = time.perf_counter()
        FPS = round(1.0 / (finish - start), 2)
        cv2.putText(frame, str(FPS), (10, 50), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 0), 2)
        # 6. show the frame
        cv2.imshow("frame", frame)

        # 7. press 'q' to exit
        if cv2.waitKey(1) == ord('q'):
            break

    # 8. release the capture and destroy all windows
    cap.release()
    cv2.destroyAllWindows()

if __name__ == '__main__':
    # there are 4 steps to use opencv dnn module to inference onnx model exported by yolov5 and show the result

    # 1. load the model
    model_path = Path("weights/yolov5s.onnx")
    net = build_model(str(model_path))
    # 2. load the video capture
    # video_source = 0
    video_source = 'rtsp://admin:aoto12345@192.168.8.204:554/h264/ch1/main/av_stream'
    video_detector(video_source)

    exit(0)

3.2 封裝成類調(diào)用

from typing import List

import onnx
from torchvision import transforms

from torchvision.ops import nms,box_convert
import cv2
import time
import numpy as np
import onnxruntime as ort
import torch

INPUT_WIDTH = 640
INPUT_HEIGHT = 640

def wrap_detection(
        input_image: np.ndarray,
        output_data: np.ndarray,
        labels: List[str],
        confidence_threshold: float = 0.6
) -> (List[int], List[float], List[List[int]]):
    # the shape of the output_data is (25200,5+nc),
    # the first 5 elements are [x, y, w, h, confidence], the rest are prediction scores of each class

    image_width, image_height, _ = input_image.shape
    x_factor = image_width / INPUT_WIDTH
    y_factor = image_height / INPUT_HEIGHT

    # transform the output_data[:, 0:4] from (x, y, w, h) to (x_min, y_min, x_max, y_max)
    # output_data[:, 0:4] = np.apply_along_axis(xywh_to_xyxy, 1, output_data[:, 0:4], image_width, image_height)

    nms_start = time.perf_counter()
    indices = cv2.dnn.NMSBoxes(output_data[:, 0:4].tolist(), output_data[:, 4].tolist(), 0.6, 0.4)
    nms_finish = time.perf_counter()
    print(f"nms time: {nms_finish - nms_start}")
    # print(indices)
    raw_boxes = output_data[:, 0:4][indices]
    raw_confidences = output_data[:, 4][indices]
    raw_class_prediction_probabilities = output_data[:, 5:][indices]

    criteria = raw_confidences > confidence_threshold
    raw_class_prediction_probabilities = raw_class_prediction_probabilities[criteria]
    raw_boxes = raw_boxes[criteria]
    raw_confidences = raw_confidences[criteria]

    bounding_boxes, confidences, class_ids = [], [], []
    for class_prediction_probability, box, confidence in zip(raw_class_prediction_probabilities, raw_boxes,
                                                             raw_confidences):
        #
        # find the least and most probable classes' indices and their probabilities
        # min_val, max_val, min_loc, mac_loc = cv2.minMaxLoc(class_prediction_probability)
        most_probable_class_index = np.argmax(class_prediction_probability)
        label = labels[most_probable_class_index]
        confidence = float(confidence)

        # bounding_boxes.append(box)
        # confidences.append(confidence)
        # class_ids.append(most_probable_class_index)

        x, y, w, h = box
        left = int((x - 0.5 * w) * x_factor)
        top = int((y - 0.5 * h) * y_factor)
        width = int(w * x_factor)
        height = int(h * y_factor)
        bounding_box = [left, top, width, height]
        bounding_boxes.append(bounding_box)
        confidences.append(confidence)
        class_ids.append(most_probable_class_index)

    return class_ids, confidences, bounding_boxes


coco_class_names = ["person", "bicycle", "car", "motorcycle", "airplane", "bus", "train", "truck", "boat",
                    "traffic light", "fire hydrant", "stop sign", "parking meter", "bench", "bird", "cat",
                    "dog", "horse", "sheep", "cow", "elephant", "bear", "zebra", "giraffe", "backpack",
                    "umbrella", "handbag", "tie", "suitcase", "frisbee", "skis", "snowboard", "sports ball",
                    "kite", "baseball bat", "baseball glove", "skateboard", "surfboard", "tennis racket",
                    "bottle", "wine glass", "cup", "fork", "knife", "spoon", "bowl", "banana", "apple",
                    "sandwich", "orange", "broccoli", "carrot", "hot dog", "pizza", "donut", "cake", "chair",
                    "couch", "potted plant", "bed", "dining table", "toilet", "tv", "laptop", "mouse",
                    "remote", "keyboard", "cell phone", "microwave", "oven", "toaster", "sink",
                    "refrigerator", "book", "clock", "vase", "scissors", "teddy bear", "hair drier",
                    "toothbrush"]

colors = np.random.uniform(0, 255, size=(len(coco_class_names), 3))
if __name__ == '__main__':
    # Load the model
    model_path = "weights/yolov5s.onnx"
    onnx_model = onnx.load(model_path)
    onnx.checker.check_model(onnx_model)
    session = ort.InferenceSession(model_path, providers=['CUDAExecutionProvider',"CPUExecutionProvider"])
    capture = cv2.VideoCapture(0)

    trans = transforms.Compose([
        transforms.Resize((640, 640)),
        transforms.ToTensor()
    ])

    from PIL import Image

    while capture.isOpened():
        success, frame = capture.read()
        start = time.perf_counter()
        if not success:
            break
        rows, cols, channels = frame.shape
        # Preprocessing
        max_size = max(rows, cols)
        input_image = np.zeros((max_size, max_size, 3), dtype=np.uint8)
        input_image[:rows, :cols, :] = frame
        input_image = cv2.cvtColor(input_image, cv2.COLOR_BGR2RGB)

        inputs = trans(Image.fromarray(input_image))
        inputs = inputs.unsqueeze(0)
        print(inputs.shape)
        # inputs.to('cuda')
        ort_inputs = {session.get_inputs()[0].name: inputs.numpy()}
        ort_outs = session.run(None, ort_inputs)
        out_prob = ort_outs[0][0]
        print(out_prob.shape)

        scores = out_prob[:, 4] # Confidence scores are in the 5th column (0-indexed)
        class_ids = out_prob[:, 5:].argmax(axis=1)  # Class labels are from the 6th column onwards
        bounding_boxes_xywh = out_prob[:, :4]  # Bounding boxes in cxcywh format

        # Filter out boxes based on confidence threshold
        confidence_threshold = 0.7
        mask = scores >= confidence_threshold
        class_ids = class_ids[mask]
        bounding_boxes_xywh = bounding_boxes_xywh[mask]
        scores = scores[mask]

        bounding_boxes_xywh = torch.tensor(bounding_boxes_xywh, dtype=torch.float32)

        # Convert bounding boxes from xywh to xyxy format
        bounding_boxes_xyxy = box_convert(bounding_boxes_xywh, in_fmt='cxcywh', out_fmt='xyxy')

        # Perform Non-Maximum Suppression to filter candidate boxes


        scores = torch.tensor(scores, dtype=torch.float32)
        bounding_boxes_xyxy.to('cuda')
        scores.to('cuda')
        nms_start = time.perf_counter()
        keep_indices = nms(bounding_boxes_xyxy, scores, 0.4)
        nms_end = time.perf_counter()
        print(f"NMS took {nms_end - nms_start} seconds")
        class_ids = class_ids[keep_indices]
        confidences = scores[keep_indices]
        bounding_boxes = bounding_boxes_xyxy[keep_indices]

        # class_ids, confidences, bounding_boxes = wrap_detection(input_image, out_prob[0], coco_class_names, 0.6)
        # break

        for i in range(len(keep_indices)):
            try:
                class_id = class_ids[i]
            except IndexError as e:
                print(e)
                print(class_ids,i, len(keep_indices))
                break
            confidence = confidences[i]
            box = bounding_boxes[i]
            color = colors[int(class_id) % len(colors)]
            label = coco_class_names[int(class_id)]

            # cv2.rectangle(frame, box, color, 2)

            print(type(box), box[0], box[1], box[2], box[3], box)
            xmin, ymin, xmax, ymax = int(box[0]), int(box[1]), int(box[2]), int(box[3])
            cv2.rectangle(frame, (xmin, ymin), (xmax, ymax), color, 2)
            # cv2.rectangle(frame, box, color, 2)
            # cv2.rectangle(frame, [box[0], box[1], box[2], box[3]], color, thickness=2)

            cv2.rectangle(frame, (xmin, ymin - 20), (xmin + 100, ymin), color, -1)
            cv2.putText(frame, str(label), (xmin, ymin - 5), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 0), 2)

        finish = time.perf_counter()
        FPS = round(1.0 / (finish - start), 2)
        cv2.putText(frame, f"FPS: {str(FPS)}", (10, 50), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 0), 2)
        # 6. show the frame
        cv2.imshow("frame", frame)

        # 7. press 'q' to exit
        if cv2.waitKey(1) == ord('q'):
            break

    # 8. release the capture and destroy all windows
    capture.release()
    cv2.destroyAllWindows()

    exit(0)