ImagetoHashtags / app.py
rajag0pal's picture
Update app.py
487c5e5
raw
history blame
2.61 kB
from transformers import GPT2TokenizerFast, ViTImageProcessor, VisionEncoderDecoderModel
import requests
import torch
from PIL import Image
import os
from tqdm import tqdm
import openai
import warnings
warnings.filterwarnings('ignore')
model_raw = VisionEncoderDecoderModel.from_pretrained("nlpconnect/vit-gpt2-image-captioning")
image_processor = ViTImageProcessor.from_pretrained("nlpconnect/vit-gpt2-image-captioning")
tokenizer = GPT2TokenizerFast.from_pretrained("nlpconnect/vit-gpt2-image-captioning")
def Image_to_caption(image,url, greedy = True, model = model_raw):
try:
img = Image.open(requests.get(url, stream =True).raw)
pixel_values = image_processor(img, return_tensors ="pt").pixel_values
except:
pixel_values = image_processor(image, return_tensors ="pt").pixel_values
# plt.imshow(np.asarray(image))
# plt.show()
if greedy:
generated_ids = model.generate(pixel_values, max_new_tokens = 30)
else:
generated_ids = model.generate(
pixel_values,
do_sample=True,
max_new_tokens = 30,
top_k=5)
generated_text = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)[0]
openai.api_key = os.environ['API_KEY']
completion = openai.ChatCompletion.create(
model= "gpt-3.5-turbo",
messages = [{"role":"user","content":f"provide me the top trending hashtags based this text {generated_text} in twitter"}]
)
list1 = []
for words in completion.choices[0].message.content.split():
if words.startswith("#"):
list1.append(words)
return '\n'.join(list1)
import gradio as gr
inputs = [ gr.inputs.Image(type="pil", label="Original Image"), gr.inputs.Textbox(label="Image URL")]
outputs = [ gr.outputs.Textbox(label = 'Caption')]
title = "Image to Hashtags"
description = "This AI tool uses cutting-edge technology to generate captions and relevant hashtags for images. By combining a state-of-the-art ViT-GPT2 image captioning model with OpenAI's GPT-3.5-Turbo API this tool can suggest popular and relevant hashtags. "
article = " <a href='https://huggingface.co/nlpconnect/vit-gpt2-image-captioning'>Model Repo on Hugging Face Model Hub</a>"
examples = [['Screenshot 2023-02-03 at 3.58.03 PM.png','Screenshot 2023-02-03 at 3.57.20 PM.png','Screenshot 2023-02-03 at 3.56.22 PM.png']]
gr.Interface(
Image_to_caption,
inputs,
outputs,
title=title,
description=description,
article=article,
examples=examples,
theme="huggingface",
).launch(debug=True, enable_queue=True)