cxc/image_identification
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

push streame

Browse Source
This commit is contained in:
quantulr
2023-10-07 14:13:49 +08:00
parent aee62b9e1e
commit c1cb289d44
2 changed files with 64 additions and 29 deletions
Download Patch File Download Diff File Expand all files Collapse all files

92
image_identification/controllers/image.py
View File

@ -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:

1
image_identification/main.py
View File

@ -42,6 +42,7 @@ async def analyze_video():
ret, frame = cap.read()
if not ret:
print("无法读取视频帧")
os.chdir("..")
break
# 是否保存帧
if True: