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 os |
|||
# from model.plugins.ModelBase import ModelBase |
|||
# import importlib.util |
|||
import threading |
|||
import time |
|||
import numpy as np |
|||
import ipywidgets as widgets |
|||
from IPython.display import display |
|||
import torch |
|||
from skvideo.io import vreader, FFmpegWriter |
|||
import IPython.display |
|||
from ais_bench.infer.interface import InferSession |
|||
|
|||
from det_utils import letterbox, scale_coords, nms |
|||
|
|||
def preprocess_image(image, cfg, bgr2rgb=True): |
|||
"""图片预处理""" |
|||
img, scale_ratio, pad_size = letterbox(image, new_shape=cfg['input_shape']) |
|||
if bgr2rgb: |
|||
img = img[:, :, ::-1] |
|||
img = img.transpose(2, 0, 1) # HWC2CHW |
|||
img = np.ascontiguousarray(img, dtype=np.float32) |
|||
return img, scale_ratio, pad_size |
|||
|
|||
|
|||
def draw_bbox(bbox, img0, color, wt, names): |
|||
"""在图片上画预测框""" |
|||
det_result_str = '' |
|||
for idx, class_id in enumerate(bbox[:, 5]): |
|||
if float(bbox[idx][4] < float(0.05)): |
|||
continue |
|||
img0 = cv2.rectangle(img0, (int(bbox[idx][0]), int(bbox[idx][1])), (int(bbox[idx][2]), int(bbox[idx][3])), |
|||
color, wt) |
|||
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) |
|||
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]) |
|||
return img0 |
|||
|
|||
|
|||
def get_labels_from_txt(path): |
|||
"""从txt文件获取图片标签""" |
|||
labels_dict = dict() |
|||
with open(path) as f: |
|||
for cat_id, label in enumerate(f.readlines()): |
|||
labels_dict[cat_id] = label.strip() |
|||
return labels_dict |
|||
|
|||
|
|||
def draw_prediction(pred, image, labels): |
|||
"""在图片上画出预测框并进行可视化展示""" |
|||
imgbox = widgets.Image(format='jpg', height=720, width=1280) |
|||
img_dw = draw_bbox(pred, image, (0, 255, 0), 2, labels) |
|||
imgbox.value = cv2.imencode('.jpg', img_dw)[1].tobytes() |
|||
display(imgbox) |
|||
|
|||
|
|||
def infer_image(img_path, model, class_names, cfg): |
|||
"""图片推理""" |
|||
# 图片载入 |
|||
image = cv2.imread(img_path) |
|||
# 数据预处理 |
|||
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() |
|||
from quart import Quart, render_template, websocket, jsonify |
|||
from quart.helpers import Response |
|||
|
|||
# def test_acl_verify(): |
|||
# '''根据路径，动态导入模块''' |
|||
# model_path = "/mnt/zfbox/model/plugins/RYRQ_ACL/RYRQ_Model_ACL.py" |
|||
# model_name = "测试" |
|||
# if os.path.exists(model_path): |
|||
# module_spec = importlib.util.spec_from_file_location(model_name, model_path) |
|||
# if module_spec is None: |
|||
# print(f"{model_path} 加载错误") |
|||
# return None |
|||
# module = importlib.util.module_from_spec(module_spec) |
|||
# module_spec.loader.exec_module(module) |
|||
# md = getattr(module, "Model")(model_path,0.5) # 实例化类 |
|||
# if not isinstance(md, ModelBase) or not md.init_ok: |
|||
# print("{} not zf_model".format(md)) |
|||
# return None |
|||
# else: |
|||
# print("{}文件不存在".format(model_path)) |
|||
# return None |
|||
# print("开执行检测") |
|||
# #return md |
|||
# img = cv2.imread("/mnt/zfbox/model/plugins/RYRQ_ACL/world_cup.jpg", cv2.IMREAD_COLOR) # 读入图片 |
|||
# data = [0,1,"[(100, 100), (50, 200), (100, 300), (300, 300), (300, 100)]"] |
|||
# boxout,ret,strret = md.verify(img,data) |
|||
# #print(boxout) |
|||
# #释放相关资源 |
|||
# md.release() |
|||
|
|||
|
|||
|
|||
# 加载YOLOv5模型 |
|||
#model = yolov5.load('yolov5s') |
|||
|
|||
app = Quart(__name__) |
|||
|
|||
class VideoCaptureWithYOLO: |
|||
def __init__(self, source): |
|||
self.source = source |
|||
self.cap = cv2.VideoCapture(self._get_gstreamer_pipeline(), cv2.CAP_GSTREAMER) |
|||
self.width = int(self.cap.get(cv2.CAP_PROP_FRAME_WIDTH)) |
|||
self.height = int(self.cap.get(cv2.CAP_PROP_FRAME_HEIGHT)) |
|||
self.frame = None |
|||
self.running = True |
|||
self.read_lock = threading.Lock() |
|||
self.thread = threading.Thread(target=self.update) |
|||
self.thread.start() |
|||
|
|||
def _get_gstreamer_pipeline(self): |
|||
return ( |
|||
f"rtspsrc location={self.source} protocols=udp latency=0 ! " |
|||
"rtph264depay ! h264parse ! avdec_h264 ! videoconvert ! appsink" |
|||
) |
|||
|
|||
def update(self): |
|||
while self.running: |
|||
ret, frame = self.cap.read() |
|||
if not ret: |
|||
break |
|||
# 对视频帧进行推理 |
|||
image_pred = infer_frame_with_vis(img_frame, model, labels_dict, cfg, bgr2rgb=True) |
|||
image_widget.value = img2bytes(image_pred) |
|||
|
|||
|
|||
def infer_camera(model, labels_dict, cfg): |
|||
"""外设摄像头实时推理""" |
|||
def find_camera_index(): |
|||
max_index_to_check = 10 # Maximum index to check for camera |
|||
|
|||
for index in range(max_index_to_check): |
|||
cap = cv2.VideoCapture(index) |
|||
if cap.read()[0]: |
|||
cap.release() |
|||
return index |
|||
|
|||
# If no camera is found |
|||
raise ValueError("No camera found.") |
|||
|
|||
# 获取摄像头 |
|||
camera_index = find_camera_index() |
|||
cap = cv2.VideoCapture(camera_index) |
|||
# 初始化可视化对象 |
|||
image_widget = widgets.Image(format='jpeg', width=1280, height=720) |
|||
display(image_widget) |
|||
continue |
|||
# YOLOv5分析 |
|||
annotated_frame = frame |
|||
|
|||
with self.read_lock: |
|||
self.frame = annotated_frame |
|||
time.sleep(0.01) # 控制帧率 |
|||
|
|||
def read(self): |
|||
with self.read_lock: |
|||
frame = self.frame.copy() if self.frame is not None else None |
|||
return frame |
|||
|
|||
def release(self): |
|||
self.running = False |
|||
self.thread.join() |
|||
self.cap.release() |
|||
|
|||
# 创建多个视频流 |
|||
streams = { |
|||
'stream1': VideoCaptureWithYOLO('rtsp://192.168.3.44/live1'), |
|||
'stream2': VideoCaptureWithYOLO('rtsp://192.168.3.44/live1'), |
|||
'stream3': VideoCaptureWithYOLO('rtsp://192.168.3.44/live1'), |
|||
'stream4': VideoCaptureWithYOLO('rtsp://192.168.3.44/live1'), |
|||
'stream5': VideoCaptureWithYOLO('rtsp://192.168.3.44/live1'), |
|||
'stream6': VideoCaptureWithYOLO('rtsp://192.168.3.44/live1'), |
|||
'stream7': VideoCaptureWithYOLO('rtsp://192.168.3.44/live1'), |
|||
'stream8': VideoCaptureWithYOLO('rtsp://192.168.3.44/live1'), |
|||
'stream9': VideoCaptureWithYOLO('rtsp://192.168.3.44/live1') |
|||
} |
|||
|
|||
@app.route('/') |
|||
async def index(): |
|||
return await render_template('index.html', streams=streams.keys()) |
|||
|
|||
async def generate_frames(stream_id): |
|||
video_capture = streams[stream_id] |
|||
while True: |
|||
# 对摄像头每一帧进行推理和可视化 |
|||
_, img_frame = cap.read() |
|||
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' |
|||
label_path = './coco_names.txt' |
|||
# 初始化推理模型 |
|||
model = InferSession(0, model_path) |
|||
labels_dict = get_labels_from_txt(label_path) |
|||
|
|||
infer_mode = 'video' |
|||
|
|||
if infer_mode == 'image': |
|||
img_path = 'world_cup.jpg' |
|||
infer_image(img_path, model, labels_dict, cfg) |
|||
elif infer_mode == 'camera': |
|||
infer_camera(model, labels_dict, cfg) |
|||
elif infer_mode == 'video': |
|||
video_path = 'racing.mp4' |
|||
infer_video(video_path, model, labels_dict, cfg) |
|||
frame = video_capture.read() |
|||
if frame is not None: |
|||
_, buffer = cv2.imencode('.jpg', frame) |
|||
frame = buffer.tobytes() |
|||
yield (b'--frame\r\n' |
|||
b'Content-Type: image/jpeg\r\n\r\n' + frame + b'\r\n') |
|||
|
|||
@app.route('/video_feed/<stream_id>') |
|||
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