commit

app.py

@@ -0,0 +1,97 @@
+import streamlit as st
+import requests
+import io
+
+
+# Designing the interface
+st.title("🖼️ Image Captioning Demo 📝")
+st.write("[Yih-Dar SHIEH](https://huggingface.co/ydshieh)")
+
+st.sidebar.markdown(
+    """
+    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
+    """
+)
+
+with st.spinner('Loading and compiling ViT-GPT2 model ...'):
+    from model import *
+
+random_image_id = get_random_image_id()
+
+st.sidebar.title("Select a sample image")
+sample_image_id = st.sidebar.selectbox(
+    "Please choose a sample image",
+    sample_image_ids
+)
+
+if st.sidebar.button("Random COCO 2017 (val) images"):
+    random_image_id = get_random_image_id()
+    sample_image_id = "None"
+
+bytes_data = None
+with st.sidebar.form("file-uploader-form", clear_on_submit=True):
+    uploaded_file = st.file_uploader("Choose a file")
+    submitted = st.form_submit_button("Upload")
+    if submitted and uploaded_file is not None:
+        bytes_data = io.BytesIO(uploaded_file.getvalue())
+
+if (bytes_data is None) and submitted:
+
+    st.write("No file is selected to upload")
+
+else:
+
+    image_id = random_image_id
+    if sample_image_id != "None":
+        assert type(sample_image_id) == int
+        image_id = sample_image_id
+
+    sample_name = f"COCO_val2017_{str(image_id).zfill(12)}.jpg"
+    sample_path = os.path.join(sample_dir, sample_name)
+
+    if bytes_data is not None:
+        image = Image.open(bytes_data)
+    elif 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)
+
+    width, height = image.size
+    resized = image.resize(size=(width, height))
+    if height > 384:
+        width = int(width / height * 384)
+        height = 384
+        resized = resized.resize(size=(width, height))
+    width, height = resized.size
+    if width > 512:
+        width = 512
+        height = int(height / width * 512)
+        resized = resized.resize(size=(width, height))
+
+    if bytes_data is None:
+        st.markdown(f"[{str(image_id).zfill(12)}.jpg](http://images.cocodataset.org/val2017/{str(image_id).zfill(12)}.jpg)")
+    show = st.image(resized)
+    show.image(resized, '\n\nSelected Image')
+    resized.close()
+
+    # For newline
+    st.sidebar.write('\n')
+
+    with st.spinner('Generating image caption ...'):
+
+        caption = predict(image)
+
+        caption_en = caption
+        st.header(f'Predicted caption:\n\n')
+        st.subheader(caption_en)
+
+    st.sidebar.header("ViT-GPT2 predicts: ")
+    st.sidebar.write(f"{caption}")
+
+    image.close()

model.py

@@ -0,0 +1,82 @@
+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
+
+
+# create target model directory
+model_dir = './models/'
+os.makedirs(model_dir, exist_ok=True)
+
+files_to_download = [
+    "config.json",
+    "flax_model.msgpack",
+    "merges.txt",
+    "special_tokens_map.json",
+    "tokenizer.json",
+    "tokenizer_config.json",
+    "vocab.json",
+    "preprocessor_config.json",
+]
+
+# copy files from checkpoint hub:
+for fn in files_to_download:
+    file_path = hf_hub_download("ydshieh/vit-gpt2-coco-en-ckpts", f"ckpt_epoch_3_step_6900/{fn}")
+    shutil.copyfile(file_path, os.path.join(model_dir, fn))
+
+model = FlaxVisionEncoderDecoderModel.from_pretrained(model_dir)
+feature_extractor = ViTFeatureExtractor.from_pretrained(model_dir)
+tokenizer = AutoTokenizer.from_pretrained(model_dir)
+
+max_length = 16
+num_beams = 4
+gen_kwargs = {"max_length": max_length, "num_beams": num_beams}
+
+
+@jax.jit
+def generate(pixel_values):
+    output_ids = model.generate(pixel_values, **gen_kwargs).sequences
+    return output_ids
+
+
+def predict(image):
+
+    if image.mode != "RGB":
+        image = image.convert(mode="RGB")
+
+    pixel_values = feature_extractor(images=image, return_tensors="np").pixel_values
+
+    output_ids = generate(pixel_values)
+    preds = tokenizer.batch_decode(output_ids, skip_special_tokens=True)
+    preds = [pred.strip() for pred in preds]
+
+    return preds[0]
+
+
+def _compile():
+
+    image_path = 'samples/val_000000039769.jpg'
+    image = Image.open(image_path)
+    predict(image)
+    image.close()
+
+
+_compile()
+
+
+sample_dir = './samples/'
+sample_image_ids = tuple(["None"] + [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

requirements.txt

@@ -0,0 +1,8 @@
+streamlit==0.84.1
+Pillow
+jax[cpu]
+flax
+transformers
+huggingface_hub
+googletrans==4.0.0-rc1
+protobuf==3.20