cmvr-es/python/vision_servo/target_position.py

import cv2
import numpy as np
import pyrealsense2 as rs

def init_realsense():
    """
    初始化 RealSense 管道并返回 pipeline 和 align 对象
    """
    pipeline = rs.pipeline()
    config = rs.config()
    config.enable_stream(rs.stream.depth, 848, 480, rs.format.z16, 30)
    config.enable_stream(rs.stream.color, 1280, 720, rs.format.bgr8, 30)
    profile = pipeline.start(config)
    align = rs.align(rs.stream.color)
    return pipeline, align


def get_closest_red_point(pipeline, align, vis_path = None,warmup=10):
    """
    获取最近红色圆点的 3D 坐标。
    返回 (x, y, z) 或 None（如果没检测到）。
    """
    # 丢掉前几帧，让相机稳定
    try:
        for _ in range(warmup):
            pipeline.wait_for_frames()

        frames = pipeline.wait_for_frames()
        aligned_frames = align.process(frames)

        depth_frame = aligned_frames.get_depth_frame()
        color_frame = aligned_frames.get_color_frame()
        if not depth_frame or not color_frame:
            return None

        depth_image = np.asanyarray(depth_frame.get_data())
        color_image = np.asanyarray(color_frame.get_data())
        depth_intrin = depth_frame.profile.as_video_stream_profile().intrinsics

        hsv = cv2.cvtColor(color_image, cv2.COLOR_BGR2HSV)
        lower_red1 = np.array([0, 100, 100])
        upper_red1 = np.array([10, 255, 255])
        lower_red2 = np.array([160, 100, 100])
        upper_red2 = np.array([179, 255, 255])
        mask1 = cv2.inRange(hsv, lower_red1, upper_red1)
        mask2 = cv2.inRange(hsv, lower_red2, upper_red2)
        mask = cv2.bitwise_or(mask1, mask2)

        mask_blur = cv2.GaussianBlur(mask, (9, 9), 2)
        circles = cv2.HoughCircles(mask_blur, cv2.HOUGH_GRADIENT, dp=1.2, minDist=20,
                                   param1=50, param2=15, minRadius=5, maxRadius=50)

        closest_point = None
        min_depth = float('inf')

        if circles is not None:
            circles = np.uint16(np.around(circles))
            for i in circles[0, :]:
                u, v, r = i
                depth = depth_frame.get_distance(u, v)
                if 0 < depth < min_depth:
                    min_depth = depth
                    closest_point = rs.rs2_deproject_pixel_to_point(depth_intrin, [u, v], depth)

        if vis_path is not None and closest_point is not None:
            cv2.circle(color_image, (u, v), 6, (0,0,255), -1)
            cv2.imwrite(vis_path, color_image)
        return closest_point
    except Exception as e:
        print(e)
        return None


def black_point_position(color_frame, depth_frame, depth_intr, vis_path=None, min_radius=10, avg_window=3):
    """
    检测白底黑圆的圆心坐标，返回相机系 (X, Y, Z) [m]

    参数
    ----
    color_frame / depth_frame : 对齐后的 RealSense frame
    depth_intr               : 深度流 intrinsics (rs.intrinsics)
    vis_path                 : 若给定，则保存可视化图片
    min_radius               : HoughCircles/轮廓的最小半径，像素
    avg_window               : 深度均值窗口半径（像素）
    """
    color_img = np.asanyarray(color_frame.get_data()).copy()
    gray = cv2.cvtColor(color_img, cv2.COLOR_BGR2GRAY)

    # 1. 二值化（寻找黑色区域）
    # Otsu 自动阈值 + 取反 => 黑圆为白，背景为黑
    _, mask = cv2.threshold(gray, 0, 255, cv2.THRESH_BINARY_INV + cv2.THRESH_OTSU)

    # 2. 轮廓检测，取面积最大的圆形候选
    cnts, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    if not cnts:
        raise RuntimeError("未检测到任何黑色区域")

    # 取最大面积
    cnt = max(cnts, key=cv2.contourArea)
    (u, v), radius = cv2.minEnclosingCircle(cnt)

    if radius < min_radius:
        raise RuntimeError(f"检测到圆半径过小 ({radius:.1f}px)，请检查 min_radius 设置或图像质量")

    # 3. 取邻域深度中值
    width, height = depth_frame.get_width(), depth_frame.get_height()
    depths = [
        depth_frame.get_distance(int(round(u + du)), int(round(v + dv)))
        for du in range(-avg_window, avg_window + 1)
        for dv in range(-avg_window, avg_window + 1)
        if 0 <= int(round(u + du)) < width and
           0 <= int(round(v + dv)) < height
    ]
    depths = [d for d in depths if d > 0]
    if not depths:
        raise RuntimeError("圆心处深度无效 (0)")

    depth = float(np.median(depths))

    # 4. 像素 -> 相机坐标
    x, y, z = rs.rs2_deproject_pixel_to_point(
        depth_intr, [u, v], depth
    )
    pos = np.array([x, y, z], dtype=np.float32)

    # 5. 可视化
    if vis_path is not None:
        cx, cy = depth_intr.ppx, depth_intr.ppy
        cv2.drawMarker(color_img, (int(cx), int(cy)), (0, 255, 0), markerType=cv2.MARKER_CROSS, markerSize=20, thickness=2)
        cv2.circle(color_img, (int(u), int(v)), int(radius), (0, 0, 255), 2)
        cv2.circle(color_img, (int(u), int(v)), 5, (0, 0, 255), -1)
        cv2.imwrite(vis_path, color_img)
        print(f"✓ 已保存标记图到 {vis_path}")

    return pos


def red_point_position(color_frame, depth_frame, depth_intr, vis_path=None, min_radius=3, avg_window=3):
    """
    检测白底红圆的圆心坐标，返回相机系 (X,Y,Z) [m]
    """
    # --- 1. 取彩色帧 ----
    color_img = np.asanyarray(color_frame.get_data()).copy()
    hsv = cv2.cvtColor(color_img, cv2.COLOR_BGR2HSV)

    # --- 2. 阈值分割：红色有两个 Hue 区间 (0-10)∪(170-180) ---
    lower_red1 = np.array([0,   100, 100])
    upper_red1 = np.array([10,  255, 255])
    lower_red2 = np.array([170, 100, 100])
    upper_red2 = np.array([180, 255, 255])

    mask1 = cv2.inRange(hsv, lower_red1, upper_red1)
    mask2 = cv2.inRange(hsv, lower_red2, upper_red2)
    mask  = cv2.bitwise_or(mask1, mask2)

    # 可选：形态学开闭运算去噪
    kernel = np.ones((5,5), np.uint8)
    mask = cv2.morphologyEx(mask, cv2.MORPH_OPEN, kernel)
    mask = cv2.morphologyEx(mask, cv2.MORPH_CLOSE, kernel)

    # --- 3. 轮廓取最大圆 ---
    cnts, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    if not cnts:
        cv2.imwrite(vis_path, color_img)
        raise RuntimeError("未检测到任何红色区域")

    cnt = max(cnts, key=cv2.contourArea)
    (u, v), radius = cv2.minEnclosingCircle(cnt)
    if radius < min_radius:
        raise RuntimeError(f"检测到圆半径过小 ({radius:.1f}px)，请检查 min_radius 或图像质量")

    # --- 4. 深度中值 ---
    width, height = depth_frame.get_width(), depth_frame.get_height()
    depths = [
        depth_frame.get_distance(int(round(u + du)), int(round(v + dv)))
        for du in range(-avg_window, avg_window + 1)
        for dv in range(-avg_window, avg_window + 1)
        if 0 <= int(round(u + du)) < width and
           0 <= int(round(v + dv)) < height
    ]
    depths = [d for d in depths if d > 0]
    if not depths:
        raise RuntimeError("圆心处深度无效 (0)")
    depth = float(np.median(depths))

    # --- 5. 反投影到 3-D ---
    x, y, z = rs.rs2_deproject_pixel_to_point(depth_intr, [u, v], depth)
    pos = np.array([x, y, z], dtype=np.float32)

    # --- 6. 可视化保存 ---
    if vis_path is not None:
        cx, cy = depth_intr.ppx, depth_intr.ppy
        cv2.drawMarker(color_img, (int(cx), int(cy)), (0, 255, 0), markerType=cv2.MARKER_CROSS, markerSize=20, thickness=2)
        cv2.circle(color_img, (int(u), int(v)), int(radius), (255, 0, 0), 2)  # 蓝圈标红圆
        cv2.circle(color_img, (int(u), int(v)), 5, (255, 0, 0), -1)
        cv2.imwrite(vis_path, color_img)
        print(f"✓ 已保存标记图到 {vis_path}")

    return pos