diff --git a/image_identification/controllers/image.py b/image_identification/controllers/image.py index e4dfe87..ccf516d 100644 --- a/image_identification/controllers/image.py +++ b/image_identification/controllers/image.py @@ -11,7 +11,7 @@ import cv2 import fastdeploy as fd import numpy as np from fastapi import APIRouter, HTTPException, UploadFile, WebSocket, WebSocketDisconnect -from fastapi.responses import StreamingResponse +from fastapi.responses import RedirectResponse, StreamingResponse from fastdeploy.serving.utils import cv2_to_base64 from image_identification.model import model_instance @@ -60,19 +60,16 @@ async def analyze(image: UploadFile): @router.websocket("/analyze_video") async def analyze_video(websocket: WebSocket): rtmp_url = "rtmp://192.168.0.136/live/livestream" - output_url = "rtmp://192.168.0.136/live/livestream_output" + push_url = "rtmp://192.168.0.136/live/livestream_output" + output_url = "http://192.168.0.136/live/livestream_output.flv" 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") - # os.mkdir("framers") - print(cap.get(3)) - print(cap.get(4)) - os.chdir("framers") - video = None + # fourcc = cv2.VideoWriter.fourcc("M", "P", "4", "V") + # os.chdir("framers") + # video = None ffmpeg_command = None ffmpeg_process = None try: @@ -82,24 +79,22 @@ async def analyze_video(websocket: WebSocket): if not ret: print("无法读取视频帧") break - # 是否保存帧 - if True: - # if frame_number % 4 == 0: - # 处理每一帧的操作,这里只是保存为图像文件 - # cv2.imwrite(frame_filename, frame) + print(f"\n第{frame_number}帧") + # 是否保存帧 + # if True: + if frame_number % 4 == 0: 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, - (width, height), - ) + # if video is None: + # video = cv2.VideoWriter( + # "__output__.mp4", + # fourcc, + # 30.0, + # (width, height), + # ) if ffmpeg_command is None: # 创建ffmpeg进程以推送视频流 ffmpeg_cmd = [ @@ -125,21 +120,60 @@ async def analyze_video(websocket: WebSocket): "ultrafast", "-f", "flv", - output_url, + push_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(f"第{frame_number}帧") + # video.write(vis_im) print(result) - await websocket.send_text(fd.vision.fd_result_to_json(result=result)) - # await websocket.send_text(frame_filename) + await websocket.send_json( + data={ + "output": output_url, + "result": json.loads( + fd.vision.fd_result_to_json(result=result) + ), + } + ) else: - print(frame.shape) - video.write(frame) + width = frame.shape[1] + height = frame.shape[0] + 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", + push_url, + ] + if ffmpeg_process is None: + ffmpeg_process = subprocess.Popen(ffmpeg_cmd, stdin=subprocess.PIPE) + ffmpeg_process.stdin.write(frame.tobytes()) + + await websocket.send_json(data={"output": output_url}) + # video.write(frame) end = time.perf_counter() cost = end - start + print(fps) print(f"time spend {cost}") frame_number += 1 except WebSocketDisconnect: diff --git a/image_identification/main.py b/image_identification/main.py index b2d2fb3..4784c36 100644 --- a/image_identification/main.py +++ b/image_identification/main.py @@ -42,6 +42,7 @@ async def analyze_video(): ret, frame = cap.read() if not ret: print("无法读取视频帧") + os.chdir("..") break # 是否保存帧 if True: