from docx import Document
from PIL import Image
from io import BytesIO

def extract_images_from_docx(docx_file, target_titles):
    """
    从指定的 docx 文件中提取目标标题下的图片。
    :param docx_file: str - docx 文件路径
    :param target_titles: list - 目标标题列表
    :return: list - 提取的图片（以 PIL Image 格式返回）
    """
    document = Document(docx_file)
    images = []
    capture_images = False

    for paragraph in document.paragraphs:
        # 判断当前段落是否为目标标题
        if paragraph.style.name.startswith("Heading") and paragraph.text.strip() in target_titles:
            capture_images = True
        elif paragraph.style.name.startswith("Heading"):
            capture_images = False

        # 如果在目标标题下，提取图片
        if capture_images:
            for run in paragraph.runs:
                if run.element.xpath('.//w:drawing'):
                    for drawing in run.element.xpath('.//w:blip/@r:embed'):
                        rel = document.part.rels[drawing]
                        image_data = rel.target_part.blob
                        image = Image.open(BytesIO(image_data))
                        images.append(image)

    return images

from paddleocr import PaddleOCR
import cv2
import numpy as np

# 初始化 PaddleOCR
ocr = PaddleOCR(use_angle_cls=True, lang="ch")  # 中英文混合识别

def perform_ocr_with_paddle(images):
    """
    使用 PaddleOCR 对图片进行文字识别。
    :param images: list - PIL Image 对象列表
    :return: list - 每张图片的识别结果，按行分组的文本内容
    """
    results = []
    for image in images:
        # OCR 识别
        try:
            # 将 PIL Image 转换为 OpenCV 格式
            cv_image = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR)
            ocr_result = ocr.ocr(cv_image, cls=True)

            # 提取文字
            text_lines = [line[1][0] for line in ocr_result[0]]  # line[1][0] 是识别的文字
            results.append("\n".join(text_lines))  # 每张图片的文本按行拼接
        except Exception as e:
            results.append(f"OCR 识别失败: {e}")
    return results

def main():
    # 配置输入文件路径和目标标题
    docx_file = "需求文档.docx"  # 替换为实际文件路径
    target_titles = ["功能需求", "测试用例"]  # 替换为实际标题
    api_url = "https://deepseek.aliyun.com/api/v1/r1"  # 替换为实际 deepseek 接口地址
    api_key = "your_api_key_here"  # 替换为实际 API Key

    # 提取图片
    images = extract_images_from_docx(docx_file, target_titles)
    if not images:
        print("未找到任何图片，请检查标题是否正确。")
        return

    # OCR 识别
    ocr_texts = perform_ocr_with_paddle(images)

    # 发送到 deepseek 接口
    deepseek_results = send_to_deepseek(ocr_texts, api_url, api_key)

    # 输出结果
    for i, (ocr_text, result) in enumerate(zip(ocr_texts, deepseek_results), 1):
        print(f"\n图片 {i} 识别的文字内容：")
        print(ocr_text)
        print(f"\n图片 {i} 的 deepseek 返回结果：")
        print(json.dumps(result, indent=2, ensure_ascii=False))