import os

os.environ["FLAGS_use_mkldnn"] = "0"
os.environ["FLAGS_enable_pir_api"] = "1"
os.environ["FLAGS_allocator_strategy"] = "auto_growth"
os.environ["OMP_NUM_THREADS"] = "1"


import gradio as gr
from paddleocr import PaddleOCR
from PIL import Image
import numpy as np

# Inisialisasi OCR (load sekali saja)
ocr = PaddleOCR(
    use_angle_cls=False,
    lang="en"
)

def ocr_image(image: Image.Image):
    if image is None:
        return "Silakan upload gambar terlebih dahulu.", None

    try:
        image = image.convert("RGB")
        image_np = np.array(image)
        width, height = image.size

        # Jalankan OCR
        result = ocr.ocr(image_np)

        texts = []
        output = []

        if result and result[0]:
            for line in result[0]:
                box = line[0]
                text = line[1][0]
                score = float(line[1][1])

                texts.append(text)
                output.append({
                    "bounding_box": box,
                    "text": text,
                    "confidence": score
                })

        summary = {
            "resolution": {
                "width": width,
                "height": height
            },
            "total_text": len(output),
            "results": output
        }

        return "\n".join(texts), summary

    except Exception as e:
        return f"Error: {str(e)}", None


with gr.Blocks(title="OCR App Paddle 3") as demo:
    gr.Markdown("## 📝 OCR Image Reader (PaddleOCR + Paddle 3)")

    with gr.Row():
        image_input = gr.Image(type="pil", label="Upload Gambar")

    with gr.Row():
        text_output = gr.Textbox(label="Hasil OCR (Text)", lines=12)
        json_output = gr.JSON(label="Detail Output")

    btn = gr.Button("Proses OCR")

    btn.click(
        fn=ocr_image,
        inputs=image_input,
        outputs=[text_output, json_output]
    )

if __name__ == "__main__":
    demo.launch()