Soumen commited on
Commit
8b090f5
1 Parent(s): d23083e

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +27 -30
app.py CHANGED
@@ -22,35 +22,32 @@ print("RESNET MODEL LOADED")
22
  def load_image(img):
23
  im = Image.open(img)
24
  return im
25
- activities = ["Generate ","About"]
26
- choice = st.sidebar.selectbox("Select Activty",activities)
27
  uploaded_photo = c2.file_uploader("Upload Image",type=['jpg','png','jpeg'], on_change=change_photo_state)
28
  camera_photo = c2.camera_input("Take a photo", on_change=change_photo_state)
29
- if choice == 'Generate':
30
- st.subheader("Generate")
31
- if st.session_state["photo"]=="done":
32
- if uploaded_photo:
33
- our_image= load_image(uploaded_photo)
34
- elif camera_photo:
35
- our_image= load_image(camera_photo)
36
- elif uploaded_photo==None and camera_photo==None:
37
- our_image= load_image('image.jpg')
38
- device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
39
- model.to(device)
40
- max_length = 16
41
- num_beams = 4
42
- gen_kwargs = {"max_length": max_length, "num_beams": num_beams}
43
- def predict_step(our_image):
44
- if our_image.mode != "RGB":
45
- our_image = our_image.convert(mode="RGB")
46
- pixel_values = feature_extractor(images=our_image, return_tensors="pt").pixel_values
47
- pixel_values = pixel_values.to(device)
48
- output_ids = model.generate(pixel_values, **gen_kwargs)
49
- preds = tokenizer.batch_decode(output_ids, skip_special_tokens=True)
50
- preds = [pred.strip() for pred in preds]
51
- return preds
52
- st.success(predict_step(our_image))
53
- elif choice == 'About':
54
- st.subheader("About Image Captioning App")
55
- st.markdown("Built with Streamlit by [Soumen Sarker](https://soumen-sarker-personal-site.streamlit.app/)")
56
- st.markdown("Demo applicaton of the following model [credit](https://huggingface.co/nlpconnect/vit-gpt2-image-captioning/)")
 
22
  def load_image(img):
23
  im = Image.open(img)
24
  return im
 
 
25
  uploaded_photo = c2.file_uploader("Upload Image",type=['jpg','png','jpeg'], on_change=change_photo_state)
26
  camera_photo = c2.camera_input("Take a photo", on_change=change_photo_state)
27
+ st.subheader("Caption is below!")
28
+ if st.session_state["photo"]=="done":
29
+ if uploaded_photo:
30
+ our_image= load_image(uploaded_photo)
31
+ elif camera_photo:
32
+ our_image= load_image(camera_photo)
33
+ elif uploaded_photo==None and camera_photo==None:
34
+ our_image= load_image('image.jpg')
35
+ device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
36
+ model.to(device)
37
+ max_length = 16
38
+ num_beams = 4
39
+ gen_kwargs = {"max_length": max_length, "num_beams": num_beams}
40
+ def predict_step(our_image):
41
+ if our_image.mode != "RGB":
42
+ our_image = our_image.convert(mode="RGB")
43
+ pixel_values = feature_extractor(images=our_image, return_tensors="pt").pixel_values
44
+ pixel_values = pixel_values.to(device)
45
+ output_ids = model.generate(pixel_values, **gen_kwargs)
46
+ preds = tokenizer.batch_decode(output_ids, skip_special_tokens=True)
47
+ preds = [pred.strip() for pred in preds]
48
+ return preds
49
+ st.success(predict_step(our_image))
50
+ if st.checkbox('About'):
51
+ st.subheader("About Image Captioning App")
52
+ st.markdown("Built with Streamlit by [Soumen Sarker](https://soumen-sarker-personal-site.streamlit.app/)")
53
+ st.markdown("Demo applicaton of the following model [credit](https://huggingface.co/nlpconnect/vit-gpt2-image-captioning/)")