MuniVis-App / app.py
Rahaf-2050's picture
Upload app.py
68e82ca verified
Raw
History Blame Contribute Delete
1.26 kB
import gradio as gr
from transformers import BlipProcessor, BlipForConditionalGeneration, MarianMTModel, MarianTokenizer
from PIL import Image
import torch
# Load BLIP model for image captioning
processor = BlipProcessor.from_pretrained("Salesforce/blip-image-captioning-base")
model = BlipForConditionalGeneration.from_pretrained("Salesforce/blip-image-captioning-base")
# Load MarianMT model for English to Arabic translation
translator_model_ar = MarianMTModel.from_pretrained("Helsinki-NLP/opus-mt-en-ar")
translator_tokenizer_ar = MarianTokenizer.from_pretrained("Helsinki-NLP/opus-mt-en-ar")
# Prediction function
def predict(image):
inputs = processor(image, return_tensors="pt")
out = model.generate(**inputs)
english_caption = processor.decode(out[0], skip_special_tokens=True)
# Translate to Arabic
tokens = translator_tokenizer_ar.prepare_seq2seq_batch([english_caption], return_tensors="pt")
translated = translator_model_ar.generate(**tokens)
arabic_caption = translator_tokenizer_ar.decode(translated[0], skip_special_tokens=True)
return arabic_caption
# Gradio interface
demo = gr.Interface(fn=predict, inputs=gr.Image(type="pil"), outputs="text", title="Image Captioning (English to Arabic)")
demo.launch()