django部署yolov5训练好的模型并显示检测后的图片,请给出view.py、models.py和前端页面
时间: 2023-06-14 13:06:19 浏览: 79
下面是一个简单的 Django 项目,可以将 YOLOv5 训练好的模型部署到 Django 上,并在前端页面上显示检测后的图片。
首先,我们需要在 Django 项目中创建一个 app,我这里将其命名为 `yolov5_detection`。
在 `models.py` 中,我们定义了一个 `DetectionModel` 模型,用于加载 YOLOv5 训练好的模型和进行目标检测。
```python
import torch
import cv2
from django.db import models
class DetectionModel:
def __init__(self, model_path):
self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
self.model = torch.hub.load('ultralytics/yolov5', 'custom', path=model_path).to(self.device).autoshape()
def detect(self, image_path):
img = cv2.imread(image_path)
results = self.model([img.to(self.device)])
return results.xyxy[0].cpu().numpy()
```
在 `view.py` 中,我们定义了一个 `detect` 视图函数,用于处理前端上传的图片并进行目标检测。在这个函数中,我们使用 `DetectionModel` 加载模型并进行目标检测,将检测结果保存到一个列表中,并将其传递给前端页面。
```python
from django.shortcuts import render
from django.http import JsonResponse
from django.views.decorators.csrf import csrf_exempt
from .models import DetectionModel
@csrf_exempt
def detect(request):
if request.method == 'POST':
model_path = 'path/to/model.pt'
detection_model = DetectionModel(model_path)
image = request.FILES['image']
image_path = 'path/to/image.jpg'
with open(image_path, 'wb+') as f:
f.write(image.read())
results = detection_model.detect(image_path)
detections = []
for res in results:
detections.append({
'class': res[5],
'confidence': res[4],
'bbox': [int(res[0]), int(res[1]), int(res[2]), int(res[3])]
})
return JsonResponse({'detections': detections})
return render(request, 'detect.html')
```
在前端页面上,我们使用 HTML 和 JavaScript 来实现上传图片并显示检测结果。在这个页面中,我们使用 `fetch` 函数向后端发送 POST 请求,上传图片并接收检测结果。在接收到结果后,将其绘制到 Canvas 上。
```html
<!DOCTYPE html>
<html>
<head>
<meta charset="UTF-8">
<title>YOLOv5 Detection</title>
</head>
<body>
<input type="file" id="image" accept="image/*">
<button id="detect">Detect</button>
<canvas id="canvas"></canvas>
<script>
const image = document.getElementById('image');
const detectButton = document.getElementById('detect');
const canvas = document.getElementById('canvas');
detectButton.addEventListener('click', async () => {
const formData = new FormData();
formData.append('image', image.files[0]);
const response = await fetch('/detect', {
method: 'POST',
body: formData
});
const data = await response.json();
const detections = data.detections;
const ctx = canvas.getContext('2d');
const img = new Image();
img.onload = () => {
canvas.width = img.width;
canvas.height = img.height;
ctx.drawImage(img, 0, 0);
ctx.lineWidth = 2;
ctx.font = '20px Arial';
for (const detection of detections) {
const [x1, y1, x2, y2] = detection.bbox;
ctx.strokeStyle = 'red';
ctx.strokeRect(x1, y1, x2 - x1, y2 - y1);
ctx.fillStyle = 'red';
ctx.fillText(`${detection.class} (${detection.confidence.toFixed(2)})`, x1, y1 - 5);
}
};
img.src = URL.createObjectURL(image.files[0]);
});
</script>
</body>
</html>
```
最后,在 `urls.py` 中将 `detect` 视图函数和前端页面进行 URL 映射。
```python
from django.urls import path
from .views import detect
urlpatterns = [
path('', detect, name='detect'),
]
```
完整代码可以参考以下链接:
- [Django-YOLOv5-Detection](https://github.com/leemengtaiwan/Django-YOLOv5-Detection)