Image_Translation / src /streamlit_app.py
JaganathC's picture
Update src/streamlit_app.py
1742599 verified
Raw
History Blame Contribute Delete
3.83 kB
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()