97 lines
3.1 KiB
Python
97 lines
3.1 KiB
Python
# -*- coding: utf-8 -*-
|
|
"""
|
|
离线视频检测与告警验证脚本
|
|
|
|
用法示例：
|
|
python test_video.py --video sample.mp4
|
|
python test_video.py --video sample.mp4 --output result.mp4
|
|
|
|
按下 `q` 可中途退出。
|
|
"""
|
|
|
|
import argparse
|
|
import cv2
|
|
from datetime import datetime
|
|
from fish_tracker import FishTracker
|
|
from history_manager import FishHistoryManager
|
|
|
|
|
|
def run(video_path: str, output_path: str | None, report_path: str | None):
|
|
tracker = FishTracker()
|
|
history = FishHistoryManager()
|
|
|
|
cap = cv2.VideoCapture(video_path)
|
|
if not cap.isOpened():
|
|
raise RuntimeError(f"无法打开视频文件：{video_path}")
|
|
|
|
writer = None
|
|
if output_path:
|
|
fourcc = cv2.VideoWriter_fourcc(*"mp4v")
|
|
fps = cap.get(cv2.CAP_PROP_FPS) or 25
|
|
w = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
|
|
h = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
|
|
writer = cv2.VideoWriter(output_path, fourcc, fps, (w, h))
|
|
|
|
try:
|
|
while True:
|
|
ret, frame = cap.read()
|
|
# 读取失败：可能是到结尾，也可能是坏帧，先尝试继续几次
|
|
if not ret or frame is None:
|
|
# 到达文件末尾：直接退出
|
|
if not ret:
|
|
break
|
|
# 否则跳过本帧
|
|
print("警告：读取到空帧，跳过。")
|
|
continue
|
|
|
|
# 防御性检查：帧尺寸异常也跳过
|
|
h, w = frame.shape[:2]
|
|
if h == 0 or w == 0:
|
|
print("警告：帧尺寸异常，跳过。")
|
|
continue
|
|
|
|
# 个别帧解码失败时，YOLO/DeepSORT 可能抛异常，这里捕获并跳过
|
|
try:
|
|
output_frame, alerts, track_stats = tracker.process_frame(frame)
|
|
except Exception as e:
|
|
print(f"警告：处理当前帧时出错，已跳过。错误信息：{e}")
|
|
continue
|
|
|
|
now = datetime.now()
|
|
history.update(track_stats, now)
|
|
|
|
# 显示结果
|
|
if alerts:
|
|
print("; ".join(alerts))
|
|
cv2.imshow("fish monitor (press q to exit)", output_frame)
|
|
|
|
if writer:
|
|
writer.write(output_frame)
|
|
|
|
if cv2.waitKey(1) & 0xFF == ord("q"):
|
|
break
|
|
finally:
|
|
cap.release()
|
|
if writer:
|
|
writer.release()
|
|
if report_path:
|
|
# 结束后生成一次离线报告
|
|
out = history.generate_report()
|
|
print(f"报告已生成：{out}")
|
|
cv2.destroyAllWindows()
|
|
|
|
|
|
def parse_args():
|
|
parser = argparse.ArgumentParser(description="离线视频检测验证脚本")
|
|
parser.add_argument("--video", required=True, help="待检测的视频文件路径")
|
|
parser.add_argument("--output", help="可选，保存标注结果的视频路径")
|
|
parser.add_argument("--report", action="store_true", help="视频结束后生成报告（reports/ 下）")
|
|
return parser.parse_args()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
args = parse_args()
|
|
report_path = "reports/fish_report_offline.txt" if args.report else None
|
|
run(args.video, args.output, report_path)
|
|
|