import cv2
import numpy as np
import os

def detect_yellow_distance_from_bottom(image_path, visualize=False):
    """
    增强可视化版本的鱼眼图像黄色点检测
    """
    def show_step(step_name, img, delay=500):
        """显示处理步骤"""
        if visualize:
            cv2.imshow(step_name, img)
            cv2.waitKey(delay)
            cv2.imwrite(f"saved_images/step_{step_name}.jpg", img)
    
    # 1. 读取图像
    img = cv2.imread(image_path)
    if img is None:
        raise ValueError("无法读取图像")
    show_step("01_original", img)
    
    # 2. 鱼眼矫正（简化版，实际需用真实相机参数）
    h, w = img.shape[:2]
    K = np.array([[w, 0, w//2], [0, h, h//2], [0, 0, 1]])  # 假设的相机矩阵
    D = np.array([-0.3, 0.1, 0, 0])  # 假设的畸变系数
    undistorted = cv2.fisheye.undistortImage(img, K, D)
    show_step("02_undistorted", undistorted)
    
    # 3. 颜色空间转换
    hsv = cv2.cvtColor(undistorted, cv2.COLOR_BGR2HSV)
    show_step("03_hsv", cv2.cvtColor(hsv, cv2.COLOR_HSV2BGR))  # 转回BGR便于显示
    
    # 4. 颜色阈值处理
    lower_yellow = np.array([20, 150, 150])  # 调整后的阈值
    upper_yellow = np.array([30, 255, 255])
    color_mask = cv2.inRange(hsv, lower_yellow, upper_yellow)
    show_step("04_color_mask", color_mask)
    
    # 5. 形态学处理
    kernel = np.ones((5,5), np.uint8)
    opened = cv2.morphologyEx(color_mask, cv2.MORPH_OPEN, kernel)
    closed = cv2.morphologyEx(opened, cv2.MORPH_CLOSE, kernel)
    show_step("05_morphology", closed)
    
    # 6. 创建垂直区域掩模
    center_x = w // 2
    line_width = max(5, int(w * 0.02))  # 2%宽度
    vertical_mask = np.zeros((h,w), dtype=np.uint8)
    cv2.rectangle(vertical_mask, (center_x-line_width,0), (center_x+line_width,h), 255, -1)
    show_step("06_vertical_mask", vertical_mask)
    
    # 7. 组合掩模
    final_mask = cv2.bitwise_and(closed, vertical_mask)
    show_step("07_final_mask", final_mask)
    
    # 8. 轮廓检测
    contours, _ = cv2.findContours(final_mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    contour_img = cv2.cvtColor(final_mask, cv2.COLOR_GRAY2BGR)
    cv2.drawContours(contour_img, contours, -1, (0,255,0), 2)
    show_step("08_contours", contour_img)
    
    if not contours:
        print("未检测到黄色区域")
        return None, None, final_mask
    
    # 9. 寻找最低点
    lowest_point = max(
        [tuple(cnt[cnt[:,:,1].argmax()][0]) for cnt in contours],
        key=lambda p: p[1]
    )
    distance = h - lowest_point[1]
    
    # 10. 最终可视化
    result_img = undistorted.copy()
    cv2.line(result_img, (center_x,0), (center_x,h), (0,255,0), 1)
    cv2.circle(result_img, lowest_point, 10, (0,0,255), -1)
    cv2.line(result_img, lowest_point, (lowest_point[0],h), (0,0,255), 2)
    cv2.putText(result_img, f"Distance: {distance}px", (10,30), 
               cv2.FONT_HERSHEY_SIMPLEX, 1, (0,0,255), 2)
    show_step("09_final_result", result_img)
    
    return distance, center_x, final_mask

# 使用示例
try:
    os.makedirs("saved_images", exist_ok=True)
    distance, center_x, _ = detect_yellow_distance_from_bottom(
        "/home/mi-task/saved_images/right_eg.jpg", 
        visualize=True
    )
    
    if distance is not None:
        print(f"中垂线最低黄点距离底部: {distance}像素 (x={center_x})")
    else:
        print("未检测到目标")
        
except Exception as e:
    print(f"错误: {str(e)}")