写一段python代码, 用yolov5实现对某个图片检测, 返回坐标和置信度
时间: 2024-12-18 15:17:15 浏览: 8
当然可以,YOLOv5是一个流行的实时物体检测模型,你可以使用Python的`pytorch-yolo-v5`库来加载模型并进行预测。以下是一个简单的例子,展示如何在YOLOv5中检测一张图片,并返回每个检测到的对象的坐标和置信度:
```python
# 首先,确保已安装所需的库
!pip install torch torchvision fastai pytorch-lightning
import torch
from PIL import Image
from yolov5.models.common import Detect
from yolov5.utils.general import non_max_suppression
# 加载预训练的YOLOv5模型
model = torch.hub.load('ultralytics/yolov5', 'yolov5s')
# 使用图片路径
image_path = 'path_to_your_image.jpg'
# 将图片转换成Tensor并进行前向传播
img = Image.open(image_path).convert('RGB')
tensor_img = img.resize((640, 640)).unsqueeze(0) # YOLOv5默认输入尺寸
preds = model(tensor_img)
# 进行非极大抑制,得到最终的检测结果
results = non_max_suppression(preds, conf_thres=0.5, iou_thres=0.5) # 设置阈值
for box in results[0]: # 如果有检测到对象
x1, y1, x2, y2 = box['bbox'] # (左上角x, 左上角y, 右下角x, 右下角y)
confidence = box['confidence']
print(f"坐标: ({x1:.2f}, {y1:.2f}), ({x2:.2f}, {y2:.2f}) 置信度: {confidence:.2f}")
阅读全文