Update app.py

app.py

@@ -4,10 +4,32 @@ import PIL.Image
 import skimage.io as io
 import streamlit as st
 from transformers import GPT2Tokenizer, GPT2LMHeadModel, AdamW, get_linear_schedule_with_warmup
+from transformers import GPT2TokenizerFast, ViTImageProcessor, VisionEncoderDecoderModel
 from model import generate2,ClipCaptionModel
 from engine import inference
 
-#model loading code 
+
+model_trained = VisionEncoderDecoderModel.from_pretrained("nlpconnect/vit-gpt2-image-captioning")
+model_trained.load_state_dict(torch.load('model_trained.pth',map_location=torch.device('cpu')))
+image_processor = ViTImageProcessor.from_pretrained("nlpconnect/vit-gpt2-image-captioning")
+tokenizer       = GPT2TokenizerFast.from_pretrained("nlpconnect/vit-gpt2-image-captioning")
+
+def show_n_generate(img, greedy = True, model = model_raw):
+    image = Image.open(img)
+    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]
+    returned generated_text
 
 device =  "cpu"
 clip_model, preprocess = clip.load("ViT-B/32", device=device, jit=False)
@@ -35,7 +57,7 @@ def ui():
         pil_image = PIL.Image.fromarray(image)
         image = preprocess(pil_image).unsqueeze(0).to(device)
 
-        option = st.selectbox('Please select the Model',('Model', 'COCO Model','PreTrained Model'))
+        option = st.selectbox('Please select the Model',('Model', 'COCO Model','PreTrained Model','Fine Tuned Model'))
 
         if option=='Model':
             with torch.no_grad():
@@ -60,6 +82,12 @@ def ui():
             st.image(uploaded_file, width = 500, channels = 'RGB')
             st.markdown("**PREDICTION:** " + out)
 
+        elif option=='Fine Tuned Model': 
+            out=show_n_generate(uploaded_file, greedy = False, model = model_trained)
+            st.image(uploaded_file, width = 500, channels = 'RGB')
+            st.markdown("**PREDICTION:** " + out)
+            
+
 
 if __name__ == '__main__':
     ui()