帮我写yolov5测试的时候画质心点的代码
时间: 2023-08-28 20:20:19 浏览: 149
好的,以下是画出YOLOv5检测出的目标框和质心点的代码示例:
```python
import cv2
import torch
from models.experimental import attempt_load
from utils.general import non_max_suppression, scale_coords
from utils.plots import plot_one_box
# 加载YOLOv5模型
weights_path = 'path/to/weights.pt'
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model = attempt_load(weights_path, map_location=device)
# 定义类别名称和颜色
class_names = ['class1', 'class2', 'class3'] # 替换成你的类别名称
colors = [[0, 255, 0], [0, 0, 255], [255, 0, 0]] # 替换成你的类别对应的颜色
# 读取测试图片
img_path = 'path/to/test/img.jpg'
img = cv2.imread(img_path)
# 对图片进行预处理
img_size = 640 # 替换成你的模型输入尺寸
img = cv2.resize(img, (img_size, img_size))
img = img[..., ::-1].transpose(2, 0, 1) # BGR to RGB,HWC to CHW
img = torch.from_numpy(img).to(device).float() / 255.0
img = img.unsqueeze(0)
# 进行目标检测
pred = model(img)[0]
pred = non_max_suppression(pred, conf_thres=0.5, iou_thres=0.5)[0] # 替换成你的目标检测参数
# 画出目标框和质心点
if pred is not None:
pred[:, :4] = scale_coords(img.shape[2:], pred[:, :4], img.shape[2:]).round()
for det in pred:
x1, y1, x2, y2, conf, cls = det.tolist()
label = class_names[int(cls)]
color = colors[int(cls)]
plot_one_box([x1, y1, x2, y2], img[0], color=color, label=label)
cx, cy = (x1 + x2) // 2, (y1 + y2) // 2
cv2.circle(img[0], (cx, cy), 3, color, -1)
# 显示结果
img = (img[0].cpu().numpy() * 255.0).astype('uint8').transpose(1, 2, 0)[..., ::-1] # CHW to HWC,RGB to BGR
cv2.imshow('result', img)
cv2.waitKey()
cv2.destroyAllWindows()
```
在上面的代码中,`plot_one_box` 函数用于画出目标框,`cv2.circle` 函数用于画出质心点。你需要替换代码中的 `weights_path`、`class_names`、`colors`、`img_path`、`img_size` 和目标检测参数等参数,以适应你的应用场景。
阅读全文