手术室耗材篮子识别包(离线 + 推流)
段内流程:手检(≥2 手 union)→ 好坏帧门控 → 耗材分类;推流在每个接触段内与耗材并行医生识别;离线在全片结束后追加一行医生信息。
与 configs/default_config.yaml 当前参数一致(imgsz_det: 1920、contact+3~+10 等)。
环境要求
- Python 3.10+(推荐 conda 环境
yolo) - NVIDIA GPU + CUDA
python3-tk(每次运行首帧弹窗框选篮子 ROI)ffmpeg/ffprobe(HEVC 视频建议先转 H.264,见下文)
快速安装
cd /path/to/6.1
bash setup.sh
# 若用 conda(推荐)
conda activate yolo
pip install -r requirements.txt
三个入口
| 脚本 | 用途 |
|---|---|
main_basket.py |
离线:全片篮子接触分段 → Phase2 → gap 合并 → 医生识别 |
main_basket_stream.py |
推流/本地 MP4 模拟推流:逐帧触发 → 段内耗材+医生并行 → 实时写 15 列 TSV |
main_segments_offline.py |
按 TSV 时间段对离线 MP4 重跑段内识别(校验用) |
1. 离线跑视频
conda activate yolo
cd /path/to/6.1
python main_basket.py \
--video /path/to/your.mp4 \
--excel input/视频中的商品信息表.xlsx \
--out output/result_offline.txt \
--save-basket-roi output/basket_roi.json \
--config configs/default_config.yaml
运行后弹窗框选篮子 ROI,然后自动全片扫描 + 段内识别。
2. 推流(或本地 MP4 测试)
python main_basket_stream.py \
--rtsp /path/to/your.mp4 \
--excel input/视频中的商品信息表.xlsx \
--out output/result_stream.txt \
--save-basket-roi output/basket_roi_stream.json \
--config configs/default_config.yaml
- 本地 MP4:
stream.infer_source: file→ 段内回源 4K(与离线一致) - 真 RTSP:无法 seek 时回退 JPEG 缓存(
cache_max_width: 1920) - 医生识别:
doctor_identity.stream_enabled: true时,每段并行识别,TSV 追加doctor_id/doctor_name/doctor_conf
3. HEVC 视频
4K HEVC 可能导致 OpenCV 解码不稳定,建议先转码:
bash scripts/remux_hevc.sh /path/to/source.mp4
# 输出: input/remuxed/<name>_h264.mp4
配置说明(configs/default_config.yaml)
| 段 | 关键参数 |
|---|---|
phase2 |
imgsz_det: 1920,pad_bottom_ratio: 0.5,det_conf: 0.6 |
classification |
好帧 0.9,耗材 0.8,重试 0.7 / 0.5 |
basket |
iou_on: 0.03,confirm: 0.1,cooldown: 10,窗口 contact+3~+10 |
stream |
段窗口与 basket 一致;infer_source: file |
io |
use_whitelist: false(全 41 类) |
doctor_identity |
enabled / stream_enabled;推流用 segment_sample_fps 在段 [start,end] 内采样 |
模型文件(weights/)
hand_detect.pt— 手部检测goodbad_frame.pt— 好坏帧门控haocai_classify.pt— 耗材分类doctor_identity_package/doctor_info.pth— 医生 ReID(需同目录train_reid_contrastive.py)
输出格式
- 推流:12 列耗材 TSV;各段内在 居中 3s 窗跑医生 ReID 后 投票取 Top1,末尾一行
医生信息:...(doctor_identity.segment_window_sec)。末段时长不足 7s 判为误触丢弃 - 离线:12 列 TSV + 末尾一行
医生信息:...(全片中间窗口识别)
可视化演示(标注 MP4)
跑完 main_basket.py 后,可用独立脚本把 视频 + TSV 合成带框演示片(段内复现手检框,Top3/医生文字直接读 TSV,不重跑耗材模型):
conda activate yolo
cd /path/to/6.3
python scripts/visualize_pipeline.py \
--video /path/to/your.mp4 \
--tsv output/result_offline.txt \
--basket-roi output/basket_roi.json \
--config configs/default_config.yaml \
--out output/demo_vis.mp4
| 叠加层 | 说明 |
|---|---|
| 青色虚线框 | 篮子 ROI(需 --basket-roi,与 --save-basket-roi 配套) |
| 绿色框 | 段内手部检测(与 configs/default_config.yaml 的 hand.backend 一致;默认 MediaPipe weights/hand_landmarker.task,标签为「手 mp」) |
| 黄色粗框 | 双手 union ROI(与 Phase2 一致) |
| 顶部信息条 | TSV 该段时间段的 rank、Top3 或失败原因 |
| 片头 | 视频/TSV 路径 + 离线 医生信息: 汇总 |
性能建议(长片/4K):--preview-width 1280 --det-stride 5。依赖系统已安装 ffmpeg。
中文显示:叠加文字使用 Pillow + 系统 CJK 字体(默认 NotoSansCJK-Regular.ttc)。若出现方框/乱码,请安装 fonts-noto-cjk,或通过 --font /path/to/font.ttc / 环境变量 VIS_CJK_FONT 指定字体。
手部后端:与 main_basket 共用 hand 配置段。默认 hand.backend: mediapipe + hand.mediapipe_task: weights/hand_landmarker.task(每帧最多 2 只手)。对比 YOLO 旧行为:--hand-backend yolo。
篮筐附近手框与 ROI:提供 --basket-roi 时,默认只绘制靠近篮子的手(篮子框外扩 20% 后 IoU > contact_iou_on),黄色 ROI 由其中与篮子 IoU 最高的两只手合并。背景手不再绘制。关闭过滤用 --no-hand-basket-filter;贴边漏检可试 --basket-expand-frac 0.3 或略降 --hand-basket-min-iou 0.02。
本地 smoke(无真实手术视频时):
conda run -n yolo python scripts/make_smoke_fixture.py
conda run -n yolo python scripts/visualize_pipeline.py \
--video output/smoke/smoke.mp4 \
--tsv output/smoke/smoke_result.txt \
--basket-roi output/smoke/smoke_basket_roi.json \
--out output/smoke/demo_vis.mp4 \
--det-stride 2 --preview-width 640 --title-sec 1
目录结构
6.1/
├── main_basket.py # 离线入口
├── main_basket_stream.py # 推流入口
├── main_segments_offline.py # TSV 段内重跑
├── configs/default_config.yaml
├── weights/ # 3 个 YOLO 权重
├── input/视频中的商品信息表.xlsx
├── doctor_identity_package/ # 医生识别(离线整片 + 推流段内)
├── src/ code/ # 编排与算法
├── output/ # 结果输出目录
├── scripts/visualize_pipeline.py # TSV → 标注演示 MP4
├── setup.sh requirements.txt
└── README.md
常见问题
- CUDA OOM:勿将
imgsz_det设为 3840;当前 1920 在 8GB 显卡可用。 - 无 GUI:无法弹窗框选 ROI,需另备 ROI JSON(当前入口每次仍弹窗)。
- 分段过多:可调大
basket.confirm_seconds或cooldown_seconds。