Update app.py

app.py

@@ -5,6 +5,7 @@ import skimage.io as io
 import streamlit as st
 from transformers import GPT2Tokenizer, GPT2LMHeadModel, AdamW, get_linear_schedule_with_warmup
 from model import generate2,ClipCaptionModel
+from engine import inference
 
 #model loading code 
 
@@ -25,8 +26,6 @@ coco_model.load_state_dict(torch.load('COCO_model.h5',map_location=torch.device(
 model = model.eval()  
 
 
-
-
 def ui():
     st.markdown("# Image Captioning")
     uploaded_file = st.file_uploader("Upload an Image", type=['png', 'jpeg', 'jpg'])
@@ -36,10 +35,9 @@ 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'))
+        option = st.selectbox('Please select the Model',('Model', 'COCO Model','PreTrained Model'))
 
         if option=='Model':
-
             with torch.no_grad():
                 prefix = clip_model.encode_image(image).to(device, dtype=torch.float32)
                 prefix_embed = model.clip_project(prefix).reshape(1, prefix_length, -1)
@@ -57,6 +55,11 @@ def ui():
             st.image(uploaded_file, width = 500, channels = 'RGB')
             st.markdown("**PREDICTION:** " + generated_text_prefix)
 
+        elif option=='PreTrained Model': 
+            out = inference(uploaded_file)
+            st.image(uploaded_file, width = 500, channels = 'RGB')
+            st.markdown("**PREDICTION:** " + out)
+
 
 if __name__ == '__main__':
     ui()