Update app.py

app.py

@@ -4,26 +4,10 @@ from PIL import Image
 from transformers import VisionEncoderDecoderModel, ViTFeatureExtractor, AutoTokenizer
 #pickle.load(open('energy_model.pkl', 'rb'))
 #vocab = np.load('w2i.p', allow_pickle=True)
-print("="*150)
-print("MODEL LOADED")
 st.title("img_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")
-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.text("Build with Streamlit and OpenCV")
 if "photo" not in st.session_state:
 	st.session_state["photo"]="not done"
@@ -38,7 +22,20 @@ uploaded_photo = c2.file_uploader("Upload Image",type=['jpg','png','jpeg'], on_c
 camera_photo = c2.camera_input("Take a photo", on_change=change_photo_state)
 #st.subheader("Detection")
 if st.checkbox("Generate_Caption"):
-   our_image=None
+   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
    if st.session_state["photo"]=="done":
       if uploaded_photo:
          our_image= load_image(uploaded_photo)
@@ -46,7 +43,7 @@ if st.checkbox("Generate_Caption"):
          our_image= load_image(camera_photo)
       elif uploaded_photo==None and camera_photo==None:
          our_image= load_image('image.jpg')
-   st.success(predict_step(our_image))
+      st.success(predict_step(our_image))
 elif st.checkbox("About"):
    st.subheader("About Image Captioning App")
    st.markdown("Built with Streamlit by [Soumen Sarker](https://soumen-sarker-personal-website.streamlit.app/)")