from rapidocr_onnxruntime import RapidOCR
from PIL import Image
import os, time, json

ocr = RapidOCR()
img_dir = r"D:\workspace\membank\售价"
images = sorted([f for f in os.listdir(img_dir) if f.lower().endswith(('.jpeg', '.jpg', '.png'))])

all_results = {}

for img_name in images:
    img_path = os.path.join(img_dir, img_name)
    img = Image.open(img_path)
    w, h = img.size
    print(f"\n=== {img_name} ({w}x{h}) ===")

    # Use chunking for images taller than 2000px
    if h > 2000:
        chunk_size = 1200
        all_text = []
        for i in range(0, h, chunk_size):
            bottom = min(i + chunk_size, h)
            chunk = img.crop((0, i, w, bottom))
            tmp_path = os.path.join(os.environ['TEMP'], f'ocr_chunk_{i}.png')
            chunk.save(tmp_path)
            result, _ = ocr(tmp_path)
            if result:
                for line in result:
                    text = line[1]
                    conf = line[2]
                    bbox = line[0]
                    y_center = (bbox[0][1] + bbox[2][1]) / 2 + i  # absolute y position
                    all_text.append((y_center, text, conf))
            if os.path.exists(tmp_path):
                os.remove(tmp_path)
        # Sort by vertical position
        all_text.sort(key=lambda x: x[0])
        lines = [(t[1], t[2]) for t in all_text]
    else:
        result, _ = ocr(img_path)
        if result:
            lines = [(line[1], line[2]) for line in result]
        else:
            lines = []

    all_results[img_name] = lines
    for text, conf in lines:
        print(f"  [{conf}] {text}")

# Save raw results as JSON for later processing
with open(os.path.join(img_dir, "ocr_results.json"), "w", encoding="utf-8") as f:
    json.dump(all_results, f, ensure_ascii=False, indent=2)

print("\n\n=== OCR完成，结果已保存到 ocr_results.json ===")