Image_Describer / app.py
kusumakar's picture
Create app.py
5149f5a
raw
history blame
1.57 kB
import numpy as np
from PIL import Image
import streamlit as st
from transformers import AutoTokenizer, ViTFeatureExtractor, VisionEncoderDecoderModel
# 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
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
# 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)
# run the app
if __name__ == '__main__':
app()