File size: 2,338 Bytes
551a5b0
 
 
 
d23083e
cabe6ed
e5e1dfb
 
d23083e
 
2b101ad
 
 
d23083e
2b101ad
 
 
d23083e
 
 
2b101ad
 
 
 
 
 
8b090f5
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
import streamlit as st
import torch
from PIL import Image
from transformers import VisionEncoderDecoderModel, ViTFeatureExtractor, AutoTokenizer
st.title("Image_Captioning_App")
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")
#pickle.load(open('energy_model.pkl', 'rb'))
#vocab = np.load('w2i.p', allow_pickle=True)
#st.text("Build with Streamlit and OpenCV")
if "photo" not in st.session_state:
	st.session_state["photo"]="not done"

c2, c3 = st.columns([2,1])
def change_photo_state():
	st.session_state["photo"]="done"
print("="*150)
print("RESNET MODEL LOADED")

@st.cache
def load_image(img):
	im = Image.open(img)
	return im
uploaded_photo = c2.file_uploader("Upload Image",type=['jpg','png','jpeg'], on_change=change_photo_state)
camera_photo = c2.camera_input("Take a photo", on_change=change_photo_state)
st.subheader("Caption is below!") 
if st.session_state["photo"]=="done":
   if uploaded_photo:
      our_image= load_image(uploaded_photo)
   elif camera_photo:
      our_image= load_image(camera_photo)
   elif uploaded_photo==None and camera_photo==None:
      our_image= load_image('image.jpg')
   device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
   model.to(device)
   max_length = 16
   num_beams = 4
   gen_kwargs = {"max_length": max_length, "num_beams": num_beams}
   def predict_step(our_image):
       if our_image.mode != "RGB":
          our_image = our_image.convert(mode="RGB")
       pixel_values = feature_extractor(images=our_image, return_tensors="pt").pixel_values
       pixel_values = pixel_values.to(device)
       output_ids = model.generate(pixel_values, **gen_kwargs)
       preds = tokenizer.batch_decode(output_ids, skip_special_tokens=True)
       preds = [pred.strip() for pred in preds]
       return preds
   st.success(predict_step(our_image))
if st.checkbox('About'):
   st.subheader("About Image Captioning App")
   st.markdown("Built with Streamlit by [Soumen Sarker](https://soumen-sarker-personal-site.streamlit.app/)")
   st.markdown("Demo applicaton of the following model [credit](https://huggingface.co/nlpconnect/vit-gpt2-image-captioning/)")