import gradio as gr
import easyocr
from transformers import pipeline
import numpy as np
from PIL import Image

# Initialize EasyOCR reader (English by default, can add more languages)
reader = easyocr.Reader(['en'], gpu=False)

# Initialize summarization pipeline from Hugging Face
summarizer = pipeline("summarization", model="facebook/bart-large-cnn")

def extract_and_summarize(image):
    """
    Extract text from image using EasyOCR and summarize using BART
    
    Args:
        image: PIL Image or numpy array
    
    Returns:
        tuple: (extracted_text, summary)
    """
    try:
        # Convert PIL Image to numpy array if needed
        if isinstance(image, Image.Image):
            image = np.array(image)
        
        # Extract text using EasyOCR
        results = reader.readtext(image)
        
        # Combine all extracted text
        extracted_text = " ".join([result[1] for result in results])
        
        if not extracted_text.strip():
            return "No text detected in the image.", "No text to summarize."
        
        # Check if text is long enough to summarize
        word_count = len(extracted_text.split())
        
        if word_count < 30:
            return extracted_text, "Text is too short to summarize. Minimum 30 words required."
        
        # Summarize the extracted text
        # Adjust max_length and min_length based on input length
        max_length = min(150, word_count)
        min_length = min(30, word_count // 2)
        
        summary = summarizer(
            extracted_text,
            max_length=max_length,
            min_length=min_length,
            do_sample=False
        )
        
        summary_text = summary[0]['summary_text']
        
        return extracted_text, summary_text
    
    except Exception as e:
        return f"Error: {str(e)}", "Could not generate summary due to error."

# Create Gradio interface
with gr.Blocks(title="OCR & Text Summarizer") as demo:
    gr.Markdown(
        """
        # 📝 OCR & Text Summarizer
        Upload an image containing text, and this app will:
        1. Extract the text using EasyOCR
        2. Summarize the extracted text using AI (BART model)
        
        **Note:** Works best with clear, readable text. Minimum 30 words required for summarization.
        """
    )
    
    with gr.Row():
        with gr.Column():
            image_input = gr.Image(
                type="pil",
                label="Upload Image"
            )
            submit_btn = gr.Button("Extract & Summarize", variant="primary")
        
        with gr.Column():
            extracted_output = gr.Textbox(
                label="Extracted Text",
                lines=10,
                placeholder="Extracted text will appear here..."
            )
            summary_output = gr.Textbox(
                label="Summary",
                lines=5,
                placeholder="Summary will appear here..."
            )
    
    gr.Examples(
        examples=[],
        inputs=image_input,
        label="Example Images (Add your own)"
    )
    
    submit_btn.click(
        fn=extract_and_summarize,
        inputs=image_input,
        outputs=[extracted_output, summary_output]
    )

# Launch the app
if __name__ == "__main__":
    demo.launch()