OperationRoomMonitor/main_basket.py

#!/usr/bin/env python3
"""篮子接触分段入口：OpenCV 框选篮子 → 手篮接触上升沿 → Phase2（跳过 ActionFormer）。"""
from __future__ import annotations

import argparse
import os
import sys
from pathlib import Path

PACK_ROOT = Path(__file__).resolve().parent
sys.path.insert(0, str(PACK_ROOT / "src"))

from paths import ensure_code_on_path

ensure_code_on_path(PACK_ROOT)

from config import load_run_config
from orchestrator import run_basket_pipeline


def main() -> int:
    os.environ.setdefault("OPENCV_FFMPEG_LOGLEVEL", "8")
    ap = argparse.ArgumentParser(
        description="手术室耗材篮子接触分段主流程（跳过 ActionFormer）"
    )
    ap.add_argument("--video", type=Path, required=True, help="输入 MP4")
    ap.add_argument(
        "--excel",
        type=Path,
        required=True,
        help="商品表 Excel（C 列白名单 + 产品编码）",
    )
    ap.add_argument("--out", type=Path, required=True, help="输出 TSV")
    ap.add_argument(
        "--config",
        type=Path,
        default=PACK_ROOT / "configs" / "default_config.yaml",
        help="继承 weights / phase2 / classification / doctor / basket 的 YAML",
    )
    ap.add_argument(
        "--save-basket-roi",
        type=Path,
        default=None,
        help="框选后将 ROI 保存为 JSON（可选；每次运行仍会先弹窗标框）",
    )
    ap.add_argument(
        "--det-conf",
        type=float,
        default=None,
        help="篮子扫描手部检测置信度（默认读 yaml basket.det_conf）",
    )
    ap.add_argument(
        "--contact-iou-on",
        type=float,
        default=None,
        help="篮子接触 IoU 进入阈值（默认读 yaml basket.contact_iou_on）",
    )
    ap.add_argument(
        "--contact-iou-off",
        type=float,
        default=None,
        help="篮子接触 IoU 退出阈值（默认读 yaml basket.contact_iou_off）",
    )
    ap.add_argument(
        "--confirm-seconds",
        type=float,
        default=None,
        help="连续接触确认时长（秒，默认 0.4）",
    )
    ap.add_argument(
        "--cooldown-seconds",
        type=float,
        default=None,
        help="触发后绝对冷却时长（秒，默认 5.0）",
    )
    ap.add_argument(
        "--contact-iou-threshold",
        type=float,
        default=None,
        help="手框与篮子 IoU 阈值（默认读 yaml basket.contact_iou_threshold）",
    )
    ap.add_argument(
        "--segment-start-offset-sec",
        type=float,
        default=None,
        help="段起点相对接触时刻偏移（秒，默认 1 → contact+1）",
    )
    ap.add_argument(
        "--segment-end-offset-sec",
        type=float,
        default=None,
        help="段终点相对接触时刻偏移（秒，默认 5 → contact+5）",
    )
    ap.add_argument(
        "--min-segment-sec",
        type=float,
        default=None,
        help="截断后段长不足此值则丢弃（秒，默认 4.0；0 表示不过滤）",
    )
    ap.add_argument(
        "--scan-frame-stride",
        type=int,
        default=None,
        help="全片接触扫描帧步长（默认 1）",
    )
    args = ap.parse_args()

    cfg_path = args.config.resolve()
    if not cfg_path.is_file():
        print("找不到配置:", cfg_path, file=sys.stderr)
        return 1

    run_cfg = load_run_config(PACK_ROOT, cfg_path)
    run_cfg.video = args.video.resolve()
    run_cfg.excel = args.excel.resolve()
    run_cfg.out = args.out.resolve()

    # 每次运行均在首帧弹窗标框，不从 JSON / yaml 复用 ROI
    run_cfg.basket_load_roi_json = None
    run_cfg.basket_skip_roi_select = False
    if args.save_basket_roi is not None:
        run_cfg.basket_save_roi_json = args.save_basket_roi.resolve()
    if args.det_conf is not None:
        run_cfg.basket_det_conf = float(args.det_conf)
    if args.contact_iou_on is not None:
        run_cfg.basket_contact_iou_on = float(args.contact_iou_on)
    if args.contact_iou_off is not None:
        run_cfg.basket_contact_iou_off = float(args.contact_iou_off)
    if args.confirm_seconds is not None:
        run_cfg.basket_confirm_seconds = float(args.confirm_seconds)
    if args.cooldown_seconds is not None:
        run_cfg.basket_cooldown_seconds = float(args.cooldown_seconds)
    if args.contact_iou_threshold is not None:
        run_cfg.basket_contact_iou_threshold = float(args.contact_iou_threshold)
    if args.segment_start_offset_sec is not None:
        run_cfg.basket_segment_start_offset_sec = float(args.segment_start_offset_sec)
    if args.segment_end_offset_sec is not None:
        run_cfg.basket_segment_end_offset_sec = float(args.segment_end_offset_sec)
    if args.min_segment_sec is not None:
        run_cfg.basket_min_segment_sec = float(args.min_segment_sec)
    if args.scan_frame_stride is not None:
        run_cfg.basket_scan_frame_stride = int(args.scan_frame_stride)

    return int(run_basket_pipeline(run_cfg))


if __name__ == "__main__":
    raise SystemExit(main())