项目简介
该项目使用YOLOv5深度学习框架来检测图像或视频中的猫和狗。YOLOv5(You Only Look Once v5)是一种高效的物体检测模型,能够快速准确地识别出图像中的目标。本项目具有以下特点:
- 图像检测:用户可以通过上传图片,系统将识别并标记出图像中的猫和狗的位置。
- 视频检测:支持实时视频流或本地视频文件的猫狗检测,并实时显示检测结果。
技术栈
- YOLOv5:用于目标检测的核心框架。
- Python:主要编程语言。
- PyTorch:YOLOv5基于此深度学习框架构建。
- OpenCV:用于视频处理和图像显示。
- Flask:用于构建Web前端,让用户能够上传图片或视频。
关键代码示例
安装YOLOv5
首先需要安装YOLOv5及其依赖项。可以通过以下命令安装:
1pip install -r requirements.txt
假设requirements.txt
文件中包含了YOLOv5所需的依赖包。
下载预训练模型
下载YOLOv5预训练模型,可以使用以下命令:
1wget https://github.com/ultralytics/yolov5/releases/download/v5.0/yolov5s.pt
图像检测
接下来是使用YOLOv5进行图像检测的代码示例:
1import torch
2from PIL import Image
3
4# 加载预训练的模型
5model = torch.hub.load('ultralytics/yolov5', 'custom', path='yolov5s.pt')
6
7# 加载图片
8img = Image.open('path/to/image.jpg')
9
10# 进行检测
11results = model(img)
12
13# 显示结果
14results.show()
视频检测
对于视频检测,你可以使用以下代码:
1import cv2
2from yolov5.detect import detect
3
4def video_detection(video_path):
5 cap = cv2.VideoCapture(video_path)
6 while cap.isOpened():
7 ret, frame = cap.read()
8 if not ret:
9 break
10
11 # 转换为PIL图像格式
12 img = Image.fromarray(frame)
13
14 # 进行检测
15 results = detect(model, img)
16
17 # 显示检测结果
18 # 注意:你需要自己实现detect函数来适配YOLOv5的输出
19 # 也可以使用YOLOv5自带的display方法
20 results.render() # 更新图像框
21 cv2.imshow('Detection Results', results.ims[0])
22
23 if cv2.waitKey(1) & 0xFF == ord('q'):
24 break
25
26 cap.release()
27 cv2.destroyAllWindows()
28
29# 使用视频检测函数
30video_detection('path/to/video.mp4')
Web界面
为了搭建一个用户可以上传图片或视频的Web界面,我们可以使用Flask框架。以下是一个简单的Web应用示例:
1from flask import Flask, render_template, request
2import os
3import torch
4from PIL import Image
5import cv2
6
7app = Flask(__name__)
8
9# 加载预训练的模型
10model = torch.hub.load('ultralytics/yolov5', 'custom', path='yolov5s.pt')
11
12@app.route('/', methods=['GET', 'POST'])
13def index():
14 if request.method == 'POST':
15 if 'file' not in request.files:
16 return 'No file part'
17 file = request.files['file']
18 if file.filename == '':
19 return 'No selected file'
20 if file:
21 filename = file.filename
22 file_path = os.path.join('uploads', filename)
23 file.save(file_path)
24
25 # 图像检测
26 img = Image.open(file_path)
27 results = model(img)
28
29 # 保存检测后的图像
30 results.save(save_dir='static/detections', exist_ok=True)
31 detection_path = os.path.join('static', 'detections', filename)
32
33 return render_template('index.html', detection=detection_path)
34
35 return render_template('index.html')
36
37if __name__ == '__main__':
38 app.run(debug=True)
注意事项
- 模型训练:为了使模型能够准确检测猫和狗,你可能需要使用包含猫和狗的数据集对模型进行微调。
- 环境配置:确保你的环境已经正确配置了PyTorch和其他必要的库。
- 性能优化:可以根据实际需求调整模型参数以优化性能。