|
|
|
from transformers import BlipProcessor, BlipForConditionalGeneration, MarianMTModel, MarianTokenizer |
|
import gradio as gr |
|
from PIL import Image |
|
|
|
|
|
processor = BlipProcessor.from_pretrained("Salesforce/blip-image-captioning-base") |
|
model = BlipForConditionalGeneration.from_pretrained("Salesforce/blip-image-captioning-base") |
|
|
|
|
|
tokenizer = MarianTokenizer.from_pretrained("Helsinki-NLP/opus-mt-en-ar") |
|
model_mt = MarianMTModel.from_pretrained("Helsinki-NLP/opus-mt-en-ar") |
|
|
|
|
|
def translate_to_arabic(text): |
|
inputs = tokenizer(text, return_tensors="pt", padding=True) |
|
translated = model_mt.generate(**inputs) |
|
translated_text = [tokenizer.decode(t, skip_special_tokens=True) for t in translated] |
|
return translated_text[0] |
|
|
|
|
|
def generate_caption(image, language): |
|
|
|
inputs = processor(images=image, return_tensors="pt") |
|
|
|
outputs = model.generate(**inputs) |
|
caption = processor.decode(outputs[0], skip_special_tokens=True) |
|
|
|
|
|
if language == "Arabic": |
|
caption = translate_to_arabic(caption) |
|
|
|
return caption |
|
|
|
|
|
iface = gr.Interface( |
|
fn=generate_caption, |
|
inputs=[gr.inputs.Image(type="pil"), gr.inputs.Dropdown(choices=["English", "Arabic"])], |
|
outputs="text", |
|
title="Image Captioning System", |
|
description="Upload an image and choose the language for caption generation." |
|
) |
|
|
|
|
|
iface.launch() |
|
|
|
|