zamborg commited on
Commit
e1cdab4
1 Parent(s): 67da685

added sub conditioning

Browse files
Files changed (1) hide show
  1. app.py +18 -9
app.py CHANGED
@@ -2,8 +2,16 @@ import streamlit as st
2
  import io
3
  import sys
4
  import time
 
5
  sys.path.append("./virtex/")
6
 
 
 
 
 
 
 
 
7
  st.title("Image Captioning Demo from Redcaps")
8
  st.sidebar.markdown(
9
  """
@@ -12,15 +20,12 @@ st.sidebar.markdown(
12
  )
13
 
14
  with st.spinner("Loading Model"):
15
- st.write("DEBUG PRINTING ==========")
16
- start = time.time()
17
  from model import *
18
- st.write(f"Import TIME: {time.time()-start}")
19
  sample_images = get_samples()
20
- start = time.time()
21
  virtexModel = VirTexModel()
22
  imageLoader = ImageLoader()
23
- st.write(f"model TIME: {time.time()-start}")
 
24
 
25
  random_image = get_rand_img(sample_images)
26
 
@@ -40,6 +45,13 @@ with st.sidebar.form("file-uploader-form", clear_on_submit=True):
40
  submitted = st.form_submit_button("Submit")
41
  if uploaded_file is not None and submitted:
42
  uploaded_image = Image.open(io.BytesIO(uploaded_file.getvalue()))
 
 
 
 
 
 
 
43
 
44
  if uploaded_image is None and submitted:
45
  st.write("Please select a file to upload")
@@ -56,10 +68,7 @@ else:
56
  show = st.image(image)
57
  show.image(image, "Your Image")
58
 
59
- with st.spinner("Generating Caption"):
60
- subreddit, caption = virtexModel.predict(image_dict)
61
- st.header("Predicted Caption:\n\n")
62
- st.subheader(f"r/{subreddit}:\t{caption}\n")
63
 
64
  image.close()
65
 
 
2
  import io
3
  import sys
4
  import time
5
+ import json
6
  sys.path.append("./virtex/")
7
 
8
+ def gen_show_caption():
9
+ with st.spinner("Generating Caption"):
10
+ subreddit, caption = virtexModel.predict(image_dict)
11
+ st.header("Predicted Caption:\n\n")
12
+ st.subheader(f"r/{subreddit}:\t{caption}\n")
13
+
14
+
15
  st.title("Image Captioning Demo from Redcaps")
16
  st.sidebar.markdown(
17
  """
 
20
  )
21
 
22
  with st.spinner("Loading Model"):
 
 
23
  from model import *
 
24
  sample_images = get_samples()
 
25
  virtexModel = VirTexModel()
26
  imageLoader = ImageLoader()
27
+ valid_subs = json.load(open(VALID_SUBREDDITS_PATH))
28
+ valid_subs.insert(0, None)
29
 
30
  random_image = get_rand_img(sample_images)
31
 
 
45
  submitted = st.form_submit_button("Submit")
46
  if uploaded_file is not None and submitted:
47
  uploaded_image = Image.open(io.BytesIO(uploaded_file.getvalue()))
48
+
49
+ st.sidebar.title("Select a Subreddit")
50
+ sub = st.sidebar.selectbox(
51
+ "Select None for a Predicted Subreddit",
52
+ valid_subs
53
+ )
54
+
55
 
56
  if uploaded_image is None and submitted:
57
  st.write("Please select a file to upload")
 
68
  show = st.image(image)
69
  show.image(image, "Your Image")
70
 
71
+ gen_show_caption()
 
 
 
72
 
73
  image.close()
74