ydshieh commited on
Commit
144ec50
1 Parent(s): 6f0178d

update logic

Browse files
Files changed (2) hide show
  1. app.py +22 -8
  2. model.py +1 -1
app.py CHANGED
@@ -20,21 +20,23 @@ st.sidebar.markdown(
20
  with st.spinner('Loading and compiling ViT-GPT2 model ...'):
21
  from model import *
22
 
 
23
 
24
  st.sidebar.title("Select a sample image")
25
- image_id = st.sidebar.selectbox(
26
  "Please choose a sample image",
27
  sample_image_ids
28
  )
29
 
30
- random_image_id = None
31
  if st.sidebar.button("Random COCO 2017 (val) images"):
32
  random_image_id = get_random_image_id()
 
33
 
34
- if random_image_id is not None:
35
- image_id = random_image_id
 
 
36
 
37
- st.write(image_id)
38
 
39
  sample_name = f"COCO_val2017_{str(image_id).zfill(12)}.jpg"
40
  sample_path = os.path.join(sample_dir, sample_name)
@@ -45,9 +47,21 @@ else:
45
  url = f"http://images.cocodataset.org/val2017/{str(image_id).zfill(12)}.jpg"
46
  image = Image.open(requests.get(url, stream=True).raw)
47
 
48
- resized = image.resize(size=(384, 384))
49
- show = st.image(resized, width=384)
50
- show.image(resized, '\n\nSelected Image', width=384)
 
 
 
 
 
 
 
 
 
 
 
 
51
  resized.close()
52
 
53
  # For newline
 
20
  with st.spinner('Loading and compiling ViT-GPT2 model ...'):
21
  from model import *
22
 
23
+ random_image_id = get_random_image_id()
24
 
25
  st.sidebar.title("Select a sample image")
26
+ sample_image_id = st.sidebar.selectbox(
27
  "Please choose a sample image",
28
  sample_image_ids
29
  )
30
 
 
31
  if st.sidebar.button("Random COCO 2017 (val) images"):
32
  random_image_id = get_random_image_id()
33
+ sample_image_id = "None"
34
 
35
+ image_id = random_image_id
36
+ if sample_image_id != "None":
37
+ assert type(sample_image_id) == int
38
+ image_id = sample_image_id
39
 
 
40
 
41
  sample_name = f"COCO_val2017_{str(image_id).zfill(12)}.jpg"
42
  sample_path = os.path.join(sample_dir, sample_name)
 
47
  url = f"http://images.cocodataset.org/val2017/{str(image_id).zfill(12)}.jpg"
48
  image = Image.open(requests.get(url, stream=True).raw)
49
 
50
+ width, height = image.size
51
+ resized = image
52
+ if height > 384:
53
+ width = int(width / height * 384)
54
+ height = 384
55
+ resized = resized.resize(size=(width, height))
56
+ if width > 512:
57
+ width = 512
58
+ height = int(height / width * 512)
59
+ resized = resized.resize(size=(width, height))
60
+
61
+
62
+ st.markdown(f"[{str(image_id).zfill(12)}.jpg](http://images.cocodataset.org/val2017/{str(image_id).zfill(12)}.jpg)")
63
+ show = st.image(resized)
64
+ show.image(resized, '\n\nSelected Image')
65
  resized.close()
66
 
67
  # For newline
model.py CHANGED
@@ -66,7 +66,7 @@ _compile()
66
 
67
 
68
  sample_dir = './samples/'
69
- sample_image_ids = tuple([int(f.replace('COCO_val2017_', '').replace('.jpg', '')) for f in os.listdir(sample_dir) if f.startswith('COCO_val2017_')])
70
 
71
  with open(os.path.join(sample_dir, "coco-val2017-img-ids.json"), "r", encoding="UTF-8") as fp:
72
  coco_2017_val_image_ids = json.load(fp)
 
66
 
67
 
68
  sample_dir = './samples/'
69
+ sample_image_ids = tuple(["None"] + [int(f.replace('COCO_val2017_', '').replace('.jpg', '')) for f in os.listdir(sample_dir) if f.startswith('COCO_val2017_')])
70
 
71
  with open(os.path.join(sample_dir, "coco-val2017-img-ids.json"), "r", encoding="UTF-8") as fp:
72
  coco_2017_val_image_ids = json.load(fp)