TractorVision初期移植

lib/alg/track_detection.py

@@ -0,0 +1,173 @@
+# 文件: tractor_vision/alg/track_detection.py
+import cv2
+import numpy as np
+import math
+from collections import deque
+
+# ——— 相机标定参数 ———
+FX = 474.0  # X轴焦距
+FY = 505.0  # Y轴焦距
+CX = 320.0  # CX 中心x轴中标
+CY = 240.0  # CY 中心y轴座标
+CAMERA_TILT_DEG = 10.0  # 向下为正
+
+# ——— ROI & 阈值 ———
+ROI = (190, 40, 250, 420)
+MIN_DEPTH = 1100
+ANGLE_TOL_DEG = 45
+
+
+def preprocess_depth(depth, blur_ksize=5, kernel_size=7):
+    """
+    - 中值滤波消噪
+    - 形态学闭运算填补空洞
+    """
+    d = cv2.medianBlur(depth, blur_ksize)
+    kernel = np.ones((kernel_size, kernel_size), np.uint8)
+    closed = cv2.morphologyEx(d, cv2.MORPH_CLOSE, kernel)
+    d[d == 0] = closed[d == 0]
+    return d
+
+
+def offset_after_rotation(line):
+    (x1, y1), (x2, y2) = line
+    phi = math.degrees(math.atan2(y2 - y1, x2 - x1))
+    rot = math.radians(90.0 - phi)
+    mx, my = (x1 + x2) / 2, (y1 + y2) / 2
+    dx, dy = mx - CX, my - CY
+    x_rot = dx * math.cos(rot) - dy * math.sin(rot)
+    return -x_rot, (int(mx), int(my))
+
+
+def compute_metrics(depth):
+    # NOTE: 先强行对数据进行转换，保持数据为uint16的副本
+    depth_mm = (
+        depth.copy().astype(np.uint16) if depth.dtype != np.uint16 else depth.copy()
+    )
+    # NOTE: 中值滤波消噪
+    den = preprocess_depth(depth_mm)
+    # NOTE: 裁剪出感兴趣的ROI区域进行分析
+    x, y, w, h = ROI
+    crop = den[y : y + h, x : x + w]
+    if not (crop >= MIN_DEPTH).any():
+        return None, None, None, None, depth_mm, ROI, None
+
+    # NOTE: 将深度裁剪区域 crop 归一化到0~255区间，并转换为 uint8 类型，方便后续边缘检测。
+    # 归一化是为了适配 OpenCV 的边缘检测算法对8位图像的需求。
+    norm = cv2.normalize(crop, None, 0, 255, cv2.NORM_MINMAX).astype("uint8")
+    # NOTE: 对归一化后的图像执行 Canny 边缘检测，提取边缘轮廓。
+    # 50和150是边缘检测的低高阈值。
+    edges = cv2.Canny(norm, 50, 150)
+    lines = cv2.HoughLinesP(edges, 1, np.pi / 180, 30, minLineLength=50, maxLineGap=10)
+    # NOTE: 如果没检测边缘直线，则返回空
+    if lines is None:
+        return None, None, None, None, depth_mm, ROI, None
+    # NOTE: 将摄像机俯仰角（单位度）转为弧度，后续角度计算时用。
+    tilt = math.radians(CAMERA_TILT_DEG)
+    # NOTE:初始化空列表，存放符合角度条件的候选轨道线。
+    cands = []
+    # NOTE: 遍历霍夫检测到的所有直线，每条线由两个端点坐标 (x1_, y1_) 和 (x2_, y2_) 表示。
+    for x1_, y1_, x2_, y2_ in lines.reshape(-1, 4):
+        # NOTE:将裁剪区域的坐标偏移加回，转换为深度图整体坐标系的位置。
+        # 因为霍夫检测是在裁剪区域内，需补偿ROI偏移。
+        p1 = (x1_ + x, y1_ + y)
+        p2 = (x2_ + x, y2_ + y)
+        # NOTE: 计算线段水平和垂直方向的向量分量
+        dx, dy = p2[0] - p1[0], p2[1] - p1[1]
+        # NOTE: 计算线段与垂直方向的角度（弧度)
+        raw = -math.atan2(dx, -dy)
+        # NOTE: 考虑摄像头俯仰角 tilt，修正线段的角度。把线段角度投影到水平面，得到真实的轨道角度。
+        grd = math.asin(math.sin(raw) * math.cos(tilt))
+        # NOTE: 角度转换成度数，并取绝对值。
+        ang = abs(math.degrees(grd))
+        # NOTE: 如果角度小于预设阈值 ANGLE_TOL_DEG（即近似垂直轨道的线），认为是候选轨道线，添加到列表。
+        if ang <= ANGLE_TOL_DEG:
+            cands.append({"line": (p1, p2), "mid_x": (p1[0] + p2[0]) / 2})
+        # NOTE: 如果没找到符合角度的候选线，返回 None。
+    if not cands:
+        return None, None, None, None, depth_mm, ROI, None
+    # NOTE: 选择最左边的轨道线作为“最佳”轨道线，依据中点横坐标最小。 
+    best = min(cands, key=lambda c: c["mid_x"])
+    # NOTE: 拆包出最佳轨道线的两个端点坐标。
+    line = best["line"]
+    (x1, y1), (x2, y2) = line
+
+    # 交点
+    dxl, dyl = x2 - x1, y2 - y1
+    # NOTE: 计算轨道线与摄像头中心线（垂直线 x=CX）交点坐标。
+    # 如果轨道线不是垂直的，则通过线性插值求出交点 y 坐标。
+    # 如果轨道线垂直（dxl=0），交点为中点 y 坐标。
+    if dxl != 0:
+        t = (CX - x1) / dxl
+        y_int = y1 + t * dyl
+        inter = (int(CX), int(y_int))
+    else:
+        inter = (int(CX), int((y1 + y2) / 2))
+
+    # NOTE: 通过各种公式计算出角度和偏移亮
+    p_up = (x1, y1) if y1 < y2 else (x2, y2)
+    vx, vy = p_up[0] - inter[0], p_up[1] - inter[1]
+    raw2 = math.atan2(vx, -vy)
+    grd2 = math.asin(math.sin(raw2) * math.cos(tilt))
+    signed = math.degrees(grd2)
+
+    off_px, midpt = offset_after_rotation(line)
+    pts = [
+        np.clip(pt, [0, 0], [depth_mm.shape[1] - 1, depth_mm.shape[0] - 1])
+        for pt in line
+    ]
+    Zs = [float(depth_mm[pt[1], pt[0]]) for pt in pts]
+    Zavg = np.mean([z for z in Zs if z > 0]) if any(z > 0 for z in Zs) else MIN_DEPTH
+    off_mm = off_px * Zavg / FX * math.cos(tilt)
+
+    return signed, off_mm, line, midpt, depth_mm, ROI, inter
+
+
+class TrackDetector:
+    def __init__(self, camera_list, alpha=0.6, win_raw=5, win_smooth=5, interval=0.2):
+        self.alpha, self.interval = alpha, interval
+        self.ang_raw = {sn: deque(maxlen=win_raw) for sn in camera_list}
+        self.off_raw = {sn: deque(maxlen=win_raw) for sn in camera_list}
+        self.ang_smooth = {sn: deque(maxlen=win_smooth) for sn in camera_list}
+        self.off_smooth = {sn: deque(maxlen=win_smooth) for sn in camera_list}
+        self.last_ang = {sn: None for sn in camera_list}
+        self.last_off = {sn: None for sn in camera_list}
+        self.last_t = {sn: 0.0 for sn in camera_list}
+
+    def process(self, depth, sn):
+        # TODO: 先进行对当前时钟判断，如果小于interval则返回空
+        now = cv2.getTickCount() / cv2.getTickFrequency()
+        if now - self.last_t[sn] < self.interval:
+            return None
+
+        self.last_t[sn] = now
+        res = compute_metrics(depth)
+        if res[0] is None:
+            return None
+        ang, off, line, midpt, _, _, inter = res
+
+        # TODO: 平滑
+        self.ang_raw[sn].append(ang)
+        self.off_raw[sn].append(off)
+        med_ang = float(np.median(self.ang_raw[sn]))
+        med_off = float(np.median(self.off_raw[sn]))
+        pa, po = self.last_ang[sn], self.last_off[sn]
+        ang_exp = (
+            med_ang if pa is None else self.alpha * med_ang + (1 - self.alpha) * pa
+        )
+        off_exp = (
+            med_off if po is None else self.alpha * med_off + (1 - self.alpha) * po
+        )
+        self.last_ang[sn], self.last_off[sn] = ang_exp, off_exp
+        self.ang_smooth[sn].append(ang_exp)
+        self.off_smooth[sn].append(off_exp)
+        disp_ang = float(np.median(self.ang_smooth[sn]))
+        disp_off = float(np.median(self.off_smooth[sn]))
+
+        return {
+            "angle": -disp_ang,  # 左正右负
+            "offset": disp_off,
+            "line": line,
+            "midpoint": midpt,
+            "intersection": inter,
+        }