Image_Describer / app.py
kusumakar's picture
Update app.py
6d700e4
raw
history blame
2.28 kB
import numpy as np
from PIL import Image
import streamlit as st
from transformers import AutoTokenizer, ViTFeatureExtractor, VisionEncoderDecoderModel, GPT2Tokenizer, GPT2LMHeadModel
# Directory path to the saved model on Google Drive
model = VisionEncoderDecoderModel.from_pretrained("nlpconnect/vit-gpt2-image-captioning")
feature_extractor = ViTFeatureExtractor.from_pretrained("nlpconnect/vit-gpt2-image-captioning")
tokenizer = AutoTokenizer.from_pretrained("nlpconnect/vit-gpt2-image-captioning")
# Load the pre-trained model and tokenizer
model_name = "gpt2"
tokenizer = GPT2Tokenizer.from_pretrained(model_name)
model = GPT2LMHeadModel.from_pretrained(model_name)
def generate_captions(image):
image = Image.open(image).convert("RGB")
generated_caption = tokenizer.decode(model.generate(feature_extractor(image, return_tensors="pt").pixel_values.to("cpu"))[0])
sentence = generated_caption
text_to_remove = "<|endoftext|>"
generated_caption = sentence.replace(text_to_remove, "")
return generated_caption
# Define the Streamlit app
def generate_paragraph(prompt):
# Tokenize the prompt
input_ids = tokenizer.encode(prompt, return_tensors="pt")
# Generate the paragraph
output = model.generate(input_ids, max_length=200, num_return_sequences=1, early_stopping=True)
# Decode the generated output into text
paragraph = tokenizer.decode(output[0], skip_special_tokens=True)
return paragraph
# create the Streamlit app
def app():
st.title('Image from your Side, Trending Hashtags from our Side')
st.write('Upload an image to see what we have in store.')
# create file uploader
uploaded_file = st.file_uploader("Got You Covered, Upload your wish!, magic on the Way! ", type=["jpg", "jpeg", "png"])
# check if file has been uploaded
if uploaded_file is not None:
# load the image
image = Image.open(uploaded_file).convert("RGB")
# Image Captions
string = generate_captions(uploaded_file)
st.image(image, caption='The Uploaded File')
st.write("First is first captions for your Photo : ", string)
generated_paragraph = generate_paragraph(string)
st.write(generated_paragraph)
# run the app
if __name__ == '__main__':
app()