TractorVision初期移植

lib/alg/distance_respect_tof.py

@@ -0,0 +1,55 @@
+# distance_detector.py
+import numpy as np
+
+# 每个相机SN对应的ROI配置（中心坐标和基准尺寸）
+# 画框测距
+ROI_SETTINGS = {
+    '233800055': {'center': (220, 340), 'base_size': (200, 200)},
+    '234000035': {'center': (220, 340), 'base_size': (200, 200)},
+    # 可根据需要添加更多相机
+}
+# 默认基准深度（毫米）
+DEFAULT_BASE_DEPTH_MM = 2000
+
+class DistanceDetector:
+    """
+    基于 TOF 深度图的动态 ROI 距离检测器，ROI 参数内置到算法中。
+
+    算法流程：
+      1. 使用相机 SN 对应的预设 ROI 中心和大小。
+      2. 在像素中心读取深度值 d。
+      3. 根据基准深度与当前深度的比例动态计算 ROI 大小。
+      4. 在 ROI 内提取非零深度值并计算平均值作为目标距离（毫米）。
+    """
+    def __init__(self, sn, base_depth_mm=DEFAULT_BASE_DEPTH_MM):
+        if sn not in ROI_SETTINGS:
+            raise ValueError(f"No ROI settings for camera SN: {sn}")
+        cfg = ROI_SETTINGS[sn]
+        self.cx, self.cy = cfg['center']
+        self.base_w, self.base_h = cfg['base_size']
+        self.base_depth = base_depth_mm
+
+    def _compute_dynamic_roi(self, depth_image):
+        h, w = depth_image.shape
+        d = float(depth_image[int(self.cy), int(self.cx)])
+        if d <= 0:
+            return None
+        scale = self.base_depth / d
+        rw = int(self.base_w * scale)
+        rh = int(self.base_h * scale)
+        x1 = max(0, int(self.cx - rw // 2))
+        y1 = max(0, int(self.cy - rh // 2))
+        x2 = min(w, x1 + rw)
+        y2 = min(h, y1 + rh)
+        return x1, y1, x2, y2
+
+    def compute_distance(self, depth_image):
+        coords = self._compute_dynamic_roi(depth_image)
+        if coords is None:
+            return None
+        x1, y1, x2, y2 = coords
+        roi = depth_image[y1:y2, x1:x2]
+        valid = roi[roi > 0]
+        if valid.size == 0:
+            return None
+        return float(np.mean(valid))