yolov5使用flask部署至前端實現(xiàn)照片\視頻識別功能
前言
初學yolo flask時,需要此功能,Csdn、Github、B站找到許多代碼,效果并不滿意。
近期,再度嘗試,實現(xiàn)簡單功能。
實現(xiàn)功能:
上傳圖片并識別,可以點擊圖片放大查看
上傳視頻并識別
識別后的文件下載功能
效果圖如下

文件結構如下:
project/
static/空
templates/
index.html
app.py
相關代碼:
app.py
import cv2
import numpy as np
import torch
from flask import Flask, request, jsonify, render_template
import base64
import os
from datetime import datetime
app = Flask(__name__)
# 全局變量:模型
model = None
# 提前加載模型
def load_model():
global model
model = torch.hub.load('', 'custom', path='yolov5s.pt', source='local')
# 路由處理圖片檢測請求
@app.route('/predict_image', methods=['POST'])
def predict_image():
global model
# 獲取圖像文件
file = request.files['image']
# 讀取圖像數(shù)據(jù)并轉換為RGB格式
image_data = file.read()
nparr = np.frombuffer(image_data, np.uint8)
image = cv2.imdecode(nparr, cv2.IMREAD_UNCHANGED)
results = model(image)
image = results.render()[0]
# 將圖像轉換為 base64 編碼的字符串
_, buffer = cv2.imencode('.png', image)
image_str = base64.b64encode(buffer).decode('utf-8')
# 獲取當前時間,并將其格式化為字符串
current_time = datetime.now().strftime('%Y%m%d%H%M%S')
# 構建保存路徑
save_dir = 'static'
if not os.path.exists(save_dir):
os.makedirs(save_dir)
filename, extension = os.path.splitext(file.filename) # 獲取上傳文件的文件名和擴展名
save_filename = f'{filename}_{current_time}{extension}'
save_path = os.path.join(save_dir, save_filename)
cv2.imwrite(save_path, image)
return jsonify({'image': image_str})
# 函數(shù)用于在視頻幀上繪制檢測結果
def detect_objects(frame, model):
results = model(frame)
detections = results.xyxy[0].cpu().numpy() # 獲取檢測結果
# 在幀上繪制檢測結果
for det in detections:
# 獲取邊界框信息
x1, y1, x2, y2, conf, class_id = det[:6]
# 在幀上繪制邊界框
cv2.rectangle(frame, (int(x1), int(y1)), (int(x2), int(y2)), (0, 255, 0), 2)
# 在幀上繪制類別和置信度
label = f'{model.names[int(class_id)]} {conf:.2f}'
cv2.putText(frame, label, (int(x1), int(y1) - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2)
return frame
# 路由處理視頻檢測請求
@app.route("/predict_video", methods=["POST"])
def predict_video():
global model
# 從請求中獲取視頻文件
video_file = request.files["video"]
# 保存視頻到臨時文件
temp_video_path = "temp_video.mp4"
video_file.save(temp_video_path)
# 逐幀讀取視頻
video = cv2.VideoCapture(temp_video_path)
# 獲取視頻的幀率和尺寸
fps = video.get(cv2.CAP_PROP_FPS)
width = int(video.get(cv2.CAP_PROP_FRAME_WIDTH))
height = int(video.get(cv2.CAP_PROP_FRAME_HEIGHT))
# 視頻寫入對象
output_video_filename = f"output_video_{datetime.now().strftime('%Y%m%d%H%M%S')}.mp4"
output_video_path = os.path.join("static", output_video_filename)
fourcc = cv2.VideoWriter_fourcc(*"avc1")
out_video = cv2.VideoWriter(output_video_path, fourcc, fps, (width, height))
# 逐幀處理視頻并進行目標檢測
while True:
ret, frame = video.read()
if not ret:
break
# 進行目標檢測
detection_result = detect_objects(frame, model)
# 將處理后的幀寫入輸出視頻
out_video.write(detection_result)
# 釋放視頻對象
video.release()
out_video.release()
return jsonify({"output_video_path": output_video_filename})
@app.route('/')
def index():
return render_template('index.html')
# 初始加載模型
load_model()
if __name__ == '__main__':
app.run(debug=True)
index.html
<!DOCTYPE html>
<html lang="en">
<head>
<meta charset="UTF-8">
<meta name="viewport" content="width=device-width, initial-scale=1.0">
<title>Object Detection</title>
<style>
body {
font-family: Arial, sans-serif;
margin: 0;
padding: 0;
background-color: #f3f3f3;
display: flex;
justify-content: center;
align-items: center;
height: 100vh;
flex-direction: column;
}
#content {
text-align: center;
max-width: 820px;
margin-top: 20px;
}
h1 {
color: #333;
}
h2 {
color: #666;
}
input[type="file"] {
margin-bottom: 10px;
}
.media-container {
display: flex;
max-width: 100%;
margin-bottom: 20px;
}
.media-container:first-child {
margin-right: 20px; /* 在第一個容器的右側添加間隔 */
}
.media-container img,
.media-container video {
max-width: 100%;
height: auto;
}
.original {
width: 400px;
overflow: hidden;
}
.processed {
flex: 2; /* 右邊容器占據(jù)剩余空間 */
}
button {
padding: 10px 20px;
background-color: #007bff;
color: #fff;
border: none;
border-radius: 5px;
cursor: pointer;
margin-bottom: 10px;
}
/* 新增樣式:模態(tài)框 */
.modal {
display: none; /* 默認隱藏 */
position: fixed;
z-index: 1;
left: 0;
top: 0;
width: 100%;
height: 100%;
overflow: auto;
background-color: rgba(0, 0, 0, 0.9); /* 半透明黑色背景 */
}
.modal-content {
margin: auto;
display: block;
width: 80%;
max-width: 800px;
position: absolute;
left: 50%;
top: 50%;
transform: translate(-50%, -50%);
text-align: center; /* 居中顯示圖片 */
}
.close {
color: #ccc;
font-size: 36px;
font-weight: bold;
cursor: pointer;
position: absolute;
top: 10px;
right: 10px;
}
.close:hover,
.close:focus {
color: #fff;
text-decoration: none;
}
#downloadButton {
padding: 10px 20px;
background-color: #007bff;
color: #fff;
border: none;
border-radius: 5px;
cursor: pointer;
margin-bottom: 10px;
}
/* 新增樣式:響應式圖片 */
.modal-content img,
.modal-content video {
max-width: 100%;
height: auto;
}
</style>
</head>
<body>
<!-- 新增模態(tài)框 -->
<div id="myModal" class="modal" onclick="closeModal()">
<div class="modal-content" id="modalContent" onclick="stopPropagation(event)">
<!-- 放大后的圖片或視頻將在這里顯示 -->
<span class="close" onclick="closeModal()">×</span>
</div>
</div>
<div id="content">
<h1>照片/視頻檢測</h1>
<!-- 上傳圖片 -->
<h2>上傳圖片</h2>
<input type="file" id="imageFile" accept="image/*" onchange="displaySelectedImage()">
<button onclick="uploadImage()">上傳</button>
<button id="downloadImageButton" onclick="downloadProcessedImage()">下載</button>
<br>
<div class="media-container">
<div class="original media-container" onclick="enlargeImage()">
<img id="uploadedImage" src="#" alt="Uploaded Image" style="display:none;">
<button id="zoomInButton" style="display:none;">Zoom In</button>
</div>
<div class="processed media-container" onclick="enlargeImage2()">
<img id="processedImage" src="#" alt="Processed Image" style="display:none;">
</div>
</div>
<br>
<!-- 上傳視頻 -->
<h2>上傳視頻</h2>
<input type="file" id="videoFile" accept="video/mp4,video/x-m4v,video/*" onchange="displaySelectedVideo()">
<button onclick="uploadVideo()">上傳</button>
<button id="downloadButton" onclick="downloadProcessedVideo()">下載</button>
<br>
<div class="media-container">
<div class="original media-container" >
<video id="uploadedVideo" src="#" controls style="display:none;"></video>
</div>
<div class="processed media-container">
<video id="processedVideo" controls style="display:none;"></video>
</div>
</div>
<br>
</div>
<script>
// 顯示選擇的權重文件
// 顯示選擇的圖片并添加點擊放大功能
function displaySelectedImage() {
var fileInput = document.getElementById('imageFile');
var file = fileInput.files[0];
var imageElement = document.getElementById('uploadedImage');
imageElement.src = URL.createObjectURL(file);
imageElement.style.display = 'inline';
document.getElementById('zoomInButton').style.display = 'inline';
}
// 顯示模態(tài)框并放大圖片
function enlargeImage() {
var modal = document.getElementById('myModal');
var modalImg = document.getElementById('modalContent');
var img = document.getElementById('uploadedImage');
modal.style.display = 'block';
modalImg.innerHTML = '<img src="' + img.src + '">';
}
// 顯示模態(tài)框并放大圖片
function enlargeImage2() {
var modal = document.getElementById('myModal');
var modalImg = document.getElementById('modalContent');
var img = document.getElementById('processedImage');
modal.style.display = 'block';
modalImg.innerHTML = '<img src="' + img.src + '">';
}
// 顯示選擇的視頻并添加點擊放大功能
function displaySelectedVideo() {
var fileInput = document.getElementById('videoFile');
var file = fileInput.files[0];
var videoElement = document.getElementById('uploadedVideo');
videoElement.src = URL.createObjectURL(file);
videoElement.style.display = 'block';
}
// 上傳圖片并向后端發(fā)送請求
function uploadImage() {
var fileInput = document.getElementById('imageFile');
var file = fileInput.files[0];
var formData = new FormData();
formData.append('image', file);
fetch('/predict_image', {
method: 'POST',
body: formData
})
.then(response => response.json())
.then(data => {
var imageElement = document.getElementById('processedImage');
imageElement.src = 'data:image/png;base64,' + data.image;
imageElement.style.display = 'inline';
document.getElementById('downloadImageButton').style.display = 'inline';
})
.catch(error => console.error('Error:', error));
}
// 下載處理后的圖片
function downloadProcessedImage() {
var imageElement = document.getElementById('processedImage');
var url = imageElement.src;
var a = document.createElement('a');
a.href = url;
a.download = 'processed_image.png';
document.body.appendChild(a);
a.click();
document.body.removeChild(a);
}
// 上傳視頻并向后端發(fā)送請求
function uploadVideo() {
var fileInput = document.getElementById('videoFile');
var file = fileInput.files[0];
var formData = new FormData();
formData.append('video', file);
fetch('/predict_video', {
method: 'POST',
body: formData
})
.then(response => response.json())
.then(data => {
var videoElement = document.getElementById('processedVideo');
// 修改路徑為正確的 Flask url_for 生成的路徑
videoElement.src = '{{ url_for("static", filename="") }}' + data.output_video_path;
videoElement.style.display = 'block';
var downloadButton = document.getElementById('downloadButton');
downloadButton.style.display = 'block';
})
.catch(error => console.error('Error:', error));
}
// 下載處理后的視頻
function downloadProcessedVideo() {
var videoElement = document.getElementById('processedVideo');
var url = videoElement.src;
var a = document.createElement('a');
a.href = url;
a.download = 'processed_video.mp4';
document.body.appendChild(a);
a.click();
document.body.removeChild(a);
}
// 關閉模態(tài)框
function closeModal() {
var modal = document.getElementById('myModal');
modal.style.display = 'none';
}
</script>
</body>
</html>使用說明:
index.html放入templates文件夾中
運行app.py
注:此處加載模型路徑更改為自己的
model = torch.hub.load('', 'custom', path='yolov5s.pt', source='local')如果模型讀取不到,顯示

FileNotFoundError: [Errno 2] No such file or directory: 'hubconf.py'
去yolov5官網,下載yolov5-master到項目文件夾

并將yolov5s.pt文件復制到y(tǒng)olov5-master文件夾中,修改model路徑
model = torch.hub.load('yolov5-master', 'custom', path='yolov5s.pt', source='local')總結
到此這篇關于yolov5使用flask部署至前端實現(xiàn)照片\視頻識別功能的文章就介紹到這了,更多相關yolov5 flask部署前端照片視頻識別內容請搜索腳本之家以前的文章或繼續(xù)瀏覽下面的相關文章希望大家以后多多支持腳本之家!
相關文章
Python使用Beautiful?Soup(BS4)庫解析HTML和XML
這篇文章介紹了Python使用Beautiful?Soup(BS4)庫解析HTML和XML的方法,對大家的學習或者工作具有一定的參考學習價值,需要的朋友們下面隨著小編來一起學習學習吧2022-06-06

