149 lines
4.9 KiB
Python
149 lines
4.9 KiB
Python
#!/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())
|