import streamlit as st
import cv2
import pytesseract
from googletrans import Translator
from PIL import Image, ImageDraw, ImageFont
import numpy as np

# Step 1: Load the image
def load_image(image_file):
    # Read the image using OpenCV
    img = cv2.imdecode(np.fromstring(image_file.read(), np.uint8), cv2.IMREAD_COLOR)
    return img

# Step 2: Extract text and coordinates using pytesseract
def extract_text_with_position(image):
    # Convert the image to grayscale for better OCR performance
    gray_image = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
    
    # Use pytesseract to get OCR results along with bounding boxes
    details = pytesseract.image_to_data(gray_image, output_type=pytesseract.Output.DICT)
    
    text_data = []
    for i in range(len(details['text'])):
        if details['text'][i].strip() != '':
            text_data.append({
                'text': details['text'][i],
                'left': details['left'][i],
                'top': details['top'][i],
                'width': details['width'][i],
                'height': details['height'][i]
            })
    
    return text_data

# Step 3: Translate the extracted text using Google Translate
def translate_text(text, target_language='en'):
    translator = Translator()
    translation = translator.translate(text, dest=target_language)
    return translation.text

# Step 4: Recreate the image with translated text
def recreate_image_with_translated_text(original_img, text_data, target_language):
    # Convert OpenCV image to Pillow image for easier manipulation
    pil_img = Image.fromarray(cv2.cvtColor(original_img, cv2.COLOR_BGR2RGB))
    draw = ImageDraw.Draw(pil_img)
    
    try:
        # Load a default font (you can use a specific font if you want)
        font = ImageFont.load_default()
    except IOError:
        font = ImageFont.truetype("/usr/share/fonts/truetype/dejavu/DejaVuSans-Bold.ttf", 30)

    for item in text_data:
        # Translate each piece of text
        translated_text = translate_text(item['text'], target_language)
        
        # Draw the translated text on the new image at the same position
        draw.text((item['left'], item['top']), translated_text, font=font, fill="black")

    # Convert the image back to an OpenCV format
    output_img = cv2.cvtColor(np.array(pil_img), cv2.COLOR_RGB2BGR)
    return output_img

# Step 5: Display or download the translated image
def display_or_download_image(output_img):
    # Convert image to bytes to allow for download
    _, img_bytes = cv2.imencode('.png', output_img)
    st.image(output_img, channels="BGR", caption="Translated Image", use_column_width=True)
    
    # Provide a button to download the image
    st.download_button(
        label="Download Translated Image",
        data=img_bytes.tobytes(),
        file_name="translated_image.png",
        mime="image/png"
    )

# Streamlit Interface
def main():
    st.title("Image Translation App")

    # Upload image
    image_file = st.file_uploader("Upload an image", type=["jpg", "png", "jpeg"])

    if image_file is not None:
        # Load the image and show it
        img = load_image(image_file)
        st.image(img, caption="Uploaded Image", channels="BGR", use_column_width=True)

        # Get the target language from the user
        target_language = st.text_input("Enter the target language code (e.g., 'es' for Spanish):")

        if target_language:
            # Extract text from image
            text_data = extract_text_with_position(img)

            # Recreate the image with translated text
            translated_img = recreate_image_with_translated_text(img, text_data, target_language)

            # Display or allow download of the translated image
            display_or_download_image(translated_img)

if __name__ == "__main__":
    main()