import gradio as gr
from transformers import VisionEncoderDecoderModel, TrOCRProcessor,AutoTokenizer,ViTFeatureExtractor
from PIL import Image
import torch

def preprocess_image(image):
    # Resize while maintaining aspect ratio
    target_size = (224, 224)
    original_size = image.size

    # Calculate the new size while maintaining aspect ratio
    aspect_ratio = original_size[0] / original_size[1]
    if aspect_ratio > 1:  # Width is greater than height
        new_width = target_size[0]
        new_height = int(target_size[0] / aspect_ratio)
    else:  # Height is greater than width
        new_height = target_size[1]
        new_width = int(target_size[1] * aspect_ratio)

    # Resize the image
    resized_img = image.resize((new_width, new_height))

    # Calculate padding values
    padding_width = target_size[0] - new_width
    padding_height = target_size[1] - new_height

    # Apply padding to center the resized image
    pad_left = padding_width // 2
    pad_top = padding_height // 2
    pad_image = Image.new('RGB', target_size, (255, 255, 255))  # White background
    pad_image.paste(resized_img, (pad_left, pad_top))
    return pad_image


# Load model directly
from transformers import AutoTokenizer, AutoModel,ViTFeatureExtractor,TrOCRProcessor,VisionEncoderDecoderModel

tokenizer = AutoTokenizer.from_pretrained("syubraj/TrOCR_Nepali")
model1 = VisionEncoderDecoderModel.from_pretrained("syubraj/TrOCR_Nepali")
feature_extractor1 = ViTFeatureExtractor.from_pretrained("syubraj/TrOCR_Nepali")

processor1 = TrOCRProcessor(feature_extractor=feature_extractor1, tokenizer=tokenizer)


# tokenizer = AutoTokenizer.from_pretrained("paudelanil/trocr-devanagari")
# model = VisionEncoderDecoderModel.from_pretrained("paudelanil/trocr-devanagari")
# feature_extractor = ViTFeatureExtractor.from_pretrained('google/vit-base-patch16-224-in21k')
device = 'cuda' if torch.cuda.is_available() else 'cpu'

model1.to(device)
def predict(image):
    # Preprocess the image
    image = Image.open(image).convert("RGB")
    image = preprocess_image(image)
    pixel_values = processor1(image, return_tensors="pt").pixel_values.to(device)
    
    # Generate text from the image
    generated_ids = model1.generate(pixel_values)
    generated_text = processor1.batch_decode(generated_ids, skip_special_tokens=True)[0]
    
    return generated_text

# Create the Gradio interface
interface = gr.Interface(
    fn=predict,
    inputs=gr.Image(type="filepath"),
    outputs="text",
    title="Devanagari OCR with TrOCR",
    description="Upload an image with Devanagari script and get the text prediction using a pre-trained Vision-Text model."
)

# Launch the interface
interface.launch(share=True)