#!/usr/bin/env python3
# -*- coding: utf-8 -*-
import os
import sys
import datetime
import cv2
import numpy as np
import easyocr
# ---------- 圖像預(yù)處理 ----------
def preprocess_image(image_path):
    """
    使用 OpenCV 對圖片進(jìn)行一系列增強，以提高 OCR 識別率。
    返回處理后的灰度圖（numpy 數(shù)組），可直接送入 EasyOCR。
    """
    # 讀取圖片（支持中文路徑）
    with open(image_path, 'rb') as f:
        data = np.frombuffer(f.read(), dtype=np.uint8)
    img = cv2.imdecode(data, cv2.IMREAD_COLOR)
    if img is None:
        raise ValueError(f"無法讀取圖片：{image_path}")
    # 1. 轉(zhuǎn)灰度
    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
    # 2. 去噪（非局部均值去噪，保留邊緣）
    denoised = cv2.fastNlMeansDenoising(gray, None, h=10, templateWindowSize=7, searchWindowSize=21)
    # 3. 對比度增強（CLAHE）
    clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8, 8))
    enhanced = clahe.apply(denoised)
    # 4. 輕微銳化（可選，改善邊緣）
    kernel_sharpen = np.array([[0, -1, 0],
                               [-1, 5, -1],
                               [0, -1, 0]])
    sharpened = cv2.filter2D(enhanced, -1, kernel_sharpen)
    # 5. 尺寸調(diào)整：如果圖片過小，放大到至少 800 像素寬（保持比例）
    height, width = sharpened.shape[:2]
    if width < 800:
        scale = 800.0 / width
        new_w = int(width * scale)
        new_h = int(height * scale)
        sharpened = cv2.resize(sharpened, (new_w, new_h), interpolation=cv2.INTER_CUBIC)
    return sharpened
# ---------- 段落合并 ----------
def group_into_paragraphs(results, y_gap_ratio=1.5):
    """
    將 EasyOCR 的檢測結(jié)果（bbox, text, confidence）按垂直距離合并為段落。
    - 首先按 top 坐標(biāo)排序所有文本行。
    - 如果當(dāng)前行的 top 與上一段落最后一行的 bottom 的距離小于
      該段落平均行高的 y_gap_ratio 倍，則認(rèn)為屬于同一段落。
    - 每個段落內(nèi)的行按照從上到下的順序用空格連接。
    """
    if not results:
        return []
    # 給每個框計算 top, bottom, center_y, center_x
    processed = []
    for (bbox, text, conf) in results:
        # bbox 是四個點 [[x1,y1],[x2,y2],[x3,y3],[x4,y4]]
        pts = np.array(bbox)
        top = np.min(pts[:, 1])
        bottom = np.max(pts[:, 1])
        center_y = (top + bottom) / 2.0
        center_x = (np.min(pts[:, 0]) + np.max(pts[:, 0])) / 2.0
        height = bottom - top
        processed.append({
            'bbox': bbox,
            'text': text,
            'conf': conf,
            'top': top,
            'bottom': bottom,
            'center_y': center_y,
            'center_x': center_x,
            'height': height
        })
    # 按 top 排序（從上到下）
    processed.sort(key=lambda x: x['top'])
    paragraphs = []
    current_para = []
    para_avg_height = 0.0
    for item in processed:
        if not current_para:
            current_para.append(item)
            para_avg_height = item['height']
        else:
            # 上一行的 bottom 與當(dāng)前行的 top 之間的間隙
            last_bottom = current_para[-1]['bottom']
            gap = item['top'] - last_bottom
            # 使用當(dāng)前段落平均行高計算閾值
            threshold = para_avg_height * y_gap_ratio if para_avg_height > 0 else 20
            if gap < threshold:
                current_para.append(item)
                # 更新平均行高
                heights = [x['height'] for x in current_para]
                para_avg_height = sum(heights) / len(heights)
            else:
                # 保存當(dāng)前段落，開始新段落
                paragraphs.append(current_para)
                current_para = [item]
                para_avg_height = item['height']
    if current_para:
        paragraphs.append(current_para)
    # 將每個段落內(nèi)的文本按順序（已經(jīng)是 top 順序）用空格合并
    merged_paragraphs = []
    for para in para_inner:
        # 同一段落內(nèi)可能同一行有多個框（按 x 排序）
        # 這里簡單按 top 歸為一組視為一行，然后行間用空格連接
        # 實際上 EasyOCR 一般返回的是行級框，這里保留后續(xù)擴(kuò)展能力
        lines = []
        while para:
            # 取第一行（最小 top）
            first = para[0]
            line_top = first['top']
            line_bottom = first['bottom']
            # 把所有 top 相近的框歸為同一行
            same_line = [x for x in para if abs(x['top'] - line_top) < first['height'] * 0.5]
            # 同一行內(nèi)按 center_x 排序
            same_line.sort(key=lambda x: x['center_x'])
            line_text = ' '.join([x['text'] for x in same_line])
            lines.append(line_text)
            # 移除已處理的行
            para = [x for x in para if x not in same_line]
        # 段落文本：行之間用空格連接（可改為換行符，看你需要）
        merged_paragraphs.append(' '.join(lines))
    return merged_paragraphs
# ---------- 主程序 ----------
def main():
    # 圖片擴(kuò)展名（常見格式，可自行添加）
    IMAGE_EXTENSIONS = ('.png', '.jpg', '.jpeg', '.bmp', '.tiff', '.tif', '.webp')
    # 初始化 EasyOCR（中英文混合，開啟 GPU 加速）
    print("正在初始化 EasyOCR，請稍候...")
    try:
        reader = easyocr.Reader(['ch_sim', 'en'], gpu=True)
    except Exception:
        print("GPU 初始化失敗，切換為 CPU 模式。")
        reader = easyocr.Reader(['ch_sim', 'en'], gpu=False)
    # 獲取當(dāng)前工作目錄
    root_dir = os.getcwd()
    print(f"開始遍歷目錄：{root_dir}")
    # 存儲結(jié)果：{文件路徑: [段落1, 段落2, ...]}
    ocr_results = {}
    for dirpath, _, filenames in os.walk(root_dir):
        for fname in filenames:
            if not fname.lower().endswith(IMAGE_EXTENSIONS):
                continue
            full_path = os.path.join(dirpath, fname)
            rel_path = os.path.relpath(full_path, root_dir)  # 相對路徑便于閱讀
            print(f"正在處理：{rel_path}")
            try:
                # 1. 圖像增強
                processed_img = preprocess_image(full_path)
                # 2. OCR 識別（返回行級結(jié)果，方便后續(xù)段落合并）
                results = reader.readtext(processed_img, paragraph=False)
                if not results:
                    print(f"  -> 未檢測到文本")
                    ocr_results[rel_path] = []
                    continue
                # 3. 段落合并
                paragraphs = group_into_paragraphs(results, y_gap_ratio=1.5)
                ocr_results[rel_path] = paragraphs
                print(f"  -> 檢測到 {len(results)} 個文本行，合并為 {len(paragraphs)} 個段落")
            except Exception as e:
                print(f"  -> 處理出錯：{e}")
                ocr_results[rel_path] = []
    # ---------- 寫入結(jié)果文件 ----------
    timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
    output_filename = f"ocr_results_{timestamp}.txt"
    with open(output_filename, 'w', encoding='utf-8') as f:
        for img_path, paragraphs in ocr_results.items():
            f.write(f"===== {img_path} =====\n")
            if paragraphs:
                for i, para in enumerate(paragraphs, 1):
                    f.write(f"[段落 {i}]\n{para}\n\n")
            else:
                f.write("(未識別到文本)\n\n")
            f.write("\n")  # 不同圖片之間再空一行
    print(f"\n所有處理完成！結(jié)果已保存至：{output_filename}")
if __name__ == "__main__":
    main()

pip install easyocr opencv-python numpy