demooo

原创于 2023-03-27 19:44:36 发布 · 152 阅读

0 ·

CC 4.0 BY-SA版权

文章标签：

#python

Python 专栏收录该内容

4 篇文章

订阅专栏

该代码段展示了如何加载YOLOv5模型，对一张图片进行推理，应用非极大值抑制去除冗余框，并计算出物体类别和边界框信息。

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

import torch
import cv2
from models.experimental import attempt_load
from utils.general import non_max_suppression

# Load YOLOv5 model
model = attempt_load('yolov5s.pt', map_location=torch.device('cpu'))

# Load image
img = cv2.imread('test.jpg')

# Inference
results = model(img)

# Non-maximum suppression to remove redundant boxes
results = non_max_suppression(results, conf_thres=0.5, iou_thres=0.5)

# Get class labels and box coordinates
class_ids = []
boxes = []
for result in results:
    if result is not None:
        for detection in result:
            class_ids.append(detection[-1])
            boxes.append(detection[:4])

# Calculate class count
class_count = {}
for class_id in class_ids:
    if class_id in class_count:
        class_count[class_id] += 1
    else:
        class_count[class_id] = 1

# Calculate box information
box_info = []
for box in boxes:
    box_info.append((box[0], box[1], box[2] - box[0], box[3] - box[1]))

# Calculate using your specific formula
# ...