20 changed files with 653 additions and 385 deletions
@ -0,0 +1,15 @@ |
|||||
|
<?xml version="1.0" encoding="UTF-8"?> |
||||
|
<project version="4"> |
||||
|
<component name="PublishConfigData" autoUpload="Always" serverName="root@192.168.3.48:22" remoteFilesAllowedToDisappearOnAutoupload="false"> |
||||
|
<serverData> |
||||
|
<paths name="root@192.168.3.48:22"> |
||||
|
<serverdata> |
||||
|
<mappings> |
||||
|
<mapping deploy="/mnt/zfbox" local="$PROJECT_DIR$" /> |
||||
|
</mappings> |
||||
|
</serverdata> |
||||
|
</paths> |
||||
|
</serverData> |
||||
|
<option name="myAutoUpload" value="ALWAYS" /> |
||||
|
</component> |
||||
|
</project> |
||||
@ -1,162 +1,120 @@ |
|||||
# 导入代码依赖 |
# 导入代码依赖 |
||||
import cv2 |
import cv2 |
||||
|
# import os |
||||
|
# from model.plugins.ModelBase import ModelBase |
||||
|
# import importlib.util |
||||
|
import threading |
||||
|
import time |
||||
import numpy as np |
import numpy as np |
||||
import ipywidgets as widgets |
from quart import Quart, render_template, websocket, jsonify |
||||
from IPython.display import display |
from quart.helpers import Response |
||||
import torch |
|
||||
from skvideo.io import vreader, FFmpegWriter |
# def test_acl_verify(): |
||||
import IPython.display |
# '''根据路径，动态导入模块''' |
||||
from ais_bench.infer.interface import InferSession |
# model_path = "/mnt/zfbox/model/plugins/RYRQ_ACL/RYRQ_Model_ACL.py" |
||||
|
# model_name = "测试" |
||||
from det_utils import letterbox, scale_coords, nms |
# if os.path.exists(model_path): |
||||
|
# module_spec = importlib.util.spec_from_file_location(model_name, model_path) |
||||
def preprocess_image(image, cfg, bgr2rgb=True): |
# if module_spec is None: |
||||
"""图片预处理""" |
# print(f"{model_path} 加载错误") |
||||
img, scale_ratio, pad_size = letterbox(image, new_shape=cfg['input_shape']) |
# return None |
||||
if bgr2rgb: |
# module = importlib.util.module_from_spec(module_spec) |
||||
img = img[:, :, ::-1] |
# module_spec.loader.exec_module(module) |
||||
img = img.transpose(2, 0, 1) # HWC2CHW |
# md = getattr(module, "Model")(model_path,0.5) # 实例化类 |
||||
img = np.ascontiguousarray(img, dtype=np.float32) |
# if not isinstance(md, ModelBase) or not md.init_ok: |
||||
return img, scale_ratio, pad_size |
# print("{} not zf_model".format(md)) |
||||
|
# return None |
||||
|
# else: |
||||
def draw_bbox(bbox, img0, color, wt, names): |
# print("{}文件不存在".format(model_path)) |
||||
"""在图片上画预测框""" |
# return None |
||||
det_result_str = '' |
# print("开执行检测") |
||||
for idx, class_id in enumerate(bbox[:, 5]): |
# #return md |
||||
if float(bbox[idx][4] < float(0.05)): |
# img = cv2.imread("/mnt/zfbox/model/plugins/RYRQ_ACL/world_cup.jpg", cv2.IMREAD_COLOR) # 读入图片 |
||||
continue |
# data = [0,1,"[(100, 100), (50, 200), (100, 300), (300, 300), (300, 100)]"] |
||||
img0 = cv2.rectangle(img0, (int(bbox[idx][0]), int(bbox[idx][1])), (int(bbox[idx][2]), int(bbox[idx][3])), |
# boxout,ret,strret = md.verify(img,data) |
||||
color, wt) |
# #print(boxout) |
||||
img0 = cv2.putText(img0, str(idx) + ' ' + names[int(class_id)], (int(bbox[idx][0]), int(bbox[idx][1] + 16)), |
# #释放相关资源 |
||||
cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 255), 1) |
# md.release() |
||||
img0 = cv2.putText(img0, '{:.4f}'.format(bbox[idx][4]), (int(bbox[idx][0]), int(bbox[idx][1] + 32)), |
|
||||
cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 255), 1) |
|
||||
det_result_str += '{} {} {} {} {} {}\n'.format( |
|
||||
names[bbox[idx][5]], str(bbox[idx][4]), bbox[idx][0], bbox[idx][1], bbox[idx][2], bbox[idx][3]) |
# 加载YOLOv5模型 |
||||
return img0 |
#model = yolov5.load('yolov5s') |
||||
|
|
||||
|
app = Quart(__name__) |
||||
def get_labels_from_txt(path): |
|
||||
"""从txt文件获取图片标签""" |
class VideoCaptureWithYOLO: |
||||
labels_dict = dict() |
def __init__(self, source): |
||||
with open(path) as f: |
self.source = source |
||||
for cat_id, label in enumerate(f.readlines()): |
self.cap = cv2.VideoCapture(self._get_gstreamer_pipeline(), cv2.CAP_GSTREAMER) |
||||
labels_dict[cat_id] = label.strip() |
self.width = int(self.cap.get(cv2.CAP_PROP_FRAME_WIDTH)) |
||||
return labels_dict |
self.height = int(self.cap.get(cv2.CAP_PROP_FRAME_HEIGHT)) |
||||
|
self.frame = None |
||||
|
self.running = True |
||||
def draw_prediction(pred, image, labels): |
self.read_lock = threading.Lock() |
||||
"""在图片上画出预测框并进行可视化展示""" |
self.thread = threading.Thread(target=self.update) |
||||
imgbox = widgets.Image(format='jpg', height=720, width=1280) |
self.thread.start() |
||||
img_dw = draw_bbox(pred, image, (0, 255, 0), 2, labels) |
|
||||
imgbox.value = cv2.imencode('.jpg', img_dw)[1].tobytes() |
def _get_gstreamer_pipeline(self): |
||||
display(imgbox) |
return ( |
||||
|
f"rtspsrc location={self.source} protocols=udp latency=0 ! " |
||||
|
"rtph264depay ! h264parse ! avdec_h264 ! videoconvert ! appsink" |
||||
def infer_image(img_path, model, class_names, cfg): |
) |
||||
"""图片推理""" |
|
||||
# 图片载入 |
def update(self): |
||||
image = cv2.imread(img_path) |
while self.running: |
||||
# 数据预处理 |
ret, frame = self.cap.read() |
||||
img, scale_ratio, pad_size = preprocess_image(image, cfg) |
|
||||
# 模型推理 |
|
||||
output = model.infer([img])[0] |
|
||||
|
|
||||
output = torch.tensor(output) |
|
||||
# 非极大值抑制后处理 |
|
||||
boxout = nms(output, conf_thres=cfg["conf_thres"], iou_thres=cfg["iou_thres"]) |
|
||||
pred_all = boxout[0].numpy() |
|
||||
# 预测坐标转换 |
|
||||
scale_coords(cfg['input_shape'], pred_all[:, :4], image.shape, ratio_pad=(scale_ratio, pad_size)) |
|
||||
# 图片预测结果可视化 |
|
||||
draw_prediction(pred_all, image, class_names) |
|
||||
|
|
||||
|
|
||||
def infer_frame_with_vis(image, model, labels_dict, cfg, bgr2rgb=True): |
|
||||
# 数据预处理 |
|
||||
img, scale_ratio, pad_size = preprocess_image(image, cfg, bgr2rgb) |
|
||||
# 模型推理 |
|
||||
output = model.infer([img])[0] |
|
||||
|
|
||||
output = torch.tensor(output) |
|
||||
# 非极大值抑制后处理 |
|
||||
boxout = nms(output, conf_thres=cfg["conf_thres"], iou_thres=cfg["iou_thres"]) |
|
||||
pred_all = boxout[0].numpy() |
|
||||
# 预测坐标转换 |
|
||||
scale_coords(cfg['input_shape'], pred_all[:, :4], image.shape, ratio_pad=(scale_ratio, pad_size)) |
|
||||
# 图片预测结果可视化 |
|
||||
img_vis = draw_bbox(pred_all, image, (0, 255, 0), 2, labels_dict) |
|
||||
return img_vis |
|
||||
|
|
||||
|
|
||||
def img2bytes(image): |
|
||||
"""将图片转换为字节码""" |
|
||||
return bytes(cv2.imencode('.jpg', image)[1]) |
|
||||
|
|
||||
|
|
||||
def infer_video(video_path, model, labels_dict, cfg): |
|
||||
"""视频推理""" |
|
||||
image_widget = widgets.Image(format='jpeg', width=800, height=600) |
|
||||
display(image_widget) |
|
||||
|
|
||||
# 读入视频 |
|
||||
cap = cv2.VideoCapture(video_path) |
|
||||
while True: |
|
||||
ret, img_frame = cap.read() |
|
||||
if not ret: |
if not ret: |
||||
break |
continue |
||||
# 对视频帧进行推理 |
# YOLOv5分析 |
||||
image_pred = infer_frame_with_vis(img_frame, model, labels_dict, cfg, bgr2rgb=True) |
annotated_frame = frame |
||||
image_widget.value = img2bytes(image_pred) |
|
||||
|
with self.read_lock: |
||||
|
self.frame = annotated_frame |
||||
def infer_camera(model, labels_dict, cfg): |
time.sleep(0.01) # 控制帧率 |
||||
"""外设摄像头实时推理""" |
|
||||
def find_camera_index(): |
def read(self): |
||||
max_index_to_check = 10 # Maximum index to check for camera |
with self.read_lock: |
||||
|
frame = self.frame.copy() if self.frame is not None else None |
||||
for index in range(max_index_to_check): |
return frame |
||||
cap = cv2.VideoCapture(index) |
|
||||
if cap.read()[0]: |
def release(self): |
||||
cap.release() |
self.running = False |
||||
return index |
self.thread.join() |
||||
|
self.cap.release() |
||||
# If no camera is found |
|
||||
raise ValueError("No camera found.") |
# 创建多个视频流 |
||||
|
streams = { |
||||
# 获取摄像头 |
'stream1': VideoCaptureWithYOLO('rtsp://192.168.3.44/live1'), |
||||
camera_index = find_camera_index() |
'stream2': VideoCaptureWithYOLO('rtsp://192.168.3.44/live1'), |
||||
cap = cv2.VideoCapture(camera_index) |
'stream3': VideoCaptureWithYOLO('rtsp://192.168.3.44/live1'), |
||||
# 初始化可视化对象 |
'stream4': VideoCaptureWithYOLO('rtsp://192.168.3.44/live1'), |
||||
image_widget = widgets.Image(format='jpeg', width=1280, height=720) |
'stream5': VideoCaptureWithYOLO('rtsp://192.168.3.44/live1'), |
||||
display(image_widget) |
'stream6': VideoCaptureWithYOLO('rtsp://192.168.3.44/live1'), |
||||
while True: |
'stream7': VideoCaptureWithYOLO('rtsp://192.168.3.44/live1'), |
||||
# 对摄像头每一帧进行推理和可视化 |
'stream8': VideoCaptureWithYOLO('rtsp://192.168.3.44/live1'), |
||||
_, img_frame = cap.read() |
'stream9': VideoCaptureWithYOLO('rtsp://192.168.3.44/live1') |
||||
image_pred = infer_frame_with_vis(img_frame, model, labels_dict, cfg) |
|
||||
image_widget.value = img2bytes(image_pred) |
|
||||
|
|
||||
if __name__ == "__main__": |
|
||||
cfg = { |
|
||||
'conf_thres': 0.4, # 模型置信度阈值，阈值越低，得到的预测框越多 |
|
||||
'iou_thres': 0.5, # IOU阈值，高于这个阈值的重叠预测框会被过滤掉 |
|
||||
'input_shape': [640, 640], # 模型输入尺寸 |
|
||||
} |
} |
||||
|
|
||||
model_path = 'yolo.om' |
@app.route('/') |
||||
label_path = './coco_names.txt' |
async def index(): |
||||
# 初始化推理模型 |
return await render_template('index.html', streams=streams.keys()) |
||||
model = InferSession(0, model_path) |
|
||||
labels_dict = get_labels_from_txt(label_path) |
async def generate_frames(stream_id): |
||||
|
video_capture = streams[stream_id] |
||||
infer_mode = 'video' |
while True: |
||||
|
frame = video_capture.read() |
||||
if infer_mode == 'image': |
if frame is not None: |
||||
img_path = 'world_cup.jpg' |
_, buffer = cv2.imencode('.jpg', frame) |
||||
infer_image(img_path, model, labels_dict, cfg) |
frame = buffer.tobytes() |
||||
elif infer_mode == 'camera': |
yield (b'--frame\r\n' |
||||
infer_camera(model, labels_dict, cfg) |
b'Content-Type: image/jpeg\r\n\r\n' + frame + b'\r\n') |
||||
elif infer_mode == 'video': |
|
||||
video_path = 'racing.mp4' |
@app.route('/video_feed/<stream_id>') |
||||
infer_video(video_path, model, labels_dict, cfg) |
async def video_feed(stream_id): |
||||
|
return Response(generate_frames(stream_id), |
||||
|
mimetype='multipart/x-mixed-replace; boundary=frame') |
||||
|
|
||||
|
if __name__ == '__main__': |
||||
|
app.run(host='0.0.0.0', port=5000) |
||||
|
|
||||
|
|||||
@ -0,0 +1,49 @@ |
|||||
|
<!DOCTYPE html> |
||||
|
<html lang="en"> |
||||
|
<head> |
||||
|
<meta charset="UTF-8"> |
||||
|
<meta name="viewport" content="width=device-width, initial-scale=1.0"> |
||||
|
<title>YOLOv5 Video Streams</title> |
||||
|
<style> |
||||
|
body { |
||||
|
display: flex; |
||||
|
justify-content: center; |
||||
|
align-items: center; |
||||
|
height: 100vh; |
||||
|
margin: 0; |
||||
|
background-color: #f0f0f0; |
||||
|
} |
||||
|
.grid-container { |
||||
|
display: grid; |
||||
|
grid-template-columns: repeat(3, 1fr); |
||||
|
grid-template-rows: repeat(3, 1fr); |
||||
|
gap: 10px; |
||||
|
width: 90vw; |
||||
|
height: 90vh; |
||||
|
} |
||||
|
.grid-item { |
||||
|
position: relative; |
||||
|
width: 100%; |
||||
|
padding-top: 75%; /* 4:3 aspect ratio */ |
||||
|
} |
||||
|
.grid-item img { |
||||
|
position: absolute; |
||||
|
top: 0; |
||||
|
left: 0; |
||||
|
width: 100%; |
||||
|
height: 100%; |
||||
|
object-fit: cover; |
||||
|
border-radius: 8px; |
||||
|
} |
||||
|
</style> |
||||
|
</head> |
||||
|
<body> |
||||
|
<div class="grid-container"> |
||||
|
{% for stream_id in streams %} |
||||
|
<div class="grid-item"> |
||||
|
<img src="{{ url_for('video_feed', stream_id=stream_id) }}" alt="Stream {{ stream_id }}"> |
||||
|
</div> |
||||
|
{% endfor %} |
||||
|
</div> |
||||
|
</body> |
||||
|
</html> |
||||
@ -1,12 +0,0 @@ |
|||||
<!DOCTYPE html> |
|
||||
<html lang="en"> |
|
||||
<head> |
|
||||
<meta charset="UTF-8"> |
|
||||
<meta name="viewport" content="width=device-width, initial-scale=1.0"> |
|
||||
<title>YOLOv5 Real-Time Detection</title> |
|
||||
</head> |
|
||||
<body> |
|
||||
<h1>YOLOv5 Real-Time Detection</h1> |
|
||||
<img src="{{ url_for('api.video_feed') }}" width="640" height="480"> |
|
||||
</body> |
|
||||
</html> |
|
||||
Binary file not shown.
Loading…
Reference in new issue