import gradio as gr
import pytesseract
from PIL import Image
import fitz  # PyMuPDF
import io
import cv2
import numpy as np
import os

pytesseract.pytesseract.tesseract_cmd = "tesseract"  # Hugging Face内はこれでOK

def preprocess_image_for_ocr(image):
    gray = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2GRAY)
    clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8, 8))
    gray = clahe.apply(gray)
    blurred = cv2.bilateralFilter(gray, 9, 75, 75)
    sharp = cv2.addWeighted(blurred, 1.5, cv2.GaussianBlur(blurred, (0, 0), 5), -0.5, 0)
    denoised = cv2.fastNlMeansDenoising(sharp, h=20)
    adaptive = cv2.adaptiveThreshold(denoised, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C,
                                     cv2.THRESH_BINARY, 35, 11)
    _, otsu = cv2.threshold(denoised, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU)
    binary = cv2.bitwise_and(adaptive, otsu)
    opened = cv2.morphologyEx(binary, cv2.MORPH_OPEN, np.ones((1, 1), np.uint8))
    h, w = opened.shape
    opened = cv2.resize(opened, (w * 2, h * 2), interpolation=cv2.INTER_CUBIC)
    return Image.fromarray(opened)

def ocr_from_file(file):
    ext = os.path.splitext(file.name)[1].lower()
    result_text = ""
    try:
        if ext in ['.png', '.jpg', '.jpeg', '.bmp', '.tiff']:
            image = Image.open(file)
            processed = preprocess_image_for_ocr(image)
            result_text = pytesseract.image_to_string(processed, lang="jpn+eng")
        elif ext == '.pdf':
            doc = fitz.open(file.name)
            for page_num in range(min(len(doc), 3)):
                page = doc.load_page(page_num)
                pix = page.get_pixmap(matrix=fitz.Matrix(4, 4))
                img = Image.open(io.BytesIO(pix.tobytes()))
                processed = preprocess_image_for_ocr(img)
                result_text += pytesseract.image_to_string(processed, lang="jpn+eng") + "\n"
        else:
            return "非対応形式です。"
    except Exception as e:
        return f"OCRエラー: {e}"
    return result_text.strip()

gr.Interface(
    fn=ocr_from_file,
    inputs=gr.File(label="画像またはPDFをアップロード"),
    outputs=gr.Textbox(label="OCR結果"),
    title="🚀 高精度OCR Lite",
    description="日本語・英語対応の高精度OCRツール（Lite体験版）"
).launch()