yolov5显示预测框的精确坐标的代码
时间: 2024-01-25 12:04:47 浏览: 29
以下是一个使用Python和OpenCV库来显示Yolov5预测框精确坐标的示例代码:
```python
import cv2
import torch
from models.experimental import attempt_load
from utils.general import non_max_suppression
# 加载模型
weights = 'yolov5s.pt' # 模型权重文件路径
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model = attempt_load(weights, map_location=device)
# 加载图像
image_path = 'image.jpg' # 图像文件路径
img = cv2.imread(image_path)
# 图像预处理
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) # 将图像从BGR格式转换为RGB格式
img = torch.from_numpy(img.transpose(2, 0, 1)).float().div(255.0).unsqueeze(0) # 转换为模型输入所需的张量格式
# 模型推理
pred = model(img.to(device))[0] # 进行模型推理
# 非最大抑制获取预测框
pred = non_max_suppression(pred, conf_thres=0.5, iou_thres=0.5)[0] # 使用自定义的非最大抑制方法获取预测框
# 显示预测框和精确坐标
if len(pred) > 0:
pred[:, :4] = pred[:, :4].clamp(min=0, max=1) # 将预测框的坐标限制在0~1的范围内
for x1, y1, x2, y2, conf, cls in pred:
x1 *= img.shape[3] # 将相对于缩放后图像的坐标转换为相对于原始图像的坐标
y1 *= img.shape[2]
x2 *= img.shape[3]
y2 *= img.shape[2]
x1, y1, x2, y2 = int(x1), int(y1), int(x2), int(y2) # 转换为整型坐标
cv2.rectangle(img, (x1, y1), (x2, y2), (255, 0, 0), 2) # 在图像上绘制预测框
cv2.putText(img, f'{cls}', (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.9, (255, 0, 0), 2) # 在图像上绘制类别标签
cv2.imshow('image', img)
cv2.waitKey(0)
cv2.destroyAllWindows()
```
该代码首先加载Yolov5模型和图像,然后进行图像预处理。接下来,使用模型进行推理,获得预测结果。最后,将预测框的相对坐标转换为原始图像的精确坐标,并在图像上绘制预测框和类别标签。
相关推荐
![rar](https://img-home.csdnimg.cn/images/20210720083606.png)
![zip](https://img-home.csdnimg.cn/images/20210720083736.png)
![](https://csdnimg.cn/download_wenku/file_type_ask_c1.png)
![](https://csdnimg.cn/download_wenku/file_type_ask_c1.png)
![](https://csdnimg.cn/download_wenku/file_type_ask_c1.png)
![](https://csdnimg.cn/download_wenku/file_type_ask_c1.png)
![](https://csdnimg.cn/download_wenku/file_type_ask_c1.png)
![](https://csdnimg.cn/download_wenku/file_type_ask_c1.png)
![](https://csdnimg.cn/download_wenku/file_type_ask_c1.png)
![](https://csdnimg.cn/download_wenku/file_type_ask_c1.png)
![](https://csdnimg.cn/download_wenku/file_type_ask_c1.png)
![](https://csdnimg.cn/download_wenku/file_type_ask_c1.png)
![](https://csdnimg.cn/download_wenku/file_type_ask_c1.png)
![](https://csdnimg.cn/download_wenku/file_type_ask_c1.png)
![](https://csdnimg.cn/download_wenku/file_type_ask_c1.png)
![](https://csdnimg.cn/download_wenku/file_type_ask_c1.png)
![](https://csdnimg.cn/download_wenku/file_type_ask_c1.png)