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()