import gradio as gr
import PIL
import numpy as np
from PIL import Image, ImageDraw, ImageFont
from translate import Translator

def translate_text(text, src_lang='en', dest_lang='vi'):
    translator = Translator(to_lang=dest_lang, from_lang=src_lang)
    translation = translator.translate(text)
    return translation

def process_image(image):
    # Convert the uploaded image data to a PIL Image
    image = Image.open(image).convert('RGB')
    image.save('temp.jpg')

    # Sample code to perform OCR and draw translated text (replace with your logic)
    # Simulated results for demonstration purposes
    result = [
        ([(100, 100), (200, 100), (200, 200), (100, 200)], ("Hello, how are you?", 0.95))
    ]

    draw = PIL.ImageDraw.Draw(image)
    for i, box in enumerate(result):
        text = translate_text(box[1][0])
        box = np.array(box[0]).astype(np.int32)
        xmin = min(box[:, 0])
        ymin = min(box[:, 1])
        xmax = max(box[:, 0])
        ymax = max(box[:, 1])
        draw.rectangle((xmin, ymin, xmax, ymax), outline="red", width=1)
        draw.text((xmin + 40, ymin + 40), f"{text}", fill="black", font=ImageFont.load_default())

    # Save the processed image
    processed_image_path = 'processed_image.jpg'
    image.save(processed_image_path)
    return processed_image_path

iface = gr.Interface(
    fn=process_image, 
    inputs=gr.Image(type="filepath", label="Upload Image"),
    outputs="image",
    title="OCR Translation App",
    description="Upload an image and see the processed image with translated text."
)

if __name__ == "__main__":
    iface.launch(share=True)