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

# Initialize PaddleOCR
# use_angle_cls=True allows it to detect rotated text
# lang='en' for English. Change to 'ch', 'fr', 'german', etc. as needed
ocr = PaddleOCR(use_angle_cls=True, lang='en')

def run_ocr(input_image):
    if input_image is None:
        return None, "No image uploaded"

    # Convert PIL Image to numpy array (OpenCV format)
    image_np = np.array(input_image)
    
    # PaddleOCR expects RGB usually, but OpenCV uses BGR for rendering. 
    # The OCR inference handles the channels, but for drawing we ensure consistency.
    
    # Run OCR
    # result structure: [ [ [box], [text, score] ], ... ]
    result = ocr.ocr(image_np, cls=True)

    # Handle cases where no text is found
    if result is None or result[0] is None:
        return input_image, "No text detected."

    result = result[0] # Access the first result list
    
    detected_texts = []
    
    # Visualization setup
    viz_image = image_np.copy()
    
    for line in result:
        box = line[0]
        text_content = line[1][0]
        score = line[1][1]
        
        # Format text for display
        detected_texts.append(f"{text_content} (Conf: {score:.2f})")
        
        # Convert box points to integer for OpenCV
        box = np.array(box).astype(np.int32).reshape((-1, 1, 2))
        
        # Draw the bounding box (Yellow, thickness 2)
        cv2.polylines(viz_image, [box], isClosed=True, color=(0, 255, 255), thickness=2)

    # Join all detected text
    full_text_output = "\n".join(detected_texts)
    
    return viz_image, full_text_output

# Create Gradio Interface
with gr.Blocks(title="PaddleOCR Demo") as demo:
    gr.Markdown("## 🚀 PaddleOCR CPU Demo")
    gr.Markdown("Upload an image to detect and recognize text using PaddleOCR running on CPU.")
    
    with gr.Row():
        with gr.Column():
            input_img = gr.Image(type="pil", label="Input Image")
            submit_btn = gr.Button("Run OCR", variant="primary")
        
        with gr.Column():
            output_img = gr.Image(label="Visualization")
            output_text = gr.Textbox(label="Detected Text", lines=10)

    submit_btn.click(fn=run_ocr, inputs=input_img, outputs=[output_img, output_text])

if __name__ == "__main__":
    demo.launch(server_name="0.0.0.0", server_port=7860)