import gradio as gr
import spaces
from PIL import Image
from .atlasocr_model import AtlasOCR

# Load model and processor

atlas_ocr=AtlasOCR()



@spaces.GPU
def perform_ocr(image):
    output_text = atlas_ocr(image)
    return output_text

# Create Gradio interface
with gr.Blocks(title="AtlasOCR") as demo:
    gr.Markdown("# AtlasOCR")
    gr.Markdown("Upload an image to extract Darija text in real-time. This model is specialized for Darija document OCR.")
    
    with gr.Row():
        with gr.Column(scale=1):
            # Input image
            image_input = gr.Image(type="numpy", label="Upload Image")
            
            # Example gallery
            gr.Examples(
                examples=[
                    ["2.jpg"],
                    ["3.jpg"]
                ],
                inputs=image_input,
                label="Example Images",
                examples_per_page=4
            )
            
            # Submit button
            submit_btn = gr.Button("Extract Text")
        
        with gr.Column(scale=1):
            # Output text
            output = gr.Textbox(label="Extracted Text", lines=20, show_copy_button=True)
            
            # Model details
            with gr.Accordion("Model Information", open=False):
                gr.Markdown("""
                **Model:** AtlasOCR-v0
                **Description:** Darija OCR model
                **Size:** 3B parameters
                **Context window:** Supports up to 2000 output tokens
                """)
    
    # Set up processing flow
    submit_btn.click(fn=perform_ocr, inputs=image_input, outputs=output)
    image_input.change(fn=perform_ocr, inputs=image_input, outputs=output)

demo.launch()