dsh/dual_camera_system/capture_and_detect.py

"""
抓取一枪机实时帧，使用当前配置的人体检测模型检测并保存标记图。
用于在 RK3588 上快速验证模型输出是否正常。
"""

import os
import sys
import cv2
import time
import numpy as np
from pathlib import Path

sys.path.insert(0, str(Path(__file__).parent))

from panorama_camera import ObjectDetector
from config import DETECTION_CONFIG

RTSP_URL = 'rtsp://admin:QAZwsx12@192.168.8.2:554/Streaming/Channels/101'
OUT_DIR = Path('/home/admin/dsh/detection_check')
OUT_DIR.mkdir(parents=True, exist_ok=True)


def main():
    print(f"加载模型: {DETECTION_CONFIG['model_path']}")
    detector = ObjectDetector(
        model_path=DETECTION_CONFIG['model_path'],
        use_gpu=DETECTION_CONFIG.get('use_gpu', True),
        model_type=DETECTION_CONFIG.get('model_type', 'auto')
    )

    print(f"打开 RTSP: {RTSP_URL}")
    cap = cv2.VideoCapture(RTSP_URL, cv2.CAP_FFMPEG)
    cap.set(cv2.CAP_PROP_BUFFERSIZE, 1)
    if not cap.isOpened():
        print("无法打开 RTSP 流")
        return

    # 读取几帧清缓存
    for _ in range(10):
        cap.read()

    ret, frame = cap.read()
    cap.release()
    if not ret or frame is None:
        print("未获取到帧")
        return

    h, w = frame.shape[:2]
    print(f"帧尺寸: {w}x{h}")

    # 检测
    dets = detector.detect(frame)
    persons = [d for d in dets if d.class_name == 'person']
    print(f"检测到 {len(persons)} 个人 (原始目标 {len(dets)})")
    for i, p in enumerate(persons):
        print(f"  {i}: conf={p.confidence:.3f} bbox={p.bbox} center={p.center}")

    # 保存原图
    orig_path = OUT_DIR / 'capture_original.jpg'
    cv2.imwrite(str(orig_path), frame)

    # 绘制标记图
    marked = frame.copy()
    for p in persons:
        x, y, bw, bh = p.bbox
        cv2.rectangle(marked, (x, y), (x + bw, y + bh), (0, 255, 0), 2)
        label = f"person {p.confidence:.2f}"
        cv2.putText(marked, label, (x, max(y - 5, 20)),
                    cv2.FONT_HERSHEY_SIMPLEX, 0.8, (0, 255, 0), 2)

    marked_path = OUT_DIR / 'capture_marked.jpg'
    cv2.imwrite(str(marked_path), marked)
    print(f"已保存: {orig_path}, {marked_path}")

    detector.release()


if __name__ == '__main__':
    main()