155 lines
5.6 KiB
Python
155 lines
5.6 KiB
Python
|
import sophon_chakcy.sail as sail
|
||
|
|
import threading
|
||
|
|
import time
|
||
|
|
|
||
|
|
|
||
|
|
class DetectVideoStream:
|
||
|
|
def __init__(self,
|
||
|
|
detectors,
|
||
|
|
video_path: str,
|
||
|
|
max_frames: int = 0,
|
||
|
|
frame_skip: int = 0,
|
||
|
|
reconnect_attempts: int = 3,
|
||
|
|
reconnect_delay: int = 5,
|
||
|
|
is_save: bool = False):
|
||
|
|
"""
|
||
|
|
处理视频流（视频文件或RTSP流）
|
||
|
|
|
||
|
|
Args:
|
||
|
|
detector: 检测器
|
||
|
|
video_path: 视频文件路径或 RTSP URL
|
||
|
|
max_frames: 最大处理帧数, 0表示处理所有帧
|
||
|
|
frame_skip: 跳帧处理, 0表示处理每一帧
|
||
|
|
reconnect_attempts: RTSP断流重连尝试次数
|
||
|
|
reconnect_delay: 重连延迟时间 (秒)
|
||
|
|
|
||
|
|
Returns:
|
||
|
|
处理的帧数, 检测到的物体总数
|
||
|
|
"""
|
||
|
|
self.detectors = detectors
|
||
|
|
self.dev_id = self.detectors[0].dev_id
|
||
|
|
self.handle = self.detectors[0].handle
|
||
|
|
self.video_path = video_path
|
||
|
|
self.max_frames = max_frames
|
||
|
|
self.frame_skip = frame_skip
|
||
|
|
if not video_path.startswith("rtsp"):
|
||
|
|
self.reconnect_attempts = 0
|
||
|
|
else:
|
||
|
|
self.reconnect_attempts = reconnect_attempts
|
||
|
|
self.reconnect_delay = reconnect_delay
|
||
|
|
self.decoder = sail.Decoder(video_path, True, self.dev_id)
|
||
|
|
self.stop_event = threading.Event()
|
||
|
|
self.is_save = is_save
|
||
|
|
if self.is_save:
|
||
|
|
import os
|
||
|
|
from pathlib import Path
|
||
|
|
import shutil
|
||
|
|
|
||
|
|
result_dir = r"results"
|
||
|
|
# 清理旧的目录
|
||
|
|
if os.path.exists(result_dir):
|
||
|
|
print(f"清理旧的构建目录: {result_dir}")
|
||
|
|
shutil.rmtree(result_dir, ignore_errors=True)
|
||
|
|
|
||
|
|
# 创建目录
|
||
|
|
result_dir_path = Path(result_dir)
|
||
|
|
if not result_dir_path.exists():
|
||
|
|
result_dir_path.mkdir(parents=True, exist_ok=True)
|
||
|
|
|
||
|
|
enc_params = f"width=1920:height=1080:bitrate=2000:gop=32:gop_preset=2:framerate=25"
|
||
|
|
self.encoder = sail.Encoder("results/output.mp4", self.dev_id, 'h264_bm', 'NV12', enc_params, 10)
|
||
|
|
|
||
|
|
def run_task(self):
|
||
|
|
if not self.decoder.is_opened():
|
||
|
|
print(f"错误: 无法打开视频流 {self.video_path}")
|
||
|
|
return 0, 0
|
||
|
|
|
||
|
|
frame_shape = self.decoder.get_frame_shape()
|
||
|
|
print(f"视频帧形状: {frame_shape}")
|
||
|
|
|
||
|
|
frame_count = 0
|
||
|
|
reconnect_count = 0
|
||
|
|
|
||
|
|
try:
|
||
|
|
while not self.stop_event.is_set():
|
||
|
|
if self.max_frames > 0 and frame_count >= self.max_frames:
|
||
|
|
print(f"\n达到最大帧数限制: {self.max_frames}")
|
||
|
|
break
|
||
|
|
|
||
|
|
# 读取一帧
|
||
|
|
bmimg = sail.BMImage()
|
||
|
|
ret = self.decoder.read(self.handle, bmimg)
|
||
|
|
if ret != 0:
|
||
|
|
print(f"读取帧失败，错误码: {ret}")
|
||
|
|
reconnect_count += 1
|
||
|
|
|
||
|
|
if reconnect_count > self.reconnect_attempts:
|
||
|
|
print("视频流无法继续读取")
|
||
|
|
break
|
||
|
|
|
||
|
|
print(f"尝试重连 ({reconnect_count}/{self.reconnect_attempts})...")
|
||
|
|
time.sleep(self.reconnect_delay)
|
||
|
|
self.decoder.release()
|
||
|
|
decoder = sail.Decoder(self.video_path, True, self.dev_id)
|
||
|
|
if decoder.is_opened():
|
||
|
|
reconnect_count = 0
|
||
|
|
print("重连成功，继续处理...")
|
||
|
|
continue
|
||
|
|
else:
|
||
|
|
print("重连失败")
|
||
|
|
|
||
|
|
if self.frame_skip > 0 and frame_count % (self.frame_skip + 1) != 0:
|
||
|
|
frame_count += 1
|
||
|
|
if self.is_save:
|
||
|
|
self.encoder.video_write(bmimg)
|
||
|
|
continue
|
||
|
|
|
||
|
|
# 开始处理当前帧
|
||
|
|
start_time = time.time()
|
||
|
|
frame_count += 1
|
||
|
|
|
||
|
|
if frame_count % 10 == 0:
|
||
|
|
print(f"处理第 {frame_count} 帧...")
|
||
|
|
|
||
|
|
# 检测
|
||
|
|
for detector in self.detectors:
|
||
|
|
detections = detector.detect(bmimg)
|
||
|
|
|
||
|
|
if len(detections) > 0:
|
||
|
|
detections = detections[0]
|
||
|
|
|
||
|
|
filtered_detections = []
|
||
|
|
for det in detections:
|
||
|
|
if len(det) >= 6:
|
||
|
|
x1, y1, x2, y2, conf, cls_id = det[:6]
|
||
|
|
if conf > detector.conf_thresh:
|
||
|
|
filtered_detections.append([x1, y1, x2, y2, conf, cls_id])
|
||
|
|
|
||
|
|
if len(filtered_detections) > 0:
|
||
|
|
result_img = detector.draw_detections(bmimg, filtered_detections, None)
|
||
|
|
else:
|
||
|
|
result_img = bmimg
|
||
|
|
else:
|
||
|
|
result_img = bmimg
|
||
|
|
|
||
|
|
if self.is_save:
|
||
|
|
self.encoder.video_write(result_img)
|
||
|
|
|
||
|
|
except KeyboardInterrupt:
|
||
|
|
print("\n检测到键盘中断，停止处理")
|
||
|
|
except Exception as e:
|
||
|
|
print(f"处理视频流时发生错误: {e}")
|
||
|
|
|
||
|
|
finally:
|
||
|
|
# 清理资源
|
||
|
|
self.decoder.release()
|
||
|
|
if self.is_save:
|
||
|
|
self.encoder.release()
|
||
|
|
print("视频解码器已释放")
|
||
|
|
|
||
|
|
def stop_task(self):
|
||
|
|
print("提前关闭")
|
||
|
|
self.stop_event.set()
|
||
|
|
|
||
|
|
|