orisys-sdk/tests/dd06.py

"""
DD06 single-camera pipeline with polygon ROI rectification.

Flow:
1. Load sensor config + nested polygon ROI
2. Capture raw frames from the video/camera source
3. Perspective-warp the ROI into a rectangle sized to camera.resolution_stitch
4. Feed the rectified frame into Sensor.get_img(..., already_processed=True)
5. Compute deformation/contact and visualize

SDK version: 0.4.4
"""
from __future__ import annotations

import argparse
import os
import sys
from pathlib import Path

os.environ["OPENCV_VIDEOIO_MSMF_ENABLE_HW_TRANSFORMS"] = "0"

import cv2
import numpy as np
import orisys

_REPO_ROOT = Path(__file__).resolve().parents[1]
_EXAMPLES_DIR = _REPO_ROOT / "examples"
if str(_EXAMPLES_DIR) not in sys.path:
    sys.path.insert(0, str(_EXAMPLES_DIR))

from gui.sensors.roi_geometry import (  # noqa: E402
    build_roi_warp,
    load_json_config,
    load_roi_polygon,
    open_video_source,
    rectified_output_size,
    warp_roi_to_rectangle,
)


def main():
    parser = argparse.ArgumentParser(
        description="DD06 pipeline: polygon-ROI rectify -> Sensor.get_img(already_processed=True)"
    )
    parser.add_argument(
        "--video",
        "-v",
        type=str,
        default="./data/dd06/0726.mp4",
        help="视频源：摄像头编号或视频文件路径",
    )
    parser.add_argument(
        "--config",
        "-c",
        type=str,
        default="./config/dd06-ov.json",
        help="配置文件名称或路径",
    )
    parser.add_argument(
        "--verbose",
        "-verbose",
        action="store_true",
        default=True,
        help="是否输出详细日志",
    )
    args = parser.parse_args()

    config = load_json_config(args.config)
    roi_points, roi_image_size = load_roi_polygon(config)
    out_w, out_h = rectified_output_size(config, roi_points)

    sensor = orisys.Sensor("external", config_name=args.config, verbose=args.verbose)
    cap, source = open_video_source(args.video)

    warp_matrix = None
    warp_src = None
    last_frame_size = None

    print(
        f"ROI points={len(roi_points)} source_size={roi_image_size} "
        f"output={out_w}x{out_h} video={source}"
    )
    print("按键说明：'q' 退出程序，'r' 重置追踪器。")

    try:
        while True:
            ok, frame = cap.read()
            if not ok or frame is None:
                print("Failed to read frame from video source.")
                break

            frame_h, frame_w = frame.shape[:2]
            frame_size = (frame_w, frame_h)
            if warp_matrix is None or last_frame_size != frame_size:
                warp_matrix, warp_src, _ = build_roi_warp(
                    roi_points,
                    roi_image_size,
                    frame_size,
                    (out_w, out_h),
                )
                last_frame_size = frame_size
                if args.verbose:
                    print(f"Built ROI warp for frame {frame_w}x{frame_h}:\n{warp_src}")

            rectified = warp_roi_to_rectangle(frame, warp_matrix, (out_w, out_h))
            sensor.get_img(frame=rectified, already_processed=True)

            is_contact_now = sensor.compute_deformation(check_motion=True, threshold=0)
            _ = is_contact_now
            sensor.compute_contact()

            fps, flow, vnormal, img, contour, centroid, depth_map, fnormal, fshearx, fsheary = (
                sensor.read_info(
                    sensor.info.FPS,
                    sensor.info.VRAW,
                    sensor.info.VNORMAL,
                    sensor.info.IMG,
                    sensor.info.CONTOUR,
                    sensor.info.CENTROID,
                    sensor.info.DEPTH,
                    sensor.info.FNORMAL,
                    sensor.info.FSHEARX,
                    sensor.info.FSHEARY,
                )
            )

            preview = frame.copy()
            cv2.polylines(
                preview,
                [np.rint(warp_src).astype(np.int32).reshape(-1, 1, 2)],
                True,
                (0, 255, 255),
                2,
                cv2.LINE_AA,
            )

            arrows = orisys.util.draw_arrows(
                img,
                flow,
                threshold=1,
                grid_spacing=20,
                arrow_scale=1.0,
            )
            cv2.imshow("形变矢量场", arrows)

            div_abs = depth_map
            if div_abs.max() > div_abs.min():
                div_normalized = (
                    (div_abs - div_abs.min()) / (div_abs.max() - div_abs.min()) * 255
                ).astype(np.uint8)
            else:
                div_normalized = np.zeros_like(div_abs, dtype=np.uint8)
            cv2.imshow("depth", cv2.applyColorMap(div_normalized, cv2.COLORMAP_JET))

            print(
                f"FPS={fps:.2f}, 法向力={fnormal:.4f}, 切向力X={fshearx:.4f}, "
                f"切向力Y={fsheary:.4f}, 深度图尺寸={depth_map.shape}, "
                f"原始分辨率={sensor.img_size_raw}, 矫正={rectified.shape[1]}x{rectified.shape[0]}"
            )

            key = cv2.waitKey(1) & 0xFF
            if key == ord("q"):
                print("\n正在退出示例程序...")
                break
            if key == ord("r"):
                sensor.reset()
                print("追踪器已重置。")
    finally:
        cap.release()
        sensor.disconnect()
        cv2.destroyAllWindows()


if __name__ == "__main__":
    main()