mi-task/utils/decode_arrow.py

import cv2
import numpy as np

def detect_arrow_direction(image, observe=False, delay=500):
    """
    从图像中提取绿色箭头并判断其指向方向（左或右）
    
    参数:
        image: 输入图像，可以是文件路径或者已加载的图像数组
        observe: 是否输出中间状态信息和可视化结果，默认为False
        delay: 展示每个步骤的等待时间(毫秒)，默认为500ms
        
    返回:
        direction: 字符串，"left"表示左箭头，"right"表示右箭头，"unknown"表示无法确定
    """
    # 如果输入是字符串（文件路径），则加载图像
    if isinstance(image, str):
        img = cv2.imread(image)
    else:
        img = image.copy()
    
    if img is None:
        print("无法加载图像")
        return "unknown"
    
    if observe:
        print("步骤1: 原始图像已加载")
        cv2.imshow("原始图像", img)
        cv2.waitKey(delay)
    
    # 转换到HSV颜色空间以便更容易提取绿色
    hsv = cv2.cvtColor(img, cv2.COLOR_BGR2HSV)
    
    if observe:
        print("步骤2: 转换到HSV颜色空间")
        cv2.imshow("HSV图像", hsv)
        cv2.waitKey(delay)
    
    # 绿色的HSV范围
    # 调整这些值以匹配图像中绿色的具体色调··
    lower_green = np.array([40, 50, 50])
    upper_green = np.array([80, 255, 255])
    
    # 创建绿色的掩码
    mask = cv2.inRange(hsv, lower_green, upper_green)
    
    if observe:
        print("步骤3: 创建绿色掩码")
        cv2.imshow("绿色掩码", mask)
        cv2.waitKey(delay)
    
    # 应用掩码，只保留绿色部分
    green_only = cv2.bitwise_and(img, img, mask=mask)
    
    if observe:
        print("步骤4: 提取绿色部分")
        cv2.imshow("只保留绿色", green_only)
        cv2.waitKey(delay)
    
    # 将掩码转为灰度图
    gray = mask.copy()
    
    # 查找轮廓
    contours, _ = cv2.findContours(gray, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    
    # 如果没有找到轮廓，返回未知
    if not contours:
        if observe:
            print("未找到轮廓")
        return "unknown"
    
    if observe:
        print(f"步骤5: 找到 {len(contours)} 个轮廓")
        contour_img = img.copy()
        cv2.drawContours(contour_img, contours, -1, (0, 255, 0), 2)
        cv2.imshow("所有轮廓", contour_img)
        cv2.waitKey(delay)
    
    # 找到最大的轮廓（假设是箭头）
    max_contour = max(contours, key=cv2.contourArea)
    
    if observe:
        print("步骤6: 提取最大轮廓")
        max_contour_img = img.copy()
        cv2.drawContours(max_contour_img, [max_contour], -1, (0, 0, 255), 2)
        cv2.imshow("最大轮廓", max_contour_img)
        cv2.waitKey(delay)
    
    # 使用多边形近似轮廓
    epsilon = 0.02 * cv2.arcLength(max_contour, True)
    approx = cv2.approxPolyDP(max_contour, epsilon, True)
    
    if observe:
        print(f"步骤7: 多边形近似，顶点数: {len(approx)}")
        approx_img = img.copy()
        cv2.drawContours(approx_img, [approx], -1, (255, 0, 0), 2)
        cv2.imshow("多边形近似", approx_img)
        cv2.waitKey(delay)
    
    # 计算凸包
    hull = cv2.convexHull(max_contour)
    
    if observe:
        print("步骤8: 计算凸包")
        hull_img = img.copy()
        cv2.drawContours(hull_img, [hull], -1, (0, 255, 0), 2)
        cv2.imshow("凸包", hull_img)
        cv2.waitKey(delay)
    
    # 计算凸缺陷
    if len(max_contour) > 3:
        defects = cv2.convexityDefects(max_contour, cv2.convexHull(max_contour, returnPoints=False))
    else:
        if observe:
            print("轮廓点太少，无法准确判断")
        return "unknown"  # 轮廓点太少，无法准确判断
    
    if defects is None:
        if observe:
            print("未找到凸缺陷")
        return "unknown"
    
    if observe:
        print(f"步骤9: 计算凸缺陷，找到 {defects.shape[0]} 个缺陷点")
        defect_img = img.copy()
        for i in range(defects.shape[0]):
            s, e, f, d = defects[i, 0]
            start = tuple(max_contour[s][0])
            end = tuple(max_contour[e][0])
            far = tuple(max_contour[f][0])
            cv2.circle(defect_img, far, 5, (0, 0, 255), -1)
        cv2.imshow("凸缺陷", defect_img)
        cv2.waitKey(delay)
        
    # 获取轮廓的最小外接矩形
    rect = cv2.minAreaRect(max_contour)
    box = cv2.boxPoints(rect)
    box = np.int0(box)
    
    if observe:
        print("步骤10: 获取最小外接矩形")
        rect_img = img.copy()
        cv2.drawContours(rect_img, [box], 0, (255, 0, 0), 2)
        cv2.imshow("最小外接矩形", rect_img)
        cv2.waitKey(delay)
    
    # 获取中心点和角度
    center = rect[0]
    angle = rect[2]
    
    if observe:
        print(f"矩形中心: {center}, 角度: {angle}")
    
    # 计算轮廓的矩，用于确定箭头方向
    M = cv2.moments(max_contour)
    
    # 避免除以零
    if M["m00"] != 0:
        cx = int(M["m10"] / M["m00"])
        cy = int(M["m01"] / M["m00"])
    else:
        cx, cy = int(center[0]), int(center[1])
    
    if observe:
        print(f"步骤11: 计算质心 - 坐标: ({cx}, {cy})")
        center_img = img.copy()
        cv2.circle(center_img, (cx, cy), 5, (255, 0, 0), -1)
        cv2.imshow("质心", center_img)
        cv2.waitKey(delay)
    
    # 改进的箭头尖端检测算法
    # 1. 找到所有凸缺陷点
    defect_points = []
    for i in range(defects.shape[0]):
        s, e, f, d = defects[i, 0]
        start = tuple(max_contour[s][0])
        end = tuple(max_contour[e][0])
        far = tuple(max_contour[f][0])
        
        # 计算缺陷点到中心的距离
        dist = np.sqrt((far[0] - cx) ** 2 + (far[1] - cy) ** 2)
        
        if observe:
            print(f"缺陷点 {i}: 位置 {far}, 到中心距离: {dist:.2f}")
        
        # 记录凸缺陷点及其距离
        defect_points.append({
            'point': far,
            'distance': dist,
            'start': start,
            'end': end
        })
    
    # 没有缺陷点，使用矩形判断
    if not defect_points:
        # 判断逻辑将在后面处理
        arrow_tip = None
    else:
        # 2. 按距离对缺陷点排序
        defect_points.sort(key=lambda x: x['distance'], reverse=True)
        
        # 3. 获取最远的几个缺陷点（可能的尖端候选）
        top_n = min(3, len(defect_points))
        candidates = defect_points[:top_n]
        
        # 4. 分析这些点的位置分布
        left_candidates = [p for p in candidates if p['point'][0] < cx]
        right_candidates = [p for p in candidates if p['point'][0] >= cx]
        
        # 5. 根据候选点的分布判断箭头朝向
        if len(left_candidates) > len(right_candidates):
            # 左侧候选点更多，箭头可能指向左边
            arrow_tip = max(left_candidates, key=lambda x: x['distance'])['point']
            arrow_direction = "left"
        elif len(right_candidates) > len(left_candidates):
            # 右侧候选点更多，箭头可能指向右边
            arrow_tip = max(right_candidates, key=lambda x: x['distance'])['point']
            arrow_direction = "right"
        else:
            # 候选点分布均衡，使用最远的点
            arrow_tip = candidates[0]['point']
            # 根据最远点的位置判断
            if arrow_tip[0] < cx:
                arrow_direction = "left"
            else:
                arrow_direction = "right"
    
    if observe and arrow_tip:
        print(f"步骤12: 找到可能的箭头尖端 - 位置: {arrow_tip}")
        tip_img = img.copy()
        cv2.circle(tip_img, arrow_tip, 8, (0, 255, 255), -1)
        cv2.line(tip_img, (cx, cy), arrow_tip, (255, 0, 255), 2)
        cv2.imshow("箭头尖端", tip_img)
        cv2.waitKey(delay)
    
    # 使用多种特征综合判断箭头方向
    if arrow_tip is None:
        # 如果没有找到明显的箭头尖端，使用最小外接矩形来判断
        width = rect[1][0]
        height = rect[1][1]
        
        if observe:
            print(f"未找到明显尖端，使用矩形判断 - 宽: {width}, 高: {height}")
        
        # 矩形的方向向量
        if width > height:  # 水平箭头
            # 考虑角度的定义
            if angle > 45:
                arrow_direction = "left"
            else:
                arrow_direction = "right"
        else:  # 垂直箭头，不在我们的考虑范围内
            if observe:
                print("垂直箭头，不在考虑范围")
            return "unknown"
    else:
        # 已经在上面的代码中确定了方向
        pass
    
    # 附加检查：计算轮廓边界上的点，确认箭头的形状特征
    x, y, w, h = cv2.boundingRect(max_contour)
    aspect_ratio = float(w) / h
    
    # 计算轮廓周长和面积
    perimeter = cv2.arcLength(max_contour, True)
    area = cv2.contourArea(max_contour)
    
    # 计算轮廓的形状复杂度
    complexity = perimeter / (4 * np.sqrt(area))
    
    if observe:
        print(f"轮廓分析 - 宽高比: {aspect_ratio:.2f}, 复杂度: {complexity:.2f}")
    
    # 箭头形状特征检查
    # 一般来说，箭头的复杂度会在一定范围内
    if 1.1 < complexity < 3.0:
        # 根据形状特征进一步确认方向判断
        # 使用质心位置相对于轮廓边界的偏移
        relative_cx = (cx - x) / w
        
        if observe:
            print(f"质心相对位置: {relative_cx:.2f}")
        
        # 如果质心偏向左侧，箭头可能指向右侧
        # 如果质心偏向右侧，箭头可能指向左侧
        # 这是一个启发式规则，可能需要根据具体箭头形状调整
        if relative_cx < 0.4:
            # 质心在左侧，可能是右箭头
            if arrow_direction == "left":
                # 当前方法判断为左，但质心位置特征表明可能是右
                # 增加额外的检查
                if aspect_ratio > 1.5:  # 宽大于高
                    arrow_direction = "right"
        elif relative_cx > 0.6:
            # 质心在右侧，可能是左箭头
            if arrow_direction == "right":
                # 当前方法判断为右，但质心位置特征表明可能是左
                if aspect_ratio > 1.5:  # 宽大于高
                    arrow_direction = "left"
    
    if observe:
        print(f"基于综合特征判断: {arrow_direction}箭头")
    
    return arrow_direction

def visualize_arrow_detection(image, save_path=None, observe=False, delay=500):
    """
    可视化箭头检测过程，显示中间结果
    
    参数:
        image: 输入图像，可以是文件路径或者已加载的图像数组
        save_path: 保存结果图像的路径（可选）
        observe: 是否输出中间状态信息和可视化结果，默认为False
        delay: 展示每个步骤的等待时间(毫秒)，默认为500ms
    """
    # 如果输入是字符串（文件路径），则加载图像
    if isinstance(image, str):
        img = cv2.imread(image)
    else:
        img = image.copy()
    
    if img is None:
        print("无法加载图像")
        return
    
    if observe:
        print("\n开始可视化箭头检测过程")
    
    # 转换到HSV颜色空间
    hsv = cv2.cvtColor(img, cv2.COLOR_BGR2HSV)
    
    # 绿色的HSV范围
    lower_green = np.array([40, 50, 50])
    upper_green = np.array([80, 255, 255])
    
    # 创建绿色的掩码
    mask = cv2.inRange(hsv, lower_green, upper_green)
    
    # 应用掩码，只保留绿色部分
    green_only = cv2.bitwise_and(img, img, mask=mask)
    
    # 查找轮廓
    contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    
    # 创建输出图像
    output = img.copy()
    
    # 如果找到轮廓，绘制最大轮廓
    if contours:
        max_contour = max(contours, key=cv2.contourArea)
        cv2.drawContours(output, [max_contour], -1, (0, 0, 255), 2)
        
        # 获取轮廓的最小外接矩形
        rect = cv2.minAreaRect(max_contour)
        box = cv2.boxPoints(rect)
        box = np.int0(box)
        cv2.drawContours(output, [box], 0, (255, 0, 0), 2)
        
        # 计算轮廓的矩
        M = cv2.moments(max_contour)
        
        # 避免除以零
        if M["m00"] != 0:
            cx = int(M["m10"] / M["m00"])
            cy = int(M["m01"] / M["m00"])
            
            # 绘制中心点
            cv2.circle(output, (cx, cy), 5, (255, 0, 0), -1)
        
        # 绘制凸包
        hull = cv2.convexHull(max_contour)
        cv2.drawContours(output, [hull], 0, (0, 255, 0), 2)
        
        # 绘制凸缺陷
        if len(max_contour) > 3:
            defects = cv2.convexityDefects(max_contour, cv2.convexHull(max_contour, returnPoints=False))
            if defects is not None:
                for i in range(defects.shape[0]):
                    s, e, f, d = defects[i, 0]
                    start = tuple(max_contour[s][0])
                    end = tuple(max_contour[e][0])
                    far = tuple(max_contour[f][0])
                    cv2.circle(output, far, 5, (0, 0, 255), -1)  # 在凸缺陷点绘制红色圆点
    
    # 获取箭头方向
    direction = detect_arrow_direction(img, observe=observe, delay=delay)
    
    # 在图像上添加方向文本
    cv2.putText(output, f"Direction: {direction}", (10, 30), 
                cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 2)
    
    # 如果提供了保存路径，保存结果图像
    if save_path:
        cv2.imwrite(save_path, output)
        if observe:
            print(f"结果已保存到: {save_path}")
    
    # 创建一个包含所有图像的窗口
    result = np.hstack((img, green_only, output))
    
    # 调整大小以便查看
    scale_percent = 50  # 缩放到原来的50%
    width = int(result.shape[1] * scale_percent / 100)
    height = int(result.shape[0] * scale_percent / 100)
    dim = (width, height)
    resized = cv2.resize(result, dim, interpolation=cv2.INTER_AREA)
    
    # 显示结果
    cv2.imshow('Arrow Detection Process', resized)
    cv2.waitKey(0)
    cv2.destroyAllWindows()

# 用法示例
if __name__ == "__main__":
    # 替换为实际图像路径
    image_path = "path/to/arrow/image.png"
    
    # 检测箭头方向，使用较长的延迟时间（1500毫秒）
    direction = detect_arrow_direction(image_path, observe=True, delay=1500)
    print(f"检测到的箭头方向: {direction}")
    
    # 可视化检测过程，使用较长的延迟时间
    visualize_arrow_detection(image_path, observe=True, delay=1500)