update UI and samples

app.py

@@ -1,4 +1,5 @@
 import streamlit as st
+import requests
 
 
 # Designing the interface
@@ -7,56 +8,60 @@ st.write("[Yih-Dar SHIEH](https://huggingface.co/ydshieh)")
 
 st.sidebar.markdown(
     """
-    An image captioning model [ViT-GPT2](https://huggingface.co/flax-community/vit-gpt2) by combining the ViT model with the GPT2 model.
-    [Part of the [Huggingface JAX/Flax event](https://discuss.huggingface.co/t/open-to-the-community-community-week-using-jax-flax-for-nlp-cv/).]\n
-    The encoder (ViT) and decoder (GPT2) are combined using Hugging Face transformers' `FlaxVisionEncoderDecoderModel`.
+    An image captioning model by combining ViT model with GPT2 model.
+    The encoder (ViT) and decoder (GPT2) are combined using Hugging Face transformers' [Vision-To-Text Encoder-Decoder
+    framework](https://huggingface.co/transformers/master/model_doc/visionencoderdecoder.html).
     The pretrained weights of both models are loaded, with a set of randomly initialized cross-attention weights.
     The model is trained on the COCO 2017 dataset for about 6900 steps (batch_size=256).
+    [Follow-up work of [Huggingface JAX/Flax event](https://discuss.huggingface.co/t/open-to-the-community-community-week-using-jax-flax-for-nlp-cv/).]\n
     """
 )
 
-#image = Image.open('samples/val_000000039769.jpg')
-#show = st.image(image, use_column_width=True)
-#show.image(image, 'Preloaded Image', use_column_width=True)
-
-
 with st.spinner('Loading and compiling ViT-GPT2 model ...'):
-
     from model import *
-    # st.sidebar.write(f'Vit-GPT2 model loaded :)')
 
-st.sidebar.title("Select a sample image")
 
-sample_name = st.sidebar.selectbox(
-    "Please choose an image",
-    sample_fns
+st.sidebar.title("Select a sample image")
+image_id = st.sidebar.selectbox(
+    "Please choose a sample image",
+    sample_image_ids
 )
 
-sample_name = f"COCO_val2014_{sample_name.replace('.jpg', '').zfill(12)}.jpg"
+random_image_id = None
+if st.sidebar.button("Random COCO 2017 (val) images"):
+    random_image_id = get_random_image_id()
+
+if random_image_id is not None:
+    image_id = random_image_id
+
+st.write(image_id)
+
+sample_name = f"COCO_val2017_{str(image_id).zfill(12)}.jpg"
 sample_path = os.path.join(sample_dir, sample_name)
 
-image = Image.open(sample_path)
-show = st.image(image, width=480)
-show.image(image, '\n\nSelected Image', width=480)
+if os.path.isfile(sample_path):
+    image = Image.open(sample_path)
+else:
+    url = f"http://images.cocodataset.org/val2017/{str(image_id).zfill(12)}.jpg"
+    image = Image.open(requests.get(url, stream=True).raw)
+
+resized = image.resize(size=(384, 384))
+show = st.image(resized, width=384)
+show.image(resized, '\n\nSelected Image', width=384)
+resized.close()
 
 # For newline
 st.sidebar.write('\n')
 
-
 with st.spinner('Generating image caption ...'):
 
     caption = predict(image)
 
     caption_en = caption
-    st.header(f'**Prediction (in English)**: {caption_en}')
-    
-    # caption_en = translator.translate(caption, src='fr', dest='en').text
-    # st.header(f'**Prediction (in French) **{caption}')
-    # st.header(f'**English Translation**: {caption_en}')
-
+    st.header(f'Predicted caption:\n\n')
+    st.subheader(caption_en)
 
 st.sidebar.header("ViT-GPT2 predicts:")
 st.sidebar.write(f"**English**: {caption}")
 
-
 image.close()

model.py

@@ -1,12 +1,13 @@
+import json
 import os, shutil
+import random
+
+
 from PIL import Image
 import jax
 from transformers import FlaxVisionEncoderDecoderModel, ViTFeatureExtractor, AutoTokenizer
 from huggingface_hub import hf_hub_download
 
-from googletrans import Translator
-translator = Translator()
-
 
 # create target model directory
 model_dir = './models/'
@@ -65,4 +66,13 @@ _compile()
 
 
 sample_dir = './samples/'
-sample_fns = tuple([f"{int(f.replace('COCO_val2014_', '').replace('.jpg', ''))}.jpg" for f in os.listdir(sample_dir) if f.startswith('COCO_val2014_')])
+sample_image_ids = tuple([int(f.replace('COCO_val2017_', '').replace('.jpg', '')) for f in os.listdir(sample_dir) if f.startswith('COCO_val2017_')])
+
+with open(os.path.join(sample_dir, "coco-val2017-img-ids.json"), "r", encoding="UTF-8") as fp:
+    coco_2017_val_image_ids = json.load(fp)
+
+
+def get_random_image_id():
+
+    image_id = random.sample(coco_2017_val_image_ids, k=1)[0]
+    return image_id