mi-task/utils/detect_track.py

import cv2
import numpy as np
from sklearn import linear_model

def detect_horizontal_track_edge(image, observe=False, delay=1000):
    """
    检测正前方横向黄色赛道的边缘，并返回y值最大的边缘点
    
    参数:
        image: 输入图像，可以是文件路径或者已加载的图像数组
        observe: 是否输出中间状态信息和可视化结果，默认为False
            delay: 展示每个步骤的等待时间(毫秒)
        
    返回:
        edge_point: 赛道前方边缘点的坐标 (x, y)
        edge_info: 边缘信息字典
    """
    # 如果输入是字符串（文件路径），则加载图像
    if isinstance(image, str):
        img = cv2.imread(image)
    else:
        img = image.copy()
    
    if img is None:
        print("无法加载图像")
        return None, None
    
    # 获取图像尺寸
    height, width = img.shape[:2]
    
    # 计算图像中间区域的范围（用于专注于正前方的赛道）
    center_x = width // 2
    search_width = int(width * 2/3)  # 搜索区域宽度为图像宽度的2/3
    search_height = height  # 搜索区域高度为图像高度的1/1
    left_bound = center_x - search_width // 2
    right_bound = center_x + search_width // 2
    bottom_bound = height
    top_bound = height - search_height
    
    if observe:
        print("步骤1: 原始图像已加载")
        search_region_img = img.copy()
        # 绘制搜索区域
        cv2.rectangle(search_region_img, (left_bound, top_bound), (right_bound, bottom_bound), (255, 0, 0), 2)
        cv2.line(search_region_img, (center_x, 0), (center_x, height), (0, 0, 255), 2)  # 中线
        cv2.imshow("搜索区域", search_region_img)
        cv2.waitKey(delay)
    
    # 转换到HSV颜色空间以便更容易提取黄色
    hsv = cv2.cvtColor(img, cv2.COLOR_BGR2HSV)
    
    # 黄色的HSV范围
    lower_yellow = np.array([20, 100, 100])
    upper_yellow = np.array([30, 255, 255])
    
    # 创建黄色的掩码
    mask = cv2.inRange(hsv, lower_yellow, upper_yellow)
    
    # 添加形态学操作以改善掩码
    kernel = np.ones((3, 3), np.uint8)
    mask = cv2.dilate(mask, kernel, iterations=1)
    
    if observe:
        print("步骤2: 创建黄色掩码")
        cv2.imshow("黄色掩码", mask)
        cv2.waitKey(delay)
    
    # 应用掩码，只保留黄色部分
    yellow_only = cv2.bitwise_and(img, img, mask=mask)
    
    if observe:
        print("步骤3: 提取黄色部分")
        cv2.imshow("只保留黄色", yellow_only)
        cv2.waitKey(delay)
    
    # 裁剪掩码到搜索区域
    search_mask = mask[top_bound:bottom_bound, left_bound:right_bound]
    
    # 找到掩码在搜索区域中最底部的非零点位置
    bottom_points = []
    non_zero_cols = np.where(np.any(search_mask, axis=0))[0]
    
    # 寻找每列的最底部点
    for col in non_zero_cols:
        col_points = np.where(search_mask[:, col] > 0)[0]
        if len(col_points) > 0:
            bottom_row = np.max(col_points)
            bottom_points.append((left_bound + col, top_bound + bottom_row))
    
    if len(bottom_points) < 3:
        # 如果找不到足够的底部点，使用canny+霍夫变换
        edges = cv2.Canny(mask, 50, 150, apertureSize=3)
        
        if observe:
            print("步骤3.1: 边缘检测")
            cv2.imshow("边缘检测", edges)
            cv2.waitKey(delay)
        
        # 使用霍夫变换检测直线 - 调低阈值以检测短线段
        lines = cv2.HoughLinesP(edges, 1, np.pi/180, threshold=30, 
                               minLineLength=width*0.1, maxLineGap=30)
        
        if lines is None or len(lines) == 0:
            if observe:
                print("未检测到直线")
            return None, None
        
        if observe:
            print(f"步骤4: 检测到 {len(lines)} 条直线")
            lines_img = img.copy()
            for line in lines:
                x1, y1, x2, y2 = line[0]
                cv2.line(lines_img, (x1, y1), (x2, y2), (0, 255, 0), 2)
            cv2.imshow("检测到的直线", lines_img)
            cv2.waitKey(delay)
        
        # 筛选水平线，但放宽斜率条件
        horizontal_lines = []
        for line in lines:
            x1, y1, x2, y2 = line[0]
            
            # 计算斜率 (避免除零错误)
            if abs(x2 - x1) < 5:  # 几乎垂直的线
                continue
                
            slope = (y2 - y1) / (x2 - x1)
            
            # 筛选接近水平的线 (斜率接近0)，但容许更大的倾斜度
            if abs(slope) < 0.3:
                # 确保线在搜索区域内
                if ((left_bound <= x1 <= right_bound and top_bound <= y1 <= bottom_bound) or
                    (left_bound <= x2 <= right_bound and top_bound <= y2 <= bottom_bound)):
                    # 计算线的中点y坐标
                    mid_y = (y1 + y2) / 2
                    line_length = np.sqrt((x2-x1)**2 + (y2-y1)**2)
                    # 保存线段、其y坐标和长度
                    horizontal_lines.append((line[0], mid_y, slope, line_length))
        
        if not horizontal_lines:
            if observe:
                print("未检测到水平线")
            return None, None
        
        if observe:
            print(f"步骤4.1: 找到 {len(horizontal_lines)} 条水平线")
            h_lines_img = img.copy()
            for line_info in horizontal_lines:
                line, _, slope, _ = line_info
                x1, y1, x2, y2 = line
                cv2.line(h_lines_img, (x1, y1), (x2, y2), (0, 255, 255), 2)
                # 显示斜率
                cv2.putText(h_lines_img, f"{slope:.2f}", ((x1+x2)//2, (y1+y2)//2), 
                            cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 0, 0), 1)
            cv2.imshow("水平线", h_lines_img)
            cv2.waitKey(delay)
        
        # 按y坐标排序 (从大到小，底部的线排在前面)
        horizontal_lines.sort(key=lambda x: x[1], reverse=True)
        
        # 取最靠近底部且足够长的线作为横向赛道线
        selected_line = None
        selected_slope = 0
        for line_info in horizontal_lines:
            line, _, slope, length = line_info
            if length > width * 0.1:  # 确保线足够长
                selected_line = line
                selected_slope = slope
                break
        
        if selected_line is None and horizontal_lines:
            # 如果没有足够长的线，就取最靠近底部的线
            selected_line = horizontal_lines[0][0]
            selected_slope = horizontal_lines[0][2]
        
        if selected_line is None:
            if observe:
                print("无法选择合适的线段")
            return None, None
            
        x1, y1, x2, y2 = selected_line
    else:
        # 使用底部点拟合直线
        if observe:
            bottom_points_img = img.copy()
            for point in bottom_points:
                cv2.circle(bottom_points_img, point, 3, (0, 255, 0), -1)
            cv2.imshow("底部边缘点", bottom_points_img)
            cv2.waitKey(delay)
        
        # 使用RANSAC拟合直线以去除异常值
        x_points = np.array([p[0] for p in bottom_points]).reshape(-1, 1)
        y_points = np.array([p[1] for p in bottom_points])
        
        # 如果点过少或分布不够宽，返回None
        if len(bottom_points) < 3 or np.max(x_points) - np.min(x_points) < width * 0.1:
            if observe:
                print("底部点太少或分布不够宽")
            return None, None
        
        ransac = linear_model.RANSACRegressor(residual_threshold=5.0)
        ransac.fit(x_points, y_points)
        
        # 获取拟合参数
        selected_slope = ransac.estimator_.coef_[0]
        intercept = ransac.estimator_.intercept_
        
        # 检查斜率是否在合理范围内
        if abs(selected_slope) > 0.3:
            if observe:
                print(f"拟合斜率过大: {selected_slope:.4f}")
            return None, None
        
        # 使用拟合的直线参数计算线段端点
        x1 = left_bound
        y1 = int(selected_slope * x1 + intercept)
        x2 = right_bound
        y2 = int(selected_slope * x2 + intercept)
        
        if observe:
            fitted_line_img = img.copy()
            cv2.line(fitted_line_img, (x1, y1), (x2, y2), (0, 255, 255), 2)
            cv2.imshow("拟合线段", fitted_line_img)
            cv2.waitKey(delay)
    
    # 确保x1 < x2
    if x1 > x2:
        x1, x2 = x2, x1
        y1, y2 = y2, y1
    
    # 找到线上y值最大的点作为边缘点(最靠近相机的点)
    if y1 > y2:
        bottom_edge_point = (x1, y1)
    else:
        bottom_edge_point = (x2, y2)
    
    # 获取线上的更多点
    selected_points = []
    step = 5  # 每5个像素取一个点
    for x in range(max(left_bound, int(min(x1, x2))), min(right_bound, int(max(x1, x2)) + 1), step):
        y = int(selected_slope * (x - x1) + y1)
        if top_bound <= y <= bottom_bound:
            selected_points.append((x, y))
    
    if observe:
        print(f"步骤5: 找到边缘点 {bottom_edge_point}")
        edge_img = img.copy()
        # 画线
        cv2.line(edge_img, (x1, y1), (x2, y2), (0, 255, 0), 2)
        # 绘制所有点
        for point in selected_points:
            cv2.circle(edge_img, point, 3, (255, 0, 0), -1)
        # 标记边缘点
        cv2.circle(edge_img, bottom_edge_point, 10, (0, 0, 255), -1)
        cv2.imshow("选定的横向线和边缘点", edge_img)
        cv2.waitKey(delay)
    
    # 计算这个点到中线的距离
    distance_to_center = bottom_edge_point[0] - center_x
    
    # 计算中线与检测到的横向线的交点
    # 横向线方程: y = slope * (x - x1) + y1
    # 中线方程: x = center_x
    # 解这个方程组得到交点坐标
    intersection_x = center_x
    intersection_y = selected_slope * (center_x - x1) + y1
    intersection_point = (int(intersection_x), int(intersection_y))
    
    # 计算交点到图像底部的距离（以像素为单位）
    distance_to_bottom = height - intersection_y
    
    if observe:
        slope_img = img.copy()
        # 画出检测到的线
        cv2.line(slope_img, (x1, y1), (x2, y2), (0, 255, 0), 2)
        # 标记边缘点
        cv2.circle(slope_img, bottom_edge_point, 10, (0, 0, 255), -1)
        # 画出中线
        cv2.line(slope_img, (center_x, 0), (center_x, height), (0, 0, 255), 2)
        # 标记中线与横向线的交点
        cv2.circle(slope_img, intersection_point, 12, (255, 0, 255), -1)
        cv2.circle(slope_img, intersection_point, 5, (255, 255, 255), -1)
        # 画出交点到底部的距离线
        cv2.line(slope_img, intersection_point, (intersection_x, height), (255, 255, 0), 2)
        
        cv2.putText(slope_img, f"Slope: {selected_slope:.4f}", (10, 30), 
                    cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 2)
        cv2.putText(slope_img, f"Distance to center: {distance_to_center}px", (10, 70), 
                    cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 2)
        cv2.putText(slope_img, f"Distance to bottom: {distance_to_bottom:.1f}px", (10, 110), 
                    cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 2)
        cv2.putText(slope_img, f"中线交点: ({intersection_point[0]}, {intersection_point[1]})", (10, 150), 
                    cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 2)
        cv2.imshow("边缘斜率和中线交点", slope_img)
        cv2.imwrite("res/path/test/edge_img.png", slope_img)
        cv2.waitKey(delay)
    
    # 创建边缘信息字典
    edge_info = {
        "x": bottom_edge_point[0],
        "y": bottom_edge_point[1],
        "distance_to_center": distance_to_center,
        "slope": selected_slope,
        "is_horizontal": abs(selected_slope) < 0.05,  # 判断边缘是否接近水平
        "points_count": len(selected_points),  # 该组中点的数量
        "intersection_point": intersection_point,  # 中线与横向线的交点
        "distance_to_bottom": distance_to_bottom,  # 交点到图像底部的距离
        "points": selected_points  # 添加选定的点组
    }
    
    return bottom_edge_point, edge_info

# 用法示例
if __name__ == "__main__":
    pass