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

ocr = PaddleOCR(use_angle_cls=True, lang='en')

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, size):
    # Convert the uploaded image data to a PIL Image
    result = ocr.ocr(image, cls=True)

    image = Image.open(image).convert('RGB')

    draw = PIL.ImageDraw.Draw(image)
    for i, box in enumerate(result[0]):
        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="white", fill="white", width=1)
        draw.text((xmin, ymin), f"{text}", fill="black",
                font=ImageFont.truetype("font/SVN-Arial 2 bold.ttf", size=size))

    # 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"),
            gr.Slider(minimum=10, maximum=100, step=5, label="Text Size")],
    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)