import cv2
import numpy as np

def detect_arrow_direction(image):
    """
    从图像中提取绿色箭头并判断其指向方向（左或右）
    
    参数:
        image: 输入图像，可以是文件路径或者已加载的图像数组
        
    返回:
        direction: 字符串，"left"表示左箭头，"right"表示右箭头，"unknown"表示无法确定
    """
    # 如果输入是字符串（文件路径），则加载图像
    if isinstance(image, str):
        img = cv2.imread(image)
    else:
        img = image.copy()
    
    if img is None:
        print("无法加载图像")
        return "unknown"
    
    # 转换到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)
    
    # 将掩码转为灰度图
    gray = mask.copy()
    
    # 查找轮廓
    contours, _ = cv2.findContours(gray, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    
    # 如果没有找到轮廓，返回未知
    if not contours:
        return "unknown"
    
    # 找到最大的轮廓（假设是箭头）
    max_contour = max(contours, key=cv2.contourArea)
    
    # 获取轮廓的最小外接矩形
    x, y, w, h = cv2.boundingRect(max_contour)
    
    # 计算轮廓的矩，用于确定箭头方向
    M = cv2.moments(max_contour)
    
    # 避免除以零
    if M["m00"] != 0:
        cx = int(M["m10"] / M["m00"])
        cy = int(M["m01"] / M["m00"])
    else:
        cx, cy = x + w//2, y + h//2
    
    # 将图像分为左右两部分
    left_region = gray[y:y+h, x:cx]
    right_region = gray[y:y+h, cx:x+w]
    
    # 计算每个区域中白色像素的数量（箭头部分）
    left_pixels = cv2.countNonZero(left_region)
    right_pixels = cv2.countNonZero(right_region)
    
    # 计算每个区域中白色像素的密度
    left_density = left_pixels / (left_region.size + 1e-10)
    right_density = right_pixels / (right_region.size + 1e-10)
    
    # 根据左右区域的像素密度确定箭头方向
    # 如果箭头指向右侧，右侧区域的箭头头部应该有更多的像素密度
    # 如果箭头指向左侧，左侧区域的箭头头部应该有更多的像素密度
    if right_density > left_density:
        return "right"
    else:
        return "left"

def visualize_arrow_detection(image, save_path=None):
    """
    可视化箭头检测过程，显示中间结果
    
    参数:
        image: 输入图像，可以是文件路径或者已加载的图像数组
        save_path: 保存结果图像的路径（可选）
    """
    # 如果输入是字符串（文件路径），则加载图像
    if isinstance(image, str):
        img = cv2.imread(image)
    else:
        img = image.copy()
    
    if img is None:
        print("无法加载图像")
        return
    
    # 转换到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)
        
        # 获取轮廓的最小外接矩形
        x, y, w, h = cv2.boundingRect(max_contour)
        cv2.rectangle(output, (x, y), (x + w, y + h), (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)
            
            # 绘制分割线
            cv2.line(output, (cx, y), (cx, y + h), (0, 255, 255), 2)
    
    # 获取箭头方向
    direction = detect_arrow_direction(img)
    
    # 在图像上添加方向文本
    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)
    
    # 创建一个包含所有图像的窗口
    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"
    
    # 检测箭头方向
    direction = detect_arrow_direction(image_path)
    print(f"检测到的箭头方向: {direction}")
    
    # 可视化检测过程
    visualize_arrow_detection(image_path)