MrSalman's picture
Create app.py
a5d977a
raw
history blame
1.67 kB
# impoprt packages
import torch
import requests
from PIL import Image
from transformers import BlipProcessor, BlipForConditionalGeneration, AutoTokenizer, pipeline
import sentencepiece
import gradio as gr
# Image captioning model
processor = BlipProcessor.from_pretrained("Salesforce/blip-image-captioning-base")
model = BlipForConditionalGeneration.from_pretrained("Salesforce/blip-image-captioning-base")
# Translate en to ar
model_translater = pipeline("translation", model="Helsinki-NLP/opus-mt-tc-big-en-ar")
# conditional image captioning (with prefix-)
def image_captioning(image_url, prefix="a "):
""" Return text (As str) to describe an image """
# Get the image by image_url and convert it to RGB
raw_image = Image.open(requests.get(image_url, stream=True).raw).convert('RGB')
# Process the image
inputs = processor(raw_image, prefix, return_tensors="pt")
# Generate text to describe the image
output = model.generate(**inputs)
# Decode the output
output = processor.decode(output[0], skip_special_tokens=True, max_length=80)
return output
def translate_text(text, to="ar"):
""" Return translated text """
translated_text = model_translater(str(text))
return translated_text[0]['translation_text']
def image_captioning_ar(image_url, prefix = "a "):
if image_url:
text = image_captioning(image_url, prefix=prefix)
return translate_text(text)
return null
imageCaptioning_interface = gr.Interface(
fn = image_captioning_ar,
inputs=gr.inputs.Textbox(lines = 7, label = 'Image url'),
outputs=gr.outputs.Textbox(label="Caption"),
title = 'Image captioning',
)
imageCaptioning_interface.launch()