import cv2
import numpy as np
import os
from pathlib import Path

def preprocess_image(image):
    """
    预处理图像：转换为HSV色彩空间，检测白色区域
    """
    # 转换为HSV色彩空间
    hsv = cv2.cvtColor(image, cv2.COLOR_BGR2HSV)
    
    # 定义白色的HSV范围 - 放宽范围
    # 白色在HSV中：饱和度低，明度高
    # 原来的范围：lower_white = np.array([0, 0, 200]), upper_white = np.array([180, 30, 255])
    # 放宽后的范围：
    lower_white = np.array([0, 0, 150])    # 降低明度下限，从200降到150
    upper_white = np.array([180, 60, 255]) # 提高饱和度上限，从30提高到60
    
    # 创建白色掩码
    white_mask = cv2.inRange(hsv, lower_white, upper_white)
    
    # 形态学操作：开运算去除小噪点，闭运算填充小孔
    kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (5, 5))
    cleaned = cv2.morphologyEx(white_mask, cv2.MORPH_OPEN, kernel)
    cleaned = cv2.morphologyEx(cleaned, cv2.MORPH_CLOSE, kernel)
    
    return cleaned

def find_white_rectangles(binary_image):
    """
    查找白色矩形区域
    """
    contours, _ = cv2.findContours(
        binary_image, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE
    )
    
    # 筛选合适的矩形轮廓
    valid_rectangles = []
    for contour in contours:
        area = cv2.contourArea(contour)
        if area > 200:  # 过滤太小的区域，从500降低到200
            # 计算轮廓的边界矩形
            x, y, w, h = cv2.boundingRect(contour)
            
            # 计算轮廓的近似多边形
            epsilon = 0.02 * cv2.arcLength(contour, True)
            approx = cv2.approxPolyDP(contour, epsilon, True)
            
            # 检查是否为四边形（矩形）
            if len(approx) == 4:
                # 计算宽高比
                aspect_ratio = w / h
                # 矩形应该有合理的宽高比（不是太细长）
                if 0.3 < aspect_ratio < 3.0:
                    valid_rectangles.append((contour, (x, y, w, h), approx))
    
    return valid_rectangles

def analyze_white_region(roi, original_roi):
    """
    分析白色区域的特征
    """
    # 计算白色像素比例
    total_pixels = roi.shape[0] * roi.shape[1]
    white_pixels = np.sum(roi == 255)
    white_ratio = white_pixels / total_pixels
    
    # 计算区域的形状特征
    height, width = roi.shape
    aspect_ratio = width / height
    
    # 计算边缘强度（白色区域应该有清晰的边缘）
    edges = cv2.Canny(original_roi, 50, 150)
    edge_density = np.sum(edges > 0) / total_pixels
    
    return {
        'white_ratio': white_ratio,
        'aspect_ratio': aspect_ratio,
        'edge_density': edge_density,
        'area': total_pixels
    }

def detect_white_rectangles_in_image(image_path):
    """
    检测图片中的白色矩形区域
    """
    # 读取图像
    image = cv2.imread(image_path)
    if image is None:
        print(f"无法读取图像: {image_path}")
        return []
    
    # 预处理
    white_mask = preprocess_image(image)
    
    # 查找白色矩形
    valid_rectangles = find_white_rectangles(white_mask)
    
    results = []
    
    for contour, (x, y, w, h), approx in valid_rectangles:
        # 提取ROI
        roi_mask = white_mask[y:y+h, x:x+w]
        roi_original = image[y:y+h, x:x+w]
        
        # 分析特征
        features = analyze_white_region(roi_mask, roi_original)
        
        # 判断是否为有效的白色矩形
        is_valid_white_rectangle = (
            features['white_ratio'] > 0.4 and  # 白色像素比例要求降低，从0.6降到0.4
            features['edge_density'] > 0.005    # 边缘密度要求降低，从0.01降到0.005
        )
        
        if is_valid_white_rectangle:
            results.append({
                'type': 'white_rectangle',
                'position': (x, y, w, h),
                'features': features,
                'confidence': 'high' if features['white_ratio'] > 0.8 else 'medium'
            })
            
            # 在图像上绘制结果
            cv2.rectangle(image, (x, y), (x + w, y + h), (0, 255, 0), 2)
            cv2.putText(image, f"白色矩形", (x, y - 10), 
                       cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0, 255, 0), 2)
            
            # 绘制轮廓点
            cv2.drawContours(image, [approx], -1, (255, 0, 0), 2)
    
    return image, results

def main():
    """
    主函数：测试白色矩形检测
    """
    # 图片路径
    img_dir = Path("imgs")
    img1_path = img_dir / "1.jpg"
    img2_path = img_dir / "2.jpg"
    
    print("开始检测图片中的白色矩形区域...")
    
    # 检测第一张图片
    if img1_path.exists():
        print(f"\n检测图片: {img1_path}")
        result_img1, results1 = detect_white_rectangles_in_image(str(img1_path))
        
        print("检测结果:")
        for result in results1:
            print(f"  类型: {result['type']}, 位置: {result['position']}, 置信度: {result['confidence']}")
            print(f"  特征: 白色比例={result['features']['white_ratio']:.2f}, "
                  f"宽高比={result['features']['aspect_ratio']:.2f}")
        
        # 保存结果图片
        output_path1 = img_dir / "result_1.jpg"
        cv2.imwrite(str(output_path1), result_img1)
        print(f"结果图片已保存到: {output_path1}")
    
    # 检测第二张图片
    if img2_path.exists():
        print(f"\n检测图片: {img2_path}")
        result_img2, results2 = detect_white_rectangles_in_image(str(img2_path))
        
        print("检测结果:")
        for result in results2:
            print(f"  类型: {result['type']}, 位置: {result['position']}, 置信度: {result['confidence']}")
            print(f"  特征: 白色比例={result['features']['white_ratio']:.2f}, "
                  f"宽高比={result['features']['aspect_ratio']:.2f}")
        
        # 保存结果图片
        output_path2 = img_dir / "result_2.jpg"
        cv2.imwrite(str(output_path2), result_img2)
        print(f"结果图片已保存到: {output_path2}")
    
    print("\n检测完成！")

if __name__ == "__main__":
    main()