File size: 3,482 Bytes
2477303 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 | import gradio as gr
from gtts import gTTS
import os
from PIL import Image
from transformers import pipeline, AutoModelForSeq2SeqLM, AutoTokenizer, BlipProcessor, BlipForConditionalGeneration
# Define translation models per language pair
translation_models = {
"French": "Helsinki-NLP/opus-mt-en-fr",
"Spanish": "Helsinki-NLP/opus-mt-en-es",
"Mandarin": "Helsinki-NLP/opus-mt-en-zh",
"Vietnamese": "Helsinki-NLP/opus-mt-en-vi",
"Tagalog": "Helsinki-NLP/opus-mt-en-tl"
}
# Directory to store local models
local_model_dir = "./models"
os.makedirs(local_model_dir, exist_ok=True)
# Load translation models and tokenizers locally
translation_pipelines = {}
for lang, model_name in translation_models.items():
local_path = os.path.join(local_model_dir, model_name.replace('/', '_'))
if not os.path.exists(local_path):
model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
tokenizer = AutoTokenizer.from_pretrained(model_name)
model.save_pretrained(local_path)
tokenizer.save_pretrained(local_path)
else:
model = AutoModelForSeq2SeqLM.from_pretrained(local_path)
tokenizer = AutoTokenizer.from_pretrained(local_path)
translation_pipelines[lang] = pipeline("translation", model=model, tokenizer=tokenizer)
# Load image captioning model
caption_processor = BlipProcessor.from_pretrained("Salesforce/blip-image-captioning-base")
caption_model = BlipForConditionalGeneration.from_pretrained("Salesforce/blip-image-captioning-base")
def generate_caption(image):
inputs = caption_processor(image, return_tensors="pt")
output = caption_model.generate(**inputs)
caption = caption_processor.decode(output[0], skip_special_tokens=True)
return caption
# Function to translate text using locally stored models
def translate_text(sentence, target_lang):
if target_lang in translation_pipelines:
model = translation_pipelines[target_lang]
translation = model(sentence, max_length=100, truncation=True)
return translation[0]['translation_text']
return sentence # Return original if no model found
# Language options
language_options = {
"English": "en", "French": "fr", "Spanish": "es", "Mandarin": "zh", "Vietnamese": "vi", "Tagalog": "tl"
}
# Main function for Gradio
def process_image(image, l1, l2):
if not image:
return "Please upload an image.", None, None
caption_en = generate_caption(image)
caption_l2 = translate_text(caption_en, l2) if l2 != "English" else caption_en
# Generate TTS
tts = gTTS(text=caption_l2, lang=language_options[l2])
audio_path = "caption_audio.mp3"
tts.save(audio_path)
return caption_l2, audio_path, caption_en
# Gradio Interface
interface = gr.Interface(
fn=process_image,
inputs=[
gr.Image(type="pil", sources=["upload"]), # Only allows upload and drag-and-drop
gr.Dropdown(list(language_options.keys()), label="Your Native Language (L1)"),
gr.Dropdown(list(language_options.keys()), label="Target Language (L2)")
],
outputs=[
gr.Textbox(label="Generated Caption (L2)"),
gr.Audio(label="Listen to the Caption"),
gr.Textbox(label="Generated Caption (L1)")
],
title="Image Caption Generator for Language Learning",
description="Upload an image, select your native language (L1) and target language (L2), and get a generated caption with text-to-speech audio."
)
# Launch Gradio
interface.launch()
|