195 lines
7.3 KiB
Python
195 lines
7.3 KiB
Python
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（如果没检测到）。
|
||
"""
|
||
# 丢掉前几帧，让相机稳定
|
||
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
|
||
|
||
|
||
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 |