import cv2
import numpy as np
import os

def detect_yellow_distance_from_bottom(image_path, visualize=False):
    """
    检测鱼眼图像中垂线上最靠近下方的黄点到图像底部的距离
    
    参数:
        image_path: 图像路径
        visualize: 是否显示检测过程可视化结果
        
    返回:
        distance: 黄点到图像底部的距离(像素)
        center_x: 黄点的x坐标(用于垂线参考)
        mask: 黄色区域掩模(可视化时使用)
    """
    # 1. 读取图像并转换色彩空间
    img = cv2.imread(image_path)
    if img is None:
        raise ValueError(f"无法读取图像，请检查路径: {image_path}")
    
    hsv = cv2.cvtColor(img, cv2.COLOR_BGR2HSV)
    height, width = img.shape[:2]
    
    # 2. 定义黄色颜色范围 (考虑不同光照条件)
    lower_yellow = np.array([20, 100, 100])
    upper_yellow = np.array([30, 255, 255])
    
    # 3. 创建黄色区域掩模
    mask = cv2.inRange(hsv, lower_yellow, upper_yellow)
    
    # 4. 形态学处理去除噪声
    kernel = np.ones((5,5), np.uint8)
    mask = cv2.morphologyEx(mask, cv2.MORPH_OPEN, kernel)
    mask = cv2.morphologyEx(mask, cv2.MORPH_CLOSE, kernel)
    
    # 5. 寻找轮廓
    contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    
    if not contours:
        print("未检测到黄色区域")
        return None, None, mask
    
    # 6. 找到所有黄色区域的中心点
    yellow_points = []
    for cnt in contours:
        M = cv2.moments(cnt)
        if M["m00"] > 100:  # 忽略太小的区域
            cx = int(M["m10"] / M["m00"])
            cy = int(M["m01"] / M["m00"])
            yellow_points.append((cx, cy))
    
    if not yellow_points:
        print("未找到有效的黄色中心点")
        return None, None, mask
    
    # 7. 计算图像中心垂线 (考虑鱼眼畸变，使用图像中心作为参考)
    center_x = width // 2
    vertical_line_threshold = width * 0.1  # 垂线左右10%的容差范围
    
    # 8. 筛选在垂线附近的黄点
    vertical_points = [p for p in yellow_points if abs(p[0] - center_x) < vertical_line_threshold]
    
    if not vertical_points:
        # 如果没有完全垂直的点，选择最接近垂线的点
        vertical_points = sorted(yellow_points, key=lambda p: abs(p[0] - center_x))[:1]
        print(f"警告: 没有严格垂直的点，使用最接近垂线的点: {vertical_points[0]}")
    
    # 9. 找出最下方的黄点
    lowest_point = max(vertical_points, key=lambda p: p[1])
    
    # 10. 计算到图像底部的距离
    distance = height - lowest_point[1]
    
    # 可视化结果
    if visualize:
        vis = img.copy()
        # 标记所有黄点
        for (cx, cy) in yellow_points:
            cv2.circle(vis, (cx, cy), 5, (0, 255, 255), -1)
        # 标记垂线区域
        cv2.line(vis, (center_x, 0), (center_x, height), (0, 255, 0), 1)
        cv2.line(vis, (int(center_x - vertical_line_threshold), 0), 
                (int(center_x - vertical_line_threshold), height), (0, 255, 0), 1)
        cv2.line(vis, (int(center_x + vertical_line_threshold), 0), 
                (int(center_x + vertical_line_threshold), height), (0, 255, 0), 1)
        # 标记最低黄点
        cv2.circle(vis, lowest_point, 10, (0, 0, 255), -1)
        cv2.line(vis, (lowest_point[0], lowest_point[1]), 
                (lowest_point[0], height), (0, 0, 255), 2)
        cv2.putText(vis, f"Distance: {distance}px", (10, 30), 
                   cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 255), 2)
        
        # 确保保存目录存在
        os.makedirs("saved_images", exist_ok=True)
        
        # 保存原始图像和结果图像
        cv2.imwrite("saved_images/Original_Image.jpg", img)
        cv2.imwrite("saved_images/Detection_Result.jpg", vis)
        
    
    return distance, center_x, mask

# 使用示例
try:
    distance, center_x, _ = detect_yellow_distance_from_bottom(
        "/home/mi-task/saved_images/right_20250820_120541_263023.jpg", 
        visualize=True
    )
    
    if distance is not None:
        print(f"黄点到图像底部的距离: {distance} 像素")
        print(f"参考垂线x坐标: {center_x}")
    else:
        print("未能检测到有效的黄色点")
        
except Exception as e:
    print(f"发生错误: {str(e)}")