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
# ...
demooo
最新推荐文章于 2024-09-05 11:10:46 发布