From aee62b9e1e74755e8c11f80a877981b8c264fbf0 Mon Sep 17 00:00:00 2001 From: quantulr <35954003+quantulr@users.noreply.github.com> Date: Thu, 28 Sep 2023 16:11:57 +0800 Subject: [PATCH] push streame --- image_identification/controllers/image.py | 48 +++++++++++++++++++++-- image_identification/model/model.py | 4 +- poetry.lock | 14 +++---- pyproject.toml | 3 +- 4 files changed, 56 insertions(+), 13 deletions(-) diff --git a/image_identification/controllers/image.py b/image_identification/controllers/image.py index 2ae57dd..e4dfe87 100644 --- a/image_identification/controllers/image.py +++ b/image_identification/controllers/image.py @@ -3,6 +3,7 @@ import io import json import os import pathlib +import subprocess import time import aiohttp @@ -58,10 +59,12 @@ async def analyze(image: UploadFile): @router.websocket("/analyze_video") async def analyze_video(websocket: WebSocket): - rtmp_url = "rtmp://localhost/live/livestream" + rtmp_url = "rtmp://192.168.0.136/live/livestream" + output_url = "rtmp://192.168.0.136/live/livestream_output" cap = cv2.VideoCapture(rtmp_url) frame_number = 0 frame_rate = 30 + fps = int(cap.get(cv2.CAP_PROP_FPS)) await websocket.accept() fourcc = cv2.VideoWriter.fourcc("M", "P", "4", "V") @@ -70,8 +73,11 @@ async def analyze_video(websocket: WebSocket): print(cap.get(4)) os.chdir("framers") video = None + ffmpeg_command = None + ffmpeg_process = None try: while True: + start = time.perf_counter() ret, frame = cap.read() if not ret: print("无法读取视频帧") @@ -85,20 +91,56 @@ async def analyze_video(websocket: WebSocket): result = model_instance.predict(frame) vis_im = fd.vision.vis_detection(frame, result, score_threshold=0.5) # cv2.imwrite(frame_filename, vis_im) + width = vis_im.shape[1] + height = vis_im.shape[0] if video is None: video = cv2.VideoWriter( "__output__.mp4", fourcc, 30.0, - (vis_im.shape[1], vis_im.shape[0]), + (width, height), ) + if ffmpeg_command is None: + # 创建ffmpeg进程以推送视频流 + ffmpeg_cmd = [ + "ffmpeg", + "-y", + "-f", + "rawvideo", + "-vcodec", + "rawvideo", + "-s", + f"{width}x{height}", + "-pix_fmt", + "bgr24", + "-r", + str(fps), + "-i", + "-", + "-c:v", + "libx264", + "-pix_fmt", + "yuv420p", + "-preset", + "ultrafast", + "-f", + "flv", + output_url, + ] + if ffmpeg_process is None: + ffmpeg_process = subprocess.Popen(ffmpeg_cmd, stdin=subprocess.PIPE) + ffmpeg_process.stdin.write(vis_im.tobytes()) video.write(vis_im) - print(vis_im.shape) + print(f"第{frame_number}帧") + print(result) await websocket.send_text(fd.vision.fd_result_to_json(result=result)) # await websocket.send_text(frame_filename) else: print(frame.shape) video.write(frame) + end = time.perf_counter() + cost = end - start + print(f"time spend {cost}") frame_number += 1 except WebSocketDisconnect: print("close") diff --git a/image_identification/model/model.py b/image_identification/model/model.py index e53ff5d..2e22ae3 100644 --- a/image_identification/model/model.py +++ b/image_identification/model/model.py @@ -4,8 +4,8 @@ import fastdeploy as fd # Configurations model_dir = "yolov5s_infer" -device = "gpu" -use_trt = True +device = "cpu" +use_trt = False # Prepare model model_file = os.path.join(model_dir, "model.pdmodel") diff --git a/poetry.lock b/poetry.lock index e989024..83e3130 100644 --- a/poetry.lock +++ b/poetry.lock @@ -420,13 +420,13 @@ url = "https://pypi.tuna.tsinghua.edu.cn/simple" reference = "mirrors" [[package]] -name = "fastdeploy-gpu-python" +name = "fastdeploy-python" version = "1.0.7" description = "Deploy Kit Tool For Deeplearning models." optional = false python-versions = "*" files = [ - {file = "fastdeploy_gpu_python-1.0.7-cp310-cp310-manylinux1_x86_64.whl", hash = "sha256:f92e9414e69b56b092d1327118912e8c5fa0acc32a0d9f4e971145e888f8d72d"}, + {file = "fastdeploy_python-1.0.7-cp39-cp39-macosx_10_14_x86_64.whl", hash = "sha256:6f1b9c6eeefed4b63b9ba28335d857d6dd4fd9b6e9090a6e9688c60b553a4c42"}, ] [package.dependencies] @@ -444,7 +444,7 @@ mypy = ["mypy (==0.600)"] [package.source] type = "url" -url = "https://bj.bcebos.com/fastdeploy/release/wheels/fastdeploy_gpu_python-1.0.7-cp310-cp310-manylinux1_x86_64.whl" +url = "https://bj.bcebos.com/fastdeploy/release/wheels/fastdeploy_python-1.0.7-cp39-cp39-macosx_10_14_x86_64.whl" [[package]] name = "fastdeploy-tools" @@ -737,13 +737,13 @@ reference = "mirrors" [[package]] name = "pydantic" -version = "2.4.1" +version = "2.4.2" description = "Data validation using Python type hints" optional = false python-versions = ">=3.7" files = [ - {file = "pydantic-2.4.1-py3-none-any.whl", hash = "sha256:2b2240c8d54bb8f84b88e061fac1bdfa1761c2859c367f9d3afe0ec2966deddc"}, - {file = "pydantic-2.4.1.tar.gz", hash = "sha256:b172505886028e4356868d617d2d1a776d7af1625d1313450fd51bdd19d9d61f"}, + {file = "pydantic-2.4.2-py3-none-any.whl", hash = "sha256:bc3ddf669d234f4220e6e1c4d96b061abe0998185a8d7855c0126782b7abc8c1"}, + {file = "pydantic-2.4.2.tar.gz", hash = "sha256:94f336138093a5d7f426aac732dcfe7ab4eb4da243c88f891d65deb4a2556ee7"}, ] [package.dependencies] @@ -1336,4 +1336,4 @@ reference = "mirrors" [metadata] lock-version = "2.0" python-versions = ">=3.9,<3.13" -content-hash = "f1ddca71df771166e502ffab7d603041816df3290b39513a7e5bacf7cb343f8a" +content-hash = "b1d0db5e3e467c63b7794237b977baadd64738b17490597f63dba35ded184000" diff --git a/pyproject.toml b/pyproject.toml index 9429b83..24bc50c 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -7,7 +7,8 @@ readme = "README.md" [tool.poetry.dependencies] python = ">=3.9,<3.13" -fastdeploy-gpu-python = { url = "https://bj.bcebos.com/fastdeploy/release/wheels/fastdeploy_gpu_python-1.0.7-cp310-cp310-manylinux1_x86_64.whl" } +# fastdeploy-gpu-python = { url = "https://bj.bcebos.com/fastdeploy/release/wheels/fastdeploy_gpu_python-1.0.7-cp310-cp310-manylinux1_x86_64.whl" } +fastdeploy-python = { url = "https://bj.bcebos.com/fastdeploy/release/wheels/fastdeploy_python-1.0.7-cp39-cp39-macosx_10_14_x86_64.whl" } numpy = "^1.26.0" opencv-python = "^4.8.0.76" python-multipart = "^0.0.6"