基于yolov5的行人检测代码
时间: 2023-11-04 09:55:11 浏览: 36
以下是基于YoloV5的行人检测代码:
1. 首先,您需要下载YoloV5代码库。您可以从 https://github.com/ultralytics/yolov5 下载它。
2. 接下来,您需要安装所需的依赖项。您可以使用以下命令在终端中安装依赖项:
```
pip install -r requirements.txt
```
3. 下载预训练权重文件,可以从 https://github.com/ultralytics/yolov5/releases 下载。将下载的权重文件保存在该目录下。
4. 创建一个Python文件,例如“pedestrian_detection.py”,并使用以下代码:
```
import torch
import cv2
from models.experimental import attempt_load
from utils.general import non_max_suppression
def pedestrian_detection(image_path):
# 加载权重
weights = 'path_to_weights_file/best.pt'
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model = attempt_load(weights, map_location=device) # 加载模型
model.eval()
# 加载图像
img0 = cv2.imread(image_path) # 读取图像
img = img0.copy()
# 将图像转换为Tensor
img = torch.from_numpy(img).to(device)
img = img.permute(2, 0, 1).float().unsqueeze(0)
# 进行预测
pred = model(img, augment=False)[0]
pred = non_max_suppression(pred, conf_thres=0.4, iou_thres=0.45)
# 处理预测结果
for det in pred[0]:
# 只保留行人类别的预测
if det[-1] == 0:
# 绘制边界框
x1, y1, x2, y2 = map(int, det[:4])
cv2.rectangle(img0, (x1, y1), (x2, y2), (0, 255, 0), 2)
# 显示图像
cv2.imshow('Pedestrian detection', img0)
cv2.waitKey(0)
cv2.destroyAllWindows()
if __name__ == '__main__':
# 图像路径
image_path = 'path_to_image_file/image.jpg'
# 进行行人检测
pedestrian_detection(image_path)
```
5. 运行Python脚本,将显示带有行人边界框的图像。
请注意,此代码仅适用于单张图像的行人检测。如果您想处理多个图像或视频,则需要进行适当的更改。