import torchimport cv2import numpy as npfrom models.experimental import attempt_loadfrom utils.general import non_max_suppressionclass YoloV5Detector: def __init__(self, model_path, conf_thresh=0.25, iou_thresh=0.45): self.conf_thresh = conf_thresh self.iou_thresh = iou_thresh self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu") self.model = attempt_load(model_path, map_location=self.device) self.model.eval() def detect(self, image_path): img = cv2.imread(image_path) img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) img = img.transpose(2, 0, 1) img = np.ascontiguousarray(img) img = torch.from_numpy(img).to(self.device).float() / 255.0 # Run inference with torch.no_grad(): results = self.model(img, size=img.shape[-2:]) results = non_max_suppression(results, conf_thres=self.conf_thresh, iou_thres=self.iou_thresh) return results
时间: 2023-06-26 12:09:43 浏览: 233
这是一个使用 YOLOv5 模型进行目标检测的 Python 代码。该类 YoloV5Detector 包含了模型加载、图片预处理和推理的功能。具体来说,会使用 OpenCV 库读取图片,将其转换为 RGB 格式,然后转换为 PyTorch 的 Tensor 格式,并将其送入 YOLOv5 模型中进行推理。最后,使用非极大值抑制算法(NMS)筛选出检测出来的物体,并返回结果。其中,conf_thresh 和 iou_thresh 分别表示置信度和重叠阈值,可以通过调整这两个参数来控制检测结果的准确率和召回率。
相关问题
Django YOLOv5
### 集成YOLOv5到Django项目
为了在Django项目中集成并使用YOLOv5实现目标检测功能,可以遵循如下方法:
#### 准备工作
确保安装了必要的依赖库。对于YOLOv5模型的操作通常需要`torch`和`torchvision`等PyTorch相关包来加载预训练权重文件以及执行推理过程[^2]。
```bash
pip install torch torchvision torchaudio
```
另外还需要克隆YOLOv5官方GitHub仓库至本地环境中以便后续调用其内部函数完成预测任务[^3]。
```bash
git clone https://github.com/ultralytics/yolov5 # 克隆远程版本库
cd yolov5
pip install -r requirements.txt # 安装运行环境所需Python模块列表
```
#### 创建Django应用
启动一个新的Django应用程序用于承载图像上传表单页面及其对应的后台逻辑处理程序[^4]。
```python
django-admin startproject mysite .
python manage.py startapp detector
```
编辑`detector/views.py`定义视图函数接收客户端提交的数据流形式图片资源并通过调用YOLOv5实例返回标注后的结果给前端展示。
```python
from django.http import JsonResponse, HttpResponseBadRequest
import cv2
from PIL import Image
import numpy as np
import io
import base64
from yolov5.models.experimental import attempt_load
from yolov5.utils.general import non_max_suppression
def detect_image(request):
if request.method != 'POST':
return HttpResponseBadRequest('Invalid Request Method')
try:
image_file = request.FILES['image']
img_bytes = image_file.read()
img_np_arr = np.frombuffer(img_bytes, dtype=np.uint8)
frame = cv2.imdecode(img_np_arr, flags=1)
model = attempt_load("yolov5s.pt", map_location="cpu") # 加载模型
results = model(frame) # 执行推断操作
detections = non_max_suppressions(results.xyxy[0], conf_thres=0.25)[0]
annotated_frame = frame.copy() # 复制原帧以绘制边界框
for *box, conf, cls in reversed(detections): # 绘制每个物体位置矩形边框
p1, p2 = (int(box[0]), int(box[1])), (int(box[2]), int(box[3]))
cv2.rectangle(annotated_frame, p1, p2, color=(0, 255, 0), thickness=2)
_, buffer = cv2.imencode('.jpg', annotated_frame)
encoded_img = base64.b64encode(buffer).decode()
response_data = {"status": "success",
"message": "",
"data": {
"detected_objects_count": len(detections),
"processed_image_b64": f"data:image/jpeg;base64,{encoded_img}"
}}
return JsonResponse(response_data)
except Exception as e:
error_message = str(e)
response_data = {"status": "error",
"message": error_message,
"data": None}
return JsonResponse(response_data)
urlpatterns = [
path('detect/', views.detect_image, name='detect'),
]
```
上述代码片段展示了如何创建一个简单的API端点接受来自用户的JPEG格式照片输入,并运用预先下载好的轻量级YOLOv5s模型对其进行分析最后反馈带有标记信息的新版图形数据回浏览器显示出来。
pyqt5 yolov5目标检测
### 实现PyQt5与YOLOv5集成
为了实现在PyQt5中集成YOLOv5进行目标检测的功能,可以遵循以下方法:
#### 创建主窗口类
创建一个继承自`QMainWindow`的类来作为应用程序的主要界面。在此类中初始化UI组件并设置布局。
```python
from PyQt5.QtWidgets import QApplication, QMainWindow, QLabel, QVBoxLayout, QWidget
import sys
class MainWindow(QMainWindow):
def __init__(self):
super().__init__()
# 设置窗口标题和大小
self.setWindowTitle('YOLOv5 Object Detection')
self.setGeometry(100, 100, 800, 600)
layout = QVBoxLayout()
label = QLabel("Welcome to YOLOv5 Detector", self)
layout.addWidget(label)
container = QWidget(self)
container.setLayout(layout)
self.setCentralWidget(container)
if __name__ == '__main__':
app = QApplication(sys.argv)
window = MainWindow()
window.show()
sys.exit(app.exec_())
```
#### 集成YOLOv5模型加载逻辑
在主窗口类内部定义用于加载预训练好的YOLOv5权重文件的方法,并确保该方法可以在GUI线程之外运行以防止阻塞用户输入[^1]。
```python
import torch
from models.experimental import attempt_load
def load_model(weights_path='yolov5s.pt'):
device = 'cuda' if torch.cuda.is_available() else 'cpu'
model = attempt_load(weights_path, map_location=device) # 加载模型
return model.eval().to(device), device
```
#### 开发检测功能模块
编写专门负责执行实际对象识别任务的部分,这通常涉及到调用YOLOv5库来进行推理操作并将结果反馈给前端显示出来[^3]。
```python
import cv2
from utils.general import non_max_suppression, scale_coords
from utils.plots import plot_one_box
def perform_detection(frame, model, img_size=640, conf_thres=0.25, iou_thres=0.45):
stride = int(model.stride.max())
names = model.module.names if hasattr(model, 'module') else model.names
h, w = frame.shape[:2]
resized_img = letterbox(frame, new_shape=img_size)[0]
input_tensor = torch.from_numpy(resized_img).permute((2, 0, 1)).float().div_(255.).unsqueeze(dim=0).contiguous()
pred = model(input_tensor.to(next(model.parameters()).device))[0]
det = non_max_suppression(pred, conf_thres, iou_thres)[0]
gn = torch.tensor([w, h, w, h])
detections = []
if len(det):
det[:, :4] = scale_coords(input_tensor.shape[2:], det[:, :4], (h, w)).round()
for *xyxy, conf, cls in reversed(det):
c = int(cls)
label = f'{names[c]} {conf:.2f}'
detection_info = {
"label": label,
"confidence": float(conf),
"bbox": tuple(map(int, xyxy))
}
detections.append(detection_info)
plot_one_box(xyxy, frame, label=label, color=(0, 255, 0))
return frame, detections
```
#### 连接前后端
最后一步就是把前面提到的所有部分连接起来,在主线程中定期获取摄像头帧或读取视频文件的数据传递给`perform_detection()`函数处理后再更新到界面上展示给用户查看[^2]。
```python
timer = QTimer(window)
timer.timeout.connect(lambda: update_frame())
timer.start(30)
video_capture = None
def start_video_feed():
global video_capture
video_capture = cv2.VideoCapture(0)
def stop_video_feed():
global video_capture
if video_capture is not None and video_capture.isOpened():
video_capture.release()
video_capture = None
def update_frame():
ret, frame = video_capture.read()
if ret:
processed_image, _ = perform_detection(frame, loaded_model)
showImage(processed_image)
start_video_feed()
```
阅读全文
相关推荐
















