demooo

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
# ...


你可能感兴趣的:(Python,python)