import cv2
import numpy as np
from PIL import Image
import pytesseract

# 配置 Tesseract 路径（仅在 Windows 上需要）
# pytesseract.pytesseract.tesseract_cmd = r'C:\Program Files\Tesseract-OCR\tesseract.exe'

def detect_red_borders(image):
    """
    检测图像中的红色虚线边框，返回边框的轮廓。
    """
    # 转换为 HSV 颜色空间
    hsv = cv2.cvtColor(image, cv2.COLOR_BGR2HSV)

    # 调整红色的 HSV 范围
    lower_red1 = np.array([0, 100, 100])
    upper_red1 = np.array([10, 255, 255])
    lower_red2 = np.array([160, 100, 100])
    upper_red2 = np.array([180, 255, 255])

    # 创建红色掩模
    mask1 = cv2.inRange(hsv, lower_red1, upper_red1)
    mask2 = cv2.inRange(hsv, lower_red2, upper_red2)
    mask = cv2.bitwise_or(mask1, mask2)

    # 使用更大的核进行形态学操作，连接虚线形成完整边框
    kernel = np.ones((5,5), np.uint8)  # 适当调整核大小
    mask = cv2.morphologyEx(mask, cv2.MORPH_CLOSE, kernel, iterations=2)  # 调整迭代次数

    # 使用 Canny 边缘检测辅助检测虚线
    edges = cv2.Canny(mask, 50, 150)

    # 查找轮廓
    contours, hierarchy = cv2.findContours(edges, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)

    # 过滤和存储边框轮廓
    red_borders = []
    min_area = 5000  # 调整最小面积阈值，根据实际情况
    for cnt in contours:
        area = cv2.contourArea(cnt)
        if area > min_area:
            x, y, w, h = cv2.boundingRect(cnt)
            red_borders.append((x, y, w, h))
            # 在原图上画出检测到的区域
            cv2.rectangle(image, (x, y), (x+w, y+h), (0, 255, 0), 2)
    
    # 保存中间结果用于调试
    cv2.imwrite('debug_mask.png', mask)
    cv2.imwrite('debug_edges.png', edges)
    cv2.imwrite('debug_contours.png', image)

    return red_borders

def detect_block_labels(image, red_borders):
    """
    在红色边框区域内检测"Block"标签的位置。
    返回每个边框对应的区域坐标。
    """
    # 将图像转换为灰度图
    gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)

    # 应用二值化
    _, thresh = cv2.threshold(gray, 150, 255, cv2.THRESH_BINARY_INV)

    # OCR 识别
    pil_img = Image.fromarray(thresh)
    data = pytesseract.image_to_data(pil_img, config='--psm 6', output_type=pytesseract.Output.DICT)

    n_boxes = len(data['level'])
    block_labels = []

    for i in range(n_boxes):
        text = data['text'][i].strip()
        if 'block' in text.lower():
            (x, y, w, h) = (data['left'][i], data['top'][i], data['width'][i], data['height'][i])
            block_labels.append((x, y, w, h))

    return block_labels

def extract_blocks(image, red_borders):
    """
    根据检测到的红色边框，从图像中提取对应的区块。
    返回一个包含区块图像的列表。
    """
    blocks = []
    for idx, (x, y, w, h) in enumerate(red_borders):
        # 提取区块区域
        block = image[y:y+h, x:x+w]
        blocks.append(block)

        # 可选：保存每个区块到文件
        # cv2.imwrite(f'block_{idx+1}.png', block)

    return blocks

def main():
    # 加载长截图
    image_path = 'img_1.png'  # 替换为您的截图路径
    image = cv2.imread(image_path)

    if image is None:
        print(f"无法加载图像: {image_path}")
        return
    else:
        print(f"成功加载图像: {image_path}, 尺寸: {image.shape}")

    # 使用修改后的检测函数
    red_borders = detect_red_borders(image)
    print(f"检测到 {len(red_borders)} 个红色边框。")

    if len(red_borders) == 0:
        print("未检测到任何红色边框，请检查边框颜色范围和截图是否正确。")
        # 继续执行以查看调试图片

    # 查看生成的调试图片：
    # - debug_mask.png：红色掩模
    # - debug_edges.png：Canny 边缘图
    # - debug_contours.png：检测到的轮廓标注在原图上

    # 可视化检测到的红色边框
    cv2.imshow('Detected Red Borders', image)
    cv2.waitKey(0)
    cv2.destroyAllWindows()

    # 可选：检测"Block"标签
    blocks = detect_block_labels(image, red_borders)
    print("检测到的标签:", blocks)

    # 提取区块
    extracted_blocks = extract_blocks(image, red_borders)
    print(f"提取了 {len(extracted_blocks)} 个区块。")

    # 保存提取的区块
    for idx, block in enumerate(extracted_blocks):
        block_path = f'block_{idx+1}.png'
        cv2.imwrite(block_path, block)
        print(f"保存区块到: {block_path}")

if __name__ == "__main__":
    main()