python脚本运行yolov5
时间: 2024-03-29 09:33:43 浏览: 19
Yolov5是一个用于目标检测的开源深度学习模型,而Python是一种常用的编程语言,可以用于编写和运行Yolov5的脚本。下面是一个简单的Python脚本示例,用于运行Yolov5进行目标检测:
```python
import torch
from PIL import Image
from torchvision import transforms
from models.experimental import attempt_load
from utils.general import non_max_suppression, scale_coords
from utils.plots import plot_one_box
# 加载Yolov5模型
model = attempt_load('yolov5s.pt', map_location=torch.device('cpu'))
# 设置输入图像的大小和变换
img_size = 640
transform = transforms.Compose([
transforms.Resize(img_size),
transforms.ToTensor(),
transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
])
# 加载并预处理图像
image_path = 'image.jpg'
image = Image.open(image_path).convert('RGB')
image_tensor = transform(image).unsqueeze(0)
# 运行目标检测
with torch.no_grad():
detections = model(image_tensor)
# 对检测结果进行后处理
results = non_max_suppression(detections, conf_thres=0.3, iou_thres=0.5)
# 绘制检测结果
if results[0] is not None:
for x1, y1, x2, y2, conf, cls in results[0]:
plot_one_box((x1, y1, x2, y2), image, label=int(cls), color=(0, 255, 0))
# 显示结果图像
image.show()
```
请注意,上述示例假设你已经下载了Yolov5的预训练模型文件(yolov5s.pt)和一张待检测的图像(image.jpg)。你需要将这些文件放在与脚本相同的目录下,并根据实际情况进行相应的修改。