259 lines
8.7 KiB
Python
259 lines
8.7 KiB
Python
import cv2
|
||
import subprocess
|
||
from threading import Thread, Lock, Event
|
||
import time
|
||
import queue
|
||
import numpy as np
|
||
from ultralytics import YOLO
|
||
import datetime
|
||
import os
|
||
import uuid
|
||
|
||
# ---------- 会话管理器 ----------
|
||
stream_sessions = {}
|
||
|
||
|
||
def start_video_session(video_path, output_url, model_path, cls, confidence,cls2=[],push = False):
|
||
"""
|
||
启动一个视频流处理会话，返回 session_id。
|
||
"""
|
||
session_id = str(uuid.uuid4())
|
||
stop_event = Event()
|
||
frame_q = queue.Queue(maxsize=5)
|
||
processed_q = queue.Queue(maxsize=5)
|
||
model_container = {
|
||
'model': YOLO(model_path),
|
||
'lock': Lock()
|
||
}
|
||
|
||
# 启动后台线程
|
||
thread = Thread(
|
||
target=stream_worker,
|
||
args= (video_path, output_url, cls, confidence, stop_event, frame_q, processed_q, model_container,cls2,push),
|
||
daemon=True
|
||
)
|
||
thread.start()
|
||
|
||
stream_sessions[session_id] = {
|
||
'thread': thread,
|
||
'stop_event': stop_event,
|
||
'model_container': model_container
|
||
}
|
||
print(f"会话 {session_id} 已启动，视频源: {video_path}")
|
||
return session_id
|
||
|
||
|
||
def stop_video_session(session_id):
|
||
"""
|
||
停止指定会话。
|
||
"""
|
||
sess = stream_sessions.get(session_id)
|
||
if not sess:
|
||
print(f"会话 {session_id} 不存在。")
|
||
return
|
||
sess['stop_event'].set()
|
||
sess['thread'].join()
|
||
print(f"会话 {session_id} 已停止。")
|
||
del stream_sessions[session_id]
|
||
|
||
|
||
def switch_model_session(session_id, new_model_path):
|
||
"""
|
||
动态切换指定会话的模型权重。
|
||
"""
|
||
sess = stream_sessions.get(session_id)
|
||
if not sess:
|
||
print(f"会话 {session_id} 不存在，无法切换模型。")
|
||
return
|
||
with sess['model_container']['lock']:
|
||
sess['model_container']['model'] = YOLO(new_model_path)
|
||
print(f"会话 {session_id} 模型已切换为: {new_model_path}")
|
||
|
||
def _read_frames(cap, frame_q, stop_event):
|
||
""" 读取帧，保持30fps，丢弃延迟帧，若连续2秒无有效帧则触发停止 """
|
||
last = time.time()
|
||
target = 1.0 / 30.0
|
||
no_frame_start = None # 开始无法读取帧的时间点
|
||
|
||
while not stop_event.is_set():
|
||
now = time.time()
|
||
if now - last < target:
|
||
time.sleep(target - (now - last))
|
||
continue
|
||
if frame_q.qsize() < frame_q.maxsize:
|
||
ret, frm = cap.read()
|
||
if not ret:
|
||
if no_frame_start is None:
|
||
no_frame_start = time.time()
|
||
elif time.time() - no_frame_start >= 2:
|
||
print("⚠️ 超过2秒未获取到帧，自动停止会话。")
|
||
stop_event.set()
|
||
break
|
||
continue
|
||
else:
|
||
no_frame_start = None # 成功读取帧，重置计时
|
||
frame_q.put((frm, time.time()))
|
||
last = time.time()
|
||
else:
|
||
time.sleep(0.001)
|
||
|
||
|
||
|
||
def _process_frames(frame_q, processed_q, cls, confidence, model_container, stop_event):
|
||
""" 检测并绘制 """
|
||
skip_counter = 0
|
||
while not stop_event.is_set():
|
||
try:
|
||
frm, ts = frame_q.get(timeout=1)
|
||
if time.time() - ts > 0.5:
|
||
continue
|
||
skip_counter = (skip_counter + 1) % 2
|
||
if skip_counter != 0:
|
||
continue
|
||
|
||
with model_container['lock']:
|
||
results = model_container['model'].track(frm, persist=True, classes=cls, conf=confidence)
|
||
counts = {}
|
||
for r in results:
|
||
if hasattr(r, 'boxes'):
|
||
for b in r.boxes:
|
||
cid = int(b.cls[0])
|
||
counts[cid] = counts.get(cid, 0) + 1
|
||
ann = results[0].plot(conf=False, line_width=1, font_size=1.5)
|
||
y_off = 30
|
||
for cid, cnt in counts.items():
|
||
txt = f"Class {cid}: {cnt}"
|
||
cv2.putText(ann, txt, (10, y_off), cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0,255,0), 2)
|
||
y_off += 25
|
||
processed_q.put((ann, ts))
|
||
except queue.Empty:
|
||
continue
|
||
except Exception as e:
|
||
print(f"_process_frames 错误: {e}")
|
||
|
||
|
||
def _write_frames(processed_q, pipe, size, stop_event):
|
||
""" 推流写入 """
|
||
last = time.time()
|
||
target = 1.0 / 30.0
|
||
while not stop_event.is_set():
|
||
try:
|
||
frm, ts = processed_q.get(timeout=1)
|
||
now = time.time()
|
||
if now - last < target:
|
||
time.sleep(target - (now - last))
|
||
img = cv2.resize(frm, size, interpolation=cv2.INTER_LINEAR)
|
||
pipe.stdin.write(img.tobytes())
|
||
last = time.time()
|
||
except queue.Empty:
|
||
continue
|
||
except Exception as e:
|
||
print(f"_write_frames 错误: {e}")
|
||
break
|
||
|
||
|
||
def _cls2_find(video_path, confidence, save_dir, stop_event):
|
||
"""
|
||
指定类别检测并上传数据库，检测到每个跟踪目标只保存一次后续忽略。
|
||
"""
|
||
model = YOLO("gdaq.pt")
|
||
cls2 = [2, 4]
|
||
cap = cv2.VideoCapture(video_path)
|
||
fps = cap.get(cv2.CAP_PROP_FPS) or 30
|
||
|
||
# 存储已经保存过的 track_id
|
||
seen_tracks = set()
|
||
|
||
while cap.isOpened() and not stop_event.is_set():
|
||
ret, frame = cap.read()
|
||
if not ret:
|
||
break
|
||
|
||
# 使用 track 来获取每个目标的 track_id
|
||
results = model.track(frame, persist=True, classes=cls2, conf=confidence, show=False)
|
||
for res in results:
|
||
for box in res.boxes:
|
||
tid = int(box.id[0]) # 跟踪ID
|
||
cls_id = int(box.cls[0])
|
||
if cls_id in cls2 and tid not in seen_tracks:
|
||
seen_tracks.add(tid)
|
||
fn = os.path.join(
|
||
save_dir,
|
||
f"track_{tid}_frame_{int(cap.get(cv2.CAP_PROP_POS_FRAMES))}.jpg"
|
||
)
|
||
cv2.imwrite(fn, frame)
|
||
print(f"保存: {fn}")
|
||
# TODO: 在这里做上传和数据库写入
|
||
# 小睡一会，防止 CPU 占用过高
|
||
time.sleep(0.01)
|
||
|
||
cap.release()
|
||
|
||
|
||
# ---------- 核心处理逻辑 ----------
|
||
|
||
def stream_worker(video_path, output_url, cls, confidence, stop_event, frame_q, processed_q, model_container, cls2=[],push = False):
|
||
"""
|
||
会话主函数：打开流、预热模型、启动子线程、等待停止。
|
||
"""
|
||
# 创建保存目录（cls2 功能）
|
||
date_str = datetime.datetime.now().strftime("%Y%m%d")
|
||
save_dir = f"AIResults_{date_str}"
|
||
os.makedirs(save_dir, exist_ok=True)
|
||
|
||
# 打开视频流
|
||
cap = cv2.VideoCapture(video_path, cv2.CAP_FFMPEG)
|
||
cap.set(cv2.CAP_PROP_BUFFERSIZE, 2)
|
||
cap.set(cv2.CAP_PROP_FOURCC, cv2.VideoWriter_fourcc(*'H264'))
|
||
cap.set(cv2.CAP_PROP_FRAME_WIDTH, 1920)
|
||
cap.set(cv2.CAP_PROP_FRAME_HEIGHT, 1080)
|
||
|
||
# 预加载并预热模型
|
||
print(f"[{video_path}] 预加载YOLO模型...")
|
||
with model_container['lock']:
|
||
dummy = np.zeros((1080, 1920, 3), dtype=np.uint8)
|
||
model_container['model'].track(dummy, classes=cls, conf=confidence, show=False)
|
||
|
||
# 构造 FFmpeg 推流命令
|
||
size = (1920, 1080)
|
||
size_str = f"{size[0]}x{size[1]}"
|
||
command = [
|
||
'ffmpeg', '-y', '-f', 'rawvideo', '-vcodec', 'rawvideo',
|
||
'-pix_fmt', 'bgr24', '-s', size_str, '-r', '30', '-probesize', '32', '-analyzeduration', '0',
|
||
'-i', '-', '-c:v', 'h264', '-pix_fmt', 'yuv420p',
|
||
'-preset', 'ultrafast', '-tune', 'zerolatency', '-f', 'flv',
|
||
'-g', '30', '-bufsize', '1000k', '-maxrate', '2000k',
|
||
'-x264opts', 'no-scenecut:keyint=30:min-keyint=30',
|
||
'-flvflags', 'no_duration_filesize', output_url
|
||
]
|
||
pipe = subprocess.Popen(command, stdin=subprocess.PIPE)
|
||
|
||
# 启动子线程
|
||
threads = []
|
||
threads.append(Thread(target=_read_frames, args=(cap, frame_q, stop_event), daemon=True))
|
||
threads.append(Thread(target=_process_frames, args=(frame_q, processed_q, cls, confidence, model_container, stop_event), daemon=True))
|
||
if push:
|
||
threads.append(Thread(target=_write_frames, args=(processed_q, pipe, size, stop_event), daemon=True))
|
||
|
||
if cls2:
|
||
threads.append(Thread(target=_cls2_find, args=(video_path, confidence, save_dir, stop_event), daemon=True))
|
||
|
||
for t in threads:
|
||
t.start()
|
||
print(f"[{video_path}] 开始处理流...")
|
||
|
||
# 等待停止信号
|
||
while not stop_event.is_set() and cap.isOpened():
|
||
time.sleep(0.1)
|
||
|
||
# 清理
|
||
print(f"[{video_path}] 停止中，正在清理...")
|
||
cap.release()
|
||
pipe.stdin.close()
|
||
pipe.wait()
|
||
for q in (frame_q, processed_q):
|
||
while not q.empty():
|
||
try: q.get_nowait()
|
||
except queue.Empty: break
|
||
print(f"[{video_path}] 会话结束。")
|