diff --git a/.gitattributes b/.gitattributes index 5f609e2..35acbc3 100644 --- a/.gitattributes +++ b/.gitattributes @@ -1,5 +1,4 @@ -*.sh text eol=lf -*.engine -text -*.ko -text -/orin/** -text -/config/door_mask.json -text +*.engine binary +*.onnx binary +*.so binary +*.ko binary diff --git a/.gitignore b/.gitignore index 4097644..85a87c1 100644 --- a/.gitignore +++ b/.gitignore @@ -1,4 +1,41 @@ -/build/ -/engines/ -/logs/ -*.tmp +# Python caches and local environments +__pycache__/ +*.py[cod] +.venv/ +venv/ + +# Build outputs (keep the supplied vendor libraries in lib/) +build/ +/build-*/ +*.o +*.a +/vision-anti-pinch + +# Models and test media +models/ +*.engine +*.onnx +*.pt +*.mp4 +*.avi + +# Runtime outputs +logs/ +outputs/ +out/ +frames/ +*.log +*.pid + +# Local credentials and legacy SSH authentication helper +.env +.env.* +!.env.example +ssh_expect.exp + +# The checked-in Orin deployment bundle intentionally includes its runtime assets. +!/orin/build/ +!/orin/build/libvision_native.so +!/orin/models/ +!/orin/models/engines/ +!/orin/models/engines/*.engine diff --git a/BRANCHES.md b/BRANCHES.md new file mode 100644 index 0000000..fa6ce16 --- /dev/null +++ b/BRANCHES.md @@ -0,0 +1,9 @@ +# 分支说明 + +| 分支 | 内容 | +|---|---| +| main | CMS v9.10 原始 YOLO + RTMW + risk 部署包，入口 run_risk.py，含 ONNX 模型 | +| use-can | 0923 V4L2/CAN + TensorRT + CMS 风险集成，提供最新结果变量、可选可视化及 orin/run.sh 可运行包 | +| inference-only | 原 main 的纯 C++ TensorRT 摄像头推理版本，保留原提交 26c5bfb88ddf9919ea46efcf5b12df6e09ebdfd9 | + +main 和 use-can 均在原 main 历史上追加提交；原 main 历史未改写。 diff --git a/CMakeLists.txt b/CMakeLists.txt index 1a5dede..48ab824 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -1,23 +1,18 @@ cmake_minimum_required(VERSION 3.24) -if(NOT DEFINED CMAKE_CUDA_ARCHITECTURES) - set(CMAKE_CUDA_ARCHITECTURES native) -endif() -project(camera_risk LANGUAGES CXX CUDA) +project(vision_anti_pinch_0923 LANGUAGES CXX CUDA) set(CMAKE_CXX_STANDARD 17) -set(CMAKE_CXX_STANDARD_REQUIRED ON) set(CMAKE_CUDA_STANDARD 17) +set(CMAKE_POSITION_INDEPENDENT_CODE ON) +set(CMAKE_CUDA_ARCHITECTURES 87) find_package(CUDAToolkit REQUIRED) -find_package(OpenCV REQUIRED COMPONENTS core imgproc imgcodecs highgui) +find_package(OpenCV REQUIRED COMPONENTS core imgproc) find_package(Threads REQUIRED) -find_path(TENSORRT_INCLUDE_DIR NvInfer.h PATH_SUFFIXES aarch64-linux-gnu x86_64-linux-gnu REQUIRED) -find_library(TENSORRT_LIBRARY nvinfer REQUIRED) -add_library(camera_risk_core - src/trt_engine.cpp src/pose_engine.cu src/gpu_preprocess.cu src/uyvy_converter.cu - src/detector.cpp src/door_risk.cpp src/camera.cpp src/camera_pipeline.cpp src/config.cpp) -target_include_directories(camera_risk_core PUBLIC include third_party ${TENSORRT_INCLUDE_DIR} ${OpenCV_INCLUDE_DIRS}) -target_link_libraries(camera_risk_core PUBLIC ${TENSORRT_LIBRARY} ${OpenCV_LIBS} Threads::Threads CUDA::cudart CUDA::nppc CUDA::nppig CUDA::nppidei CUDA::nppicc) -target_compile_options(camera_risk_core PRIVATE $<$:-Wall;-Wextra;-Wpedantic>) -add_library(camera_risk_visualizer src/visualizer.cpp) -target_link_libraries(camera_risk_visualizer PUBLIC camera_risk_core) -add_executable(camera_risk src/main.cpp) -target_link_libraries(camera_risk PRIVATE camera_risk_visualizer) +find_path(TRT_INCLUDE NvInfer.h PATH_SUFFIXES aarch64-linux-gnu REQUIRED) +find_library(TRT_LIB nvinfer REQUIRED) +add_library(vision_native SHARED native/native_api.cpp v4l2_capture.cpp can_comm.cpp + accelerator/src/trt_engine.cpp accelerator/src/pose_engine.cu + accelerator/src/gpu_preprocess.cu accelerator/src/uyvy_converter.cu accelerator/src/detector.cpp) +target_include_directories(vision_native PRIVATE . include accelerator/include ${TRT_INCLUDE} ${OpenCV_INCLUDE_DIRS}) +target_link_libraries(vision_native PRIVATE ${TRT_LIB} ${OpenCV_LIBS} Threads::Threads CUDA::cudart CUDA::nppc CUDA::nppig CUDA::nppidei CUDA::nppicc + "${CMAKE_SOURCE_DIR}/lib/libusbcanfd.so" "${CMAKE_SOURCE_DIR}/lib/libusb-1.0.so") +set_target_properties(vision_native PROPERTIES BUILD_WITH_INSTALL_RPATH TRUE INSTALL_RPATH "$ORIGIN/../lib;/usr/local/cuda/lib64") diff --git a/Makefile b/Makefile new file mode 100644 index 0000000..d665629 --- /dev/null +++ b/Makefile @@ -0,0 +1,11 @@ +.PHONY: all run test baseline clean +all: + ./scripts/build.sh +run: all + ./run_camera.sh +test: + python3 -m unittest discover -s tests -v +baseline: + $(MAKE) -f Makefile.baseline +clean: + cmake --build build --target clean diff --git a/Makefile.baseline b/Makefile.baseline new file mode 100644 index 0000000..cb54c55 --- /dev/null +++ b/Makefile.baseline @@ -0,0 +1,71 @@ +# Makefile - V4L2 MMAP 取图 demo (OpenCV + CAN 收发) +# +# 用法: +# make # 编译 +# make clean # 清理 +# make run # 编译并运行 (取 100 帧) +# make capture # 编译并取 5 帧保存到 ./out +# make no-can # 同 capture, 但关闭 CAN 收发 +# +# 说明: 本 Makefile 在 Orin 上本地编译。 +# CAN 模块依赖 include/zcan.h + lib/libusbcanfd.so + lib/libusb-1.0.so +# CAN 收发默认开启, 程序运行时加 --no-can 关闭 + +# ---- 工具链 ---- +CXX ?= g++ +CXXFLAGS ?= -std=c++17 -O3 -Wall -Wextra -Wno-unused-parameter -Wno-unused-result +LDFLAGS ?= + +# ---- OpenCV 检测 (优先 opencv4, 回退 opencv) ---- +OPENCV_PC := $(shell pkg-config --exists opencv4 && echo opencv4 || echo opencv) +OPENCV_CFLAGS := $(shell pkg-config --cflags $(OPENCV_PC)) +OPENCV_LIBS := $(shell pkg-config --libs $(OPENCV_PC)) + +# ---- 本地头文件/库目录 (zcan.h, libusbcanfd.so, libusb-1.0.so) ---- +INC_DIR := include +LIB_DIR := lib +IDIRS := -I$(INC_DIR) $(OPENCV_CFLAGS) +LDIRS := -L$(LIB_DIR) +RPATH := -Wl,-rpath,$$(pwd)/$(LIB_DIR) + +CXXFLAGS += $(IDIRS) +LDLIBS += $(OPENCV_LIBS) -lm -lpthread -lusbcanfd -lusb-1.0 + +# ---- 目标与源文件 ---- +BIN := vision-anti-pinch +SOURCES := main.cpp v4l2_capture.cpp can_comm.cpp +HEADERS := v4l2_capture.h can_comm.h $(INC_DIR)/zcan.h +OBJECTS := $(SOURCES:.cpp=.o) + +# ---- 默认目标 ---- +.PHONY: all clean run capture + +all: $(BIN) + +$(BIN): $(OBJECTS) + $(CXX) $(LDFLAGS) -o $@ $^ $(LDIRS) $(RPATH) $(LDLIBS) + +# ---- 依赖规则 (头文件改动则全部重编) ---- +$(OBJECTS): $(HEADERS) + +%.o: %.cpp + $(CXX) $(CXXFLAGS) -c -o $@ $< + +# ---- 便捷目标 ---- +run: $(BIN) + ./$(BIN) -d /dev/video5 -n 100 -v + +capture: $(BIN) + mkdir -p out + ./$(BIN) -d /dev/video5 -n 100 -s 5 -t jpg -o ./out + +# CAN 收发默认开启, 与 capture 相同 +can: capture + +# 关闭 CAN 收发的取图 +no-can: $(BIN) + mkdir -p out + ./$(BIN) -d /dev/video5 -n 100 -s 5 -t jpg -o ./out --no-can -v + +clean: + rm -f $(OBJECTS) $(BIN) diff --git a/PROVENANCE.json b/PROVENANCE.json new file mode 100644 index 0000000..7a083b2 --- /dev/null +++ b/PROVENANCE.json @@ -0,0 +1,61 @@ +{ + "baseline": "orin@192.168.195.100:/home/orin/vision-anti-pinch", + "baseline_sha256": [ + "a6f884d241b3b345fbdc3bbb2d0119467f6a01d6828786a7d33d7f97ad101f6e main.cpp", + "ba7ce790ec4b43403fd6b9ccc4693148414b7314060c29d9f9f92dee931509cf can_comm.cpp", + "5aaeaba6f0f25a6202ff2cf61c626b61bda2350f5bd9b36e36c828b5b8fec994 can_comm.h", + "f5c6658eb3e8362569c18d528fd4c8a2e311e5f779d9687e8ff10b79a11dc5b6 v4l2_capture.cpp", + "dd7a6f92e4d3335368691dd587715cec3d24816e9a96d44ff11c91296f5082eb v4l2_capture.h", + "65b16fe3d4549ad218679f510afe4049ad0f197912832060daaca13acc0f5f4f Makefile" + ], + "risk_source": "/home/thor/文档/cms_v9_10_yolo_rtmw_risk_deploy_20260923", + "risk_pipeline": { + "occlusion_aware_pose_v2.py": { + "source_sha256": "81088b8bc072162bec13f57c65eb0d396e6d1e1074a83f92bcf29e393e27e11f", + "integrated_sha256": "81088b8bc072162bec13f57c65eb0d396e6d1e1074a83f92bcf29e393e27e11f" + }, + "fuse_rtmw_l_temporal.py": { + "source_sha256": "bc74a400ce5c5608c2fa49a69ccd8b41ec8633cd6a8239d9d594b8d9e5c3732a", + "integrated_sha256": "bc74a400ce5c5608c2fa49a69ccd8b41ec8633cd6a8239d9d594b8d9e5c3732a" + }, + "rtmw_l_person_mask_ablation.py": { + "source_sha256": "d983c9431237afcd7a786c6c4610be6d57a8b4d9ea8f2e9d1889fc0b0ca1d5df", + "integrated_sha256": "bf8bb754cd841a52b4410f958a651b005de36ac4d17cc0992fbe1c90ed30caca" + }, + "door_mask_risk.py": { + "source_sha256": "6d6d14a972bc4e5e9f8256514169e73ba0392a301cacd1221eaa26146aad00d2", + "integrated_sha256": "6d6d14a972bc4e5e9f8256514169e73ba0392a301cacd1221eaa26146aad00d2" + }, + "live_pose_engine.py": { + "source_sha256": "a15baaedf5a25b9896ff898fc1a095388483dc52b1b7000b56d5506ff73bf67e", + "integrated_sha256": "cc8ea36d543896cd90f7a00822f06fc67771a1811859a4087995ac372df30276" + }, + "predict_optimized_pose_videos.py": { + "source_sha256": "36d203f257444e57f890cdc3826dad0c7a46725b938e5f8507e357a70ae2cef9", + "integrated_sha256": "059e2a67c17ce63354986d3242a5d361103ee4241c5391f18140eeb19e47c939" + }, + "door_zone_risk.py": { + "source_sha256": "46d70e1e90a1f36d544b7492ef1fbf08821377a1c460a6683f42d4d7006055c7", + "integrated_sha256": "46d70e1e90a1f36d544b7492ef1fbf08821377a1c460a6683f42d4d7006055c7" + }, + "rtmw_risk_pipeline.py": { + "source_sha256": "1bb58d851ba5d684527b8d1e1940db6241d2b0b2862687fa3b151899a65320d2", + "integrated_sha256": "1bb58d851ba5d684527b8d1e1940db6241d2b0b2862687fa3b151899a65320d2" + }, + "compare_rtmw_ml.py": { + "source_sha256": "5de11436bb888f2f1c4c0e7ad0853e3864193fe2625cc65294b40683ebe8729a", + "integrated_sha256": "fd3ae92255abf95ef037a1060ee75897f9705b4040f9394c3130f6ab860667e1" + } + }, + "engines_sha256": { + "rtmw_l_384.engine": "ebcd09e8c2e35008adbb0c1d222e580b17e0df1c8093df3481ba3374cd9a8174", + "yolo11s.engine": "f6a65b926ac1144484f4aa93cd26d7997d32d62310ab93fe867f2bd4412222e9" + }, + "engine_origin": "orin@192.168.195.100:/home/orin/vision-anti-pinch-0920/orin/engines", + "can_mapping_user_confirmed": [ + "left_front", + "right_front", + "left_rear", + "right_rear" + ] +} \ No newline at end of file diff --git a/README-risk-upstream.md b/README-risk-upstream.md new file mode 100644 index 0000000..5af9397 --- /dev/null +++ b/README-risk-upstream.md @@ -0,0 +1,330 @@ +# CMS v9.10 YOLO + RTMW 门区风险部署包 + +该目录是 v9.10 的精简部署版本，只保留运行 `YOLO + RTMW + risk` 所需的代码、模型、标定文件和依赖清单。历史推理结果、研发测试、标注网页、参考图片、缓存与离线分析入口均未放入部署包。 + +## 1. 功能链路 + +```text +视频帧 + -> YOLO11s 人体检测与轨迹 ID + -> RTMW-L COCO-WholeBody 133 点关键点 + -> 手臂/手部时序融合 + -> 骨骼位置和尺度为主、人体框为辅的前后门归属 + -> CAN 活跃门过滤 + -> 手部区域（无有效手部时使用人体框兜底） + -> 与固定车门危险区 mask 计算二维风险 + -> causal_results.jsonl +``` + +本版本不包含深度估计和 EdgeTAM。输出是二维图像风险，不是物理距离，也不是安全认证结论。 + +## 2. 目录结构 + +```text +run_risk.py 主入口 +requirements-realtime-gpu.txt CUDA 推理依赖 +requirements-tensorrt-windows.txt Windows TensorRT 可选依赖 +models/ + yolo11s.onnx 人体检测模型 + rtmw-l-384.onnx WholeBody 关键点模型 +annotations/ + constant_door_mask.json 当前固定机位车门危险区 + door_zone_multifeature.json 当前前门/后门多特征标定 +pipeline/ 推理、融合、门区和风险计算代码 +``` + +`annotations` 中的两个 JSON 都与摄像头机位相关。当前参数来自研发临时机位；换成实车机位后必须重新标定和验证，不能直接将临时机位结论作为量产参数。 + +## 3. 环境安装 + +推荐 Python 3.10～3.12 和 NVIDIA CUDA GPU。先安装与目标机 CUDA 或 JetPack 匹配的 PyTorch，再安装其余依赖： + +```powershell +python -m venv .venv +.\.venv\Scripts\Activate.ps1 +python -m pip install --upgrade pip +pip install -r requirements-realtime-gpu.txt +``` + +不要同时安装 CPU 版 `onnxruntime` 和 `onnxruntime-gpu`。Jetson Orin 应优先使用 NVIDIA 针对对应 JetPack 发布的 PyTorch/torchvision，而不是强制使用桌面 wheel。 + +如需在 Windows 上使用 RTMW TensorRT Execution Provider，再执行： + +```powershell +pip install -r requirements-tensorrt-windows.txt +``` + +## 4. 运行 + +### 4.1 固定车门危险区 + +以后门 CAN=0、前门未进入电吸范围为例： + +```powershell +python run_risk.py ` + --video D:\data\segment_008.mp4 ` + --constant-door-mask annotations\constant_door_mask.json ` + --door-zone-config annotations\door_zone_multifeature.json ` + --active-door rear ` + --output outputs\seg_08_rear ` + --det-device 0 ` + --pose-device cuda ` + --write-video +``` + +`--active-door` 对应当前 CAN=0 的车门： + +- `front`：只有前门进入电吸范围； +- `rear`：只有后门进入电吸范围； +- `both`：前后门都进入电吸范围； +- `none`：两门都不进行风险判断。 + +当前命令行参数对整个视频保持不变，用于离线部署验证。实车 CAN 随帧变化时，集成层需要把每帧实际 CAN 状态传给门级风险选择逻辑。 + +不需要叠加视频时不要传 `--write-video`，可减少编码和磁盘 I/O。完整参数说明： + +```powershell +python run_risk.py --help +``` + +### 4.2 TensorRT RTMW + +```powershell +python run_risk.py ` + --video D:\data\test.mp4 ` + --constant-door-mask annotations\constant_door_mask.json ` + --active-door rear ` + --output outputs\test_trt ` + --pose-provider tensorrt ` + --pose-device cuda +``` + +首次运行会在 `models/rtmw_l_trt_cache` 生成 TensorRT 缓存，该缓存与目标 GPU、TensorRT 版本和运行环境相关，不随部署包预置。 + +## 5. 输出文件 + +输出目录包含： + +- `causal_results.jsonl`：逐帧因果结果，一行一个 JSON 对象，是 Demo 对接的主要接口； +- `summary.json`：模型、帧数、吞吐、延迟及风险统计； +- `rtmw_risk_overlay.mp4`：仅在传入 `--write-video` 时生成。 + +`causal_results.jsonl` 每行顶层字段： + +| 字段 | 含义 | +|---|---| +| `frame_index` | 从 0 开始的视频帧号 | +| `track_ids` | 当前帧人体轨迹 ID 数组 | +| `boxes` | 与 `track_ids` 同序的人体框，格式 `[x1,y1,x2,y2]` | +| `skeletons` | 与检测人体对齐的 133 点融合骨骼 | +| `active_door` | 本次运行指定的 CAN=0 门：`front/rear/both/none` | +| `door_assignments` | 每个轨迹的前门/后门归属及特征解释 | +| `door_risks` | `front` 和 `rear` 两扇门各自的最终风险 | +| `risk` | 经过 CAN 和门区过滤后供业务使用的当前风险 | +| `ungated_risk` | 未经过前后门/CAN 过滤的原始二维风险，仅用于诊断 | + +数组关联规则：`boxes[i]`、`track_ids[i]`、`skeletons[i].person_index == i`、`door_assignments[i].person_index == i` 表示同一个人。跨帧关联应使用 `track_id`，不要依赖数组下标。 + +## 6. `risk` 的含义 + +### 6.1 应读取哪个字段 + +单门 CAN=0 时，业务防夹建议读取： + +```text +risk.risk_type +risk.risk_level +risk.risk_score +risk.decision +``` + +需要同时展示前后门时分别读取： + +```text +door_risks.front +door_risks.rear +``` + +`ungated_risk` 只用于排查“图像上是否靠近危险区”，不能直接代替门级最终结论，否则会绕过前后门归属和 CAN 过滤。 + +### 6.2 风险标签 + +`risk_type` 是带门位置的最终风险标签： + +| `risk_type` | `risk_level` | 含义 | +|---|---|---| +| `FRONT_CLEAR` / `REAR_CLEAR` | `clear` | 对应活跃门没有风险 | +| `FRONT_WATCH` / `REAR_WATCH` | `watch` | 进入关注距离，但未达到 warning | +| `FRONT_WARNING` / `REAR_WARNING` | `warning` | 距离达到 warning 阈值 | +| `FRONT_DANGER` / `REAR_DANGER` | `danger` | 人体风险区域与车门危险区相交 | +| `FRONT_UNKNOWN` / `REAR_UNKNOWN` | `unknown` | 人物处于 `transition`，无法可靠确定前后门 | +| `FRONT_INACTIVE` / `REAR_INACTIVE` | `inactive` | 该门 CAN 不活跃，不计算防夹风险 | +| `NO_ACTIVE_DOOR` | `inactive` | 前后门均不活跃 | + +对应的 `decision`： + +| `decision` | 含义 | +|---|---| +| `DANGER_INTERSECTION` | 风险区域与车门 mask 有像素相交 | +| `WARNING_NEAR` | 相对距离小于等于 warning 阈值 | +| `WATCH` | 相对距离处于 warning 与 release 之间 | +| `CLEAR` | 相对距离大于 release 阈值 | +| `CLEAR_NO_PERSON_ASSIGNED_TO_ACTIVE_DOOR` | 没有人被分配给该活跃门 | +| `EVIDENCE_UNAVAILABLE_TRANSITION_ZONE` | 存在 `transition` 人员，门归属不确定 | +| `DOOR_CAN_NOT_ACTIVE` | 该门 CAN 不活跃 | +| `NO_DOOR_CAN_ACTIVE` | 没有活跃门 | + +默认阈值： + +```text +warning_distance_relative = 0.30 +release_distance_relative = 0.45 +``` + +对每个风险区域： + +```text +distance_relative = min_distance_px / sqrt(region_area_px) +risk_score = clip(1 - distance_relative / release_distance_relative, 0, 1) +``` + +只要 `door_overlap_px > 0`，就强制 `risk_score=1.0` 且等级为 `danger`。因此 `risk_score` 是归一化二维接近程度，不是概率，也不是厘米/毫米距离。 + +### 6.3 `regions` 与证据来源 + +每个门风险中的 `regions` 保存独立人体区域的解释信息，常用字段包括： + +- `track_id`、`side`：所属轨迹和左右手； +- `source`：`RTMW_HAND_KEYPOINTS` 或 `YOLO_BODY_BOX_FALLBACK`； +- `region_area_px`：风险区域面积； +- `door_overlap_px`、`region_overlap_ratio`：与车门 mask 的相交量； +- `min_distance_px`：区域到车门 mask 的最小像素距离； +- `region_scale_px`：`sqrt(region_area_px)`； +- `distance_relative`：尺度归一化距离； +- `risk_score`、`risk_level`、`decision`：该区域自己的风险结果。 + +存在可用手部时，每只手用 21 点凸包计算风险；一个人没有任何可用手部时，才使用整个人体框兜底。人体框兜底更保守，也更容易产生误报。 + +### 6.4 `transition` 的当前行为 + +`door_assignments[].door_zone` 的位置标签为： + +- `front`：归属前门； +- `rear`：归属后门； +- `transition`：前后门证据不足或位于过渡区； +- `unassigned`：人体尺度低于有效阈值，不参与门级风险。 + +当前 v9.10 中，`transition` 人员即使在 `ungated_risk` 中距离很近，也不会直接变成 `FRONT_WARNING/REAR_WARNING` 或 `DANGER`，而是让活跃门输出 `*_UNKNOWN`，其 `risk_score` 为 `null`。Demo 必须显式展示或处理 `unknown`，不能把它当成 `clear`。这是当前版本行为，不代表最终量产安全策略。 + +## 7. `skeletons` 的含义与可视化对接 + +### 7.1 数据结构 + +每个人的骨骼对象如下： + +```json +{ + "track_id": 1, + "person_index": 0, + "layout": "coco_wholebody_133", + "keypoints": [[x0, y0], [x1, y1]], + "metadata": { + "left": {}, + "right": {} + } +} +``` + +- `keypoints` 实际固定为 `133 x 2`，示例中省略了其余点； +- 坐标单位是原视频像素，原点在左上角，x 向右、y 向下； +- 坐标保留 3 位小数； +- 这里没有逐点置信度，`metadata` 提供的是手臂/手部融合后的状态； +- `keypoints` 使用的正是当前叠加视频 `draw_fused` 使用的融合结果。 + +### 7.2 COCO-WholeBody 133 点索引 + +| 索引 | 含义 | +|---|---| +| `0` | nose | +| `1/2` | left/right eye | +| `3/4` | left/right ear | +| `5/6` | left/right shoulder | +| `7/8` | left/right elbow | +| `9/10` | left/right wrist | +| `11/12` | left/right hip | +| `13/14` | left/right knee | +| `15/16` | left/right ankle | +| `17/18/19` | left big toe / small toe / heel | +| `20/21/22` | right big toe / small toe / heel | +| `23–90` | 68 个面部关键点 | +| `91–111` | 左手 21 点 | +| `112–132` | 右手 21 点 | + +每只手的局部 21 点顺序： + +```text +0 wrist +1-4 thumb +5-8 index finger +9-12 middle finger +13-16 ring finger +17-20 little finger +``` + +左手局部索引需要加 91，右手局部索引需要加 112，才能得到 `keypoints` 中的全局索引。 + +### 7.3 与当前叠加视频一致的绘制规则 + +当前风险叠加视频只绘制手臂和双手，不绘制面部、躯干和腿部。手臂连线： + +```text +(5,7), (7,9) 左肩-左肘-左腕 +(6,8), (8,10) 右肩-右肘-右腕 +``` + +每只手使用以下局部连线，分别作用到 `91–111` 和 `112–132`： + +```text +(0,1),(1,2),(2,3),(3,4) +(0,5),(5,6),(6,7),(7,8) +(0,9),(9,10),(10,11),(11,12) +(0,13),(13,14),(14,15),(15,16) +(0,17),(17,18),(18,19),(19,20) +``` + +为了与现有视频表现一致： + +- `metadata[side].limb_state == "consensus"`：手臂画实线，否则画虚线； +- `metadata[side].hand_state == "rejected"`：不画该手； +- 手未被拒绝且 `distance_eligible == true`：手部画实线； +- 手未被拒绝但 `distance_eligible == false`：手部画虚线； +- `alarm_eligible == true` 表示该手允许进入 risk 区域计算。 + +Demo 如果只需要复现当前可视化，读取上述 48 个点即可；如果后续需要完整人体展示，可以使用全部 133 点，但应自行补充完整 WholeBody 连线定义。 + +## 8. `door_assignments` 对接说明 + +常用字段： + +- `track_id`、`person_index`：与检测框和骨骼关联； +- `door_zone`：经过最多 5 帧因果中值后的 `front/rear/transition/unassigned`； +- `raw_door_zone`：当前帧未经时序中值稳定的标签； +- `front_probability`：当前帧前门概率； +- `smoothed_front_probability`：因果时序中值； +- `confidence`：相对 0.5 的归属置信度，不是人体检测置信度； +- `pose_feature_count`：本帧可用骨骼特征数量； +- `features`、`feature_frontness`：标定特征值和各特征前门贡献； +- `history_frames`：当前轨迹已用于中值的历史帧数。 + +当前标定阈值为：`smoothed_front_probability >= 0.72` 判为前门，`<= 0.28` 判为后门，中间为 `transition`。这些参数应随实车机位重新验证。 + +## 9. Demo 集成注意事项 + +1. 每帧先按 `track_id` 关联检测框、骨骼、门区和风险，避免多人时按数组顺序长期跟踪。 +2. 主界面风险展示使用 `risk`；需要前后门双状态时使用 `door_risks.front/rear`。 +3. `unknown`、`inactive` 和 `clear` 是三种不同状态，不能合并显示。 +4. `risk_score == null` 表示该状态没有可用数值风险，不应转换成 0。 +5. `ungated_risk` 只用于诊断叠加，不能作为最终防夹指令。 +6. 更换分辨率或裁剪策略后必须核对车门 mask、骨骼坐标与视频坐标是否仍一致。 +7. 当前系统可能把玻璃倒影或反光误检成人，部署 Demo 应保留原视频、检测框、`track_id` 和风险解释字段用于问题回溯。 + diff --git a/README.md b/README.md index d1a7025..3af6e70 100644 --- a/README.md +++ b/README.md @@ -1,73 +1,205 @@ -# Vision Anti-Pinch +# vision-anti-pinch-0923 -C++ / TensorRT 实时摄像头风险推理，算法来自最新 run_risk.py。 -持续取图、推理并更新内存中的 FrameResult；可视化可选。 +以 Orin `~/vision-anti-pinch` 的 V4L2 MMAP 摄像头及 USBCANFD 模块为基线，接入 CMS v9.10 `run_risk.py` 的完整风险流程。模型及 GPU 推理代码来自远端 `~/vision-anti-pinch-0920`。 -## Orin 直接运行 +``` +基线 V4L2 MMAP → NPP UYVY 转 BGR → 最新帧槽位 + → YOLO11s TensorRT → RTMW-L TensorRT（含小人遮挡隔离推理） + → Python 原版轨迹 / 手臂手部融合 / 全局手关联 + → 骨骼主导的前后门归属 → 固定 mask 风险 → 逐帧 CAN 活跃门过滤 + → 内存中的最新只读结果（可选 JSONL / 浏览器预览） +``` -适用环境：Jetson AGX Orin、JetPack 5.1.2 / L4T R35.4.1、CUDA 11.4、TensorRT 8.5.2、OpenCV 4.5.4。 -仓库已包含该板端编译的程序和重新构建的 TensorRT 引擎。 +## 仓库分支与 Orin 直接运行包 + +- `main`：原 CMS v9.10 Python/ONNX 部署包。 +- `inference-only`：保留原 main 的纯 C++ TensorRT 推理版本。 +- `use-can`（本分支）：0923 CAN 集成版本。 + +本分支已提供可单独复制的 `orin/` 运行包，包含预编译库和 TensorRT 模型： ```bash -git clone https://git.kendeji.fun/zmv/vision-anti-pinch.git +git clone --branch use-can https://git.kendeji.fun/zmv/vision-anti-pinch.git cd vision-anti-pinch/orin -./scripts/initialize_camera.sh # 开机后初始化 /dev/video5 摄像头链路 -./run.sh # 持续推理，结果只保存在内存 +./scripts/initialize_camera.sh # 开机后需要初始化摄像头时执行 +./run.sh # 默认无可视化，结果保存在内存 +./run.sh --visualize # 可选预览 ``` -可选板端本地可视化： - -```bash -./run.sh --visualize -``` - -Ctrl+C / SIGTERM 停止程序。 -关闭预览窗口或按 q/Esc 只关闭预览，推理继续。 - -摄像头为 /dev/video5，1920×1536 UYVY，30 FPS。 -当前 door mask 是指定的 constant_door_mask.json 副本，标定状态为 false；装车后应更换与实际画面匹配的门区并更新标定状态。 - -## 下游获取结果 - -main.cpp 注释中也包含以下示例。下游通过自己的调用位置或线程读取： - -```cpp -auto result = pipeline.latest(); -if (result) { - const auto& risk = result->risk; - const auto& people = result->persons; - // result->frame_index / track_ids / boxes - // result->image / door_mask_pixels / capture / timing - // people[i].selected / smoothed / temporal：各 133 个关键点 -} -``` - -latest() 返回 shared_ptr；首帧尚未完成时为空。 -公共定义位于 include/camera_pipeline.hpp 和 include/frame_result.hpp。 - -## 从源码构建 +`orin/run.sh` 通过相对目录定位运行库、模型与标定，使用系统 Python，无需现场编译。目标环境及下游调用说明见 `orin/README.md`。 +源码变更后，在目标 Orin 构建并执行 `scripts/package_orin.sh` 重新生成运行包；`ENGINE_DIR` 和 `DRIVER_DIR` 可指定模型/驱动目录。 + +## 两端目录和环境 + +- 本机：`/home/thor/文档/vision-anti-pinch-0923` +- Orin：`/home/orin/vision-anti-pinch-0923` +- 目标推理环境：Orin，JetPack R35.4.1、CUDA 11.4、系统 Python 3.8、原生 TensorRT。 +- 本机保存完整源码、测试报告、模型副本；本机是 Thor / R38，未发现 `/dev/video*`。Orin 的序列化 engine 不能直接作为 Thor 的本机模型使用，需在 Thor 上另行转换。本次实时采集/推理在 Orin，浏览器可在本机查看。 +- 远端 `models/engines` 软链接到 `../../vision-anti-pinch-0920/orin/engines`，避免重复占用有限磁盘。不要删除被引用的 0920 模型。 +- 原有远端 `vision-anti-pinch` / `vision-anti-pinch-0920` 源码均未修改。测试时已停止占用 `/dev/video5` 的旧 0920 演示进程，0923 可独立启动。 + +## 编译和运行（在 Orin） ```bash +cd ~/vision-anti-pinch-0923 ./scripts/build.sh +# 开机后如 /dev/video5 不存在，加载已部署的 0920 摄像头驱动： ./scripts/initialize_camera.sh -./run_camera.sh +# 编译也可使用 make +./run_camera.sh --active-door can --can-side left +# 等价入口：python3 run_risk.py --active-door can --can-side left ``` -源码构建要求 CMake 3.24+、CUDA/NPP、TensorRT 和 OpenCV 开发依赖。 -构建脚本会优先使用 Orin 上已有的 ~/miniconda3/envs/yolo/bin/cmake。 -源码配置直接使用 orin/engines/ 中的板端引擎。该引擎仅适用于上述 Orin 环境；仓库不包含 ONNX 或模型转换脚本。 -独立库目标为 camera_risk_core；可视化在 camera_risk_visualizer。 +默认只在内存中发布结果，不创建结果目录，不绘制或编码图像，不启动 HTTP 服务。 -## 目录 +后台启动、停止： -- src/、include/：采集、推理、结果变量及可选绘制。 -- config/：源码构建程序的配置和固定门区。 -- scripts/：构建和摄像头初始化脚本。 -- orin/：已编译程序、TensorRT 引擎、配置、当前 GMSL 驱动及启动脚本。 -- third_party/：nlohmann JSON 及其许可证。 +```bash +./scripts/start.sh # 无可视化 +./scripts/stop.sh +./scripts/start.sh --visualize # 可选浏览器预览 +``` -NPP 用于 UYVY 转换、YOLO 缩放填充、RTMW 仿射裁剪；CUDA 用于通道转换与归一化。 -风险几何计算仍包含 CPU/OpenCV 操作；算法内部的临时 JSON 对象只在内存中使用。 +默认从 `/dev/video5` 采集 UYVY 1920×1536，驱动当前为 30 FPS。驱动不支持 `VIDIOC_G_PARM`，程序会打印警告并继续采集，命令行 fps 不代表驱动一定支持修改。 -当前摄像头驱动启动时偶发首帧 ERROR，应用会丢弃并计数；持续错误、采集超时或推理错误会明确报错并退出。 -现有 GMSL 驱动停止采集时可产生 videobuf2 缓冲区告警，v4l2-ctl 同样复现；需在驱动侧完成量产修复和验证。 +- 开启 `--visualize` 后，预览：[http://192.168.195.100:8765](http://192.168.195.100:8765) +- 开启 `--visualize` 后，最新结果：[http://192.168.195.100:8765/result.json](http://192.168.195.100:8765/result.json) +- 开启 `--visualize` 后，最新图像：[http://192.168.195.100:8765/frame.jpg](http://192.168.195.100:8765/frame.jpg) +- 后台进程日志：`logs/live.log`，PID：`logs/camera.pid`；加保存开关后，后台结果写入 `outputs/live/`。 +- `--visualize` 开启预览；仅设置 `--preview-port` 不会开启预览。监听地址默认 `0.0.0.0:8765`，可用 `--preview-host 127.0.0.1` 限制为本机。可视化在独立观察线程中绘制和编码，只保留最新待绘制帧。端口占用或绘制失败不会终止推理；关闭浏览器也不会停止推理。 + +摄像头限定时长、模拟门状态： + +```bash +./run_camera.sh --no-can --active-door both --duration 30 --save-results --output outputs/demo +``` + +有人视频回放（同一 TensorRT/融合/风险链路）： + +```bash +./run_camera.sh --video tests/person-test.mp4 --no-can --active-door both \ + --max-frames 150 --save-results --save-pose-records --output outputs/replay +``` + +`tests/person-test.mp4` 来自用户提供的 `segment_009.mp4` 第 10 秒起的 150 帧片段。 + +## 下游直接读取结果变量 + +参考 0920 的 `LatestResult` / `CameraPipeline.latest()`，0923 提供同进程 Python 接口： + +```python +from run_camera import CameraPipeline, parse_args + +args = parse_args([]) # 默认 CAN 模式；不打开可视化，不保存文件 +with CameraPipeline(args) as pipeline: + result = pipeline.wait_for_result(timeout=10) # 等待第一帧，可选 + # 在下游自己的循环/线程中读取当前完整快照： + result = pipeline.latest_result # 等价于 pipeline.latest() + if result is not None: + risk = result["risk"] + door_risks = result["door_risks"] + skeletons = result["skeletons"] + can = result["can"] + print(result["frame_index"], risk["risk_type"], risk["risk_score"]) +``` + +完整持续读取示例：`examples/read_latest.py`，在 Orin 执行。示例会自己启动推理并占用摄像头，运行前先停止已有后台进程： + +```bash +./scripts/stop.sh +python3 examples/read_latest.py +# 或开启可视化： +python3 examples/read_latest.py --visualize +``` + +- `latest_result` 是单个最新结果变量，首帧前为 `None`。每帧完成风险计算后发布，早于任何可选绘制或文件写入。 +- 快照深度只读：字典使用只读映射，数组使用 tuple；字段和索引含义与原 JSON 一致。旧快照不会随下一帧变化，多个读取线程互不消费或修改结果。 +- 只保留最新一帧，没有结果队列；下游处理慢可跳过中间帧。需要每帧历史时显式开启 `--save-results`。 +- `pipeline.wait_for_result(after_frame_index, timeout)` 是可选等待接口，返回更新的最新结果；超时或结束且无新帧时返回 `None`。 +- `pipeline.request_stop()` 请求停止，`pipeline.join()` 等待退出并向调用方传播推理异常；停止后可读取最后快照及 `pipeline.summary`。使用 `with` 自动停止并回收资源。 +- 停止后的快照只是最后结果；持续控制时应同时检查 `pipeline.running()` 及帧序号/时间戳，避免使用旧结果。 +- 需要修改副本或 JSON 序列化时使用 `from run_camera import result_to_dict`，再调用 `result_to_dict(result)`。 +- 变量属于创建 `CameraPipeline` 的 Python 进程。另一个进程或 C++ 程序不能直接 import 读取后台服务的变量，需要另行使用 IPC/HTTP 接口；当前没有新增 C++ 风险结果 ABI。 + +## 可选输出开关 + +| 参数 | 行为 | +|---|---| +| 无输出参数 | 只发布内存结果，不绘制、不保存结果文件 | +| `--visualize` | 开启独立可视化线程和 HTTP 预览 | +| `--save-results` | 保存逐帧 JSONL 及约每 10 秒的 latest.json | +| `--save-snapshots` | 约每 10 秒绘制并保存 latest.jpg，无 HTTP 服务 | +| `--save-pose-records` | 保存诊断姿态 JSONL | +| `--write-video` | 绘制并保存叠加 MP4，无需打开预览 | + +输出开关可组合。`--output` 仅指定保存目录，不会自动开启保存。未开启保存时，旧输出文件保留但不会更新。任一保存开关开启后，停止时额外保存 summary.json；默认内存模式的统计保存在 `pipeline.summary` 并打印到终端。 + +恢复此前的预览与文件输出行为： + +```bash +./run_camera.sh --visualize --save-results --save-snapshots --output outputs/demo +``` + +## CAN 约定 + +按用户确认的 `0x391` 数据布局： + +| 字节 | 车门 | +|---|---| +| data[0] | 左前 | +| data[1] | 右前 | +| data[2] | 左后 | +| data[3] | 右后 | + +- `--can-side left` 使用 data[0]/data[2]；`right` 使用 data[1]/data[3]。默认 left，请按摄像头安装侧选择。 +- 值为 0 时对应门活跃，每帧读取最新 CAN 快照。有效报文至少需要 4 字节。 +- 未收到有效报文或超过 `--can-timeout-ms`（默认 1000 ms），整体输出 `CAN_UNKNOWN`，门级输出 `FRONT_UNKNOWN/REAR_UNKNOWN`，分数为 null。 +- 此时 `active_door=none` 仅为内部无有效活跃门输入，业务必须读取 `risk` 和 `can.valid`，不能将其解释为已确认两门关闭。 +- 显式 `--active-door front/rear/both/none` 是模拟模式，输出 `can.source=simulation`，预览显示 CAN SIMULATION。`--no-can` 要求显式选择模拟状态。 +- 本项目提供最终风险结果；基线未提供防夹执行器的 CAN 指令协议，因此没有自行添加执行器报文。 +- 本次 USB-CAN `VCI_OpenDevice` 返回失败，`lsusb` 中未见 USB-CAN 设备；真实 0x391 报文链路尚未验证。连接适配器后重新启动程序；字节映射及门状态组合通过软件测试。 + +## 风险行为与机位标定 + +风险算法、默认 warning=0.30 / release=0.45、133 点骨骼、track_id、前后门归属、transition→unknown、无有效手部时人体框兜底均沿用原版。详见 `README-risk-upstream.md`。 + +`annotations/constant_door_mask.json` 和 `door_zone_multifeature.json` 原样复制，属于原研发机位。当前室内摄像头位置与其不一致，默认输出 `calibration_verified=false`，预览显示 CALIBRATION UNVERIFIED。实车验证需更换标定文件；`--calibration-verified` 仅标记经人工验证的标定，不会自动标定。 + +当前实时场景没有人，约 29 FPS 是该场景的实测值，不能当作多人满负载帧率。 + +## 优化与输出 + +- 重用 0920 的常驻 TensorRT engine、CUDA/NPP 预处理、显存及锁页缓存、GPU UYVY 转换。 +- 独立取帧线程只保留最新帧，主动丢弃旧帧，避免积压；不像 0920 那样单独并行检测线程，完整 Python 融合阶段按结果顺序执行。 +- 风险 mask / 距离变换按原版缓存；开启预览时默认每 3 帧提交一次待绘制图像；开启快照保存时每 10 秒一次。 +- 长期运行清理过期轨迹，计时统计仅保留最近 3000 帧。 +- 开启 `--save-results` 后，`causal_results.jsonl` 默认每 32 MiB 轮换，保留 `.1` 和 `.2`；`--jsonl-max-mb 0` 可禁用轮换。读取全部历史时按 `.2 → .1 → 当前文件` 顺序处理。 +- `--save-pose-records` 额外保存同样轮换的诊断姿态记录；`--write-video` 录制叠加 MP4（视频不轮换，建议配合 `--duration`）。 +- 文件输出遵循上面的独立开关；关闭可视化仍可完整读取内存结果。 +- 新增 `capture`、`can`、`capture_to_result_ms`、`calibration_verified` 字段，原版核心 JSON 字段保留。 +- 驱动原始时间戳存在时钟偏移，保存在 `capture.driver_timestamp_ms` 供诊断。延迟采用主机单调时钟在 DQBUF 后取样，表示“主机取帧到风险结果”，不包含曝光到驱动交付的时间。 + +## 文件关系与测试 + +- `main.cpp`, `v4l2_capture.*`, `can_comm.*`, `include/zcan.h`, `lib/`：基线工程。新工程修正错误 V4L2 buffer 处理；旧 main 同步修正 CAN 字节解析。 +- `native/native_api.cpp`：采集线程、CAN 快照、TensorRT 的 ctypes C 接口。 +- `accelerator/`：远端 0920 源码快照，当前构建只选取 TensorRT/预处理/检测/姿态源文件。 +- `pipeline/`：原版 Python 算法，推理适配在 `native_backend.py`；ONNX/Ultralytics 的可选导入移到旧离线入口，原版风险公式不变。 +- `Makefile.baseline` 和 `make baseline`：旧基线演示程序的编译方式；默认 `make` 编译新的原生共享库。 +- `pipeline/latest_result.py`：线程安全的单结果发布槽位及不可变快照。 +- `examples/read_latest.py`：同进程下游读取示例。 +- `tests/test_latest_result.py`：结果所有权、线程生命周期、可选输出、独立可视化及异常传播测试。 +- `tests/test_risk_integration.py`：CAN 左右门组合、失效状态、手部/人体兜底、transition、轨迹回收、日志轮换测试。 +- `TEST_REPORT.md`：实测结果和限制。 + +```bash +python3 -m unittest discover -s tests -v +``` + +本机重新部署源码（会使用 SSH 交互认证，不保存密码）： + +```bash +./scripts/deploy.sh orin@192.168.195.100 +``` + +重新编译共享库前应先停止 0923 运行进程，再启动新进程加载新库。 diff --git a/TEST_REPORT.md b/TEST_REPORT.md new file mode 100644 index 0000000..51b00e4 --- /dev/null +++ b/TEST_REPORT.md @@ -0,0 +1,56 @@ +# 2026-09-23 集成验证 + +执行设备：192.168.195.100 / Orin；摄像头 /dev/video5，UYVY 1920×1536。当前机位朝向室内设备，没有测试人员；用户确认暂不调整机位。 + +| 验证 | 结果 | +|---|---| +| 原生 C++/CUDA 共享库 | Orin Release 编译通过 | +| 摄像头 + CAN 模式，30 秒 | 882 帧，29.399 FPS，无运行异常 | +| 主机取帧到风险结果 | 均值 42.369 ms，P95 57.368 ms | +| latest-frame 策略 | 30 秒期间丢弃 16 个积压帧 | +| 真实 CAN | USB-CAN 打开失败，lsusb 未见该设备；882 帧均为 CAN_UNKNOWN，score=null | +| 有人视频 150 帧回放 | 32.898 FPS，31 帧有 1 人，119 帧无有效人体框 | +| RTMW 骨骼 | 31 个人体记录，每个 133 点 | +| 风险证据 | 22 个手部区域，9 个人体框兜底区域 | +| 视频门级风险 | clear 132，watch 5，warning 2，unknown 11 | +| 原始 run_risk.py 对照 | 同一姿态记录重放 150 帧，核心输出完全一致 | +| 行为测试 | 本机与 Orin 的 8 项测试均通过 | +| HTTP 实时预览 | 本机访问远端 /result.json 成功，连续返回新帧 | + +对照方式：将远端 `outputs/person-replay/pose_records.jsonl` 下载到本机，用用户原始工程的 `run_risk.py --pose-jsonl` 运行，逐帧比较原始 JSON 的所有顶层字段，包括 skeletons、door_assignments、door_risks、risk、ungated_risk。仅排除不同工程目录导致的绝对 mask 路径差异；150 帧无差异。**此项对照验证相同姿态输入下的风险集成，不是 TensorRT 与 ONNX 输出的数值误差测试。** + +证据： + +- `outputs/camera-can/summary.json` +- `outputs/person-replay/summary.json` 与 `causal_results.jsonl` +- `outputs/reference-replay/summary.json` 与 `causal_results.jsonl` +- `outputs/reference-parity.json` +- `logs/build.log`, `logs/camera-can.log`, `logs/person-replay.log` + +`outputs/camera-smoke/` 是初次调通记录，使用了不可靠的驱动时间戳，延迟统计无效。已修正计时来源，性能结论以 `camera-can` 为准。 + +未完成的硬件验证：有效 CAN 报文到达与门切换、真实摄像头下人员/手部靠近门区、实车机位 mask 与前后门参数标定。当前无法据此宣称已完成实车防夹效果验收。 + +本机没有摄像头设备节点；本机完成源码/模型归档、风险对照及行为测试，通过网络预览 Orin 的实时结果。 + + +## 最新结果变量与可选可视化改造 + +已对照重启后远端 0920 的 `include/latest_result.hpp` 和 `src/main.cpp`，采用默认关闭可视化、单个不可变最新快照、start/latest/request_stop/join 生命周期。 + +- 本机与 Orin：17 项测试通过，覆盖原有风险测试及只读快照、旧快照保持、跨线程等待/停止、错误传播、默认无绘制/HTTP/结果文件、JSON 与内存结果一致、预览端口失败不影响推理、独立可视化线程。 +- Orin 真实 TensorRT 视频回放：下游示例读取 90 帧，27.321 FPS；未开启可视化和文件保存。保留 watch/warning/unknown 等风险输出。 +- Orin 真实摄像头：下游示例读取 90 帧，28.858 FPS；主机取帧至结果均值 32.341 ms；关闭可视化，没有创建指定输出目录。 +- `--visualize` 实机运行 15 秒：成功通过 HTTP 读取逐帧结果及 JPEG 图像；未开启文件保存，没有创建指定输出目录。 +- 后台运行入口默认无可视化；开启预览需显式加 `--visualize`。 +- 重启后摄像头节点未自动生成。0920 初始化脚本因驱动缺少 vermagic 元数据而判定不匹配；0923 新增初始化脚本，存在 vermagic 时检查内核版本，缺失时交给正常 insmod 的内核检查，未使用强制加载参数。加载原有 0920 驱动后 /dev/video5 恢复。 + +日志：`logs/tests-optional-local.log`、`logs/tests-optional-orin.log`、`logs/result-api-video.log`、`logs/result-api-camera.log`、`logs/optional-preview.log`、`logs/optional-preview-probe.json`。 + +## use-can Orin 可运行包验证 + +- 独立运行包位于 `orin/`：包含预编译 ARM64 原生库、TensorRT 引擎、Python 风险代码、CAN USB 库、GMSL 驱动和 ./run.sh。 +- 在 Orin 新建的独立临时目录复制整个运行包，SHA256SUMS 全部通过。 +- 从临时目录执行 `./run.sh --max-frames 90`，实际摄像头 90 帧完成，29.080 FPS，主机取帧至风险结果平均 36.310 ms；无结果文件输出。 +- 原生库 RUNPATH 为 `$ORIGIN/../lib:/usr/local/cuda/lib64`，启动脚本按自身位置定位所有项目资产。 +- 系统依赖仍为目标 JetPack 的 CUDA / TensorRT / OpenCV 与 Python numpy/cv2，不包含操作系统或完整 Python 环境。 diff --git a/accelerator/.gitattributes b/accelerator/.gitattributes new file mode 100644 index 0000000..5f609e2 --- /dev/null +++ b/accelerator/.gitattributes @@ -0,0 +1,5 @@ +*.sh text eol=lf +*.engine -text +*.ko -text +/orin/** -text +/config/door_mask.json -text diff --git a/accelerator/.gitignore b/accelerator/.gitignore new file mode 100644 index 0000000..4097644 --- /dev/null +++ b/accelerator/.gitignore @@ -0,0 +1,4 @@ +/build/ +/engines/ +/logs/ +*.tmp diff --git a/accelerator/CMakeLists.txt b/accelerator/CMakeLists.txt new file mode 100644 index 0000000..1a5dede --- /dev/null +++ b/accelerator/CMakeLists.txt @@ -0,0 +1,23 @@ +cmake_minimum_required(VERSION 3.24) +if(NOT DEFINED CMAKE_CUDA_ARCHITECTURES) + set(CMAKE_CUDA_ARCHITECTURES native) +endif() +project(camera_risk LANGUAGES CXX CUDA) +set(CMAKE_CXX_STANDARD 17) +set(CMAKE_CXX_STANDARD_REQUIRED ON) +set(CMAKE_CUDA_STANDARD 17) +find_package(CUDAToolkit REQUIRED) +find_package(OpenCV REQUIRED COMPONENTS core imgproc imgcodecs highgui) +find_package(Threads REQUIRED) +find_path(TENSORRT_INCLUDE_DIR NvInfer.h PATH_SUFFIXES aarch64-linux-gnu x86_64-linux-gnu REQUIRED) +find_library(TENSORRT_LIBRARY nvinfer REQUIRED) +add_library(camera_risk_core + src/trt_engine.cpp src/pose_engine.cu src/gpu_preprocess.cu src/uyvy_converter.cu + src/detector.cpp src/door_risk.cpp src/camera.cpp src/camera_pipeline.cpp src/config.cpp) +target_include_directories(camera_risk_core PUBLIC include third_party ${TENSORRT_INCLUDE_DIR} ${OpenCV_INCLUDE_DIRS}) +target_link_libraries(camera_risk_core PUBLIC ${TENSORRT_LIBRARY} ${OpenCV_LIBS} Threads::Threads CUDA::cudart CUDA::nppc CUDA::nppig CUDA::nppidei CUDA::nppicc) +target_compile_options(camera_risk_core PRIVATE $<$:-Wall;-Wextra;-Wpedantic>) +add_library(camera_risk_visualizer src/visualizer.cpp) +target_link_libraries(camera_risk_visualizer PUBLIC camera_risk_core) +add_executable(camera_risk src/main.cpp) +target_link_libraries(camera_risk PRIVATE camera_risk_visualizer) diff --git a/accelerator/README.md b/accelerator/README.md new file mode 100644 index 0000000..1f815a6 --- /dev/null +++ b/accelerator/README.md @@ -0,0 +1,76 @@ +# Vision Anti-Pinch + +C++ / TensorRT 实时摄像头风险推理，算法来自最新 run_risk.py。 +持续取图、推理并更新内存中的 FrameResult；可视化可选。 + +## Orin 直接运行 + +适用环境：Jetson AGX Orin、JetPack 5.1.2 / L4T R35.4.1、CUDA 11.4、TensorRT 8.5.2、OpenCV 4.5.4。 +仓库已包含该板端编译的程序和重新构建的 TensorRT 引擎，无需先编译或转换。 + +```bash +git clone https://git.kendeji.fun/zmv/vision-anti-pinch.git +cd vision-anti-pinch/orin +./scripts/initialize_camera.sh # 开机后初始化 /dev/video5 摄像头链路 +./run.sh # 持续推理，结果只保存在内存 +``` + +可选板端本地可视化： + +```bash +./run.sh --visualize +``` + +Ctrl+C / SIGTERM 停止程序。没有有限帧选项；正常运行和退出不输出统计。 +关闭预览窗口或按 q/Esc 只关闭预览，推理继续。没有图像网络服务、端口转发、JSONL 或视频输出。 + +摄像头为 /dev/video5，1920×1536 UYVY，30 FPS。 +当前 door mask 是指定的 constant_door_mask.json 副本，标定状态为 false；装车后应更换与实际画面匹配的门区并更新标定状态。 +即使不开启可视化，风险计算仍需要门区。 + +## 下游获取结果 + +main.cpp 注释中也包含以下示例。下游通过自己的调用位置或线程读取： + +```cpp +auto result = pipeline.latest(); +if (result) { + const auto& risk = result->risk; + const auto& people = result->persons; + // result->frame_index / track_ids / boxes + // result->image / door_mask_pixels / capture / timing + // people[i].selected / smoothed / temporal：各 133 个关键点 +} +``` + +latest() 返回 shared_ptr；首帧尚未完成时为空。 +读取不删除结果，也不产生消费状态。推理不关心是否有下游读取。 +下游持有的结果保持有效，后续更新不会改变旧快照；历史保存由下游自行负责。 +公共定义位于 include/camera_pipeline.hpp 和 include/frame_result.hpp。 + +## 从源码构建 + +```bash +./scripts/build.sh +./scripts/initialize_camera.sh +./run_camera.sh +``` + +源码构建要求 CMake 3.24+、CUDA/NPP、TensorRT 和 OpenCV 开发依赖。 +构建脚本会优先使用 Orin 上已有的 ~/miniconda3/envs/yolo/bin/cmake。 +源码配置直接使用 orin/engines/ 中的板端引擎。该引擎仅适用于上述 Orin 环境；仓库不包含 ONNX 或模型转换脚本。 +独立库目标为 camera_risk_core；可视化在 camera_risk_visualizer。 + +## 目录 + +- src/、include/：采集、推理、结果变量及可选绘制。 +- config/：源码构建程序的配置和固定门区。 +- scripts/：构建和摄像头初始化脚本。 +- orin/：已编译程序、TensorRT 引擎、配置、当前 GMSL 驱动及启动脚本。 +- third_party/：nlohmann JSON 及其许可证。 + +NPP 用于 UYVY 转换、YOLO 缩放填充、RTMW 仿射裁剪；CUDA 用于通道转换与归一化。 +风险几何计算仍包含 CPU/OpenCV 操作；算法内部的临时 JSON 对象只在内存中使用。 + +当前摄像头驱动启动时偶发首帧 ERROR，应用会丢弃并计数；持续错误、采集超时或推理错误会明确报错并退出。 +现有 GMSL 驱动停止采集时可产生 videobuf2 缓冲区告警，v4l2-ctl 同样复现；需在驱动侧完成量产修复和验证。 diff --git a/config/camera.json b/accelerator/config/camera.json similarity index 100% rename from config/camera.json rename to accelerator/config/camera.json diff --git a/config/door_mask.json b/accelerator/config/door_mask.json similarity index 100% rename from config/door_mask.json rename to accelerator/config/door_mask.json diff --git a/include/bounded_queue.hpp b/accelerator/include/bounded_queue.hpp similarity index 100% rename from include/bounded_queue.hpp rename to accelerator/include/bounded_queue.hpp diff --git a/include/camera.hpp b/accelerator/include/camera.hpp similarity index 100% rename from include/camera.hpp rename to accelerator/include/camera.hpp diff --git a/include/camera_pipeline.hpp b/accelerator/include/camera_pipeline.hpp similarity index 100% rename from include/camera_pipeline.hpp rename to accelerator/include/camera_pipeline.hpp diff --git a/include/common.hpp b/accelerator/include/common.hpp similarity index 100% rename from include/common.hpp rename to accelerator/include/common.hpp diff --git a/include/config.hpp b/accelerator/include/config.hpp similarity index 100% rename from include/config.hpp rename to accelerator/include/config.hpp diff --git a/include/detector.hpp b/accelerator/include/detector.hpp similarity index 100% rename from include/detector.hpp rename to accelerator/include/detector.hpp diff --git a/include/door_risk.hpp b/accelerator/include/door_risk.hpp similarity index 100% rename from include/door_risk.hpp rename to accelerator/include/door_risk.hpp diff --git a/include/frame_result.hpp b/accelerator/include/frame_result.hpp similarity index 100% rename from include/frame_result.hpp rename to accelerator/include/frame_result.hpp diff --git a/include/gpu_preprocess.hpp b/accelerator/include/gpu_preprocess.hpp similarity index 100% rename from include/gpu_preprocess.hpp rename to accelerator/include/gpu_preprocess.hpp diff --git a/include/image.hpp b/accelerator/include/image.hpp similarity index 100% rename from include/image.hpp rename to accelerator/include/image.hpp diff --git a/include/latest_result.hpp b/accelerator/include/latest_result.hpp similarity index 100% rename from include/latest_result.hpp rename to accelerator/include/latest_result.hpp diff --git a/include/pipeline_types.hpp b/accelerator/include/pipeline_types.hpp similarity index 100% rename from include/pipeline_types.hpp rename to accelerator/include/pipeline_types.hpp diff --git a/include/pose_engine.hpp b/accelerator/include/pose_engine.hpp similarity index 100% rename from include/pose_engine.hpp rename to accelerator/include/pose_engine.hpp diff --git a/include/records.hpp b/accelerator/include/records.hpp similarity index 100% rename from include/records.hpp rename to accelerator/include/records.hpp diff --git a/include/result_conversion.hpp b/accelerator/include/result_conversion.hpp similarity index 100% rename from include/result_conversion.hpp rename to accelerator/include/result_conversion.hpp diff --git a/include/stop_signals.hpp b/accelerator/include/stop_signals.hpp similarity index 100% rename from include/stop_signals.hpp rename to accelerator/include/stop_signals.hpp diff --git a/include/trt_engine.hpp b/accelerator/include/trt_engine.hpp similarity index 100% rename from include/trt_engine.hpp rename to accelerator/include/trt_engine.hpp diff --git a/include/uyvy_converter.hpp b/accelerator/include/uyvy_converter.hpp similarity index 100% rename from include/uyvy_converter.hpp rename to accelerator/include/uyvy_converter.hpp diff --git a/include/visualizer.hpp b/accelerator/include/visualizer.hpp similarity index 100% rename from include/visualizer.hpp rename to accelerator/include/visualizer.hpp diff --git a/accelerator/run_camera.sh b/accelerator/run_camera.sh new file mode 100755 index 0000000..c153d56 --- /dev/null +++ b/accelerator/run_camera.sh @@ -0,0 +1,4 @@ +#!/usr/bin/env bash +set -euo pipefail +root="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)" +exec "$root/build/camera_risk" --config "$root/config/camera.json" "$@" diff --git a/accelerator/scripts/build.sh b/accelerator/scripts/build.sh new file mode 100755 index 0000000..5b0fa8d --- /dev/null +++ b/accelerator/scripts/build.sh @@ -0,0 +1,7 @@ +#!/usr/bin/env bash +set -euo pipefail +root="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")/.." && pwd)" +cmake_bin="${CMAKE_COMMAND:-cmake}" +if [[ -x "$HOME/miniconda3/envs/yolo/bin/cmake" ]]; then cmake_bin="${CMAKE_COMMAND:-$HOME/miniconda3/envs/yolo/bin/cmake}"; fi +"$cmake_bin" -S "$root" -B "$root/build" -DCMAKE_BUILD_TYPE=Release -DCMAKE_CUDA_COMPILER="${CUDACXX:-/usr/local/cuda/bin/nvcc}" "$@" +"$cmake_bin" --build "$root/build" -j4 diff --git a/accelerator/scripts/initialize_camera.sh b/accelerator/scripts/initialize_camera.sh new file mode 100755 index 0000000..8da09bd --- /dev/null +++ b/accelerator/scripts/initialize_camera.sh @@ -0,0 +1,4 @@ +#!/usr/bin/env bash +set -euo pipefail +root="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")/.." && pwd)" +exec "$root/orin/scripts/initialize_camera.sh" "$@" diff --git a/src/camera.cpp b/accelerator/src/camera.cpp similarity index 100% rename from src/camera.cpp rename to accelerator/src/camera.cpp diff --git a/src/camera_pipeline.cpp b/accelerator/src/camera_pipeline.cpp similarity index 100% rename from src/camera_pipeline.cpp rename to accelerator/src/camera_pipeline.cpp diff --git a/src/config.cpp b/accelerator/src/config.cpp similarity index 100% rename from src/config.cpp rename to accelerator/src/config.cpp diff --git a/src/detector.cpp b/accelerator/src/detector.cpp similarity index 100% rename from src/detector.cpp rename to accelerator/src/detector.cpp diff --git a/src/door_risk.cpp b/accelerator/src/door_risk.cpp similarity index 100% rename from src/door_risk.cpp rename to accelerator/src/door_risk.cpp diff --git a/src/gpu_preprocess.cu b/accelerator/src/gpu_preprocess.cu similarity index 100% rename from src/gpu_preprocess.cu rename to accelerator/src/gpu_preprocess.cu diff --git a/src/main.cpp b/accelerator/src/main.cpp similarity index 100% rename from src/main.cpp rename to accelerator/src/main.cpp diff --git a/src/pose_engine.cu b/accelerator/src/pose_engine.cu similarity index 100% rename from src/pose_engine.cu rename to accelerator/src/pose_engine.cu diff --git a/src/trt_engine.cpp b/accelerator/src/trt_engine.cpp similarity index 100% rename from src/trt_engine.cpp rename to accelerator/src/trt_engine.cpp diff --git a/src/uyvy_converter.cu b/accelerator/src/uyvy_converter.cu similarity index 100% rename from src/uyvy_converter.cu rename to accelerator/src/uyvy_converter.cu diff --git a/src/visualizer.cpp b/accelerator/src/visualizer.cpp similarity index 100% rename from src/visualizer.cpp rename to accelerator/src/visualizer.cpp diff --git a/third_party/nlohmann/LICENSE.MIT b/accelerator/third_party/nlohmann/LICENSE.MIT similarity index 100% rename from third_party/nlohmann/LICENSE.MIT rename to accelerator/third_party/nlohmann/LICENSE.MIT diff --git a/third_party/nlohmann/json.hpp b/accelerator/third_party/nlohmann/json.hpp similarity index 100% rename from third_party/nlohmann/json.hpp rename to accelerator/third_party/nlohmann/json.hpp diff --git a/orin/config/door_mask.json b/annotations/constant_door_mask.json similarity index 100% rename from orin/config/door_mask.json rename to annotations/constant_door_mask.json diff --git a/annotations/door_zone_multifeature.json b/annotations/door_zone_multifeature.json new file mode 100644 index 0000000..3432c66 --- /dev/null +++ b/annotations/door_zone_multifeature.json @@ -0,0 +1,47 @@ +{ + "version": "v9.10-research-calibration-1", + "description": "Temporary-camera calibration from seg_01 and seg_08; recalibrate for the production camera.", + "front_probability_min": 0.72, + "rear_probability_max": 0.28, + "minimum_person_height": 0.4, + "minimum_pose_features": 4, + "history_frames": 5, + "stale_track_frames": 120, + "features": { + "shoulder_center_x": { + "rear": 0.53, + "front": 0.64, + "weight": 1.2 + }, + "shoulder_center_y": { + "rear": 0.54, + "front": 0.46, + "weight": 1.3 + }, + "shoulder_width": { + "rear": 0.14, + "front": 0.2, + "weight": 1.4 + }, + "torso_height": { + "rear": 0.35, + "front": 0.47, + "weight": 1.5 + }, + "hip_center_y": { + "rear": 0.89, + "front": 0.93, + "weight": 0.6 + }, + "box_center_x": { + "rear": 0.48, + "front": 0.58, + "weight": 0.35 + }, + "box_height": { + "rear": 0.7, + "front": 0.82, + "weight": 0.45 + } + } +} diff --git a/build_and_run.sh b/build_and_run.sh new file mode 100644 index 0000000..5ec936a --- /dev/null +++ b/build_and_run.sh @@ -0,0 +1,101 @@ +#!/bin/bash +# build_and_run.sh - 在 Orin 上构建并运行 V4L2 取图 demo +# +# 使用方法: +# ./build_and_run.sh [选项] +# +# 选项: +# --deploy 将源码复制到 Orin 并构建 +# --run 在 Orin 上运行已构建的程序 +# --both 部署并运行 (默认) +# +# 示例: +# ./build_and_run.sh --both +# ./build_and_run.sh --deploy +# ./build_and_run.sh --run --args "-n 100 -v" + +set -e + +# 配置 +ORIN_HOST="192.168.195.100" +ORIN_USER="orin" +ORIN_PASS="${ORIN_PASS:?Set ORIN_PASS locally, or use scripts/deploy.sh for interactive SSH}" +export SSHPASS="$ORIN_PASS" +LOCAL_DIR="$(cd "$(dirname "$0")" && pwd)" +BUILD_DIR="$LOCAL_DIR/build" +BIN_NAME="vision-anti-pinch" + +# 默认参数 +ACTION="both" +RUN_ARGS="-n 100" + +# 解析参数 +while [[ $# -gt 0 ]]; do + case $1 in + --deploy) ACTION="deploy"; shift ;; + --run) ACTION="run"; shift ;; + --both) ACTION="both"; shift ;; + --args) + RUN_ARGS="$2" + shift 2 + ;; + *) + echo "Unknown option: $1" + exit 1 + ;; + esac +done + +# 部署函数 +do_deploy() { + echo "=== 部署源码到 Orin ===" + + # 创建临时目录用于传输 + TEMP_DIR=$(mktemp -d) + cp -r "$LOCAL_DIR"/* "$TEMP_DIR/" + + # 使用 sshpass 复制 (如果没有 sshpass, 手动复制) + if command -v sshpass &> /dev/null; then + echo "复制源码到 Orin..." + sshpass -e scp -o StrictHostKeyChecking=no -r "$TEMP_DIR"/* "${ORIN_USER}@${ORIN_HOST}:~/vision-anti-pinch/" + else + echo "警告: 未安装 sshpass, 请手动复制源码到 Orin:" + echo " scp -r $LOCAL_DIR/* ${ORIN_USER}@${ORIN_HOST}:~/vision-anti-pinch/" + return 1 + fi + + rm -rf "$TEMP_DIR" + + echo "=== 在 Orin 上构建 ===" + sshpass -e ssh -o StrictHostKeyChecking=no "${ORIN_USER}@${ORIN_HOST}" << 'EOF' +cd ~/vision-anti-pinch +make clean +make -j$(nproc) +echo "构建完成" +EOF + + echo "部署完成" +} + +# 运行函数 +do_run() { + echo "=== 在 Orin 上运行 ===" + sshpass -e ssh -o StrictHostKeyChecking=no "${ORIN_USER}@${ORIN_HOST}" << EOF +cd ~/vision-anti-pinch +./${BIN_NAME} ${RUN_ARGS} +EOF +} + +# 主流程 +echo "V4L2 Capture Demo - Build & Run" +echo "================================" + +if [[ "$ACTION" == "deploy" ]] || [[ "$ACTION" == "both" ]]; then + do_deploy +fi + +if [[ "$ACTION" == "run" ]] || [[ "$ACTION" == "both" ]]; then + do_run +fi + +echo "完成" diff --git a/camera_reset.sh b/camera_reset.sh new file mode 100644 index 0000000..9ce7aaa --- /dev/null +++ b/camera_reset.sh @@ -0,0 +1,110 @@ +#!/bin/bash +# camera_reset.sh - 修复 GMSL2 相机"第二次运行全绿"问题 +# +# 问题原因： +# quick_bring_up_64g.sh 第二次运行时，仅重新设置 sensor_mode/trig_pin， +# 但 sgx_yuv_gmsl2 链路状态未重新初始化，导致相机输出全绿。 +# +# 模块依赖关系 (modinfo 确认): +# sgx-yuv-gmsl2.ko depends: max9296, max9295 +# 必须按顺序加载: max9295.ko -> max9296.ko -> sgx-yuv-gmsl2.ko +# 否则 insmod sgx-yuv-gmsl2.ko 报 "Unknown symbol in module" +# +# 本脚本: 依赖模块仅在未加载时 insmod (不重载, 避免链路抖动), +# sgx_yuv_gmsl2 先 rmmod 再 insmod 完成链路复位。 +# 全绿图像多由 gmsl_mode 错误引起 (本工程固定 2,2,2,2); +# 若复位后仍异常, 请重启 Orin 后运行 quick_bring_up_64g.sh。 +# +# 使用方法： +# sudo ./camera_reset.sh [camera_type] [port] +# +# camera_type: 0-17 对应 quick_bring_up_64g.sh 中的相机类型 (默认 8) +# port: 0-7 对应 video 设备号 (默认 5) + +set -e + +SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)" +KO_DIR="/home/orin/SG8A_ORIN_GMSL2-F_V2_AGX_Orin_YUV_JP5.1.2_L4TR35.4.1/ko" +VIDEO_PORT="${2:-5}" +CAMERA_TYPE="${1:-5}" # 默认 camera_type=5 (SG3S-ISX031C-GMSL2F) + +echo "=========================================================" +echo "Camera Reset Script" +echo "=========================================================" +echo "Camera type: $CAMERA_TYPE" +echo "Video port: $VIDEO_PORT (/dev/video${VIDEO_PORT})" +echo "" + +# Step 1: 停止占用 video 设备的进程 +echo "[1/4] Stopping processes using video devices..." +pkill -9 gst-launch-1.0 2>/dev/null || true +pkill -9 vision-anti-pinch 2>/dev/null || true +sleep 1 + +# Step 2: 加载依赖模块 max9295/max9296 (未加载时), 再重载 sgx_yuv_gmsl2 +echo "[2/4] Loading modules (max9295 -> max9296 -> sgx_yuv_gmsl2)..." +if ! lsmod | grep -q '^max9295'; then + sudo insmod "$KO_DIR/max9295.ko" || { echo "ERROR: insmod max9295 failed"; exit 1; } + echo " loaded max9295.ko" +else + echo " max9295 already loaded (skip)" +fi +if ! lsmod | grep -q '^max9296'; then + sudo insmod "$KO_DIR/max9296.ko" || { echo "ERROR: insmod max9296 failed"; exit 1; } + echo " loaded max9296.ko" +else + echo " max9296 already loaded (skip)" +fi +if lsmod | grep -q '^sgx_yuv_gmsl2'; then + sudo rmmod sgx_yuv_gmsl2 + sleep 1 +fi +sudo insmod "$KO_DIR/sgx-yuv-gmsl2.ko" gmsl_mode=2,2,2,2 || { echo "ERROR: insmod sgx-yuv-gmsl2 failed"; exit 1; } +echo " loaded sgx-yuv-gmsl2.ko (gmsl_mode=2,2,2,2)" +sleep 2 + +# Step 3: 验证设备节点 +echo "[3/4] Verifying device nodes..." +if [ ! -c "/dev/video${VIDEO_PORT}" ]; then + echo "Error: /dev/video${VIDEO_PORT} not found!" + echo "Please reboot the system and run quick_bring_up_64g.sh again." + exit 1 +fi +echo " /dev/video${VIDEO_PORT} exists" + +# Step 4: 根据 camera_type 配置 sensor_mode 和 trig_pin +echo "[4/4] Configuring sensor..." +case "$CAMERA_TYPE" in + 0) v4l2-ctl -d /dev/video${VIDEO_PORT} -c sensor_mode=0,trig_pin=0xffff0007 ;; + 1) v4l2-ctl -d /dev/video${VIDEO_PORT} -c sensor_mode=0,trig_pin=0xffff0007 ;; + 2) v4l2-ctl -d /dev/video${VIDEO_PORT} -c sensor_mode=0,trig_pin=0xffff0007 ;; + 3) v4l2-ctl -d /dev/video${VIDEO_PORT} -c sensor_mode=1,trig_pin=0xffff0007 ;; + 4) v4l2-ctl -d /dev/video${VIDEO_PORT} -c sensor_mode=0,trig_pin=0xffff0007 ;; + 5) v4l2-ctl -d /dev/video${VIDEO_PORT} -c sensor_mode=1,trig_pin=0xffff0007 ;; + 6) v4l2-ctl -d /dev/video${VIDEO_PORT} -c sensor_mode=1,trig_pin=0xffff0007 ;; + 7) v4l2-ctl -d /dev/video${VIDEO_PORT} -c sensor_mode=2,trig_pin=0xffff0008,trig_mode=1 ;; + 8) v4l2-ctl -d /dev/video${VIDEO_PORT} -c sensor_mode=3,trig_pin=0xffff0007 ;; + 9) v4l2-ctl -d /dev/video${VIDEO_PORT} -c sensor_mode=3,trig_pin=0xffff0008 ;; + 10) v4l2-ctl -d /dev/video${VIDEO_PORT} -c sensor_mode=3,trig_pin=0xffff0007 ;; + 11) v4l2-ctl -d /dev/video${VIDEO_PORT} -c sensor_mode=3,trig_pin=0xffff0008 ;; + 12) v4l2-ctl -d /dev/video${VIDEO_PORT} -c sensor_mode=5,trig_pin=0xffff0007,trig_mode=3 ;; + 13) v4l2-ctl -d /dev/video${VIDEO_PORT} -c sensor_mode=4,trig_pin=0xffff0007,trig_mode=3 ;; + 14) v4l2-ctl -d /dev/video${VIDEO_PORT} -c sensor_mode=6,trig_pin=0xffff0007 ;; + 15) v4l2-ctl -d /dev/video${VIDEO_PORT} -c sensor_mode=8 ;; + 16) v4l2-ctl -d /dev/video${VIDEO_PORT} -c sensor_mode=9 ;; + 17) v4l2-ctl -d /dev/video${VIDEO_PORT} -c sensor_mode=0 ;; + *) echo "Unknown camera_type: $CAMERA_TYPE"; exit 1 ;; +esac +sleep 1 + +echo "" +echo "=========================================================" +echo "Camera reset complete!" +echo "" +echo "Now run: ./vision-anti-pinch -d /dev/video5 -W 1920 -H 1536 -n 100 -s 5 -t jpg" +echo "" +echo "⚠️ If image is still all-green, the GMSL2 link is damaged." +echo " Please REBOOT the Orin system to recover:" +echo " sudo reboot" +echo " (then run quick_bring_up_64g.sh after reboot)" +echo "=========================================================" diff --git a/can_comm.cpp b/can_comm.cpp new file mode 100644 index 0000000..84b1760 --- /dev/null +++ b/can_comm.cpp @@ -0,0 +1,338 @@ +// can_comm.cpp - USBCANFD CAN/CANFD 收发封装实现 +// +// 代码迁移自 vision-anti-pinch_bak/main.cpp 的 CAN 模块部分。 + +#include "can_comm.h" + +#include +#include +#include +#include + +#include "zcan.h" + +/* ===================================================================== + * 类型/宏/全局变量 + * ===================================================================== */ + +typedef uint32_t UINT; +typedef uint32_t U32; +#define msleep(ms) usleep((ms) * 1000) + +#define USBCANFD 33 +#define MAX_CHANNELS 2 // 最大通道数量 +#define RX_WAIT_TIME 100 +#define RX_BUFF_SIZE 1000 + +// 接收线程上下文 +typedef struct +{ + int dev_type; // 设备类型 + int dev_idx; // 设备索引 + int chn_idx; // 通道号 + int total; // 接收总数 + int stop; // 线程结束标志 +} THREAD_CTX; + +/* CAN 运行标志, 由信号处理函数(Ctrl+C)或外部置位 */ +static volatile bool g_can_running = true; + +static int g_receive_count = 0; + +/* ===== 仅供 can_start/can_stop 共享的文件内静态变量 ===== */ +static int s_DevType = USBCANFD; +static int s_DevIdx = 0; +static THREAD_CTX s_rx_ctx[MAX_CHANNELS]; +static pthread_t s_rx_threads[MAX_CHANNELS]; +static bool s_can_opened = false; + +namespace CanComm { + +bool is_open() { return s_can_opened; } + +int receive_count() { return g_receive_count; } + +/* ===== 报文解析回调 (由 main.cpp 通过 set_msg_callback 注册) ===== */ +static MsgCallback s_msg_cb = nullptr; + +void set_msg_callback(MsgCallback cb) { s_msg_cb = cb; } + +} // namespace CanComm + +/* ===================================================================== + * 帧构建与收发 + * ===================================================================== */ + +namespace CanComm { + +// 构建 CAN 帧 +static void construct_can_frame(ZCAN_20_MSG *can_msg, UINT id, int chn_idx, int pad) +{ + memset(can_msg, 0, sizeof(ZCAN_20_MSG)); + can_msg->hdr.inf.txm = 0; // 0-正常发送 + can_msg->hdr.inf.fmt = 0; // 0-CAN + can_msg->hdr.inf.sdf = 0; // 0-数据帧, 1-远程帧 + can_msg->hdr.inf.sef = 0; // 0-标准帧, 1-扩展帧 + + can_msg->hdr.id = id; // ID + can_msg->hdr.chn = chn_idx; // 通道 + can_msg->hdr.len = 8; // 数据长度 + + if (pad > 0) + { + can_msg->hdr.pad = pad; // 发送后延迟 pad ms + can_msg->hdr.inf.qsend = 1; // 队列发送帧,仅判断首帧 + can_msg->hdr.inf.qsend_100us = 0; // 队列发送单位,0-ms,1-100us + } + + for (int i = 0; i < can_msg->hdr.len; i++) + can_msg->dat[i] = i; +} + +// 构建 CANFD 帧 +static void construct_canfd_frame(ZCAN_FD_MSG *canfd_msg, UINT id, int chn_idx, int pad) +{ + memset(canfd_msg, 0, sizeof(ZCAN_FD_MSG)); + canfd_msg->hdr.inf.txm = 0; // 0-正常发送 + canfd_msg->hdr.inf.fmt = 1; // 0-CAN, 1-CANFD + canfd_msg->hdr.inf.sdf = 0; // 0-数据帧 CANFD只有数据帧! + canfd_msg->hdr.inf.sef = 1; // 0-标准帧, 1-扩展帧 + canfd_msg->hdr.inf.brs = 1; // canfd 加速 + + canfd_msg->hdr.id = id; // ID + canfd_msg->hdr.chn = chn_idx; // 通道 + canfd_msg->hdr.len = 64; // 数据长度 + + if (pad > 0) + { + canfd_msg->hdr.pad = 100; // 发送后延迟100ms + canfd_msg->hdr.inf.qsend = 1; // 队列发送帧,仅判断首帧 + canfd_msg->hdr.inf.qsend_100us = 0; // 队列发送单位,0-ms,1-100us + } + + for (int i = 0; i < canfd_msg->hdr.len; i++) + canfd_msg->dat[i] = i; +} + +// 分发报文到 main.cpp 注册的解析回调 +static inline void dispatch_msg(uint32_t ts, uint32_t msg_id, const uint8_t *data, uint32_t len) +{ + if (s_msg_cb) + s_msg_cb(ts, msg_id, data, len); +} + +// 接收线程 +static void *rx_thread(void *data) +{ + THREAD_CTX *ctx = (THREAD_CTX *)data; + int DevType = ctx->dev_type; + int DevIdx = ctx->dev_idx; + int chn_idx = ctx->chn_idx; + + ZCAN_20_MSG can_data[RX_BUFF_SIZE]; + ZCAN_FD_MSG canfd_data[RX_BUFF_SIZE]; + + while (!ctx->stop) + { + memset(can_data, 0, sizeof(can_data)); + memset(canfd_data, 0, sizeof(canfd_data)); + + int rcount = VCI_Receive(DevType, DevIdx, chn_idx, can_data, RX_BUFF_SIZE, RX_WAIT_TIME); // CAN + for (int i = 0; i < rcount; ++i) + { + if (can_data[i].hdr.inf.sdf == 0) + { // 数据帧 + dispatch_msg(can_data->hdr.ts, can_data[i].hdr.id, can_data[i].dat, can_data[i].hdr.len); + } + } + ctx->total += rcount; + + rcount = VCI_ReceiveFD(DevType, DevIdx, chn_idx, canfd_data, RX_BUFF_SIZE, RX_WAIT_TIME); // CANFD + for (int i = 0; i < rcount; ++i) + { + dispatch_msg(canfd_data->hdr.ts, canfd_data[i].hdr.id, canfd_data[i].dat, canfd_data[i].hdr.len); + } + ctx->total += rcount; + msleep(10); + } + g_receive_count = ctx->total; + printf("chn: %d receive %d\n", chn_idx, ctx->total); + pthread_exit(0); +} + +// 合并接收线程 +static void *rx_merge_thread(void *data) +{ + THREAD_CTX *ctx = (THREAD_CTX *)data; + int DevType = ctx->dev_type; + int DevIdx = ctx->dev_idx; + + ZCANDataObj obj_data[RX_BUFF_SIZE]; + + while (!ctx->stop) + { + memset(obj_data, 0, sizeof(obj_data)); + int rcount = VCI_ReceiveData(DevType, DevIdx, 0, obj_data, RX_BUFF_SIZE, RX_WAIT_TIME); // 第三个参数固定为0就好 + for (int i = 0; i < rcount; ++i) + { + if (obj_data[i].dataType == ZCAN_DT_ZCAN_CAN_DATA) + { // CAN + ZCAN_20_MSG *can_data = &obj_data[i].data.zcanCANData; + if (can_data->hdr.inf.sdf == 0) + { // 数据帧 + dispatch_msg(can_data->hdr.ts, can_data->hdr.id, can_data->dat, can_data->hdr.len); + } + } + else if (obj_data[i].dataType == ZCAN_DT_ZCAN_CANFD_DATA) + { // CANFD + ZCAN_FD_MSG *canfd_data = &obj_data[i].data.zcanCANFDData; + dispatch_msg(canfd_data->hdr.ts, canfd_data->hdr.id, canfd_data->dat, canfd_data->hdr.len); + } + } + ctx->total += rcount; + msleep(10); + } + g_receive_count = ctx->total; + printf("merge receive %d\n", ctx->total); + pthread_exit(0); +} + +/* + * 初始化并启动 CAN 收发。 + * 返回 true 成功; false 失败(主流程可继续仅取图)。 + */ +bool can_start() +{ + g_can_running = true; + int IsMerge = 1; // 是否合并接收 + + // 打开设备 + if (!VCI_OpenDevice(s_DevType, s_DevIdx, 0)) + { + printf("Open device fail\n"); + return false; + } + s_can_opened = true; + printf("Open device success\n"); + + // 初始化,启动通道 + for (int i = 0; i < MAX_CHANNELS; i++) + { + ZCAN_INIT init; // 波特率结构体,数据根据zcanpro的波特率计算器得出 + init.clk = 60000000; // clock: 60M(V1.01) 80M(V1.03即以上) + init.mode = 0; // 0-正常 + + init.aset.tseg1 = 14; // 仲裁域 500kbps + init.aset.tseg2 = 3; + init.aset.sjw = 2; + init.aset.smp = 0; + init.aset.brp = 5; + + init.dset.tseg1 = 10; // 数据域 2000kbps + init.dset.tseg2 = 2; + init.dset.sjw = 2; + init.dset.smp = 0; + init.dset.brp = 1; + + if (!VCI_InitCAN(s_DevType, s_DevIdx, i, &init)) // 初始化通道 + { + printf("InitCAN(%d) fail\n", i); + return false; + } + printf("InitCAN(%d) success\n", i); + + U32 on = 1; + if (!VCI_SetReference(s_DevType, s_DevIdx, i, CMD_CAN_TRES, &on)) // 终端电阻 + { + printf("CMD_CAN_TRES fail\n"); + } + + if (i == 0) + { + if (!VCI_SetReference(s_DevType, s_DevIdx, i, ZCAN_CMD_SET_CHNL_RECV_MERGE, &IsMerge)) // 合并接收 + { + printf("ZCAN_CMD_SET_CHNL_RECV_MERGE fail\n"); + } + } + + if (!VCI_StartCAN(s_DevType, s_DevIdx, i)) // 启动通道 + { + printf("StartCAN(%d) fail\n", i); + return false; + } + printf("StartCAN(%d) success\n", i); + + s_rx_ctx[i].dev_type = s_DevType; + s_rx_ctx[i].dev_idx = s_DevIdx; + s_rx_ctx[i].chn_idx = i; + s_rx_ctx[i].total = 0; + s_rx_ctx[i].stop = 0; + + if (IsMerge && i == 0) + { + pthread_create(&s_rx_threads[i], NULL, rx_merge_thread, &s_rx_ctx[i]); // 合并接收线程 + } + else + { + pthread_create(&s_rx_threads[i], NULL, rx_thread, &s_rx_ctx[i]); // 创建接收线程 + } + } + return true; +} + +/* + * 停止 CAN 接收并关闭设备。 + */ +void can_stop() +{ + if (!s_can_opened) + return; + + g_can_running = false; + + for (int i = 0; i < MAX_CHANNELS; i++) + { + s_rx_ctx[i].stop = 1; + pthread_join(s_rx_threads[i], NULL); // 等待接收线程退出 + + if (!VCI_ResetCAN(s_DevType, s_DevIdx, i)) // 复位通道 + printf("ResetCAN(%d) fail\n", i); + else + printf("ResetCAN(%d) success!\n", i); + } + + if (!VCI_CloseDevice(s_DevType, s_DevIdx)) + printf("CloseDevice fail\n"); + else + printf("CloseDevice success\n"); + + s_can_opened = false; +} + +uint32_t can_send(int chn_idx, uint32_t id, const uint8_t data[8]) +{ + if (!s_can_opened) + return 0; + + ZCAN_20_MSG msg; + construct_can_frame(&msg, id, chn_idx, 0); + for (int i = 0; i < 8; ++i) + msg.dat[i] = data[i]; + + return VCI_Transmit(s_DevType, s_DevIdx, chn_idx, &msg, 1); +} + +uint32_t canfd_send(int chn_idx, uint32_t id, const uint8_t data[64]) +{ + if (!s_can_opened) + return 0; + + ZCAN_FD_MSG msg; + construct_canfd_frame(&msg, id, chn_idx, 0); + for (int i = 0; i < 64; ++i) + msg.dat[i] = data[i]; + + return VCI_TransmitFD(s_DevType, s_DevIdx, chn_idx, &msg, 1); +} + +} // namespace CanComm diff --git a/can_comm.h b/can_comm.h new file mode 100644 index 0000000..42415e3 --- /dev/null +++ b/can_comm.h @@ -0,0 +1,54 @@ +// can_comm.h - USBCANFD CAN/CANFD 收发封装 +// +// 整合自 vision-anti-pinch_bak/main.cpp 的 CAN 模块, 提供: +// can_start() - 打开设备, 初始化并启动双通道接收线程 +// can_stop() - 停止接收并关闭设备 +// can_send() - 发送一帧 CAN (8 字节) +// canfd_send() - 发送一帧 CANFD (64 字节) +// +// 依赖 zcan.h (include/) 和 libusbcanfd.so + libusb-1.0.so (lib/) + +#pragma once + +#include +#include + +namespace CanComm { + +// 初始化并启动 CAN 收发(双通道, 仲裁域 500kbps / 数据域 2000kbps)。 +// 返回 true 成功; false 失败(主流程可继续仅取图)。 +bool can_start(); + +// 停止 CAN 接收线程并关闭设备。 +void can_stop(); + +// 注册报文解析回调。收到每帧 CAN/CANFD 数据帧时调用。 +// time: 时间戳(us) +// msg_id: CAN-ID +// data: 数据指针(回调内有效, 不要保存) +// len: 数据长度 +// 传 nullptr 取消注册。 +using MsgCallback = void (*)(int64_t time, uint32_t msg_id, const uint8_t *data, uint32_t len); +void set_msg_callback(MsgCallback cb); + +// 发送一帧 CAN2.0 数据帧(8 字节)。 +// chn_idx: 通道号 0/1 +// id: CAN-ID +// data: 8 字节数据(不足补 0, 超出截断) +// 返回实际发送帧数(0 失败)。 +uint32_t can_send(int chn_idx, uint32_t id, const uint8_t data[8]); + +// 发送一帧 CANFD 数据帧(64 字节, 扩展帧+加速)。 +// chn_idx: 通道号 0/1 +// id: CAN-ID +// data: 64 字节数据(不足补 0, 超出截断) +// 返回实际发送帧数(0 失败)。 +uint32_t canfd_send(int chn_idx, uint32_t id, const uint8_t data[64]); + +// 是否已打开并运行 +bool is_open(); + +// 累计接收帧数(线程安全) +int receive_count(); + +} // namespace CanComm diff --git a/examples/read_latest.py b/examples/read_latest.py new file mode 100644 index 0000000..2a3f0bd --- /dev/null +++ b/examples/read_latest.py @@ -0,0 +1,36 @@ +#!/usr/bin/env python3 +"""Same-process downstream example. Run on Orin, which owns the camera/GPU.""" +import sys +from pathlib import Path +sys.path.insert(0, str(Path(__file__).resolve().parents[1])) +from run_camera import CameraPipeline, parse_args + + +def main(): + # No visualization or disk output by default. CLI flags can enable either. + pipeline = CameraPipeline(parse_args()) + last_frame = -1 + try: + pipeline.start() + while True: + # Reading the variable is enough for a consumer with its own clock: + # result = pipeline.latest_result # equivalent to pipeline.latest() + # This sample waits to avoid busy polling; still only the latest + # result is returned, without queuing/acknowledging individual frames. + result = pipeline.wait_for_result(last_frame, timeout=1.0) + if result is not None: + last_frame = result['frame_index'] + print(last_frame, result['risk']['risk_type'], + result['risk']['risk_score'], flush=True) + # Downstream code uses result['door_risks'], ['skeletons'], ['can']. + elif not pipeline.running(): + break + except KeyboardInterrupt: + pass + finally: + pipeline.request_stop() + pipeline.join() # Also propagates worker errors to the caller. + + +if __name__ == '__main__': + main() diff --git a/include/zcan.h b/include/zcan.h new file mode 100644 index 0000000..9f9d935 --- /dev/null +++ b/include/zcan.h @@ -0,0 +1,443 @@ +#pragma once + +#ifdef __cplusplus +#define EXTERN_C extern "C" +#else +#define EXTERN_C +#endif + +#define ZCAN_API +#define INVALID_HANDLE_VALUE ((HANDLE)(-1)) + +#define CMD_CAN_FILTER 0x14 // 滤波 +#define CMD_CAN_TTX 0x16 // 定时发送 +#define CMD_CAN_TTX_CTL 0x17 // 使能定时发送 +#define CMD_CAN_TRES 0x18 // CAN终端电阻 + +#define ZCAN_CMD_SET_CHNL_RECV_MERGE 0x32 // 设置合并接收 0:不合并接收;1:合并接收 +#define ZCAN_CMD_GET_CHNL_RECV_MERGE 0x33 // 获取是否开启合并接收 0:不合并接收;1:合并接收 + +#define CMD_SET_SN 0x42 // 获取SN号 +#define CMD_GET_SN 0x43 // 设置SN号 +#define CMD_CAN_TX_TIMEOUT 0x44 // 发送超时 + +#define ZCAN_CMD_GET_SEND_QUEUE_SIZE 0x100 // 获取队列大小，uint32_t +#define ZCAN_CMD_GET_SEND_QUEUE_SPACE 0x101 // 获取队列剩余空间, uint32_t +#define ZCAN_CMD_SET_SEND_QUEUE_CLR 0x102 // 清空发送队列,1：清空 +#define ZCAN_CMD_SET_SEND_QUEUE_EN 0x103 // 开启发送队列,1：使能 + +typedef unsigned char U8; +typedef unsigned short U16; +typedef unsigned int U32; +typedef unsigned long long U64; + +enum eZCANDataDEF +{ + //数据类型 + ZCAN_DT_ZCAN_CAN_DATA = 1, // CAN/CANFD数据 + ZCAN_DT_ZCAN_CANFD_DATA = 2, // CAN/CANFD数据 + ZCAN_DT_ZCAN_ERROR_DATA = 3, // CAN/CANFD错误数据 + ZCAN_DT_ZCAN_LIN_DATA = 4, // LIN数据 + ZCAN_DT_ZCAN_LIN_ERROR_DATA = 5, // LIN错误数据 +}; +#pragma pack(push, 1) + +typedef enum { + ZCAN_TX_NORM = 0, /**< normal transmission */ + ZCAN_TX_ONCE = 1, /**< single-shot transmission */ + ZCAN_SR_NORM = 2, /**< self reception */ + ZCAN_SR_ONCE = 3, /**< single-shot transmission & self reception */ +} ZCAN_TX_MODE; + +/** CAN filter configuration */ +typedef struct { + U8 type; /**< 0-std_frame, 1-ext_frame */ + U8 pad[3]; + U32 sid; /**< start-id */ + U32 eid; /**< end-id */ +} ZCAN_FILTER; + +typedef struct { + U32 size; + ZCAN_FILTER table[64]; +} ZCAN_FILTER_TABLE; + +/** controller initialization */ +typedef struct { + U32 clk; // 时钟频率 + U32 mode; /**< bit0-normal/listen_only, bit1-ISO/BOSCH */ + struct { + U8 tseg1; + U8 tseg2; + U8 sjw; + U8 smp; + U16 brp; + } aset; + struct { + U8 tseg1; + U8 tseg2; + U8 sjw; + U8 smp; + U16 brp; + } dset; +} ZCAN_INIT; + +/** CAN message info */ +typedef struct { + U32 txm : 4; /**< TX-mode, @see ZCAN_TX_MODE */ + U32 fmt : 4; /**< 0-CAN2.0, 1-CANFD */ + U32 sdf : 1; /**< 0-data_frame, 1-remote_frame */ + U32 sef : 1; /**< 0-std_frame, 1-ext_frame */ + U32 err : 1; /**< error flag */ + U32 brs : 1; /**< bit-rate switch */ + U32 est : 1; /**< error state */ + U32 tx : 1; /**< received valid, tx frame */ + U32 echo : 1; /**< tx valid, echo frame */ + U32 qsend_100us:1; /**< queue send delay unit, 1-100us, 0-ms */ + U32 qsend:1; /**< send valid, queue send frame */ + U32 pad : 15; +} ZCAN_MSG_INF; + +/** CAN message header */ +typedef struct { + U32 ts; /**< timestamp */ + U32 id; /**< CAN-ID */ + ZCAN_MSG_INF inf; /**< @see ZCAN_MSG_INF */ + U16 pad; /**< send queue delay */ + U8 chn; /**< channel */ + U8 len; /**< data length */ +} ZCAN_MSG_HDR; + +/** CAN2.0-frame */ +typedef struct { + ZCAN_MSG_HDR hdr; + U8 dat[8]; +} ZCAN_20_MSG; + +/** CANFD-frame */ +typedef struct { + ZCAN_MSG_HDR hdr; + U8 dat[64]; +} ZCAN_FD_MSG; + +// 错误帧 +typedef struct { + ZCAN_MSG_HDR hdr; + U8 dat[8]; +} ZCAN_ERR_MSG; + +// 定时发送结构体 +typedef struct { + U32 interval; // 定时发送周期，单位 100us + U16 repeat; // 发送次数，0等于循环发 + U8 index; // 定时发送列表的帧索引号，也就是第几条定时发送报文 + U8 flags; // 0-此帧禁用定时发送，1-此帧使能定时发送 + ZCAN_FD_MSG msg; // CANFD帧结构体 +} ZCAN_TTX; + +typedef struct { + U32 size; // 实际生效的数组的长度（用来控制实际生效多少条定时发送帧结构体） + ZCAN_TTX table[8]; // 最大设置8条 +} ZCAN_TTX_CFG; + +// 设备信息 +typedef struct { + U16 hwv; /**< hardware version */ + U16 fwv; /**< firmware version */ + U16 drv; /**< driver version */ + U16 api; /**< API version */ + U16 irq; /**< IRQ */ + U8 chn; /**< channels */ + U8 sn[20]; /**< serial number */ + U8 id[40]; /**< card id */ + U16 pad[4]; +} ZCAN_DEV_INF; + +/** controller status */ +typedef struct { + U8 IR; /**< not used(for backward compatibility) */ + U8 MOD; /**< not used */ + U8 SR; /**< not used */ + U8 ALC; /**< not used */ + U8 ECC; /**< not used */ + U8 EWL; /**< not used */ + U8 RXE; /**< RX errors */ + U8 TXE; /**< TX errors */ + U32 PAD; +} ZCAN_STAT; + +//lin begin +typedef struct _VCI_LIN_INIT_CONFIG +{ + U8 linMode; // 是否作为主机，0-从机，1-主机 + U8 chkSumMode; // 校验方式，1-经典校验 2-增强校验 3-自动(对应eZLINChkSumMode的模式) + U16 reserved; // 保留 + U32 linBaud; // 波特率，取值1000~20000 +}ZCAN_LIN_INIT_CONFIG, *PZCAN_LIN_INIT_CONFIG; + +typedef struct tagZCANLINData +{ + union{ + struct { + U8 ID : 6; // 帧ID + U8 Parity : 2; // 帧ID校验 + }unionVal; + U8 rawVal; // 受保护的ID原始值 + } PID; // 受保护的ID + struct + { + U64 timeStamp; // 时间戳，单位微秒(us) + U8 dataLen; // 数据长度 + U8 dir; // 传输方向，0-接收 1-发送 + U8 chkSum; // 数据校验, 部分设备不支持校验数据的获取 + U8 reserved[13]; // 保留 + U8 data[8]; // 数据 + }RxData; // 仅接收数据时有效 + U8 reserved[7]; // 保留 +}ZCANLINData; + +typedef struct tagZCANLINErrData +{ + U64 timeStamp; // 时间戳, 单位微秒(us) + union { + struct { + U8 ID : 6; // 帧ID + U8 Parity : 2; // 帧ID校验 + }unionVal; + U8 rawVal; // 受保护的ID原始值 + } PID; // 受保护的ID + U8 dataLen; + U8 data[8]; + union + { + struct + { + U16 errStage : 4; // 错误阶段 + U16 errReason : 4; // 错误原因 + U16 reserved : 8; // 保留 + }; + U16 unionErrData; + }errData; + U8 dir; // 传输方向 + U8 chkSum; // 数据校验, 部分设备不支持校验数据的获取 + U8 reserved[10]; // 保留 +}ZCANLINErrData; + +typedef struct _VCI_LIN_MSG{ + U8 chnl; // 数据通道 + U8 dataType; // 数据类型, 0-LIN数据 1-LIN错误数据 + union + { + ZCANLINData zcanLINData; // LIN数据 + ZCANLINErrData zcanLINErrData; // LIN错误数据 + U8 raw[46]; // RAW数据 + } data; // 实际数据, 联合体，有效成员根据 dataType 字段而定 +}ZCAN_LIN_MSG, *PZCAN_LIN_MSG; + +enum eZLINChkSumMode +{ + DEFAULT = 0, // 默认，启动时配置 + CLASSIC_CHKSUM, // 经典校验 + ENHANCE_CHKSUM, // 增强校验 + AUTOMATIC, // 自动，设备自动识别校验方式（仅ZCAN_SetLINSubscribe时有效） +}; +enum +{ + MAX_LIN_ID_COUNT = 63, + MIN_LIN_DLC = 1, + MAX_LIN_DLC = 8, + AUTO_LIN_DLC = 255, +}; +typedef struct _VCI_LIN_PUBLISH_CFG +{ + U8 ID; // 受保护的ID（ID取值范围为0-63） + U8 dataLen; // dataLen范围为1-8 + U8 data[8]; + U8 chkSumMode; // 校验方式，0-默认，启动时配置 1-经典校验 2-增强校验(对应eZLINChkSumMode的模式) + U8 reserved[5]; // 保留 +}ZCAN_LIN_PUBLISH_CFG, *PZCAN_LIN_PUBLISH_CFG; + +typedef struct _VCI_LIN_SUBSCIBE_CFG +{ + U8 ID; // 受保护的ID（ID取值范围为0-63） + U8 dataLen; // dataLen范围为1-8 当为255（0xff）则表示设备自动识别报文长度 + U8 chkSumMode; // 校验方式，0-默认，启动时配置 1-经典校验 2-增强校验 3-自动(对应eZLINChkSumMode的模式) + U8 reserved[5]; // 保留 +}ZCAN_LIN_SUBSCIBE_CFG, *PZCAN_LIN_SUBSCIBE_CFG; + +// 合并接收数据数据结构，支持CAN/CANFD/LIN/错误等不同类型数据 +typedef struct tagZCANDataObj +{ + U8 dataType; // 数据类型, 参考eZCANDataDEF中 数据类型 部分定义 + U8 chnl; // 数据通道 + union{ + struct{ + U16 reserved : 16; // 保留 + }unionVal; + U16 rawVal; + }flag; // 标志信息, 暂未使用 + U8 extraData[4]; // 额外数据, 暂未使用 + union + { + ZCAN_20_MSG zcanCANData; // CAN数据 + ZCAN_FD_MSG zcanCANFDData; // CANFD数据 + ZCAN_ERR_MSG zcanErrData; // 错误数据 + ZCANLINData zcanLINData; // LIN数据 + ZCANLINErrData zcanLINErrData; // LIN错误数据 + U8 raw[92]; // RAW数据 + } data; // 实际数据, 联合体，有效成员根据 dataType 字段而定 +}ZCANDataObj; + + +// UDS传输协议版本 +typedef U8 ZCAN_UDS_TRANS_VER; +#define ZCAN_UDS_TRANS_VER_0 0 // ISO15765-2(2004版本) +#define ZCAN_UDS_TRANS_VER_1 1 // ISO15765-2(2016版本) + +// 帧类型 +typedef U8 ZCAN_UDS_FRAME_TYPE; +#define ZCAN_UDS_FRAME_CAN 0 // CAN帧 +#define ZCAN_UDS_FRAME_CANFD 1 // CANFD帧 +#define ZCAN_UDS_FRAME_CANFD_BRS 2 // CANFD加速帧 + +// UDS填充模式 +#define FILL_MODE_SHORT 0 // 小于8字节填充至8字节，大于8字节时按DLC就近填充 +#define FILL_MODE_NONE 1 // 不填充 +#define FILL_MODE_MAX 2 // 填充至最大数据长度 (不建议) + +// CAN UDS请求数据 +typedef struct _ZCAN_UDS_REQUEST +{ + U32 req_id; // 请求事务ID，范围0~65535，本次请求的唯一标识 + U8 channel; // 设备通道索引 0~255 + ZCAN_UDS_FRAME_TYPE frame_type; // 帧类型 + U8 reserved0[2]; // 保留 + U32 src_addr; // 请求地址 + U32 dst_addr; // 响应地址 + U8 suppress_response; // 1:抑制响应 + U8 sid; // 请求服务id + U8 reserved1[6]; // 保留 + struct { + U32 timeout; // 响应超时时间(ms)。因PC定时器误差，建议设置不小于200ms + U32 enhanced_timeout; // 收到消极响应错误码为0x78后的超时时间(ms)。因PC定时器误差，建议设置不小于200ms + U8 check_any_negative_response:1; // 接收到非本次请求服务的消极响应时是否需要判定为响应错误 + U8 wait_if_suppress_response:1; // 抑制响应时是否需要等待消极响应，等待时长为响应超时时间 + U8 flag:6; // 保留 + U8 reserved0[7]; // 保留 + } session_param; // 会话层参数 + struct { + ZCAN_UDS_TRANS_VER version; // 传输协议版本, VERSION_0, VERSION_1 + U8 max_data_len; // 单帧最大数据长度, can:8, canfd:64 + U8 local_st_min; // 本程序发送流控时用，连续帧之间的最小间隔, 0x00-0x7F(0ms~127ms), 0xF1-0xF9(100us~900us) + U8 block_size; // 流控帧的块大小 + U8 fill_byte; // 无效字节的填充数据 + U8 ext_frame; // 0:标准帧 1:扩展帧 + U8 is_modify_ecu_st_min; // 是否忽略ECU返回流控的STmin，强制使用本程序设置的 remote_st_min + U8 remote_st_min; // 发送多帧时用, is_ignore_ecu_st_min = 1 时有效, 0x00-0x7F(0ms~127ms), 0xF1-0xF9(100us~900us) + U32 fc_timeout; // 接收流控超时时间(ms), 如发送首帧后需要等待回应流控帧 + U8 fill_mode; // 数据长度填充模式, FILL_MODE_SHORT, FILL_MODE_NONE, FILL_MODE_MAX + U8 reserved0[3]; // 保留 + } trans_param; // 传输层参数 + U8 *data; // 数据数组(不包含SID) + U32 data_len; // 数据数组的长度 + U32 reserved2; // 保留 +} ZCAN_UDS_REQUEST; + +// UDS错误码 +typedef U8 ZCAN_UDS_ERROR; +#define ZCAN_UDS_ERROR_OK 0 // 没错误 +#define ZCAN_UDS_ERROR_TIMEOUT 1 // 响应超时 +#define ZCAN_UDS_ERROR_TRANSPORT 2 // 发送数据失败 +#define ZCAN_UDS_ERROR_CANCEL 3 // 取消请求 +#define ZCAN_UDS_ERROR_SUPPRESS_RESPONSE 4 // 抑制响应 +#define ZCAN_UDS_ERROR_BUSY 5 // 忙碌中 +#define ZCAN_UDS_ERROR_REQ_PARAM 6 // 请求参数错误 +#define ZCAN_UDS_ERROR_OTHTER 100 + +typedef U8 ZCAN_UDS_RESPONSE_TYPE; +#define ZCAN_UDS_RT_NEGATIVE 0 // 消极响应 +#define ZCAN_UDS_RT_POSITIVE 1 // 积极响应 + +// UDS响应数据 +typedef struct _ZCAN_UDS_RESPONSE +{ + ZCAN_UDS_ERROR status; // 响应状态 + U8 reserved[6]; // 保留 + ZCAN_UDS_RESPONSE_TYPE type; // 响应类型 + union { + struct { + U8 sid; // 响应服务id + U32 data_len; // 数据长度(不包含SID), 数据存放在接口传入的dataBuf中 + } positive; + struct { + U8 neg_code; // 固定为0x7F + U8 sid; // 请求服务id + U8 error_code; // 错误码 + } negative; + U8 raw[8]; + }; +} ZCAN_UDS_RESPONSE; + +// UDS控制类型 +typedef U32 ZCAN_UDS_CTRL_CODE; +#define ZCAN_UDS_CTRL_STOP_REQ 0 // 停止UDS请求 + +// UDS控制请求 +typedef struct _ZCAN_UDS_CTRL_REQ +{ + U32 reqID; // 请求事务ID，指明要操作哪一条请求 + ZCAN_UDS_CTRL_CODE cmd; // 控制类型 + U8 reserved[8]; // 保留 +} ZCAN_UDS_CTRL_REQ; + +// UDS控制结果 +typedef U32 ZCAN_UDS_CTRL_RESULT; +#define ZCAN_UDS_CTRL_RESULT_OK 0 // 成功 +#define ZCAN_UDS_CTRL_RESULT_ERR 1 // 失败 + +// UDS控制响应数据 +typedef struct _ZCAN_UDS_CTRL_RESP +{ + ZCAN_UDS_CTRL_RESULT result; // 操作结果 + U8 reserved[8]; // 保留 +} ZCAN_UDS_CTRL_RESP; + +#pragma pack(pop) + +// +EXTERN_C U32 ZCAN_API VCI_OpenDevice(U32 Type, U32 Card, U32 Reserved); +EXTERN_C U32 ZCAN_API VCI_CloseDevice(U32 Type, U32 Card); +EXTERN_C U32 ZCAN_API VCI_InitCAN(U32 Type, U32 Card, U32 Port, ZCAN_INIT *pInit); +EXTERN_C U32 ZCAN_API VCI_ReadBoardInfo(U32 Type, U32 Card, ZCAN_DEV_INF *pInfo); +EXTERN_C U32 ZCAN_API VCI_ReadErrInfo(U32 Type, U32 Card, U32 Port, ZCAN_ERR_MSG *pErr); +EXTERN_C U32 ZCAN_API VCI_ReadCANStatus(U32 Type, U32 Card, U32 Port, ZCAN_STAT *pStat); +EXTERN_C U32 ZCAN_API VCI_GetReference(U32 Type, U32 Card, U32 Port, U32 Ref, void *pData); +EXTERN_C U32 ZCAN_API VCI_SetReference(U32 Type, U32 Card, U32 Port, U32 Ref, void *pData); +EXTERN_C U32 ZCAN_API VCI_GetReceiveNum(U32 Type, U32 Card, U32 Port); +EXTERN_C U32 ZCAN_API VCI_ClearBuffer(U32 Type, U32 Card, U32 Port); +EXTERN_C U32 ZCAN_API VCI_StartCAN(U32 Type, U32 Card, U32 Port); +EXTERN_C U32 ZCAN_API VCI_ResetCAN(U32 Type, U32 Card, U32 Port); +EXTERN_C U32 ZCAN_API VCI_Transmit(U32 Type, U32 Card, U32 Port, ZCAN_20_MSG *pData, U32 Count); +EXTERN_C U32 ZCAN_API VCI_TransmitFD(U32 Type, U32 Card, U32 Port, ZCAN_FD_MSG *pData, U32 Count); +EXTERN_C U32 ZCAN_API VCI_Receive(U32 Type, U32 Card, U32 Port, ZCAN_20_MSG *pData, U32 Count, U32 Time); +EXTERN_C U32 ZCAN_API VCI_ReceiveFD(U32 Type, U32 Card, U32 Port, ZCAN_FD_MSG *pData, U32 Count, U32 Time); +EXTERN_C U32 ZCAN_API VCI_Debug(U32 Debug); + +// LIN +EXTERN_C U32 VCI_InitLIN(U32 Type, U32 Card, U32 LinChn, PZCAN_LIN_INIT_CONFIG pLINInitConfig); +EXTERN_C U32 VCI_StartLIN(U32 Type, U32 Card, U32 LinChn); +EXTERN_C U32 VCI_ResetLIN(U32 Type, U32 Card, U32 LinChn); +EXTERN_C U32 VCI_TransmitLIN(U32 Type, U32 Card, U32 LinChn, PZCAN_LIN_MSG pSend, U32 Len); +EXTERN_C U32 VCI_GetLINReceiveNum(U32 Type, U32 Card, U32 LinChn); +EXTERN_C U32 VCI_ClearLINBuffer(U32 Type, U32 Card, U32 LinChn); +EXTERN_C U32 VCI_ReceiveLIN(U32 Type, U32 Card, U32 LinChn, PZCAN_LIN_MSG pReceive, U32 Len,int WaitTime); +EXTERN_C U32 VCI_SetLINSubscribe(U32 Type, U32 Card, U32 LinChn, PZCAN_LIN_SUBSCIBE_CFG pSend, U32 nSubscribeCount); +EXTERN_C U32 VCI_SetLINPublish(U32 Type, U32 Card, U32 LinChn, PZCAN_LIN_PUBLISH_CFG pSend, U32 nPublishCount); + +// 合并接收 +EXTERN_C U32 VCI_TransmitData(unsigned Type, unsigned Card, unsigned Port, ZCANDataObj *pData, unsigned Count); +EXTERN_C U32 VCI_ReceiveData(unsigned Type, unsigned Card, unsigned Port, ZCANDataObj *pData, unsigned Count, unsigned Time); + +// UDS +EXTERN_C U32 VCI_UDS_Request(unsigned Type, unsigned Card, const ZCAN_UDS_REQUEST *req, ZCAN_UDS_RESPONSE *resp, U8 *dataBuf, U32 dataBufSize); +EXTERN_C U32 VCI_UDS_Control(unsigned Type, unsigned Card, const ZCAN_UDS_CTRL_REQ *ctrl, ZCAN_UDS_CTRL_RESP *resp); diff --git a/lib/libusb-1.0.so b/lib/libusb-1.0.so new file mode 100644 index 0000000..0824ee0 Binary files /dev/null and b/lib/libusb-1.0.so differ diff --git a/lib/libusb-1.0.so.0 b/lib/libusb-1.0.so.0 new file mode 120000 index 0000000..4970d7f --- /dev/null +++ b/lib/libusb-1.0.so.0 @@ -0,0 +1 @@ +libusb-1.0.so \ No newline at end of file diff --git a/lib/libusbcanfd.so b/lib/libusbcanfd.so new file mode 100644 index 0000000..72b9370 Binary files /dev/null and b/lib/libusbcanfd.so differ diff --git a/lib/libusbcanfd.so.1.0.12 b/lib/libusbcanfd.so.1.0.12 new file mode 120000 index 0000000..2ee2cf4 --- /dev/null +++ b/lib/libusbcanfd.so.1.0.12 @@ -0,0 +1 @@ +libusbcanfd.so \ No newline at end of file diff --git a/main.cpp b/main.cpp new file mode 100644 index 0000000..04921ea --- /dev/null +++ b/main.cpp @@ -0,0 +1,344 @@ +// vision-anti-pinch - V4L2 MMAP 取图 + USBCANFD CAN 收发 (OpenCV 版本) +// +// 使用 OpenCV 进行 UYVY/YUYV -> BGR 转换，避免手动转换公式问题 +// CAN 收发默认开启, --no-can 关闭 +// +// 用法: +// ./vision-anti-pinch -d /dev/video5 -n 10 -s 2 -t jpg -o ./out + +#include "v4l2_capture.h" +#include "can_comm.h" +#include "send.hpp" + +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +std::atomic g_Feedback_Of_The_Actual_Position_Of_The_Left_Front_Door_391_S; +std::atomic g_Feedback_Of_The_Actual_Position_Of_The_Right_Front_Door_391_S; +std::atomic g_Feedback_Of_The_Actual_Position_Of_The_Left_Rear_Door_391_S; +std::atomic g_Feedback_Of_The_Actual_Position_Of_The_Right_Rear_Door_391_S; + +// Sensing GMSL2 驱动自定义控制 ID +#define CID_SENSOR_MODE 0x009a2008 +#define CID_TRIG_PIN 0x009a2084 +#define CID_TRIG_MODE 0x009a2085 +#define CID_GAIN 0x009a2009 +#define CID_EXPOSURE 0x009a200a + +static volatile sig_atomic_t g_stop = 0; +static void onSignal(int) { g_stop = 1; } + +static void printUsage(const char *prog) { + printf( + "Usage: %s [options]\n" + "V4L2 MMAP capture demo for Sensing GMSL2 camera on Jetson Orin.\n" + "\n" + "Capture options:\n" + " -d V4L2 device node (default: /dev/video5)\n" + " -W frame width (default: 1920)\n" + " -H frame height (default: 1536)\n" + " -f target frame rate (default: 30)\n" + " -F pixel format: UYVY | YUYV (default: UYVY)\n" + " -m sensor_mode control value\n" + " -p trig_pin control value (hex/dec)\n" + " --gain sensor gain\n" + " --exposure sensor exposure\n" + "\n" + "Output options:\n" + " -n frames to capture, 0 = run forever\n" + " -s frames to save\n" + " -o output directory (default: ./frames)\n" + " -t file type: jpg | png | bmp | raw (default: jpg)\n" + " --quality <1..100> JPEG quality (default: 95)\n" + " -v verbose output\n" + " -h print this help\n" + "\n" + "CAN options:\n" + " --no-can disable USBCANFD CAN/CANFD receive (enabled by default)\n" + "\n" + "Example:\n" + " %s -d /dev/video5 -n 100 -s 5 -t jpg -o ./out\n" + " %s -d /dev/video5 -n 0 -v\n", + prog, prog, prog); +} + +struct Args { + std::string device = "/dev/video5"; + uint32_t width = 1920; + uint32_t height = 1536; // SG3S-ISX031C-GMSL2F 默认分辨率 + uint32_t fps = 30; + std::string format = "UYVY"; + long sensor_mode = -1; + long trig_pin = -1; + long trig_mode = -1; + long gain = -1; + long exposure = -1; + long count = 0; + long save = 0; + std::string out_dir = "./frames"; + std::string save_type = "jpg"; + int quality = 95; + bool verbose = false; + bool can_enabled = true; // 默认启用 CAN 收发, --no-can 关闭 +}; + +// 手动解析参数 +static bool parseArgs(int argc, char **argv, Args &a) { + auto need_val = [&](int &i) -> std::string { + if (i + 1 >= argc) { + printf("missing value for %s\n", argv[i]); + std::exit(1); + } + return argv[++i]; + }; + + for (int i = 1; i < argc; ++i) { + std::string arg = argv[i]; + if (arg == "-d") a.device = need_val(i); + else if (arg == "-W") a.width = (uint32_t)std::atoi(need_val(i).c_str()); + else if (arg == "-H") a.height = (uint32_t)std::atoi(need_val(i).c_str()); + else if (arg == "-f") a.fps = (uint32_t)std::atoi(need_val(i).c_str()); + else if (arg == "-F") a.format = need_val(i); + else if (arg == "-m") a.sensor_mode = std::strtol(need_val(i).c_str(), nullptr, 0); + else if (arg == "-p") a.trig_pin = std::strtol(need_val(i).c_str(), nullptr, 0); + else if (arg == "-M") a.trig_mode = std::strtol(need_val(i).c_str(), nullptr, 0); + else if (arg == "-n") a.count = std::strtol(need_val(i).c_str(), nullptr, 0); + else if (arg == "-s") a.save = std::strtol(need_val(i).c_str(), nullptr, 0); + else if (arg == "-o") a.out_dir = need_val(i); + else if (arg == "-t") a.save_type = need_val(i); + else if (arg == "-v") a.verbose = true; + else if (arg == "-h") { printUsage(argv[0]); std::exit(0); } + else if (arg == "--gain") a.gain = std::strtol(need_val(i).c_str(), nullptr, 0); + else if (arg == "--exposure") a.exposure = std::strtol(need_val(i).c_str(), nullptr, 0); + else if (arg == "--quality") a.quality = std::atoi(need_val(i).c_str()); + else if (arg == "--no-can") a.can_enabled = false; + else { + printf("unknown option: %s\n", arg.c_str()); + printUsage(argv[0]); + return false; + } + } + return true; +} + +// 使用 OpenCV 将 V4L2 帧 (UYVY/YUYV) 转为 BGR Mat +static cv::Mat frameToBgr(const V4l2Capture::Frame &frame) { + // OpenCV 原生支持 UYVY/YUYV -> BGR 转换 + // UYVY = COLOR_YUV2BGR_UYVY, YUYV = COLOR_YUV2BGR_YUYV + int code = (frame.pixelformat == V4L2_PIX_FMT_UYVY) + ? cv::COLOR_YUV2BGR_UYVY + : cv::COLOR_YUV2BGR_YUYV; + cv::Mat yuv((int)frame.height, (int)frame.width, CV_8UC2, (void *)frame.data); + cv::Mat bgr; + cv::cvtColor(yuv, bgr, code); + return bgr; +} + +/* ===================================================================== + * CAN 报文接收解析 (由 can_comm 接收线程回调) + * ===================================================================== */ + +// 512 位报文位域, deal_msg 中按 msg_id 拷贝后可按位解析信号 +static std::bitset<512> g_bitMsg; + +static uint64_t get_bit_value(int star_bit, int bit_len) +{ + uint64_t signal_value = 0; + for (int i = 0; i < bit_len; ++i) + { + signal_value |= g_bitMsg[star_bit + i] << i; + } + return signal_value; +} + +// 收到报文的回调点: 目前仅解析 0x391, 后续信号解析在此扩展 +static void deal_msg(int64_t time, uint32_t msg_id, const uint8_t *data, uint32_t len) +{ + if (msg_id == 0x391 && len >= 4 && len <= 64) + { + memcpy(&g_bitMsg, data, len); + g_Feedback_Of_The_Actual_Position_Of_The_Left_Front_Door_391_S = data[0]; + g_Feedback_Of_The_Actual_Position_Of_The_Right_Front_Door_391_S = data[1]; + g_Feedback_Of_The_Actual_Position_Of_The_Left_Rear_Door_391_S = data[2]; + g_Feedback_Of_The_Actual_Position_Of_The_Right_Rear_Door_391_S = data[3]; + // Rear_BCM_0x391,0x391,64,1.0-1.7,8,Feedback_Of_The_Actual_Position_Of_The_Left_Front_Door_391_S + // Rear_BCM_0x391,0x391,64,2.0-2.7,8,Feedback_Of_The_Actual_Position_Of_The_Right_Front_Door_391_S + // Rear_BCM_0x391,0x391,64,3.0-3.7,8,Feedback_Of_The_Actual_Position_Of_The_Left_Rear_Door_391_S + // Rear_BCM_0x391,0x391,64,4.0-4.7,8,Feedback_Of_The_Actual_Position_Of_The_Right_Rear_Door_391_S + } +} + +int main(int argc, char **argv) { + Args args; + if (!parseArgs(argc, argv, args)) { + return 1; + } + + std::signal(SIGINT, onSignal); + std::signal(SIGTERM, onSignal); + + uint32_t pixfmt = (args.format == "UYVY") ? V4L2_PIX_FMT_UYVY : V4L2_PIX_FMT_YUYV; + + // 启动 CAN 收发(失败不致命, 可继续仅取图) + if (args.can_enabled) { + CanComm::set_msg_callback(deal_msg); // 注册报文解析回调 + if (!CanComm::can_start()) { + printf("CAN start failed, continue with camera only.\n"); + args.can_enabled = false; + } + } + + V4l2Capture cap; + + if (!cap.open(args.device)) { + return 1; + } + + // 设置相机控制 + if (args.sensor_mode >= 0) cap.setControl(CID_SENSOR_MODE, (int32_t)args.sensor_mode); + if (args.trig_pin >= 0) cap.setControl(CID_TRIG_PIN, (int32_t)args.trig_pin); + if (args.trig_mode >= 0) cap.setControl(CID_TRIG_MODE, (int32_t)args.trig_mode); + if (args.gain >= 0) cap.setControl(CID_GAIN, (int32_t)args.gain); + if (args.exposure >= 0) cap.setControl(CID_EXPOSURE, (int32_t)args.exposure); + + if (!cap.setFormat(args.width, args.height, pixfmt)) { + return 1; + } + cap.setFps(args.fps); + + if (!cap.initMmap(4)) { + return 1; + } + if (!cap.startStreaming()) { + return 1; + } + + printf("=========================================================\n"); + printf("Capturing %s %ux%u @ %u fps from %s\n", + args.format.c_str(), cap.width(), cap.height(), args.fps, args.device.c_str()); + printf("Press Ctrl+C to stop.\n"); + printf("=========================================================\n"); + + // 创建输出目录 + if (args.save > 0) { + std::string mkcmd = "mkdir -p " + args.out_dir; + std::system(mkcmd.c_str()); + } + + using clock = std::chrono::steady_clock; + auto t_start = clock::now(); + auto t_last_report = t_start; + uint64_t frame_count = 0; + uint64_t save_count = 0; + uint32_t timeout_consecutive = 0; + + while (!g_stop) { + if (args.count > 0 && frame_count >= (uint64_t)args.count) { + break; + } + + V4l2Capture::Frame frame; + if (!cap.getFrame(frame, 2000)) { + timeout_consecutive++; + if (timeout_consecutive > 5) { + printf("getFrame: timeout/empty for %u consecutive tries, aborting\n", + timeout_consecutive); + break; + } + continue; + } + timeout_consecutive = 0; + frame_count++; + + // 保存帧 + if (save_count < (uint64_t)args.save) { + bool ok = false; + char fname[512]; + + // 使用 OpenCV 转换并保存 + cv::Mat bgr = frameToBgr(frame); + if (!bgr.empty()) + { + std::snprintf(fname, sizeof(fname), "%s/frame_%06lu.%s", + args.out_dir.c_str(), save_count, args.save_type.c_str()); + + std::vector params; + if (args.save_type == "jpg" || args.save_type == "jpeg") + { + params = {cv::IMWRITE_JPEG_QUALITY, args.quality}; + } + else if (args.save_type == "png") + { + params = {cv::IMWRITE_PNG_COMPRESSION, 3}; + } + ok = cv::imwrite(fname, bgr, params); + } + + if (ok) { + printf("[saved] %s\n", fname); + save_count++; + } else { + printf("[save failed] frame %lu\n", save_count); + } + } + + DebugInfo *pData = new DebugInfo(); + // pData->CameraId = (ctx->cam_no == 1 ? 0 : 1); + // pData->TimeStampMs = epochTimeStamp_ms; + memcpy(pData->Img, frame.data, frame.size); + bool res = send_debuginfo(pData); + //printf("send %ld %d %d\n", pData->TimeStampMs, pData->CameraId, res); + delete pData; + pData = nullptr; + + cap.releaseFrame(frame); + + // if (args.verbose) { + // printf("frame #%lu seq=%u size=%zu\n", + // frame_count, frame.sequence, frame.size); + // } + + auto now = clock::now(); + auto since_report = std::chrono::duration_cast(now - t_last_report); + if (since_report.count() >= 1000) { + double total_sec = std::chrono::duration_cast(now - t_start).count() / 1000.0; + double avg_fps = total_sec > 0 ? frame_count / total_sec : 0; + printf("[stats] frames=%lu avg_fps=%.2f door_position[%ld %ld %ld %ld]\n", frame_count, avg_fps, g_Feedback_Of_The_Actual_Position_Of_The_Left_Front_Door_391_S.load(), g_Feedback_Of_The_Actual_Position_Of_The_Right_Front_Door_391_S.load(), g_Feedback_Of_The_Actual_Position_Of_The_Left_Rear_Door_391_S.load(), g_Feedback_Of_The_Actual_Position_Of_The_Right_Rear_Door_391_S.load()); + t_last_report = now; + } + } + + auto t_end = clock::now(); + double total_sec = std::chrono::duration_cast(t_end - t_start).count() / 1000.0; + + printf("=========================================================\n"); + printf("Capture finished. frames=%lu duration=%.2fs avg_fps=%.2f\n", + frame_count, total_sec, total_sec > 0 ? frame_count / total_sec : 0); + if (args.save > 0) { + printf("Saved %lu frames to %s/\n", save_count, args.out_dir.c_str()); + } + if (args.can_enabled && CanComm::is_open()) { + printf("CAN received: %d frames\n", CanComm::receive_count()); + } + printf("=========================================================\n"); + + cap.stopStreaming(); + cap.close(); + + if (args.can_enabled) { + CanComm::can_stop(); + } + return 0; +} diff --git a/native/native_api.cpp b/native/native_api.cpp new file mode 100644 index 0000000..f256516 --- /dev/null +++ b/native/native_api.cpp @@ -0,0 +1,115 @@ +// Baseline V4L2/CAN + 0920 TensorRT/NPP, with a narrow ctypes ABI. +#include "v4l2_capture.h" +#include "can_comm.h" +#include "detector.hpp" +#include "uyvy_converter.hpp" +#include +#include +#include +#include +#include +#include +#include + +namespace { +using Clock=std::chrono::steady_clock; +double nowMs(){return std::chrono::duration(Clock::now().time_since_epoch()).count();} +thread_local std::string error; +std::mutex canMutex; +int positions[4]={-1,-1,-1,-1}; +double canTime=0; +void receiveCan(int64_t, uint32_t id, const uint8_t* data, uint32_t len){ + if(id!=0x391||len<4)return; + std::lock_guard lock(canMutex); + // Baseline comments specify bytes 1.0..4.7 (not overlapping bits 0,1,2,3). + for(int i=0;i<4;++i)positions[i]=data[i]; + canTime=nowMs(); +} +struct Engine { + fangjia::TrtLogger logger; + std::unique_ptr detector; + std::unique_ptr pose; + cudaStream_t stream{}; + int width,height,maxPeople; + Engine(const char* dir,int w,int h,int n,float confidence):width(w),height(h),maxPeople(n){ + detector=std::make_unique(std::filesystem::path(dir)/"yolo11s.engine",logger,confidence,n); + pose=std::make_unique(std::filesystem::path(dir)/"rtmw_l_384.engine",logger,w,h,n); + CUDA_CHECK(cudaStreamCreateWithFlags(&stream,cudaStreamNonBlocking)); + try {pose->warmup(stream);detector->infer(cv::Mat(h,w,CV_8UC3,cv::Scalar()),stream);}catch(...){cudaStreamDestroy(stream);throw;} + } + ~Engine(){if(stream)cudaStreamDestroy(stream);} +}; +struct Camera { + V4l2Capture capture; + std::unique_ptr converter; + std::thread worker; + std::mutex mutex; + std::condition_variable ready; + std::atomic stop{false}; + std::shared_ptr latest; + uint64_t generation=0, consumed=0, sequence=0; + double timestamp=0,driverTimestamp=0; + int width,height; + bool canEnabled=false; + std::string failure; + Camera(const char* device,int w,int h,int fps,bool can):width(w),height(h){ + if(!capture.open(device)||!capture.setFormat(w,h,V4L2_PIX_FMT_UYVY))throw std::runtime_error("Camera open/format failed"); + if(capture.width()!=uint32_t(w)||capture.height()!=uint32_t(h))throw std::runtime_error("Camera negotiated unexpected dimensions"); + capture.setFps(fps); + if(!capture.initMmap(4)||!capture.startStreaming())throw std::runtime_error("Camera streaming failed"); + converter=std::make_unique(w,h); + if(can){CanComm::set_msg_callback(receiveCan);canEnabled=CanComm::can_start();} + try {worker=std::thread([this]{loop();});}catch(...){if(canEnabled)CanComm::can_stop();throw;} + } + ~Camera(){stop=true;ready.notify_all();if(worker.joinable())worker.join();capture.close();if(canEnabled)CanComm::can_stop();} + void loop(){ + try { + int misses=0; + while(!stop){ + V4l2Capture::Frame frame; + if(!capture.getFrame(frame,1000)){if(++misses>=5)throw std::runtime_error("Camera timed out for 5 seconds");continue;} + double received=nowMs(); + misses=0; + auto stride=frame.bytes_per_line?frame.bytes_per_line:width*2; + if(frame.size image; + try {image=converter->convert(frame.data,stride);}catch(...){capture.releaseFrame(frame);throw;} + auto seq=frame.sequence;double driverTs=frame.timestamp_us/1000.0; + capture.releaseFrame(frame); + {std::lock_guard lock(mutex);latest=std::move(image);sequence=seq;timestamp=received;driverTimestamp=driverTs;++generation;} + ready.notify_one(); + } + }catch(std::exception const&e){std::lock_guard lock(mutex);failure=e.what();stop=true;ready.notify_all();} + } +}; +} +extern "C" { +const char* vap_error(){return error.c_str();} +void* vap_engine(const char* dir,int w,int h,int n,float confidence){try {return new Engine(dir,w,h,n,confidence);}catch(std::exception const&e){error=e.what();return nullptr;}} +void vap_engine_close(void* p){delete static_cast(p);} +int vap_detect(void* p,unsigned char* pixels,int w,int h,float* out){try { + auto& e=*static_cast(p);auto boxes=e.detector->infer(cv::Mat(h,w,CV_8UC3,pixels),e.stream); + for(size_t i=0;i(p); + if(count<0||count>e.maxPeople||w>e.width||h>e.height)throw std::runtime_error("Pose input exceeds configured capacity"); + std::vector input;for(int i=0;iinfer(cv::Mat(h,w,CV_8UC3,pixels),input,e.stream); + for(size_t i=0;i(p);} +int vap_read(void* p,unsigned char* out,double* meta){try { + auto& c=*static_cast(p);std::shared_ptr image; + {std::unique_lock lock(c.mutex); + if(!c.ready.wait_for(lock,std::chrono::seconds(6),[&]{return c.generation>c.consumed||c.stop;}))throw std::runtime_error("No camera frame"); + if(!c.failure.empty())throw std::runtime_error(c.failure); + if(c.stop)return 0; + image=c.latest;meta[0]=c.sequence;meta[1]=c.timestamp;meta[2]=c.generation-c.consumed-1;meta[3]=c.driverTimestamp;c.consumed=c.generation;} + std::memcpy(out,image->pixels.data(),image->pixels.size());return 1; +}catch(std::exception const&e){error=e.what();return -1;}} +void vap_can(double* out){std::lock_guard lock(canMutex);for(int i=0;i<4;++i)out[i]=positions[i];out[4]=canTime?nowMs()-canTime:-1;} +} diff --git a/orin/README.md b/orin/README.md index ca62f85..9bbe3f5 100644 --- a/orin/README.md +++ b/orin/README.md @@ -1,22 +1,39 @@ -# Orin 可直接运行目录 +# Orin 可直接运行包 — use-can -本目录包含实际在 Jetson AGX Orin 上编译的 aarch64 程序和 TensorRT FP16 引擎。 -目标环境：L4T R35.4.1 / JetPack 5.1.2，CUDA 11.4，TensorRT 8.5.2，OpenCV 4.5.4。 -使用系统安装的 CUDA、TensorRT、OpenCV 共享库；不适用于任意 JetPack 版本或 x86 主机。 +这个目录可以单独复制到 Orin 的任意目录，不依赖兄弟目录 0920/0923，不需要重新编译。 +包含板端编译的 TensorRT/CUDA 原生库、YOLO/RTMW engine、Python 风险代码、CAN USB 库、标定文件和原有 GMSL 驱动。 + +环境：Jetson AGX Orin，JetPack 5.1.2 / L4T R35.4.1，Linux 5.10.120-tegra，CUDA 11.4，TensorRT 8.5.x，系统 OpenCV 4.5.4 库；Python 3.8+，numpy、cv2。系统 CUDA/TensorRT/OpenCV 运行库由 JetPack 提供，未重复打包。 ```bash cd orin sha256sum -c SHA256SUMS -./scripts/initialize_camera.sh -./run.sh +./scripts/initialize_camera.sh # 开机后无 /dev/video5 时执行；需要 sudo 权限 +./run.sh # 默认 CAN 左侧、无可视化，结果保存在内存 +./run.sh --visualize # 可选浏览器预览，默认端口 8765 ``` -可视化可选：`./run.sh --visualize`，需要板端本地桌面 DISPLAY。 -默认持续推理直到 Ctrl+C / SIGTERM；正常退出没有统计输出。 -预览窗口关闭后推理继续。结果通过 C++ CameraPipeline::latest() 获取，不通过端口或结果文件传递。 +Ctrl+C / SIGTERM 退出。仅有一个进程可以占用同一个摄像头。 -程序为 build/camera_risk，配置为 config/camera.json，引擎为 engines/*.engine。 -./run.sh 从所在目录解析路径，也可以在其他工作目录通过绝对路径调用。 +```bash +./run.sh --no-can --active-door both --duration 10 +./run.sh --save-results --output ./outputs/session +./run.sh --visualize --save-results --save-snapshots --output ./outputs/session +``` -bin/drivers 中 GMSL 模块对应内核 5.10.120-tegra。初始化脚本检查内核版本；加载模块需要 root 或现有 sudo 权限。 -门区保留当前室内联调设置，door.calibration_verified=false。 +CAN 0x391 的 data[0..3] 顺序为左前、右前、左后、右后。`--can-side right` 切换右侧。未收到有效 CAN 或报文过期时输出 CAN_UNKNOWN，不发送执行器控制报文。 +标定文件仍为研发机位，默认 calibration_verified=false。 + +同进程下游接口： + +```python +from run_camera import CameraPipeline, parse_args +with CameraPipeline(parse_args([])) as pipeline: + result = pipeline.wait_for_result(timeout=10) + result = pipeline.latest_result # 等价于 pipeline.latest() + if result is not None: + print(result['risk'], result['door_risks']) +``` + +持续读取示例：`python3 examples/read_latest.py`（先停止其他采集进程）。`latest_result` 是单个只读快照；它不是跨进程共享变量。 +TensorRT engine 绑定目标 GPU/TRT 环境，不能直接用于 Thor、其他 JetPack 版本或桌面 GPU。 diff --git a/orin/SHA256SUMS b/orin/SHA256SUMS index 6034f9e..9588b90 100644 --- a/orin/SHA256SUMS +++ b/orin/SHA256SUMS @@ -1,11 +1,27 @@ -926a5165a4ebffe9568cf694018afb097cfb69b6e696a19939e5d364a9e3dfa5 README.md +4171567ecbea2f15b3c945b8d8363362d9f422b8c732d33ee1e15461b40014da README.md +9ba768d559fd3b64e833b50c62c70f2adfd6c169620aea5fba32cfb1b4c5b0e9 annotations/constant_door_mask.json +3544fb1c7b1b27b39db27d8f849d5e8652f5d09a950a59c5dd0c7703086b4f6b annotations/door_zone_multifeature.json 6d5e874fd27fcfd75d303d7e35d9672028af71041f52dcc82a1e83676faefa65 bin/drivers/max9295.ko 8a78d04ef8be44e9b62a5e4ee315d3d6d679e34aedff02a234bbf4d40e222933 bin/drivers/max9296.ko 3da628587477cc35d5e3f8aee86e718e2ddf16808e893e00ba6e164b2307e843 bin/drivers/sgx-yuv-gmsl2.ko -b5102fe0a529349bb8e9700ae1d8e8ddeaa7c4279cb28e7826c5cc32dce2bb1d build/camera_risk -f2334c43e6a48480260c770d46a3ab0edb855f6f52bc8b91d1559e69edfad043 config/camera.json -9ba768d559fd3b64e833b50c62c70f2adfd6c169620aea5fba32cfb1b4c5b0e9 config/door_mask.json -ebcd09e8c2e35008adbb0c1d222e580b17e0df1c8093df3481ba3374cd9a8174 engines/rtmw_l_384.engine -f6a65b926ac1144484f4aa93cd26d7997d32d62310ab93fe867f2bd4412222e9 engines/yolo11s.engine -7f42b0e7f9396e0d82313a892fc62c153261dab1206a206a466bce6376f58a36 run.sh -54e1d63e3aa72ad9b056a224d55968d4947ca6703a8d1c405570de5707a7e7ef scripts/initialize_camera.sh +cd75632beaaf81df2f59ce3e83fb335c21cc745068ef2b721c74cc3b2b03db26 build/libvision_native.so +49cc1067c1fa9f82a3d57dbfa77a55a54e5e5062ceae7ae84144b77745db79cf examples/read_latest.py +8dae8c557deadbe6bd7241c0dcd7025cb007077ee39a79b115e953a2b22cd9c1 lib/libusb-1.0.so +ab8b7f7249d5fb73d1d33d14fa7532a522088ffb6c25f37ba334daf4c6b4d313 lib/libusbcanfd.so +ebcd09e8c2e35008adbb0c1d222e580b17e0df1c8093df3481ba3374cd9a8174 models/engines/rtmw_l_384.engine +f6a65b926ac1144484f4aa93cd26d7997d32d62310ab93fe867f2bd4412222e9 models/engines/yolo11s.engine +fd3ae92255abf95ef037a1060ee75897f9705b4040f9394c3130f6ab860667e1 pipeline/compare_rtmw_ml.py +6d6d14a972bc4e5e9f8256514169e73ba0392a301cacd1221eaa26146aad00d2 pipeline/door_mask_risk.py +46d70e1e90a1f36d544b7492ef1fbf08821377a1c460a6683f42d4d7006055c7 pipeline/door_zone_risk.py +bc74a400ce5c5608c2fa49a69ccd8b41ec8633cd6a8239d9d594b8d9e5c3732a pipeline/fuse_rtmw_l_temporal.py +bcf2cc99dd27c8749801a5b743656767d9d64ef03d568d227cb654c573d70b00 pipeline/latest_result.py +cc8ea36d543896cd90f7a00822f06fc67771a1811859a4087995ac372df30276 pipeline/live_pose_engine.py +da299c11bd25127a09ffa738876f1de6e0d53d27add2840aeb6092ab0fe7016b pipeline/native_backend.py +81088b8bc072162bec13f57c65eb0d396e6d1e1074a83f92bcf29e393e27e11f pipeline/occlusion_aware_pose_v2.py +059e2a67c17ce63354986d3242a5d361103ee4241c5391f18140eeb19e47c939 pipeline/predict_optimized_pose_videos.py +bf8bb754cd841a52b4410f958a651b005de36ac4d17cc0992fbe1c90ed30caca pipeline/rtmw_l_person_mask_ablation.py +1bb58d851ba5d684527b8d1e1940db6241d2b0b2862687fa3b151899a65320d2 pipeline/rtmw_risk_pipeline.py +8a10e035334ed97b7480f558aff7b0464c255ac6138dd99f55a08ade68853e0e run.sh +4393f367b4ac7e5a00908043ffd21d5573cc12ed19b956fd101897ada22afd45 run_camera.py +11b5fbc3c17c0610ec0514a0a26cf5d92588d15c0e50e7086d9db5a5c0306e99 scripts/initialize_camera.sh +4cc40e244495e1de60a0429a8fc58cdf383ccda471955e53360a1c36a278c822 scripts/initialize_camera_impl.sh diff --git a/orin/annotations/constant_door_mask.json b/orin/annotations/constant_door_mask.json new file mode 100644 index 0000000..135b2a1 --- /dev/null +++ b/orin/annotations/constant_door_mask.json @@ -0,0 +1,77 @@ +{ + "version": "x-anylabeling", + "flags": {}, + "shapes": [ + { + "label": "door", + "points": [ + [ + 212.86, + 145.919 + ], + [ + 378.891, + 851.857 + ], + [ + 468.293, + 996.446 + ], + [ + 536.408, + 1294.131 + ], + [ + 553.437, + 1404.7 + ], + [ + 561.951, + 1532.279 + ], + [ + 51.086, + 1532.279 + ], + [ + 8.514, + 1515.268 + ], + [ + 0.0, + 1468.489 + ], + [ + 8.514, + 26.845 + ], + [ + 25.543, + 5.582 + ], + [ + 131.973, + -2.924 + ], + [ + 191.574, + 1.329 + ], + [ + 208.603, + 145.919 + ] + ], + "group_id": null, + "shape_type": "polygon", + "flags": {} + } + ], + "imagePath": "constant_door_mask_reference.jpg", + "imageData": null, + "imageHeight": 1536, + "imageWidth": 1920, + "sourceVideo": "D:\\huang.xuewei\\CMS\\video\\20260921104502\\segment_009.mp4", + "sourceFrameIndex": 244, + "constantDoorMask": true +} \ No newline at end of file diff --git a/orin/annotations/door_zone_multifeature.json b/orin/annotations/door_zone_multifeature.json new file mode 100644 index 0000000..3432c66 --- /dev/null +++ b/orin/annotations/door_zone_multifeature.json @@ -0,0 +1,47 @@ +{ + "version": "v9.10-research-calibration-1", + "description": "Temporary-camera calibration from seg_01 and seg_08; recalibrate for the production camera.", + "front_probability_min": 0.72, + "rear_probability_max": 0.28, + "minimum_person_height": 0.4, + "minimum_pose_features": 4, + "history_frames": 5, + "stale_track_frames": 120, + "features": { + "shoulder_center_x": { + "rear": 0.53, + "front": 0.64, + "weight": 1.2 + }, + "shoulder_center_y": { + "rear": 0.54, + "front": 0.46, + "weight": 1.3 + }, + "shoulder_width": { + "rear": 0.14, + "front": 0.2, + "weight": 1.4 + }, + "torso_height": { + "rear": 0.35, + "front": 0.47, + "weight": 1.5 + }, + "hip_center_y": { + "rear": 0.89, + "front": 0.93, + "weight": 0.6 + }, + "box_center_x": { + "rear": 0.48, + "front": 0.58, + "weight": 0.35 + }, + "box_height": { + "rear": 0.7, + "front": 0.82, + "weight": 0.45 + } + } +} diff --git a/orin/build/camera_risk b/orin/build/camera_risk deleted file mode 100755 index 276b22c..0000000 Binary files a/orin/build/camera_risk and /dev/null differ diff --git a/orin/build/libvision_native.so b/orin/build/libvision_native.so new file mode 100755 index 0000000..db26404 Binary files /dev/null and b/orin/build/libvision_native.so differ diff --git a/orin/config/camera.json b/orin/config/camera.json deleted file mode 100644 index 6cd1134..0000000 --- a/orin/config/camera.json +++ /dev/null @@ -1,22 +0,0 @@ -{ - "camera": { - "device": "/dev/video5", - "width": 1920, - "height": 1536, - "fps": 30, - "buffers": 4, - "timeout_ms": 2000 - }, - "inference": { - "engine_directory": "../engines", - "detector_confidence": 0.30, - "max_people": 3, - "queue_capacity": 2 - }, - "door": { - "mask": "door_mask.json", - "warning_relative": 0.30, - "release_relative": 0.45, - "calibration_verified": false - } -} diff --git a/orin/examples/read_latest.py b/orin/examples/read_latest.py new file mode 100644 index 0000000..2a3f0bd --- /dev/null +++ b/orin/examples/read_latest.py @@ -0,0 +1,36 @@ +#!/usr/bin/env python3 +"""Same-process downstream example. Run on Orin, which owns the camera/GPU.""" +import sys +from pathlib import Path +sys.path.insert(0, str(Path(__file__).resolve().parents[1])) +from run_camera import CameraPipeline, parse_args + + +def main(): + # No visualization or disk output by default. CLI flags can enable either. + pipeline = CameraPipeline(parse_args()) + last_frame = -1 + try: + pipeline.start() + while True: + # Reading the variable is enough for a consumer with its own clock: + # result = pipeline.latest_result # equivalent to pipeline.latest() + # This sample waits to avoid busy polling; still only the latest + # result is returned, without queuing/acknowledging individual frames. + result = pipeline.wait_for_result(last_frame, timeout=1.0) + if result is not None: + last_frame = result['frame_index'] + print(last_frame, result['risk']['risk_type'], + result['risk']['risk_score'], flush=True) + # Downstream code uses result['door_risks'], ['skeletons'], ['can']. + elif not pipeline.running(): + break + except KeyboardInterrupt: + pass + finally: + pipeline.request_stop() + pipeline.join() # Also propagates worker errors to the caller. + + +if __name__ == '__main__': + main() diff --git a/orin/lib/libusb-1.0.so b/orin/lib/libusb-1.0.so new file mode 100644 index 0000000..0824ee0 Binary files /dev/null and b/orin/lib/libusb-1.0.so differ diff --git a/orin/lib/libusb-1.0.so.0 b/orin/lib/libusb-1.0.so.0 new file mode 120000 index 0000000..4970d7f --- /dev/null +++ b/orin/lib/libusb-1.0.so.0 @@ -0,0 +1 @@ +libusb-1.0.so \ No newline at end of file diff --git a/orin/lib/libusbcanfd.so b/orin/lib/libusbcanfd.so new file mode 100644 index 0000000..72b9370 Binary files /dev/null and b/orin/lib/libusbcanfd.so differ diff --git a/orin/lib/libusbcanfd.so.1.0.12 b/orin/lib/libusbcanfd.so.1.0.12 new file mode 120000 index 0000000..2ee2cf4 --- /dev/null +++ b/orin/lib/libusbcanfd.so.1.0.12 @@ -0,0 +1 @@ +libusbcanfd.so \ No newline at end of file diff --git a/orin/engines/rtmw_l_384.engine b/orin/models/engines/rtmw_l_384.engine similarity index 100% rename from orin/engines/rtmw_l_384.engine rename to orin/models/engines/rtmw_l_384.engine diff --git a/orin/engines/yolo11s.engine b/orin/models/engines/yolo11s.engine similarity index 100% rename from orin/engines/yolo11s.engine rename to orin/models/engines/yolo11s.engine diff --git a/orin/pipeline/compare_rtmw_ml.py b/orin/pipeline/compare_rtmw_ml.py new file mode 100644 index 0000000..7d92a4a --- /dev/null +++ b/orin/pipeline/compare_rtmw_ml.py @@ -0,0 +1,384 @@ +"""Compare RTMW-M 256x192 and RTMW-L 384x288 on cached YOLO boxes. + +This is an experiment script. It intentionally keeps the candidate weights +outside the formal models/ directory and does not change deployment code. +""" + +from __future__ import annotations + +import argparse +import json +import math +import time +from collections import defaultdict +from pathlib import Path + +import cv2 +import numpy as np + + +BODY_EDGES = ((5, 7), (7, 9), (6, 8), (8, 10)) +HAND_EDGES_LOCAL = ( + (0, 1), (1, 2), (2, 3), (3, 4), + (0, 5), (5, 6), (6, 7), (7, 8), + (0, 9), (9, 10), (10, 11), (11, 12), + (0, 13), (13, 14), (14, 15), (15, 16), + (0, 17), (17, 18), (18, 19), (19, 20), +) +HAND_SLICES = {"left": (91, 112), "right": (112, 133)} +BODY_WRISTS = {"left": 9, "right": 10} + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser() + parser.add_argument("--video", type=Path, required=True) + parser.add_argument("--boxes", type=Path, required=True) + parser.add_argument("--model-m", type=Path, required=True) + parser.add_argument("--model-l", type=Path, required=True) + parser.add_argument("--output", type=Path, required=True) + parser.add_argument("--max-people", type=int, default=2) + return parser.parse_args() + + +def load_boxes(path: Path, max_people: int) -> dict[int, list[list[float]]]: + frames: dict[int, list[list[float]]] = {} + with path.open("r", encoding="utf-8") as handle: + for line in handle: + item = json.loads(line) + boxes = [entry["xyxy"] for entry in item["boxes"]] + frames[int(item["frame_index"])] = sorted(boxes, key=lambda b: b[0])[:max_people] + return frames + + +def point(pt: np.ndarray) -> tuple[int, int]: + return int(round(float(pt[0]))), int(round(float(pt[1]))) + + +def draw_pose( + image: np.ndarray, + boxes: list[list[float]], + keypoints: np.ndarray, + title: str, +) -> np.ndarray: + canvas = image.copy() + for person_idx, (bbox, kpts) in enumerate(zip(boxes, keypoints), start=1): + x1, y1, x2, y2 = (int(round(value)) for value in bbox) + cv2.rectangle(canvas, (x1, y1), (x2, y2), (0, 220, 255), 2) + cv2.putText( + canvas, + f"P{person_idx}", + (x1 + 4, max(22, y1 + 24)), + cv2.FONT_HERSHEY_SIMPLEX, + 0.75, + (0, 220, 255), + 2, + cv2.LINE_AA, + ) + + for edge_idx, (a, b) in enumerate(BODY_EDGES): + color = (40, 40, 255) if edge_idx < 2 else (30, 230, 30) + cv2.line(canvas, point(kpts[a]), point(kpts[b]), color, 3, cv2.LINE_AA) + for idx in (5, 6, 7, 8, 9, 10): + color = (40, 40, 255) if idx in (5, 7, 9) else (30, 230, 30) + cv2.circle(canvas, point(kpts[idx]), 5, color, -1, cv2.LINE_AA) + + for side, (start, end) in HAND_SLICES.items(): + hand = kpts[start:end] + color = (255, 80, 220) if side == "left" else (255, 190, 40) + for a, b in HAND_EDGES_LOCAL: + cv2.line(canvas, point(hand[a]), point(hand[b]), color, 2, cv2.LINE_AA) + for hand_point in hand: + cv2.circle(canvas, point(hand_point), 2, color, -1, cv2.LINE_AA) + + cv2.rectangle(canvas, (0, 0), (610, 45), (0, 0, 0), -1) + cv2.putText( + canvas, + title, + (10, 31), + cv2.FONT_HERSHEY_SIMPLEX, + 0.8, + (255, 255, 255), + 2, + cv2.LINE_AA, + ) + return canvas + + +def normalized_keypoints(kpts: np.ndarray, bbox: list[float]) -> np.ndarray: + x1, y1, x2, y2 = bbox + width = max(x2 - x1, 1.0) + height = max(y2 - y1, 1.0) + result = kpts.copy().astype(np.float64) + result[:, 0] = (result[:, 0] - x1) / width + result[:, 1] = (result[:, 1] - y1) / height + return result + + +def hand_span(kpts: np.ndarray, bbox: list[float], side: str) -> float: + start, end = HAND_SLICES[side] + hand = kpts[start:end] + span = math.hypot( + float(np.ptp(hand[:, 0])), + float(np.ptp(hand[:, 1])), + ) + x1, y1, x2, y2 = bbox + return span / max(math.hypot(x2 - x1, y2 - y1), 1.0) + + +def summarize(values: list[float]) -> dict[str, float | int | None]: + if not values: + return {"count": 0, "median": None, "p95": None, "max": None} + array = np.asarray(values, dtype=np.float64) + return { + "count": int(array.size), + "median": round(float(np.median(array)), 5), + "p95": round(float(np.percentile(array, 95)), 5), + "max": round(float(np.max(array)), 5), + } + + +def temporal_metrics(records: list[dict]) -> dict: + by_model_person: dict[tuple[str, int], list[dict]] = defaultdict(list) + for record in records: + for model_name in ("M", "L"): + for person_idx in range(len(record[model_name]["keypoints"])): + by_model_person[(model_name, person_idx)].append( + { + "frame": record["frame_index"], + "bbox": record["boxes"][person_idx], + "kpts": np.asarray( + record[model_name]["keypoints"][person_idx], dtype=np.float64 + ), + } + ) + + result: dict[str, dict] = {"M": {}, "L": {}} + for (model_name, person_idx), sequence in by_model_person.items(): + upper_steps: list[float] = [] + wrist_steps = {"left": [], "right": []} + hand_steps = {"left": [], "right": []} + hand_spans = {"left": [], "right": []} + bone_lengths = {"left_upper": [], "left_forearm": [], + "right_upper": [], "right_forearm": []} + + previous: np.ndarray | None = None + for entry in sequence: + normalized = normalized_keypoints(entry["kpts"], entry["bbox"]) + if previous is not None: + steps = np.linalg.norm(normalized - previous, axis=1) + upper_steps.extend(float(steps[idx]) for idx in (5, 6, 7, 8, 9, 10)) + for side, idx in BODY_WRISTS.items(): + wrist_steps[side].append(float(steps[idx])) + for side, (start, end) in HAND_SLICES.items(): + hand_steps[side].extend(float(value) for value in steps[start:end]) + previous = normalized + + hand_spans["left"].append(hand_span(entry["kpts"], entry["bbox"], "left")) + hand_spans["right"].append(hand_span(entry["kpts"], entry["bbox"], "right")) + bone_lengths["left_upper"].append(float(np.linalg.norm(normalized[5] - normalized[7]))) + bone_lengths["left_forearm"].append(float(np.linalg.norm(normalized[7] - normalized[9]))) + bone_lengths["right_upper"].append(float(np.linalg.norm(normalized[6] - normalized[8]))) + bone_lengths["right_forearm"].append(float(np.linalg.norm(normalized[8] - normalized[10]))) + + span_summary = {} + for side, values in hand_spans.items(): + median = float(np.median(values)) + gross_limit = max(0.2, median * 2.5) + span_summary[side] = { + **summarize(values), + "gross_limit": round(gross_limit, 5), + "gross_outlier_frames": int(sum(value > gross_limit for value in values)), + } + + bone_summary = {} + for name, values in bone_lengths.items(): + array = np.asarray(values, dtype=np.float64) + mean = float(np.mean(array)) + bone_summary[name] = { + **summarize(values), + "coefficient_of_variation": round(float(np.std(array) / max(mean, 1e-9)), 5), + } + + result[model_name][f"P{person_idx + 1}"] = { + "upper_body_frame_step": summarize(upper_steps), + "wrist_frame_step": {side: summarize(values) for side, values in wrist_steps.items()}, + "hand_landmark_frame_step": { + side: summarize(values) for side, values in hand_steps.items() + }, + "hand_span_relative_to_person": span_summary, + "bone_length_relative_to_box": bone_summary, + } + return result + + +def model_disagreement(records: list[dict]) -> tuple[dict, list[int]]: + values: dict[str, list[float]] = defaultdict(list) + per_frame: list[tuple[float, int]] = [] + for record in records: + frame_max = 0.0 + for person_idx, bbox in enumerate(record["boxes"]): + m = normalized_keypoints( + np.asarray(record["M"]["keypoints"][person_idx]), bbox + ) + l = normalized_keypoints( + np.asarray(record["L"]["keypoints"][person_idx]), bbox + ) + distance = np.linalg.norm(m - l, axis=1) + values[f"P{person_idx + 1}_upper"].extend( + float(distance[idx]) for idx in (5, 6, 7, 8, 9, 10) + ) + for side, idx in BODY_WRISTS.items(): + values[f"P{person_idx + 1}_{side}_wrist"].append(float(distance[idx])) + frame_max = max(frame_max, float(np.max(distance[[5, 6, 7, 8, 9, 10]]))) + per_frame.append((frame_max, record["frame_index"])) + worst_frames = [frame for _, frame in sorted(per_frame, reverse=True)[:12]] + return {name: summarize(items) for name, items in values.items()}, worst_frames + + +def write_montage( + video_path: Path, + records: list[dict], + selected_frames: list[int], + output_path: Path, + titles: tuple[str, str] = ("RTMW-M 256x192", "RTMW-L 384x288"), +) -> None: + selected = set(selected_frames) + record_map = {record["frame_index"]: record for record in records} + cap = cv2.VideoCapture(str(video_path)) + rows: list[np.ndarray] = [] + frame_idx = 0 + while True: + ok, frame = cap.read() + if not ok: + break + if frame_idx in selected: + record = record_map[frame_idx] + m = draw_pose( + frame, + record["boxes"], + np.asarray(record["M"]["keypoints"]), + f"{titles[0]} | frame {frame_idx}", + ) + l = draw_pose( + frame, + record["boxes"], + np.asarray(record["L"]["keypoints"]), + f"{titles[1]} | frame {frame_idx}", + ) + m = cv2.resize(m, (640, 360), interpolation=cv2.INTER_AREA) + l = cv2.resize(l, (640, 360), interpolation=cv2.INTER_AREA) + rows.append(np.hstack([m, l])) + frame_idx += 1 + cap.release() + if rows: + cv2.imwrite(str(output_path), np.vstack(rows)) + + +def main() -> None: + from rtmlib import RTMPose + args = parse_args() + args.output.mkdir(parents=True, exist_ok=True) + cached_boxes = load_boxes(args.boxes, args.max_people) + + pose_m = RTMPose( + str(args.model_m), model_input_size=(192, 256), backend="onnxruntime", device="cpu" + ) + pose_l = RTMPose( + str(args.model_l), model_input_size=(288, 384), backend="onnxruntime", device="cpu" + ) + + cap = cv2.VideoCapture(str(args.video)) + fps = float(cap.get(cv2.CAP_PROP_FPS)) + width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)) + height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)) + video_output = args.output / "reference_video_rtmw_m_vs_l.mp4" + writer = cv2.VideoWriter( + str(video_output), + cv2.VideoWriter_fourcc(*"mp4v"), + fps, + (width, height // 2), + ) + + records: list[dict] = [] + timings = {"M": [], "L": []} + frame_idx = 0 + while True: + ok, frame = cap.read() + if not ok: + break + boxes = cached_boxes[frame_idx] + + started = time.perf_counter() + kpts_m, scores_m = pose_m(frame, bboxes=boxes) + timings["M"].append((time.perf_counter() - started) * 1000.0) + + started = time.perf_counter() + kpts_l, scores_l = pose_l(frame, bboxes=boxes) + timings["L"].append((time.perf_counter() - started) * 1000.0) + + records.append( + { + "frame_index": frame_idx, + "boxes": boxes, + "M": { + "keypoints": np.round(kpts_m, 3).tolist(), + "raw_scores": np.round(scores_m, 4).tolist(), + }, + "L": { + "keypoints": np.round(kpts_l, 3).tolist(), + "raw_scores": np.round(scores_l, 4).tolist(), + }, + } + ) + + left = draw_pose(frame, boxes, kpts_m, f"RTMW-M 256x192 | frame {frame_idx}") + right = draw_pose(frame, boxes, kpts_l, f"RTMW-L 384x288 | frame {frame_idx}") + left = cv2.resize(left, (width // 2, height // 2), interpolation=cv2.INTER_AREA) + right = cv2.resize(right, (width // 2, height // 2), interpolation=cv2.INTER_AREA) + writer.write(np.hstack([left, right])) + + if frame_idx % 50 == 0: + print(f"processed frame {frame_idx}", flush=True) + frame_idx += 1 + + cap.release() + writer.release() + + jsonl_path = args.output / "reference_video_rtmw_m_vs_l.jsonl" + with jsonl_path.open("w", encoding="utf-8") as handle: + for record in records: + handle.write(json.dumps(record, ensure_ascii=False) + "\n") + + disagreement, worst_frames = model_disagreement(records) + fixed_frames = [0, 60, 120, 180, 240, 252, 300, 360, 411] + montage_frames = sorted(set(fixed_frames + worst_frames)) + write_montage( + args.video, + records, + montage_frames, + args.output / "rtmw_m_vs_l_full_montage.jpg", + ) + + summary = { + "frames": len(records), + "input_sizes_wh": {"M": [192, 256], "L": [288, 384]}, + "raw_score_note": ( + "M and L raw SimCC peak scores use different scales and must not be compared " + "or thresholded with the same numeric value." + ), + "timing_ms_two_people_cpu": { + name: summarize(values[5:] if len(values) > 5 else values) + for name, values in timings.items() + }, + "temporal_metrics_normalized_to_yolo_box": temporal_metrics(records), + "m_l_position_disagreement_normalized_to_yolo_box": disagreement, + "largest_upper_body_disagreement_frames": worst_frames, + "montage_frames": montage_frames, + } + with (args.output / "summary.json").open("w", encoding="utf-8") as handle: + json.dump(summary, handle, ensure_ascii=False, indent=2) + print(json.dumps(summary, ensure_ascii=False, indent=2)) + + +if __name__ == "__main__": + main() diff --git a/orin/pipeline/door_mask_risk.py b/orin/pipeline/door_mask_risk.py new file mode 100644 index 0000000..49eeec5 --- /dev/null +++ b/orin/pipeline/door_mask_risk.py @@ -0,0 +1,488 @@ +"""Per-frame door-mask loading and YOLO/RTMW-only 2-D risk scoring.""" + +from __future__ import annotations + +import json +import math +import re +import time +from collections import Counter +from dataclasses import dataclass +from pathlib import Path + +import cv2 +import numpy as np + + +RASTER_SUFFIXES = {".png", ".bmp", ".tif", ".tiff"} +DEFAULT_DOOR_LABELS = ("door", "车门") + + +@dataclass +class LoadedDoorMask: + status: str + annotation_index: int + path: Path | None + mask: np.ndarray | None + source_format: str | None + + def metadata(self) -> dict: + return { + "status": self.status, + "annotation_index": self.annotation_index, + "path": str(self.path.resolve()) if self.path is not None else None, + "format": self.source_format, + "area_px": int(np.count_nonzero(self.mask)) if self.mask is not None else 0, + } + + +def _frame_index_from_path(path: Path) -> int | None: + matches = re.findall(r"\d+", path.stem) + return int(matches[-1]) if matches else None + + +class DoorMaskSequence: + """Index a directory of X-AnyLabeling JSON polygons or raster masks.""" + + def __init__( + self, + directory: Path | None = None, + mask_format: str = "auto", + labels: tuple[str, ...] = DEFAULT_DOOR_LABELS, + frame_offset: int = 0, + require_every_frame: bool = False, + constant_mask: Path | None = None, + ): + if (directory is None) == (constant_mask is None): + raise ValueError("Provide exactly one of directory or constant_mask") + self.constant_path = ( + Path(constant_mask).expanduser().resolve() + if constant_mask is not None + else None + ) + self.directory = ( + self.constant_path.parent if self.constant_path is not None else Path(directory) + ) + if self.constant_path is not None and not self.constant_path.is_file(): + raise ValueError(f"Constant door mask does not exist: {self.constant_path}") + if self.constant_path is None and not self.directory.is_dir(): + raise ValueError(f"Door mask directory does not exist: {self.directory}") + self.labels = {label.casefold() for label in labels} + self.frame_offset = int(frame_offset) + self.require_every_frame = require_every_frame + + if self.constant_path is not None: + if mask_format == "auto": + selected_format = ( + "xanylabeling" + if self.constant_path.suffix.lower() == ".json" + else "raster" + ) + else: + selected_format = mask_format + self.source_format = selected_format + self.paths = {0: self.constant_path} + return + + json_paths = sorted(self.directory.rglob("*.json")) + raster_paths = sorted( + path + for path in self.directory.rglob("*") + if path.is_file() and path.suffix.lower() in RASTER_SUFFIXES + ) + if mask_format == "auto": + selected_format = "xanylabeling" if json_paths else "raster" + else: + selected_format = mask_format + candidates = json_paths if selected_format == "xanylabeling" else raster_paths + if not candidates: + expected = "JSON" if selected_format == "xanylabeling" else "PNG/BMP/TIFF" + raise ValueError(f"No {expected} door masks found in {self.directory}") + + self.source_format = selected_format + self.paths: dict[int, Path] = {} + for path in candidates: + index = _frame_index_from_path(path) + if index is None and selected_format == "xanylabeling": + data = json.loads(path.read_text(encoding="utf-8-sig")) + explicit = data.get("frame_index", data.get("frameIndex")) + index = int(explicit) if explicit is not None else None + if index is None: + continue + if index in self.paths: + raise ValueError( + f"Multiple door masks map to annotation frame {index}: " + f"{self.paths[index]} and {path}" + ) + self.paths[index] = path + if not self.paths: + raise ValueError( + "Door mask filenames need a frame number, for example 000000.json" + ) + + def load(self, video_frame_index: int, frame_shape: tuple[int, int]) -> LoadedDoorMask: + if self.constant_path is not None: + annotation_index = video_frame_index + path = self.constant_path + else: + annotation_index = video_frame_index + self.frame_offset + path = self.paths.get(annotation_index) + if path is None: + if self.require_every_frame: + raise RuntimeError( + f"No door mask for video frame {video_frame_index} " + f"(expected annotation index {annotation_index})" + ) + return LoadedDoorMask("missing", annotation_index, None, None, self.source_format) + + if self.source_format == "xanylabeling": + mask = self._load_xanylabeling(path, frame_shape) + else: + mask = self._load_raster(path, frame_shape) + status = "loaded" if np.any(mask) else "empty" + return LoadedDoorMask(status, annotation_index, path, mask, self.source_format) + + @staticmethod + def _load_raster(path: Path, frame_shape: tuple[int, int]) -> np.ndarray: + image = cv2.imread(str(path), cv2.IMREAD_UNCHANGED) + if image is None: + raise RuntimeError(f"Cannot read door mask: {path}") + if image.ndim == 3: + image = np.max(image[:, :, :3], axis=2) + height, width = frame_shape + if image.shape != (height, width): + image = cv2.resize(image, (width, height), interpolation=cv2.INTER_NEAREST) + return image > 0 + + def _load_xanylabeling( + self, path: Path, frame_shape: tuple[int, int] + ) -> np.ndarray: + data = json.loads(path.read_text(encoding="utf-8-sig")) + height, width = frame_shape + source_width = float(data.get("imageWidth") or width) + source_height = float(data.get("imageHeight") or height) + scale = np.asarray( + [width / max(source_width, 1.0), height / max(source_height, 1.0)], + dtype=np.float32, + ) + mask = np.zeros((height, width), dtype=np.uint8) + for shape in data.get("shapes", []): + if str(shape.get("label", "")).casefold() not in self.labels: + continue + points = np.asarray(shape.get("points", []), dtype=np.float32).reshape(-1, 2) + if len(points) < 2: + continue + points = np.rint(points * scale).astype(np.int32) + shape_type = str(shape.get("shape_type", "polygon")).lower() + if shape_type == "rectangle" and len(points) >= 2: + x1, y1 = points[0] + x2, y2 = points[1] + cv2.rectangle(mask, (x1, y1), (x2, y2), 1, -1) + elif shape_type == "circle" and len(points) >= 2: + radius = int(round(float(np.linalg.norm(points[1] - points[0])))) + cv2.circle(mask, tuple(points[0]), radius, 1, -1) + elif len(points) >= 3: + cv2.fillPoly(mask, [points], 1) + return mask > 0 + + +def _fit_mask(mask: np.ndarray, frame_shape: tuple[int, int]) -> np.ndarray: + height, width = frame_shape + result = np.asarray(mask) > 0 + if result.shape != (height, width): + result = cv2.resize( + result.astype(np.uint8), (width, height), interpolation=cv2.INTER_NEAREST + ) > 0 + return result + + +def _box_mask(box: list[float] | None, frame_shape: tuple[int, int]) -> np.ndarray | None: + if box is None or len(box) < 4: + return None + height, width = frame_shape + x1, y1, x2, y2 = [int(round(float(value))) for value in box[:4]] + x1, x2 = sorted((max(0, min(width - 1, x1)), max(0, min(width - 1, x2)))) + y1, y2 = sorted((max(0, min(height - 1, y1)), max(0, min(height - 1, y2)))) + if x2 <= x1 or y2 <= y1: + return None + mask = np.zeros((height, width), dtype=bool) + mask[y1 : y2 + 1, x1 : x2 + 1] = True + return mask + + +def _hand_mask(points: list, frame_shape: tuple[int, int]) -> np.ndarray | None: + hand = np.asarray(points, dtype=np.float32).reshape(-1, 2) + if len(hand) < 3 or not np.isfinite(hand).all(): + return None + hull = cv2.convexHull(np.rint(hand).astype(np.int32)) + mask = np.zeros(frame_shape, dtype=np.uint8) + cv2.fillConvexPoly(mask, hull, 1) + return mask > 0 if np.any(mask) else None + + +def _decision(distance_relative: float, overlap_px: int, warning: float, release: float) -> tuple[str, str]: + if overlap_px > 0: + return "danger", "DANGER_INTERSECTION" + if distance_relative <= warning: + return "warning", "WARNING_NEAR" + if distance_relative <= release: + return "watch", "WATCH" + return "clear", "CLEAR" + + +class DoorRiskEvaluator: + """Score RTMW hand regions, with one YOLO body-box fallback per person.""" + + def __init__( + self, + sequence: DoorMaskSequence, + warning_relative: float = 0.30, + release_relative: float = 0.45, + ): + if warning_relative < 0 or release_relative <= warning_relative: + raise ValueError("Risk thresholds require 0 <= warning < release") + self.sequence = sequence + self.warning_relative = float(warning_relative) + self.release_relative = float(release_relative) + self.mask_status_counts: Counter = Counter() + self.level_counts: Counter = Counter() + self.source_counts: Counter = Counter() + self.maximum_score = 0.0 + self.timing_seconds: Counter = Counter() + self.timing_counts: Counter = Counter() + self._constant_cache_shape: tuple[int, int] | None = None + self._constant_loaded: LoadedDoorMask | None = None + self._constant_distance_map: np.ndarray | None = None + self.distance_map_cache_builds = 0 + + def _record_timing(self, name: str, started: float) -> None: + self.timing_seconds[name] += time.perf_counter() - started + self.timing_counts[name] += 1 + + def _measure_region( + self, + region: np.ndarray, + door_mask: np.ndarray, + distance_map: np.ndarray, + source: str, + track_id: int, + side: str, + state: str | None, + ) -> dict: + area_px = int(np.count_nonzero(region)) + overlap_px = int(np.count_nonzero(region & door_mask)) + distances = distance_map[region] + distance_px = float(np.min(distances)) if len(distances) else math.inf + scale_px = max(math.sqrt(area_px), 1.0) + distance_relative = distance_px / scale_px + score = float(np.clip(1.0 - distance_relative / self.release_relative, 0.0, 1.0)) + if overlap_px: + score = 1.0 + level, decision = _decision( + distance_relative, + overlap_px, + self.warning_relative, + self.release_relative, + ) + result = { + "status": "calculated", + "risk_score": round(score, 5), + "risk_level": level, + "decision": decision, + "source": source, + "track_id": int(track_id), + "side": side, + "state": state, + "region_area_px": area_px, + "door_overlap_px": overlap_px, + "region_overlap_ratio": round(overlap_px / max(area_px, 1), 6), + "min_distance_px": round(distance_px, 3), + "region_scale_px": round(scale_px, 3), + "distance_relative": round(distance_relative, 6), + } + self.level_counts[level] += 1 + self.source_counts[source] += 1 + self.maximum_score = max(self.maximum_score, score) + return result + + @staticmethod + def _unavailable(source: str, track_id: int, side: str, status: str) -> dict: + return { + "status": status, + "risk_score": None, + "risk_level": "unavailable", + "decision": "EVIDENCE_UNAVAILABLE", + "source": source, + "track_id": int(track_id), + "side": side, + } + + def evaluate( + self, + frame_index: int, + record: dict, + frame_shape: tuple[int, int], + ) -> tuple[dict, np.ndarray | None]: + total_started = time.perf_counter() + stage_started = time.perf_counter() + use_constant_cache = self.sequence.constant_path is not None + if use_constant_cache and self._constant_cache_shape == frame_shape: + cached = self._constant_loaded + assert cached is not None + loaded = LoadedDoorMask( + cached.status, + frame_index, + cached.path, + cached.mask, + cached.source_format, + ) + else: + loaded = self.sequence.load(frame_index, frame_shape) + if use_constant_cache: + self._constant_cache_shape = frame_shape + self._constant_loaded = loaded + self._constant_distance_map = None + self._record_timing("door_mask_load", stage_started) + self.mask_status_counts[loaded.status] += 1 + mask_ready = loaded.status == "loaded" and loaded.mask is not None + door_mask = loaded.mask if mask_ready else None + distance_map = None + if door_mask is not None and use_constant_cache and self._constant_distance_map is not None: + distance_map = self._constant_distance_map + elif door_mask is not None: + stage_started = time.perf_counter() + distance_map = cv2.distanceTransform( + (~door_mask).astype(np.uint8), cv2.DIST_L2, cv2.DIST_MASK_PRECISE + ) + self._record_timing("distance_transform", stage_started) + if use_constant_cache: + self._constant_distance_map = distance_map + self.distance_map_cache_builds += 1 + + stage_started = time.perf_counter() + regions: list[tuple[dict, np.ndarray, str, int, str, str | None]] = [] + track_ids = [int(value) for value in record.get("track_ids", [])] + person_positions_with_hands: set[int] = set() + for hand in record.get("global_hands", []): + position = int(hand.get("person_id", 0)) - 1 + track_id = track_ids[position] if 0 <= position < len(track_ids) else position + 1 + side = str(hand.get("side", "unknown")) + region = _hand_mask(hand.get("keypoints", []), frame_shape) + if region is None: + continue + person_positions_with_hands.add(position) + regions.append( + ( + hand, + region, + "RTMW_HAND_KEYPOINTS", + track_id, + side, + str(hand.get("state", "unknown")), + ) + ) + + boxes = record.get("boxes", []) + body_fallbacks = [] + for position, box in enumerate(boxes): + if position in person_positions_with_hands: + continue + region = _box_mask(box, frame_shape) + if region is None: + continue + track_id = track_ids[position] if position < len(track_ids) else position + 1 + fallback = { + "track_id": track_id, + "side": "body", + "state": "rtmw_hand_unavailable", + "body_zone_box": box, + } + body_fallbacks.append(fallback) + regions.append( + ( + fallback, + region, + "YOLO_BODY_BOX_FALLBACK", + track_id, + "body", + "rtmw_hand_unavailable", + ) + ) + record["body_fallbacks"] = body_fallbacks + self._record_timing("region_build", stage_started) + + stage_started = time.perf_counter() + measurements = [] + unavailable_status = f"door_mask_{loaded.status}" + for owner, region, source, track_id, side, state in regions: + if door_mask is None or distance_map is None: + measurement = self._unavailable( + source, track_id, side, unavailable_status + ) + else: + measurement = self._measure_region( + region, door_mask, distance_map, source, track_id, side, state + ) + owner["risk"] = measurement + measurements.append(measurement) + self._record_timing("region_measure", stage_started) + + calculated = [item for item in measurements if item["risk_score"] is not None] + nearest = max(calculated, key=lambda item: item["risk_score"]) if calculated else None + if loaded.status != "loaded": + status = unavailable_status + elif not calculated: + status = "person_region_missing" + else: + status = "calculated" + frame_risk = { + "status": status, + "risk_score": nearest["risk_score"] if nearest else None, + "risk_level": nearest["risk_level"] if nearest else "unavailable", + "decision": nearest["decision"] if nearest else "EVIDENCE_UNAVAILABLE", + "source": nearest["source"] if nearest else None, + "track_id": nearest["track_id"] if nearest else None, + "side": nearest["side"] if nearest else None, + "door_mask": loaded.metadata(), + "warning_distance_relative": self.warning_relative, + "release_distance_relative": self.release_relative, + "regions": measurements, + } + self._record_timing("risk_total", total_started) + return frame_risk, loaded.mask + + def summary(self) -> dict: + return { + "mask_directory": str(self.sequence.directory.resolve()), + "constant_mask": ( + str(self.sequence.constant_path) + if self.sequence.constant_path is not None + else None + ), + "constant_for_all_frames": self.sequence.constant_path is not None, + "mask_format": self.sequence.source_format, + "door_labels": sorted(self.sequence.labels), + "frame_offset": self.sequence.frame_offset, + "warning_distance_relative": self.warning_relative, + "release_distance_relative": self.release_relative, + "mask_status_counts": dict(self.mask_status_counts), + "risk_level_counts": dict(self.level_counts), + "source_counts": dict(self.source_counts), + "maximum_risk_score": round(self.maximum_score, 5), + "constant_distance_map_cached": self.sequence.constant_path is not None, + "distance_map_cache_builds": self.distance_map_cache_builds, + "timing_ms": { + name: { + "count": int(self.timing_counts[name]), + "total": round(float(seconds) * 1000.0, 3), + "mean": round( + float(seconds) * 1000.0 + / max(int(self.timing_counts[name]), 1), + 3, + ), + } + for name, seconds in self.timing_seconds.items() + }, + "score_semantics": "2-D image overlap/proximity research score; not millimeters or safety certification", + "region_policy": "RTMW alarm-eligible hand hulls; YOLO body box only when a person has no usable RTMW hand region", + } diff --git a/orin/pipeline/door_zone_risk.py b/orin/pipeline/door_zone_risk.py new file mode 100644 index 0000000..972cfb7 --- /dev/null +++ b/orin/pipeline/door_zone_risk.py @@ -0,0 +1,358 @@ +"""Multi-feature person-to-door assignment and active-CAN risk gating.""" + +from __future__ import annotations + +import json +import math +from collections import Counter, defaultdict, deque +from dataclasses import dataclass +from pathlib import Path +from statistics import median + + +DOORS = ("front", "rear") +POSE_FEATURES = { + "shoulder_center_x", + "shoulder_center_y", + "shoulder_width", + "torso_height", + "hip_center_y", +} + + +@dataclass(frozen=True) +class FeatureCalibration: + rear: float + front: float + weight: float + + def frontness(self, value: float) -> float: + span = self.front - self.rear + if span == 0: + return 0.5 + return min(1.0, max(0.0, (value - self.rear) / span)) + + +@dataclass(frozen=True) +class DoorZoneConfig: + features: dict[str, FeatureCalibration] + front_probability_min: float + rear_probability_max: float + minimum_person_height: float + minimum_pose_features: int + history_frames: int + stale_track_frames: int + source_path: Path | None = None + + @classmethod + def load(cls, path: Path) -> "DoorZoneConfig": + source = Path(path).expanduser().resolve() + data = json.loads(source.read_text(encoding="utf-8-sig")) + features = { + name: FeatureCalibration( + rear=float(values["rear"]), + front=float(values["front"]), + weight=float(values["weight"]), + ) + for name, values in data["features"].items() + } + return cls( + features=features, + front_probability_min=float(data.get("front_probability_min", 0.72)), + rear_probability_max=float(data.get("rear_probability_max", 0.28)), + minimum_person_height=float(data.get("minimum_person_height", 0.40)), + minimum_pose_features=int(data.get("minimum_pose_features", 4)), + history_frames=int(data.get("history_frames", 5)), + stale_track_frames=int(data.get("stale_track_frames", 120)), + source_path=source, + ) + + def classify(self, probability: float) -> str: + if probability >= self.front_probability_min: + return "front" + if probability <= self.rear_probability_max: + return "rear" + return "transition" + + +def _point(keypoints: list, index: int) -> tuple[float, float] | None: + if index >= len(keypoints) or len(keypoints[index]) < 2: + return None + x, y = float(keypoints[index][0]), float(keypoints[index][1]) + return (x, y) if math.isfinite(x) and math.isfinite(y) else None + + +def extract_person_features( + record: dict, + position: int, + frame_shape: tuple[int, int], +) -> tuple[dict[str, float], int]: + """Extract fixed-camera pose features, with box geometry as low-weight fallback.""" + + height, width = frame_shape + box = record.get("boxes", [])[position] + x1, y1, x2, y2 = [float(value) for value in box[:4]] + features = { + "box_center_x": ((x1 + x2) * 0.5) / max(width, 1), + "box_height": (y2 - y1) / max(height, 1), + } + + all_keypoints = record.get("temporal_selected_keypoints", []) + if position >= len(all_keypoints): + return features, 0 + keypoints = all_keypoints[position] + left_shoulder = _point(keypoints, 5) + right_shoulder = _point(keypoints, 6) + left_hip = _point(keypoints, 11) + right_hip = _point(keypoints, 12) + if None in (left_shoulder, right_shoulder, left_hip, right_hip): + return features, 0 + + shoulder_x = (left_shoulder[0] + right_shoulder[0]) * 0.5 + shoulder_y = (left_shoulder[1] + right_shoulder[1]) * 0.5 + hip_y = (left_hip[1] + right_hip[1]) * 0.5 + features.update( + { + "shoulder_center_x": shoulder_x / max(width, 1), + "shoulder_center_y": shoulder_y / max(height, 1), + "shoulder_width": math.hypot( + left_shoulder[0] - right_shoulder[0], + left_shoulder[1] - right_shoulder[1], + ) + / max(width, 1), + "torso_height": abs(hip_y - shoulder_y) / max(height, 1), + "hip_center_y": hip_y / max(height, 1), + } + ) + return features, len(POSE_FEATURES) + + +class DoorZoneAssigner: + """Fuse pose-dominant evidence and smooth front probability per track.""" + + def __init__(self, config: DoorZoneConfig): + self.config = config + self.history: dict[int, deque[float]] = defaultdict( + lambda: deque(maxlen=self.config.history_frames) + ) + self.last_seen: dict[int, int] = {} + self.zone_counts: Counter = Counter() + + def _score(self, features: dict[str, float]) -> tuple[float, dict[str, float]]: + weighted_sum = 0.0 + total_weight = 0.0 + contributions = {} + for name, calibration in self.config.features.items(): + if name not in features: + continue + frontness = calibration.frontness(float(features[name])) + contributions[name] = round(frontness, 6) + weighted_sum += frontness * calibration.weight + total_weight += calibration.weight + probability = weighted_sum / total_weight if total_weight else 0.5 + return probability, contributions + + def update( + self, + frame_index: int, + record: dict, + frame_shape: tuple[int, int], + ) -> list[dict]: + track_ids = [int(value) for value in record.get("track_ids", [])] + assignments = [] + for position, _box in enumerate(record.get("boxes", [])): + track_id = track_ids[position] if position < len(track_ids) else position + 1 + features, pose_feature_count = extract_person_features( + record, position, frame_shape + ) + probability, contributions = self._score(features) + enough_pose = pose_feature_count >= self.config.minimum_pose_features + if features["box_height"] < self.config.minimum_person_height: + raw_zone = "unassigned" + door_zone = "unassigned" + elif not enough_pose: + raw_zone = "transition" + door_zone = "transition" + else: + raw_zone = self.config.classify(probability) + values = self.history[track_id] + values.append(probability) + door_zone = self.config.classify(float(median(values))) + + values = self.history.get(track_id, ()) + smoothed_probability = ( + float(median(values)) if values else float(probability) + ) + self.last_seen[track_id] = int(frame_index) + self.zone_counts[door_zone] += 1 + assignments.append( + { + "track_id": track_id, + "person_index": position, + "door_zone": door_zone, + "raw_door_zone": raw_zone, + "front_probability": round(float(probability), 6), + "smoothed_front_probability": round(smoothed_probability, 6), + "confidence": round(abs(smoothed_probability - 0.5) * 2.0, 6), + "evidence_source": "pose_primary_bbox_secondary", + "pose_feature_count": pose_feature_count, + "features": {name: round(value, 6) for name, value in features.items()}, + "feature_frontness": contributions, + "history_frames": len(values), + } + ) + + expired = [ + track_id + for track_id, seen_at in self.last_seen.items() + if frame_index - seen_at > self.config.stale_track_frames + ] + for track_id in expired: + self.history.pop(track_id, None) + self.last_seen.pop(track_id, None) + return assignments + + def summary(self) -> dict: + return { + "policy": "pose-primary weighted front probability with causal median", + "config": str(self.config.source_path) if self.config.source_path else None, + "history_frames": self.config.history_frames, + "front_probability_min": self.config.front_probability_min, + "rear_probability_max": self.config.rear_probability_max, + "minimum_person_height": self.config.minimum_person_height, + "minimum_pose_features": self.config.minimum_pose_features, + "features": { + name: { + "rear": calibration.rear, + "front": calibration.front, + "weight": calibration.weight, + } + for name, calibration in self.config.features.items() + }, + "zone_counts": dict(self.zone_counts), + } + + +def _inactive_risk(door: str) -> dict: + return { + "door_id": door, + "status": "inactive", + "risk_score": None, + "risk_level": "inactive", + "decision": "DOOR_CAN_NOT_ACTIVE", + "risk_type": f"{door.upper()}_INACTIVE", + "source": None, + "track_id": None, + "side": None, + "regions": [], + } + + +def _clear_risk(door: str) -> dict: + return { + "door_id": door, + "status": "calculated", + "risk_score": 0.0, + "risk_level": "clear", + "decision": "CLEAR_NO_PERSON_ASSIGNED_TO_ACTIVE_DOOR", + "risk_type": f"{door.upper()}_CLEAR", + "source": None, + "track_id": None, + "side": None, + "regions": [], + } + + +def _unknown_risk(door: str, transition_regions: list[dict]) -> dict: + return { + "door_id": door, + "status": "door_zone_ambiguous", + "risk_score": None, + "risk_level": "unknown", + "decision": "EVIDENCE_UNAVAILABLE_TRANSITION_ZONE", + "risk_type": f"{door.upper()}_UNKNOWN", + "source": None, + "track_id": None, + "side": None, + "regions": transition_regions, + } + + +def build_door_risks( + raw_risk: dict, + assignments: list[dict], + active_door: str, +) -> tuple[dict[str, dict], dict]: + """Build per-door risk and select the CAN-active door result.""" + + zone_by_track = { + int(assignment["track_id"]): str(assignment["door_zone"]) + for assignment in assignments + } + annotated_regions = [] + for original in raw_risk.get("regions", []): + region = dict(original) + region["door_zone"] = zone_by_track.get(int(region.get("track_id", -1)), "unknown") + annotated_regions.append(region) + + active_doors = set(DOORS) if active_door == "both" else {active_door} + if active_door == "none": + active_doors.clear() + + door_risks: dict[str, dict] = {} + for door in DOORS: + if door not in active_doors: + door_risks[door] = _inactive_risk(door) + continue + exact = [region for region in annotated_regions if region["door_zone"] == door] + transition = [ + region for region in annotated_regions if region["door_zone"] == "transition" + ] + calculated = [region for region in exact if region.get("risk_score") is not None] + if calculated: + nearest = max(calculated, key=lambda item: float(item["risk_score"])) + level = str(nearest["risk_level"]) + if transition and level in ("clear", "watch"): + door_risks[door] = _unknown_risk(door, transition) + else: + door_risks[door] = { + "door_id": door, + "status": "calculated", + "risk_score": nearest["risk_score"], + "risk_level": level, + "decision": nearest["decision"], + "risk_type": f"{door.upper()}_{level.upper()}", + "source": nearest["source"], + "track_id": nearest["track_id"], + "side": nearest["side"], + "regions": exact, + } + elif transition: + door_risks[door] = _unknown_risk(door, transition) + else: + door_risks[door] = _clear_risk(door) + + if len(active_doors) == 1: + selected = dict(door_risks[next(iter(active_doors))]) + elif not active_doors: + selected = { + "door_id": None, + "status": "inactive", + "risk_score": None, + "risk_level": "inactive", + "decision": "NO_DOOR_CAN_ACTIVE", + "risk_type": "NO_ACTIVE_DOOR", + "source": None, + "track_id": None, + "side": None, + "regions": [], + } + else: + severity = {"danger": 4, "warning": 3, "watch": 2, "unknown": 1, "clear": 0} + selected = dict( + max(door_risks.values(), key=lambda item: severity.get(item["risk_level"], -1)) + ) + selected["door_mask"] = raw_risk.get("door_mask") + selected["warning_distance_relative"] = raw_risk.get("warning_distance_relative") + selected["release_distance_relative"] = raw_risk.get("release_distance_relative") + return door_risks, selected diff --git a/orin/pipeline/fuse_rtmw_l_temporal.py b/orin/pipeline/fuse_rtmw_l_temporal.py new file mode 100644 index 0000000..f212589 --- /dev/null +++ b/orin/pipeline/fuse_rtmw_l_temporal.py @@ -0,0 +1,855 @@ +"""Fuse raw and person-isolated RTMW-L poses with temporal consistency. + +The fusion keeps the two inference views as hypotheses. It selects body limbs +and hands separately, rejects implausible hand geometry, adds source hysteresis +and constant-velocity consistency, then applies a motion-adaptive EMA in +person-box coordinates. Occluded/conflicting limbs remain labelled as such so +they are not accidentally consumed as precise anti-pinch distance evidence. +""" + +from __future__ import annotations + +import argparse +import json +import math +from collections import Counter, defaultdict +from dataclasses import dataclass, field +from pathlib import Path + +import cv2 +import numpy as np + +from compare_rtmw_ml import ( + BODY_EDGES, + BODY_WRISTS, + HAND_EDGES_LOCAL, + HAND_SLICES, + draw_pose, + normalized_keypoints, + summarize, + temporal_metrics, +) + + +BODY_SIDE_INDICES = {"left": (5, 7, 9), "right": (6, 8, 10)} +SIDE_COLORS = {"left": (40, 40, 255), "right": (30, 230, 30)} +HAND_COLORS = {"left": (255, 80, 220), "right": (255, 190, 40)} + + +@dataclass +class Candidate: + source: str + keypoints: np.ndarray + scores: np.ndarray + normalized: np.ndarray + limb_score: float + hand_score: float + upper_length: float + forearm_length: float + hand_span: float + wrist_link: float + limb_valid: bool + hand_valid: bool + + +@dataclass +class SideTrack: + limb_previous: np.ndarray | None = None + limb_velocity: np.ndarray | None = None + hand_previous: np.ndarray | None = None + hand_velocity: np.ndarray | None = None + limb_source: str | None = None + hand_source: str | None = None + previous_hand_eligible: bool = False + + +@dataclass +class PersonTrack: + sides: dict[str, SideTrack] = field( + default_factory=lambda: {"left": SideTrack(), "right": SideTrack()} + ) + smoothed: np.ndarray | None = None + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser() + parser.add_argument("--video", type=Path, required=True) + parser.add_argument("--raw-jsonl", type=Path, required=True) + parser.add_argument("--isolated-jsonl", type=Path, required=True) + parser.add_argument("--output", type=Path, required=True) + return parser.parse_args() + + +def bbox_diagonal(bbox: list[float]) -> float: + return max(math.hypot(bbox[2] - bbox[0], bbox[3] - bbox[1]), 1.0) + + +def build_candidate( + source: str, + keypoints: np.ndarray, + scores: np.ndarray, + bbox: list[float], + side: str, +) -> Candidate: + normalized = normalized_keypoints(keypoints, bbox) + shoulder, elbow, wrist = BODY_SIDE_INDICES[side] + hand_start, hand_end = HAND_SLICES[side] + hand = keypoints[hand_start:hand_end] + diagonal = bbox_diagonal(bbox) + upper_length = float(np.linalg.norm(normalized[shoulder] - normalized[elbow])) + forearm_length = float(np.linalg.norm(normalized[elbow] - normalized[wrist])) + hand_span = math.hypot(float(np.ptp(hand[:, 0])), float(np.ptp(hand[:, 1]))) / diagonal + wrist_link = float(np.linalg.norm(keypoints[wrist] - hand[0])) / diagonal + limb_score = float(np.mean(scores[[shoulder, elbow, wrist]])) + hand_score = float(np.mean(scores[hand_start:hand_end])) + limb_valid = 0.025 <= upper_length <= 0.8 and 0.025 <= forearm_length <= 0.8 + hand_valid = 0.008 <= hand_span <= 0.35 and wrist_link <= 0.12 + return Candidate( + source=source, + keypoints=keypoints, + scores=scores, + normalized=normalized, + limb_score=limb_score, + hand_score=hand_score, + upper_length=upper_length, + forearm_length=forearm_length, + hand_span=hand_span, + wrist_link=wrist_link, + limb_valid=limb_valid, + hand_valid=hand_valid, + ) + + +def predicted(previous: np.ndarray | None, velocity: np.ndarray | None) -> np.ndarray | None: + if previous is None: + return None + if velocity is None: + return previous + return previous + 0.65 * velocity + + +def robust_distance(points: np.ndarray, target: np.ndarray | None) -> float: + if target is None: + return 0.0 + distances = np.linalg.norm(points - target, axis=1) + return float(np.median(distances)) + + +def select_limb( + candidates: list[Candidate], + side: str, + track: SideTrack | None, +) -> tuple[Candidate, dict[str, float]]: + indices = np.asarray(BODY_SIDE_INDICES[side]) + target = predicted(track.limb_previous, track.limb_velocity) if track else None + scored: list[tuple[float, Candidate, float]] = [] + for candidate in candidates: + temporal_cost = robust_distance(candidate.normalized[indices], target) + geometry_bonus = 1.2 if candidate.limb_valid else -5.0 + switch_bonus = 0.35 if track and candidate.source == track.limb_source else 0.0 + score = candidate.limb_score + geometry_bonus + switch_bonus - 7.0 * temporal_cost + scored.append((score, candidate, temporal_cost)) + score, selected, temporal_cost = max(scored, key=lambda item: item[0]) + return selected, {"score": score, "temporal_cost": temporal_cost} + + +def hand_link_to_limb( + candidate: Candidate, + selected_limb: Candidate, + bbox: list[float], + side: str, +) -> float: + hand_start, _ = HAND_SLICES[side] + wrist_idx = BODY_WRISTS[side] + return float( + np.linalg.norm(candidate.keypoints[hand_start] - selected_limb.keypoints[wrist_idx]) + ) / bbox_diagonal(bbox) + + +def select_hand( + candidates: list[Candidate], + selected_limb: Candidate, + bbox: list[float], + side: str, + track: SideTrack | None, +) -> tuple[Candidate, bool, dict[str, float]]: + start, end = HAND_SLICES[side] + target = predicted(track.hand_previous, track.hand_velocity) if track else None + scored: list[tuple[float, Candidate, bool, float, float]] = [] + for candidate in candidates: + temporal_cost = robust_distance(candidate.normalized[start:end], target) + link = hand_link_to_limb(candidate, selected_limb, bbox, side) + valid = candidate.hand_valid + geometry_bonus = 3.0 if valid else -7.0 + switch_bonus = 0.45 if track and candidate.source == track.hand_source else 0.0 + limb_source_bonus = 0.15 if candidate.source == selected_limb.source else 0.0 + # A newly visible endpoint may disagree with the currently selected + # (occlusion-driven) limb. Keep this as a soft cost so strong image + # evidence can switch both the hand and its limb together. + link_penalty = 1.5 * min(max(link - 0.12, 0.0), 0.2) + score = ( + candidate.hand_score + + geometry_bonus + + switch_bonus + + limb_source_bonus + - link_penalty + - 8.0 * min(temporal_cost, 0.15) + ) + scored.append((score, candidate, valid, temporal_cost, link)) + score, selected, valid, temporal_cost, link = max(scored, key=lambda item: item[0]) + return selected, valid, { + "score": score, + "temporal_cost": temporal_cost, + "link_to_selected_limb": link, + } + + +def candidate_disagreement( + raw: Candidate, + isolated: Candidate, + side: str, +) -> tuple[float, float]: + shoulder, elbow, wrist = BODY_SIDE_INDICES[side] + start, end = HAND_SLICES[side] + limb = float( + np.median( + np.linalg.norm( + raw.normalized[[shoulder, elbow, wrist]] + - isolated.normalized[[shoulder, elbow, wrist]], + axis=1, + ) + ) + ) + hand_distances = np.linalg.norm( + raw.normalized[start:end] - isolated.normalized[start:end], axis=1 + ) + hand = float(np.percentile(hand_distances, 90)) + return limb, hand + + +def suppress_duplicate_hands( + fused: np.ndarray, + metadata: dict, + bbox: list[float], +) -> None: + left_start, left_end = HAND_SLICES["left"] + right_start, right_end = HAND_SLICES["right"] + left_center = np.mean(fused[left_start:left_end], axis=0) + right_center = np.mean(fused[right_start:right_end], axis=0) + separation = float(np.linalg.norm(left_center - right_center)) / bbox_diagonal(bbox) + if separation >= 0.08: + return + weaker = min( + ("left", "right"), + key=lambda side: metadata[side]["selected_hand_score"], + ) + metadata[weaker]["hand_state"] = "duplicate" + metadata[weaker]["distance_eligible"] = False + metadata[weaker]["alarm_eligible"] = False + metadata[weaker]["duplicate_hand_separation"] = round(separation, 5) + + +def fuse_person( + raw_keypoints: np.ndarray, + raw_scores: np.ndarray, + isolated_keypoints: np.ndarray, + isolated_scores: np.ndarray, + bbox: list[float], + track: PersonTrack | None, +) -> tuple[np.ndarray, dict]: + fused = raw_keypoints.copy() + metadata: dict[str, dict] = {} + + for side in ("left", "right"): + side_track = track.sides[side] if track else None + raw = build_candidate("raw", raw_keypoints, raw_scores, bbox, side) + isolated = build_candidate( + "isolated", isolated_keypoints, isolated_scores, bbox, side + ) + candidates = [raw, isolated] + limb, limb_detail = select_limb(candidates, side, side_track) + hand, hand_valid, hand_detail = select_hand( + candidates, limb, bbox, side, side_track + ) + + other_hand = isolated if hand.source == "raw" else raw + endpoint_advantage = hand.hand_score - other_hand.hand_score + if hand.source != limb.source and hand_valid and endpoint_advantage >= 0.6: + limb = hand + limb_detail["visible_endpoint_override"] = 1.0 + hand_detail["link_to_selected_limb"] = hand.wrist_link + + body_indices = np.asarray(BODY_SIDE_INDICES[side]) + hand_start, hand_end = HAND_SLICES[side] + fused[body_indices] = limb.keypoints[body_indices] + fused[hand_start:hand_end] = hand.keypoints[hand_start:hand_end] + + limb_disagreement, hand_disagreement = candidate_disagreement(raw, isolated, side) + if limb_disagreement <= 0.05: + limb_state = "consensus" + else: + limb_state = "inferred" + + score_advantage = raw.hand_score - isolated.hand_score + raw_temporally_stable = hand_detail["temporal_cost"] <= 0.08 + if not hand_valid: + hand_state = "rejected" + distance_eligible = False + elif hand_disagreement <= 0.06: + hand_state = "consensus" + distance_eligible = True + elif ( + hand.source == "raw" + and score_advantage >= 0.6 + and raw_temporally_stable + ): + hand_state = "visible_endpoint" + distance_eligible = True + elif hand.source == "isolated": + hand_state = "inferred" + distance_eligible = False + else: + hand_state = "conflict" + distance_eligible = False + fast_alarm_endpoint = ( + hand_valid and hand.source == "raw" and score_advantage >= 0.6 + ) + alarm_eligible = distance_eligible or fast_alarm_endpoint + + metadata[side] = { + "limb_source": limb.source, + "hand_source": hand.source, + "limb_state": limb_state, + "hand_state": hand_state, + "distance_eligible": distance_eligible, + "alarm_eligible": alarm_eligible, + "limb_disagreement": round(limb_disagreement, 5), + "hand_disagreement": round(hand_disagreement, 5), + "raw_hand_score": round(raw.hand_score, 4), + "isolated_hand_score": round(isolated.hand_score, 4), + "selected_hand_span": round(hand.hand_span, 5), + "selected_hand_score": round(hand.hand_score, 4), + "selected_hand_valid": hand_valid, + "limb_selection": {key: round(value, 5) for key, value in limb_detail.items()}, + "hand_selection": {key: round(value, 5) for key, value in hand_detail.items()}, + } + + if side_track is not None: + limb_now = limb.normalized[body_indices] + hand_now = hand.normalized[hand_start:hand_end] + if side_track.limb_previous is not None: + delta = limb_now - side_track.limb_previous + side_track.limb_velocity = ( + delta + if side_track.limb_velocity is None + else 0.65 * side_track.limb_velocity + 0.35 * delta + ) + if hand_valid: + if side_track.hand_previous is not None: + delta = hand_now - side_track.hand_previous + side_track.hand_velocity = ( + delta + if side_track.hand_velocity is None + else 0.65 * side_track.hand_velocity + 0.35 * delta + ) + side_track.hand_previous = hand_now + side_track.hand_source = hand.source + side_track.limb_previous = limb_now + side_track.limb_source = limb.source + side_track.previous_hand_eligible = distance_eligible + + suppress_duplicate_hands(fused, metadata, bbox) + return fused, metadata + + +def adaptive_smooth( + keypoints: np.ndarray, + bbox: list[float], + track: PersonTrack, +) -> np.ndarray: + normalized = normalized_keypoints(keypoints, bbox) + if track.smoothed is None: + track.smoothed = normalized + else: + displacement = np.linalg.norm(normalized - track.smoothed, axis=1) + alpha = 0.32 + 0.58 * np.clip(displacement / 0.08, 0.0, 1.0) + track.smoothed = alpha[:, None] * normalized + (1.0 - alpha[:, None]) * track.smoothed + + x1, y1, x2, y2 = bbox + result = track.smoothed.copy() + result[:, 0] = result[:, 0] * max(x2 - x1, 1.0) + x1 + result[:, 1] = result[:, 1] * max(y2 - y1, 1.0) + y1 + return result + + +def build_global_hands( + keypoints: list[np.ndarray], + metadata: list[dict], +) -> list[dict]: + candidates = [] + for person_idx, (person_keypoints, person_meta) in enumerate( + zip(keypoints, metadata), start=1 + ): + for side, (start, end) in HAND_SLICES.items(): + meta = person_meta[side] + if not meta["alarm_eligible"]: + continue + hand = person_keypoints[start:end] + span_px = math.hypot(float(np.ptp(hand[:, 0])), float(np.ptp(hand[:, 1]))) + candidates.append( + { + "person_id": person_idx, + "side": side, + "source": meta["hand_source"], + "state": meta["hand_state"], + "distance_eligible": meta["distance_eligible"], + "alarm_eligible": meta["alarm_eligible"], + "score": meta["selected_hand_score"], + "center": np.mean(hand, axis=0), + "span_px": span_px, + "keypoints": hand, + "aliases": [], + } + ) + + kept: list[dict] = [] + for candidate in sorted(candidates, key=lambda item: item["score"], reverse=True): + duplicate = None + for existing in kept: + threshold = float(np.clip(0.45 * max(candidate["span_px"], existing["span_px"]), 18, 70)) + if float(np.linalg.norm(candidate["center"] - existing["center"])) <= threshold: + duplicate = existing + break + if duplicate is None: + kept.append(candidate) + else: + duplicate["aliases"].append( + {"person_id": candidate["person_id"], "side": candidate["side"]} + ) + + return [ + { + "person_id": item["person_id"], + "side": item["side"], + "source": item["source"], + "state": item["state"], + "distance_eligible": item["distance_eligible"], + "alarm_eligible": item["alarm_eligible"], + "score": item["score"], + "center": np.round(item["center"], 3).tolist(), + "span_px": round(item["span_px"], 3), + "keypoints": np.round(item["keypoints"], 3).tolist(), + "aliases": item["aliases"], + } + for item in kept + ] + + +def draw_dashed_line( + image: np.ndarray, + start: np.ndarray, + end: np.ndarray, + color: tuple[int, int, int], + thickness: int = 2, +) -> None: + a = np.asarray(start, dtype=np.float64) + b = np.asarray(end, dtype=np.float64) + length = float(np.linalg.norm(b - a)) + if length < 1.0: + return + direction = (b - a) / length + for offset in np.arange(0.0, length, 14.0): + segment_start = a + direction * offset + segment_end = a + direction * min(offset + 8.0, length) + cv2.line( + image, + tuple(np.round(segment_start).astype(int)), + tuple(np.round(segment_end).astype(int)), + color, + thickness, + cv2.LINE_AA, + ) + + +def draw_fused( + image: np.ndarray, + boxes: list[list[float]], + keypoints: list[np.ndarray], + metadata: list[dict], + title: str, +) -> np.ndarray: + canvas = image.copy() + for person_idx, (bbox, kpts, person_meta) in enumerate( + zip(boxes, keypoints, metadata), start=1 + ): + x1, y1, x2, y2 = (int(round(value)) for value in bbox) + cv2.rectangle(canvas, (x1, y1), (x2, y2), (0, 220, 255), 2) + cv2.putText( + canvas, + f"P{person_idx}", + (x1 + 4, max(22, y1 + 24)), + cv2.FONT_HERSHEY_SIMPLEX, + 0.72, + (0, 220, 255), + 2, + cv2.LINE_AA, + ) + for side in ("left", "right"): + shoulder, elbow, wrist = BODY_SIDE_INDICES[side] + side_meta = person_meta[side] + color = SIDE_COLORS[side] + for a, b in ((shoulder, elbow), (elbow, wrist)): + if side_meta["limb_state"] == "consensus": + cv2.line( + canvas, + tuple(np.round(kpts[a]).astype(int)), + tuple(np.round(kpts[b]).astype(int)), + color, + 3, + cv2.LINE_AA, + ) + else: + draw_dashed_line(canvas, kpts[a], kpts[b], (0, 140, 255), 3) + + start, end = HAND_SLICES[side] + hand = kpts[start:end] + if side_meta["hand_state"] != "rejected": + hand_color = HAND_COLORS[side] + if not side_meta["distance_eligible"]: + hand_color = (0, 140, 255) + for a, b in HAND_EDGES_LOCAL: + if side_meta["distance_eligible"]: + cv2.line( + canvas, + tuple(np.round(hand[a]).astype(int)), + tuple(np.round(hand[b]).astype(int)), + hand_color, + 2, + cv2.LINE_AA, + ) + else: + draw_dashed_line(canvas, hand[a], hand[b], hand_color, 2) + label = ( + f"{side[0].upper()}:{side_meta['hand_state'][:4]}" + f"/{side_meta['hand_source'][0]}" + ) + label_y = max(45, y1 + (48 if side == "left" else 70)) + cv2.putText( + canvas, + label, + (x1 + 4, label_y), + cv2.FONT_HERSHEY_SIMPLEX, + 0.5, + HAND_COLORS[side] if side_meta["distance_eligible"] else (0, 140, 255), + 1, + cv2.LINE_AA, + ) + + cv2.rectangle(canvas, (0, 0), (650, 45), (0, 0, 0), -1) + cv2.putText( + canvas, + title, + (10, 31), + cv2.FONT_HERSHEY_SIMPLEX, + 0.8, + (255, 255, 255), + 2, + cv2.LINE_AA, + ) + return canvas + + +def make_montage( + video_path: Path, + records: list[dict], + selected_frames: list[int], + output_path: Path, +) -> None: + selected = set(selected_frames) + by_frame = {record["frame_index"]: record for record in records} + cap = cv2.VideoCapture(str(video_path)) + rows = [] + frame_idx = 0 + while True: + ok, frame = cap.read() + if not ok: + break + if frame_idx in selected: + record = by_frame[frame_idx] + boxes = record["boxes"] + raw = draw_pose( + frame, + boxes, + np.asarray(record["raw_keypoints"]), + f"RTMW-L raw | frame {frame_idx}", + ) + spatial = draw_fused( + frame, + boxes, + [np.asarray(points) for points in record["spatial_keypoints"]], + record["spatial_metadata"], + f"spatial fusion | frame {frame_idx}", + ) + temporal = draw_fused( + frame, + boxes, + [np.asarray(points) for points in record["temporal_keypoints"]], + record["temporal_metadata"], + f"temporal fusion | frame {frame_idx}", + ) + panels = [ + cv2.resize(panel, (480, 270), interpolation=cv2.INTER_AREA) + for panel in (raw, spatial, temporal) + ] + rows.append(np.hstack(panels)) + frame_idx += 1 + cap.release() + if rows: + cv2.imwrite(str(output_path), np.vstack(rows)) + + +def active_hand_temporal_metrics(records: list[dict], variant: str) -> dict: + values: dict[str, list[float]] = defaultdict(list) + previous: dict[tuple[int, str], tuple[int, np.ndarray]] = {} + for record in records: + for person_idx, bbox in enumerate(record["boxes"]): + kpts = normalized_keypoints( + np.asarray(record[f"{variant}_keypoints"][person_idx]), bbox + ) + meta = record[f"{variant}_metadata"][person_idx] + for side, (start, end) in HAND_SLICES.items(): + key = (person_idx, side) + if not meta[side]["distance_eligible"]: + previous.pop(key, None) + continue + prior = previous.get(key) + if prior is not None and prior[0] == record["frame_index"] - 1: + step = np.linalg.norm(kpts[start:end] - prior[1], axis=1) + values[f"P{person_idx + 1}_{side}"].extend(float(v) for v in step) + previous[key] = (record["frame_index"], kpts[start:end]) + return {name: summarize(items) for name, items in values.items()} + + +def metadata_summary(records: list[dict], variant: str) -> dict: + state_counts: Counter = Counter() + source_counts: Counter = Counter() + eligible_counts: Counter = Counter() + alarm_counts: Counter = Counter() + switches: Counter = Counter() + previous_source: dict[tuple[int, str, str], str] = {} + for record in records: + for person_idx, person_meta in enumerate(record[f"{variant}_metadata"]): + for side in ("left", "right"): + meta = person_meta[side] + key_prefix = f"P{person_idx + 1}_{side}" + state_counts[f"{key_prefix}_{meta['hand_state']}"] += 1 + source_counts[f"{key_prefix}_{meta['hand_source']}"] += 1 + if meta["distance_eligible"]: + eligible_counts[key_prefix] += 1 + if meta["alarm_eligible"]: + alarm_counts[key_prefix] += 1 + for kind in ("limb", "hand"): + source = meta[f"{kind}_source"] + key = (person_idx, side, kind) + if key in previous_source and previous_source[key] != source: + switches[f"{key_prefix}_{kind}"] += 1 + previous_source[key] = source + return { + "hand_state_counts": dict(sorted(state_counts.items())), + "source_counts": dict(sorted(source_counts.items())), + "distance_eligible_frames": dict(sorted(eligible_counts.items())), + "alarm_eligible_frames": dict(sorted(alarm_counts.items())), + "source_switches": dict(sorted(switches.items())), + } + + +def main() -> None: + args = parse_args() + args.output.mkdir(parents=True, exist_ok=True) + raw_records = [ + json.loads(line) for line in args.raw_jsonl.read_text(encoding="utf-8").splitlines() + ] + isolated_records = [ + json.loads(line) + for line in args.isolated_jsonl.read_text(encoding="utf-8").splitlines() + ] + + cap = cv2.VideoCapture(str(args.video)) + fps = float(cap.get(cv2.CAP_PROP_FPS)) + width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)) + height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)) + writer = cv2.VideoWriter( + str(args.output / "rtmw_l_raw_spatial_temporal.mp4"), + cv2.VideoWriter_fourcc(*"mp4v"), + fps, + (1920, 360), + ) + + tracks = [PersonTrack(), PersonTrack()] + output_records: list[dict] = [] + frame_idx = 0 + while True: + ok, frame = cap.read() + if not ok: + break + raw_record = raw_records[frame_idx] + isolated_record = isolated_records[frame_idx] + boxes = raw_record["boxes"][:2] + raw_keypoints = np.asarray(raw_record["L"]["keypoints"][:2], dtype=np.float64) + raw_scores = np.asarray(raw_record["L"]["raw_scores"][:2], dtype=np.float64) + isolated_keypoints = np.asarray( + isolated_record["L"]["keypoints"][:2], dtype=np.float64 + ) + isolated_scores = np.asarray( + isolated_record["masked_raw_scores"][:2], dtype=np.float64 + ) + + spatial_keypoints: list[np.ndarray] = [] + spatial_metadata: list[dict] = [] + temporal_keypoints: list[np.ndarray] = [] + temporal_selected_keypoints: list[np.ndarray] = [] + temporal_metadata: list[dict] = [] + for person_idx in range(2): + spatial, spatial_meta = fuse_person( + raw_keypoints[person_idx], + raw_scores[person_idx], + isolated_keypoints[person_idx], + isolated_scores[person_idx], + boxes[person_idx], + track=None, + ) + temporal_unfiltered, temporal_meta = fuse_person( + raw_keypoints[person_idx], + raw_scores[person_idx], + isolated_keypoints[person_idx], + isolated_scores[person_idx], + boxes[person_idx], + track=tracks[person_idx], + ) + temporal = adaptive_smooth( + temporal_unfiltered, boxes[person_idx], tracks[person_idx] + ) + spatial_keypoints.append(spatial) + spatial_metadata.append(spatial_meta) + temporal_keypoints.append(temporal) + temporal_selected_keypoints.append(temporal_unfiltered) + temporal_metadata.append(temporal_meta) + + global_hands = build_global_hands( + temporal_selected_keypoints, temporal_metadata + ) + + output_record = { + "frame_index": frame_idx, + "boxes": boxes, + "raw_keypoints": np.round(raw_keypoints, 3).tolist(), + "spatial_keypoints": np.round(np.asarray(spatial_keypoints), 3).tolist(), + "spatial_metadata": spatial_metadata, + "temporal_keypoints": np.round(np.asarray(temporal_keypoints), 3).tolist(), + "temporal_selected_keypoints": np.round( + np.asarray(temporal_selected_keypoints), 3 + ).tolist(), + "temporal_metadata": temporal_metadata, + "global_hands": global_hands, + } + output_records.append(output_record) + + raw_panel = draw_pose(frame, boxes, raw_keypoints, f"RTMW-L raw | frame {frame_idx}") + spatial_panel = draw_fused( + frame, + boxes, + spatial_keypoints, + spatial_metadata, + f"spatial fusion | frame {frame_idx}", + ) + temporal_panel = draw_fused( + frame, + boxes, + temporal_keypoints, + temporal_metadata, + f"temporal fusion | frame {frame_idx}", + ) + panels = [ + cv2.resize(panel, (640, 360), interpolation=cv2.INTER_AREA) + for panel in (raw_panel, spatial_panel, temporal_panel) + ] + writer.write(np.hstack(panels)) + frame_idx += 1 + + cap.release() + writer.release() + + with (args.output / "fusion.jsonl").open("w", encoding="utf-8") as handle: + for record in output_records: + handle.write(json.dumps(record, ensure_ascii=False) + "\n") + + def comparison_records(variant: str) -> list[dict]: + return [ + { + "frame_index": record["frame_index"], + "boxes": record["boxes"], + "M": {"keypoints": record["raw_keypoints"]}, + "L": {"keypoints": record[f"{variant}_keypoints"]}, + } + for record in output_records + ] + + summary = { + "frames": len(output_records), + "fusion_policy": { + "candidates": ["RTMW-L raw", "RTMW-L smaller-person isolated"], + "selection": [ + "limb and hand selected separately", + "geometry validity", + "constant-velocity temporal consistency", + "source-switch hysteresis", + "motion-adaptive EMA in person-box coordinates", + ], + "distance_rule": ( + "Only consensus or temporally stable raw visible endpoints are eligible. " + "Inferred/conflict/rejected hands are excluded from precise distance." + ), + "alarm_rule": ( + "Alarm attack is asymmetric: a geometrically valid raw endpoint whose score " + "is clearly stronger than the isolated view is admitted immediately. Alarm " + "release should still require temporal confirmation and hysteresis." + ), + "risk_distance_note": ( + "Use global_hands and temporal_selected_keypoints (unfiltered) for alarm entry " + "to avoid smoothing delay. Use distance_eligible hands and temporal_keypoints " + "for stable reporting. When the door edge is available, a conservative " + "implementation can use the smaller of the unfiltered and smoothed distances." + ), + }, + "spatial": metadata_summary(output_records, "spatial"), + "temporal": metadata_summary(output_records, "temporal"), + "active_hand_frame_step": { + "spatial": active_hand_temporal_metrics(output_records, "spatial"), + "temporal": active_hand_temporal_metrics(output_records, "temporal"), + }, + "raw_vs_spatial_temporal_metrics": temporal_metrics(comparison_records("spatial")), + "raw_vs_temporal_temporal_metrics": temporal_metrics(comparison_records("temporal")), + "global_hand_count_frames": dict( + sorted(Counter(len(record["global_hands"]) for record in output_records).items()) + ), + "cross_person_hand_aliases_total": sum( + len(hand["aliases"]) + for record in output_records + for hand in record["global_hands"] + ), + } + with (args.output / "summary.json").open("w", encoding="utf-8") as handle: + json.dump(summary, handle, ensure_ascii=False, indent=2) + + key_frames = [0, 45, 50, 55, 120, 132, 136, 180, 240, 252, 295, 300, 304, 360, 411] + make_montage( + args.video, + output_records, + key_frames, + args.output / "rtmw_l_raw_spatial_temporal_montage.jpg", + ) + print(json.dumps(summary, ensure_ascii=False, indent=2)) + + +if __name__ == "__main__": + main() diff --git a/orin/pipeline/latest_result.py b/orin/pipeline/latest_result.py new file mode 100644 index 0000000..1ab68ae --- /dev/null +++ b/orin/pipeline/latest_result.py @@ -0,0 +1,64 @@ +"""One owned, immutable latest result, mirroring the 0920 LatestResult contract.""" +from __future__ import annotations + +from collections.abc import Mapping +import threading +from types import MappingProxyType + + +def freeze(value): + if isinstance(value, Mapping): + return MappingProxyType({key: freeze(item) for key, item in value.items()}) + if isinstance(value, (list, tuple)): + return tuple(freeze(item) for item in value) + return value + + +def result_to_dict(value): + """Return a mutable JSON-compatible copy when a consumer needs serialization.""" + if isinstance(value, Mapping): + return {key: result_to_dict(item) for key, item in value.items()} + if isinstance(value, tuple): + return [result_to_dict(item) for item in value] + return value + + +class LatestResult: + """Readers share a stable snapshot; reading never removes or acknowledges it. + + Only the latest result is retained here. Slow readers can skip frames, and + references they already hold remain valid after subsequent publications. + """ + + def __init__(self): + self._condition = threading.Condition() + self._value = None + self._closed = False + + def publish(self, result): + if result is None: + raise ValueError("Cannot publish a null result") + snapshot = freeze(result) + with self._condition: + if self._closed: + raise RuntimeError("Result publisher is closed") + self._value = snapshot + self._condition.notify_all() + + def load(self): + with self._condition: + return self._value + + def wait(self, after_frame_index=-1, timeout=None): + """Wait for a newer result; return None on timeout/end without a new frame.""" + with self._condition: + def has_new_frame(): + return self._value is not None and self._value['frame_index'] > after_frame_index + + self._condition.wait_for(lambda: has_new_frame() or self._closed, timeout) + return self._value if has_new_frame() else None + + def close(self): + with self._condition: + self._closed = True + self._condition.notify_all() diff --git a/orin/pipeline/live_pose_engine.py b/orin/pipeline/live_pose_engine.py new file mode 100644 index 0000000..b0140e5 --- /dev/null +++ b/orin/pipeline/live_pose_engine.py @@ -0,0 +1,248 @@ +"""Resident YOLO11s + RTMW-L pose stage extracted from the v9.9 pipeline.""" + +from __future__ import annotations + +import argparse +import os +import time +from collections import defaultdict +from pathlib import Path + +import numpy as np + + +SIDES = ("left", "right") + + +class LivePoseEngine: + """Run the v9.9 YOLO/RTMW pose and temporal-fusion path without EdgeTAM.""" + + def __init__(self, args: argparse.Namespace): + self._trt_dll_directory = None + if args.pose_provider == "tensorrt" and os.name == "nt": + try: + import tensorrt_libs + except ImportError as exc: + raise RuntimeError( + "--pose-provider tensorrt on Windows requires the " + "tensorrt-cu12-libs package in the active Python environment." + ) from exc + + trt_lib_dir = Path(tensorrt_libs.__file__).resolve().parent + self._trt_dll_directory = os.add_dll_directory(str(trt_lib_dir)) + os.environ["PATH"] = ( + str(trt_lib_dir) + os.pathsep + os.environ.get("PATH", "") + ) + + if args.pose_device.lower().startswith("cuda"): + import onnxruntime as ort + + if "CUDAExecutionProvider" not in ort.get_available_providers(): + raise RuntimeError( + "--pose-device cuda requires onnxruntime-gpu. Uninstall " + "onnxruntime and install onnxruntime-gpu==1.22.0." + ) + + from rtmlib import RTMPose + from ultralytics import YOLO + from predict_optimized_pose_videos import BoxTracker + + self.args = args + self.detector = YOLO(str(args.det_model)) + self.pose = RTMPose( + str(args.pose_model), + model_input_size=(288, 384), + backend="onnxruntime", + device=args.pose_device, + ) + if args.pose_provider == "tensorrt": + import onnxruntime as ort + + available = ort.get_available_providers() + if "TensorrtExecutionProvider" not in available: + raise RuntimeError( + "--pose-provider tensorrt requires ONNX Runtime with " + "TensorrtExecutionProvider. " + f"Available providers: {available}" + ) + cache_dir = args.pose_trt_cache.expanduser().resolve() + cache_dir.mkdir(parents=True, exist_ok=True) + device_id = ( + int(str(args.pose_device).split(":")[-1]) + if ":" in str(args.pose_device) + else 0 + ) + trt_options = { + "device_id": device_id, + "trt_fp16_enable": "True", + "trt_engine_cache_enable": "True", + "trt_engine_cache_path": str(cache_dir), + "trt_timing_cache_enable": "True", + "trt_timing_cache_path": str(cache_dir), + } + self.pose.session = ort.InferenceSession( + str(args.pose_model.resolve()), + providers=[ + ("TensorrtExecutionProvider", trt_options), + ("CUDAExecutionProvider", {"device_id": device_id}), + ], + ) + active_providers = self.pose.session.get_providers() + if not active_providers or active_providers[0] != "TensorrtExecutionProvider": + raise RuntimeError( + "RTMW TensorRT initialization failed instead of using the " + "requested provider. Active providers: " + f"{active_providers}" + ) + print(f"RTMW ONNX Runtime providers: {active_providers}") + + self.box_tracker = BoxTracker(max_age=30) + self.people = {} + self.timings = defaultdict(list) + self.isolated_inferences = 0 + + def infer(self, frame: np.ndarray, frame_idx: int) -> dict: + from compare_rtmw_ml import HAND_SLICES + from fuse_rtmw_l_temporal import adaptive_smooth, build_global_hands, fuse_person + from occlusion_aware_pose_v2 import Measurement, resolve_global_hand_assignments + from predict_optimized_pose_videos import ( + LivePersonState, + filter_person_boxes, + mask_other_people, + update_live_v2, + ) + + infer_started = time.perf_counter() + height, width = frame.shape[:2] + started = time.perf_counter() + detection_height = max( + 1, min(height, round(height * self.args.det_roi_y_max_ratio)) + ) + if hasattr(self, "native"): + boxes, confidences = self.native.detect(frame[:detection_height]) + self.timings["detector_ms"].append((time.perf_counter() - started) * 1000) + else: + detection = self.detector.predict( + frame[:detection_height], + imgsz=self.args.det_imgsz, + conf=self.args.det_conf, + iou=0.30, + classes=[0], + device=self.args.det_device, + verbose=False, + )[0] + self.timings["detector_ms"].append((time.perf_counter() - started) * 1000) + boxes, confidences, _ = filter_person_boxes( + detection, width, height, self.args.max_people + ) + track_ids = self.box_tracker.update(boxes, frame_idx, (height, width)) + + # Retire disappeared people in long-running camera sessions. + for stale_id in list(self.people): + if stale_id not in self.box_tracker.tracks: + del self.people[stale_id] + if boxes: + started = time.perf_counter() + raw_keypoints, raw_scores = self.pose(frame, bboxes=boxes) + self.timings["raw_pose_ms"].append((time.perf_counter() - started) * 1000) + isolated_keypoints = raw_keypoints.copy() + isolated_scores = raw_scores.copy() + started = time.perf_counter() + for person_idx, bbox in enumerate(boxes): + target_area = max((bbox[2] - bbox[0]) * (bbox[3] - bbox[1]), 1.0) + needs_isolation = any( + idx != person_idx + and max((other[2] - other[0]) * (other[3] - other[1]), 1.0) + < 0.5 * target_area + for idx, other in enumerate(boxes) + ) + if not needs_isolation: + continue + masked = mask_other_people(frame, boxes, person_idx) + keypoints, scores = self.pose(masked, bboxes=[bbox]) + isolated_keypoints[person_idx] = keypoints[0] + isolated_scores[person_idx] = scores[0] + self.isolated_inferences += 1 + self.timings["isolated_extra_ms"].append( + (time.perf_counter() - started) * 1000 + ) + else: + raw_keypoints = np.empty((0, 133, 2), dtype=np.float64) + raw_scores = np.empty((0, 133), dtype=np.float64) + isolated_keypoints = raw_keypoints.copy() + isolated_scores = raw_scores.copy() + + selected_list, smoothed_list, metadata_list, states = [], [], [], [] + for person_idx, (bbox, track_id) in enumerate(zip(boxes, track_ids)): + state = self.people.setdefault(track_id, LivePersonState()) + selected, metadata = fuse_person( + raw_keypoints[person_idx], + raw_scores[person_idx], + isolated_keypoints[person_idx], + isolated_scores[person_idx], + bbox, + state.fusion, + ) + selected_list.append(selected) + smoothed_list.append(adaptive_smooth(selected, bbox, state.fusion)) + metadata_list.append(metadata) + states.append(state) + + measurements, endpoint_tracks = {}, {} + for person_idx, (selected, metadata, state) in enumerate( + zip(selected_list, metadata_list, states) + ): + for side, (start, end) in HAND_SLICES.items(): + meta = metadata[side] + measurements[(person_idx, side)] = ( + Measurement( + keypoints=selected[start:end], + precise=bool(meta["distance_eligible"]), + provisional=not bool(meta["distance_eligible"]), + quality=float(meta.get("selected_hand_score", 0.0)), + ) + if meta["alarm_eligible"] + else None + ) + endpoint_tracks[(person_idx, side)] = state.endpoints[side] + measurements, assignments = resolve_global_hand_assignments( + measurements, endpoint_tracks, boxes, 30 + ) + + v2_list, v2_meta_list = [], [] + for person_idx, (bbox, selected, smoothed, metadata, state) in enumerate( + zip(boxes, selected_list, smoothed_list, metadata_list, states) + ): + v2, v2_meta = update_live_v2( + frame_idx, + bbox, + selected, + smoothed, + metadata, + state, + {side: measurements[(person_idx, side)] for side in SIDES}, + {side: assignments.get((person_idx, side)) for side in SIDES}, + 12, + 30, + ) + v2_list.append(v2) + v2_meta_list.append(v2_meta) + + self.timings["pose_pipeline_total_ms"].append( + (time.perf_counter() - infer_started) * 1000 + ) + return { + "frame_index": frame_idx, + "track_ids": track_ids, + "boxes": np.round(np.asarray(boxes), 3).tolist() if boxes else [], + "det_confidences": [round(value, 4) for value in confidences], + "raw_keypoints": np.round(raw_keypoints, 3).tolist(), + "temporal_selected_keypoints": np.round( + np.asarray(selected_list), 3 + ).tolist(), + "temporal_keypoints": np.round(np.asarray(smoothed_list), 3).tolist(), + "temporal_metadata": metadata_list, + "v2_keypoints": np.round(np.asarray(v2_list), 3).tolist(), + "v2_metadata": v2_meta_list, + "global_hands": build_global_hands(selected_list, metadata_list), + } diff --git a/orin/pipeline/native_backend.py b/orin/pipeline/native_backend.py new file mode 100644 index 0000000..91895f5 --- /dev/null +++ b/orin/pipeline/native_backend.py @@ -0,0 +1,75 @@ +"""Native Jetson TensorRT adapter; temporal/risk policy stays in upstream Python.""" +from __future__ import annotations +import ctypes as C +from collections import defaultdict, deque +from pathlib import Path +import numpy as np +from live_pose_engine import LivePoseEngine +from predict_optimized_pose_videos import BoxTracker + +ROOT = Path(__file__).resolve().parents[1] + +class NativeBackend: + def __init__(self, args): + self.lib = C.CDLL(str(ROOT / 'build/libvision_native.so')) + pointer = C.c_void_p + signatures = { + 'vap_error': ([], C.c_char_p), + 'vap_engine': ([C.c_char_p,C.c_int,C.c_int,C.c_int,C.c_float],pointer), + 'vap_engine_close': ([pointer],None), + 'vap_detect': ([pointer,pointer,C.c_int,C.c_int,pointer],C.c_int), + 'vap_pose': ([pointer,pointer,C.c_int,C.c_int,pointer,C.c_int,pointer],C.c_int), + 'vap_camera': ([C.c_char_p,C.c_int,C.c_int,C.c_int,C.c_int],pointer), + 'vap_camera_close': ([pointer],None), + 'vap_read': ([pointer,pointer,pointer],C.c_int), + 'vap_can': ([pointer],None), + } + for name,(inputs,output) in signatures.items(): + fn=getattr(self.lib,name);fn.argtypes=inputs;fn.restype=output + self.engine = self.lib.vap_engine(str(args.engines.resolve()).encode(),args.width,args.height,args.max_people,args.det_conf) + if not self.engine: self.fail() + self.camera=None + self.args=args + + def fail(self): + raise RuntimeError(self.lib.vap_error().decode()) + + def open_camera(self): + a=self.args + self.camera=self.lib.vap_camera(a.camera.encode(),a.width,a.height,a.fps,not a.no_can) + if not self.camera:self.fail() + + def read(self): + image=np.empty((self.args.height,self.args.width,3),np.uint8) + meta=np.zeros(4,np.float64) + result=self.lib.vap_read(self.camera,image.ctypes.data,meta.ctypes.data) + if result<0:self.fail() + if result==0:return None,None + return image,{'sequence':int(meta[0]),'capture_monotonic_ms':float(meta[1]),'dropped_before_inference':int(meta[2]),'driver_timestamp_ms':float(meta[3]),'latency_clock':'host_monotonic_at_dequeue'} + + def can_snapshot(self): + data=np.zeros(5,np.float64);self.lib.vap_can(data.ctypes.data) + return {'positions':[int(x) for x in data[:4]],'age_ms':float(data[4])} + + def detect(self, frame): + frame=np.ascontiguousarray(frame);out=np.empty((self.args.max_people,5),np.float32) + n=self.lib.vap_detect(self.engine,frame.ctypes.data,frame.shape[1],frame.shape[0],out.ctypes.data) + if n<0:self.fail() + return out[:n,:4].tolist(),out[:n,4].tolist() + + def pose(self,frame,bboxes): + frame=np.ascontiguousarray(frame);boxes=np.ascontiguousarray(bboxes,dtype=np.float32) + out=np.empty((len(boxes),133,3),np.float32) + n=self.lib.vap_pose(self.engine,frame.ctypes.data,frame.shape[1],frame.shape[0],boxes.ctypes.data,len(boxes),out.ctypes.data) + if n<0:self.fail() + return out[:n,:,:2].astype(np.float64),out[:n,:,2].astype(np.float64) + + def close(self): + if self.camera:self.lib.vap_camera_close(self.camera);self.camera=None + if self.engine:self.lib.vap_engine_close(self.engine);self.engine=None + +class NativePoseEngine(LivePoseEngine): + def __init__(self,args,native): + self.args=args;self.native=native;self.pose=native.pose + self.box_tracker=BoxTracker(max_age=30);self.people={} + self.timings=defaultdict(lambda:deque(maxlen=3000));self.isolated_inferences=0 diff --git a/orin/pipeline/occlusion_aware_pose_v2.py b/orin/pipeline/occlusion_aware_pose_v2.py new file mode 100644 index 0000000..5b7c73f --- /dev/null +++ b/orin/pipeline/occlusion_aware_pose_v2.py @@ -0,0 +1,1130 @@ +"""Occlusion-aware temporal pose reconstruction on top of RTMW-L fusion. + +Visible endpoints are filtered and optionally connected with a two-bone IK +solution. Short gaps are interpolated/predicted, medium gaps use a local motion +corridor, and long occlusions retain an undrawn shoulder sector for downstream +door-band intersection instead of emitting a fake precise wrist. +""" + +from __future__ import annotations + +import argparse +import json +import math +from collections import Counter, defaultdict +from dataclasses import dataclass +from pathlib import Path + +import cv2 +import numpy as np + +from compare_rtmw_ml import HAND_EDGES_LOCAL, HAND_SLICES, normalized_keypoints, summarize +from fuse_rtmw_l_temporal import BODY_SIDE_INDICES, draw_dashed_line, draw_fused + + +STATE_COLORS = { + "observed": (40, 220, 40), + "provisional": (255, 190, 40), + "recovered": (220, 180, 40), + "reassigned": (220, 120, 40), + "interpolated": (0, 220, 255), + "predicted_short": (0, 140, 255), + "reachable_medium": (0, 120, 255), + "reachable_long": (30, 80, 255), +} + + +@dataclass +class Measurement: + keypoints: np.ndarray + precise: bool + provisional: bool + source: str = "model" + quality: float = 0.0 + + +@dataclass +class EndpointTrack: + wrist: np.ndarray | None = None + velocity: np.ndarray | None = None + hand_shape: np.ndarray | None = None + last_observed_frame: int | None = None + missed_frames: int = 0 + last_elbow: np.ndarray | None = None + last_direction: np.ndarray | None = None + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser() + parser.add_argument("--video", type=Path, required=True) + parser.add_argument("--fusion-jsonl", type=Path, required=True) + parser.add_argument("--output", type=Path, required=True) + parser.add_argument("--short-occlusion-frames", type=int, default=12) + parser.add_argument("--medium-occlusion-frames", type=int, default=30) + parser.add_argument("--lookahead-frames", type=int, default=3) + return parser.parse_args() + + +def bbox_diagonal(bbox: list[float]) -> float: + return max(math.hypot(bbox[2] - bbox[0], bbox[3] - bbox[1]), 1.0) + + +def point_distance_ratio(a: np.ndarray, b: np.ndarray, bbox: list[float]) -> float: + return float(np.linalg.norm(a - b)) / bbox_diagonal(bbox) + + +def collect_measurements( + records: list[dict], + person_idx: int, + side: str, + lookahead_frames: int, +) -> list[Measurement | None]: + start, end = HAND_SLICES[side] + measurements: list[Measurement | None] = [None] * len(records) + observed_frames = [] + for frame_idx, record in enumerate(records): + meta = record["temporal_metadata"][person_idx][side] + if not meta["alarm_eligible"]: + continue + keypoints = np.asarray( + record["temporal_selected_keypoints"][person_idx][start:end], + dtype=np.float64, + ) + measurements[frame_idx] = Measurement( + keypoints=keypoints, + precise=bool(meta["distance_eligible"]), + provisional=not bool(meta["distance_eligible"]), + quality=float(meta.get("selected_hand_score", 0.0)), + ) + observed_frames.append(frame_idx) + + # A small future window is allowed for stable reporting. The alarm fast + # path still uses the unmodified current-frame candidates from V1. + for left, right in zip(observed_frames, observed_frames[1:]): + gap = right - left - 1 + if gap <= 0 or gap > lookahead_frames: + continue + left_points = measurements[left].keypoints + right_points = measurements[right].keypoints + for offset in range(1, gap + 1): + ratio = offset / (gap + 1) + points = (1.0 - ratio) * left_points + ratio * right_points + measurements[left + offset] = Measurement( + keypoints=points, + precise=False, + provisional=False, + source="fixed_lag_interpolation", + quality=min(measurements[left].quality, measurements[right].quality), + ) + return measurements + + +def estimate_bone_ratios(records: list[dict]) -> dict[tuple[int, str], tuple[float, float]]: + values: dict[tuple[int, str], list[tuple[float, float]]] = defaultdict(list) + for record in records: + for person_idx, bbox in enumerate(record["boxes"]): + keypoints = np.asarray( + record["temporal_selected_keypoints"][person_idx], dtype=np.float64 + ) + diagonal = bbox_diagonal(bbox) + for side, (shoulder, elbow, wrist) in BODY_SIDE_INDICES.items(): + meta = record["temporal_metadata"][person_idx][side] + if meta["limb_state"] != "consensus": + continue + upper = float(np.linalg.norm(keypoints[shoulder] - keypoints[elbow])) / diagonal + forearm = float(np.linalg.norm(keypoints[elbow] - keypoints[wrist])) / diagonal + if 0.04 <= upper <= 0.65 and 0.04 <= forearm <= 0.65: + values[(person_idx, side)].append((upper, forearm)) + + estimates: dict[tuple[int, str], tuple[float, float]] = {} + person_count = len(records[0]["boxes"]) + for person_idx in range(person_count): + for side in ("left", "right"): + samples = values[(person_idx, side)] + if len(samples) < 20: + opposite = "right" if side == "left" else "left" + samples = values[(person_idx, opposite)] + if samples: + array = np.asarray(samples) + estimates[(person_idx, side)] = ( + float(np.median(array[:, 0])), + float(np.median(array[:, 1])), + ) + else: + estimates[(person_idx, side)] = (0.24, 0.22) + return estimates + + +def solve_two_bone_elbow( + shoulder: np.ndarray, + wrist: np.ndarray, + upper_length: float, + forearm_length: float, + reference_elbow: np.ndarray, +) -> np.ndarray: + direction = wrist - shoulder + distance = float(np.linalg.norm(direction)) + if distance < 1e-6: + return reference_elbow.copy() + unit = direction / distance + + if distance >= upper_length + forearm_length: + return shoulder + unit * upper_length + if distance <= abs(upper_length - forearm_length): + sign = 1.0 if upper_length >= forearm_length else -1.0 + return shoulder + unit * upper_length * sign + + along = ( + upper_length * upper_length + - forearm_length * forearm_length + + distance * distance + ) / (2.0 * distance) + height = math.sqrt(max(upper_length * upper_length - along * along, 0.0)) + base = shoulder + unit * along + perpendicular = np.asarray([-unit[1], unit[0]]) + candidates = (base + perpendicular * height, base - perpendicular * height) + return min(candidates, key=lambda candidate: float(np.linalg.norm(candidate - reference_elbow))) + + +def update_endpoint_track( + track: EndpointTrack, + measurement: Measurement | None, + frame_idx: int, + short_occlusion_frames: int, + medium_occlusion_frames: int, +) -> tuple[np.ndarray | None, np.ndarray | None, str]: + if measurement is not None: + measured_wrist = measurement.keypoints[0] + measured_shape = measurement.keypoints - measured_wrist + if track.wrist is None or track.missed_frames > short_occlusion_frames: + track.wrist = measured_wrist.copy() + track.velocity = np.zeros(2, dtype=np.float64) + track.hand_shape = measured_shape.copy() + else: + prediction = track.wrist + track.velocity + residual = measured_wrist - prediction + alpha = 0.78 if measurement.precise else 0.88 + beta = 0.18 if measurement.precise else 0.24 + track.wrist = prediction + alpha * residual + track.velocity = 0.82 * track.velocity + beta * residual + shape_alpha = 0.55 if measurement.precise else 0.7 + track.hand_shape = ( + shape_alpha * measured_shape + (1.0 - shape_alpha) * track.hand_shape + ) + track.last_observed_frame = frame_idx + track.missed_frames = 0 + if measurement.source == "global_reassignment": + state = "reassigned" + elif measurement.source == "temporal_recovery": + state = "recovered" + elif measurement.precise: + state = "observed" + elif measurement.provisional: + state = "provisional" + else: + state = "interpolated" + return track.wrist.copy(), track.wrist + track.hand_shape, state + + track.missed_frames += 1 + if track.wrist is None: + return None, None, "reachable_long" + if track.missed_frames <= short_occlusion_frames: + track.wrist = track.wrist + track.velocity + track.velocity *= 0.86 + return track.wrist.copy(), track.wrist + track.hand_shape, "predicted_short" + if track.missed_frames <= medium_occlusion_frames: + return None, None, "reachable_medium" + return None, None, "reachable_long" + + +def temporal_recovery_measurement( + keypoints: np.ndarray, + metadata: dict, + side: str, + bbox: list[float], + track: EndpointTrack, + medium_occlusion_frames: int, +) -> Measurement | None: + """Recover a weak current-frame hand only when history strongly supports it.""" + if track.wrist is None or track.missed_frames > medium_occlusion_frames: + return None + if metadata.get("hand_state") not in {"conflict", "inferred"}: + return None + if not metadata.get("selected_hand_valid", False): + return None + if float(metadata.get("selected_hand_score", 0.0)) < 2.0: + return None + link = float(metadata.get("hand_selection", {}).get("link_to_selected_limb", 1.0)) + if link > 0.08: + return None + + start, end = HAND_SLICES[side] + hand = np.asarray(keypoints[start:end], dtype=np.float64) + if hand.shape != (end - start, 2) or not np.isfinite(hand).all(): + return None + prediction = track.wrist + track.velocity + diagonal = bbox_diagonal(bbox) + gate = min(0.20, 0.05 + 0.008 * max(track.missed_frames, 1)) * diagonal + if float(np.linalg.norm(hand[0] - prediction)) > gate: + return None + return Measurement( + keypoints=hand, + precise=False, + provisional=True, + source="temporal_recovery", + quality=float(metadata.get("selected_hand_score", 0.0)), + ) + + +def hand_span_px(hand: np.ndarray) -> float: + return math.hypot(float(np.ptp(hand[:, 0])), float(np.ptp(hand[:, 1]))) + + +def bbox_area(bbox: list[float]) -> float: + return max((bbox[2] - bbox[0]) * (bbox[3] - bbox[1]), 1.0) + + +def reset_endpoint_track(track: EndpointTrack) -> None: + track.wrist = None + track.velocity = None + track.hand_shape = None + track.last_observed_frame = None + track.missed_frames = 0 + track.last_elbow = None + track.last_direction = None + + +def resolve_global_hand_assignments( + frame_measurements: dict[tuple[int, str], Measurement | None], + tracks: dict[tuple[int, str], EndpointTrack], + boxes: list[list[float]], + medium_occlusion_frames: int, +) -> tuple[dict[tuple[int, str], Measurement | None], dict[tuple[int, str], dict]]: + """Resolve overlapping-person hands before binding them to body sides. + + Large foreground top-down crops may include a smaller background person. + The same physical hand can therefore be emitted for both people, or only by + the smaller crop even though it belongs to the foreground person. Direct + duplicates are kept on the foreground track when their apparent scales + agree. A missing foreground side can then reclaim a nearby candidate using + its recent hand trajectory and historical hand scale. + """ + resolved = dict(frame_measurements) + events: dict[tuple[int, str], dict] = {} + keys = list(resolved) + + # First collapse direct cross-person duplicates. Keep this deliberately + # asymmetric: a substantially larger foreground crop must also carry the + # substantially larger historical hand shape. + for index, first_key in enumerate(keys): + first = resolved[first_key] + if first is None: + continue + for second_key in keys[index + 1 :]: + second = resolved[second_key] + if second is None or first_key[0] == second_key[0]: + continue + first_hand = first.keypoints + second_hand = second.keypoints + first_span = hand_span_px(first_hand) + second_span = hand_span_px(second_hand) + center_distance = float( + np.linalg.norm(np.mean(first_hand, axis=0) - np.mean(second_hand, axis=0)) + ) + if center_distance > max(70.0, 0.5 * max(first_span, second_span)): + continue + + first_area = bbox_area(boxes[first_key[0]]) + second_area = bbox_area(boxes[second_key[0]]) + if first_area >= 2.5 * second_area and ( + first_span >= 1.45 * second_span + or first.quality >= second.quality + 0.6 + ): + owner_key, released_key = first_key, second_key + elif second_area >= 2.5 * first_area and ( + second_span >= 1.45 * first_span + or second.quality >= first.quality + 0.6 + ): + owner_key, released_key = second_key, first_key + else: + continue + resolved[released_key] = None + reset_endpoint_track(tracks[released_key]) + events[owner_key] = { + "method": "foreground_duplicate_owner", + "released_person_id": released_key[0] + 1, + "released_side": released_key[1], + } + events[released_key] = { + "method": "reassigned_away", + "assigned_person_id": owner_key[0] + 1, + "assigned_side": owner_key[1], + } + + # Then recover missing foreground hands from another crop. The candidate + # must fall inside the target's recent motion gate and be much smaller than + # the target's historical foreground hand shape. + available_sources = [key for key, value in resolved.items() if value is not None] + target_candidates = [] + for target_key, target_measurement in resolved.items(): + if target_measurement is not None: + continue + target_track = tracks[target_key] + if ( + target_track.wrist is None + or target_track.hand_shape is None + or target_track.missed_frames > medium_occlusion_frames + ): + continue + target_area = bbox_area(boxes[target_key[0]]) + predicted_center = target_track.wrist + np.mean(target_track.hand_shape, axis=0) + target_span = hand_span_px(target_track.hand_shape) + gate = 0.20 * bbox_diagonal(boxes[target_key[0]]) + for source_key in available_sources: + if source_key[0] == target_key[0]: + continue + source = resolved[source_key] + if source is None or source.source != "model": + continue + source_area = bbox_area(boxes[source_key[0]]) + if target_area < 2.5 * source_area: + continue + source_center = np.mean(source.keypoints, axis=0) + source_span = hand_span_px(source.keypoints) + distance = float(np.linalg.norm(source_center - predicted_center)) + if distance > gate or target_span < 1.45 * max(source_span, 1.0): + continue + target_candidates.append( + (distance / gate, target_key, source_key, source_center) + ) + + used_targets = set() + used_sources = set() + for _, target_key, source_key, source_center in sorted(target_candidates): + if target_key in used_targets or source_key in used_sources: + continue + target_track = tracks[target_key] + shape_center = np.mean(target_track.hand_shape, axis=0) + reassigned_wrist = source_center - shape_center + reassigned_hand = reassigned_wrist + target_track.hand_shape + resolved[target_key] = Measurement( + keypoints=reassigned_hand, + precise=False, + provisional=True, + source="global_reassignment", + quality=resolved[source_key].quality, + ) + resolved[source_key] = None + reset_endpoint_track(tracks[source_key]) + events[target_key] = { + "method": "foreground_temporal_reassignment", + "source_person_id": source_key[0] + 1, + "source_side": source_key[1], + "observed_center": np.round(source_center, 3).tolist(), + } + events[source_key] = { + "method": "reassigned_away", + "assigned_person_id": target_key[0] + 1, + "assigned_side": target_key[1], + } + used_targets.add(target_key) + used_sources.add(source_key) + return resolved, events + + +def track_direction(track: EndpointTrack) -> np.ndarray: + if track.last_direction is not None: + return track.last_direction / max(float(np.linalg.norm(track.last_direction)), 1e-6) + if track.velocity is not None and float(np.linalg.norm(track.velocity)) > 1e-6: + return track.velocity / np.linalg.norm(track.velocity) + return np.asarray([1.0, 0.0], dtype=np.float64) + + +def ellipse_polygon( + center: np.ndarray, + major_axis: float, + minor_axis: float, + direction: np.ndarray, + point_count: int = 64, +) -> np.ndarray: + perpendicular = np.asarray([-direction[1], direction[0]], dtype=np.float64) + angles = np.linspace(0.0, 2.0 * math.pi, point_count, endpoint=False) + return ( + center + + np.cos(angles)[:, None] * major_axis * direction + + np.sin(angles)[:, None] * minor_axis * perpendicular + ) + + +def clip_polygon_from_anchor_to_circle( + points: np.ndarray, + anchor: np.ndarray, + circle_center: np.ndarray, + radius: float, +) -> np.ndarray: + """Clip radial polygon samples to a shoulder-centered anatomical limit.""" + clipped = [] + offset = anchor - circle_center + for point in points: + if float(np.linalg.norm(point - circle_center)) <= radius: + clipped.append(point) + continue + ray = point - anchor + a = float(np.dot(ray, ray)) + b = 2.0 * float(np.dot(offset, ray)) + c = float(np.dot(offset, offset) - radius * radius) + discriminant = max(b * b - 4.0 * a * c, 0.0) + roots = [ + value + for value in ( + (-b - math.sqrt(discriminant)) / max(2.0 * a, 1e-6), + (-b + math.sqrt(discriminant)) / max(2.0 * a, 1e-6), + ) + if 0.0 <= value <= 1.0 + ] + clipped.append(anchor + (max(roots) if roots else 0.0) * ray) + return np.asarray(clipped, dtype=np.float64) + + +def medium_motion_region( + shoulder: np.ndarray, + upper_length: float, + forearm_length: float, + bbox: list[float], + track: EndpointTrack, + short_occlusion_frames: int, +) -> dict: + diagonal = bbox_diagonal(bbox) + direction = track_direction(track) + hand_shape = ( + track.hand_shape + if track.hand_shape is not None + else np.zeros((21, 2), dtype=np.float64) + ) + shape_center = np.mean(hand_shape, axis=0) + centered_shape = hand_shape - shape_center + perpendicular = np.asarray([-direction[1], direction[0]], dtype=np.float64) + hand_major = float(np.max(np.abs(centered_shape @ direction))) + hand_minor = float(np.max(np.abs(centered_shape @ perpendicular))) + + extra_frames = max(track.missed_frames - short_occlusion_frames, 1) + motion_major = min(0.20, 0.04 + 0.007 * extra_frames) * diagonal + motion_minor = min(0.13, 0.03 + 0.0045 * extra_frames) * diagonal + major_axis = max(hand_major + motion_major, 0.08 * diagonal) + minor_axis = max(hand_minor + motion_minor, 0.06 * diagonal) + + center = track.wrist + shape_center + tracked_hand_extent = float(np.max(np.linalg.norm(centered_shape, axis=1))) + hand_reach = max(tracked_hand_extent, 0.10 * diagonal) + anatomical_reach = upper_length + forearm_length + hand_reach + 0.02 * diagonal + shoulder_to_center = center - shoulder + center_distance = float(np.linalg.norm(shoulder_to_center)) + if center_distance > 0.96 * anatomical_reach: + center = shoulder + shoulder_to_center / center_distance * 0.96 * anatomical_reach + + polygon = ellipse_polygon(center, major_axis, minor_axis, direction) + polygon = clip_polygon_from_anchor_to_circle( + polygon, + center, + shoulder, + anatomical_reach, + ) + direction_deg = math.degrees(math.atan2(float(direction[1]), float(direction[0]))) + return { + "geometry": "motion_corridor", + "center": np.round(center, 3).tolist(), + "predicted_wrist": np.round(track.wrist, 3).tolist(), + "major_axis_px": round(major_axis, 3), + "minor_axis_px": round(minor_axis, 3), + "outer_radius_px": round(major_axis, 3), + "direction_deg": round(direction_deg, 3), + "max_reach_from_shoulder_px": round(anatomical_reach, 3), + "polygon": np.round(polygon, 3).tolist(), + "requires_door_intersection": True, + "render_full_region": True, + "distance_lower_bound_rule": "minimum distance from motion-corridor polygon to door_edge", + } + + +def long_reachable_region( + shoulder: np.ndarray, + upper_length: float, + forearm_length: float, + bbox: list[float], + track: EndpointTrack, + medium_occlusion_frames: int, +) -> dict: + diagonal = bbox_diagonal(bbox) + height = max(bbox[3] - bbox[1], 1.0) + arm_reach = upper_length + forearm_length + hand_reach = min(0.13 * diagonal, 0.17 * height) + safety_margin = 0.025 * diagonal + outer_radius = arm_reach + hand_reach + safety_margin + inner_radius = max(abs(upper_length - forearm_length) - 0.04 * diagonal, 0.0) + if track.last_direction is None: + direction_deg = 0.0 + half_angle_deg = 180.0 + else: + direction = track_direction(track) + direction_deg = math.degrees(math.atan2(float(direction[1]), float(direction[0]))) + half_angle_deg = min( + 180.0, + 65.0 + 2.0 * max(track.missed_frames - medium_occlusion_frames, 0), + ) + return { + "geometry": "shoulder_sector", + "center": np.round(shoulder, 3).tolist(), + "inner_radius_px": round(inner_radius, 3), + "outer_radius_px": round(outer_radius, 3), + "arm_reach_px": round(arm_reach, 3), + "hand_reach_px": round(hand_reach, 3), + "safety_margin_px": round(safety_margin, 3), + "direction_deg": round(direction_deg, 3), + "half_angle_deg": round(half_angle_deg, 3), + "requires_door_intersection": True, + "render_full_region": False, + "distance_lower_bound_rule": "intersect shoulder sector with door danger band before alarming", + } + + +def reachable_region( + shoulder: np.ndarray, + upper_length: float, + forearm_length: float, + bbox: list[float], + track: EndpointTrack, + state: str = "reachable_long", + short_occlusion_frames: int = 12, + medium_occlusion_frames: int = 30, +) -> dict: + if state == "reachable_medium" and track.wrist is not None: + return medium_motion_region( + shoulder, + upper_length, + forearm_length, + bbox, + track, + short_occlusion_frames, + ) + return long_reachable_region( + shoulder, + upper_length, + forearm_length, + bbox, + track, + medium_occlusion_frames, + ) + + +def calibrated_bone_lengths( + upper_ratio: float, + forearm_ratio: float, + bbox: list[float], +) -> tuple[float, float]: + diagonal = bbox_diagonal(bbox) + height = max(bbox[3] - bbox[1], 1.0) + upper = upper_ratio * diagonal + forearm = forearm_ratio * diagonal + total = upper + forearm + # Keep only a loose sanity cap. The previous 0.55 * bbox-height cap was + # unsafe for close, clipped CMS views: a valid projected arm can occupy + # almost the full detection height. The ratios are already robust medians, + # so this cap only limits catastrophic calibration outliers. + maximum_total = 0.95 * diagonal + if total > maximum_total: + scale = maximum_total / total + upper *= scale + forearm *= scale + return upper, forearm + + +def process_sequence( + records: list[dict], + short_occlusion_frames: int, + medium_occlusion_frames: int, + lookahead_frames: int, +) -> list[dict]: + person_count = len(records[0]["boxes"]) + measurements = { + (person_idx, side): collect_measurements( + records, person_idx, side, lookahead_frames + ) + for person_idx in range(person_count) + for side in ("left", "right") + } + bone_ratios = estimate_bone_ratios(records) + tracks = { + (person_idx, side): EndpointTrack() + for person_idx in range(person_count) + for side in ("left", "right") + } + + output = [] + for frame_idx, record in enumerate(records): + boxes = record["boxes"] + base_keypoints = np.asarray(record["temporal_keypoints"], dtype=np.float64) + unfiltered_keypoints = np.asarray( + record["temporal_selected_keypoints"], dtype=np.float64 + ) + v2_keypoints = base_keypoints.copy() + frame_measurements = { + (person_idx, side): measurements[(person_idx, side)][frame_idx] + for person_idx in range(person_count) + for side in ("left", "right") + } + frame_measurements, assignment_events = resolve_global_hand_assignments( + frame_measurements, + tracks, + boxes, + medium_occlusion_frames, + ) + people_meta = [] + for person_idx, bbox in enumerate(boxes): + diagonal = bbox_diagonal(bbox) + person_meta = {} + for side, (shoulder_idx, elbow_idx, wrist_idx) in BODY_SIDE_INDICES.items(): + track = tracks[(person_idx, side)] + measurement = frame_measurements[(person_idx, side)] + if measurement is None: + measurement = temporal_recovery_measurement( + unfiltered_keypoints[person_idx], + record["temporal_metadata"][person_idx][side], + side, + bbox, + track, + medium_occlusion_frames, + ) + returning_after_long_occlusion = ( + measurement is not None + and track.missed_frames > short_occlusion_frames + ) + wrist, hand, state = update_endpoint_track( + track, + measurement, + frame_idx, + short_occlusion_frames, + medium_occlusion_frames, + ) + shoulder = base_keypoints[person_idx, shoulder_idx] + upper_ratio, forearm_ratio = bone_ratios[(person_idx, side)] + upper_length, forearm_length = calibrated_bone_lengths( + upper_ratio, forearm_ratio, bbox + ) + reachable = None + + if wrist is not None and hand is not None: + reference_elbow = ( + unfiltered_keypoints[person_idx, elbow_idx] + if returning_after_long_occlusion or track.last_elbow is None + else track.last_elbow + ) + use_ik = ( + record["temporal_metadata"][person_idx][side]["limb_state"] + != "consensus" + or state in ( + "provisional", + "recovered", + "reassigned", + "interpolated", + "predicted_short", + ) + ) + if use_ik: + elbow = solve_two_bone_elbow( + shoulder, + wrist, + upper_length, + forearm_length, + reference_elbow, + ) + else: + elbow = base_keypoints[person_idx, elbow_idx] + track.last_elbow = elbow.copy() + direction = wrist - shoulder + if float(np.linalg.norm(direction)) > 1e-6: + track.last_direction = direction / np.linalg.norm(direction) + + v2_keypoints[person_idx, shoulder_idx] = shoulder + v2_keypoints[person_idx, elbow_idx] = elbow + v2_keypoints[person_idx, wrist_idx] = wrist + start, end = HAND_SLICES[side] + v2_keypoints[person_idx, start:end] = hand + uncertainty_radius = ( + 0.025 * diagonal + if state == "observed" + else (0.04 + 0.015 * track.missed_frames) * diagonal + ) + precise_distance = state == "observed" + fast_alarm = state in ( + "observed", + "provisional", + "reassigned", + "interpolated", + "predicted_short", + ) + else: + reachable = reachable_region( + shoulder, + upper_length, + forearm_length, + bbox, + track, + state, + short_occlusion_frames, + medium_occlusion_frames, + ) + uncertainty_radius = reachable["outer_radius_px"] + precise_distance = False + fast_alarm = True + + person_meta[side] = { + "state": state, + "precise_distance": precise_distance, + "fast_alarm_candidate": fast_alarm, + "risk_mode": ( + "precise_hand" + if state == "observed" + else ( + "motion_corridor_intersection" + if state == "reachable_medium" + else ( + "door_intersection_only" + if state == "reachable_long" + else ( + "reassigned_hand" + if state == "reassigned" + else "temporal_hand" + ) + ) + ) + ), + "requires_door_intersection": state + in ("reachable_medium", "reachable_long"), + "missed_frames": track.missed_frames, + "uncertainty_radius_px": round(float(uncertainty_radius), 3), + "bone_lengths_px": { + "upper": round(upper_length, 3), + "forearm": round(forearm_length, 3), + }, + "identity_assignment": assignment_events.get( + (person_idx, side) + ), + "reachable_region": reachable, + } + people_meta.append(person_meta) + + output.append( + { + "frame_index": frame_idx, + "boxes": boxes, + "v1_keypoints": record["temporal_keypoints"], + "v1_metadata": record["temporal_metadata"], + "v2_keypoints": np.round(v2_keypoints, 3).tolist(), + "v2_metadata": people_meta, + } + ) + return output + + +def sector_points(region: dict) -> np.ndarray: + center = np.asarray(region["center"], dtype=np.float64) + direction = region["direction_deg"] + half_angle = region["half_angle_deg"] + angles = np.linspace(direction - half_angle, direction + half_angle, 48) + radius = region["outer_radius_px"] + arc = np.column_stack( + [ + center[0] + radius * np.cos(np.deg2rad(angles)), + center[1] + radius * np.sin(np.deg2rad(angles)), + ] + ) + if half_angle < 179.0: + return np.vstack([center, arc, center]).astype(np.int32) + return arc.astype(np.int32) + + +def draw_v2( + image: np.ndarray, + boxes: list[list[float]], + keypoints: np.ndarray, + metadata: list[dict], + title: str, + show_reachable: bool = True, +) -> np.ndarray: + canvas = image.copy() + overlay = canvas.copy() + for person_idx, (bbox, kpts, person_meta) in enumerate( + zip(boxes, keypoints, metadata), start=1 + ): + x1, y1, x2, y2 = (int(round(value)) for value in bbox) + cv2.rectangle(canvas, (x1, y1), (x2, y2), (0, 220, 255), 2) + cv2.putText( + canvas, + f"P{person_idx}", + (x1 + 4, max(22, y1 + 24)), + cv2.FONT_HERSHEY_SIMPLEX, + 0.72, + (0, 220, 255), + 2, + cv2.LINE_AA, + ) + for side, (shoulder, elbow, wrist) in BODY_SIDE_INDICES.items(): + meta = person_meta[side] + state = meta["state"] + color = STATE_COLORS[state] + if state in ("reachable_medium", "reachable_long"): + region = meta["reachable_region"] + if state == "reachable_medium": + if show_reachable: + polygon = np.round(region["polygon"]).astype(np.int32) + cv2.fillPoly(overlay, [polygon], color) + cv2.polylines(canvas, [polygon], True, color, 3, cv2.LINE_AA) + center = tuple(np.round(region["center"]).astype(int)) + wrist_point = tuple( + np.round(region["predicted_wrist"]).astype(int) + ) + cv2.line(canvas, center, wrist_point, color, 2, cv2.LINE_AA) + label = f"{side[0].upper()}:corridor" + else: + center = tuple(np.round(region["center"]).astype(int)) + cv2.circle(canvas, center, 8, color, 2, cv2.LINE_AA) + if region["half_angle_deg"] >= 179.0: + cv2.putText( + canvas, + "?", + (center[0] + 11, center[1] + 9), + cv2.FONT_HERSHEY_SIMPLEX, + 0.65, + color, + 2, + cv2.LINE_AA, + ) + else: + direction = math.radians(region["direction_deg"]) + direction_end = ( + int(round(center[0] + 55 * math.cos(direction))), + int(round(center[1] + 55 * math.sin(direction))), + ) + cv2.arrowedLine( + canvas, + center, + direction_end, + color, + 2, + cv2.LINE_AA, + tipLength=0.25, + ) + if show_reachable and region.get("render_full_region", False): + polygon = sector_points(region) + if region["half_angle_deg"] < 179.0: + cv2.fillPoly(overlay, [polygon], color) + cv2.polylines( + canvas, [polygon], True, color, 3, cv2.LINE_AA + ) + else: + radius = int(round(region["outer_radius_px"])) + cv2.circle(canvas, center, radius, color, 3, cv2.LINE_AA) + label = f"{side[0].upper()}:unknown" + else: + thickness = 3 + if state in ("reassigned", "recovered"): + draw_dashed_line( + canvas, + kpts[shoulder], + kpts[elbow], + color, + thickness, + ) + draw_dashed_line( + canvas, + kpts[elbow], + kpts[wrist], + color, + thickness, + ) + else: + cv2.line( + canvas, + tuple(np.round(kpts[shoulder]).astype(int)), + tuple(np.round(kpts[elbow]).astype(int)), + color, + thickness, + cv2.LINE_AA, + ) + cv2.line( + canvas, + tuple(np.round(kpts[elbow]).astype(int)), + tuple(np.round(kpts[wrist]).astype(int)), + color, + thickness, + cv2.LINE_AA, + ) + start, end = HAND_SLICES[side] + hand = kpts[start:end] + for a, b in HAND_EDGES_LOCAL: + cv2.line( + canvas, + tuple(np.round(hand[a]).astype(int)), + tuple(np.round(hand[b]).astype(int)), + color, + 2, + cv2.LINE_AA, + ) + label = f"{side[0].upper()}:{state[:5]}" + cv2.putText( + canvas, + label, + (x1 + 4, max(45, y1 + (48 if side == "left" else 70))), + cv2.FONT_HERSHEY_SIMPLEX, + 0.48, + color, + 1, + cv2.LINE_AA, + ) + canvas = cv2.addWeighted(overlay, 0.06, canvas, 0.94, 0.0) + cv2.rectangle(canvas, (0, 0), (670, 45), (0, 0, 0), -1) + cv2.putText( + canvas, + title, + (10, 31), + cv2.FONT_HERSHEY_SIMPLEX, + 0.8, + (255, 255, 255), + 2, + cv2.LINE_AA, + ) + return canvas + + +def endpoint_metrics(records: list[dict]) -> dict: + states: Counter = Counter() + assignments: Counter = Counter() + steps: dict[str, list[float]] = defaultdict(list) + previous: dict[tuple[int, str], tuple[int, np.ndarray]] = {} + for record in records: + for person_idx, bbox in enumerate(record["boxes"]): + kpts = np.asarray(record["v2_keypoints"][person_idx], dtype=np.float64) + for side, (_, _, wrist) in BODY_SIDE_INDICES.items(): + meta = record["v2_metadata"][person_idx][side] + key_name = f"P{person_idx + 1}_{side}" + states[f"{key_name}_{meta['state']}"] += 1 + identity_assignment = meta.get("identity_assignment") + if identity_assignment is not None: + assignments[identity_assignment["method"]] += 1 + if meta["state"] in ("reachable_medium", "reachable_long"): + previous.pop((person_idx, side), None) + continue + normalized = normalized_keypoints(kpts, bbox)[wrist] + prior = previous.get((person_idx, side)) + if prior is not None and prior[0] == record["frame_index"] - 1: + steps[key_name].append(float(np.linalg.norm(normalized - prior[1]))) + previous[(person_idx, side)] = (record["frame_index"], normalized) + return { + "state_counts": dict(sorted(states.items())), + "identity_assignment_counts": dict(sorted(assignments.items())), + "wrist_frame_step": {name: summarize(items) for name, items in steps.items()}, + } + + +def write_video_and_montage( + video_path: Path, + records: list[dict], + output_dir: Path, +) -> None: + cap = cv2.VideoCapture(str(video_path)) + fps = float(cap.get(cv2.CAP_PROP_FPS)) + writer = cv2.VideoWriter( + str(output_dir / "rtmw_l_temporal_v1_vs_occlusion_v4.mp4"), + cv2.VideoWriter_fourcc(*"mp4v"), + fps, + (1920, 540), + ) + key_frames = {0, 50, 120, 126, 132, 138, 180, 211, 220, 240, 252, + 292, 296, 300, 304, 320, 360, 411} + montage_rows = [] + frame_idx = 0 + while True: + ok, frame = cap.read() + if not ok: + break + record = records[frame_idx] + v1 = draw_fused( + frame, + record["boxes"], + [np.asarray(points) for points in record["v1_keypoints"]], + record["v1_metadata"], + f"temporal fusion V1 | frame {frame_idx}", + ) + v2 = draw_v2( + frame, + record["boxes"], + np.asarray(record["v2_keypoints"]), + record["v2_metadata"], + f"occlusion V4 global assignment | frame {frame_idx}", + ) + left = cv2.resize(v1, (960, 540), interpolation=cv2.INTER_AREA) + right = cv2.resize(v2, (960, 540), interpolation=cv2.INTER_AREA) + writer.write(np.hstack([left, right])) + if frame_idx in key_frames: + montage_rows.append( + np.hstack( + [ + cv2.resize(v1, (640, 360), interpolation=cv2.INTER_AREA), + cv2.resize(v2, (640, 360), interpolation=cv2.INTER_AREA), + ] + ) + ) + frame_idx += 1 + cap.release() + writer.release() + cv2.imwrite( + str(output_dir / "rtmw_l_temporal_v1_vs_occlusion_v4_montage.jpg"), + np.vstack(montage_rows), + ) + + +def main() -> None: + args = parse_args() + args.output.mkdir(parents=True, exist_ok=True) + records = [ + json.loads(line) + for line in args.fusion_jsonl.read_text(encoding="utf-8").splitlines() + ] + output = process_sequence( + records, + short_occlusion_frames=args.short_occlusion_frames, + medium_occlusion_frames=args.medium_occlusion_frames, + lookahead_frames=args.lookahead_frames, + ) + with (args.output / "occlusion_v4.jsonl").open("w", encoding="utf-8") as handle: + for record in output: + handle.write(json.dumps(record, ensure_ascii=False) + "\n") + + summary = { + "algorithm_version": "occlusion_v4_global_assignment", + "frames": len(output), + "short_occlusion_frames": args.short_occlusion_frames, + "medium_occlusion_frames": args.medium_occlusion_frames, + "lookahead_frames": args.lookahead_frames, + "policy": { + "observed": "filtered visible endpoint; precise distance allowed", + "provisional": "strong current-frame endpoint; fast alarm only", + "recovered": "weak current-frame endpoint accepted by temporal gating", + "reassigned": "global hand candidate transferred using foreground scale and trajectory", + "interpolated": "fixed-lag short-gap interpolation; fast alarm only", + "predicted_short": "damped-velocity prediction; uncertainty grows each frame", + "reachable_medium": "track-centered motion corridor; intersect with door band", + "reachable_long": "no exact wrist; keep shoulder sector data for door intersection only", + "ik": "two-bone shoulder/wrist reconstruction using robust visible bone lengths", + "ownership": "one-to-one global hand reassignment using foreground scale and trajectory", + }, + "metrics": endpoint_metrics(output), + } + with (args.output / "summary.json").open("w", encoding="utf-8") as handle: + json.dump(summary, handle, ensure_ascii=False, indent=2) + write_video_and_montage(args.video, output, args.output) + print(json.dumps(summary, ensure_ascii=False, indent=2)) + + +if __name__ == "__main__": + main() diff --git a/orin/pipeline/predict_optimized_pose_videos.py b/orin/pipeline/predict_optimized_pose_videos.py new file mode 100644 index 0000000..dd15617 --- /dev/null +++ b/orin/pipeline/predict_optimized_pose_videos.py @@ -0,0 +1,683 @@ +"""End-to-end optimized RTMW-L prediction for CMS videos. + +Pipeline: +YOLO11s person -> fixed-camera false-positive filters -> lightweight track IDs +-> RTMW-L raw/smaller-person-isolated views -> temporal fusion -> causal +occlusion-aware V4 global hand assignment. Reachable regions stay in JSON and are hidden in the main +comparison video. +""" + +from __future__ import annotations + +import argparse +import json +import math +import time +from collections import Counter, defaultdict, deque +from dataclasses import dataclass, field +from pathlib import Path + +import cv2 +import numpy as np + +from compare_rtmw_ml import HAND_SLICES, draw_pose, summarize +from fuse_rtmw_l_temporal import ( + BODY_SIDE_INDICES, + PersonTrack, + adaptive_smooth, + build_global_hands, + draw_fused, + fuse_person, +) +from occlusion_aware_pose_v2 import ( + EndpointTrack, + Measurement, + calibrated_bone_lengths, + draw_v2, + reachable_region, + resolve_global_hand_assignments, + solve_two_bone_elbow, + temporal_recovery_measurement, + update_endpoint_track, +) +from rtmw_l_person_mask_ablation import mask_other_people + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser() + parser.add_argument("--videos", type=Path, nargs="+", required=True) + parser.add_argument("--det-model", type=Path, required=True) + parser.add_argument("--pose-model", type=Path, required=True) + parser.add_argument("--output", type=Path, required=True) + parser.add_argument( + "--det-device", + default="0", + help="Ultralytics YOLO device, for example 0, cuda:0, or cpu.", + ) + parser.add_argument( + "--pose-device", + default="cpu", + help="RTMW ONNX Runtime device: cpu or cuda.", + ) + parser.add_argument("--det-imgsz", type=int, default=960) + parser.add_argument("--det-conf", type=float, default=0.30) + parser.add_argument("--det-roi-y-max-ratio", type=float, default=1.0) + parser.add_argument("--short-occlusion-frames", type=int, default=12) + parser.add_argument("--medium-occlusion-frames", type=int, default=30) + parser.add_argument("--max-people", type=int, default=3) + return parser.parse_args() + + +def bbox_iou(a: np.ndarray, b: np.ndarray) -> float: + x1 = max(float(a[0]), float(b[0])) + y1 = max(float(a[1]), float(b[1])) + x2 = min(float(a[2]), float(b[2])) + y2 = min(float(a[3]), float(b[3])) + intersection = max(x2 - x1, 0.0) * max(y2 - y1, 0.0) + union = ( + max(float(a[2] - a[0]), 0.0) * max(float(a[3] - a[1]), 0.0) + + max(float(b[2] - b[0]), 0.0) * max(float(b[3] - b[1]), 0.0) + - intersection + ) + return intersection / max(union, 1.0) + + +class BoxTracker: + def __init__(self, max_age: int = 30): + self.max_age = max_age + self.next_id = 1 + self.tracks: dict[int, dict] = {} + + def update( + self, + boxes: list[list[float]], + frame_idx: int, + image_shape: tuple[int, int], + ) -> list[int]: + height, width = image_shape + image_diagonal = math.hypot(width, height) + active_ids = [ + track_id + for track_id, track in self.tracks.items() + if frame_idx - track["last_frame"] <= self.max_age + ] + self.tracks = {track_id: self.tracks[track_id] for track_id in active_ids} + pairs = [] + for detection_idx, box in enumerate(boxes): + box_array = np.asarray(box, dtype=np.float64) + center = (box_array[:2] + box_array[2:]) * 0.5 + for track_id in active_ids: + prior = np.asarray(self.tracks[track_id]["box"], dtype=np.float64) + prior_center = (prior[:2] + prior[2:]) * 0.5 + iou = bbox_iou(box_array, prior) + center_distance = float(np.linalg.norm(center - prior_center)) / image_diagonal + if iou >= 0.05 or center_distance <= 0.18: + score = 2.0 * iou + max(0.0, 1.0 - center_distance / 0.18) + pairs.append((score, detection_idx, track_id)) + + assigned_detections = set() + assigned_tracks = set() + result: list[int | None] = [None] * len(boxes) + for _, detection_idx, track_id in sorted(pairs, reverse=True): + if detection_idx in assigned_detections or track_id in assigned_tracks: + continue + assigned_detections.add(detection_idx) + assigned_tracks.add(track_id) + result[detection_idx] = track_id + + for detection_idx, box in enumerate(boxes): + track_id = result[detection_idx] + if track_id is None: + track_id = self.next_id + self.next_id += 1 + result[detection_idx] = track_id + self.tracks[track_id] = {"box": box, "last_frame": frame_idx} + return [int(track_id) for track_id in result] + + +@dataclass +class LivePersonState: + fusion: PersonTrack = field(default_factory=PersonTrack) + endpoints: dict[str, EndpointTrack] = field( + default_factory=lambda: {"left": EndpointTrack(), "right": EndpointTrack()} + ) + bone_samples: dict[str, deque] = field( + default_factory=lambda: { + "left": deque(maxlen=120), + "right": deque(maxlen=120), + } + ) + last_frame: int = -1 + + +def bbox_diagonal(bbox: list[float]) -> float: + return max(math.hypot(bbox[2] - bbox[0], bbox[3] - bbox[1]), 1.0) + + +def update_bone_samples( + state: LivePersonState, + keypoints: np.ndarray, + metadata: dict, + bbox: list[float], +) -> None: + diagonal = bbox_diagonal(bbox) + for side, (shoulder, elbow, wrist) in BODY_SIDE_INDICES.items(): + meta = metadata[side] + if meta["limb_state"] != "consensus" or not meta["distance_eligible"]: + continue + upper = float(np.linalg.norm(keypoints[shoulder] - keypoints[elbow])) / diagonal + forearm = float(np.linalg.norm(keypoints[elbow] - keypoints[wrist])) / diagonal + if 0.04 <= upper <= 0.65 and 0.04 <= forearm <= 0.65: + state.bone_samples[side].append((upper, forearm)) + + +def current_bone_ratios(state: LivePersonState, side: str) -> tuple[float, float]: + samples = state.bone_samples[side] + if len(samples) < 10: + opposite = "right" if side == "left" else "left" + if len(state.bone_samples[opposite]) >= 10: + samples = state.bone_samples[opposite] + if not samples: + return 0.24, 0.22 + values = np.asarray(samples, dtype=np.float64) + return float(np.median(values[:, 0])), float(np.median(values[:, 1])) + + +def update_live_v2( + frame_idx: int, + bbox: list[float], + selected_keypoints: np.ndarray, + smoothed_keypoints: np.ndarray, + fusion_metadata: dict, + state: LivePersonState, + assigned_measurements: dict[str, Measurement | None] | None, + identity_assignments: dict[str, dict] | None, + short_occlusion_frames: int, + medium_occlusion_frames: int, +) -> tuple[np.ndarray, dict]: + frame_gap = 0 if state.last_frame < 0 else max(frame_idx - state.last_frame - 1, 0) + for endpoint in state.endpoints.values(): + endpoint.missed_frames += frame_gap + state.last_frame = frame_idx + update_bone_samples(state, selected_keypoints, fusion_metadata, bbox) + + output = smoothed_keypoints.copy() + output_metadata = {} + diagonal = bbox_diagonal(bbox) + for side, (shoulder_idx, elbow_idx, wrist_idx) in BODY_SIDE_INDICES.items(): + endpoint = state.endpoints[side] + meta = fusion_metadata[side] + start, end = HAND_SLICES[side] + measurement = ( + assigned_measurements.get(side) + if assigned_measurements is not None + else None + ) + if assigned_measurements is None and meta["alarm_eligible"]: + measurement = Measurement( + keypoints=selected_keypoints[start:end], + precise=bool(meta["distance_eligible"]), + provisional=not bool(meta["distance_eligible"]), + quality=float(meta.get("selected_hand_score", 0.0)), + ) + if measurement is None: + measurement = temporal_recovery_measurement( + selected_keypoints, + meta, + side, + bbox, + endpoint, + medium_occlusion_frames, + ) + returning_after_long = ( + measurement is not None + and endpoint.missed_frames > short_occlusion_frames + ) + wrist, hand, endpoint_state = update_endpoint_track( + endpoint, + measurement, + frame_idx, + short_occlusion_frames, + medium_occlusion_frames, + ) + shoulder = smoothed_keypoints[shoulder_idx] + upper_ratio, forearm_ratio = current_bone_ratios(state, side) + upper_length, forearm_length = calibrated_bone_lengths( + upper_ratio, forearm_ratio, bbox + ) + region = None + if wrist is not None and hand is not None: + reference_elbow = ( + selected_keypoints[elbow_idx] + if returning_after_long or endpoint.last_elbow is None + else endpoint.last_elbow + ) + use_ik = ( + meta["limb_state"] != "consensus" + or endpoint_state + in ("provisional", "recovered", "reassigned", "predicted_short") + ) + elbow = ( + solve_two_bone_elbow( + shoulder, + wrist, + upper_length, + forearm_length, + reference_elbow, + ) + if use_ik + else smoothed_keypoints[elbow_idx] + ) + endpoint.last_elbow = elbow.copy() + direction = wrist - shoulder + if float(np.linalg.norm(direction)) > 1e-6: + endpoint.last_direction = direction / np.linalg.norm(direction) + output[shoulder_idx] = shoulder + output[elbow_idx] = elbow + output[wrist_idx] = wrist + output[start:end] = hand + uncertainty = ( + 0.025 * diagonal + if endpoint_state == "observed" + else (0.04 + 0.015 * endpoint.missed_frames) * diagonal + ) + precise = endpoint_state == "observed" + else: + region = reachable_region( + shoulder, + upper_length, + forearm_length, + bbox, + endpoint, + endpoint_state, + short_occlusion_frames, + medium_occlusion_frames, + ) + uncertainty = region["outer_radius_px"] + precise = False + output_metadata[side] = { + "state": endpoint_state, + "precise_distance": precise, + "fast_alarm_candidate": True, + "risk_mode": ( + "precise_hand" + if endpoint_state == "observed" + else ( + "motion_corridor_intersection" + if endpoint_state == "reachable_medium" + else ( + "door_intersection_only" + if endpoint_state == "reachable_long" + else ( + "reassigned_hand" + if endpoint_state == "reassigned" + else "temporal_hand" + ) + ) + ) + ), + "requires_door_intersection": endpoint_state + in ("reachable_medium", "reachable_long"), + "identity_assignment": ( + identity_assignments.get(side) + if identity_assignments is not None + else None + ), + "missed_frames": endpoint.missed_frames, + "uncertainty_radius_px": round(float(uncertainty), 3), + "bone_lengths_px": { + "upper": round(upper_length, 3), + "forearm": round(forearm_length, 3), + }, + "reachable_region": region, + } + return output, output_metadata + + +def filter_person_boxes( + result, + width: int, + height: int, + max_people: int, +) -> tuple[list[list[float]], list[float], Counter]: + boxes = result.boxes.xyxy.cpu().numpy() + scores = result.boxes.conf.cpu().numpy() + rejected = Counter() + kept = [] + for box, score in zip(boxes, scores): + center_x = float(box[0] + box[2]) * 0.5 + box_height = float(box[3] - box[1]) + if center_x > 0.80 * width: + rejected["right_reflection_roi"] += 1 + continue + if box_height < 0.20 * height: + rejected["too_short"] += 1 + continue + kept.append((box.tolist(), float(score))) + kept.sort(key=lambda item: item[1], reverse=True) + if len(kept) > max_people: + rejected["over_max_people"] += len(kept) - max_people + kept = kept[:max_people] + kept.sort(key=lambda item: item[0][0]) + return [item[0] for item in kept], [item[1] for item in kept], rejected + + +def process_video( + video_path: Path, + output_dir: Path, + detector: YOLO, + pose: RTMPose, + det_device: str, + det_imgsz: int, + det_conf: float, + det_roi_y_max_ratio: float, + short_occlusion_frames: int, + medium_occlusion_frames: int, + max_people: int, + video_label: str, +) -> dict: + output_dir.mkdir(parents=True, exist_ok=True) + cap = cv2.VideoCapture(str(video_path)) + fps = float(cap.get(cv2.CAP_PROP_FPS)) + width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)) + height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)) + writer = cv2.VideoWriter( + str(output_dir / "raw_temporal_occlusion_v4.mp4"), + cv2.VideoWriter_fourcc(*"mp4v"), + fps, + (1920, 360), + ) + + box_tracker = BoxTracker(max_age=30) + people: dict[int, LivePersonState] = {} + records = [] + detection_counts = Counter() + rejection_counts = Counter() + state_counts = Counter() + timings = defaultdict(list) + isolated_inferences = 0 + frame_idx = 0 + + while True: + ok, frame = cap.read() + if not ok: + break + started = time.perf_counter() + detection_height = max(1, min(height, round(height * det_roi_y_max_ratio))) + detection = detector.predict( + frame[:detection_height], + imgsz=det_imgsz, + conf=det_conf, + iou=0.30, + classes=[0], + device=det_device, + verbose=False, + )[0] + timings["detector_ms"].append((time.perf_counter() - started) * 1000.0) + boxes, confidences, rejected = filter_person_boxes( + detection, width, height, max_people + ) + rejection_counts.update(rejected) + detection_counts[len(boxes)] += 1 + track_ids = box_tracker.update(boxes, frame_idx, (height, width)) + + if boxes: + started = time.perf_counter() + raw_keypoints, raw_scores = pose(frame, bboxes=boxes) + timings["raw_pose_ms"].append((time.perf_counter() - started) * 1000.0) + isolated_keypoints = raw_keypoints.copy() + isolated_scores = raw_scores.copy() + started = time.perf_counter() + for person_idx, bbox in enumerate(boxes): + target_area = max((bbox[2] - bbox[0]) * (bbox[3] - bbox[1]), 1.0) + needs_isolation = any( + idx != person_idx + and max((other[2] - other[0]) * (other[3] - other[1]), 1.0) + < 0.5 * target_area + for idx, other in enumerate(boxes) + ) + if not needs_isolation: + continue + masked = mask_other_people(frame, boxes, person_idx) + kpts, scores = pose(masked, bboxes=[bbox]) + isolated_keypoints[person_idx] = kpts[0] + isolated_scores[person_idx] = scores[0] + isolated_inferences += 1 + timings["isolated_extra_ms"].append((time.perf_counter() - started) * 1000.0) + else: + raw_keypoints = np.empty((0, 133, 2), dtype=np.float64) + raw_scores = np.empty((0, 133), dtype=np.float64) + isolated_keypoints = raw_keypoints.copy() + isolated_scores = raw_scores.copy() + + selected_list = [] + smoothed_list = [] + fusion_meta_list = [] + person_states = [] + for person_idx, (bbox, track_id) in enumerate(zip(boxes, track_ids)): + state = people.setdefault(track_id, LivePersonState()) + selected, fusion_meta = fuse_person( + raw_keypoints[person_idx], + raw_scores[person_idx], + isolated_keypoints[person_idx], + isolated_scores[person_idx], + bbox, + state.fusion, + ) + smoothed = adaptive_smooth(selected, bbox, state.fusion) + selected_list.append(selected) + smoothed_list.append(smoothed) + fusion_meta_list.append(fusion_meta) + person_states.append(state) + + frame_measurements = {} + endpoint_tracks = {} + for person_idx, (selected, fusion_meta, state) in enumerate( + zip(selected_list, fusion_meta_list, person_states) + ): + for side, (start, end) in HAND_SLICES.items(): + meta = fusion_meta[side] + frame_measurements[(person_idx, side)] = ( + Measurement( + keypoints=selected[start:end], + precise=bool(meta["distance_eligible"]), + provisional=not bool(meta["distance_eligible"]), + quality=float(meta.get("selected_hand_score", 0.0)), + ) + if meta["alarm_eligible"] + else None + ) + endpoint_tracks[(person_idx, side)] = state.endpoints[side] + frame_measurements, assignment_events = resolve_global_hand_assignments( + frame_measurements, + endpoint_tracks, + boxes, + medium_occlusion_frames, + ) + + v2_list = [] + v2_meta_list = [] + for person_idx, (bbox, selected, smoothed, fusion_meta, state) in enumerate( + zip( + boxes, + selected_list, + smoothed_list, + fusion_meta_list, + person_states, + ) + ): + v2, v2_meta = update_live_v2( + frame_idx, + bbox, + selected, + smoothed, + fusion_meta, + state, + { + side: frame_measurements[(person_idx, side)] + for side in ("left", "right") + }, + { + side: assignment_events.get((person_idx, side)) + for side in ("left", "right") + }, + short_occlusion_frames, + medium_occlusion_frames, + ) + v2_list.append(v2) + v2_meta_list.append(v2_meta) + for side in ("left", "right"): + state_counts[f"{v2_meta[side]['state']}"] += 1 + + global_hands = build_global_hands(selected_list, fusion_meta_list) + records.append( + { + "frame_index": frame_idx, + "timestamp_ms": round(frame_idx / fps * 1000.0, 3), + "track_ids": track_ids, + "boxes": np.round(np.asarray(boxes), 3).tolist() if boxes else [], + "det_confidences": [round(value, 4) for value in confidences], + "raw_keypoints": np.round(raw_keypoints, 3).tolist(), + "temporal_selected_keypoints": np.round(np.asarray(selected_list), 3).tolist(), + "temporal_keypoints": np.round(np.asarray(smoothed_list), 3).tolist(), + "temporal_metadata": fusion_meta_list, + "v2_keypoints": np.round(np.asarray(v2_list), 3).tolist(), + "v2_metadata": v2_meta_list, + "global_hands": global_hands, + } + ) + + raw_panel = draw_pose( + frame, + boxes, + raw_keypoints, + f"{video_label} RTMW-L raw | frame {frame_idx}", + ) + temporal_panel = draw_fused( + frame, + boxes, + smoothed_list, + fusion_meta_list, + f"{video_label} temporal fusion | frame {frame_idx}", + ) + v2_panel = draw_v2( + frame, + boxes, + np.asarray(v2_list), + v2_meta_list, + f"{video_label} occlusion V4 | frame {frame_idx}", + show_reachable=False, + ) + panels = [ + cv2.resize(panel, (640, 360), interpolation=cv2.INTER_AREA) + for panel in (raw_panel, temporal_panel, v2_panel) + ] + writer.write(np.hstack(panels)) + if frame_idx == 0 and video_path.suffix.lower() in { + ".jpg", + ".jpeg", + ".png", + ".bmp", + ".webp", + }: + cv2.imwrite(str(output_dir / "raw_pose.jpg"), raw_panel) + cv2.imwrite(str(output_dir / "spatial_fusion.jpg"), temporal_panel) + cv2.imwrite(str(output_dir / "occlusion_v2.jpg"), v2_panel) + panel_width = 960 + panel_height = round(height * panel_width / width) + image_panels = [ + cv2.resize( + panel, + (panel_width, panel_height), + interpolation=cv2.INTER_AREA, + ) + for panel in (raw_panel, temporal_panel, v2_panel) + ] + cv2.imwrite( + str(output_dir / "raw_spatial_occlusion_v2.jpg"), + np.hstack(image_panels), + ) + if frame_idx % 100 == 0: + print(f"{video_label}: processed frame {frame_idx}", flush=True) + frame_idx += 1 + + cap.release() + writer.release() + with (output_dir / "predictions.jsonl").open("w", encoding="utf-8") as handle: + for record in records: + handle.write(json.dumps(record, ensure_ascii=False) + "\n") + + summary = { + "occlusion_algorithm": "v4_global_assignment", + "video": str(video_path.resolve()), + "frames": len(records), + "fps": round(fps, 3), + "resolution": [width, height], + "person_filter": { + "center_x_max_ratio": 0.80, + "minimum_box_height_ratio": 0.20, + "max_people": max_people, + }, + "detector_settings": { + "imgsz": det_imgsz, + "confidence": det_conf, + "roi_y_max_ratio": det_roi_y_max_ratio, + }, + "filtered_person_count_frames": dict(sorted(detection_counts.items())), + "rejected_detection_counts": dict(sorted(rejection_counts.items())), + "track_ids_created": box_tracker.next_id - 1, + "isolated_pose_inferences": isolated_inferences, + "short_occlusion_frames": short_occlusion_frames, + "medium_occlusion_frames": medium_occlusion_frames, + "v2_state_counts": dict(sorted(state_counts.items())), + "timing_ms_cpu": { + name: summarize(values[5:] if len(values) > 5 else values) + for name, values in timings.items() + }, + "output_note": ( + "Main video hides reachable-region arcs. Reachable-region geometry remains in " + "predictions.jsonl for later door-distance lower-bound calculation." + ), + } + with (output_dir / "summary.json").open("w", encoding="utf-8") as handle: + json.dump(summary, handle, ensure_ascii=False, indent=2) + return summary + + +def main() -> None: + from rtmlib import RTMPose + from ultralytics import YOLO + args = parse_args() + detector = YOLO(str(args.det_model)) + pose = RTMPose( + str(args.pose_model), + model_input_size=(288, 384), + backend="onnxruntime", + device=args.pose_device, + ) + all_summaries = {} + for video_idx, video_path in enumerate(args.videos, start=1): + label = f"video{video_idx}" + output_dir = args.output / label + all_summaries[label] = process_video( + video_path, + output_dir, + detector, + pose, + args.det_device, + args.det_imgsz, + args.det_conf, + args.det_roi_y_max_ratio, + args.short_occlusion_frames, + args.medium_occlusion_frames, + args.max_people, + label, + ) + args.output.mkdir(parents=True, exist_ok=True) + with (args.output / "summary.json").open("w", encoding="utf-8") as handle: + json.dump(all_summaries, handle, ensure_ascii=False, indent=2) + print(json.dumps(all_summaries, ensure_ascii=False, indent=2)) + + +if __name__ == "__main__": + main() diff --git a/orin/pipeline/rtmw_l_person_mask_ablation.py b/orin/pipeline/rtmw_l_person_mask_ablation.py new file mode 100644 index 0000000..6263500 --- /dev/null +++ b/orin/pipeline/rtmw_l_person_mask_ablation.py @@ -0,0 +1,166 @@ +"""Test whether other people inside a top-down crop cause RTMW-L mix-ups.""" + +from __future__ import annotations + +import argparse +import json +import time +from pathlib import Path + +import cv2 +import numpy as np + +from compare_rtmw_ml import ( + draw_pose, + model_disagreement, + summarize, + temporal_metrics, + write_montage, +) + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser() + parser.add_argument("--video", type=Path, required=True) + parser.add_argument("--baseline-jsonl", type=Path, required=True) + parser.add_argument("--model-l", type=Path, required=True) + parser.add_argument("--output", type=Path, required=True) + return parser.parse_args() + + +def mask_other_people( + frame: np.ndarray, + boxes: list[list[float]], + target_idx: int, +) -> np.ndarray: + masked = frame.copy() + height, width = frame.shape[:2] + mask = np.zeros((height, width), dtype=bool) + target = boxes[target_idx] + target_area = max((target[2] - target[0]) * (target[3] - target[1]), 1.0) + for idx, bbox in enumerate(boxes): + if idx == target_idx: + continue + other_area = max((bbox[2] - bbox[0]) * (bbox[3] - bbox[1]), 1.0) + # The failure in this video is the small background person being + # absorbed by the much larger foreground person's 3:4 pose crop. + # Mask only markedly smaller people. Masking a larger occluder while + # estimating a small person removes genuine target pixels at overlap. + if other_area >= 0.5 * target_area: + continue + x1, y1, x2, y2 = (int(round(value)) for value in bbox) + x1, x2 = max(0, x1), min(width - 1, x2) + y1, y2 = max(0, y1), min(height - 1, y2) + mask[y1:y2 + 1, x1:x2 + 1] = True + + masked[mask] = (124, 116, 104) + return masked + + +def main() -> None: + from rtmlib import RTMPose + args = parse_args() + args.output.mkdir(parents=True, exist_ok=True) + baseline = [ + json.loads(line) + for line in args.baseline_jsonl.read_text(encoding="utf-8").splitlines() + ] + pose_l = RTMPose( + str(args.model_l), model_input_size=(288, 384), backend="onnxruntime", device="cpu" + ) + + cap = cv2.VideoCapture(str(args.video)) + fps = float(cap.get(cv2.CAP_PROP_FPS)) + width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)) + height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)) + writer = cv2.VideoWriter( + str(args.output / "rtmw_l_raw_vs_person_masked.mp4"), + cv2.VideoWriter_fourcc(*"mp4v"), + fps, + (width, height // 2), + ) + + comparison: list[dict] = [] + timings: list[float] = [] + frame_idx = 0 + while True: + ok, frame = cap.read() + if not ok: + break + baseline_record = baseline[frame_idx] + boxes = baseline_record["boxes"][:2] + masked_keypoints = [] + masked_scores = [] + started = time.perf_counter() + for person_idx, bbox in enumerate(boxes): + masked_frame = mask_other_people(frame, boxes, person_idx) + keypoints, scores = pose_l(masked_frame, bboxes=[bbox]) + masked_keypoints.append(keypoints[0]) + masked_scores.append(scores[0]) + timings.append((time.perf_counter() - started) * 1000.0) + masked_keypoints_array = np.asarray(masked_keypoints) + + raw_keypoints = np.asarray(baseline_record["L"]["keypoints"][:2]) + comparison.append( + { + "frame_index": frame_idx, + "boxes": boxes, + "M": {"keypoints": raw_keypoints.tolist()}, + "L": {"keypoints": np.round(masked_keypoints_array, 3).tolist()}, + "masked_raw_scores": np.round(np.asarray(masked_scores), 4).tolist(), + } + ) + + left = draw_pose(frame, boxes, raw_keypoints, f"RTMW-L raw | frame {frame_idx}") + right = draw_pose( + frame, + boxes, + masked_keypoints_array, + f"RTMW-L mask other person | frame {frame_idx}", + ) + left = cv2.resize(left, (width // 2, height // 2), interpolation=cv2.INTER_AREA) + right = cv2.resize(right, (width // 2, height // 2), interpolation=cv2.INTER_AREA) + writer.write(np.hstack([left, right])) + if frame_idx % 50 == 0: + print(f"processed frame {frame_idx}", flush=True) + frame_idx += 1 + + cap.release() + writer.release() + + with (args.output / "rtmw_l_raw_vs_person_masked.jsonl").open( + "w", encoding="utf-8" + ) as handle: + for record in comparison: + handle.write(json.dumps(record, ensure_ascii=False) + "\n") + + disagreement, worst_frames = model_disagreement(comparison) + fixed_frames = [0, 60, 120, 132, 136, 180, 240, 252, 300, 360, 411] + montage_frames = sorted(set(fixed_frames + worst_frames[:8])) + write_montage( + args.video, + comparison, + montage_frames, + args.output / "rtmw_l_raw_vs_person_masked_montage.jpg", + titles=("RTMW-L raw", "RTMW-L smaller-person masked"), + ) + + summary = { + "frames": len(comparison), + "labels": { + "M": "RTMW-L raw crop", + "L": "RTMW-L with markedly smaller people masked", + }, + "timing_ms_two_people_cpu": summarize(timings[5:]), + "temporal_metrics_normalized_to_yolo_box": temporal_metrics(comparison), + "position_change_after_masking": disagreement, + "largest_change_frames": worst_frames, + "montage_frames": montage_frames, + } + with (args.output / "summary.json").open("w", encoding="utf-8") as handle: + json.dump(summary, handle, ensure_ascii=False, indent=2) + print(json.dumps(summary, ensure_ascii=False, indent=2)) + + +if __name__ == "__main__": + main() diff --git a/orin/pipeline/rtmw_risk_pipeline.py b/orin/pipeline/rtmw_risk_pipeline.py new file mode 100644 index 0000000..a61a490 --- /dev/null +++ b/orin/pipeline/rtmw_risk_pipeline.py @@ -0,0 +1,390 @@ +"""Lightweight causal door-risk inference using only YOLO11s and RTMW-L.""" + +from __future__ import annotations + +import argparse +import json +import sys +import time +from pathlib import Path + +import cv2 +import numpy as np + + +PACKAGE_DIR = Path(__file__).resolve().parents[1] +PIPELINE_DIR = Path(__file__).resolve().parent +for import_dir in (PIPELINE_DIR, PACKAGE_DIR): + if str(import_dir) not in sys.path: + sys.path.insert(0, str(import_dir)) + +from door_mask_risk import DoorMaskSequence, DoorRiskEvaluator # noqa: E402 +from door_zone_risk import ( # noqa: E402 + DoorZoneAssigner, + DoorZoneConfig, + build_door_risks, +) +from live_pose_engine import LivePoseEngine # noqa: E402 + + +def parse_image_size(value: str) -> int | tuple[int, int]: + parts = value.lower().replace(",", "x").split("x") + if len(parts) == 1: + return int(parts[0]) + if len(parts) == 2: + return int(parts[0]), int(parts[1]) + raise argparse.ArgumentTypeError( + "image size must be one integer or HEIGHTxWIDTH, for example 960 or 544x960" + ) + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--video", type=Path, required=True) + door_source = parser.add_mutually_exclusive_group(required=True) + door_source.add_argument("--door-mask-dir", type=Path) + door_source.add_argument( + "--constant-door-mask", + type=Path, + help="One X-AnyLabeling JSON or raster mask reused for every video frame.", + ) + parser.add_argument("--output", type=Path, required=True) + parser.add_argument( + "--det-model", type=Path, default=PACKAGE_DIR / "models" / "yolo11s.onnx" + ) + parser.add_argument( + "--pose-model", type=Path, default=PACKAGE_DIR / "models" / "rtmw-l-384.onnx" + ) + parser.add_argument( + "--pose-jsonl", + type=Path, + help="Diagnostic replay input; bypasses live YOLO/RTMW model execution.", + ) + parser.add_argument("--det-device", default="0") + parser.add_argument("--pose-device", default="cuda") + parser.add_argument( + "--pose-provider", + choices=("cuda", "tensorrt"), + default="cuda", + help="ONNX Runtime provider for RTMW.", + ) + parser.add_argument( + "--pose-trt-cache", + type=Path, + default=PACKAGE_DIR / "models" / "rtmw_l_trt_cache", + ) + parser.add_argument("--det-imgsz", type=parse_image_size, default=960) + parser.add_argument("--det-conf", type=float, default=0.30) + parser.add_argument("--det-roi-y-max-ratio", type=float, default=1.0) + parser.add_argument("--max-people", type=int, default=3) + parser.add_argument( + "--door-mask-format", + choices=("auto", "xanylabeling", "raster"), + default="auto", + ) + parser.add_argument("--door-label", action="append") + parser.add_argument("--door-frame-offset", type=int, default=0) + parser.add_argument("--require-door-mask", action="store_true") + parser.add_argument("--risk-warning-relative", type=float, default=0.30) + parser.add_argument("--risk-release-relative", type=float, default=0.45) + parser.add_argument( + "--active-door", + choices=("front", "rear", "both", "none"), + default="both", + help="Door(s) whose CAN state is 0. Use front/rear for deployment simulation.", + ) + parser.add_argument( + "--door-zone-config", + type=Path, + default=PACKAGE_DIR / "annotations" / "door_zone_multifeature.json", + help="Pose-primary fixed-camera front/rear calibration JSON.", + ) + parser.add_argument("--max-frames", type=int, default=0) + parser.add_argument("--write-video", action="store_true") + parser.add_argument("--progress-every", type=int, default=100) + return parser.parse_args() + + +def load_pose_records(path: Path) -> dict[int, dict]: + records = {} + with path.open("r", encoding="utf-8") as handle: + for line in handle: + if not line.strip(): + continue + record = json.loads(line) + records[int(record["frame_index"])] = record + return records + + +def percentile(values: list[float], p: float) -> float | None: + return round(float(np.percentile(values, p)), 3) if values else None + + +def timing_stats(values: list[float]) -> dict: + return { + "count": len(values), + "total": round(float(np.sum(values)), 3) if values else 0.0, + "mean": round(float(np.mean(values)), 3) if values else None, + "median": percentile(values, 50), + "p95": percentile(values, 95), + } + + +def visualization_skeletons(record: dict) -> list[dict]: + """Return the exact fused keypoints and metadata consumed by draw_fused.""" + + boxes = record.get("boxes", []) + track_ids = [int(value) for value in record.get("track_ids", [])] + keypoints = record.get("temporal_selected_keypoints", []) + metadata = record.get("temporal_metadata", []) + if not (len(boxes) == len(keypoints) == len(metadata)): + return [] + return [ + { + "track_id": track_ids[position] if position < len(track_ids) else position + 1, + "person_index": position, + "layout": "coco_wholebody_133", + "keypoints": person_keypoints, + "metadata": metadata[position], + } + for position, person_keypoints in enumerate(keypoints) + ] + + +def render_frame( + frame: np.ndarray, + record: dict, + frame_risk: dict, + door_mask: np.ndarray | None, + door_assignments: list[dict], +) -> np.ndarray: + canvas = frame.copy() + boxes = record.get("boxes", []) + keypoints = record.get("temporal_selected_keypoints", []) + metadata = record.get("temporal_metadata", []) + if len(boxes) == len(keypoints) == len(metadata) and boxes: + from fuse_rtmw_l_temporal import draw_fused + + canvas = draw_fused(canvas, boxes, keypoints, metadata, "YOLO + RTMW risk") + else: + for position, box in enumerate(boxes): + x1, y1, x2, y2 = [int(round(float(value))) for value in box[:4]] + cv2.rectangle(canvas, (x1, y1), (x2, y2), (0, 220, 255), 2) + cv2.putText( + canvas, + f"P{position + 1}", + (x1 + 4, max(22, y1 + 24)), + cv2.FONT_HERSHEY_SIMPLEX, + 0.7, + (0, 220, 255), + 2, + cv2.LINE_AA, + ) + zone_colors = { + "front": (0, 220, 255), + "rear": (255, 180, 0), + "transition": (0, 140, 255), + "unassigned": (150, 150, 150), + } + for assignment in door_assignments: + position = int(assignment["person_index"]) + if not 0 <= position < len(boxes): + continue + box = boxes[position] + x1, y1 = int(round(float(box[0]))), int(round(float(box[1]))) + zone = str(assignment["door_zone"]) + cv2.putText( + canvas, + f"T{assignment['track_id']} {zone.upper()}", + (max(0, x1 + 4), max(24, y1 - 8)), + cv2.FONT_HERSHEY_SIMPLEX, + 0.68, + zone_colors.get(zone, (160, 160, 160)), + 2, + cv2.LINE_AA, + ) + if door_mask is not None and np.any(door_mask): + overlay = np.zeros_like(canvas) + overlay[door_mask > 0] = (255, 180, 0) + canvas = cv2.addWeighted(canvas, 1.0, overlay, 0.22, 0) + + score = frame_risk.get("risk_score") + score_text = "N/A" if score is None else f"{float(score):.3f}" + level = str(frame_risk.get("risk_level", "unavailable")) + color = (0, 0, 255) if level == "danger" else (0, 180, 255) + if level in ("clear", "unavailable", "inactive", "unknown"): + color = (40, 210, 80) if level == "clear" else (160, 160, 160) + cv2.rectangle(canvas, (12, 55), (490, 95), (0, 0, 0), -1) + cv2.putText( + canvas, + f"{frame_risk.get('risk_type', 'DOOR_RISK')} score={score_text} level={level}", + (24, 83), + cv2.FONT_HERSHEY_SIMPLEX, + 0.68, + color, + 2, + cv2.LINE_AA, + ) + return canvas + + +def main() -> None: + args = parse_args() + args.output.mkdir(parents=True, exist_ok=True) + if args.pose_provider == "tensorrt" and not args.pose_device.lower().startswith("cuda"): + raise ValueError("--pose-provider tensorrt requires --pose-device cuda") + + capture = cv2.VideoCapture(str(args.video.resolve())) + if not capture.isOpened(): + raise RuntimeError(f"Cannot open video: {args.video}") + input_fps = float(capture.get(cv2.CAP_PROP_FPS)) or 30.0 + width = int(capture.get(cv2.CAP_PROP_FRAME_WIDTH)) + height = int(capture.get(cv2.CAP_PROP_FRAME_HEIGHT)) + + risk_evaluator = DoorRiskEvaluator( + DoorMaskSequence( + args.door_mask_dir, + mask_format=args.door_mask_format, + labels=tuple(args.door_label or ("door", "车门")), + frame_offset=args.door_frame_offset, + require_every_frame=args.require_door_mask, + constant_mask=args.constant_door_mask, + ), + warning_relative=args.risk_warning_relative, + release_relative=args.risk_release_relative, + ) + zone_assigner = DoorZoneAssigner(DoorZoneConfig.load(args.door_zone_config)) + print( + f"Door masks: {risk_evaluator.sequence.source_format} " + + ( + f"(constant: {risk_evaluator.sequence.constant_path})" + if risk_evaluator.sequence.constant_path is not None + else f"({len(risk_evaluator.sequence.paths)} indexed frames)" + ) + ) + + pose_records = load_pose_records(args.pose_jsonl) if args.pose_jsonl else None + load_started = time.perf_counter() + pose_engine = None if pose_records is not None else LivePoseEngine(args) + model_load_seconds = time.perf_counter() - load_started + + writer = None + if args.write_video: + writer = cv2.VideoWriter( + str(args.output / "rtmw_risk_overlay.mp4"), + cv2.VideoWriter_fourcc(*"mp4v"), + input_fps, + (width, height), + ) + + frame_latencies = [] + pose_latencies = [] + risk_latencies = [] + render_latencies = [] + started_all = time.perf_counter() + frame_index = 0 + output_path = args.output / "causal_results.jsonl" + with output_path.open("w", encoding="utf-8") as output_handle: + while True: + if args.max_frames > 0 and frame_index >= args.max_frames: + break + ok, frame = capture.read() + if not ok: + break + frame_started = time.perf_counter() + pose_started = time.perf_counter() + if pose_records is None: + record = pose_engine.infer(frame, frame_index) + else: + record = pose_records.get(frame_index) + if record is None: + raise RuntimeError(f"Pose JSONL has no record for frame {frame_index}") + pose_latencies.append((time.perf_counter() - pose_started) * 1000) + + risk_started = time.perf_counter() + assignments = zone_assigner.update( + frame_index, record, (height, width) + ) + raw_risk, door_mask = risk_evaluator.evaluate( + frame_index, record, (height, width) + ) + door_risks, frame_risk = build_door_risks( + raw_risk, assignments, args.active_door + ) + risk_latencies.append((time.perf_counter() - risk_started) * 1000) + output_record = { + "frame_index": frame_index, + "track_ids": record.get("track_ids", []), + "boxes": record.get("boxes", []), + "skeletons": visualization_skeletons(record), + "active_door": args.active_door, + "door_assignments": assignments, + "door_risks": door_risks, + "risk": frame_risk, + "ungated_risk": raw_risk, + } + output_handle.write(json.dumps(output_record, ensure_ascii=False) + "\n") + + if writer is not None: + render_started = time.perf_counter() + writer.write( + render_frame(frame, record, frame_risk, door_mask, assignments) + ) + render_latencies.append((time.perf_counter() - render_started) * 1000) + + frame_latencies.append((time.perf_counter() - frame_started) * 1000) + if args.progress_every > 0 and frame_index % args.progress_every == 0: + recent = frame_latencies[-min(len(frame_latencies), args.progress_every) :] + print( + f"frame={frame_index} recent_ms={np.mean(recent):.1f} " + f"risk={frame_risk['risk_level']}", + flush=True, + ) + frame_index += 1 + + capture.release() + if writer is not None: + writer.release() + elapsed = time.perf_counter() - started_all + summary = { + "mode": "yolo_rtmw_door_risk_multifeature_v9_10", + "video": str(args.video.resolve()), + "pose_source": "jsonl_diagnostic" if pose_records is not None else "resident_yolo_rtmw", + "models_loaded": [] if pose_records is not None else ["YOLO11s", "RTMW-L"], + "edgetam_loaded": False, + "risk_region_policy": ( + "RTMW alarm-eligible hand hulls; YOLO body box only when a person " + "has no usable RTMW hand region" + ), + "frames": frame_index, + "input_fps": round(input_fps, 3), + "elapsed_seconds": round(elapsed, 3), + "throughput_fps": round(frame_index / elapsed, 3) if elapsed else None, + "model_load_seconds": round(model_load_seconds, 3), + "latency_ms": { + "total": timing_stats(frame_latencies), + "pose": timing_stats(pose_latencies), + "risk": timing_stats(risk_latencies), + "render": timing_stats(render_latencies), + }, + "write_video": args.write_video, + "active_door": args.active_door, + "door_zone": zone_assigner.summary(), + "door_risk": risk_evaluator.summary(), + } + if pose_engine is not None: + summary["pose"] = { + "provider": args.pose_provider, + "isolated_inferences": pose_engine.isolated_inferences, + "timing_ms": { + key: timing_stats(values) for key, values in pose_engine.timings.items() + }, + } + (args.output / "summary.json").write_text( + json.dumps(summary, ensure_ascii=False, indent=2), encoding="utf-8" + ) + print(json.dumps(summary, ensure_ascii=False, indent=2)) + + +if __name__ == "__main__": + main() diff --git a/orin/run.sh b/orin/run.sh index c153d56..7f8c795 100755 --- a/orin/run.sh +++ b/orin/run.sh @@ -1,4 +1,7 @@ #!/usr/bin/env bash set -euo pipefail -root="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)" -exec "$root/build/camera_risk" --config "$root/config/camera.json" "$@" +root="$(cd "$(dirname "$0")" && pwd)" +export OPENBLAS_NUM_THREADS=1 OMP_NUM_THREADS=2 +export LD_LIBRARY_PATH="$root/lib${LD_LIBRARY_PATH:+:$LD_LIBRARY_PATH}" +export PYTHONDONTWRITEBYTECODE=1 +exec "${PYTHON:-python3}" -u "$root/run_camera.py" "$@" diff --git a/orin/run_camera.py b/orin/run_camera.py new file mode 100644 index 0000000..576304c --- /dev/null +++ b/orin/run_camera.py @@ -0,0 +1,299 @@ +#!/usr/bin/env python3 +"""Live baseline V4L2/CAN capture + native TensorRT + CMS v9.10 risk.""" +from __future__ import annotations +import argparse +from collections import Counter, deque +from contextlib import ExitStack +import copy +import json +import signal +import sys +import threading +import time +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer +from pathlib import Path +import cv2 +import numpy as np +ROOT=Path(__file__).resolve().parent +sys.path.insert(0,str(ROOT/'pipeline')) +from native_backend import NativeBackend, NativePoseEngine +from latest_result import LatestResult, result_to_dict +from door_mask_risk import DoorMaskSequence, DoorRiskEvaluator +from door_zone_risk import DoorZoneConfig, DoorZoneAssigner, build_door_risks +from rtmw_risk_pipeline import render_frame, visualization_skeletons, timing_stats + + +def select_active_door(mode, snapshot, side='left', timeout_ms=1000): + if mode!='can': return mode, True + p=snapshot['positions'];offset=0 if side=='left' else 1 + if snapshot['age_ms']<0 or snapshot['age_ms']>timeout_ms or min(p[offset],p[offset+2])<0: + return 'none',False + front,rear=p[offset]==0,p[offset+2]==0 + return ('both' if front and rear else 'front' if front else 'rear' if rear else 'none'),True + + +def gate_risks(raw, assignments, active, can_valid): + doors,risk=build_door_risks(raw,assignments,active) + if not can_valid: + for door,item in doors.items(): + item.update(risk_type=door.upper()+'_UNKNOWN',risk_level='unknown',risk_score=None,decision='CAN_UNAVAILABLE_OR_STALE') + risk=dict(risk,risk_type='CAN_UNKNOWN',risk_level='unknown',risk_score=None,decision='CAN_UNAVAILABLE_OR_STALE') + return doors,risk + + +class JsonlWriter: + """Bound the disk usage of continuous camera runs (current + two backups).""" + def __init__(self,path,max_bytes): + self.path=path;self.max_bytes=max_bytes;self.size=0 + self.handle=path.open('w',encoding='utf-8',buffering=1) + def __enter__(self):return self + def __exit__(self,*args):self.close() + def close(self):self.handle.close() + def write(self,text): + size=len(text.encode('utf-8')) + if self.max_bytes and self.size and self.size+size>self.max_bytes: + self.handle.close() + backup=self.path.with_suffix(self.path.suffix+'.1') + if backup.exists():backup.replace(self.path.with_suffix(self.path.suffix+'.2')) + self.path.replace(backup) + self.handle=self.path.open('w',encoding='utf-8',buffering=1);self.size=0 + self.handle.write(text);self.size+=size + + +def render_output(frame, record, risk, mask, assignments, status): + canvas = render_frame(frame, record, risk, mask, assignments) + cv2.putText(canvas,status,(24,frame.shape[0]-30),cv2.FONT_HERSHEY_SIMPLEX,.8,(0,180,255),2) + return canvas + + +class Preview: + def __init__(self,host,port,results): + self.jpeg=None;self.results=results;self.lock=threading.Lock() + self.pending=threading.Condition();self.packet=None;self.stopped=False + outer=self + class Handler(BaseHTTPRequestHandler): + def do_GET(self): + with outer.lock: + if self.path.split('?')[0]=='/frame.jpg': payload=outer.jpeg or b'';mime='image/jpeg' + elif self.path=='/result.json':payload=json.dumps(result_to_dict(outer.results.load()) or {},ensure_ascii=False,allow_nan=False).encode();mime='application/json' + else: + payload=b'''Vision anti-pinch 0923

Vision anti-pinch 0923 / Live camera

''';mime='text/html; charset=utf-8' + self.send_response(200 if payload else 503);self.send_header('Content-Type',mime);self.send_header('Cache-Control','no-store');self.send_header('Content-Length',str(len(payload)));self.end_headers() + try:self.wfile.write(payload) + except (BrokenPipeError,ConnectionResetError):pass + def log_message(self,*args):pass + self.server=ThreadingHTTPServer((host,port),Handler) + self.thread=threading.Thread(target=self.server.serve_forever,daemon=True);self.thread.start() + self.renderer=threading.Thread(target=self._render_loop,name='risk-preview',daemon=True) + self.renderer.start() + def submit(self, frame, record, risk, mask, assignments, status): + # All inputs belong to a finished frame; inference never modifies them + # after submission. Retain only one pending frame, replacing older work. + with self.pending: + if not self.stopped: + self.packet=(frame,record,risk,mask,assignments,status) + self.pending.notify() + + def _render_loop(self): + try: + while True: + with self.pending: + self.pending.wait_for(lambda:self.packet is not None or self.stopped) + if self.stopped:return + packet=self.packet;self.packet=None + self.update(render_output(*packet)) + except Exception as exc: + with self.pending:self.stopped=True;self.packet=None + with self.lock:self.jpeg=None + print('Visualization disabled; inference continues: '+str(exc),file=sys.stderr,flush=True) + + def update(self,image): + w=960;h=round(image.shape[0]*w/image.shape[1]);image=cv2.resize(image,(w,h)) + ok,jpeg=cv2.imencode('.jpg',image,[cv2.IMWRITE_JPEG_QUALITY,80]) + if ok: + with self.lock:self.jpeg=jpeg.tobytes() + def close(self): + with self.pending:self.stopped=True;self.packet=None;self.pending.notify_all() + self.renderer.join() + self.server.shutdown();self.server.server_close();self.thread.join() + + +def parse_args(argv=None): + p=argparse.ArgumentParser(description=__doc__) + p.add_argument('--video',type=Path,help='Diagnostic video replay through the same TensorRT/risk chain');p.add_argument('--camera',default='/dev/video5');p.add_argument('--width',type=int,default=1920);p.add_argument('--height',type=int,default=1536);p.add_argument('--fps',type=int,default=30) + p.add_argument('--engines',type=Path,default=ROOT/'models/engines') + p.add_argument('--constant-door-mask',type=Path,default=ROOT/'annotations/constant_door_mask.json') + p.add_argument('--door-zone-config',type=Path,default=ROOT/'annotations/door_zone_multifeature.json') + p.add_argument('--active-door',choices=['can','front','rear','both','none'],default='can') + p.add_argument('--can-side',choices=['left','right'],default='left');p.add_argument('--can-timeout-ms',type=float,default=1000);p.add_argument('--no-can',action='store_true') + p.add_argument('--max-people',type=int,default=3);p.add_argument('--det-conf',type=float,default=.30) + p.add_argument('--max-frames',type=int,default=0);p.add_argument('--duration',type=float,default=0) + p.add_argument('--jsonl-max-mb',type=float,default=32,help='Rotate each JSONL at this size; keep two backups; 0 disables rotation');p.add_argument('--output',type=Path,default=ROOT/'outputs'/time.strftime('%Y%m%d-%H%M%S')) + p.add_argument('--visualize',action='store_true',help='Enable optional HTTP visualization; disabled by default') + p.add_argument('--save-results',action='store_true',help='Save per-frame JSONL and periodic latest.json') + p.add_argument('--save-snapshots',action='store_true',help='Render and save latest.jpg every 10 seconds') + p.add_argument('--preview-port',type=int,default=8765);p.add_argument('--preview-host',default='0.0.0.0');p.add_argument('--preview-every',type=int,default=3) + p.add_argument('--write-video',action='store_true');p.add_argument('--save-pose-records',action='store_true');p.add_argument('--calibration-verified',action='store_true') + a=p.parse_args(argv) + if not 1<=a.max_people<=16 or not 0=a.max_frames) or (a.duration and time.monotonic()-run_started>=a.duration):break + if video: + ok,frame=video.read() + if not ok:break + capture={'sequence':count,'capture_monotonic_ms':time.monotonic()*1000,'dropped_before_inference':0,'latency_clock':'video_decode_complete'} + else:frame,capture=native.read() + if frame is None:break + tick=time.monotonic();record=pose.infer(frame,count);pose_ms.append((time.monotonic()-tick)*1000) + rt=time.monotonic();snapshot=native.can_snapshot();active,valid=select_active_door(a.active_door,snapshot,a.can_side,a.can_timeout_ms) + assignments=assigner.update(count,record,(a.height,a.width));raw,mask=evaluator.evaluate(count,record,(a.height,a.width)) + doors,risk=gate_risks(raw,assignments,active,valid);risk_ms.append((time.monotonic()-rt)*1000) + result={'frame_index':count,'track_ids':record['track_ids'],'boxes':record['boxes'],'skeletons':visualization_skeletons(record),'active_door':active,'door_assignments':assignments,'door_risks':doors,'risk':risk,'ungated_risk':raw,'capture':capture,'can':dict(snapshot,valid=valid,source='hardware' if a.active_door=='can' else 'simulation',side=a.can_side),'calibration_verified':a.calibration_verified} + result['capture_to_result_ms']=round(time.monotonic()*1000-capture['capture_monotonic_ms'],3) + results.publish(result) + if out:out.write(json.dumps(result,ensure_ascii=False,allow_nan=False)+'\n') + if pose_out:pose_out.write(json.dumps(record,ensure_ascii=False,allow_nan=False)+'\n') + now=time.monotonic() + snapshot_due=now-last_snapshot>=10 + if a.save_results and snapshot_due: + (a.output/'latest.json').write_text(json.dumps(result,ensure_ascii=False,indent=2)) + status=('CAN LIVE' if a.active_door=='can' else 'CAN SIMULATION: '+active)+' | '+('CALIBRATED' if a.calibration_verified else 'CALIBRATION UNVERIFIED') + if preview and count%a.preview_every==0: + preview.submit(frame,record,risk,mask,assignments,status) + if writer or (a.save_snapshots and snapshot_due): + canvas=render_output(frame,record,risk,mask,assignments,status) + if writer:writer.write(canvas) + if a.save_snapshots and snapshot_due: + cv2.imwrite(str(a.output/'latest.jpg'),canvas) + if snapshot_due:last_snapshot=now + count+=1;dropped+=capture['dropped_before_inference'];levels[risk['risk_level']]+=1;ages.append(result['capture_to_result_ms']);latencies.append((time.monotonic()-tick)*1000) + if now-last_report>=5: + print(json.dumps({'frames':count,'fps':round(count/(now-run_started),2),'people':len(record['boxes']),'risk':risk['risk_type'],'latency_ms':round(float(np.mean(latencies)),1),'capture_age_ms':round(float(np.mean(ages)),1),'can_valid':valid}),flush=True);last_report=now + finally: + if pose_out:pose_out.close() + except Exception as exc: + error=str(exc);raise + finally: + ended=time.monotonic() + if video:video.release() + if native:native.close() + if writer:writer.release() + if preview:preview.close() + elapsed=ended-(run_started or started) + summary={'mode':'baseline_v4l2_can_native_trt_cms_v9_10','camera':a.camera if not a.video else None,'video':str(a.video) if a.video else None,'frames':count,'elapsed_seconds':round(elapsed,3),'throughput_fps':round(count/max(elapsed,.001),3),'dropped_latest_frames':dropped,'risk_levels':dict(levels),'active_door_mode':a.active_door,'can_side':a.can_side,'calibration_verified':a.calibration_verified,'visualize_requested':a.visualize,'save_results':a.save_results,'save_snapshots':a.save_snapshots,'latency_window_frames':3000,'latency_ms':{'total':timing_stats(list(latencies)),'pose':timing_stats(list(pose_ms)),'risk':timing_stats(list(risk_ms)),'capture_to_result':timing_stats(list(ages))},'error':error} + if summary_sink is not None:summary_sink.update(summary) + if save_any:(a.output/'summary.json').write_text(json.dumps(summary,ensure_ascii=False,indent=2)) + print(json.dumps(summary,ensure_ascii=False),flush=True) + +class CameraPipeline: + """One lifecycle per instance, with the 0920 start/latest/stop/join pattern. + + The worker owns inference. Downstream threads read the latest immutable + snapshot directly, without JSON files, HTTP, callbacks or a result queue. + """ + def __init__(self, args=None): + self.args = copy.deepcopy(args if args is not None else parse_args([])) + self._results = LatestResult() + self._stop = threading.Event() + self._lifecycle_lock = threading.Lock() + self._thread = None + self._error = None + self.summary = {} + + @property + def latest_result(self): + """The latest complete result variable, or None before the first frame.""" + return self._results.load() + + def latest(self): + return self.latest_result + + def wait_for_result(self, after_frame_index=-1, timeout=None): + return self._results.wait(after_frame_index, timeout) + + def start(self): + with self._lifecycle_lock: + if self._thread is not None: + raise RuntimeError('CameraPipeline can only be started once') + self._thread = threading.Thread(target=self._run, name='vision-risk') + self._thread.start() + return self + + def _run(self): + try: + run_pipeline(self.args, self._results, self._stop, self.summary) + except BaseException as exc: + self._error = exc + finally: + self._results.close() + + def running(self): + return self._thread is not None and self._thread.is_alive() + + def request_stop(self): + self._stop.set() + + def join(self, timeout=None): + if self._thread is not None: + self._thread.join(timeout) + if not self._thread.is_alive() and self._error is not None: + raise self._error + + def __enter__(self): + return self.start() + + def __exit__(self, *args): + self.request_stop() + self.join() + + +def main(): + pipeline=CameraPipeline(parse_args()) + for sig in (signal.SIGINT,signal.SIGTERM): + signal.signal(sig,lambda *_:pipeline.request_stop()) + pipeline.start() + # Periodic joins allow Python to deliver stop signals on the main thread. + while pipeline.running(): + pipeline.join(.2) + pipeline.join() + + +if __name__=='__main__':main() diff --git a/orin/scripts/initialize_camera.sh b/orin/scripts/initialize_camera.sh index 5f70e1f..1383006 100755 --- a/orin/scripts/initialize_camera.sh +++ b/orin/scripts/initialize_camera.sh @@ -1,21 +1,5 @@ #!/usr/bin/env bash set -euo pipefail -root="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")/.." && pwd)" -device=/dev/video5 -loaded() { test -d "/sys/module/$1"; } -root_command() { if [[ $EUID -eq 0 ]]; then "$@"; else sudo -n "$@"; fi; } -if [[ -c "$device" ]] && loaded max9295 && loaded max9296 && loaded sgx_yuv_gmsl2; then - echo "Camera already initialized: $device" - exit 0 -fi -for driver in max9295 max9296 sgx-yuv-gmsl2; do - file="$root/bin/drivers/$driver.ko" - [[ -f "$file" ]] || { echo "Missing driver: $file" >&2; exit 1; } - [[ "$(modinfo -F vermagic "$file")" == "$(uname -r) "* ]] || { echo "Kernel mismatch: $file" >&2; exit 1; } -done -loaded max9295 || root_command insmod "$root/bin/drivers/max9295.ko" -loaded max9296 || root_command insmod "$root/bin/drivers/max9296.ko" -loaded sgx_yuv_gmsl2 || root_command insmod "$root/bin/drivers/sgx-yuv-gmsl2.ko" gmsl_mode=2,2,2,2 -for _ in {1..20}; do [[ -c "$device" ]] && break; sleep 0.25; done -[[ -c "$device" ]] || { echo "$device was not created; verify camera/DTB" >&2; exit 1; } -v4l2-ctl -d "$device" -c sensor_mode=1,trig_pin=0xffff0007 +root="$(cd "$(dirname "$0")/.." && pwd)" +export DRIVER_DIR="$root/bin/drivers" +exec "$root/scripts/initialize_camera_impl.sh" "$@" diff --git a/orin/scripts/initialize_camera_impl.sh b/orin/scripts/initialize_camera_impl.sh new file mode 100755 index 0000000..c840953 --- /dev/null +++ b/orin/scripts/initialize_camera_impl.sh @@ -0,0 +1,27 @@ +#!/usr/bin/env bash +# Use the previously deployed 0920 drivers; do not force module compatibility. +set -euo pipefail +root="$(cd "$(dirname "$0")/.." && pwd)" +driver_dir="${DRIVER_DIR:-$root/../vision-anti-pinch-0920/orin/bin/drivers}" +device="${1:-/dev/video5}" +kernel="$(uname -r)" +for driver in max9295 max9296 sgx-yuv-gmsl2; do + module="${driver//-/_}" + [[ -d "/sys/module/$module" ]] && continue + file="$driver_dir/$driver.ko" + [[ -f "$file" ]] || { echo "Missing driver: $file" >&2; exit 1; } + vermagic="$(modinfo -F vermagic "$file")" + if [[ -n "$vermagic" && "${vermagic%% *}" != "$kernel" ]]; then + echo "Kernel mismatch: $file ($vermagic), running $kernel" >&2; exit 1 + fi + if [[ -z "$vermagic" ]]; then + echo "No vermagic metadata: $file; using normal insmod kernel validation." + fi + options=() + [[ "$driver" == sgx-yuv-gmsl2 ]] && options=(gmsl_mode=2,2,2,2) + sudo -n insmod "$file" "${options[@]}" +done +for _ in {1..20}; do [[ -c "$device" ]] && break; sleep .25; done +[[ -c "$device" ]] || { echo "Camera not created: $device" >&2; exit 1; } +v4l2-ctl -d "$device" -c sensor_mode=1,trig_pin=0xffff0007 +echo "Camera ready: $device" diff --git a/pipeline/compare_rtmw_ml.py b/pipeline/compare_rtmw_ml.py new file mode 100644 index 0000000..7d92a4a --- /dev/null +++ b/pipeline/compare_rtmw_ml.py @@ -0,0 +1,384 @@ +"""Compare RTMW-M 256x192 and RTMW-L 384x288 on cached YOLO boxes. + +This is an experiment script. It intentionally keeps the candidate weights +outside the formal models/ directory and does not change deployment code. +""" + +from __future__ import annotations + +import argparse +import json +import math +import time +from collections import defaultdict +from pathlib import Path + +import cv2 +import numpy as np + + +BODY_EDGES = ((5, 7), (7, 9), (6, 8), (8, 10)) +HAND_EDGES_LOCAL = ( + (0, 1), (1, 2), (2, 3), (3, 4), + (0, 5), (5, 6), (6, 7), (7, 8), + (0, 9), (9, 10), (10, 11), (11, 12), + (0, 13), (13, 14), (14, 15), (15, 16), + (0, 17), (17, 18), (18, 19), (19, 20), +) +HAND_SLICES = {"left": (91, 112), "right": (112, 133)} +BODY_WRISTS = {"left": 9, "right": 10} + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser() + parser.add_argument("--video", type=Path, required=True) + parser.add_argument("--boxes", type=Path, required=True) + parser.add_argument("--model-m", type=Path, required=True) + parser.add_argument("--model-l", type=Path, required=True) + parser.add_argument("--output", type=Path, required=True) + parser.add_argument("--max-people", type=int, default=2) + return parser.parse_args() + + +def load_boxes(path: Path, max_people: int) -> dict[int, list[list[float]]]: + frames: dict[int, list[list[float]]] = {} + with path.open("r", encoding="utf-8") as handle: + for line in handle: + item = json.loads(line) + boxes = [entry["xyxy"] for entry in item["boxes"]] + frames[int(item["frame_index"])] = sorted(boxes, key=lambda b: b[0])[:max_people] + return frames + + +def point(pt: np.ndarray) -> tuple[int, int]: + return int(round(float(pt[0]))), int(round(float(pt[1]))) + + +def draw_pose( + image: np.ndarray, + boxes: list[list[float]], + keypoints: np.ndarray, + title: str, +) -> np.ndarray: + canvas = image.copy() + for person_idx, (bbox, kpts) in enumerate(zip(boxes, keypoints), start=1): + x1, y1, x2, y2 = (int(round(value)) for value in bbox) + cv2.rectangle(canvas, (x1, y1), (x2, y2), (0, 220, 255), 2) + cv2.putText( + canvas, + f"P{person_idx}", + (x1 + 4, max(22, y1 + 24)), + cv2.FONT_HERSHEY_SIMPLEX, + 0.75, + (0, 220, 255), + 2, + cv2.LINE_AA, + ) + + for edge_idx, (a, b) in enumerate(BODY_EDGES): + color = (40, 40, 255) if edge_idx < 2 else (30, 230, 30) + cv2.line(canvas, point(kpts[a]), point(kpts[b]), color, 3, cv2.LINE_AA) + for idx in (5, 6, 7, 8, 9, 10): + color = (40, 40, 255) if idx in (5, 7, 9) else (30, 230, 30) + cv2.circle(canvas, point(kpts[idx]), 5, color, -1, cv2.LINE_AA) + + for side, (start, end) in HAND_SLICES.items(): + hand = kpts[start:end] + color = (255, 80, 220) if side == "left" else (255, 190, 40) + for a, b in HAND_EDGES_LOCAL: + cv2.line(canvas, point(hand[a]), point(hand[b]), color, 2, cv2.LINE_AA) + for hand_point in hand: + cv2.circle(canvas, point(hand_point), 2, color, -1, cv2.LINE_AA) + + cv2.rectangle(canvas, (0, 0), (610, 45), (0, 0, 0), -1) + cv2.putText( + canvas, + title, + (10, 31), + cv2.FONT_HERSHEY_SIMPLEX, + 0.8, + (255, 255, 255), + 2, + cv2.LINE_AA, + ) + return canvas + + +def normalized_keypoints(kpts: np.ndarray, bbox: list[float]) -> np.ndarray: + x1, y1, x2, y2 = bbox + width = max(x2 - x1, 1.0) + height = max(y2 - y1, 1.0) + result = kpts.copy().astype(np.float64) + result[:, 0] = (result[:, 0] - x1) / width + result[:, 1] = (result[:, 1] - y1) / height + return result + + +def hand_span(kpts: np.ndarray, bbox: list[float], side: str) -> float: + start, end = HAND_SLICES[side] + hand = kpts[start:end] + span = math.hypot( + float(np.ptp(hand[:, 0])), + float(np.ptp(hand[:, 1])), + ) + x1, y1, x2, y2 = bbox + return span / max(math.hypot(x2 - x1, y2 - y1), 1.0) + + +def summarize(values: list[float]) -> dict[str, float | int | None]: + if not values: + return {"count": 0, "median": None, "p95": None, "max": None} + array = np.asarray(values, dtype=np.float64) + return { + "count": int(array.size), + "median": round(float(np.median(array)), 5), + "p95": round(float(np.percentile(array, 95)), 5), + "max": round(float(np.max(array)), 5), + } + + +def temporal_metrics(records: list[dict]) -> dict: + by_model_person: dict[tuple[str, int], list[dict]] = defaultdict(list) + for record in records: + for model_name in ("M", "L"): + for person_idx in range(len(record[model_name]["keypoints"])): + by_model_person[(model_name, person_idx)].append( + { + "frame": record["frame_index"], + "bbox": record["boxes"][person_idx], + "kpts": np.asarray( + record[model_name]["keypoints"][person_idx], dtype=np.float64 + ), + } + ) + + result: dict[str, dict] = {"M": {}, "L": {}} + for (model_name, person_idx), sequence in by_model_person.items(): + upper_steps: list[float] = [] + wrist_steps = {"left": [], "right": []} + hand_steps = {"left": [], "right": []} + hand_spans = {"left": [], "right": []} + bone_lengths = {"left_upper": [], "left_forearm": [], + "right_upper": [], "right_forearm": []} + + previous: np.ndarray | None = None + for entry in sequence: + normalized = normalized_keypoints(entry["kpts"], entry["bbox"]) + if previous is not None: + steps = np.linalg.norm(normalized - previous, axis=1) + upper_steps.extend(float(steps[idx]) for idx in (5, 6, 7, 8, 9, 10)) + for side, idx in BODY_WRISTS.items(): + wrist_steps[side].append(float(steps[idx])) + for side, (start, end) in HAND_SLICES.items(): + hand_steps[side].extend(float(value) for value in steps[start:end]) + previous = normalized + + hand_spans["left"].append(hand_span(entry["kpts"], entry["bbox"], "left")) + hand_spans["right"].append(hand_span(entry["kpts"], entry["bbox"], "right")) + bone_lengths["left_upper"].append(float(np.linalg.norm(normalized[5] - normalized[7]))) + bone_lengths["left_forearm"].append(float(np.linalg.norm(normalized[7] - normalized[9]))) + bone_lengths["right_upper"].append(float(np.linalg.norm(normalized[6] - normalized[8]))) + bone_lengths["right_forearm"].append(float(np.linalg.norm(normalized[8] - normalized[10]))) + + span_summary = {} + for side, values in hand_spans.items(): + median = float(np.median(values)) + gross_limit = max(0.2, median * 2.5) + span_summary[side] = { + **summarize(values), + "gross_limit": round(gross_limit, 5), + "gross_outlier_frames": int(sum(value > gross_limit for value in values)), + } + + bone_summary = {} + for name, values in bone_lengths.items(): + array = np.asarray(values, dtype=np.float64) + mean = float(np.mean(array)) + bone_summary[name] = { + **summarize(values), + "coefficient_of_variation": round(float(np.std(array) / max(mean, 1e-9)), 5), + } + + result[model_name][f"P{person_idx + 1}"] = { + "upper_body_frame_step": summarize(upper_steps), + "wrist_frame_step": {side: summarize(values) for side, values in wrist_steps.items()}, + "hand_landmark_frame_step": { + side: summarize(values) for side, values in hand_steps.items() + }, + "hand_span_relative_to_person": span_summary, + "bone_length_relative_to_box": bone_summary, + } + return result + + +def model_disagreement(records: list[dict]) -> tuple[dict, list[int]]: + values: dict[str, list[float]] = defaultdict(list) + per_frame: list[tuple[float, int]] = [] + for record in records: + frame_max = 0.0 + for person_idx, bbox in enumerate(record["boxes"]): + m = normalized_keypoints( + np.asarray(record["M"]["keypoints"][person_idx]), bbox + ) + l = normalized_keypoints( + np.asarray(record["L"]["keypoints"][person_idx]), bbox + ) + distance = np.linalg.norm(m - l, axis=1) + values[f"P{person_idx + 1}_upper"].extend( + float(distance[idx]) for idx in (5, 6, 7, 8, 9, 10) + ) + for side, idx in BODY_WRISTS.items(): + values[f"P{person_idx + 1}_{side}_wrist"].append(float(distance[idx])) + frame_max = max(frame_max, float(np.max(distance[[5, 6, 7, 8, 9, 10]]))) + per_frame.append((frame_max, record["frame_index"])) + worst_frames = [frame for _, frame in sorted(per_frame, reverse=True)[:12]] + return {name: summarize(items) for name, items in values.items()}, worst_frames + + +def write_montage( + video_path: Path, + records: list[dict], + selected_frames: list[int], + output_path: Path, + titles: tuple[str, str] = ("RTMW-M 256x192", "RTMW-L 384x288"), +) -> None: + selected = set(selected_frames) + record_map = {record["frame_index"]: record for record in records} + cap = cv2.VideoCapture(str(video_path)) + rows: list[np.ndarray] = [] + frame_idx = 0 + while True: + ok, frame = cap.read() + if not ok: + break + if frame_idx in selected: + record = record_map[frame_idx] + m = draw_pose( + frame, + record["boxes"], + np.asarray(record["M"]["keypoints"]), + f"{titles[0]} | frame {frame_idx}", + ) + l = draw_pose( + frame, + record["boxes"], + np.asarray(record["L"]["keypoints"]), + f"{titles[1]} | frame {frame_idx}", + ) + m = cv2.resize(m, (640, 360), interpolation=cv2.INTER_AREA) + l = cv2.resize(l, (640, 360), interpolation=cv2.INTER_AREA) + rows.append(np.hstack([m, l])) + frame_idx += 1 + cap.release() + if rows: + cv2.imwrite(str(output_path), np.vstack(rows)) + + +def main() -> None: + from rtmlib import RTMPose + args = parse_args() + args.output.mkdir(parents=True, exist_ok=True) + cached_boxes = load_boxes(args.boxes, args.max_people) + + pose_m = RTMPose( + str(args.model_m), model_input_size=(192, 256), backend="onnxruntime", device="cpu" + ) + pose_l = RTMPose( + str(args.model_l), model_input_size=(288, 384), backend="onnxruntime", device="cpu" + ) + + cap = cv2.VideoCapture(str(args.video)) + fps = float(cap.get(cv2.CAP_PROP_FPS)) + width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)) + height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)) + video_output = args.output / "reference_video_rtmw_m_vs_l.mp4" + writer = cv2.VideoWriter( + str(video_output), + cv2.VideoWriter_fourcc(*"mp4v"), + fps, + (width, height // 2), + ) + + records: list[dict] = [] + timings = {"M": [], "L": []} + frame_idx = 0 + while True: + ok, frame = cap.read() + if not ok: + break + boxes = cached_boxes[frame_idx] + + started = time.perf_counter() + kpts_m, scores_m = pose_m(frame, bboxes=boxes) + timings["M"].append((time.perf_counter() - started) * 1000.0) + + started = time.perf_counter() + kpts_l, scores_l = pose_l(frame, bboxes=boxes) + timings["L"].append((time.perf_counter() - started) * 1000.0) + + records.append( + { + "frame_index": frame_idx, + "boxes": boxes, + "M": { + "keypoints": np.round(kpts_m, 3).tolist(), + "raw_scores": np.round(scores_m, 4).tolist(), + }, + "L": { + "keypoints": np.round(kpts_l, 3).tolist(), + "raw_scores": np.round(scores_l, 4).tolist(), + }, + } + ) + + left = draw_pose(frame, boxes, kpts_m, f"RTMW-M 256x192 | frame {frame_idx}") + right = draw_pose(frame, boxes, kpts_l, f"RTMW-L 384x288 | frame {frame_idx}") + left = cv2.resize(left, (width // 2, height // 2), interpolation=cv2.INTER_AREA) + right = cv2.resize(right, (width // 2, height // 2), interpolation=cv2.INTER_AREA) + writer.write(np.hstack([left, right])) + + if frame_idx % 50 == 0: + print(f"processed frame {frame_idx}", flush=True) + frame_idx += 1 + + cap.release() + writer.release() + + jsonl_path = args.output / "reference_video_rtmw_m_vs_l.jsonl" + with jsonl_path.open("w", encoding="utf-8") as handle: + for record in records: + handle.write(json.dumps(record, ensure_ascii=False) + "\n") + + disagreement, worst_frames = model_disagreement(records) + fixed_frames = [0, 60, 120, 180, 240, 252, 300, 360, 411] + montage_frames = sorted(set(fixed_frames + worst_frames)) + write_montage( + args.video, + records, + montage_frames, + args.output / "rtmw_m_vs_l_full_montage.jpg", + ) + + summary = { + "frames": len(records), + "input_sizes_wh": {"M": [192, 256], "L": [288, 384]}, + "raw_score_note": ( + "M and L raw SimCC peak scores use different scales and must not be compared " + "or thresholded with the same numeric value." + ), + "timing_ms_two_people_cpu": { + name: summarize(values[5:] if len(values) > 5 else values) + for name, values in timings.items() + }, + "temporal_metrics_normalized_to_yolo_box": temporal_metrics(records), + "m_l_position_disagreement_normalized_to_yolo_box": disagreement, + "largest_upper_body_disagreement_frames": worst_frames, + "montage_frames": montage_frames, + } + with (args.output / "summary.json").open("w", encoding="utf-8") as handle: + json.dump(summary, handle, ensure_ascii=False, indent=2) + print(json.dumps(summary, ensure_ascii=False, indent=2)) + + +if __name__ == "__main__": + main() diff --git a/pipeline/door_mask_risk.py b/pipeline/door_mask_risk.py new file mode 100644 index 0000000..49eeec5 --- /dev/null +++ b/pipeline/door_mask_risk.py @@ -0,0 +1,488 @@ +"""Per-frame door-mask loading and YOLO/RTMW-only 2-D risk scoring.""" + +from __future__ import annotations + +import json +import math +import re +import time +from collections import Counter +from dataclasses import dataclass +from pathlib import Path + +import cv2 +import numpy as np + + +RASTER_SUFFIXES = {".png", ".bmp", ".tif", ".tiff"} +DEFAULT_DOOR_LABELS = ("door", "车门") + + +@dataclass +class LoadedDoorMask: + status: str + annotation_index: int + path: Path | None + mask: np.ndarray | None + source_format: str | None + + def metadata(self) -> dict: + return { + "status": self.status, + "annotation_index": self.annotation_index, + "path": str(self.path.resolve()) if self.path is not None else None, + "format": self.source_format, + "area_px": int(np.count_nonzero(self.mask)) if self.mask is not None else 0, + } + + +def _frame_index_from_path(path: Path) -> int | None: + matches = re.findall(r"\d+", path.stem) + return int(matches[-1]) if matches else None + + +class DoorMaskSequence: + """Index a directory of X-AnyLabeling JSON polygons or raster masks.""" + + def __init__( + self, + directory: Path | None = None, + mask_format: str = "auto", + labels: tuple[str, ...] = DEFAULT_DOOR_LABELS, + frame_offset: int = 0, + require_every_frame: bool = False, + constant_mask: Path | None = None, + ): + if (directory is None) == (constant_mask is None): + raise ValueError("Provide exactly one of directory or constant_mask") + self.constant_path = ( + Path(constant_mask).expanduser().resolve() + if constant_mask is not None + else None + ) + self.directory = ( + self.constant_path.parent if self.constant_path is not None else Path(directory) + ) + if self.constant_path is not None and not self.constant_path.is_file(): + raise ValueError(f"Constant door mask does not exist: {self.constant_path}") + if self.constant_path is None and not self.directory.is_dir(): + raise ValueError(f"Door mask directory does not exist: {self.directory}") + self.labels = {label.casefold() for label in labels} + self.frame_offset = int(frame_offset) + self.require_every_frame = require_every_frame + + if self.constant_path is not None: + if mask_format == "auto": + selected_format = ( + "xanylabeling" + if self.constant_path.suffix.lower() == ".json" + else "raster" + ) + else: + selected_format = mask_format + self.source_format = selected_format + self.paths = {0: self.constant_path} + return + + json_paths = sorted(self.directory.rglob("*.json")) + raster_paths = sorted( + path + for path in self.directory.rglob("*") + if path.is_file() and path.suffix.lower() in RASTER_SUFFIXES + ) + if mask_format == "auto": + selected_format = "xanylabeling" if json_paths else "raster" + else: + selected_format = mask_format + candidates = json_paths if selected_format == "xanylabeling" else raster_paths + if not candidates: + expected = "JSON" if selected_format == "xanylabeling" else "PNG/BMP/TIFF" + raise ValueError(f"No {expected} door masks found in {self.directory}") + + self.source_format = selected_format + self.paths: dict[int, Path] = {} + for path in candidates: + index = _frame_index_from_path(path) + if index is None and selected_format == "xanylabeling": + data = json.loads(path.read_text(encoding="utf-8-sig")) + explicit = data.get("frame_index", data.get("frameIndex")) + index = int(explicit) if explicit is not None else None + if index is None: + continue + if index in self.paths: + raise ValueError( + f"Multiple door masks map to annotation frame {index}: " + f"{self.paths[index]} and {path}" + ) + self.paths[index] = path + if not self.paths: + raise ValueError( + "Door mask filenames need a frame number, for example 000000.json" + ) + + def load(self, video_frame_index: int, frame_shape: tuple[int, int]) -> LoadedDoorMask: + if self.constant_path is not None: + annotation_index = video_frame_index + path = self.constant_path + else: + annotation_index = video_frame_index + self.frame_offset + path = self.paths.get(annotation_index) + if path is None: + if self.require_every_frame: + raise RuntimeError( + f"No door mask for video frame {video_frame_index} " + f"(expected annotation index {annotation_index})" + ) + return LoadedDoorMask("missing", annotation_index, None, None, self.source_format) + + if self.source_format == "xanylabeling": + mask = self._load_xanylabeling(path, frame_shape) + else: + mask = self._load_raster(path, frame_shape) + status = "loaded" if np.any(mask) else "empty" + return LoadedDoorMask(status, annotation_index, path, mask, self.source_format) + + @staticmethod + def _load_raster(path: Path, frame_shape: tuple[int, int]) -> np.ndarray: + image = cv2.imread(str(path), cv2.IMREAD_UNCHANGED) + if image is None: + raise RuntimeError(f"Cannot read door mask: {path}") + if image.ndim == 3: + image = np.max(image[:, :, :3], axis=2) + height, width = frame_shape + if image.shape != (height, width): + image = cv2.resize(image, (width, height), interpolation=cv2.INTER_NEAREST) + return image > 0 + + def _load_xanylabeling( + self, path: Path, frame_shape: tuple[int, int] + ) -> np.ndarray: + data = json.loads(path.read_text(encoding="utf-8-sig")) + height, width = frame_shape + source_width = float(data.get("imageWidth") or width) + source_height = float(data.get("imageHeight") or height) + scale = np.asarray( + [width / max(source_width, 1.0), height / max(source_height, 1.0)], + dtype=np.float32, + ) + mask = np.zeros((height, width), dtype=np.uint8) + for shape in data.get("shapes", []): + if str(shape.get("label", "")).casefold() not in self.labels: + continue + points = np.asarray(shape.get("points", []), dtype=np.float32).reshape(-1, 2) + if len(points) < 2: + continue + points = np.rint(points * scale).astype(np.int32) + shape_type = str(shape.get("shape_type", "polygon")).lower() + if shape_type == "rectangle" and len(points) >= 2: + x1, y1 = points[0] + x2, y2 = points[1] + cv2.rectangle(mask, (x1, y1), (x2, y2), 1, -1) + elif shape_type == "circle" and len(points) >= 2: + radius = int(round(float(np.linalg.norm(points[1] - points[0])))) + cv2.circle(mask, tuple(points[0]), radius, 1, -1) + elif len(points) >= 3: + cv2.fillPoly(mask, [points], 1) + return mask > 0 + + +def _fit_mask(mask: np.ndarray, frame_shape: tuple[int, int]) -> np.ndarray: + height, width = frame_shape + result = np.asarray(mask) > 0 + if result.shape != (height, width): + result = cv2.resize( + result.astype(np.uint8), (width, height), interpolation=cv2.INTER_NEAREST + ) > 0 + return result + + +def _box_mask(box: list[float] | None, frame_shape: tuple[int, int]) -> np.ndarray | None: + if box is None or len(box) < 4: + return None + height, width = frame_shape + x1, y1, x2, y2 = [int(round(float(value))) for value in box[:4]] + x1, x2 = sorted((max(0, min(width - 1, x1)), max(0, min(width - 1, x2)))) + y1, y2 = sorted((max(0, min(height - 1, y1)), max(0, min(height - 1, y2)))) + if x2 <= x1 or y2 <= y1: + return None + mask = np.zeros((height, width), dtype=bool) + mask[y1 : y2 + 1, x1 : x2 + 1] = True + return mask + + +def _hand_mask(points: list, frame_shape: tuple[int, int]) -> np.ndarray | None: + hand = np.asarray(points, dtype=np.float32).reshape(-1, 2) + if len(hand) < 3 or not np.isfinite(hand).all(): + return None + hull = cv2.convexHull(np.rint(hand).astype(np.int32)) + mask = np.zeros(frame_shape, dtype=np.uint8) + cv2.fillConvexPoly(mask, hull, 1) + return mask > 0 if np.any(mask) else None + + +def _decision(distance_relative: float, overlap_px: int, warning: float, release: float) -> tuple[str, str]: + if overlap_px > 0: + return "danger", "DANGER_INTERSECTION" + if distance_relative <= warning: + return "warning", "WARNING_NEAR" + if distance_relative <= release: + return "watch", "WATCH" + return "clear", "CLEAR" + + +class DoorRiskEvaluator: + """Score RTMW hand regions, with one YOLO body-box fallback per person.""" + + def __init__( + self, + sequence: DoorMaskSequence, + warning_relative: float = 0.30, + release_relative: float = 0.45, + ): + if warning_relative < 0 or release_relative <= warning_relative: + raise ValueError("Risk thresholds require 0 <= warning < release") + self.sequence = sequence + self.warning_relative = float(warning_relative) + self.release_relative = float(release_relative) + self.mask_status_counts: Counter = Counter() + self.level_counts: Counter = Counter() + self.source_counts: Counter = Counter() + self.maximum_score = 0.0 + self.timing_seconds: Counter = Counter() + self.timing_counts: Counter = Counter() + self._constant_cache_shape: tuple[int, int] | None = None + self._constant_loaded: LoadedDoorMask | None = None + self._constant_distance_map: np.ndarray | None = None + self.distance_map_cache_builds = 0 + + def _record_timing(self, name: str, started: float) -> None: + self.timing_seconds[name] += time.perf_counter() - started + self.timing_counts[name] += 1 + + def _measure_region( + self, + region: np.ndarray, + door_mask: np.ndarray, + distance_map: np.ndarray, + source: str, + track_id: int, + side: str, + state: str | None, + ) -> dict: + area_px = int(np.count_nonzero(region)) + overlap_px = int(np.count_nonzero(region & door_mask)) + distances = distance_map[region] + distance_px = float(np.min(distances)) if len(distances) else math.inf + scale_px = max(math.sqrt(area_px), 1.0) + distance_relative = distance_px / scale_px + score = float(np.clip(1.0 - distance_relative / self.release_relative, 0.0, 1.0)) + if overlap_px: + score = 1.0 + level, decision = _decision( + distance_relative, + overlap_px, + self.warning_relative, + self.release_relative, + ) + result = { + "status": "calculated", + "risk_score": round(score, 5), + "risk_level": level, + "decision": decision, + "source": source, + "track_id": int(track_id), + "side": side, + "state": state, + "region_area_px": area_px, + "door_overlap_px": overlap_px, + "region_overlap_ratio": round(overlap_px / max(area_px, 1), 6), + "min_distance_px": round(distance_px, 3), + "region_scale_px": round(scale_px, 3), + "distance_relative": round(distance_relative, 6), + } + self.level_counts[level] += 1 + self.source_counts[source] += 1 + self.maximum_score = max(self.maximum_score, score) + return result + + @staticmethod + def _unavailable(source: str, track_id: int, side: str, status: str) -> dict: + return { + "status": status, + "risk_score": None, + "risk_level": "unavailable", + "decision": "EVIDENCE_UNAVAILABLE", + "source": source, + "track_id": int(track_id), + "side": side, + } + + def evaluate( + self, + frame_index: int, + record: dict, + frame_shape: tuple[int, int], + ) -> tuple[dict, np.ndarray | None]: + total_started = time.perf_counter() + stage_started = time.perf_counter() + use_constant_cache = self.sequence.constant_path is not None + if use_constant_cache and self._constant_cache_shape == frame_shape: + cached = self._constant_loaded + assert cached is not None + loaded = LoadedDoorMask( + cached.status, + frame_index, + cached.path, + cached.mask, + cached.source_format, + ) + else: + loaded = self.sequence.load(frame_index, frame_shape) + if use_constant_cache: + self._constant_cache_shape = frame_shape + self._constant_loaded = loaded + self._constant_distance_map = None + self._record_timing("door_mask_load", stage_started) + self.mask_status_counts[loaded.status] += 1 + mask_ready = loaded.status == "loaded" and loaded.mask is not None + door_mask = loaded.mask if mask_ready else None + distance_map = None + if door_mask is not None and use_constant_cache and self._constant_distance_map is not None: + distance_map = self._constant_distance_map + elif door_mask is not None: + stage_started = time.perf_counter() + distance_map = cv2.distanceTransform( + (~door_mask).astype(np.uint8), cv2.DIST_L2, cv2.DIST_MASK_PRECISE + ) + self._record_timing("distance_transform", stage_started) + if use_constant_cache: + self._constant_distance_map = distance_map + self.distance_map_cache_builds += 1 + + stage_started = time.perf_counter() + regions: list[tuple[dict, np.ndarray, str, int, str, str | None]] = [] + track_ids = [int(value) for value in record.get("track_ids", [])] + person_positions_with_hands: set[int] = set() + for hand in record.get("global_hands", []): + position = int(hand.get("person_id", 0)) - 1 + track_id = track_ids[position] if 0 <= position < len(track_ids) else position + 1 + side = str(hand.get("side", "unknown")) + region = _hand_mask(hand.get("keypoints", []), frame_shape) + if region is None: + continue + person_positions_with_hands.add(position) + regions.append( + ( + hand, + region, + "RTMW_HAND_KEYPOINTS", + track_id, + side, + str(hand.get("state", "unknown")), + ) + ) + + boxes = record.get("boxes", []) + body_fallbacks = [] + for position, box in enumerate(boxes): + if position in person_positions_with_hands: + continue + region = _box_mask(box, frame_shape) + if region is None: + continue + track_id = track_ids[position] if position < len(track_ids) else position + 1 + fallback = { + "track_id": track_id, + "side": "body", + "state": "rtmw_hand_unavailable", + "body_zone_box": box, + } + body_fallbacks.append(fallback) + regions.append( + ( + fallback, + region, + "YOLO_BODY_BOX_FALLBACK", + track_id, + "body", + "rtmw_hand_unavailable", + ) + ) + record["body_fallbacks"] = body_fallbacks + self._record_timing("region_build", stage_started) + + stage_started = time.perf_counter() + measurements = [] + unavailable_status = f"door_mask_{loaded.status}" + for owner, region, source, track_id, side, state in regions: + if door_mask is None or distance_map is None: + measurement = self._unavailable( + source, track_id, side, unavailable_status + ) + else: + measurement = self._measure_region( + region, door_mask, distance_map, source, track_id, side, state + ) + owner["risk"] = measurement + measurements.append(measurement) + self._record_timing("region_measure", stage_started) + + calculated = [item for item in measurements if item["risk_score"] is not None] + nearest = max(calculated, key=lambda item: item["risk_score"]) if calculated else None + if loaded.status != "loaded": + status = unavailable_status + elif not calculated: + status = "person_region_missing" + else: + status = "calculated" + frame_risk = { + "status": status, + "risk_score": nearest["risk_score"] if nearest else None, + "risk_level": nearest["risk_level"] if nearest else "unavailable", + "decision": nearest["decision"] if nearest else "EVIDENCE_UNAVAILABLE", + "source": nearest["source"] if nearest else None, + "track_id": nearest["track_id"] if nearest else None, + "side": nearest["side"] if nearest else None, + "door_mask": loaded.metadata(), + "warning_distance_relative": self.warning_relative, + "release_distance_relative": self.release_relative, + "regions": measurements, + } + self._record_timing("risk_total", total_started) + return frame_risk, loaded.mask + + def summary(self) -> dict: + return { + "mask_directory": str(self.sequence.directory.resolve()), + "constant_mask": ( + str(self.sequence.constant_path) + if self.sequence.constant_path is not None + else None + ), + "constant_for_all_frames": self.sequence.constant_path is not None, + "mask_format": self.sequence.source_format, + "door_labels": sorted(self.sequence.labels), + "frame_offset": self.sequence.frame_offset, + "warning_distance_relative": self.warning_relative, + "release_distance_relative": self.release_relative, + "mask_status_counts": dict(self.mask_status_counts), + "risk_level_counts": dict(self.level_counts), + "source_counts": dict(self.source_counts), + "maximum_risk_score": round(self.maximum_score, 5), + "constant_distance_map_cached": self.sequence.constant_path is not None, + "distance_map_cache_builds": self.distance_map_cache_builds, + "timing_ms": { + name: { + "count": int(self.timing_counts[name]), + "total": round(float(seconds) * 1000.0, 3), + "mean": round( + float(seconds) * 1000.0 + / max(int(self.timing_counts[name]), 1), + 3, + ), + } + for name, seconds in self.timing_seconds.items() + }, + "score_semantics": "2-D image overlap/proximity research score; not millimeters or safety certification", + "region_policy": "RTMW alarm-eligible hand hulls; YOLO body box only when a person has no usable RTMW hand region", + } diff --git a/pipeline/door_zone_risk.py b/pipeline/door_zone_risk.py new file mode 100644 index 0000000..972cfb7 --- /dev/null +++ b/pipeline/door_zone_risk.py @@ -0,0 +1,358 @@ +"""Multi-feature person-to-door assignment and active-CAN risk gating.""" + +from __future__ import annotations + +import json +import math +from collections import Counter, defaultdict, deque +from dataclasses import dataclass +from pathlib import Path +from statistics import median + + +DOORS = ("front", "rear") +POSE_FEATURES = { + "shoulder_center_x", + "shoulder_center_y", + "shoulder_width", + "torso_height", + "hip_center_y", +} + + +@dataclass(frozen=True) +class FeatureCalibration: + rear: float + front: float + weight: float + + def frontness(self, value: float) -> float: + span = self.front - self.rear + if span == 0: + return 0.5 + return min(1.0, max(0.0, (value - self.rear) / span)) + + +@dataclass(frozen=True) +class DoorZoneConfig: + features: dict[str, FeatureCalibration] + front_probability_min: float + rear_probability_max: float + minimum_person_height: float + minimum_pose_features: int + history_frames: int + stale_track_frames: int + source_path: Path | None = None + + @classmethod + def load(cls, path: Path) -> "DoorZoneConfig": + source = Path(path).expanduser().resolve() + data = json.loads(source.read_text(encoding="utf-8-sig")) + features = { + name: FeatureCalibration( + rear=float(values["rear"]), + front=float(values["front"]), + weight=float(values["weight"]), + ) + for name, values in data["features"].items() + } + return cls( + features=features, + front_probability_min=float(data.get("front_probability_min", 0.72)), + rear_probability_max=float(data.get("rear_probability_max", 0.28)), + minimum_person_height=float(data.get("minimum_person_height", 0.40)), + minimum_pose_features=int(data.get("minimum_pose_features", 4)), + history_frames=int(data.get("history_frames", 5)), + stale_track_frames=int(data.get("stale_track_frames", 120)), + source_path=source, + ) + + def classify(self, probability: float) -> str: + if probability >= self.front_probability_min: + return "front" + if probability <= self.rear_probability_max: + return "rear" + return "transition" + + +def _point(keypoints: list, index: int) -> tuple[float, float] | None: + if index >= len(keypoints) or len(keypoints[index]) < 2: + return None + x, y = float(keypoints[index][0]), float(keypoints[index][1]) + return (x, y) if math.isfinite(x) and math.isfinite(y) else None + + +def extract_person_features( + record: dict, + position: int, + frame_shape: tuple[int, int], +) -> tuple[dict[str, float], int]: + """Extract fixed-camera pose features, with box geometry as low-weight fallback.""" + + height, width = frame_shape + box = record.get("boxes", [])[position] + x1, y1, x2, y2 = [float(value) for value in box[:4]] + features = { + "box_center_x": ((x1 + x2) * 0.5) / max(width, 1), + "box_height": (y2 - y1) / max(height, 1), + } + + all_keypoints = record.get("temporal_selected_keypoints", []) + if position >= len(all_keypoints): + return features, 0 + keypoints = all_keypoints[position] + left_shoulder = _point(keypoints, 5) + right_shoulder = _point(keypoints, 6) + left_hip = _point(keypoints, 11) + right_hip = _point(keypoints, 12) + if None in (left_shoulder, right_shoulder, left_hip, right_hip): + return features, 0 + + shoulder_x = (left_shoulder[0] + right_shoulder[0]) * 0.5 + shoulder_y = (left_shoulder[1] + right_shoulder[1]) * 0.5 + hip_y = (left_hip[1] + right_hip[1]) * 0.5 + features.update( + { + "shoulder_center_x": shoulder_x / max(width, 1), + "shoulder_center_y": shoulder_y / max(height, 1), + "shoulder_width": math.hypot( + left_shoulder[0] - right_shoulder[0], + left_shoulder[1] - right_shoulder[1], + ) + / max(width, 1), + "torso_height": abs(hip_y - shoulder_y) / max(height, 1), + "hip_center_y": hip_y / max(height, 1), + } + ) + return features, len(POSE_FEATURES) + + +class DoorZoneAssigner: + """Fuse pose-dominant evidence and smooth front probability per track.""" + + def __init__(self, config: DoorZoneConfig): + self.config = config + self.history: dict[int, deque[float]] = defaultdict( + lambda: deque(maxlen=self.config.history_frames) + ) + self.last_seen: dict[int, int] = {} + self.zone_counts: Counter = Counter() + + def _score(self, features: dict[str, float]) -> tuple[float, dict[str, float]]: + weighted_sum = 0.0 + total_weight = 0.0 + contributions = {} + for name, calibration in self.config.features.items(): + if name not in features: + continue + frontness = calibration.frontness(float(features[name])) + contributions[name] = round(frontness, 6) + weighted_sum += frontness * calibration.weight + total_weight += calibration.weight + probability = weighted_sum / total_weight if total_weight else 0.5 + return probability, contributions + + def update( + self, + frame_index: int, + record: dict, + frame_shape: tuple[int, int], + ) -> list[dict]: + track_ids = [int(value) for value in record.get("track_ids", [])] + assignments = [] + for position, _box in enumerate(record.get("boxes", [])): + track_id = track_ids[position] if position < len(track_ids) else position + 1 + features, pose_feature_count = extract_person_features( + record, position, frame_shape + ) + probability, contributions = self._score(features) + enough_pose = pose_feature_count >= self.config.minimum_pose_features + if features["box_height"] < self.config.minimum_person_height: + raw_zone = "unassigned" + door_zone = "unassigned" + elif not enough_pose: + raw_zone = "transition" + door_zone = "transition" + else: + raw_zone = self.config.classify(probability) + values = self.history[track_id] + values.append(probability) + door_zone = self.config.classify(float(median(values))) + + values = self.history.get(track_id, ()) + smoothed_probability = ( + float(median(values)) if values else float(probability) + ) + self.last_seen[track_id] = int(frame_index) + self.zone_counts[door_zone] += 1 + assignments.append( + { + "track_id": track_id, + "person_index": position, + "door_zone": door_zone, + "raw_door_zone": raw_zone, + "front_probability": round(float(probability), 6), + "smoothed_front_probability": round(smoothed_probability, 6), + "confidence": round(abs(smoothed_probability - 0.5) * 2.0, 6), + "evidence_source": "pose_primary_bbox_secondary", + "pose_feature_count": pose_feature_count, + "features": {name: round(value, 6) for name, value in features.items()}, + "feature_frontness": contributions, + "history_frames": len(values), + } + ) + + expired = [ + track_id + for track_id, seen_at in self.last_seen.items() + if frame_index - seen_at > self.config.stale_track_frames + ] + for track_id in expired: + self.history.pop(track_id, None) + self.last_seen.pop(track_id, None) + return assignments + + def summary(self) -> dict: + return { + "policy": "pose-primary weighted front probability with causal median", + "config": str(self.config.source_path) if self.config.source_path else None, + "history_frames": self.config.history_frames, + "front_probability_min": self.config.front_probability_min, + "rear_probability_max": self.config.rear_probability_max, + "minimum_person_height": self.config.minimum_person_height, + "minimum_pose_features": self.config.minimum_pose_features, + "features": { + name: { + "rear": calibration.rear, + "front": calibration.front, + "weight": calibration.weight, + } + for name, calibration in self.config.features.items() + }, + "zone_counts": dict(self.zone_counts), + } + + +def _inactive_risk(door: str) -> dict: + return { + "door_id": door, + "status": "inactive", + "risk_score": None, + "risk_level": "inactive", + "decision": "DOOR_CAN_NOT_ACTIVE", + "risk_type": f"{door.upper()}_INACTIVE", + "source": None, + "track_id": None, + "side": None, + "regions": [], + } + + +def _clear_risk(door: str) -> dict: + return { + "door_id": door, + "status": "calculated", + "risk_score": 0.0, + "risk_level": "clear", + "decision": "CLEAR_NO_PERSON_ASSIGNED_TO_ACTIVE_DOOR", + "risk_type": f"{door.upper()}_CLEAR", + "source": None, + "track_id": None, + "side": None, + "regions": [], + } + + +def _unknown_risk(door: str, transition_regions: list[dict]) -> dict: + return { + "door_id": door, + "status": "door_zone_ambiguous", + "risk_score": None, + "risk_level": "unknown", + "decision": "EVIDENCE_UNAVAILABLE_TRANSITION_ZONE", + "risk_type": f"{door.upper()}_UNKNOWN", + "source": None, + "track_id": None, + "side": None, + "regions": transition_regions, + } + + +def build_door_risks( + raw_risk: dict, + assignments: list[dict], + active_door: str, +) -> tuple[dict[str, dict], dict]: + """Build per-door risk and select the CAN-active door result.""" + + zone_by_track = { + int(assignment["track_id"]): str(assignment["door_zone"]) + for assignment in assignments + } + annotated_regions = [] + for original in raw_risk.get("regions", []): + region = dict(original) + region["door_zone"] = zone_by_track.get(int(region.get("track_id", -1)), "unknown") + annotated_regions.append(region) + + active_doors = set(DOORS) if active_door == "both" else {active_door} + if active_door == "none": + active_doors.clear() + + door_risks: dict[str, dict] = {} + for door in DOORS: + if door not in active_doors: + door_risks[door] = _inactive_risk(door) + continue + exact = [region for region in annotated_regions if region["door_zone"] == door] + transition = [ + region for region in annotated_regions if region["door_zone"] == "transition" + ] + calculated = [region for region in exact if region.get("risk_score") is not None] + if calculated: + nearest = max(calculated, key=lambda item: float(item["risk_score"])) + level = str(nearest["risk_level"]) + if transition and level in ("clear", "watch"): + door_risks[door] = _unknown_risk(door, transition) + else: + door_risks[door] = { + "door_id": door, + "status": "calculated", + "risk_score": nearest["risk_score"], + "risk_level": level, + "decision": nearest["decision"], + "risk_type": f"{door.upper()}_{level.upper()}", + "source": nearest["source"], + "track_id": nearest["track_id"], + "side": nearest["side"], + "regions": exact, + } + elif transition: + door_risks[door] = _unknown_risk(door, transition) + else: + door_risks[door] = _clear_risk(door) + + if len(active_doors) == 1: + selected = dict(door_risks[next(iter(active_doors))]) + elif not active_doors: + selected = { + "door_id": None, + "status": "inactive", + "risk_score": None, + "risk_level": "inactive", + "decision": "NO_DOOR_CAN_ACTIVE", + "risk_type": "NO_ACTIVE_DOOR", + "source": None, + "track_id": None, + "side": None, + "regions": [], + } + else: + severity = {"danger": 4, "warning": 3, "watch": 2, "unknown": 1, "clear": 0} + selected = dict( + max(door_risks.values(), key=lambda item: severity.get(item["risk_level"], -1)) + ) + selected["door_mask"] = raw_risk.get("door_mask") + selected["warning_distance_relative"] = raw_risk.get("warning_distance_relative") + selected["release_distance_relative"] = raw_risk.get("release_distance_relative") + return door_risks, selected diff --git a/pipeline/fuse_rtmw_l_temporal.py b/pipeline/fuse_rtmw_l_temporal.py new file mode 100644 index 0000000..f212589 --- /dev/null +++ b/pipeline/fuse_rtmw_l_temporal.py @@ -0,0 +1,855 @@ +"""Fuse raw and person-isolated RTMW-L poses with temporal consistency. + +The fusion keeps the two inference views as hypotheses. It selects body limbs +and hands separately, rejects implausible hand geometry, adds source hysteresis +and constant-velocity consistency, then applies a motion-adaptive EMA in +person-box coordinates. Occluded/conflicting limbs remain labelled as such so +they are not accidentally consumed as precise anti-pinch distance evidence. +""" + +from __future__ import annotations + +import argparse +import json +import math +from collections import Counter, defaultdict +from dataclasses import dataclass, field +from pathlib import Path + +import cv2 +import numpy as np + +from compare_rtmw_ml import ( + BODY_EDGES, + BODY_WRISTS, + HAND_EDGES_LOCAL, + HAND_SLICES, + draw_pose, + normalized_keypoints, + summarize, + temporal_metrics, +) + + +BODY_SIDE_INDICES = {"left": (5, 7, 9), "right": (6, 8, 10)} +SIDE_COLORS = {"left": (40, 40, 255), "right": (30, 230, 30)} +HAND_COLORS = {"left": (255, 80, 220), "right": (255, 190, 40)} + + +@dataclass +class Candidate: + source: str + keypoints: np.ndarray + scores: np.ndarray + normalized: np.ndarray + limb_score: float + hand_score: float + upper_length: float + forearm_length: float + hand_span: float + wrist_link: float + limb_valid: bool + hand_valid: bool + + +@dataclass +class SideTrack: + limb_previous: np.ndarray | None = None + limb_velocity: np.ndarray | None = None + hand_previous: np.ndarray | None = None + hand_velocity: np.ndarray | None = None + limb_source: str | None = None + hand_source: str | None = None + previous_hand_eligible: bool = False + + +@dataclass +class PersonTrack: + sides: dict[str, SideTrack] = field( + default_factory=lambda: {"left": SideTrack(), "right": SideTrack()} + ) + smoothed: np.ndarray | None = None + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser() + parser.add_argument("--video", type=Path, required=True) + parser.add_argument("--raw-jsonl", type=Path, required=True) + parser.add_argument("--isolated-jsonl", type=Path, required=True) + parser.add_argument("--output", type=Path, required=True) + return parser.parse_args() + + +def bbox_diagonal(bbox: list[float]) -> float: + return max(math.hypot(bbox[2] - bbox[0], bbox[3] - bbox[1]), 1.0) + + +def build_candidate( + source: str, + keypoints: np.ndarray, + scores: np.ndarray, + bbox: list[float], + side: str, +) -> Candidate: + normalized = normalized_keypoints(keypoints, bbox) + shoulder, elbow, wrist = BODY_SIDE_INDICES[side] + hand_start, hand_end = HAND_SLICES[side] + hand = keypoints[hand_start:hand_end] + diagonal = bbox_diagonal(bbox) + upper_length = float(np.linalg.norm(normalized[shoulder] - normalized[elbow])) + forearm_length = float(np.linalg.norm(normalized[elbow] - normalized[wrist])) + hand_span = math.hypot(float(np.ptp(hand[:, 0])), float(np.ptp(hand[:, 1]))) / diagonal + wrist_link = float(np.linalg.norm(keypoints[wrist] - hand[0])) / diagonal + limb_score = float(np.mean(scores[[shoulder, elbow, wrist]])) + hand_score = float(np.mean(scores[hand_start:hand_end])) + limb_valid = 0.025 <= upper_length <= 0.8 and 0.025 <= forearm_length <= 0.8 + hand_valid = 0.008 <= hand_span <= 0.35 and wrist_link <= 0.12 + return Candidate( + source=source, + keypoints=keypoints, + scores=scores, + normalized=normalized, + limb_score=limb_score, + hand_score=hand_score, + upper_length=upper_length, + forearm_length=forearm_length, + hand_span=hand_span, + wrist_link=wrist_link, + limb_valid=limb_valid, + hand_valid=hand_valid, + ) + + +def predicted(previous: np.ndarray | None, velocity: np.ndarray | None) -> np.ndarray | None: + if previous is None: + return None + if velocity is None: + return previous + return previous + 0.65 * velocity + + +def robust_distance(points: np.ndarray, target: np.ndarray | None) -> float: + if target is None: + return 0.0 + distances = np.linalg.norm(points - target, axis=1) + return float(np.median(distances)) + + +def select_limb( + candidates: list[Candidate], + side: str, + track: SideTrack | None, +) -> tuple[Candidate, dict[str, float]]: + indices = np.asarray(BODY_SIDE_INDICES[side]) + target = predicted(track.limb_previous, track.limb_velocity) if track else None + scored: list[tuple[float, Candidate, float]] = [] + for candidate in candidates: + temporal_cost = robust_distance(candidate.normalized[indices], target) + geometry_bonus = 1.2 if candidate.limb_valid else -5.0 + switch_bonus = 0.35 if track and candidate.source == track.limb_source else 0.0 + score = candidate.limb_score + geometry_bonus + switch_bonus - 7.0 * temporal_cost + scored.append((score, candidate, temporal_cost)) + score, selected, temporal_cost = max(scored, key=lambda item: item[0]) + return selected, {"score": score, "temporal_cost": temporal_cost} + + +def hand_link_to_limb( + candidate: Candidate, + selected_limb: Candidate, + bbox: list[float], + side: str, +) -> float: + hand_start, _ = HAND_SLICES[side] + wrist_idx = BODY_WRISTS[side] + return float( + np.linalg.norm(candidate.keypoints[hand_start] - selected_limb.keypoints[wrist_idx]) + ) / bbox_diagonal(bbox) + + +def select_hand( + candidates: list[Candidate], + selected_limb: Candidate, + bbox: list[float], + side: str, + track: SideTrack | None, +) -> tuple[Candidate, bool, dict[str, float]]: + start, end = HAND_SLICES[side] + target = predicted(track.hand_previous, track.hand_velocity) if track else None + scored: list[tuple[float, Candidate, bool, float, float]] = [] + for candidate in candidates: + temporal_cost = robust_distance(candidate.normalized[start:end], target) + link = hand_link_to_limb(candidate, selected_limb, bbox, side) + valid = candidate.hand_valid + geometry_bonus = 3.0 if valid else -7.0 + switch_bonus = 0.45 if track and candidate.source == track.hand_source else 0.0 + limb_source_bonus = 0.15 if candidate.source == selected_limb.source else 0.0 + # A newly visible endpoint may disagree with the currently selected + # (occlusion-driven) limb. Keep this as a soft cost so strong image + # evidence can switch both the hand and its limb together. + link_penalty = 1.5 * min(max(link - 0.12, 0.0), 0.2) + score = ( + candidate.hand_score + + geometry_bonus + + switch_bonus + + limb_source_bonus + - link_penalty + - 8.0 * min(temporal_cost, 0.15) + ) + scored.append((score, candidate, valid, temporal_cost, link)) + score, selected, valid, temporal_cost, link = max(scored, key=lambda item: item[0]) + return selected, valid, { + "score": score, + "temporal_cost": temporal_cost, + "link_to_selected_limb": link, + } + + +def candidate_disagreement( + raw: Candidate, + isolated: Candidate, + side: str, +) -> tuple[float, float]: + shoulder, elbow, wrist = BODY_SIDE_INDICES[side] + start, end = HAND_SLICES[side] + limb = float( + np.median( + np.linalg.norm( + raw.normalized[[shoulder, elbow, wrist]] + - isolated.normalized[[shoulder, elbow, wrist]], + axis=1, + ) + ) + ) + hand_distances = np.linalg.norm( + raw.normalized[start:end] - isolated.normalized[start:end], axis=1 + ) + hand = float(np.percentile(hand_distances, 90)) + return limb, hand + + +def suppress_duplicate_hands( + fused: np.ndarray, + metadata: dict, + bbox: list[float], +) -> None: + left_start, left_end = HAND_SLICES["left"] + right_start, right_end = HAND_SLICES["right"] + left_center = np.mean(fused[left_start:left_end], axis=0) + right_center = np.mean(fused[right_start:right_end], axis=0) + separation = float(np.linalg.norm(left_center - right_center)) / bbox_diagonal(bbox) + if separation >= 0.08: + return + weaker = min( + ("left", "right"), + key=lambda side: metadata[side]["selected_hand_score"], + ) + metadata[weaker]["hand_state"] = "duplicate" + metadata[weaker]["distance_eligible"] = False + metadata[weaker]["alarm_eligible"] = False + metadata[weaker]["duplicate_hand_separation"] = round(separation, 5) + + +def fuse_person( + raw_keypoints: np.ndarray, + raw_scores: np.ndarray, + isolated_keypoints: np.ndarray, + isolated_scores: np.ndarray, + bbox: list[float], + track: PersonTrack | None, +) -> tuple[np.ndarray, dict]: + fused = raw_keypoints.copy() + metadata: dict[str, dict] = {} + + for side in ("left", "right"): + side_track = track.sides[side] if track else None + raw = build_candidate("raw", raw_keypoints, raw_scores, bbox, side) + isolated = build_candidate( + "isolated", isolated_keypoints, isolated_scores, bbox, side + ) + candidates = [raw, isolated] + limb, limb_detail = select_limb(candidates, side, side_track) + hand, hand_valid, hand_detail = select_hand( + candidates, limb, bbox, side, side_track + ) + + other_hand = isolated if hand.source == "raw" else raw + endpoint_advantage = hand.hand_score - other_hand.hand_score + if hand.source != limb.source and hand_valid and endpoint_advantage >= 0.6: + limb = hand + limb_detail["visible_endpoint_override"] = 1.0 + hand_detail["link_to_selected_limb"] = hand.wrist_link + + body_indices = np.asarray(BODY_SIDE_INDICES[side]) + hand_start, hand_end = HAND_SLICES[side] + fused[body_indices] = limb.keypoints[body_indices] + fused[hand_start:hand_end] = hand.keypoints[hand_start:hand_end] + + limb_disagreement, hand_disagreement = candidate_disagreement(raw, isolated, side) + if limb_disagreement <= 0.05: + limb_state = "consensus" + else: + limb_state = "inferred" + + score_advantage = raw.hand_score - isolated.hand_score + raw_temporally_stable = hand_detail["temporal_cost"] <= 0.08 + if not hand_valid: + hand_state = "rejected" + distance_eligible = False + elif hand_disagreement <= 0.06: + hand_state = "consensus" + distance_eligible = True + elif ( + hand.source == "raw" + and score_advantage >= 0.6 + and raw_temporally_stable + ): + hand_state = "visible_endpoint" + distance_eligible = True + elif hand.source == "isolated": + hand_state = "inferred" + distance_eligible = False + else: + hand_state = "conflict" + distance_eligible = False + fast_alarm_endpoint = ( + hand_valid and hand.source == "raw" and score_advantage >= 0.6 + ) + alarm_eligible = distance_eligible or fast_alarm_endpoint + + metadata[side] = { + "limb_source": limb.source, + "hand_source": hand.source, + "limb_state": limb_state, + "hand_state": hand_state, + "distance_eligible": distance_eligible, + "alarm_eligible": alarm_eligible, + "limb_disagreement": round(limb_disagreement, 5), + "hand_disagreement": round(hand_disagreement, 5), + "raw_hand_score": round(raw.hand_score, 4), + "isolated_hand_score": round(isolated.hand_score, 4), + "selected_hand_span": round(hand.hand_span, 5), + "selected_hand_score": round(hand.hand_score, 4), + "selected_hand_valid": hand_valid, + "limb_selection": {key: round(value, 5) for key, value in limb_detail.items()}, + "hand_selection": {key: round(value, 5) for key, value in hand_detail.items()}, + } + + if side_track is not None: + limb_now = limb.normalized[body_indices] + hand_now = hand.normalized[hand_start:hand_end] + if side_track.limb_previous is not None: + delta = limb_now - side_track.limb_previous + side_track.limb_velocity = ( + delta + if side_track.limb_velocity is None + else 0.65 * side_track.limb_velocity + 0.35 * delta + ) + if hand_valid: + if side_track.hand_previous is not None: + delta = hand_now - side_track.hand_previous + side_track.hand_velocity = ( + delta + if side_track.hand_velocity is None + else 0.65 * side_track.hand_velocity + 0.35 * delta + ) + side_track.hand_previous = hand_now + side_track.hand_source = hand.source + side_track.limb_previous = limb_now + side_track.limb_source = limb.source + side_track.previous_hand_eligible = distance_eligible + + suppress_duplicate_hands(fused, metadata, bbox) + return fused, metadata + + +def adaptive_smooth( + keypoints: np.ndarray, + bbox: list[float], + track: PersonTrack, +) -> np.ndarray: + normalized = normalized_keypoints(keypoints, bbox) + if track.smoothed is None: + track.smoothed = normalized + else: + displacement = np.linalg.norm(normalized - track.smoothed, axis=1) + alpha = 0.32 + 0.58 * np.clip(displacement / 0.08, 0.0, 1.0) + track.smoothed = alpha[:, None] * normalized + (1.0 - alpha[:, None]) * track.smoothed + + x1, y1, x2, y2 = bbox + result = track.smoothed.copy() + result[:, 0] = result[:, 0] * max(x2 - x1, 1.0) + x1 + result[:, 1] = result[:, 1] * max(y2 - y1, 1.0) + y1 + return result + + +def build_global_hands( + keypoints: list[np.ndarray], + metadata: list[dict], +) -> list[dict]: + candidates = [] + for person_idx, (person_keypoints, person_meta) in enumerate( + zip(keypoints, metadata), start=1 + ): + for side, (start, end) in HAND_SLICES.items(): + meta = person_meta[side] + if not meta["alarm_eligible"]: + continue + hand = person_keypoints[start:end] + span_px = math.hypot(float(np.ptp(hand[:, 0])), float(np.ptp(hand[:, 1]))) + candidates.append( + { + "person_id": person_idx, + "side": side, + "source": meta["hand_source"], + "state": meta["hand_state"], + "distance_eligible": meta["distance_eligible"], + "alarm_eligible": meta["alarm_eligible"], + "score": meta["selected_hand_score"], + "center": np.mean(hand, axis=0), + "span_px": span_px, + "keypoints": hand, + "aliases": [], + } + ) + + kept: list[dict] = [] + for candidate in sorted(candidates, key=lambda item: item["score"], reverse=True): + duplicate = None + for existing in kept: + threshold = float(np.clip(0.45 * max(candidate["span_px"], existing["span_px"]), 18, 70)) + if float(np.linalg.norm(candidate["center"] - existing["center"])) <= threshold: + duplicate = existing + break + if duplicate is None: + kept.append(candidate) + else: + duplicate["aliases"].append( + {"person_id": candidate["person_id"], "side": candidate["side"]} + ) + + return [ + { + "person_id": item["person_id"], + "side": item["side"], + "source": item["source"], + "state": item["state"], + "distance_eligible": item["distance_eligible"], + "alarm_eligible": item["alarm_eligible"], + "score": item["score"], + "center": np.round(item["center"], 3).tolist(), + "span_px": round(item["span_px"], 3), + "keypoints": np.round(item["keypoints"], 3).tolist(), + "aliases": item["aliases"], + } + for item in kept + ] + + +def draw_dashed_line( + image: np.ndarray, + start: np.ndarray, + end: np.ndarray, + color: tuple[int, int, int], + thickness: int = 2, +) -> None: + a = np.asarray(start, dtype=np.float64) + b = np.asarray(end, dtype=np.float64) + length = float(np.linalg.norm(b - a)) + if length < 1.0: + return + direction = (b - a) / length + for offset in np.arange(0.0, length, 14.0): + segment_start = a + direction * offset + segment_end = a + direction * min(offset + 8.0, length) + cv2.line( + image, + tuple(np.round(segment_start).astype(int)), + tuple(np.round(segment_end).astype(int)), + color, + thickness, + cv2.LINE_AA, + ) + + +def draw_fused( + image: np.ndarray, + boxes: list[list[float]], + keypoints: list[np.ndarray], + metadata: list[dict], + title: str, +) -> np.ndarray: + canvas = image.copy() + for person_idx, (bbox, kpts, person_meta) in enumerate( + zip(boxes, keypoints, metadata), start=1 + ): + x1, y1, x2, y2 = (int(round(value)) for value in bbox) + cv2.rectangle(canvas, (x1, y1), (x2, y2), (0, 220, 255), 2) + cv2.putText( + canvas, + f"P{person_idx}", + (x1 + 4, max(22, y1 + 24)), + cv2.FONT_HERSHEY_SIMPLEX, + 0.72, + (0, 220, 255), + 2, + cv2.LINE_AA, + ) + for side in ("left", "right"): + shoulder, elbow, wrist = BODY_SIDE_INDICES[side] + side_meta = person_meta[side] + color = SIDE_COLORS[side] + for a, b in ((shoulder, elbow), (elbow, wrist)): + if side_meta["limb_state"] == "consensus": + cv2.line( + canvas, + tuple(np.round(kpts[a]).astype(int)), + tuple(np.round(kpts[b]).astype(int)), + color, + 3, + cv2.LINE_AA, + ) + else: + draw_dashed_line(canvas, kpts[a], kpts[b], (0, 140, 255), 3) + + start, end = HAND_SLICES[side] + hand = kpts[start:end] + if side_meta["hand_state"] != "rejected": + hand_color = HAND_COLORS[side] + if not side_meta["distance_eligible"]: + hand_color = (0, 140, 255) + for a, b in HAND_EDGES_LOCAL: + if side_meta["distance_eligible"]: + cv2.line( + canvas, + tuple(np.round(hand[a]).astype(int)), + tuple(np.round(hand[b]).astype(int)), + hand_color, + 2, + cv2.LINE_AA, + ) + else: + draw_dashed_line(canvas, hand[a], hand[b], hand_color, 2) + label = ( + f"{side[0].upper()}:{side_meta['hand_state'][:4]}" + f"/{side_meta['hand_source'][0]}" + ) + label_y = max(45, y1 + (48 if side == "left" else 70)) + cv2.putText( + canvas, + label, + (x1 + 4, label_y), + cv2.FONT_HERSHEY_SIMPLEX, + 0.5, + HAND_COLORS[side] if side_meta["distance_eligible"] else (0, 140, 255), + 1, + cv2.LINE_AA, + ) + + cv2.rectangle(canvas, (0, 0), (650, 45), (0, 0, 0), -1) + cv2.putText( + canvas, + title, + (10, 31), + cv2.FONT_HERSHEY_SIMPLEX, + 0.8, + (255, 255, 255), + 2, + cv2.LINE_AA, + ) + return canvas + + +def make_montage( + video_path: Path, + records: list[dict], + selected_frames: list[int], + output_path: Path, +) -> None: + selected = set(selected_frames) + by_frame = {record["frame_index"]: record for record in records} + cap = cv2.VideoCapture(str(video_path)) + rows = [] + frame_idx = 0 + while True: + ok, frame = cap.read() + if not ok: + break + if frame_idx in selected: + record = by_frame[frame_idx] + boxes = record["boxes"] + raw = draw_pose( + frame, + boxes, + np.asarray(record["raw_keypoints"]), + f"RTMW-L raw | frame {frame_idx}", + ) + spatial = draw_fused( + frame, + boxes, + [np.asarray(points) for points in record["spatial_keypoints"]], + record["spatial_metadata"], + f"spatial fusion | frame {frame_idx}", + ) + temporal = draw_fused( + frame, + boxes, + [np.asarray(points) for points in record["temporal_keypoints"]], + record["temporal_metadata"], + f"temporal fusion | frame {frame_idx}", + ) + panels = [ + cv2.resize(panel, (480, 270), interpolation=cv2.INTER_AREA) + for panel in (raw, spatial, temporal) + ] + rows.append(np.hstack(panels)) + frame_idx += 1 + cap.release() + if rows: + cv2.imwrite(str(output_path), np.vstack(rows)) + + +def active_hand_temporal_metrics(records: list[dict], variant: str) -> dict: + values: dict[str, list[float]] = defaultdict(list) + previous: dict[tuple[int, str], tuple[int, np.ndarray]] = {} + for record in records: + for person_idx, bbox in enumerate(record["boxes"]): + kpts = normalized_keypoints( + np.asarray(record[f"{variant}_keypoints"][person_idx]), bbox + ) + meta = record[f"{variant}_metadata"][person_idx] + for side, (start, end) in HAND_SLICES.items(): + key = (person_idx, side) + if not meta[side]["distance_eligible"]: + previous.pop(key, None) + continue + prior = previous.get(key) + if prior is not None and prior[0] == record["frame_index"] - 1: + step = np.linalg.norm(kpts[start:end] - prior[1], axis=1) + values[f"P{person_idx + 1}_{side}"].extend(float(v) for v in step) + previous[key] = (record["frame_index"], kpts[start:end]) + return {name: summarize(items) for name, items in values.items()} + + +def metadata_summary(records: list[dict], variant: str) -> dict: + state_counts: Counter = Counter() + source_counts: Counter = Counter() + eligible_counts: Counter = Counter() + alarm_counts: Counter = Counter() + switches: Counter = Counter() + previous_source: dict[tuple[int, str, str], str] = {} + for record in records: + for person_idx, person_meta in enumerate(record[f"{variant}_metadata"]): + for side in ("left", "right"): + meta = person_meta[side] + key_prefix = f"P{person_idx + 1}_{side}" + state_counts[f"{key_prefix}_{meta['hand_state']}"] += 1 + source_counts[f"{key_prefix}_{meta['hand_source']}"] += 1 + if meta["distance_eligible"]: + eligible_counts[key_prefix] += 1 + if meta["alarm_eligible"]: + alarm_counts[key_prefix] += 1 + for kind in ("limb", "hand"): + source = meta[f"{kind}_source"] + key = (person_idx, side, kind) + if key in previous_source and previous_source[key] != source: + switches[f"{key_prefix}_{kind}"] += 1 + previous_source[key] = source + return { + "hand_state_counts": dict(sorted(state_counts.items())), + "source_counts": dict(sorted(source_counts.items())), + "distance_eligible_frames": dict(sorted(eligible_counts.items())), + "alarm_eligible_frames": dict(sorted(alarm_counts.items())), + "source_switches": dict(sorted(switches.items())), + } + + +def main() -> None: + args = parse_args() + args.output.mkdir(parents=True, exist_ok=True) + raw_records = [ + json.loads(line) for line in args.raw_jsonl.read_text(encoding="utf-8").splitlines() + ] + isolated_records = [ + json.loads(line) + for line in args.isolated_jsonl.read_text(encoding="utf-8").splitlines() + ] + + cap = cv2.VideoCapture(str(args.video)) + fps = float(cap.get(cv2.CAP_PROP_FPS)) + width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)) + height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)) + writer = cv2.VideoWriter( + str(args.output / "rtmw_l_raw_spatial_temporal.mp4"), + cv2.VideoWriter_fourcc(*"mp4v"), + fps, + (1920, 360), + ) + + tracks = [PersonTrack(), PersonTrack()] + output_records: list[dict] = [] + frame_idx = 0 + while True: + ok, frame = cap.read() + if not ok: + break + raw_record = raw_records[frame_idx] + isolated_record = isolated_records[frame_idx] + boxes = raw_record["boxes"][:2] + raw_keypoints = np.asarray(raw_record["L"]["keypoints"][:2], dtype=np.float64) + raw_scores = np.asarray(raw_record["L"]["raw_scores"][:2], dtype=np.float64) + isolated_keypoints = np.asarray( + isolated_record["L"]["keypoints"][:2], dtype=np.float64 + ) + isolated_scores = np.asarray( + isolated_record["masked_raw_scores"][:2], dtype=np.float64 + ) + + spatial_keypoints: list[np.ndarray] = [] + spatial_metadata: list[dict] = [] + temporal_keypoints: list[np.ndarray] = [] + temporal_selected_keypoints: list[np.ndarray] = [] + temporal_metadata: list[dict] = [] + for person_idx in range(2): + spatial, spatial_meta = fuse_person( + raw_keypoints[person_idx], + raw_scores[person_idx], + isolated_keypoints[person_idx], + isolated_scores[person_idx], + boxes[person_idx], + track=None, + ) + temporal_unfiltered, temporal_meta = fuse_person( + raw_keypoints[person_idx], + raw_scores[person_idx], + isolated_keypoints[person_idx], + isolated_scores[person_idx], + boxes[person_idx], + track=tracks[person_idx], + ) + temporal = adaptive_smooth( + temporal_unfiltered, boxes[person_idx], tracks[person_idx] + ) + spatial_keypoints.append(spatial) + spatial_metadata.append(spatial_meta) + temporal_keypoints.append(temporal) + temporal_selected_keypoints.append(temporal_unfiltered) + temporal_metadata.append(temporal_meta) + + global_hands = build_global_hands( + temporal_selected_keypoints, temporal_metadata + ) + + output_record = { + "frame_index": frame_idx, + "boxes": boxes, + "raw_keypoints": np.round(raw_keypoints, 3).tolist(), + "spatial_keypoints": np.round(np.asarray(spatial_keypoints), 3).tolist(), + "spatial_metadata": spatial_metadata, + "temporal_keypoints": np.round(np.asarray(temporal_keypoints), 3).tolist(), + "temporal_selected_keypoints": np.round( + np.asarray(temporal_selected_keypoints), 3 + ).tolist(), + "temporal_metadata": temporal_metadata, + "global_hands": global_hands, + } + output_records.append(output_record) + + raw_panel = draw_pose(frame, boxes, raw_keypoints, f"RTMW-L raw | frame {frame_idx}") + spatial_panel = draw_fused( + frame, + boxes, + spatial_keypoints, + spatial_metadata, + f"spatial fusion | frame {frame_idx}", + ) + temporal_panel = draw_fused( + frame, + boxes, + temporal_keypoints, + temporal_metadata, + f"temporal fusion | frame {frame_idx}", + ) + panels = [ + cv2.resize(panel, (640, 360), interpolation=cv2.INTER_AREA) + for panel in (raw_panel, spatial_panel, temporal_panel) + ] + writer.write(np.hstack(panels)) + frame_idx += 1 + + cap.release() + writer.release() + + with (args.output / "fusion.jsonl").open("w", encoding="utf-8") as handle: + for record in output_records: + handle.write(json.dumps(record, ensure_ascii=False) + "\n") + + def comparison_records(variant: str) -> list[dict]: + return [ + { + "frame_index": record["frame_index"], + "boxes": record["boxes"], + "M": {"keypoints": record["raw_keypoints"]}, + "L": {"keypoints": record[f"{variant}_keypoints"]}, + } + for record in output_records + ] + + summary = { + "frames": len(output_records), + "fusion_policy": { + "candidates": ["RTMW-L raw", "RTMW-L smaller-person isolated"], + "selection": [ + "limb and hand selected separately", + "geometry validity", + "constant-velocity temporal consistency", + "source-switch hysteresis", + "motion-adaptive EMA in person-box coordinates", + ], + "distance_rule": ( + "Only consensus or temporally stable raw visible endpoints are eligible. " + "Inferred/conflict/rejected hands are excluded from precise distance." + ), + "alarm_rule": ( + "Alarm attack is asymmetric: a geometrically valid raw endpoint whose score " + "is clearly stronger than the isolated view is admitted immediately. Alarm " + "release should still require temporal confirmation and hysteresis." + ), + "risk_distance_note": ( + "Use global_hands and temporal_selected_keypoints (unfiltered) for alarm entry " + "to avoid smoothing delay. Use distance_eligible hands and temporal_keypoints " + "for stable reporting. When the door edge is available, a conservative " + "implementation can use the smaller of the unfiltered and smoothed distances." + ), + }, + "spatial": metadata_summary(output_records, "spatial"), + "temporal": metadata_summary(output_records, "temporal"), + "active_hand_frame_step": { + "spatial": active_hand_temporal_metrics(output_records, "spatial"), + "temporal": active_hand_temporal_metrics(output_records, "temporal"), + }, + "raw_vs_spatial_temporal_metrics": temporal_metrics(comparison_records("spatial")), + "raw_vs_temporal_temporal_metrics": temporal_metrics(comparison_records("temporal")), + "global_hand_count_frames": dict( + sorted(Counter(len(record["global_hands"]) for record in output_records).items()) + ), + "cross_person_hand_aliases_total": sum( + len(hand["aliases"]) + for record in output_records + for hand in record["global_hands"] + ), + } + with (args.output / "summary.json").open("w", encoding="utf-8") as handle: + json.dump(summary, handle, ensure_ascii=False, indent=2) + + key_frames = [0, 45, 50, 55, 120, 132, 136, 180, 240, 252, 295, 300, 304, 360, 411] + make_montage( + args.video, + output_records, + key_frames, + args.output / "rtmw_l_raw_spatial_temporal_montage.jpg", + ) + print(json.dumps(summary, ensure_ascii=False, indent=2)) + + +if __name__ == "__main__": + main() diff --git a/pipeline/latest_result.py b/pipeline/latest_result.py new file mode 100644 index 0000000..1ab68ae --- /dev/null +++ b/pipeline/latest_result.py @@ -0,0 +1,64 @@ +"""One owned, immutable latest result, mirroring the 0920 LatestResult contract.""" +from __future__ import annotations + +from collections.abc import Mapping +import threading +from types import MappingProxyType + + +def freeze(value): + if isinstance(value, Mapping): + return MappingProxyType({key: freeze(item) for key, item in value.items()}) + if isinstance(value, (list, tuple)): + return tuple(freeze(item) for item in value) + return value + + +def result_to_dict(value): + """Return a mutable JSON-compatible copy when a consumer needs serialization.""" + if isinstance(value, Mapping): + return {key: result_to_dict(item) for key, item in value.items()} + if isinstance(value, tuple): + return [result_to_dict(item) for item in value] + return value + + +class LatestResult: + """Readers share a stable snapshot; reading never removes or acknowledges it. + + Only the latest result is retained here. Slow readers can skip frames, and + references they already hold remain valid after subsequent publications. + """ + + def __init__(self): + self._condition = threading.Condition() + self._value = None + self._closed = False + + def publish(self, result): + if result is None: + raise ValueError("Cannot publish a null result") + snapshot = freeze(result) + with self._condition: + if self._closed: + raise RuntimeError("Result publisher is closed") + self._value = snapshot + self._condition.notify_all() + + def load(self): + with self._condition: + return self._value + + def wait(self, after_frame_index=-1, timeout=None): + """Wait for a newer result; return None on timeout/end without a new frame.""" + with self._condition: + def has_new_frame(): + return self._value is not None and self._value['frame_index'] > after_frame_index + + self._condition.wait_for(lambda: has_new_frame() or self._closed, timeout) + return self._value if has_new_frame() else None + + def close(self): + with self._condition: + self._closed = True + self._condition.notify_all() diff --git a/pipeline/live_pose_engine.py b/pipeline/live_pose_engine.py new file mode 100644 index 0000000..b0140e5 --- /dev/null +++ b/pipeline/live_pose_engine.py @@ -0,0 +1,248 @@ +"""Resident YOLO11s + RTMW-L pose stage extracted from the v9.9 pipeline.""" + +from __future__ import annotations + +import argparse +import os +import time +from collections import defaultdict +from pathlib import Path + +import numpy as np + + +SIDES = ("left", "right") + + +class LivePoseEngine: + """Run the v9.9 YOLO/RTMW pose and temporal-fusion path without EdgeTAM.""" + + def __init__(self, args: argparse.Namespace): + self._trt_dll_directory = None + if args.pose_provider == "tensorrt" and os.name == "nt": + try: + import tensorrt_libs + except ImportError as exc: + raise RuntimeError( + "--pose-provider tensorrt on Windows requires the " + "tensorrt-cu12-libs package in the active Python environment." + ) from exc + + trt_lib_dir = Path(tensorrt_libs.__file__).resolve().parent + self._trt_dll_directory = os.add_dll_directory(str(trt_lib_dir)) + os.environ["PATH"] = ( + str(trt_lib_dir) + os.pathsep + os.environ.get("PATH", "") + ) + + if args.pose_device.lower().startswith("cuda"): + import onnxruntime as ort + + if "CUDAExecutionProvider" not in ort.get_available_providers(): + raise RuntimeError( + "--pose-device cuda requires onnxruntime-gpu. Uninstall " + "onnxruntime and install onnxruntime-gpu==1.22.0." + ) + + from rtmlib import RTMPose + from ultralytics import YOLO + from predict_optimized_pose_videos import BoxTracker + + self.args = args + self.detector = YOLO(str(args.det_model)) + self.pose = RTMPose( + str(args.pose_model), + model_input_size=(288, 384), + backend="onnxruntime", + device=args.pose_device, + ) + if args.pose_provider == "tensorrt": + import onnxruntime as ort + + available = ort.get_available_providers() + if "TensorrtExecutionProvider" not in available: + raise RuntimeError( + "--pose-provider tensorrt requires ONNX Runtime with " + "TensorrtExecutionProvider. " + f"Available providers: {available}" + ) + cache_dir = args.pose_trt_cache.expanduser().resolve() + cache_dir.mkdir(parents=True, exist_ok=True) + device_id = ( + int(str(args.pose_device).split(":")[-1]) + if ":" in str(args.pose_device) + else 0 + ) + trt_options = { + "device_id": device_id, + "trt_fp16_enable": "True", + "trt_engine_cache_enable": "True", + "trt_engine_cache_path": str(cache_dir), + "trt_timing_cache_enable": "True", + "trt_timing_cache_path": str(cache_dir), + } + self.pose.session = ort.InferenceSession( + str(args.pose_model.resolve()), + providers=[ + ("TensorrtExecutionProvider", trt_options), + ("CUDAExecutionProvider", {"device_id": device_id}), + ], + ) + active_providers = self.pose.session.get_providers() + if not active_providers or active_providers[0] != "TensorrtExecutionProvider": + raise RuntimeError( + "RTMW TensorRT initialization failed instead of using the " + "requested provider. Active providers: " + f"{active_providers}" + ) + print(f"RTMW ONNX Runtime providers: {active_providers}") + + self.box_tracker = BoxTracker(max_age=30) + self.people = {} + self.timings = defaultdict(list) + self.isolated_inferences = 0 + + def infer(self, frame: np.ndarray, frame_idx: int) -> dict: + from compare_rtmw_ml import HAND_SLICES + from fuse_rtmw_l_temporal import adaptive_smooth, build_global_hands, fuse_person + from occlusion_aware_pose_v2 import Measurement, resolve_global_hand_assignments + from predict_optimized_pose_videos import ( + LivePersonState, + filter_person_boxes, + mask_other_people, + update_live_v2, + ) + + infer_started = time.perf_counter() + height, width = frame.shape[:2] + started = time.perf_counter() + detection_height = max( + 1, min(height, round(height * self.args.det_roi_y_max_ratio)) + ) + if hasattr(self, "native"): + boxes, confidences = self.native.detect(frame[:detection_height]) + self.timings["detector_ms"].append((time.perf_counter() - started) * 1000) + else: + detection = self.detector.predict( + frame[:detection_height], + imgsz=self.args.det_imgsz, + conf=self.args.det_conf, + iou=0.30, + classes=[0], + device=self.args.det_device, + verbose=False, + )[0] + self.timings["detector_ms"].append((time.perf_counter() - started) * 1000) + boxes, confidences, _ = filter_person_boxes( + detection, width, height, self.args.max_people + ) + track_ids = self.box_tracker.update(boxes, frame_idx, (height, width)) + + # Retire disappeared people in long-running camera sessions. + for stale_id in list(self.people): + if stale_id not in self.box_tracker.tracks: + del self.people[stale_id] + if boxes: + started = time.perf_counter() + raw_keypoints, raw_scores = self.pose(frame, bboxes=boxes) + self.timings["raw_pose_ms"].append((time.perf_counter() - started) * 1000) + isolated_keypoints = raw_keypoints.copy() + isolated_scores = raw_scores.copy() + started = time.perf_counter() + for person_idx, bbox in enumerate(boxes): + target_area = max((bbox[2] - bbox[0]) * (bbox[3] - bbox[1]), 1.0) + needs_isolation = any( + idx != person_idx + and max((other[2] - other[0]) * (other[3] - other[1]), 1.0) + < 0.5 * target_area + for idx, other in enumerate(boxes) + ) + if not needs_isolation: + continue + masked = mask_other_people(frame, boxes, person_idx) + keypoints, scores = self.pose(masked, bboxes=[bbox]) + isolated_keypoints[person_idx] = keypoints[0] + isolated_scores[person_idx] = scores[0] + self.isolated_inferences += 1 + self.timings["isolated_extra_ms"].append( + (time.perf_counter() - started) * 1000 + ) + else: + raw_keypoints = np.empty((0, 133, 2), dtype=np.float64) + raw_scores = np.empty((0, 133), dtype=np.float64) + isolated_keypoints = raw_keypoints.copy() + isolated_scores = raw_scores.copy() + + selected_list, smoothed_list, metadata_list, states = [], [], [], [] + for person_idx, (bbox, track_id) in enumerate(zip(boxes, track_ids)): + state = self.people.setdefault(track_id, LivePersonState()) + selected, metadata = fuse_person( + raw_keypoints[person_idx], + raw_scores[person_idx], + isolated_keypoints[person_idx], + isolated_scores[person_idx], + bbox, + state.fusion, + ) + selected_list.append(selected) + smoothed_list.append(adaptive_smooth(selected, bbox, state.fusion)) + metadata_list.append(metadata) + states.append(state) + + measurements, endpoint_tracks = {}, {} + for person_idx, (selected, metadata, state) in enumerate( + zip(selected_list, metadata_list, states) + ): + for side, (start, end) in HAND_SLICES.items(): + meta = metadata[side] + measurements[(person_idx, side)] = ( + Measurement( + keypoints=selected[start:end], + precise=bool(meta["distance_eligible"]), + provisional=not bool(meta["distance_eligible"]), + quality=float(meta.get("selected_hand_score", 0.0)), + ) + if meta["alarm_eligible"] + else None + ) + endpoint_tracks[(person_idx, side)] = state.endpoints[side] + measurements, assignments = resolve_global_hand_assignments( + measurements, endpoint_tracks, boxes, 30 + ) + + v2_list, v2_meta_list = [], [] + for person_idx, (bbox, selected, smoothed, metadata, state) in enumerate( + zip(boxes, selected_list, smoothed_list, metadata_list, states) + ): + v2, v2_meta = update_live_v2( + frame_idx, + bbox, + selected, + smoothed, + metadata, + state, + {side: measurements[(person_idx, side)] for side in SIDES}, + {side: assignments.get((person_idx, side)) for side in SIDES}, + 12, + 30, + ) + v2_list.append(v2) + v2_meta_list.append(v2_meta) + + self.timings["pose_pipeline_total_ms"].append( + (time.perf_counter() - infer_started) * 1000 + ) + return { + "frame_index": frame_idx, + "track_ids": track_ids, + "boxes": np.round(np.asarray(boxes), 3).tolist() if boxes else [], + "det_confidences": [round(value, 4) for value in confidences], + "raw_keypoints": np.round(raw_keypoints, 3).tolist(), + "temporal_selected_keypoints": np.round( + np.asarray(selected_list), 3 + ).tolist(), + "temporal_keypoints": np.round(np.asarray(smoothed_list), 3).tolist(), + "temporal_metadata": metadata_list, + "v2_keypoints": np.round(np.asarray(v2_list), 3).tolist(), + "v2_metadata": v2_meta_list, + "global_hands": build_global_hands(selected_list, metadata_list), + } diff --git a/pipeline/native_backend.py b/pipeline/native_backend.py new file mode 100644 index 0000000..91895f5 --- /dev/null +++ b/pipeline/native_backend.py @@ -0,0 +1,75 @@ +"""Native Jetson TensorRT adapter; temporal/risk policy stays in upstream Python.""" +from __future__ import annotations +import ctypes as C +from collections import defaultdict, deque +from pathlib import Path +import numpy as np +from live_pose_engine import LivePoseEngine +from predict_optimized_pose_videos import BoxTracker + +ROOT = Path(__file__).resolve().parents[1] + +class NativeBackend: + def __init__(self, args): + self.lib = C.CDLL(str(ROOT / 'build/libvision_native.so')) + pointer = C.c_void_p + signatures = { + 'vap_error': ([], C.c_char_p), + 'vap_engine': ([C.c_char_p,C.c_int,C.c_int,C.c_int,C.c_float],pointer), + 'vap_engine_close': ([pointer],None), + 'vap_detect': ([pointer,pointer,C.c_int,C.c_int,pointer],C.c_int), + 'vap_pose': ([pointer,pointer,C.c_int,C.c_int,pointer,C.c_int,pointer],C.c_int), + 'vap_camera': ([C.c_char_p,C.c_int,C.c_int,C.c_int,C.c_int],pointer), + 'vap_camera_close': ([pointer],None), + 'vap_read': ([pointer,pointer,pointer],C.c_int), + 'vap_can': ([pointer],None), + } + for name,(inputs,output) in signatures.items(): + fn=getattr(self.lib,name);fn.argtypes=inputs;fn.restype=output + self.engine = self.lib.vap_engine(str(args.engines.resolve()).encode(),args.width,args.height,args.max_people,args.det_conf) + if not self.engine: self.fail() + self.camera=None + self.args=args + + def fail(self): + raise RuntimeError(self.lib.vap_error().decode()) + + def open_camera(self): + a=self.args + self.camera=self.lib.vap_camera(a.camera.encode(),a.width,a.height,a.fps,not a.no_can) + if not self.camera:self.fail() + + def read(self): + image=np.empty((self.args.height,self.args.width,3),np.uint8) + meta=np.zeros(4,np.float64) + result=self.lib.vap_read(self.camera,image.ctypes.data,meta.ctypes.data) + if result<0:self.fail() + if result==0:return None,None + return image,{'sequence':int(meta[0]),'capture_monotonic_ms':float(meta[1]),'dropped_before_inference':int(meta[2]),'driver_timestamp_ms':float(meta[3]),'latency_clock':'host_monotonic_at_dequeue'} + + def can_snapshot(self): + data=np.zeros(5,np.float64);self.lib.vap_can(data.ctypes.data) + return {'positions':[int(x) for x in data[:4]],'age_ms':float(data[4])} + + def detect(self, frame): + frame=np.ascontiguousarray(frame);out=np.empty((self.args.max_people,5),np.float32) + n=self.lib.vap_detect(self.engine,frame.ctypes.data,frame.shape[1],frame.shape[0],out.ctypes.data) + if n<0:self.fail() + return out[:n,:4].tolist(),out[:n,4].tolist() + + def pose(self,frame,bboxes): + frame=np.ascontiguousarray(frame);boxes=np.ascontiguousarray(bboxes,dtype=np.float32) + out=np.empty((len(boxes),133,3),np.float32) + n=self.lib.vap_pose(self.engine,frame.ctypes.data,frame.shape[1],frame.shape[0],boxes.ctypes.data,len(boxes),out.ctypes.data) + if n<0:self.fail() + return out[:n,:,:2].astype(np.float64),out[:n,:,2].astype(np.float64) + + def close(self): + if self.camera:self.lib.vap_camera_close(self.camera);self.camera=None + if self.engine:self.lib.vap_engine_close(self.engine);self.engine=None + +class NativePoseEngine(LivePoseEngine): + def __init__(self,args,native): + self.args=args;self.native=native;self.pose=native.pose + self.box_tracker=BoxTracker(max_age=30);self.people={} + self.timings=defaultdict(lambda:deque(maxlen=3000));self.isolated_inferences=0 diff --git a/pipeline/occlusion_aware_pose_v2.py b/pipeline/occlusion_aware_pose_v2.py new file mode 100644 index 0000000..5b7c73f --- /dev/null +++ b/pipeline/occlusion_aware_pose_v2.py @@ -0,0 +1,1130 @@ +"""Occlusion-aware temporal pose reconstruction on top of RTMW-L fusion. + +Visible endpoints are filtered and optionally connected with a two-bone IK +solution. Short gaps are interpolated/predicted, medium gaps use a local motion +corridor, and long occlusions retain an undrawn shoulder sector for downstream +door-band intersection instead of emitting a fake precise wrist. +""" + +from __future__ import annotations + +import argparse +import json +import math +from collections import Counter, defaultdict +from dataclasses import dataclass +from pathlib import Path + +import cv2 +import numpy as np + +from compare_rtmw_ml import HAND_EDGES_LOCAL, HAND_SLICES, normalized_keypoints, summarize +from fuse_rtmw_l_temporal import BODY_SIDE_INDICES, draw_dashed_line, draw_fused + + +STATE_COLORS = { + "observed": (40, 220, 40), + "provisional": (255, 190, 40), + "recovered": (220, 180, 40), + "reassigned": (220, 120, 40), + "interpolated": (0, 220, 255), + "predicted_short": (0, 140, 255), + "reachable_medium": (0, 120, 255), + "reachable_long": (30, 80, 255), +} + + +@dataclass +class Measurement: + keypoints: np.ndarray + precise: bool + provisional: bool + source: str = "model" + quality: float = 0.0 + + +@dataclass +class EndpointTrack: + wrist: np.ndarray | None = None + velocity: np.ndarray | None = None + hand_shape: np.ndarray | None = None + last_observed_frame: int | None = None + missed_frames: int = 0 + last_elbow: np.ndarray | None = None + last_direction: np.ndarray | None = None + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser() + parser.add_argument("--video", type=Path, required=True) + parser.add_argument("--fusion-jsonl", type=Path, required=True) + parser.add_argument("--output", type=Path, required=True) + parser.add_argument("--short-occlusion-frames", type=int, default=12) + parser.add_argument("--medium-occlusion-frames", type=int, default=30) + parser.add_argument("--lookahead-frames", type=int, default=3) + return parser.parse_args() + + +def bbox_diagonal(bbox: list[float]) -> float: + return max(math.hypot(bbox[2] - bbox[0], bbox[3] - bbox[1]), 1.0) + + +def point_distance_ratio(a: np.ndarray, b: np.ndarray, bbox: list[float]) -> float: + return float(np.linalg.norm(a - b)) / bbox_diagonal(bbox) + + +def collect_measurements( + records: list[dict], + person_idx: int, + side: str, + lookahead_frames: int, +) -> list[Measurement | None]: + start, end = HAND_SLICES[side] + measurements: list[Measurement | None] = [None] * len(records) + observed_frames = [] + for frame_idx, record in enumerate(records): + meta = record["temporal_metadata"][person_idx][side] + if not meta["alarm_eligible"]: + continue + keypoints = np.asarray( + record["temporal_selected_keypoints"][person_idx][start:end], + dtype=np.float64, + ) + measurements[frame_idx] = Measurement( + keypoints=keypoints, + precise=bool(meta["distance_eligible"]), + provisional=not bool(meta["distance_eligible"]), + quality=float(meta.get("selected_hand_score", 0.0)), + ) + observed_frames.append(frame_idx) + + # A small future window is allowed for stable reporting. The alarm fast + # path still uses the unmodified current-frame candidates from V1. + for left, right in zip(observed_frames, observed_frames[1:]): + gap = right - left - 1 + if gap <= 0 or gap > lookahead_frames: + continue + left_points = measurements[left].keypoints + right_points = measurements[right].keypoints + for offset in range(1, gap + 1): + ratio = offset / (gap + 1) + points = (1.0 - ratio) * left_points + ratio * right_points + measurements[left + offset] = Measurement( + keypoints=points, + precise=False, + provisional=False, + source="fixed_lag_interpolation", + quality=min(measurements[left].quality, measurements[right].quality), + ) + return measurements + + +def estimate_bone_ratios(records: list[dict]) -> dict[tuple[int, str], tuple[float, float]]: + values: dict[tuple[int, str], list[tuple[float, float]]] = defaultdict(list) + for record in records: + for person_idx, bbox in enumerate(record["boxes"]): + keypoints = np.asarray( + record["temporal_selected_keypoints"][person_idx], dtype=np.float64 + ) + diagonal = bbox_diagonal(bbox) + for side, (shoulder, elbow, wrist) in BODY_SIDE_INDICES.items(): + meta = record["temporal_metadata"][person_idx][side] + if meta["limb_state"] != "consensus": + continue + upper = float(np.linalg.norm(keypoints[shoulder] - keypoints[elbow])) / diagonal + forearm = float(np.linalg.norm(keypoints[elbow] - keypoints[wrist])) / diagonal + if 0.04 <= upper <= 0.65 and 0.04 <= forearm <= 0.65: + values[(person_idx, side)].append((upper, forearm)) + + estimates: dict[tuple[int, str], tuple[float, float]] = {} + person_count = len(records[0]["boxes"]) + for person_idx in range(person_count): + for side in ("left", "right"): + samples = values[(person_idx, side)] + if len(samples) < 20: + opposite = "right" if side == "left" else "left" + samples = values[(person_idx, opposite)] + if samples: + array = np.asarray(samples) + estimates[(person_idx, side)] = ( + float(np.median(array[:, 0])), + float(np.median(array[:, 1])), + ) + else: + estimates[(person_idx, side)] = (0.24, 0.22) + return estimates + + +def solve_two_bone_elbow( + shoulder: np.ndarray, + wrist: np.ndarray, + upper_length: float, + forearm_length: float, + reference_elbow: np.ndarray, +) -> np.ndarray: + direction = wrist - shoulder + distance = float(np.linalg.norm(direction)) + if distance < 1e-6: + return reference_elbow.copy() + unit = direction / distance + + if distance >= upper_length + forearm_length: + return shoulder + unit * upper_length + if distance <= abs(upper_length - forearm_length): + sign = 1.0 if upper_length >= forearm_length else -1.0 + return shoulder + unit * upper_length * sign + + along = ( + upper_length * upper_length + - forearm_length * forearm_length + + distance * distance + ) / (2.0 * distance) + height = math.sqrt(max(upper_length * upper_length - along * along, 0.0)) + base = shoulder + unit * along + perpendicular = np.asarray([-unit[1], unit[0]]) + candidates = (base + perpendicular * height, base - perpendicular * height) + return min(candidates, key=lambda candidate: float(np.linalg.norm(candidate - reference_elbow))) + + +def update_endpoint_track( + track: EndpointTrack, + measurement: Measurement | None, + frame_idx: int, + short_occlusion_frames: int, + medium_occlusion_frames: int, +) -> tuple[np.ndarray | None, np.ndarray | None, str]: + if measurement is not None: + measured_wrist = measurement.keypoints[0] + measured_shape = measurement.keypoints - measured_wrist + if track.wrist is None or track.missed_frames > short_occlusion_frames: + track.wrist = measured_wrist.copy() + track.velocity = np.zeros(2, dtype=np.float64) + track.hand_shape = measured_shape.copy() + else: + prediction = track.wrist + track.velocity + residual = measured_wrist - prediction + alpha = 0.78 if measurement.precise else 0.88 + beta = 0.18 if measurement.precise else 0.24 + track.wrist = prediction + alpha * residual + track.velocity = 0.82 * track.velocity + beta * residual + shape_alpha = 0.55 if measurement.precise else 0.7 + track.hand_shape = ( + shape_alpha * measured_shape + (1.0 - shape_alpha) * track.hand_shape + ) + track.last_observed_frame = frame_idx + track.missed_frames = 0 + if measurement.source == "global_reassignment": + state = "reassigned" + elif measurement.source == "temporal_recovery": + state = "recovered" + elif measurement.precise: + state = "observed" + elif measurement.provisional: + state = "provisional" + else: + state = "interpolated" + return track.wrist.copy(), track.wrist + track.hand_shape, state + + track.missed_frames += 1 + if track.wrist is None: + return None, None, "reachable_long" + if track.missed_frames <= short_occlusion_frames: + track.wrist = track.wrist + track.velocity + track.velocity *= 0.86 + return track.wrist.copy(), track.wrist + track.hand_shape, "predicted_short" + if track.missed_frames <= medium_occlusion_frames: + return None, None, "reachable_medium" + return None, None, "reachable_long" + + +def temporal_recovery_measurement( + keypoints: np.ndarray, + metadata: dict, + side: str, + bbox: list[float], + track: EndpointTrack, + medium_occlusion_frames: int, +) -> Measurement | None: + """Recover a weak current-frame hand only when history strongly supports it.""" + if track.wrist is None or track.missed_frames > medium_occlusion_frames: + return None + if metadata.get("hand_state") not in {"conflict", "inferred"}: + return None + if not metadata.get("selected_hand_valid", False): + return None + if float(metadata.get("selected_hand_score", 0.0)) < 2.0: + return None + link = float(metadata.get("hand_selection", {}).get("link_to_selected_limb", 1.0)) + if link > 0.08: + return None + + start, end = HAND_SLICES[side] + hand = np.asarray(keypoints[start:end], dtype=np.float64) + if hand.shape != (end - start, 2) or not np.isfinite(hand).all(): + return None + prediction = track.wrist + track.velocity + diagonal = bbox_diagonal(bbox) + gate = min(0.20, 0.05 + 0.008 * max(track.missed_frames, 1)) * diagonal + if float(np.linalg.norm(hand[0] - prediction)) > gate: + return None + return Measurement( + keypoints=hand, + precise=False, + provisional=True, + source="temporal_recovery", + quality=float(metadata.get("selected_hand_score", 0.0)), + ) + + +def hand_span_px(hand: np.ndarray) -> float: + return math.hypot(float(np.ptp(hand[:, 0])), float(np.ptp(hand[:, 1]))) + + +def bbox_area(bbox: list[float]) -> float: + return max((bbox[2] - bbox[0]) * (bbox[3] - bbox[1]), 1.0) + + +def reset_endpoint_track(track: EndpointTrack) -> None: + track.wrist = None + track.velocity = None + track.hand_shape = None + track.last_observed_frame = None + track.missed_frames = 0 + track.last_elbow = None + track.last_direction = None + + +def resolve_global_hand_assignments( + frame_measurements: dict[tuple[int, str], Measurement | None], + tracks: dict[tuple[int, str], EndpointTrack], + boxes: list[list[float]], + medium_occlusion_frames: int, +) -> tuple[dict[tuple[int, str], Measurement | None], dict[tuple[int, str], dict]]: + """Resolve overlapping-person hands before binding them to body sides. + + Large foreground top-down crops may include a smaller background person. + The same physical hand can therefore be emitted for both people, or only by + the smaller crop even though it belongs to the foreground person. Direct + duplicates are kept on the foreground track when their apparent scales + agree. A missing foreground side can then reclaim a nearby candidate using + its recent hand trajectory and historical hand scale. + """ + resolved = dict(frame_measurements) + events: dict[tuple[int, str], dict] = {} + keys = list(resolved) + + # First collapse direct cross-person duplicates. Keep this deliberately + # asymmetric: a substantially larger foreground crop must also carry the + # substantially larger historical hand shape. + for index, first_key in enumerate(keys): + first = resolved[first_key] + if first is None: + continue + for second_key in keys[index + 1 :]: + second = resolved[second_key] + if second is None or first_key[0] == second_key[0]: + continue + first_hand = first.keypoints + second_hand = second.keypoints + first_span = hand_span_px(first_hand) + second_span = hand_span_px(second_hand) + center_distance = float( + np.linalg.norm(np.mean(first_hand, axis=0) - np.mean(second_hand, axis=0)) + ) + if center_distance > max(70.0, 0.5 * max(first_span, second_span)): + continue + + first_area = bbox_area(boxes[first_key[0]]) + second_area = bbox_area(boxes[second_key[0]]) + if first_area >= 2.5 * second_area and ( + first_span >= 1.45 * second_span + or first.quality >= second.quality + 0.6 + ): + owner_key, released_key = first_key, second_key + elif second_area >= 2.5 * first_area and ( + second_span >= 1.45 * first_span + or second.quality >= first.quality + 0.6 + ): + owner_key, released_key = second_key, first_key + else: + continue + resolved[released_key] = None + reset_endpoint_track(tracks[released_key]) + events[owner_key] = { + "method": "foreground_duplicate_owner", + "released_person_id": released_key[0] + 1, + "released_side": released_key[1], + } + events[released_key] = { + "method": "reassigned_away", + "assigned_person_id": owner_key[0] + 1, + "assigned_side": owner_key[1], + } + + # Then recover missing foreground hands from another crop. The candidate + # must fall inside the target's recent motion gate and be much smaller than + # the target's historical foreground hand shape. + available_sources = [key for key, value in resolved.items() if value is not None] + target_candidates = [] + for target_key, target_measurement in resolved.items(): + if target_measurement is not None: + continue + target_track = tracks[target_key] + if ( + target_track.wrist is None + or target_track.hand_shape is None + or target_track.missed_frames > medium_occlusion_frames + ): + continue + target_area = bbox_area(boxes[target_key[0]]) + predicted_center = target_track.wrist + np.mean(target_track.hand_shape, axis=0) + target_span = hand_span_px(target_track.hand_shape) + gate = 0.20 * bbox_diagonal(boxes[target_key[0]]) + for source_key in available_sources: + if source_key[0] == target_key[0]: + continue + source = resolved[source_key] + if source is None or source.source != "model": + continue + source_area = bbox_area(boxes[source_key[0]]) + if target_area < 2.5 * source_area: + continue + source_center = np.mean(source.keypoints, axis=0) + source_span = hand_span_px(source.keypoints) + distance = float(np.linalg.norm(source_center - predicted_center)) + if distance > gate or target_span < 1.45 * max(source_span, 1.0): + continue + target_candidates.append( + (distance / gate, target_key, source_key, source_center) + ) + + used_targets = set() + used_sources = set() + for _, target_key, source_key, source_center in sorted(target_candidates): + if target_key in used_targets or source_key in used_sources: + continue + target_track = tracks[target_key] + shape_center = np.mean(target_track.hand_shape, axis=0) + reassigned_wrist = source_center - shape_center + reassigned_hand = reassigned_wrist + target_track.hand_shape + resolved[target_key] = Measurement( + keypoints=reassigned_hand, + precise=False, + provisional=True, + source="global_reassignment", + quality=resolved[source_key].quality, + ) + resolved[source_key] = None + reset_endpoint_track(tracks[source_key]) + events[target_key] = { + "method": "foreground_temporal_reassignment", + "source_person_id": source_key[0] + 1, + "source_side": source_key[1], + "observed_center": np.round(source_center, 3).tolist(), + } + events[source_key] = { + "method": "reassigned_away", + "assigned_person_id": target_key[0] + 1, + "assigned_side": target_key[1], + } + used_targets.add(target_key) + used_sources.add(source_key) + return resolved, events + + +def track_direction(track: EndpointTrack) -> np.ndarray: + if track.last_direction is not None: + return track.last_direction / max(float(np.linalg.norm(track.last_direction)), 1e-6) + if track.velocity is not None and float(np.linalg.norm(track.velocity)) > 1e-6: + return track.velocity / np.linalg.norm(track.velocity) + return np.asarray([1.0, 0.0], dtype=np.float64) + + +def ellipse_polygon( + center: np.ndarray, + major_axis: float, + minor_axis: float, + direction: np.ndarray, + point_count: int = 64, +) -> np.ndarray: + perpendicular = np.asarray([-direction[1], direction[0]], dtype=np.float64) + angles = np.linspace(0.0, 2.0 * math.pi, point_count, endpoint=False) + return ( + center + + np.cos(angles)[:, None] * major_axis * direction + + np.sin(angles)[:, None] * minor_axis * perpendicular + ) + + +def clip_polygon_from_anchor_to_circle( + points: np.ndarray, + anchor: np.ndarray, + circle_center: np.ndarray, + radius: float, +) -> np.ndarray: + """Clip radial polygon samples to a shoulder-centered anatomical limit.""" + clipped = [] + offset = anchor - circle_center + for point in points: + if float(np.linalg.norm(point - circle_center)) <= radius: + clipped.append(point) + continue + ray = point - anchor + a = float(np.dot(ray, ray)) + b = 2.0 * float(np.dot(offset, ray)) + c = float(np.dot(offset, offset) - radius * radius) + discriminant = max(b * b - 4.0 * a * c, 0.0) + roots = [ + value + for value in ( + (-b - math.sqrt(discriminant)) / max(2.0 * a, 1e-6), + (-b + math.sqrt(discriminant)) / max(2.0 * a, 1e-6), + ) + if 0.0 <= value <= 1.0 + ] + clipped.append(anchor + (max(roots) if roots else 0.0) * ray) + return np.asarray(clipped, dtype=np.float64) + + +def medium_motion_region( + shoulder: np.ndarray, + upper_length: float, + forearm_length: float, + bbox: list[float], + track: EndpointTrack, + short_occlusion_frames: int, +) -> dict: + diagonal = bbox_diagonal(bbox) + direction = track_direction(track) + hand_shape = ( + track.hand_shape + if track.hand_shape is not None + else np.zeros((21, 2), dtype=np.float64) + ) + shape_center = np.mean(hand_shape, axis=0) + centered_shape = hand_shape - shape_center + perpendicular = np.asarray([-direction[1], direction[0]], dtype=np.float64) + hand_major = float(np.max(np.abs(centered_shape @ direction))) + hand_minor = float(np.max(np.abs(centered_shape @ perpendicular))) + + extra_frames = max(track.missed_frames - short_occlusion_frames, 1) + motion_major = min(0.20, 0.04 + 0.007 * extra_frames) * diagonal + motion_minor = min(0.13, 0.03 + 0.0045 * extra_frames) * diagonal + major_axis = max(hand_major + motion_major, 0.08 * diagonal) + minor_axis = max(hand_minor + motion_minor, 0.06 * diagonal) + + center = track.wrist + shape_center + tracked_hand_extent = float(np.max(np.linalg.norm(centered_shape, axis=1))) + hand_reach = max(tracked_hand_extent, 0.10 * diagonal) + anatomical_reach = upper_length + forearm_length + hand_reach + 0.02 * diagonal + shoulder_to_center = center - shoulder + center_distance = float(np.linalg.norm(shoulder_to_center)) + if center_distance > 0.96 * anatomical_reach: + center = shoulder + shoulder_to_center / center_distance * 0.96 * anatomical_reach + + polygon = ellipse_polygon(center, major_axis, minor_axis, direction) + polygon = clip_polygon_from_anchor_to_circle( + polygon, + center, + shoulder, + anatomical_reach, + ) + direction_deg = math.degrees(math.atan2(float(direction[1]), float(direction[0]))) + return { + "geometry": "motion_corridor", + "center": np.round(center, 3).tolist(), + "predicted_wrist": np.round(track.wrist, 3).tolist(), + "major_axis_px": round(major_axis, 3), + "minor_axis_px": round(minor_axis, 3), + "outer_radius_px": round(major_axis, 3), + "direction_deg": round(direction_deg, 3), + "max_reach_from_shoulder_px": round(anatomical_reach, 3), + "polygon": np.round(polygon, 3).tolist(), + "requires_door_intersection": True, + "render_full_region": True, + "distance_lower_bound_rule": "minimum distance from motion-corridor polygon to door_edge", + } + + +def long_reachable_region( + shoulder: np.ndarray, + upper_length: float, + forearm_length: float, + bbox: list[float], + track: EndpointTrack, + medium_occlusion_frames: int, +) -> dict: + diagonal = bbox_diagonal(bbox) + height = max(bbox[3] - bbox[1], 1.0) + arm_reach = upper_length + forearm_length + hand_reach = min(0.13 * diagonal, 0.17 * height) + safety_margin = 0.025 * diagonal + outer_radius = arm_reach + hand_reach + safety_margin + inner_radius = max(abs(upper_length - forearm_length) - 0.04 * diagonal, 0.0) + if track.last_direction is None: + direction_deg = 0.0 + half_angle_deg = 180.0 + else: + direction = track_direction(track) + direction_deg = math.degrees(math.atan2(float(direction[1]), float(direction[0]))) + half_angle_deg = min( + 180.0, + 65.0 + 2.0 * max(track.missed_frames - medium_occlusion_frames, 0), + ) + return { + "geometry": "shoulder_sector", + "center": np.round(shoulder, 3).tolist(), + "inner_radius_px": round(inner_radius, 3), + "outer_radius_px": round(outer_radius, 3), + "arm_reach_px": round(arm_reach, 3), + "hand_reach_px": round(hand_reach, 3), + "safety_margin_px": round(safety_margin, 3), + "direction_deg": round(direction_deg, 3), + "half_angle_deg": round(half_angle_deg, 3), + "requires_door_intersection": True, + "render_full_region": False, + "distance_lower_bound_rule": "intersect shoulder sector with door danger band before alarming", + } + + +def reachable_region( + shoulder: np.ndarray, + upper_length: float, + forearm_length: float, + bbox: list[float], + track: EndpointTrack, + state: str = "reachable_long", + short_occlusion_frames: int = 12, + medium_occlusion_frames: int = 30, +) -> dict: + if state == "reachable_medium" and track.wrist is not None: + return medium_motion_region( + shoulder, + upper_length, + forearm_length, + bbox, + track, + short_occlusion_frames, + ) + return long_reachable_region( + shoulder, + upper_length, + forearm_length, + bbox, + track, + medium_occlusion_frames, + ) + + +def calibrated_bone_lengths( + upper_ratio: float, + forearm_ratio: float, + bbox: list[float], +) -> tuple[float, float]: + diagonal = bbox_diagonal(bbox) + height = max(bbox[3] - bbox[1], 1.0) + upper = upper_ratio * diagonal + forearm = forearm_ratio * diagonal + total = upper + forearm + # Keep only a loose sanity cap. The previous 0.55 * bbox-height cap was + # unsafe for close, clipped CMS views: a valid projected arm can occupy + # almost the full detection height. The ratios are already robust medians, + # so this cap only limits catastrophic calibration outliers. + maximum_total = 0.95 * diagonal + if total > maximum_total: + scale = maximum_total / total + upper *= scale + forearm *= scale + return upper, forearm + + +def process_sequence( + records: list[dict], + short_occlusion_frames: int, + medium_occlusion_frames: int, + lookahead_frames: int, +) -> list[dict]: + person_count = len(records[0]["boxes"]) + measurements = { + (person_idx, side): collect_measurements( + records, person_idx, side, lookahead_frames + ) + for person_idx in range(person_count) + for side in ("left", "right") + } + bone_ratios = estimate_bone_ratios(records) + tracks = { + (person_idx, side): EndpointTrack() + for person_idx in range(person_count) + for side in ("left", "right") + } + + output = [] + for frame_idx, record in enumerate(records): + boxes = record["boxes"] + base_keypoints = np.asarray(record["temporal_keypoints"], dtype=np.float64) + unfiltered_keypoints = np.asarray( + record["temporal_selected_keypoints"], dtype=np.float64 + ) + v2_keypoints = base_keypoints.copy() + frame_measurements = { + (person_idx, side): measurements[(person_idx, side)][frame_idx] + for person_idx in range(person_count) + for side in ("left", "right") + } + frame_measurements, assignment_events = resolve_global_hand_assignments( + frame_measurements, + tracks, + boxes, + medium_occlusion_frames, + ) + people_meta = [] + for person_idx, bbox in enumerate(boxes): + diagonal = bbox_diagonal(bbox) + person_meta = {} + for side, (shoulder_idx, elbow_idx, wrist_idx) in BODY_SIDE_INDICES.items(): + track = tracks[(person_idx, side)] + measurement = frame_measurements[(person_idx, side)] + if measurement is None: + measurement = temporal_recovery_measurement( + unfiltered_keypoints[person_idx], + record["temporal_metadata"][person_idx][side], + side, + bbox, + track, + medium_occlusion_frames, + ) + returning_after_long_occlusion = ( + measurement is not None + and track.missed_frames > short_occlusion_frames + ) + wrist, hand, state = update_endpoint_track( + track, + measurement, + frame_idx, + short_occlusion_frames, + medium_occlusion_frames, + ) + shoulder = base_keypoints[person_idx, shoulder_idx] + upper_ratio, forearm_ratio = bone_ratios[(person_idx, side)] + upper_length, forearm_length = calibrated_bone_lengths( + upper_ratio, forearm_ratio, bbox + ) + reachable = None + + if wrist is not None and hand is not None: + reference_elbow = ( + unfiltered_keypoints[person_idx, elbow_idx] + if returning_after_long_occlusion or track.last_elbow is None + else track.last_elbow + ) + use_ik = ( + record["temporal_metadata"][person_idx][side]["limb_state"] + != "consensus" + or state in ( + "provisional", + "recovered", + "reassigned", + "interpolated", + "predicted_short", + ) + ) + if use_ik: + elbow = solve_two_bone_elbow( + shoulder, + wrist, + upper_length, + forearm_length, + reference_elbow, + ) + else: + elbow = base_keypoints[person_idx, elbow_idx] + track.last_elbow = elbow.copy() + direction = wrist - shoulder + if float(np.linalg.norm(direction)) > 1e-6: + track.last_direction = direction / np.linalg.norm(direction) + + v2_keypoints[person_idx, shoulder_idx] = shoulder + v2_keypoints[person_idx, elbow_idx] = elbow + v2_keypoints[person_idx, wrist_idx] = wrist + start, end = HAND_SLICES[side] + v2_keypoints[person_idx, start:end] = hand + uncertainty_radius = ( + 0.025 * diagonal + if state == "observed" + else (0.04 + 0.015 * track.missed_frames) * diagonal + ) + precise_distance = state == "observed" + fast_alarm = state in ( + "observed", + "provisional", + "reassigned", + "interpolated", + "predicted_short", + ) + else: + reachable = reachable_region( + shoulder, + upper_length, + forearm_length, + bbox, + track, + state, + short_occlusion_frames, + medium_occlusion_frames, + ) + uncertainty_radius = reachable["outer_radius_px"] + precise_distance = False + fast_alarm = True + + person_meta[side] = { + "state": state, + "precise_distance": precise_distance, + "fast_alarm_candidate": fast_alarm, + "risk_mode": ( + "precise_hand" + if state == "observed" + else ( + "motion_corridor_intersection" + if state == "reachable_medium" + else ( + "door_intersection_only" + if state == "reachable_long" + else ( + "reassigned_hand" + if state == "reassigned" + else "temporal_hand" + ) + ) + ) + ), + "requires_door_intersection": state + in ("reachable_medium", "reachable_long"), + "missed_frames": track.missed_frames, + "uncertainty_radius_px": round(float(uncertainty_radius), 3), + "bone_lengths_px": { + "upper": round(upper_length, 3), + "forearm": round(forearm_length, 3), + }, + "identity_assignment": assignment_events.get( + (person_idx, side) + ), + "reachable_region": reachable, + } + people_meta.append(person_meta) + + output.append( + { + "frame_index": frame_idx, + "boxes": boxes, + "v1_keypoints": record["temporal_keypoints"], + "v1_metadata": record["temporal_metadata"], + "v2_keypoints": np.round(v2_keypoints, 3).tolist(), + "v2_metadata": people_meta, + } + ) + return output + + +def sector_points(region: dict) -> np.ndarray: + center = np.asarray(region["center"], dtype=np.float64) + direction = region["direction_deg"] + half_angle = region["half_angle_deg"] + angles = np.linspace(direction - half_angle, direction + half_angle, 48) + radius = region["outer_radius_px"] + arc = np.column_stack( + [ + center[0] + radius * np.cos(np.deg2rad(angles)), + center[1] + radius * np.sin(np.deg2rad(angles)), + ] + ) + if half_angle < 179.0: + return np.vstack([center, arc, center]).astype(np.int32) + return arc.astype(np.int32) + + +def draw_v2( + image: np.ndarray, + boxes: list[list[float]], + keypoints: np.ndarray, + metadata: list[dict], + title: str, + show_reachable: bool = True, +) -> np.ndarray: + canvas = image.copy() + overlay = canvas.copy() + for person_idx, (bbox, kpts, person_meta) in enumerate( + zip(boxes, keypoints, metadata), start=1 + ): + x1, y1, x2, y2 = (int(round(value)) for value in bbox) + cv2.rectangle(canvas, (x1, y1), (x2, y2), (0, 220, 255), 2) + cv2.putText( + canvas, + f"P{person_idx}", + (x1 + 4, max(22, y1 + 24)), + cv2.FONT_HERSHEY_SIMPLEX, + 0.72, + (0, 220, 255), + 2, + cv2.LINE_AA, + ) + for side, (shoulder, elbow, wrist) in BODY_SIDE_INDICES.items(): + meta = person_meta[side] + state = meta["state"] + color = STATE_COLORS[state] + if state in ("reachable_medium", "reachable_long"): + region = meta["reachable_region"] + if state == "reachable_medium": + if show_reachable: + polygon = np.round(region["polygon"]).astype(np.int32) + cv2.fillPoly(overlay, [polygon], color) + cv2.polylines(canvas, [polygon], True, color, 3, cv2.LINE_AA) + center = tuple(np.round(region["center"]).astype(int)) + wrist_point = tuple( + np.round(region["predicted_wrist"]).astype(int) + ) + cv2.line(canvas, center, wrist_point, color, 2, cv2.LINE_AA) + label = f"{side[0].upper()}:corridor" + else: + center = tuple(np.round(region["center"]).astype(int)) + cv2.circle(canvas, center, 8, color, 2, cv2.LINE_AA) + if region["half_angle_deg"] >= 179.0: + cv2.putText( + canvas, + "?", + (center[0] + 11, center[1] + 9), + cv2.FONT_HERSHEY_SIMPLEX, + 0.65, + color, + 2, + cv2.LINE_AA, + ) + else: + direction = math.radians(region["direction_deg"]) + direction_end = ( + int(round(center[0] + 55 * math.cos(direction))), + int(round(center[1] + 55 * math.sin(direction))), + ) + cv2.arrowedLine( + canvas, + center, + direction_end, + color, + 2, + cv2.LINE_AA, + tipLength=0.25, + ) + if show_reachable and region.get("render_full_region", False): + polygon = sector_points(region) + if region["half_angle_deg"] < 179.0: + cv2.fillPoly(overlay, [polygon], color) + cv2.polylines( + canvas, [polygon], True, color, 3, cv2.LINE_AA + ) + else: + radius = int(round(region["outer_radius_px"])) + cv2.circle(canvas, center, radius, color, 3, cv2.LINE_AA) + label = f"{side[0].upper()}:unknown" + else: + thickness = 3 + if state in ("reassigned", "recovered"): + draw_dashed_line( + canvas, + kpts[shoulder], + kpts[elbow], + color, + thickness, + ) + draw_dashed_line( + canvas, + kpts[elbow], + kpts[wrist], + color, + thickness, + ) + else: + cv2.line( + canvas, + tuple(np.round(kpts[shoulder]).astype(int)), + tuple(np.round(kpts[elbow]).astype(int)), + color, + thickness, + cv2.LINE_AA, + ) + cv2.line( + canvas, + tuple(np.round(kpts[elbow]).astype(int)), + tuple(np.round(kpts[wrist]).astype(int)), + color, + thickness, + cv2.LINE_AA, + ) + start, end = HAND_SLICES[side] + hand = kpts[start:end] + for a, b in HAND_EDGES_LOCAL: + cv2.line( + canvas, + tuple(np.round(hand[a]).astype(int)), + tuple(np.round(hand[b]).astype(int)), + color, + 2, + cv2.LINE_AA, + ) + label = f"{side[0].upper()}:{state[:5]}" + cv2.putText( + canvas, + label, + (x1 + 4, max(45, y1 + (48 if side == "left" else 70))), + cv2.FONT_HERSHEY_SIMPLEX, + 0.48, + color, + 1, + cv2.LINE_AA, + ) + canvas = cv2.addWeighted(overlay, 0.06, canvas, 0.94, 0.0) + cv2.rectangle(canvas, (0, 0), (670, 45), (0, 0, 0), -1) + cv2.putText( + canvas, + title, + (10, 31), + cv2.FONT_HERSHEY_SIMPLEX, + 0.8, + (255, 255, 255), + 2, + cv2.LINE_AA, + ) + return canvas + + +def endpoint_metrics(records: list[dict]) -> dict: + states: Counter = Counter() + assignments: Counter = Counter() + steps: dict[str, list[float]] = defaultdict(list) + previous: dict[tuple[int, str], tuple[int, np.ndarray]] = {} + for record in records: + for person_idx, bbox in enumerate(record["boxes"]): + kpts = np.asarray(record["v2_keypoints"][person_idx], dtype=np.float64) + for side, (_, _, wrist) in BODY_SIDE_INDICES.items(): + meta = record["v2_metadata"][person_idx][side] + key_name = f"P{person_idx + 1}_{side}" + states[f"{key_name}_{meta['state']}"] += 1 + identity_assignment = meta.get("identity_assignment") + if identity_assignment is not None: + assignments[identity_assignment["method"]] += 1 + if meta["state"] in ("reachable_medium", "reachable_long"): + previous.pop((person_idx, side), None) + continue + normalized = normalized_keypoints(kpts, bbox)[wrist] + prior = previous.get((person_idx, side)) + if prior is not None and prior[0] == record["frame_index"] - 1: + steps[key_name].append(float(np.linalg.norm(normalized - prior[1]))) + previous[(person_idx, side)] = (record["frame_index"], normalized) + return { + "state_counts": dict(sorted(states.items())), + "identity_assignment_counts": dict(sorted(assignments.items())), + "wrist_frame_step": {name: summarize(items) for name, items in steps.items()}, + } + + +def write_video_and_montage( + video_path: Path, + records: list[dict], + output_dir: Path, +) -> None: + cap = cv2.VideoCapture(str(video_path)) + fps = float(cap.get(cv2.CAP_PROP_FPS)) + writer = cv2.VideoWriter( + str(output_dir / "rtmw_l_temporal_v1_vs_occlusion_v4.mp4"), + cv2.VideoWriter_fourcc(*"mp4v"), + fps, + (1920, 540), + ) + key_frames = {0, 50, 120, 126, 132, 138, 180, 211, 220, 240, 252, + 292, 296, 300, 304, 320, 360, 411} + montage_rows = [] + frame_idx = 0 + while True: + ok, frame = cap.read() + if not ok: + break + record = records[frame_idx] + v1 = draw_fused( + frame, + record["boxes"], + [np.asarray(points) for points in record["v1_keypoints"]], + record["v1_metadata"], + f"temporal fusion V1 | frame {frame_idx}", + ) + v2 = draw_v2( + frame, + record["boxes"], + np.asarray(record["v2_keypoints"]), + record["v2_metadata"], + f"occlusion V4 global assignment | frame {frame_idx}", + ) + left = cv2.resize(v1, (960, 540), interpolation=cv2.INTER_AREA) + right = cv2.resize(v2, (960, 540), interpolation=cv2.INTER_AREA) + writer.write(np.hstack([left, right])) + if frame_idx in key_frames: + montage_rows.append( + np.hstack( + [ + cv2.resize(v1, (640, 360), interpolation=cv2.INTER_AREA), + cv2.resize(v2, (640, 360), interpolation=cv2.INTER_AREA), + ] + ) + ) + frame_idx += 1 + cap.release() + writer.release() + cv2.imwrite( + str(output_dir / "rtmw_l_temporal_v1_vs_occlusion_v4_montage.jpg"), + np.vstack(montage_rows), + ) + + +def main() -> None: + args = parse_args() + args.output.mkdir(parents=True, exist_ok=True) + records = [ + json.loads(line) + for line in args.fusion_jsonl.read_text(encoding="utf-8").splitlines() + ] + output = process_sequence( + records, + short_occlusion_frames=args.short_occlusion_frames, + medium_occlusion_frames=args.medium_occlusion_frames, + lookahead_frames=args.lookahead_frames, + ) + with (args.output / "occlusion_v4.jsonl").open("w", encoding="utf-8") as handle: + for record in output: + handle.write(json.dumps(record, ensure_ascii=False) + "\n") + + summary = { + "algorithm_version": "occlusion_v4_global_assignment", + "frames": len(output), + "short_occlusion_frames": args.short_occlusion_frames, + "medium_occlusion_frames": args.medium_occlusion_frames, + "lookahead_frames": args.lookahead_frames, + "policy": { + "observed": "filtered visible endpoint; precise distance allowed", + "provisional": "strong current-frame endpoint; fast alarm only", + "recovered": "weak current-frame endpoint accepted by temporal gating", + "reassigned": "global hand candidate transferred using foreground scale and trajectory", + "interpolated": "fixed-lag short-gap interpolation; fast alarm only", + "predicted_short": "damped-velocity prediction; uncertainty grows each frame", + "reachable_medium": "track-centered motion corridor; intersect with door band", + "reachable_long": "no exact wrist; keep shoulder sector data for door intersection only", + "ik": "two-bone shoulder/wrist reconstruction using robust visible bone lengths", + "ownership": "one-to-one global hand reassignment using foreground scale and trajectory", + }, + "metrics": endpoint_metrics(output), + } + with (args.output / "summary.json").open("w", encoding="utf-8") as handle: + json.dump(summary, handle, ensure_ascii=False, indent=2) + write_video_and_montage(args.video, output, args.output) + print(json.dumps(summary, ensure_ascii=False, indent=2)) + + +if __name__ == "__main__": + main() diff --git a/pipeline/predict_optimized_pose_videos.py b/pipeline/predict_optimized_pose_videos.py new file mode 100644 index 0000000..dd15617 --- /dev/null +++ b/pipeline/predict_optimized_pose_videos.py @@ -0,0 +1,683 @@ +"""End-to-end optimized RTMW-L prediction for CMS videos. + +Pipeline: +YOLO11s person -> fixed-camera false-positive filters -> lightweight track IDs +-> RTMW-L raw/smaller-person-isolated views -> temporal fusion -> causal +occlusion-aware V4 global hand assignment. Reachable regions stay in JSON and are hidden in the main +comparison video. +""" + +from __future__ import annotations + +import argparse +import json +import math +import time +from collections import Counter, defaultdict, deque +from dataclasses import dataclass, field +from pathlib import Path + +import cv2 +import numpy as np + +from compare_rtmw_ml import HAND_SLICES, draw_pose, summarize +from fuse_rtmw_l_temporal import ( + BODY_SIDE_INDICES, + PersonTrack, + adaptive_smooth, + build_global_hands, + draw_fused, + fuse_person, +) +from occlusion_aware_pose_v2 import ( + EndpointTrack, + Measurement, + calibrated_bone_lengths, + draw_v2, + reachable_region, + resolve_global_hand_assignments, + solve_two_bone_elbow, + temporal_recovery_measurement, + update_endpoint_track, +) +from rtmw_l_person_mask_ablation import mask_other_people + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser() + parser.add_argument("--videos", type=Path, nargs="+", required=True) + parser.add_argument("--det-model", type=Path, required=True) + parser.add_argument("--pose-model", type=Path, required=True) + parser.add_argument("--output", type=Path, required=True) + parser.add_argument( + "--det-device", + default="0", + help="Ultralytics YOLO device, for example 0, cuda:0, or cpu.", + ) + parser.add_argument( + "--pose-device", + default="cpu", + help="RTMW ONNX Runtime device: cpu or cuda.", + ) + parser.add_argument("--det-imgsz", type=int, default=960) + parser.add_argument("--det-conf", type=float, default=0.30) + parser.add_argument("--det-roi-y-max-ratio", type=float, default=1.0) + parser.add_argument("--short-occlusion-frames", type=int, default=12) + parser.add_argument("--medium-occlusion-frames", type=int, default=30) + parser.add_argument("--max-people", type=int, default=3) + return parser.parse_args() + + +def bbox_iou(a: np.ndarray, b: np.ndarray) -> float: + x1 = max(float(a[0]), float(b[0])) + y1 = max(float(a[1]), float(b[1])) + x2 = min(float(a[2]), float(b[2])) + y2 = min(float(a[3]), float(b[3])) + intersection = max(x2 - x1, 0.0) * max(y2 - y1, 0.0) + union = ( + max(float(a[2] - a[0]), 0.0) * max(float(a[3] - a[1]), 0.0) + + max(float(b[2] - b[0]), 0.0) * max(float(b[3] - b[1]), 0.0) + - intersection + ) + return intersection / max(union, 1.0) + + +class BoxTracker: + def __init__(self, max_age: int = 30): + self.max_age = max_age + self.next_id = 1 + self.tracks: dict[int, dict] = {} + + def update( + self, + boxes: list[list[float]], + frame_idx: int, + image_shape: tuple[int, int], + ) -> list[int]: + height, width = image_shape + image_diagonal = math.hypot(width, height) + active_ids = [ + track_id + for track_id, track in self.tracks.items() + if frame_idx - track["last_frame"] <= self.max_age + ] + self.tracks = {track_id: self.tracks[track_id] for track_id in active_ids} + pairs = [] + for detection_idx, box in enumerate(boxes): + box_array = np.asarray(box, dtype=np.float64) + center = (box_array[:2] + box_array[2:]) * 0.5 + for track_id in active_ids: + prior = np.asarray(self.tracks[track_id]["box"], dtype=np.float64) + prior_center = (prior[:2] + prior[2:]) * 0.5 + iou = bbox_iou(box_array, prior) + center_distance = float(np.linalg.norm(center - prior_center)) / image_diagonal + if iou >= 0.05 or center_distance <= 0.18: + score = 2.0 * iou + max(0.0, 1.0 - center_distance / 0.18) + pairs.append((score, detection_idx, track_id)) + + assigned_detections = set() + assigned_tracks = set() + result: list[int | None] = [None] * len(boxes) + for _, detection_idx, track_id in sorted(pairs, reverse=True): + if detection_idx in assigned_detections or track_id in assigned_tracks: + continue + assigned_detections.add(detection_idx) + assigned_tracks.add(track_id) + result[detection_idx] = track_id + + for detection_idx, box in enumerate(boxes): + track_id = result[detection_idx] + if track_id is None: + track_id = self.next_id + self.next_id += 1 + result[detection_idx] = track_id + self.tracks[track_id] = {"box": box, "last_frame": frame_idx} + return [int(track_id) for track_id in result] + + +@dataclass +class LivePersonState: + fusion: PersonTrack = field(default_factory=PersonTrack) + endpoints: dict[str, EndpointTrack] = field( + default_factory=lambda: {"left": EndpointTrack(), "right": EndpointTrack()} + ) + bone_samples: dict[str, deque] = field( + default_factory=lambda: { + "left": deque(maxlen=120), + "right": deque(maxlen=120), + } + ) + last_frame: int = -1 + + +def bbox_diagonal(bbox: list[float]) -> float: + return max(math.hypot(bbox[2] - bbox[0], bbox[3] - bbox[1]), 1.0) + + +def update_bone_samples( + state: LivePersonState, + keypoints: np.ndarray, + metadata: dict, + bbox: list[float], +) -> None: + diagonal = bbox_diagonal(bbox) + for side, (shoulder, elbow, wrist) in BODY_SIDE_INDICES.items(): + meta = metadata[side] + if meta["limb_state"] != "consensus" or not meta["distance_eligible"]: + continue + upper = float(np.linalg.norm(keypoints[shoulder] - keypoints[elbow])) / diagonal + forearm = float(np.linalg.norm(keypoints[elbow] - keypoints[wrist])) / diagonal + if 0.04 <= upper <= 0.65 and 0.04 <= forearm <= 0.65: + state.bone_samples[side].append((upper, forearm)) + + +def current_bone_ratios(state: LivePersonState, side: str) -> tuple[float, float]: + samples = state.bone_samples[side] + if len(samples) < 10: + opposite = "right" if side == "left" else "left" + if len(state.bone_samples[opposite]) >= 10: + samples = state.bone_samples[opposite] + if not samples: + return 0.24, 0.22 + values = np.asarray(samples, dtype=np.float64) + return float(np.median(values[:, 0])), float(np.median(values[:, 1])) + + +def update_live_v2( + frame_idx: int, + bbox: list[float], + selected_keypoints: np.ndarray, + smoothed_keypoints: np.ndarray, + fusion_metadata: dict, + state: LivePersonState, + assigned_measurements: dict[str, Measurement | None] | None, + identity_assignments: dict[str, dict] | None, + short_occlusion_frames: int, + medium_occlusion_frames: int, +) -> tuple[np.ndarray, dict]: + frame_gap = 0 if state.last_frame < 0 else max(frame_idx - state.last_frame - 1, 0) + for endpoint in state.endpoints.values(): + endpoint.missed_frames += frame_gap + state.last_frame = frame_idx + update_bone_samples(state, selected_keypoints, fusion_metadata, bbox) + + output = smoothed_keypoints.copy() + output_metadata = {} + diagonal = bbox_diagonal(bbox) + for side, (shoulder_idx, elbow_idx, wrist_idx) in BODY_SIDE_INDICES.items(): + endpoint = state.endpoints[side] + meta = fusion_metadata[side] + start, end = HAND_SLICES[side] + measurement = ( + assigned_measurements.get(side) + if assigned_measurements is not None + else None + ) + if assigned_measurements is None and meta["alarm_eligible"]: + measurement = Measurement( + keypoints=selected_keypoints[start:end], + precise=bool(meta["distance_eligible"]), + provisional=not bool(meta["distance_eligible"]), + quality=float(meta.get("selected_hand_score", 0.0)), + ) + if measurement is None: + measurement = temporal_recovery_measurement( + selected_keypoints, + meta, + side, + bbox, + endpoint, + medium_occlusion_frames, + ) + returning_after_long = ( + measurement is not None + and endpoint.missed_frames > short_occlusion_frames + ) + wrist, hand, endpoint_state = update_endpoint_track( + endpoint, + measurement, + frame_idx, + short_occlusion_frames, + medium_occlusion_frames, + ) + shoulder = smoothed_keypoints[shoulder_idx] + upper_ratio, forearm_ratio = current_bone_ratios(state, side) + upper_length, forearm_length = calibrated_bone_lengths( + upper_ratio, forearm_ratio, bbox + ) + region = None + if wrist is not None and hand is not None: + reference_elbow = ( + selected_keypoints[elbow_idx] + if returning_after_long or endpoint.last_elbow is None + else endpoint.last_elbow + ) + use_ik = ( + meta["limb_state"] != "consensus" + or endpoint_state + in ("provisional", "recovered", "reassigned", "predicted_short") + ) + elbow = ( + solve_two_bone_elbow( + shoulder, + wrist, + upper_length, + forearm_length, + reference_elbow, + ) + if use_ik + else smoothed_keypoints[elbow_idx] + ) + endpoint.last_elbow = elbow.copy() + direction = wrist - shoulder + if float(np.linalg.norm(direction)) > 1e-6: + endpoint.last_direction = direction / np.linalg.norm(direction) + output[shoulder_idx] = shoulder + output[elbow_idx] = elbow + output[wrist_idx] = wrist + output[start:end] = hand + uncertainty = ( + 0.025 * diagonal + if endpoint_state == "observed" + else (0.04 + 0.015 * endpoint.missed_frames) * diagonal + ) + precise = endpoint_state == "observed" + else: + region = reachable_region( + shoulder, + upper_length, + forearm_length, + bbox, + endpoint, + endpoint_state, + short_occlusion_frames, + medium_occlusion_frames, + ) + uncertainty = region["outer_radius_px"] + precise = False + output_metadata[side] = { + "state": endpoint_state, + "precise_distance": precise, + "fast_alarm_candidate": True, + "risk_mode": ( + "precise_hand" + if endpoint_state == "observed" + else ( + "motion_corridor_intersection" + if endpoint_state == "reachable_medium" + else ( + "door_intersection_only" + if endpoint_state == "reachable_long" + else ( + "reassigned_hand" + if endpoint_state == "reassigned" + else "temporal_hand" + ) + ) + ) + ), + "requires_door_intersection": endpoint_state + in ("reachable_medium", "reachable_long"), + "identity_assignment": ( + identity_assignments.get(side) + if identity_assignments is not None + else None + ), + "missed_frames": endpoint.missed_frames, + "uncertainty_radius_px": round(float(uncertainty), 3), + "bone_lengths_px": { + "upper": round(upper_length, 3), + "forearm": round(forearm_length, 3), + }, + "reachable_region": region, + } + return output, output_metadata + + +def filter_person_boxes( + result, + width: int, + height: int, + max_people: int, +) -> tuple[list[list[float]], list[float], Counter]: + boxes = result.boxes.xyxy.cpu().numpy() + scores = result.boxes.conf.cpu().numpy() + rejected = Counter() + kept = [] + for box, score in zip(boxes, scores): + center_x = float(box[0] + box[2]) * 0.5 + box_height = float(box[3] - box[1]) + if center_x > 0.80 * width: + rejected["right_reflection_roi"] += 1 + continue + if box_height < 0.20 * height: + rejected["too_short"] += 1 + continue + kept.append((box.tolist(), float(score))) + kept.sort(key=lambda item: item[1], reverse=True) + if len(kept) > max_people: + rejected["over_max_people"] += len(kept) - max_people + kept = kept[:max_people] + kept.sort(key=lambda item: item[0][0]) + return [item[0] for item in kept], [item[1] for item in kept], rejected + + +def process_video( + video_path: Path, + output_dir: Path, + detector: YOLO, + pose: RTMPose, + det_device: str, + det_imgsz: int, + det_conf: float, + det_roi_y_max_ratio: float, + short_occlusion_frames: int, + medium_occlusion_frames: int, + max_people: int, + video_label: str, +) -> dict: + output_dir.mkdir(parents=True, exist_ok=True) + cap = cv2.VideoCapture(str(video_path)) + fps = float(cap.get(cv2.CAP_PROP_FPS)) + width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)) + height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)) + writer = cv2.VideoWriter( + str(output_dir / "raw_temporal_occlusion_v4.mp4"), + cv2.VideoWriter_fourcc(*"mp4v"), + fps, + (1920, 360), + ) + + box_tracker = BoxTracker(max_age=30) + people: dict[int, LivePersonState] = {} + records = [] + detection_counts = Counter() + rejection_counts = Counter() + state_counts = Counter() + timings = defaultdict(list) + isolated_inferences = 0 + frame_idx = 0 + + while True: + ok, frame = cap.read() + if not ok: + break + started = time.perf_counter() + detection_height = max(1, min(height, round(height * det_roi_y_max_ratio))) + detection = detector.predict( + frame[:detection_height], + imgsz=det_imgsz, + conf=det_conf, + iou=0.30, + classes=[0], + device=det_device, + verbose=False, + )[0] + timings["detector_ms"].append((time.perf_counter() - started) * 1000.0) + boxes, confidences, rejected = filter_person_boxes( + detection, width, height, max_people + ) + rejection_counts.update(rejected) + detection_counts[len(boxes)] += 1 + track_ids = box_tracker.update(boxes, frame_idx, (height, width)) + + if boxes: + started = time.perf_counter() + raw_keypoints, raw_scores = pose(frame, bboxes=boxes) + timings["raw_pose_ms"].append((time.perf_counter() - started) * 1000.0) + isolated_keypoints = raw_keypoints.copy() + isolated_scores = raw_scores.copy() + started = time.perf_counter() + for person_idx, bbox in enumerate(boxes): + target_area = max((bbox[2] - bbox[0]) * (bbox[3] - bbox[1]), 1.0) + needs_isolation = any( + idx != person_idx + and max((other[2] - other[0]) * (other[3] - other[1]), 1.0) + < 0.5 * target_area + for idx, other in enumerate(boxes) + ) + if not needs_isolation: + continue + masked = mask_other_people(frame, boxes, person_idx) + kpts, scores = pose(masked, bboxes=[bbox]) + isolated_keypoints[person_idx] = kpts[0] + isolated_scores[person_idx] = scores[0] + isolated_inferences += 1 + timings["isolated_extra_ms"].append((time.perf_counter() - started) * 1000.0) + else: + raw_keypoints = np.empty((0, 133, 2), dtype=np.float64) + raw_scores = np.empty((0, 133), dtype=np.float64) + isolated_keypoints = raw_keypoints.copy() + isolated_scores = raw_scores.copy() + + selected_list = [] + smoothed_list = [] + fusion_meta_list = [] + person_states = [] + for person_idx, (bbox, track_id) in enumerate(zip(boxes, track_ids)): + state = people.setdefault(track_id, LivePersonState()) + selected, fusion_meta = fuse_person( + raw_keypoints[person_idx], + raw_scores[person_idx], + isolated_keypoints[person_idx], + isolated_scores[person_idx], + bbox, + state.fusion, + ) + smoothed = adaptive_smooth(selected, bbox, state.fusion) + selected_list.append(selected) + smoothed_list.append(smoothed) + fusion_meta_list.append(fusion_meta) + person_states.append(state) + + frame_measurements = {} + endpoint_tracks = {} + for person_idx, (selected, fusion_meta, state) in enumerate( + zip(selected_list, fusion_meta_list, person_states) + ): + for side, (start, end) in HAND_SLICES.items(): + meta = fusion_meta[side] + frame_measurements[(person_idx, side)] = ( + Measurement( + keypoints=selected[start:end], + precise=bool(meta["distance_eligible"]), + provisional=not bool(meta["distance_eligible"]), + quality=float(meta.get("selected_hand_score", 0.0)), + ) + if meta["alarm_eligible"] + else None + ) + endpoint_tracks[(person_idx, side)] = state.endpoints[side] + frame_measurements, assignment_events = resolve_global_hand_assignments( + frame_measurements, + endpoint_tracks, + boxes, + medium_occlusion_frames, + ) + + v2_list = [] + v2_meta_list = [] + for person_idx, (bbox, selected, smoothed, fusion_meta, state) in enumerate( + zip( + boxes, + selected_list, + smoothed_list, + fusion_meta_list, + person_states, + ) + ): + v2, v2_meta = update_live_v2( + frame_idx, + bbox, + selected, + smoothed, + fusion_meta, + state, + { + side: frame_measurements[(person_idx, side)] + for side in ("left", "right") + }, + { + side: assignment_events.get((person_idx, side)) + for side in ("left", "right") + }, + short_occlusion_frames, + medium_occlusion_frames, + ) + v2_list.append(v2) + v2_meta_list.append(v2_meta) + for side in ("left", "right"): + state_counts[f"{v2_meta[side]['state']}"] += 1 + + global_hands = build_global_hands(selected_list, fusion_meta_list) + records.append( + { + "frame_index": frame_idx, + "timestamp_ms": round(frame_idx / fps * 1000.0, 3), + "track_ids": track_ids, + "boxes": np.round(np.asarray(boxes), 3).tolist() if boxes else [], + "det_confidences": [round(value, 4) for value in confidences], + "raw_keypoints": np.round(raw_keypoints, 3).tolist(), + "temporal_selected_keypoints": np.round(np.asarray(selected_list), 3).tolist(), + "temporal_keypoints": np.round(np.asarray(smoothed_list), 3).tolist(), + "temporal_metadata": fusion_meta_list, + "v2_keypoints": np.round(np.asarray(v2_list), 3).tolist(), + "v2_metadata": v2_meta_list, + "global_hands": global_hands, + } + ) + + raw_panel = draw_pose( + frame, + boxes, + raw_keypoints, + f"{video_label} RTMW-L raw | frame {frame_idx}", + ) + temporal_panel = draw_fused( + frame, + boxes, + smoothed_list, + fusion_meta_list, + f"{video_label} temporal fusion | frame {frame_idx}", + ) + v2_panel = draw_v2( + frame, + boxes, + np.asarray(v2_list), + v2_meta_list, + f"{video_label} occlusion V4 | frame {frame_idx}", + show_reachable=False, + ) + panels = [ + cv2.resize(panel, (640, 360), interpolation=cv2.INTER_AREA) + for panel in (raw_panel, temporal_panel, v2_panel) + ] + writer.write(np.hstack(panels)) + if frame_idx == 0 and video_path.suffix.lower() in { + ".jpg", + ".jpeg", + ".png", + ".bmp", + ".webp", + }: + cv2.imwrite(str(output_dir / "raw_pose.jpg"), raw_panel) + cv2.imwrite(str(output_dir / "spatial_fusion.jpg"), temporal_panel) + cv2.imwrite(str(output_dir / "occlusion_v2.jpg"), v2_panel) + panel_width = 960 + panel_height = round(height * panel_width / width) + image_panels = [ + cv2.resize( + panel, + (panel_width, panel_height), + interpolation=cv2.INTER_AREA, + ) + for panel in (raw_panel, temporal_panel, v2_panel) + ] + cv2.imwrite( + str(output_dir / "raw_spatial_occlusion_v2.jpg"), + np.hstack(image_panels), + ) + if frame_idx % 100 == 0: + print(f"{video_label}: processed frame {frame_idx}", flush=True) + frame_idx += 1 + + cap.release() + writer.release() + with (output_dir / "predictions.jsonl").open("w", encoding="utf-8") as handle: + for record in records: + handle.write(json.dumps(record, ensure_ascii=False) + "\n") + + summary = { + "occlusion_algorithm": "v4_global_assignment", + "video": str(video_path.resolve()), + "frames": len(records), + "fps": round(fps, 3), + "resolution": [width, height], + "person_filter": { + "center_x_max_ratio": 0.80, + "minimum_box_height_ratio": 0.20, + "max_people": max_people, + }, + "detector_settings": { + "imgsz": det_imgsz, + "confidence": det_conf, + "roi_y_max_ratio": det_roi_y_max_ratio, + }, + "filtered_person_count_frames": dict(sorted(detection_counts.items())), + "rejected_detection_counts": dict(sorted(rejection_counts.items())), + "track_ids_created": box_tracker.next_id - 1, + "isolated_pose_inferences": isolated_inferences, + "short_occlusion_frames": short_occlusion_frames, + "medium_occlusion_frames": medium_occlusion_frames, + "v2_state_counts": dict(sorted(state_counts.items())), + "timing_ms_cpu": { + name: summarize(values[5:] if len(values) > 5 else values) + for name, values in timings.items() + }, + "output_note": ( + "Main video hides reachable-region arcs. Reachable-region geometry remains in " + "predictions.jsonl for later door-distance lower-bound calculation." + ), + } + with (output_dir / "summary.json").open("w", encoding="utf-8") as handle: + json.dump(summary, handle, ensure_ascii=False, indent=2) + return summary + + +def main() -> None: + from rtmlib import RTMPose + from ultralytics import YOLO + args = parse_args() + detector = YOLO(str(args.det_model)) + pose = RTMPose( + str(args.pose_model), + model_input_size=(288, 384), + backend="onnxruntime", + device=args.pose_device, + ) + all_summaries = {} + for video_idx, video_path in enumerate(args.videos, start=1): + label = f"video{video_idx}" + output_dir = args.output / label + all_summaries[label] = process_video( + video_path, + output_dir, + detector, + pose, + args.det_device, + args.det_imgsz, + args.det_conf, + args.det_roi_y_max_ratio, + args.short_occlusion_frames, + args.medium_occlusion_frames, + args.max_people, + label, + ) + args.output.mkdir(parents=True, exist_ok=True) + with (args.output / "summary.json").open("w", encoding="utf-8") as handle: + json.dump(all_summaries, handle, ensure_ascii=False, indent=2) + print(json.dumps(all_summaries, ensure_ascii=False, indent=2)) + + +if __name__ == "__main__": + main() diff --git a/pipeline/rtmw_l_person_mask_ablation.py b/pipeline/rtmw_l_person_mask_ablation.py new file mode 100644 index 0000000..6263500 --- /dev/null +++ b/pipeline/rtmw_l_person_mask_ablation.py @@ -0,0 +1,166 @@ +"""Test whether other people inside a top-down crop cause RTMW-L mix-ups.""" + +from __future__ import annotations + +import argparse +import json +import time +from pathlib import Path + +import cv2 +import numpy as np + +from compare_rtmw_ml import ( + draw_pose, + model_disagreement, + summarize, + temporal_metrics, + write_montage, +) + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser() + parser.add_argument("--video", type=Path, required=True) + parser.add_argument("--baseline-jsonl", type=Path, required=True) + parser.add_argument("--model-l", type=Path, required=True) + parser.add_argument("--output", type=Path, required=True) + return parser.parse_args() + + +def mask_other_people( + frame: np.ndarray, + boxes: list[list[float]], + target_idx: int, +) -> np.ndarray: + masked = frame.copy() + height, width = frame.shape[:2] + mask = np.zeros((height, width), dtype=bool) + target = boxes[target_idx] + target_area = max((target[2] - target[0]) * (target[3] - target[1]), 1.0) + for idx, bbox in enumerate(boxes): + if idx == target_idx: + continue + other_area = max((bbox[2] - bbox[0]) * (bbox[3] - bbox[1]), 1.0) + # The failure in this video is the small background person being + # absorbed by the much larger foreground person's 3:4 pose crop. + # Mask only markedly smaller people. Masking a larger occluder while + # estimating a small person removes genuine target pixels at overlap. + if other_area >= 0.5 * target_area: + continue + x1, y1, x2, y2 = (int(round(value)) for value in bbox) + x1, x2 = max(0, x1), min(width - 1, x2) + y1, y2 = max(0, y1), min(height - 1, y2) + mask[y1:y2 + 1, x1:x2 + 1] = True + + masked[mask] = (124, 116, 104) + return masked + + +def main() -> None: + from rtmlib import RTMPose + args = parse_args() + args.output.mkdir(parents=True, exist_ok=True) + baseline = [ + json.loads(line) + for line in args.baseline_jsonl.read_text(encoding="utf-8").splitlines() + ] + pose_l = RTMPose( + str(args.model_l), model_input_size=(288, 384), backend="onnxruntime", device="cpu" + ) + + cap = cv2.VideoCapture(str(args.video)) + fps = float(cap.get(cv2.CAP_PROP_FPS)) + width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)) + height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)) + writer = cv2.VideoWriter( + str(args.output / "rtmw_l_raw_vs_person_masked.mp4"), + cv2.VideoWriter_fourcc(*"mp4v"), + fps, + (width, height // 2), + ) + + comparison: list[dict] = [] + timings: list[float] = [] + frame_idx = 0 + while True: + ok, frame = cap.read() + if not ok: + break + baseline_record = baseline[frame_idx] + boxes = baseline_record["boxes"][:2] + masked_keypoints = [] + masked_scores = [] + started = time.perf_counter() + for person_idx, bbox in enumerate(boxes): + masked_frame = mask_other_people(frame, boxes, person_idx) + keypoints, scores = pose_l(masked_frame, bboxes=[bbox]) + masked_keypoints.append(keypoints[0]) + masked_scores.append(scores[0]) + timings.append((time.perf_counter() - started) * 1000.0) + masked_keypoints_array = np.asarray(masked_keypoints) + + raw_keypoints = np.asarray(baseline_record["L"]["keypoints"][:2]) + comparison.append( + { + "frame_index": frame_idx, + "boxes": boxes, + "M": {"keypoints": raw_keypoints.tolist()}, + "L": {"keypoints": np.round(masked_keypoints_array, 3).tolist()}, + "masked_raw_scores": np.round(np.asarray(masked_scores), 4).tolist(), + } + ) + + left = draw_pose(frame, boxes, raw_keypoints, f"RTMW-L raw | frame {frame_idx}") + right = draw_pose( + frame, + boxes, + masked_keypoints_array, + f"RTMW-L mask other person | frame {frame_idx}", + ) + left = cv2.resize(left, (width // 2, height // 2), interpolation=cv2.INTER_AREA) + right = cv2.resize(right, (width // 2, height // 2), interpolation=cv2.INTER_AREA) + writer.write(np.hstack([left, right])) + if frame_idx % 50 == 0: + print(f"processed frame {frame_idx}", flush=True) + frame_idx += 1 + + cap.release() + writer.release() + + with (args.output / "rtmw_l_raw_vs_person_masked.jsonl").open( + "w", encoding="utf-8" + ) as handle: + for record in comparison: + handle.write(json.dumps(record, ensure_ascii=False) + "\n") + + disagreement, worst_frames = model_disagreement(comparison) + fixed_frames = [0, 60, 120, 132, 136, 180, 240, 252, 300, 360, 411] + montage_frames = sorted(set(fixed_frames + worst_frames[:8])) + write_montage( + args.video, + comparison, + montage_frames, + args.output / "rtmw_l_raw_vs_person_masked_montage.jpg", + titles=("RTMW-L raw", "RTMW-L smaller-person masked"), + ) + + summary = { + "frames": len(comparison), + "labels": { + "M": "RTMW-L raw crop", + "L": "RTMW-L with markedly smaller people masked", + }, + "timing_ms_two_people_cpu": summarize(timings[5:]), + "temporal_metrics_normalized_to_yolo_box": temporal_metrics(comparison), + "position_change_after_masking": disagreement, + "largest_change_frames": worst_frames, + "montage_frames": montage_frames, + } + with (args.output / "summary.json").open("w", encoding="utf-8") as handle: + json.dump(summary, handle, ensure_ascii=False, indent=2) + print(json.dumps(summary, ensure_ascii=False, indent=2)) + + +if __name__ == "__main__": + main() diff --git a/pipeline/rtmw_risk_pipeline.py b/pipeline/rtmw_risk_pipeline.py new file mode 100644 index 0000000..a61a490 --- /dev/null +++ b/pipeline/rtmw_risk_pipeline.py @@ -0,0 +1,390 @@ +"""Lightweight causal door-risk inference using only YOLO11s and RTMW-L.""" + +from __future__ import annotations + +import argparse +import json +import sys +import time +from pathlib import Path + +import cv2 +import numpy as np + + +PACKAGE_DIR = Path(__file__).resolve().parents[1] +PIPELINE_DIR = Path(__file__).resolve().parent +for import_dir in (PIPELINE_DIR, PACKAGE_DIR): + if str(import_dir) not in sys.path: + sys.path.insert(0, str(import_dir)) + +from door_mask_risk import DoorMaskSequence, DoorRiskEvaluator # noqa: E402 +from door_zone_risk import ( # noqa: E402 + DoorZoneAssigner, + DoorZoneConfig, + build_door_risks, +) +from live_pose_engine import LivePoseEngine # noqa: E402 + + +def parse_image_size(value: str) -> int | tuple[int, int]: + parts = value.lower().replace(",", "x").split("x") + if len(parts) == 1: + return int(parts[0]) + if len(parts) == 2: + return int(parts[0]), int(parts[1]) + raise argparse.ArgumentTypeError( + "image size must be one integer or HEIGHTxWIDTH, for example 960 or 544x960" + ) + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--video", type=Path, required=True) + door_source = parser.add_mutually_exclusive_group(required=True) + door_source.add_argument("--door-mask-dir", type=Path) + door_source.add_argument( + "--constant-door-mask", + type=Path, + help="One X-AnyLabeling JSON or raster mask reused for every video frame.", + ) + parser.add_argument("--output", type=Path, required=True) + parser.add_argument( + "--det-model", type=Path, default=PACKAGE_DIR / "models" / "yolo11s.onnx" + ) + parser.add_argument( + "--pose-model", type=Path, default=PACKAGE_DIR / "models" / "rtmw-l-384.onnx" + ) + parser.add_argument( + "--pose-jsonl", + type=Path, + help="Diagnostic replay input; bypasses live YOLO/RTMW model execution.", + ) + parser.add_argument("--det-device", default="0") + parser.add_argument("--pose-device", default="cuda") + parser.add_argument( + "--pose-provider", + choices=("cuda", "tensorrt"), + default="cuda", + help="ONNX Runtime provider for RTMW.", + ) + parser.add_argument( + "--pose-trt-cache", + type=Path, + default=PACKAGE_DIR / "models" / "rtmw_l_trt_cache", + ) + parser.add_argument("--det-imgsz", type=parse_image_size, default=960) + parser.add_argument("--det-conf", type=float, default=0.30) + parser.add_argument("--det-roi-y-max-ratio", type=float, default=1.0) + parser.add_argument("--max-people", type=int, default=3) + parser.add_argument( + "--door-mask-format", + choices=("auto", "xanylabeling", "raster"), + default="auto", + ) + parser.add_argument("--door-label", action="append") + parser.add_argument("--door-frame-offset", type=int, default=0) + parser.add_argument("--require-door-mask", action="store_true") + parser.add_argument("--risk-warning-relative", type=float, default=0.30) + parser.add_argument("--risk-release-relative", type=float, default=0.45) + parser.add_argument( + "--active-door", + choices=("front", "rear", "both", "none"), + default="both", + help="Door(s) whose CAN state is 0. Use front/rear for deployment simulation.", + ) + parser.add_argument( + "--door-zone-config", + type=Path, + default=PACKAGE_DIR / "annotations" / "door_zone_multifeature.json", + help="Pose-primary fixed-camera front/rear calibration JSON.", + ) + parser.add_argument("--max-frames", type=int, default=0) + parser.add_argument("--write-video", action="store_true") + parser.add_argument("--progress-every", type=int, default=100) + return parser.parse_args() + + +def load_pose_records(path: Path) -> dict[int, dict]: + records = {} + with path.open("r", encoding="utf-8") as handle: + for line in handle: + if not line.strip(): + continue + record = json.loads(line) + records[int(record["frame_index"])] = record + return records + + +def percentile(values: list[float], p: float) -> float | None: + return round(float(np.percentile(values, p)), 3) if values else None + + +def timing_stats(values: list[float]) -> dict: + return { + "count": len(values), + "total": round(float(np.sum(values)), 3) if values else 0.0, + "mean": round(float(np.mean(values)), 3) if values else None, + "median": percentile(values, 50), + "p95": percentile(values, 95), + } + + +def visualization_skeletons(record: dict) -> list[dict]: + """Return the exact fused keypoints and metadata consumed by draw_fused.""" + + boxes = record.get("boxes", []) + track_ids = [int(value) for value in record.get("track_ids", [])] + keypoints = record.get("temporal_selected_keypoints", []) + metadata = record.get("temporal_metadata", []) + if not (len(boxes) == len(keypoints) == len(metadata)): + return [] + return [ + { + "track_id": track_ids[position] if position < len(track_ids) else position + 1, + "person_index": position, + "layout": "coco_wholebody_133", + "keypoints": person_keypoints, + "metadata": metadata[position], + } + for position, person_keypoints in enumerate(keypoints) + ] + + +def render_frame( + frame: np.ndarray, + record: dict, + frame_risk: dict, + door_mask: np.ndarray | None, + door_assignments: list[dict], +) -> np.ndarray: + canvas = frame.copy() + boxes = record.get("boxes", []) + keypoints = record.get("temporal_selected_keypoints", []) + metadata = record.get("temporal_metadata", []) + if len(boxes) == len(keypoints) == len(metadata) and boxes: + from fuse_rtmw_l_temporal import draw_fused + + canvas = draw_fused(canvas, boxes, keypoints, metadata, "YOLO + RTMW risk") + else: + for position, box in enumerate(boxes): + x1, y1, x2, y2 = [int(round(float(value))) for value in box[:4]] + cv2.rectangle(canvas, (x1, y1), (x2, y2), (0, 220, 255), 2) + cv2.putText( + canvas, + f"P{position + 1}", + (x1 + 4, max(22, y1 + 24)), + cv2.FONT_HERSHEY_SIMPLEX, + 0.7, + (0, 220, 255), + 2, + cv2.LINE_AA, + ) + zone_colors = { + "front": (0, 220, 255), + "rear": (255, 180, 0), + "transition": (0, 140, 255), + "unassigned": (150, 150, 150), + } + for assignment in door_assignments: + position = int(assignment["person_index"]) + if not 0 <= position < len(boxes): + continue + box = boxes[position] + x1, y1 = int(round(float(box[0]))), int(round(float(box[1]))) + zone = str(assignment["door_zone"]) + cv2.putText( + canvas, + f"T{assignment['track_id']} {zone.upper()}", + (max(0, x1 + 4), max(24, y1 - 8)), + cv2.FONT_HERSHEY_SIMPLEX, + 0.68, + zone_colors.get(zone, (160, 160, 160)), + 2, + cv2.LINE_AA, + ) + if door_mask is not None and np.any(door_mask): + overlay = np.zeros_like(canvas) + overlay[door_mask > 0] = (255, 180, 0) + canvas = cv2.addWeighted(canvas, 1.0, overlay, 0.22, 0) + + score = frame_risk.get("risk_score") + score_text = "N/A" if score is None else f"{float(score):.3f}" + level = str(frame_risk.get("risk_level", "unavailable")) + color = (0, 0, 255) if level == "danger" else (0, 180, 255) + if level in ("clear", "unavailable", "inactive", "unknown"): + color = (40, 210, 80) if level == "clear" else (160, 160, 160) + cv2.rectangle(canvas, (12, 55), (490, 95), (0, 0, 0), -1) + cv2.putText( + canvas, + f"{frame_risk.get('risk_type', 'DOOR_RISK')} score={score_text} level={level}", + (24, 83), + cv2.FONT_HERSHEY_SIMPLEX, + 0.68, + color, + 2, + cv2.LINE_AA, + ) + return canvas + + +def main() -> None: + args = parse_args() + args.output.mkdir(parents=True, exist_ok=True) + if args.pose_provider == "tensorrt" and not args.pose_device.lower().startswith("cuda"): + raise ValueError("--pose-provider tensorrt requires --pose-device cuda") + + capture = cv2.VideoCapture(str(args.video.resolve())) + if not capture.isOpened(): + raise RuntimeError(f"Cannot open video: {args.video}") + input_fps = float(capture.get(cv2.CAP_PROP_FPS)) or 30.0 + width = int(capture.get(cv2.CAP_PROP_FRAME_WIDTH)) + height = int(capture.get(cv2.CAP_PROP_FRAME_HEIGHT)) + + risk_evaluator = DoorRiskEvaluator( + DoorMaskSequence( + args.door_mask_dir, + mask_format=args.door_mask_format, + labels=tuple(args.door_label or ("door", "车门")), + frame_offset=args.door_frame_offset, + require_every_frame=args.require_door_mask, + constant_mask=args.constant_door_mask, + ), + warning_relative=args.risk_warning_relative, + release_relative=args.risk_release_relative, + ) + zone_assigner = DoorZoneAssigner(DoorZoneConfig.load(args.door_zone_config)) + print( + f"Door masks: {risk_evaluator.sequence.source_format} " + + ( + f"(constant: {risk_evaluator.sequence.constant_path})" + if risk_evaluator.sequence.constant_path is not None + else f"({len(risk_evaluator.sequence.paths)} indexed frames)" + ) + ) + + pose_records = load_pose_records(args.pose_jsonl) if args.pose_jsonl else None + load_started = time.perf_counter() + pose_engine = None if pose_records is not None else LivePoseEngine(args) + model_load_seconds = time.perf_counter() - load_started + + writer = None + if args.write_video: + writer = cv2.VideoWriter( + str(args.output / "rtmw_risk_overlay.mp4"), + cv2.VideoWriter_fourcc(*"mp4v"), + input_fps, + (width, height), + ) + + frame_latencies = [] + pose_latencies = [] + risk_latencies = [] + render_latencies = [] + started_all = time.perf_counter() + frame_index = 0 + output_path = args.output / "causal_results.jsonl" + with output_path.open("w", encoding="utf-8") as output_handle: + while True: + if args.max_frames > 0 and frame_index >= args.max_frames: + break + ok, frame = capture.read() + if not ok: + break + frame_started = time.perf_counter() + pose_started = time.perf_counter() + if pose_records is None: + record = pose_engine.infer(frame, frame_index) + else: + record = pose_records.get(frame_index) + if record is None: + raise RuntimeError(f"Pose JSONL has no record for frame {frame_index}") + pose_latencies.append((time.perf_counter() - pose_started) * 1000) + + risk_started = time.perf_counter() + assignments = zone_assigner.update( + frame_index, record, (height, width) + ) + raw_risk, door_mask = risk_evaluator.evaluate( + frame_index, record, (height, width) + ) + door_risks, frame_risk = build_door_risks( + raw_risk, assignments, args.active_door + ) + risk_latencies.append((time.perf_counter() - risk_started) * 1000) + output_record = { + "frame_index": frame_index, + "track_ids": record.get("track_ids", []), + "boxes": record.get("boxes", []), + "skeletons": visualization_skeletons(record), + "active_door": args.active_door, + "door_assignments": assignments, + "door_risks": door_risks, + "risk": frame_risk, + "ungated_risk": raw_risk, + } + output_handle.write(json.dumps(output_record, ensure_ascii=False) + "\n") + + if writer is not None: + render_started = time.perf_counter() + writer.write( + render_frame(frame, record, frame_risk, door_mask, assignments) + ) + render_latencies.append((time.perf_counter() - render_started) * 1000) + + frame_latencies.append((time.perf_counter() - frame_started) * 1000) + if args.progress_every > 0 and frame_index % args.progress_every == 0: + recent = frame_latencies[-min(len(frame_latencies), args.progress_every) :] + print( + f"frame={frame_index} recent_ms={np.mean(recent):.1f} " + f"risk={frame_risk['risk_level']}", + flush=True, + ) + frame_index += 1 + + capture.release() + if writer is not None: + writer.release() + elapsed = time.perf_counter() - started_all + summary = { + "mode": "yolo_rtmw_door_risk_multifeature_v9_10", + "video": str(args.video.resolve()), + "pose_source": "jsonl_diagnostic" if pose_records is not None else "resident_yolo_rtmw", + "models_loaded": [] if pose_records is not None else ["YOLO11s", "RTMW-L"], + "edgetam_loaded": False, + "risk_region_policy": ( + "RTMW alarm-eligible hand hulls; YOLO body box only when a person " + "has no usable RTMW hand region" + ), + "frames": frame_index, + "input_fps": round(input_fps, 3), + "elapsed_seconds": round(elapsed, 3), + "throughput_fps": round(frame_index / elapsed, 3) if elapsed else None, + "model_load_seconds": round(model_load_seconds, 3), + "latency_ms": { + "total": timing_stats(frame_latencies), + "pose": timing_stats(pose_latencies), + "risk": timing_stats(risk_latencies), + "render": timing_stats(render_latencies), + }, + "write_video": args.write_video, + "active_door": args.active_door, + "door_zone": zone_assigner.summary(), + "door_risk": risk_evaluator.summary(), + } + if pose_engine is not None: + summary["pose"] = { + "provider": args.pose_provider, + "isolated_inferences": pose_engine.isolated_inferences, + "timing_ms": { + key: timing_stats(values) for key, values in pose_engine.timings.items() + }, + } + (args.output / "summary.json").write_text( + json.dumps(summary, ensure_ascii=False, indent=2), encoding="utf-8" + ) + print(json.dumps(summary, ensure_ascii=False, indent=2)) + + +if __name__ == "__main__": + main() diff --git a/run_camera.py b/run_camera.py new file mode 100644 index 0000000..576304c --- /dev/null +++ b/run_camera.py @@ -0,0 +1,299 @@ +#!/usr/bin/env python3 +"""Live baseline V4L2/CAN capture + native TensorRT + CMS v9.10 risk.""" +from __future__ import annotations +import argparse +from collections import Counter, deque +from contextlib import ExitStack +import copy +import json +import signal +import sys +import threading +import time +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer +from pathlib import Path +import cv2 +import numpy as np +ROOT=Path(__file__).resolve().parent +sys.path.insert(0,str(ROOT/'pipeline')) +from native_backend import NativeBackend, NativePoseEngine +from latest_result import LatestResult, result_to_dict +from door_mask_risk import DoorMaskSequence, DoorRiskEvaluator +from door_zone_risk import DoorZoneConfig, DoorZoneAssigner, build_door_risks +from rtmw_risk_pipeline import render_frame, visualization_skeletons, timing_stats + + +def select_active_door(mode, snapshot, side='left', timeout_ms=1000): + if mode!='can': return mode, True + p=snapshot['positions'];offset=0 if side=='left' else 1 + if snapshot['age_ms']<0 or snapshot['age_ms']>timeout_ms or min(p[offset],p[offset+2])<0: + return 'none',False + front,rear=p[offset]==0,p[offset+2]==0 + return ('both' if front and rear else 'front' if front else 'rear' if rear else 'none'),True + + +def gate_risks(raw, assignments, active, can_valid): + doors,risk=build_door_risks(raw,assignments,active) + if not can_valid: + for door,item in doors.items(): + item.update(risk_type=door.upper()+'_UNKNOWN',risk_level='unknown',risk_score=None,decision='CAN_UNAVAILABLE_OR_STALE') + risk=dict(risk,risk_type='CAN_UNKNOWN',risk_level='unknown',risk_score=None,decision='CAN_UNAVAILABLE_OR_STALE') + return doors,risk + + +class JsonlWriter: + """Bound the disk usage of continuous camera runs (current + two backups).""" + def __init__(self,path,max_bytes): + self.path=path;self.max_bytes=max_bytes;self.size=0 + self.handle=path.open('w',encoding='utf-8',buffering=1) + def __enter__(self):return self + def __exit__(self,*args):self.close() + def close(self):self.handle.close() + def write(self,text): + size=len(text.encode('utf-8')) + if self.max_bytes and self.size and self.size+size>self.max_bytes: + self.handle.close() + backup=self.path.with_suffix(self.path.suffix+'.1') + if backup.exists():backup.replace(self.path.with_suffix(self.path.suffix+'.2')) + self.path.replace(backup) + self.handle=self.path.open('w',encoding='utf-8',buffering=1);self.size=0 + self.handle.write(text);self.size+=size + + +def render_output(frame, record, risk, mask, assignments, status): + canvas = render_frame(frame, record, risk, mask, assignments) + cv2.putText(canvas,status,(24,frame.shape[0]-30),cv2.FONT_HERSHEY_SIMPLEX,.8,(0,180,255),2) + return canvas + + +class Preview: + def __init__(self,host,port,results): + self.jpeg=None;self.results=results;self.lock=threading.Lock() + self.pending=threading.Condition();self.packet=None;self.stopped=False + outer=self + class Handler(BaseHTTPRequestHandler): + def do_GET(self): + with outer.lock: + if self.path.split('?')[0]=='/frame.jpg': payload=outer.jpeg or b'';mime='image/jpeg' + elif self.path=='/result.json':payload=json.dumps(result_to_dict(outer.results.load()) or {},ensure_ascii=False,allow_nan=False).encode();mime='application/json' + else: + payload=b'''Vision anti-pinch 0923

Vision anti-pinch 0923 / Live camera

''';mime='text/html; charset=utf-8' + self.send_response(200 if payload else 503);self.send_header('Content-Type',mime);self.send_header('Cache-Control','no-store');self.send_header('Content-Length',str(len(payload)));self.end_headers() + try:self.wfile.write(payload) + except (BrokenPipeError,ConnectionResetError):pass + def log_message(self,*args):pass + self.server=ThreadingHTTPServer((host,port),Handler) + self.thread=threading.Thread(target=self.server.serve_forever,daemon=True);self.thread.start() + self.renderer=threading.Thread(target=self._render_loop,name='risk-preview',daemon=True) + self.renderer.start() + def submit(self, frame, record, risk, mask, assignments, status): + # All inputs belong to a finished frame; inference never modifies them + # after submission. Retain only one pending frame, replacing older work. + with self.pending: + if not self.stopped: + self.packet=(frame,record,risk,mask,assignments,status) + self.pending.notify() + + def _render_loop(self): + try: + while True: + with self.pending: + self.pending.wait_for(lambda:self.packet is not None or self.stopped) + if self.stopped:return + packet=self.packet;self.packet=None + self.update(render_output(*packet)) + except Exception as exc: + with self.pending:self.stopped=True;self.packet=None + with self.lock:self.jpeg=None + print('Visualization disabled; inference continues: '+str(exc),file=sys.stderr,flush=True) + + def update(self,image): + w=960;h=round(image.shape[0]*w/image.shape[1]);image=cv2.resize(image,(w,h)) + ok,jpeg=cv2.imencode('.jpg',image,[cv2.IMWRITE_JPEG_QUALITY,80]) + if ok: + with self.lock:self.jpeg=jpeg.tobytes() + def close(self): + with self.pending:self.stopped=True;self.packet=None;self.pending.notify_all() + self.renderer.join() + self.server.shutdown();self.server.server_close();self.thread.join() + + +def parse_args(argv=None): + p=argparse.ArgumentParser(description=__doc__) + p.add_argument('--video',type=Path,help='Diagnostic video replay through the same TensorRT/risk chain');p.add_argument('--camera',default='/dev/video5');p.add_argument('--width',type=int,default=1920);p.add_argument('--height',type=int,default=1536);p.add_argument('--fps',type=int,default=30) + p.add_argument('--engines',type=Path,default=ROOT/'models/engines') + p.add_argument('--constant-door-mask',type=Path,default=ROOT/'annotations/constant_door_mask.json') + p.add_argument('--door-zone-config',type=Path,default=ROOT/'annotations/door_zone_multifeature.json') + p.add_argument('--active-door',choices=['can','front','rear','both','none'],default='can') + p.add_argument('--can-side',choices=['left','right'],default='left');p.add_argument('--can-timeout-ms',type=float,default=1000);p.add_argument('--no-can',action='store_true') + p.add_argument('--max-people',type=int,default=3);p.add_argument('--det-conf',type=float,default=.30) + p.add_argument('--max-frames',type=int,default=0);p.add_argument('--duration',type=float,default=0) + p.add_argument('--jsonl-max-mb',type=float,default=32,help='Rotate each JSONL at this size; keep two backups; 0 disables rotation');p.add_argument('--output',type=Path,default=ROOT/'outputs'/time.strftime('%Y%m%d-%H%M%S')) + p.add_argument('--visualize',action='store_true',help='Enable optional HTTP visualization; disabled by default') + p.add_argument('--save-results',action='store_true',help='Save per-frame JSONL and periodic latest.json') + p.add_argument('--save-snapshots',action='store_true',help='Render and save latest.jpg every 10 seconds') + p.add_argument('--preview-port',type=int,default=8765);p.add_argument('--preview-host',default='0.0.0.0');p.add_argument('--preview-every',type=int,default=3) + p.add_argument('--write-video',action='store_true');p.add_argument('--save-pose-records',action='store_true');p.add_argument('--calibration-verified',action='store_true') + a=p.parse_args(argv) + if not 1<=a.max_people<=16 or not 0=a.max_frames) or (a.duration and time.monotonic()-run_started>=a.duration):break + if video: + ok,frame=video.read() + if not ok:break + capture={'sequence':count,'capture_monotonic_ms':time.monotonic()*1000,'dropped_before_inference':0,'latency_clock':'video_decode_complete'} + else:frame,capture=native.read() + if frame is None:break + tick=time.monotonic();record=pose.infer(frame,count);pose_ms.append((time.monotonic()-tick)*1000) + rt=time.monotonic();snapshot=native.can_snapshot();active,valid=select_active_door(a.active_door,snapshot,a.can_side,a.can_timeout_ms) + assignments=assigner.update(count,record,(a.height,a.width));raw,mask=evaluator.evaluate(count,record,(a.height,a.width)) + doors,risk=gate_risks(raw,assignments,active,valid);risk_ms.append((time.monotonic()-rt)*1000) + result={'frame_index':count,'track_ids':record['track_ids'],'boxes':record['boxes'],'skeletons':visualization_skeletons(record),'active_door':active,'door_assignments':assignments,'door_risks':doors,'risk':risk,'ungated_risk':raw,'capture':capture,'can':dict(snapshot,valid=valid,source='hardware' if a.active_door=='can' else 'simulation',side=a.can_side),'calibration_verified':a.calibration_verified} + result['capture_to_result_ms']=round(time.monotonic()*1000-capture['capture_monotonic_ms'],3) + results.publish(result) + if out:out.write(json.dumps(result,ensure_ascii=False,allow_nan=False)+'\n') + if pose_out:pose_out.write(json.dumps(record,ensure_ascii=False,allow_nan=False)+'\n') + now=time.monotonic() + snapshot_due=now-last_snapshot>=10 + if a.save_results and snapshot_due: + (a.output/'latest.json').write_text(json.dumps(result,ensure_ascii=False,indent=2)) + status=('CAN LIVE' if a.active_door=='can' else 'CAN SIMULATION: '+active)+' | '+('CALIBRATED' if a.calibration_verified else 'CALIBRATION UNVERIFIED') + if preview and count%a.preview_every==0: + preview.submit(frame,record,risk,mask,assignments,status) + if writer or (a.save_snapshots and snapshot_due): + canvas=render_output(frame,record,risk,mask,assignments,status) + if writer:writer.write(canvas) + if a.save_snapshots and snapshot_due: + cv2.imwrite(str(a.output/'latest.jpg'),canvas) + if snapshot_due:last_snapshot=now + count+=1;dropped+=capture['dropped_before_inference'];levels[risk['risk_level']]+=1;ages.append(result['capture_to_result_ms']);latencies.append((time.monotonic()-tick)*1000) + if now-last_report>=5: + print(json.dumps({'frames':count,'fps':round(count/(now-run_started),2),'people':len(record['boxes']),'risk':risk['risk_type'],'latency_ms':round(float(np.mean(latencies)),1),'capture_age_ms':round(float(np.mean(ages)),1),'can_valid':valid}),flush=True);last_report=now + finally: + if pose_out:pose_out.close() + except Exception as exc: + error=str(exc);raise + finally: + ended=time.monotonic() + if video:video.release() + if native:native.close() + if writer:writer.release() + if preview:preview.close() + elapsed=ended-(run_started or started) + summary={'mode':'baseline_v4l2_can_native_trt_cms_v9_10','camera':a.camera if not a.video else None,'video':str(a.video) if a.video else None,'frames':count,'elapsed_seconds':round(elapsed,3),'throughput_fps':round(count/max(elapsed,.001),3),'dropped_latest_frames':dropped,'risk_levels':dict(levels),'active_door_mode':a.active_door,'can_side':a.can_side,'calibration_verified':a.calibration_verified,'visualize_requested':a.visualize,'save_results':a.save_results,'save_snapshots':a.save_snapshots,'latency_window_frames':3000,'latency_ms':{'total':timing_stats(list(latencies)),'pose':timing_stats(list(pose_ms)),'risk':timing_stats(list(risk_ms)),'capture_to_result':timing_stats(list(ages))},'error':error} + if summary_sink is not None:summary_sink.update(summary) + if save_any:(a.output/'summary.json').write_text(json.dumps(summary,ensure_ascii=False,indent=2)) + print(json.dumps(summary,ensure_ascii=False),flush=True) + +class CameraPipeline: + """One lifecycle per instance, with the 0920 start/latest/stop/join pattern. + + The worker owns inference. Downstream threads read the latest immutable + snapshot directly, without JSON files, HTTP, callbacks or a result queue. + """ + def __init__(self, args=None): + self.args = copy.deepcopy(args if args is not None else parse_args([])) + self._results = LatestResult() + self._stop = threading.Event() + self._lifecycle_lock = threading.Lock() + self._thread = None + self._error = None + self.summary = {} + + @property + def latest_result(self): + """The latest complete result variable, or None before the first frame.""" + return self._results.load() + + def latest(self): + return self.latest_result + + def wait_for_result(self, after_frame_index=-1, timeout=None): + return self._results.wait(after_frame_index, timeout) + + def start(self): + with self._lifecycle_lock: + if self._thread is not None: + raise RuntimeError('CameraPipeline can only be started once') + self._thread = threading.Thread(target=self._run, name='vision-risk') + self._thread.start() + return self + + def _run(self): + try: + run_pipeline(self.args, self._results, self._stop, self.summary) + except BaseException as exc: + self._error = exc + finally: + self._results.close() + + def running(self): + return self._thread is not None and self._thread.is_alive() + + def request_stop(self): + self._stop.set() + + def join(self, timeout=None): + if self._thread is not None: + self._thread.join(timeout) + if not self._thread.is_alive() and self._error is not None: + raise self._error + + def __enter__(self): + return self.start() + + def __exit__(self, *args): + self.request_stop() + self.join() + + +def main(): + pipeline=CameraPipeline(parse_args()) + for sig in (signal.SIGINT,signal.SIGTERM): + signal.signal(sig,lambda *_:pipeline.request_stop()) + pipeline.start() + # Periodic joins allow Python to deliver stop signals on the main thread. + while pipeline.running(): + pipeline.join(.2) + pipeline.join() + + +if __name__=='__main__':main() diff --git a/run_camera.sh b/run_camera.sh index c153d56..c19253a 100755 --- a/run_camera.sh +++ b/run_camera.sh @@ -1,4 +1,5 @@ #!/usr/bin/env bash set -euo pipefail -root="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)" -exec "$root/build/camera_risk" --config "$root/config/camera.json" "$@" +root="$(cd "$(dirname "$0")" && pwd)" +export OPENBLAS_NUM_THREADS=1 OMP_NUM_THREADS=2 +exec "${PYTHON:-python3}" -u "$root/run_camera.py" "$@" diff --git a/run_risk.py b/run_risk.py new file mode 100644 index 0000000..9cdf2b3 --- /dev/null +++ b/run_risk.py @@ -0,0 +1,5 @@ +"""0923 entry point: native TensorRT camera/video inference and CMS v9.10 risk.""" +from run_camera import main + +if __name__ == '__main__': + main() diff --git a/scripts/build.sh b/scripts/build.sh index 5b0fa8d..ea63c29 100755 --- a/scripts/build.sh +++ b/scripts/build.sh @@ -1,7 +1,7 @@ #!/usr/bin/env bash set -euo pipefail -root="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")/.." && pwd)" +root="$(cd "$(dirname "$0")/.." && pwd)" cmake_bin="${CMAKE_COMMAND:-cmake}" if [[ -x "$HOME/miniconda3/envs/yolo/bin/cmake" ]]; then cmake_bin="${CMAKE_COMMAND:-$HOME/miniconda3/envs/yolo/bin/cmake}"; fi -"$cmake_bin" -S "$root" -B "$root/build" -DCMAKE_BUILD_TYPE=Release -DCMAKE_CUDA_COMPILER="${CUDACXX:-/usr/local/cuda/bin/nvcc}" "$@" +"$cmake_bin" -S "$root" -B "$root/build" -DCMAKE_BUILD_TYPE=Release -DCMAKE_CUDA_COMPILER="${CUDACXX:-/usr/local/cuda/bin/nvcc}" "$cmake_bin" --build "$root/build" -j4 diff --git a/scripts/deploy.sh b/scripts/deploy.sh new file mode 100755 index 0000000..7d4455f --- /dev/null +++ b/scripts/deploy.sh @@ -0,0 +1,6 @@ +#!/usr/bin/env bash +set -euo pipefail +root="$(cd "$(dirname "$0")/.." && pwd)" +remote="${1:-orin@192.168.195.100}" +rsync -a --exclude .git --exclude build --exclude __pycache__ --exclude outputs --exclude logs --exclude models --exclude 'tests/*.mp4' "$root/" "$remote:vision-anti-pinch-0923/" +ssh "$remote" 'cd ~/vision-anti-pinch-0923 && mkdir -p models && if [ ! -e models/engines ]; then ln -s ../../vision-anti-pinch-0920/orin/engines models/engines; fi && ./scripts/build.sh' diff --git a/scripts/initialize_camera.sh b/scripts/initialize_camera.sh index 8da09bd..c840953 100755 --- a/scripts/initialize_camera.sh +++ b/scripts/initialize_camera.sh @@ -1,4 +1,27 @@ #!/usr/bin/env bash +# Use the previously deployed 0920 drivers; do not force module compatibility. set -euo pipefail -root="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")/.." && pwd)" -exec "$root/orin/scripts/initialize_camera.sh" "$@" +root="$(cd "$(dirname "$0")/.." && pwd)" +driver_dir="${DRIVER_DIR:-$root/../vision-anti-pinch-0920/orin/bin/drivers}" +device="${1:-/dev/video5}" +kernel="$(uname -r)" +for driver in max9295 max9296 sgx-yuv-gmsl2; do + module="${driver//-/_}" + [[ -d "/sys/module/$module" ]] && continue + file="$driver_dir/$driver.ko" + [[ -f "$file" ]] || { echo "Missing driver: $file" >&2; exit 1; } + vermagic="$(modinfo -F vermagic "$file")" + if [[ -n "$vermagic" && "${vermagic%% *}" != "$kernel" ]]; then + echo "Kernel mismatch: $file ($vermagic), running $kernel" >&2; exit 1 + fi + if [[ -z "$vermagic" ]]; then + echo "No vermagic metadata: $file; using normal insmod kernel validation." + fi + options=() + [[ "$driver" == sgx-yuv-gmsl2 ]] && options=(gmsl_mode=2,2,2,2) + sudo -n insmod "$file" "${options[@]}" +done +for _ in {1..20}; do [[ -c "$device" ]] && break; sleep .25; done +[[ -c "$device" ]] || { echo "Camera not created: $device" >&2; exit 1; } +v4l2-ctl -d "$device" -c sensor_mode=1,trig_pin=0xffff0007 +echo "Camera ready: $device" diff --git a/scripts/package_orin.sh b/scripts/package_orin.sh new file mode 100755 index 0000000..767a105 --- /dev/null +++ b/scripts/package_orin.sh @@ -0,0 +1,34 @@ +#!/usr/bin/env bash +# Run after building on the target Orin; outputs a relocatable ./orin/run.sh package. +set -euo pipefail +root="$(cd "$(dirname "$0")/.." && pwd)" +package="$root/orin" +engine_dir="${ENGINE_DIR:-$root/models/engines}" +driver_dir="${DRIVER_DIR:-$root/../vision-anti-pinch-0920/orin/bin/drivers}" +[[ -f "$root/build/libvision_native.so" ]] || { echo 'Build the Orin native library first.' >&2; exit 1; } +mkdir -p "$package/build" "$package/lib" "$package/pipeline" "$package/annotations" "$package/models/engines" "$package/bin/drivers" "$package/scripts" "$package/examples" +cp -p "$root/build/libvision_native.so" "$package/build/" +cp -p "$root/run_camera.py" "$package/" +cp -p "$root/pipeline/"*.py "$package/pipeline/" +cp -p "$root/annotations/"*.json "$package/annotations/" +cp -p "$root/lib/libusbcanfd.so" "$root/lib/libusb-1.0.so" "$package/lib/" +ln -sfn libusbcanfd.so "$package/lib/libusbcanfd.so.1.0.12" +ln -sfn libusb-1.0.so "$package/lib/libusb-1.0.so.0" +for model in yolo11s.engine rtmw_l_384.engine; do + [[ -f "$engine_dir/$model" ]] || { echo "Missing engine: $engine_dir/$model" >&2; exit 1; } + cp -p "$engine_dir/$model" "$package/models/engines/" +done +for driver in max9295 max9296 sgx-yuv-gmsl2; do + if [[ -f "$driver_dir/$driver.ko" ]]; then cp -p "$driver_dir/$driver.ko" "$package/bin/drivers/"; fi + [[ -f "$package/bin/drivers/$driver.ko" ]] || { echo "Missing driver: $driver.ko" >&2; exit 1; } +done +cp -p "$root/scripts/initialize_camera.sh" "$package/scripts/initialize_camera_impl.sh" +cp -p "$root/examples/read_latest.py" "$package/examples/" +python3 - "$package" <<'PY' +import hashlib,sys +from pathlib import Path +root=Path(sys.argv[1]);ignored={'__pycache__','outputs','logs'} +files=sorted(p for p in root.rglob('*') if p.is_file() and not p.is_symlink() and p.name!='SHA256SUMS' and not ignored.intersection(p.relative_to(root).parts)) +(root/'SHA256SUMS').write_text(''.join(hashlib.sha256(p.read_bytes()).hexdigest()+' '+p.relative_to(root).as_posix()+'\n' for p in files)) +PY +echo "Package ready: $package/run.sh" diff --git a/scripts/start.sh b/scripts/start.sh new file mode 100755 index 0000000..d471295 --- /dev/null +++ b/scripts/start.sh @@ -0,0 +1,21 @@ +#!/usr/bin/env bash +set -euo pipefail +root="$(cd "$(dirname "$0")/.." && pwd)" +mkdir -p "$root/logs" +pidfile="$root/logs/camera.pid" +if [[ -f "$pidfile" ]]; then + pid=$(cat "$pidfile") + if [[ "$pid" =~ ^[0-9]+$ ]] && kill -0 "$pid" 2>/dev/null && grep -zFq "$root/run_camera.py" "/proc/$pid/cmdline"; then + echo "Already running: PID $pid"; exit 0 + fi +fi +nohup "$root/run_camera.sh" --output "$root/outputs/live" "$@" > "$root/logs/live.log" 2>&1 < /dev/null & +pid=$!; echo "$pid" > "$pidfile" +sleep 2 +if ! kill -0 "$pid" 2>/dev/null; then tail -30 "$root/logs/live.log"; exit 1; fi +echo "Started PID $pid; log: $root/logs/live.log" +for option in "$@"; do + if [[ "$option" == --visualize ]]; then + echo "Visualization requested; see log for startup status (default port 8765)." + fi +done diff --git a/scripts/stop.sh b/scripts/stop.sh new file mode 100755 index 0000000..c372287 --- /dev/null +++ b/scripts/stop.sh @@ -0,0 +1,13 @@ +#!/usr/bin/env bash +set -euo pipefail +root="$(cd "$(dirname "$0")/.." && pwd)" +pidfile="$root/logs/camera.pid" +[[ -f "$pidfile" ]] || { echo 'Not running'; exit 0; } +pid=$(cat "$pidfile") +if [[ "$pid" =~ ^[0-9]+$ ]] && kill -0 "$pid" 2>/dev/null; then + grep -zFq "$root/run_camera.py" "/proc/$pid/cmdline" || { echo 'PID belongs to another process; refusing to stop'; exit 1; } + kill -TERM "$pid" + for _ in {1..30}; do kill -0 "$pid" 2>/dev/null || break; sleep .2; done +fi +rm -f "$pidfile" +echo 'Stop requested' diff --git a/send.hpp b/send.hpp new file mode 100644 index 0000000..54c0cc1 --- /dev/null +++ b/send.hpp @@ -0,0 +1,108 @@ + +#pragma once +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +struct DebugInfo +{ + int64_t TimeStampMs; + uint8_t Img[1920 * 1536 * 2]; +}; + +/*******************线程发送检测图片与结果 ***********************************/ +pthread_mutex_t g_send_mutex; +std::map g_map_debuginfo; +/*******************线程发送检测图片与结果***********************************/ + +int g_sockfd = -1; +std::string g_server_ip = "192.168.195.110"; +int g_server_port = 8002; + +void disconnect_server() +{ + if (g_sockfd >= 0) + { + close(g_sockfd); + g_sockfd = -1; + } +} + +bool reconnect_server() +{ + + disconnect_server(); + + struct sockaddr_in s_addr; + + /*建立socket*/ + g_sockfd = socket(AF_INET, SOCK_STREAM, 0); + if (g_sockfd < 0) + { + g_sockfd = -1; + return false; + } + + /*设置服务器ip*/ + bzero(&s_addr, sizeof(s_addr)); + s_addr.sin_family = AF_INET; + s_addr.sin_port = htons(g_server_port); + + if (inet_aton(g_server_ip.c_str(), (struct in_addr *)&s_addr.sin_addr.s_addr) <= 0) + { + close(g_sockfd); + g_sockfd = -1; + return false; + } + + /*开始连接服务器*/ + if (connect(g_sockfd, (struct sockaddr *)&s_addr, sizeof(struct sockaddr)) < 0) + { + close(g_sockfd); + g_sockfd = -1; + return false; + } + + return true; +} + +bool send_debuginfo(DebugInfo *pData) +{ + bool res = false; + if (g_sockfd >= 0) + { + // tcp 发送数据 + int len = 0; + if (pData) + { + len = send(g_sockfd, pData, sizeof(DebugInfo), MSG_NOSIGNAL); + } + + if (len < 0) + { + reconnect_server(); + sleep(1); + } + else + { + res = true; + } + } + else + { + reconnect_server(); + sleep(1); + } + return res; +} diff --git a/tests/test_latest_result.py b/tests/test_latest_result.py new file mode 100644 index 0000000..e3be130 --- /dev/null +++ b/tests/test_latest_result.py @@ -0,0 +1,155 @@ +"""Latest-result ownership, optional output, and asynchronous consumer contract.""" +import copy +import json +from pathlib import Path +import sys +import tempfile +import threading +import time +import unittest +from unittest.mock import patch +from urllib.request import urlopen +import numpy as np + +ROOT = Path(__file__).resolve().parents[1] +sys.path.insert(0, str(ROOT)) +import run_camera as app +from latest_result import LatestResult, result_to_dict + + +class ResultTests(unittest.TestCase): + def test_read_only_owned_snapshot_and_old_reference(self): + channel=LatestResult() + self.assertIsNone(channel.load()) + source={'frame_index':0,'risk':{'risk_score':None},'boxes':[[1,2,3,4]]} + channel.publish(source);old=channel.load() + source['boxes'][0][0]=99 + self.assertEqual(old['boxes'][0][0],1) + with self.assertRaises(TypeError):old['risk']['risk_score']=0 + with self.assertRaises(TypeError):old['boxes'][0][0]=0 + channel.publish({'frame_index':1,'risk':{'risk_score':1}}) + self.assertEqual(old['frame_index'],0) + self.assertIs(channel.load(),channel.load()) + self.assertEqual(channel.load()['frame_index'],1) + self.assertEqual(result_to_dict(old)['boxes'],[[1,2,3,4]]) + + def test_wait_replaces_old_frames_and_close_wakes_reader(self): + channel=LatestResult() + for i in range(20):channel.publish({'frame_index':i}) + self.assertEqual(channel.wait(0,.1)['frame_index'],19) + self.assertIsNone(channel.wait(19,.01)) + answer=[];worker=threading.Thread(target=lambda:answer.append(channel.wait(19,1))) + worker.start();channel.close();worker.join(1) + self.assertFalse(worker.is_alive());self.assertEqual(answer,[None]) + self.assertEqual(channel.load()['frame_index'],19) + + +class FakeNative: + instances=[] + def __init__(self,args): + self.args=args;self.index=0;self.closed=False;self.opened=False + self.frame=np.zeros((args.height,args.width,3),np.uint8) + self.__class__.instances.append(self) + def open_camera(self):self.opened=True + def read(self): + time.sleep(.005) + self.index+=1 + return self.frame,{'sequence':self.index,'capture_monotonic_ms':time.monotonic()*1000,'dropped_before_inference':0} + def can_snapshot(self):return {'positions':[-1]*4,'age_ms':-1} + def close(self):self.closed=True + + +class FakePose: + def __init__(self,*args):pass + def infer(self,frame,index): + return {'track_ids':[7],'boxes':[[0,0,frame.shape[1],frame.shape[0]]], + 'temporal_selected_keypoints':[[[20.,20.]]*133], + 'temporal_metadata':[{'left':{},'right':{}}],'global_hands':[]} + + +class PipelineTests(unittest.TestCase): + def run_fake(self,output,*extra): + args=app.parse_args(['--output',str(output),'--max-frames','3',*extra]) + pipeline=app.CameraPipeline(args) + with patch.object(app,'NativeBackend',FakeNative),patch.object(app,'NativePoseEngine',FakePose): + pipeline.start();pipeline.join(10) + self.assertFalse(pipeline.running()) + self.assertTrue(FakeNative.instances[-1].closed) + return pipeline + + def test_default_has_no_renderer_http_encoding_or_disk_output(self): + with tempfile.TemporaryDirectory() as tmp: + out=Path(tmp)/'must-not-exist' + with patch.object(app,'Preview') as preview,patch.object(app,'render_frame') as render,patch.object(app.cv2,'imencode') as encode: + pipeline=self.run_fake(out) + preview.assert_not_called();render.assert_not_called();encode.assert_not_called() + self.assertFalse(out.exists()) + result=pipeline.latest_result + self.assertEqual(result['frame_index'],2) + self.assertIs(result,pipeline.latest()) + self.assertEqual(result['risk']['risk_type'],'CAN_UNKNOWN') + self.assertIsNone(result['risk']['risk_score']) + self.assertEqual(len(result['skeletons'][0]['keypoints']),133) + self.assertEqual(pipeline.summary['frames'],3) + with self.assertRaises(RuntimeError):pipeline.start() + + def test_json_output_matches_in_memory_result_without_rendering(self): + with tempfile.TemporaryDirectory() as tmp,patch.object(app,'render_frame') as render: + out=Path(tmp)/'json' + pipeline=self.run_fake(out,'--save-results') + disk=[json.loads(line) for line in (out/'causal_results.jsonl').read_text().splitlines()] + self.assertEqual(disk[-1],result_to_dict(pipeline.latest_result)) + self.assertTrue((out/'latest.json').exists());self.assertTrue((out/'summary.json').exists()) + self.assertFalse((out/'latest.jpg').exists());render.assert_not_called() + + def test_preview_start_failure_does_not_stop_inference(self): + with tempfile.TemporaryDirectory() as tmp,patch.object(app,'Preview',side_effect=OSError('port busy')): + pipeline=self.run_fake(Path(tmp)/'none','--visualize') + self.assertEqual(pipeline.latest_result['frame_index'],2) + + def test_worker_error_reaches_downstream_and_releases_resources(self): + with tempfile.TemporaryDirectory() as tmp,patch.object(FakePose,'infer',side_effect=RuntimeError('pose failed')): + with self.assertRaisesRegex(RuntimeError,'pose failed'):self.run_fake(Path(tmp)/'none') + self.assertTrue(FakeNative.instances[-1].closed) + + def test_consumer_reads_while_running_and_can_request_stop(self): + args=app.parse_args([]);pipeline=app.CameraPipeline(args) + with patch.object(app,'NativeBackend',FakeNative),patch.object(app,'NativePoseEngine',FakePose): + pipeline.start() + try: + first=pipeline.wait_for_result(timeout=5) + self.assertIsNotNone(first) + second=pipeline.wait_for_result(first['frame_index'],5) + self.assertGreater(second['frame_index'],first['frame_index']) + self.assertTrue(pipeline.running()) + finally: + pipeline.request_stop();pipeline.join(5) + self.assertFalse(pipeline.running()) + + +class PreviewTests(unittest.TestCase): + def test_http_results_follow_memory_without_waiting_for_new_image(self): + results=LatestResult();results.publish({'frame_index':1,'risk':{'risk_score':None}}) + preview=app.Preview('127.0.0.1',0,results) + try: + url='http://127.0.0.1:%d/result.json'%preview.server.server_port + self.assertEqual(json.load(urlopen(url,timeout=2))['frame_index'],1) + results.publish({'frame_index':2,'risk':{'risk_score':None}}) + self.assertEqual(json.load(urlopen(url,timeout=2))['frame_index'],2) + finally:preview.close() + + def test_preview_renders_on_observer_thread(self): + results=LatestResult();rendered=threading.Event();worker_threads=[] + def render(*args): + worker_threads.append(threading.current_thread().name) + return np.zeros((48,64,3),np.uint8) + preview=app.Preview('127.0.0.1',0,results) + try: + with patch.object(app,'render_output',side_effect=render),patch.object(preview,'update',side_effect=lambda image:rendered.set()): + preview.submit(None,None,None,None,None,None) + self.assertTrue(rendered.wait(2)) + self.assertEqual(worker_threads,['risk-preview']) + finally:preview.close() + + +if __name__=='__main__':unittest.main() diff --git a/tests/test_risk_integration.py b/tests/test_risk_integration.py new file mode 100644 index 0000000..26e5ad1 --- /dev/null +++ b/tests/test_risk_integration.py @@ -0,0 +1,70 @@ +"""Behavior tests for CAN validity and upstream per-door evidence gating.""" +import sys +import unittest +import tempfile +from pathlib import Path +import numpy as np +ROOT=Path(__file__).resolve().parents[1] +sys.path.insert(0,str(ROOT)) +from run_camera import select_active_door,gate_risks,JsonlWriter +from door_mask_risk import DoorMaskSequence,DoorRiskEvaluator +from door_zone_risk import DoorZoneConfig,DoorZoneAssigner +from predict_optimized_pose_videos import BoxTracker + +class CanTests(unittest.TestCase): + def test_four_bytes_and_side(self): + for front in [0,20]: + for rear in [0,30]: + for side in ['left','right']: + pos=[40]*4;i=0 if side=='left' else 1;pos[i]=front;pos[i+2]=rear + expected='both' if front==rear==0 else 'front' if front==0 else 'rear' if rear==0 else 'none' + self.assertEqual(select_active_door('can',{'positions':pos,'age_ms':10},side),(expected,True)) + def test_invalid_never_becomes_clear(self): + for age in [-1,1001]: + active,valid=select_active_door('can',{'positions':[0]*4,'age_ms':age}) + self.assertFalse(valid) + doors,risk=gate_risks({'regions':[]},[],active,valid) + self.assertEqual(risk['risk_level'],'unknown');self.assertIsNone(risk['risk_score']) + self.assertTrue(all(r['risk_level']=='unknown' for r in doors.values())) + def test_real_none_is_inactive(self): + active,valid=select_active_door('can',{'positions':[10]*4,'age_ms':0}) + _,risk=gate_risks({'regions':[]},[],active,valid) + self.assertEqual(risk['risk_level'],'inactive') + +class RiskTests(unittest.TestCase): + def setUp(self): + self.evaluator=DoorRiskEvaluator(DoorMaskSequence(constant_mask=ROOT/'annotations/constant_door_mask.json')) + def test_body_fallback_intersection_and_can_filter(self): + record={'boxes':[[0,0,1920,1536]],'track_ids':[42],'global_hands':[]} + raw,_=self.evaluator.evaluate(0,record,(1536,1920)) + self.assertEqual(raw['risk_level'],'danger') + self.assertEqual(raw['regions'][0]['source'],'YOLO_BODY_BOX_FALLBACK') + assignments=[{'person_index':0,'track_id':42,'door_zone':'rear'}] + _,rear=gate_risks(raw,assignments,'rear',True) + _,front=gate_risks(raw,assignments,'front',True) + self.assertEqual(rear['risk_type'],'REAR_DANGER') + self.assertEqual(front['risk_type'],'FRONT_CLEAR') + def test_transition_keeps_unknown(self): + raw,_=self.evaluator.evaluate(0,{'boxes':[[0,0,1920,1536]],'track_ids':[42]},(1536,1920)) + _,risk=gate_risks(raw,[{'person_index':0,'track_id':42,'door_zone':'transition'}],'front',True) + self.assertEqual(risk['risk_level'],'unknown');self.assertIsNone(risk['risk_score']) + def test_hand_overrides_whole_body_fallback(self): + hand={'person_id':1,'side':'left','keypoints':[[1,1],[50,1],[50,50],[1,50]]} + raw,_=self.evaluator.evaluate(0,{'boxes':[[0,0,1920,1536]],'track_ids':[42],'global_hands':[hand]},(1536,1920)) + self.assertEqual(len(raw['regions']),1) + self.assertEqual(raw['regions'][0]['source'],'RTMW_HAND_KEYPOINTS') + def test_expired_tracker_state_pruned(self): + t=BoxTracker(max_age=30);t.update([[1,1,100,100]],0,(1536,1920));t.update([],31,(1536,1920)) + self.assertEqual(t.tracks,{}) + +class OutputTests(unittest.TestCase): + def test_rotation_keeps_complete_utf8_records(self): + with tempfile.TemporaryDirectory() as tmp: + path=Path(tmp)/'risk.jsonl' + with JsonlWriter(path,19) as writer: + for i in range(6):writer.write('{"门":%d}\n'%i) + self.assertEqual(path.read_text(),'{"门":5}\n') + self.assertEqual(Path(str(path)+'.1').read_text(),'{"门":4}\n') + self.assertEqual(Path(str(path)+'.2').read_text(),'{"门":3}\n') + +if __name__=='__main__':unittest.main() diff --git a/v4l2_capture.cpp b/v4l2_capture.cpp new file mode 100644 index 0000000..2f20786 --- /dev/null +++ b/v4l2_capture.cpp @@ -0,0 +1,380 @@ +#include "v4l2_capture.h" +#include +#include +#include +#include +#include +#include +#include +#include +#include + +// 四字符码构造宏 (使用 v4l2_fourcc 避免重复定义) +// 注意: 这里使用videodev2.h已有的v4l2_fourcc宏 + +namespace { + +// 将 v4l2_capability 的 driver[16] / card[32] 转成 C 字符串 (v4l2_capability 使用 __u8) +inline std::string to_cstr(const __u8 *arr, size_t len) { + std::string s; + s.reserve(len); + for (size_t i = 0; i < len && arr[i] != '\0'; ++i) { + char c = static_cast(arr[i]); + if (c >= 32 && c < 127) s += c; // 只保留可打印 ASCII + } + // 去除末尾空格 + while (!s.empty() && s.back() == ' ') s.pop_back(); + return s; +} + +// 打印 v4l2 格式调试信息 +void dumpFormat(struct v4l2_format &fmt, const char *msg) { + printf("%s: width=%u height=%u pixfmt=0x%08x '%c%c%c%c' bpl=%u size=%u\n", + msg, + fmt.fmt.pix.width, + fmt.fmt.pix.height, + fmt.fmt.pix.pixelformat, + (fmt.fmt.pix.pixelformat) & 0xFF, + (fmt.fmt.pix.pixelformat >> 8) & 0xFF, + (fmt.fmt.pix.pixelformat >> 16) & 0xFF, + (fmt.fmt.pix.pixelformat >> 24) & 0xFF, + fmt.fmt.pix.bytesperline, + fmt.fmt.pix.sizeimage); +} + +} // namespace + +V4l2Capture::~V4l2Capture() { + close(); +} + +bool V4l2Capture::open(const std::string &device) { + if (fd_ >= 0) { + return false; + } + fd_ = ::open(device.c_str(), O_RDWR | O_NONBLOCK, 0); + if (fd_ < 0) { + printf("V4l2Capture: open %s failed: %s\n", device.c_str(), strerror(errno)); + return false; + } + + // 检查设备能力 + struct v4l2_capability cap; + if (ioctl(fd_, VIDIOC_QUERYCAP, &cap) < 0) { + printf("V4l2Capture: VIDIOC_QUERYCAP failed: %s\n", strerror(errno)); + ::close(fd_); + fd_ = -1; + return false; + } + printf("V4l2Capture: opened %s, driver='%s' card='%s', caps=0x%x\n", + device.c_str(), + to_cstr(cap.driver, sizeof(cap.driver)).c_str(), + to_cstr(cap.card, sizeof(cap.card)).c_str(), + cap.capabilities); + + // 必须支持 V4L2_CAP_VIDEO_CAPTURE | V4L2_CAP_STREAMING + if (!(cap.capabilities & V4L2_CAP_VIDEO_CAPTURE)) { + printf("V4l2Capture: not a video capture device\n"); + ::close(fd_); + fd_ = -1; + return false; + } + if (!(cap.capabilities & V4L2_CAP_STREAMING)) { + printf("V4l2Capture: does not support streaming I/O\n"); + ::close(fd_); + fd_ = -1; + return false; + } + + return true; +} + +void V4l2Capture::close() { + if (streaming_) { + stopStreaming(); + } + buffers_.clear(); + if (fd_ >= 0) { + ::close(fd_); + fd_ = -1; + } +} + +bool V4l2Capture::setFormat(uint32_t width, uint32_t height, uint32_t pixelformat) { + if (fd_ < 0) return false; + + struct v4l2_format fmt; + std::memset(&fmt, 0, sizeof(fmt)); + fmt.type = V4L2_BUF_TYPE_VIDEO_CAPTURE; + fmt.fmt.pix.width = width; + fmt.fmt.pix.height = height; + fmt.fmt.pix.pixelformat = pixelformat; + fmt.fmt.pix.field = V4L2_FIELD_NONE; // 逐行扫描 + + if (ioctl(fd_, VIDIOC_S_FMT, &fmt) < 0) { + printf("V4l2Capture: VIDIOC_S_FMT failed: %s\n", strerror(errno)); + return false; + } + + dumpFormat(fmt, "V4l2Capture: setFormat"); + + width_ = fmt.fmt.pix.width; + height_ = fmt.fmt.pix.height; + pixelformat_ = fmt.fmt.pix.pixelformat; + bytes_per_line_ = fmt.fmt.pix.bytesperline; + + // 如果驱动不按请求设置, 给出警告 + if (width_ != width || height_ != height) { + printf("V4l2Capture: warning: requested %ux%u but got %ux%u\n", + width, height, width_, height_); + } + if (pixelformat_ != pixelformat) { + char req[5] = {0}, got[5] = {0}; + req[0] = pixelformat & 0xFF; req[1] = (pixelformat>>8)&0xFF; + req[2] = (pixelformat>>16)&0xFF; req[3] = (pixelformat>>24)&0xFF; + got[0] = pixelformat_ & 0xFF; got[1] = (pixelformat_>>8)&0xFF; + got[2] = (pixelformat_>>16)&0xFF; got[3] = (pixelformat_>>24)&0xFF; + printf("V4l2Capture: warning: requested '%s' but got '%s'\n", req, got); + } + + return true; +} + +bool V4l2Capture::setFps(uint32_t fps) { + if (fd_ < 0) return false; + + struct v4l2_streamparm parm; + std::memset(&parm, 0, sizeof(parm)); + parm.type = V4L2_BUF_TYPE_VIDEO_CAPTURE; + + if (ioctl(fd_, VIDIOC_G_PARM, &parm) < 0) { + printf("V4l2Capture: VIDIOC_G_PARM failed: %s\n", strerror(errno)); + return false; + } + + // 尝试设置 FPS. 需要 V4L2_CAP_TIMEPERFRAME 支持 + if (parm.parm.capture.capability & V4L2_CAP_TIMEPERFRAME) { + parm.parm.capture.timeperframe.numerator = 1; + parm.parm.capture.timeperframe.denominator = fps; + if (ioctl(fd_, VIDIOC_S_PARM, &parm) < 0) { + printf("V4l2Capture: VIDIOC_S_PARM fps=%u failed: %s\n", fps, strerror(errno)); + return false; + } + printf("V4l2Capture: set fps=%u (actual %u/%u)\n", + fps, + parm.parm.capture.timeperframe.denominator, + parm.parm.capture.timeperframe.numerator); + } else { + printf("V4l2Capture: device does not support V4L2_CAP_TIMEPERFRAME\n"); + } + return true; +} + +bool V4l2Capture::setControl(uint32_t id, int32_t value) { + if (fd_ < 0) return false; + + struct v4l2_control ctrl; + std::memset(&ctrl, 0, sizeof(ctrl)); + ctrl.id = id; + ctrl.value = value; + + if (ioctl(fd_, VIDIOC_S_CTRL, &ctrl) < 0) { + // 有些控制可能是只读的或不存在, 打印警告但不失败 + // 继续尝试 VIDIOC_TRY_CTRL 看是否可用 + struct v4l2_ext_controls ext_ctrls = {}; + struct v4l2_ext_control ext_ctrl = {}; + ext_ctrl.id = id; + ext_ctrl.value = value; + ext_ctrls.count = 1; + ext_ctrls.controls = &ext_ctrl; + + if (ioctl(fd_, VIDIOC_TRY_EXT_CTRLS, &ext_ctrls) < 0) { + printf("V4l2Capture: VIDIOC_S_CTRL id=0x%08x value=%d failed: %s (ignored)\n", + id, value, strerror(errno)); + return false; + } else { + printf("V4l2Capture: VIDIOC_S_CTRL id=0x%08x value=%d not supported (ignored)\n", id, value); + return false; + } + } + printf("V4l2Capture: set control id=0x%08x value=%d\n", id, value); + return true; +} + +bool V4l2Capture::initMmap(uint32_t buffer_count) { + if (fd_ < 0) return false; + + // 请求 MMAP 缓冲区 + struct v4l2_requestbuffers req; + std::memset(&req, 0, sizeof(req)); + req.count = buffer_count; + req.type = V4L2_BUF_TYPE_VIDEO_CAPTURE; + req.memory = V4L2_MEMORY_MMAP; + + if (ioctl(fd_, VIDIOC_REQBUFS, &req) < 0) { + printf("V4l2Capture: VIDIOC_REQBUFS failed: %s\n", strerror(errno)); + return false; + } + if (req.count < buffer_count) { + printf("V4l2Capture: only %u buffers available (requested %u)\n", req.count, buffer_count); + buffer_count = req.count; + } + printf("V4l2Capture: allocated %u MMAP buffers\n", (unsigned)req.count); + + buffers_.resize(buffer_count); + for (uint32_t i = 0; i < buffer_count; ++i) { + struct v4l2_buffer buf; + std::memset(&buf, 0, sizeof(buf)); + buf.type = V4L2_BUF_TYPE_VIDEO_CAPTURE; + buf.memory = V4L2_MEMORY_MMAP; + buf.index = i; + + if (ioctl(fd_, VIDIOC_QUERYBUF, &buf) < 0) { + printf("V4l2Capture: VIDIOC_QUERYBUF %u failed: %s\n", i, strerror(errno)); + return false; + } + + buffers_[i].length = buf.length; + buffers_[i].start = mmap(nullptr, buf.length, + PROT_READ | PROT_WRITE, MAP_SHARED, + fd_, buf.m.offset); + if (buffers_[i].start == MAP_FAILED) { + printf("V4l2Capture: mmap buffer %u failed: %s\n", i, strerror(errno)); + return false; + } + buffers_[i].queued = false; + printf("V4l2Capture: buffer %u: addr=%p len=%zu\n", + i, buffers_[i].start, buffers_[i].length); + } + return true; +} + +bool V4l2Capture::startStreaming() { + if (fd_ < 0 || streaming_) return false; + + // 把所有缓冲区入队 + for (uint32_t i = 0; i < buffers_.size(); ++i) { + struct v4l2_buffer buf; + std::memset(&buf, 0, sizeof(buf)); + buf.type = V4L2_BUF_TYPE_VIDEO_CAPTURE; + buf.memory = V4L2_MEMORY_MMAP; + buf.index = i; + buf.bytesused = 0; // 驱动会填充 + + if (ioctl(fd_, VIDIOC_QBUF, &buf) < 0) { + printf("V4l2Capture: VIDIOC_QBUF %u failed: %s\n", i, strerror(errno)); + return false; + } + buffers_[i].queued = true; + } + + // 启动流 + int type = V4L2_BUF_TYPE_VIDEO_CAPTURE; + if (ioctl(fd_, VIDIOC_STREAMON, &type) < 0) { + printf("V4l2Capture: VIDIOC_STREAMON failed: %s\n", strerror(errno)); + return false; + } + streaming_ = true; + printf("V4l2Capture: streaming started\n"); + return true; +} + +bool V4l2Capture::stopStreaming() { + if (fd_ < 0 || !streaming_) return false; + + int type = V4L2_BUF_TYPE_VIDEO_CAPTURE; + if (ioctl(fd_, VIDIOC_STREAMOFF, &type) < 0) { + printf("V4l2Capture: VIDIOC_STREAMOFF failed: %s\n", strerror(errno)); + return false; + } + streaming_ = false; + last_dequeued_ = -1; + printf("V4l2Capture: streaming stopped\n"); + return true; +} + +bool V4l2Capture::getFrame(Frame &frame, int timeout_ms) { + if (fd_ < 0 || !streaming_) return false; + + fd_set fds; + FD_ZERO(&fds); + FD_SET(fd_, &fds); + + struct timeval tv; + tv.tv_sec = timeout_ms / 1000; + tv.tv_usec = (timeout_ms % 1000) * 1000; + + int r = select(fd_ + 1, &fds, nullptr, nullptr, &tv); + if (r < 0) { + if (errno == EINTR) return getFrame(frame, timeout_ms); // 重试 + printf("V4l2Capture: select failed: %s\n", strerror(errno)); + return false; + } + if (r == 0) { + // 超时 + return false; + } + + struct v4l2_buffer buf; + std::memset(&buf, 0, sizeof(buf)); + buf.type = V4L2_BUF_TYPE_VIDEO_CAPTURE; + buf.memory = V4L2_MEMORY_MMAP; + + if (ioctl(fd_, VIDIOC_DQBUF, &buf) < 0) { + if (errno == EAGAIN || errno == EIO) { + return false; // 再试 + } + printf("V4l2Capture: VIDIOC_DQBUF failed: %s\n", strerror(errno)); + return false; + } + + // 如果之前有未归还的缓冲区, 先归还 + if (last_dequeued_ >= 0) { + releaseFrameByIndex(last_dequeued_); + } + + uint32_t idx = buf.index; + if (idx >= buffers_.size()) { + printf("V4l2Capture: invalid buffer index %u\n", idx); + return false; + } + + last_dequeued_ = idx; + if ((buf.flags & V4L2_BUF_FLAG_ERROR) || buf.bytesused > buffers_[idx].length) { + releaseFrameByIndex(idx); + last_dequeued_ = -1; + return false; + } + + frame.data = buffers_[idx].start; + frame.size = buf.bytesused; + frame.width = width_; + frame.height = height_; + frame.pixelformat = pixelformat_; + frame.bytes_per_line = bytes_per_line_; + frame.sequence = buf.sequence; + frame.timestamp_us = (uint64_t)buf.timestamp.tv_sec * 1000000ULL + buf.timestamp.tv_usec; + frame.index = idx; + + return true; +} + +void V4l2Capture::releaseFrame(const Frame &frame) { + if (frame.index >= 0) { + releaseFrameByIndex(frame.index); + last_dequeued_ = -1; + } +} + +void V4l2Capture::releaseFrameByIndex(int idx) { + if (idx < 0 || (uint32_t)idx >= buffers_.size()) return; + struct v4l2_buffer buf; + std::memset(&buf, 0, sizeof(buf)); + buf.type = V4L2_BUF_TYPE_VIDEO_CAPTURE; + buf.memory = V4L2_MEMORY_MMAP; + buf.index = idx; + if (ioctl(fd_, VIDIOC_QBUF, &buf) < 0) { + printf("V4l2Capture: VIDIOC_QBUF %d failed: %s\n", idx, strerror(errno)); + } + buffers_[idx].queued = true; +} diff --git a/v4l2_capture.h b/v4l2_capture.h new file mode 100644 index 0000000..5dfc2b6 --- /dev/null +++ b/v4l2_capture.h @@ -0,0 +1,91 @@ +#ifndef V4L2_CAPTURE_H +#define V4L2_CAPTURE_H + +#include +#include +#include +#include + +// V4L2 流式(MMAP)取图封装 +// 参考 quick_bring_up_64g.sh 中 v4l2src 的行为: +// - 打开 /dev/videoN +// - 协商 UYVY 像素格式与分辨率 +// - 设置相机控制(sensor_mode / trig_pin 等, 等价于 v4l2-ctl -c ...) +// - MMAP 方式申请内核缓冲区并流式采集 +class V4l2Capture { +public: + // 单帧采集结果 + struct Frame { + void *data = nullptr; // 指向 MMAP 缓冲区, 释放前有效 + size_t size = 0; // 实际帧字节数(bytesused) + uint32_t width = 0; + uint32_t height = 0; + uint32_t pixelformat = 0; // V4L2_PIX_FMT_* (如 V4L2_PIX_FMT_UYVY) + uint32_t bytes_per_line = 0; + uint32_t sequence = 0; // 驱动序号 + uint64_t timestamp_us = 0; // 单调时钟时间戳(us) + int index = -1; // 内部缓冲区索引 + }; + + V4l2Capture() = default; + ~V4l2Capture(); + + V4l2Capture(const V4l2Capture &) = delete; + V4l2Capture &operator=(const V4l2Capture &) = delete; + + // 打开设备节点, 例如 /dev/video5 + bool open(const std::string &device); + + // 关闭设备并释放所有资源 + void close(); + + bool isOpen() const { return fd_ >= 0; } + + // 协商像素格式与分辨率. pixelformat 使用 V4L2_PIX_FMT_* 四字符码 + bool setFormat(uint32_t width, uint32_t height, uint32_t pixelformat); + + // 设置帧率(通过 VIDIOC_S_PARM). fps 分子/分母均为 fps/1 + bool setFps(uint32_t fps); + + // 设置 V4L2 控制, 等价 v4l2-ctl -c =. id 为 V4L2_CID_* 之四字符码 + // 例如 sensor_mode=0x009a2008, trig_pin=0x009a2084 + bool setControl(uint32_t id, int32_t value); + + // 申请 MMAP 缓冲区, 默认 4 个 + bool initMmap(uint32_t buffer_count = 4); + + // 开始 / 停止采集 + bool startStreaming(); + bool stopStreaming(); + + // 阻塞取一帧, 超时 timeout_ms 毫秒. 成功返回 true, 需配对调用 releaseFrame() + bool getFrame(Frame &frame, int timeout_ms = 1000); + + // 归还缓冲区给驱动队列, 必须在下次 getFrame 前调用 + void releaseFrame(const Frame &frame); + + uint32_t width() const { return width_; } + uint32_t height() const { return height_; } + uint32_t pixelformat() const { return pixelformat_; } + +private: + struct MmapBuffer { + void *start = nullptr; + size_t length = 0; + bool queued = false; // 是否在驱动队列中 + }; + + // 归还指定索引的缓冲区 (内部使用) + void releaseFrameByIndex(int idx); + + int fd_ = -1; + uint32_t width_ = 0; + uint32_t height_ = 0; + uint32_t pixelformat_ = 0; + uint32_t bytes_per_line_ = 0; + bool streaming_ = false; + std::vector buffers_; + int last_dequeued_ = -1; // 当前已出队但未归还的缓冲区索引 +}; + +#endif // V4L2_CAPTURE_H