Duplicate from sujitpal/clip-rsicd-demo

Co-authored-by: Sujit Pal <[email protected]>

.gitattributes

@@ -0,0 +1,16 @@
+*.bin.* filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tar.gz filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text

README.md

@@ -0,0 +1,34 @@
+---
+title: CLIP-RSICD Demo
+emoji: 🛰️
+colorFrom: green
+colorTo: purple
+sdk: streamlit
+app_file: app.py
+pinned: false
+duplicated_from: sujitpal/clip-rsicd-demo
+---
+
+# Configuration
+
+`title`: _string_  
+Display title for the Space
+
+`emoji`: _string_  
+Space emoji (emoji-only character allowed)
+
+`colorFrom`: _string_  
+Color for Thumbnail gradient (red, yellow, green, blue, indigo, purple, pink, gray)
+
+`colorTo`: _string_  
+Color for Thumbnail gradient (red, yellow, green, blue, indigo, purple, pink, gray)
+
+`sdk`: _string_  
+Can be either `gradio` or `streamlit`
+
+`app_file`: _string_  
+Path to your main application file (which contains either `gradio` or `streamlit` Python code).  
+Path is relative to the root of the repository.
+
+`pinned`: _boolean_  
+Whether the Space stays on top of your list.

app.py

@@ -0,0 +1,26 @@
+import dashboard_text2image
+import dashboard_image2image
+import dashboard_featurefinder
+
+import streamlit as st
+
+PAGES = {
+    "Retrieve Images given Text": dashboard_text2image,
+    "Retrieve Images given Image": dashboard_image2image,
+    "Find Feature in Image": dashboard_featurefinder,
+}
+
+st.sidebar.title("CLIP-RSICD")
+st.sidebar.image("thumbnail.jpg")
+st.sidebar.markdown("""
+    We have fine-tuned the CLIP model (see [Model card](https://huggingface.co/flax-community/clip-rsicd-v2)) 
+    using remote sensing images and captions from the [RSICD dataset](https://github.com/201528014227051/RSICD_optimal). 
+    The CLIP model from OpenAI is trained in a self-supervised manner using contrastive learning to project images 
+    and caption text onto a common embedding space.
+
+    Please click here for [more information about our project](https://github.com/arampacha/CLIP-rsicd).
+
+""")
+selection = st.sidebar.radio("Go to", list(PAGES.keys()))
+page = PAGES[selection]
+page.app()

dashboard_featurefinder.py

@@ -0,0 +1,160 @@
+import jax
+import flax
+import matplotlib.pyplot as plt
+import nmslib
+import numpy as np
+import os
+import requests
+import streamlit as st
+
+from tempfile import NamedTemporaryFile
+from torchvision.transforms import Compose, Resize, ToPILImage
+from transformers import CLIPProcessor, FlaxCLIPModel
+from PIL import Image
+
+import utils
+
+BASELINE_MODEL = "openai/clip-vit-base-patch32"
+MODEL_PATH = "flax-community/clip-rsicd-v2"
+
+IMAGE_VECTOR_FILE = "./vectors/test-bs128x8-lr5e-6-adam-ckpt-1.tsv"
+
+IMAGES_DIR = "./images"
+DEMO_IMAGES_DIR = "./demo-images"
+
+
+def split_image(X):
+  num_rows = X.shape[0] // 224
+  num_cols = X.shape[1] // 224
+  Xc = X[0 : num_rows * 224, 0 : num_cols * 224, :]
+  patches = []
+  for j in range(num_rows):
+    for i in range(num_cols):
+      patches.append(Xc[j * 224 : (j + 1) * 224, 
+                        i * 224 : (i + 1) * 224,
+                        :])
+  return num_rows, num_cols, patches
+
+
+def get_patch_probabilities(patches, searched_feature, 
+                            image_preprocesor,
+                            model, processor):
+  images = [image_preprocesor(patch) for patch in patches]
+  text = "An aerial image of {:s}".format(searched_feature)
+  inputs = processor(images=images,
+                    text=text,
+                    return_tensors="jax",
+                    padding=True)
+  outputs = model(**inputs)
+  probs = jax.nn.softmax(outputs.logits_per_text, axis=-1)
+  probs_np = np.asarray(probs)[0]
+  return probs_np
+
+
+def get_image_ranks(probs):
+  temp = np.argsort(-probs)
+  ranks = np.empty_like(temp)
+  ranks[temp] = np.arange(len(probs))
+  return ranks
+
+
+def download_and_prepare_image(image_url):
+    """
+        Take input image and resize it to 672x896 
+    """
+    try:
+        image_raw = requests.get(image_url, stream=True,).raw
+        image = Image.open(image_raw).convert("RGB")
+        width, height = image.size
+        # print("WID,HGT:", width, height)
+        if width < 224 or height < 224:
+            return None
+        # take the short edge and reduce to 672
+        if width < height:
+            resize_factor = 672 / width
+            image = image.resize((672, int(height * resize_factor)))
+            image = image.crop((0, 0, 672, 896))
+        else:
+            resize_factor = 672 / height
+            image = image.resize((int(width * resize_factor), 896))
+            image = image.crop((0, 0, 896, 672))
+        return np.asarray(image)
+    except Exception as e:
+        # print(e)
+        return None
+
+
+
+def app():
+    model, processor = utils.load_model(MODEL_PATH, BASELINE_MODEL)
+
+    st.title("Find Features in Images")
+    st.markdown("""
+        This demo shows the ability of the model to find specific features
+        (specified as text queries) in the image. As an example, say you wish to
+        find the parts of the following image that contain a `beach`, `houses`, 
+        or `ships`. We partition the image into tiles of (224, 224) and report
+        how likely each of them are to contain each text features.
+    """)
+    st.image("demo-images/st_tropez_1.png")
+    st.image("demo-images/st_tropez_2.png")
+    st.markdown("""
+        For this image and the queries listed above, our model reports that the
+        two left tiles are most likely to contain a `beach`, the two top right 
+        tiles are most likely to contain `houses`, and the two bottom right tiles
+        are likely to contain `boats`.
+
+        We have provided a few representative images from [Unsplash](https://unsplash.com/s/photos/aerial-view) 
+        that you can experiment with. Use the image name to put in an initial feature
+        to look for, this will show the original image, and you will get more ideas
+        for features that you can ask the model to identify.
+    """)
+    image_file = st.selectbox(
+        "Sample Image File",
+        options=[
+            "-- select one --",
+            "St-Tropez-Port.jpg",
+            "Acopulco-Bay.jpg",
+            "Highway-through-Forest.jpg",
+            "Forest-with-River.jpg",
+            "Eagle-Bay-Coastline.jpg",
+            "Multistoreyed-Buildings.jpg",
+            "Street-View-Malayasia.jpg",
+        ])
+    image_url = st.text_input(
+        "OR provide an image URL",
+        value="https://static.eos.com/wp-content/uploads/2019/04/Main.jpg")
+    searched_feature = st.text_input("Feature to find", value="beach")
+
+    if st.button("Find"):
+        if image_file.startswith("--"):
+            image = download_and_prepare_image(image_url)
+        else:
+            image = plt.imread(os.path.join("demo-images", image_file))
+
+        if image is None:
+            st.error("Image could not be downloaded, please try another one")
+        else:
+            st.image(image, caption="Input Image")
+            st.markdown("---")
+            num_rows, num_cols, patches = split_image(image)
+            image_preprocessor = Compose([
+                ToPILImage(),
+                Resize(224)
+            ])
+            num_rows, num_cols, patches = split_image(image)
+            patch_probs = get_patch_probabilities(
+                patches,
+                searched_feature,
+                image_preprocessor,
+                model,
+                processor)
+            patch_ranks = get_image_ranks(patch_probs)
+            pid = 0
+            for i in range(num_rows):
+                cols = st.columns(num_cols)
+                for col in cols:
+                    caption = "#{:d} p({:s})={:.3f}".format(
+                        patch_ranks[pid] + 1, searched_feature, patch_probs[pid])
+                    col.image(patches[pid], caption=caption)
+                    pid += 1

dashboard_image2image.py

@@ -0,0 +1,167 @@
+import matplotlib.pyplot as plt
+import nmslib
+import numpy as np
+import os
+import requests
+import streamlit as st
+
+from PIL import Image
+from transformers import CLIPProcessor, FlaxCLIPModel
+
+import utils
+
+BASELINE_MODEL = "openai/clip-vit-base-patch32"
+MODEL_PATH = "flax-community/clip-rsicd-v2"
+IMAGE_VECTOR_FILE = "./vectors/test-bs128x8-lr5e-6-adam-ckpt-1.tsv"
+IMAGES_DIR = "./images"
+CAPTIONS_FILE = os.path.join(IMAGES_DIR, "test-captions.json")
+
+@st.cache(allow_output_mutation=True)
+def load_example_images():
+    example_images = {}
+    image_names = os.listdir(IMAGES_DIR)
+    for image_name in image_names:
+        if image_name.find("_") < 0:
+            continue
+        image_class = image_name.split("_")[0]
+        if image_class in example_images.keys():
+            example_images[image_class].append(image_name)
+        else:
+            example_images[image_class] = [image_name]
+    example_image_list = sorted([v[np.random.randint(0, len(v))] 
+                                for k, v in example_images.items()][0:10])
+    return example_image_list
+
+
+def get_image_thumbnail(image_filename):
+    image = Image.open(os.path.join(IMAGES_DIR, image_filename))
+    image = image.resize((100, 100))
+    return image
+
+
+def download_and_prepare_image(image_url):
+    try:
+        image_raw = requests.get(image_url, stream=True,).raw
+        image = Image.open(image_raw).convert("RGB")
+        width, height = image.size
+        resize_mult = width / 224 if width < height else height / 224
+        image = image.resize((int(width // resize_mult), 
+                              int(height // resize_mult)))
+        width, height = image.size
+        left = int((width - 224) // 2)
+        top = int((height - 224) // 2)
+        right = int((width + 224) // 2)
+        bottom = int((height + 224) // 2)
+        image = image.crop((left, top, right, bottom))
+        return image
+    except Exception as e:
+        return None
+
+def app():
+    filenames, index = utils.load_index(IMAGE_VECTOR_FILE)
+    model, processor = utils.load_model(MODEL_PATH, BASELINE_MODEL)
+    image2caption = utils.load_captions(CAPTIONS_FILE)
+
+    example_image_list = load_example_images()
+
+    st.title("Retrieve Images given Images")
+    st.markdown("""
+        This demo shows the image to image retrieval capabilities of this model, i.e., 
+        given an image file name as a query, we use our fine-tuned CLIP model 
+        to project the query image to the image/caption embedding space and search 
+        for nearby images (by cosine similarity) in this space.
+        
+        Our fine-tuned CLIP model was previously used to generate image vectors for 
+        our demo, and NMSLib was used for fast vector access.
+
+        Here are some randomly generated image files from our corpus, that you can 
+        find similar images for by selecting the button below it. Alternatively you
+        can upload your own image from the Internet.
+    """)
+
+    suggest_idx = -1
+    col0, col1, col2, col3, col4 = st.columns(5)
+    col0.image(get_image_thumbnail(example_image_list[0]))
+    col1.image(get_image_thumbnail(example_image_list[1]))
+    col2.image(get_image_thumbnail(example_image_list[2]))
+    col3.image(get_image_thumbnail(example_image_list[3]))
+    col4.image(get_image_thumbnail(example_image_list[4]))
+    col0t, col1t, col2t, col3t, col4t = st.columns(5)
+    with col0t:
+        if st.button("Image-1"):
+            suggest_idx = 0
+    with col1t:
+        if st.button("Image-2"):
+            suggest_idx = 1
+    with col2t:
+        if st.button("Image-3"):
+            suggest_idx = 2
+    with col3t:
+        if st.button("Image-4"):
+            suggest_idx = 3
+    with col4t:
+        if st.button("Image-5"):
+            suggest_idx = 4
+    col5, col6, col7, col8, col9 = st.columns(5)
+    col5.image(get_image_thumbnail(example_image_list[5]))
+    col6.image(get_image_thumbnail(example_image_list[6]))
+    col7.image(get_image_thumbnail(example_image_list[7]))
+    col8.image(get_image_thumbnail(example_image_list[8]))
+    col9.image(get_image_thumbnail(example_image_list[9]))
+    col5t, col6t, col7t, col8t, col9t = st.columns(5)
+    with col5t:
+        if st.button("Image-6"):
+            suggest_idx = 5
+    with col6t:
+        if st.button("Image-7"):
+            suggest_idx = 6
+    with col7t:
+        if st.button("Image-8"):
+            suggest_idx = 7
+    with col8t:
+        if st.button("Image-9"):
+            suggest_idx = 8
+    with col9t:
+        if st.button("Image-10"):
+            suggest_idx = 9
+
+    image_url = st.text_input(
+        "OR provide an image URL",
+        value="https://static.eos.com/wp-content/uploads/2019/04/Main.jpg")
+    
+    submit_button = st.button("Find Similar")
+    
+    if submit_button or suggest_idx > -1:
+        image_name = None
+        if suggest_idx > -1:
+            image_name = example_image_list[suggest_idx]
+            image = Image.fromarray(plt.imread(os.path.join(IMAGES_DIR, image_name)))
+        else:
+            image = download_and_prepare_image(image_url)
+            st.image(image, caption="Input Image")
+            st.markdown("---")
+
+        if image is None:
+            st.error("Image could not be downloaded, please try another one!")
+        else:
+            inputs = processor(images=image, return_tensors="jax", padding=True)
+            query_vec = model.get_image_features(**inputs)
+            query_vec = np.asarray(query_vec)
+            ids, distances = index.knnQuery(query_vec, k=11)
+            result_filenames = [filenames[id] for id in ids]
+            rank = 0
+            for result_filename, score in zip(result_filenames, distances):
+                if image_name is not None and result_filename == image_name:
+                    continue
+                caption = "{:s} (score: {:.3f})".format(result_filename, 1.0 - score)
+                col1, col2, col3 = st.columns([2, 10, 10])
+                col1.markdown("{:d}.".format(rank + 1))
+                col2.image(Image.open(os.path.join(IMAGES_DIR, result_filename)),
+                        caption=caption)
+                caption_text = []
+                for caption in image2caption[result_filename]:
+                    caption_text.append("* {:s}\n".format(caption))
+                col3.markdown("".join(caption_text))                       
+                rank += 1
+                st.markdown("---")
+            suggest_idx = -1

dashboard_text2image.py

@@ -0,0 +1,88 @@
+import matplotlib.pyplot as plt
+import nmslib
+import numpy as np
+import os
+import streamlit as st
+
+from PIL import Image
+from transformers import CLIPProcessor, FlaxCLIPModel
+
+import utils
+
+BASELINE_MODEL = "openai/clip-vit-base-patch32"
+MODEL_PATH = "flax-community/clip-rsicd-v2"
+IMAGE_VECTOR_FILE = "./vectors/test-bs128x8-lr5e-6-adam-ckpt-1.tsv"
+IMAGES_DIR = "./images"
+CAPTIONS_FILE = os.path.join(IMAGES_DIR, "test-captions.json")
+
+def app():
+    filenames, index = utils.load_index(IMAGE_VECTOR_FILE)
+    model, processor = utils.load_model(MODEL_PATH, BASELINE_MODEL)
+    image2caption = utils.load_captions(CAPTIONS_FILE)
+
+    st.title("Retrieve Images given Text")
+    st.markdown("""
+        This demo shows the image to text retrieval capabilities of this model, i.e., 
+        given a text query, we use our fine-tuned CLIP model to project the text query 
+        to the image/caption embedding space and search for nearby images (by 
+        cosine similarity) in this space.
+        
+        Our fine-tuned CLIP model was previously used to generate image vectors for 
+        our demo, and NMSLib was used for fast vector access.
+
+    """)
+    suggested_query = [
+        "ships",
+        "school house",
+        "military installation",
+        "mountains",
+        "beaches",
+        "airports",
+        "lakes"
+    ]
+    st.text("Some suggested queries to start you off with...")
+    col0, col1, col2, col3, col4, col5, col6 = st.columns(7)
+        # [1, 1.1, 1.3, 1.1, 1, 1, 1])
+    suggest_idx = -1
+    with col0:
+        if st.button(suggested_query[0]):
+            suggest_idx = 0
+    with col1:
+        if st.button(suggested_query[1]):
+            suggest_idx = 1
+    with col2:
+        if st.button(suggested_query[2]):
+            suggest_idx = 2
+    with col3:
+        if st.button(suggested_query[3]):
+            suggest_idx = 3
+    with col4:
+        if st.button(suggested_query[4]):
+            suggest_idx = 4
+    with col5:
+        if st.button(suggested_query[5]):
+            suggest_idx = 5
+    with col6:
+        if st.button(suggested_query[6]):
+            suggest_idx = 6
+    query = st.text_input("OR enter a text Query:")
+    query = suggested_query[suggest_idx] if suggest_idx > -1 else query
+
+    if st.button("Query") or suggest_idx > -1:
+        inputs = processor(text=[query], images=None, return_tensors="jax", padding=True)
+        query_vec = model.get_text_features(**inputs)
+        query_vec = np.asarray(query_vec)
+        ids, distances = index.knnQuery(query_vec, k=10)
+        result_filenames = [filenames[id] for id in ids]
+        for rank, (result_filename, score) in enumerate(zip(result_filenames, distances)):
+            caption = "{:s} (score: {:.3f})".format(result_filename, 1.0 - score)
+            col1, col2, col3 = st.columns([2, 10, 10])
+            col1.markdown("{:d}.".format(rank + 1))
+            col2.image(Image.open(os.path.join(IMAGES_DIR, result_filename)),
+                       caption=caption)
+            caption_text = []
+            for caption in image2caption[result_filename]:
+                caption_text.append("* {:s}\n".format(caption))
+            col3.markdown("".join(caption_text))                       
+            st.markdown("---")
+        suggest_idx = -1

demo-image-encoder.py

@@ -0,0 +1,69 @@
+import argparse
+import jax
+import jax.numpy as jnp
+import json
+import matplotlib.pyplot as plt
+import numpy as np
+import requests
+import os
+
+from PIL import Image
+from transformers import CLIPProcessor, FlaxCLIPModel
+
+
+def encode_image(image_file, model, processor):
+    image = Image.fromarray(plt.imread(os.path.join(IMAGES_DIR, image_file)))
+    inputs = processor(images=image, return_tensors="jax")
+    image_vec = model.get_image_features(**inputs)
+    return np.array(image_vec).reshape(-1)
+
+
+DATA_DIR = "/home/shared/data"
+IMAGES_DIR = os.path.join(DATA_DIR, "rsicd_images")
+CAPTIONS_FILE = os.path.join(DATA_DIR, "dataset_rsicd.json")
+VECTORS_DIR = os.path.join(DATA_DIR, "vectors")
+BASELINE_MODEL = "openai/clip-vit-base-patch32"
+
+parser = argparse.ArgumentParser()
+parser.add_argument("model_dir", help="Path to model to use for encoding")
+args = parser.parse_args()
+
+print("Loading image list...", end="")
+image2captions = {}
+with open(CAPTIONS_FILE, "r") as fcap:
+    data = json.loads(fcap.read())
+for image in data["images"]:
+    if image["split"] == "test":
+        filename = image["filename"]
+        sentences = []
+        for sentence in image["sentences"]:
+            sentences.append(sentence["raw"])
+        image2captions[filename] = sentences
+    
+print("{:d} images".format(len(image2captions)))
+
+
+print("Loading model...")
+if args.model_dir == "baseline":
+    model = FlaxCLIPModel.from_pretrained(BASELINE_MODEL)
+else:
+    model = FlaxCLIPModel.from_pretrained(args.model_dir)
+processor = CLIPProcessor.from_pretrained(BASELINE_MODEL)
+
+
+model_basename = "-".join(args.model_dir.split("/")[-2:])
+vector_file = os.path.join(VECTORS_DIR, "test-{:s}.tsv".format(model_basename))
+print("Vectors written to {:s}".format(vector_file))
+num_written = 0
+fvec = open(vector_file, "w")
+for image_file in image2captions.keys():
+    if num_written % 100 == 0:
+        print("{:d} images processed".format(num_written))
+    image_vec = encode_image(image_file, model, processor)
+    image_vec_s = ",".join(["{:.7e}".format(x) for x in image_vec])
+    fvec.write("{:s}\t{:s}\n".format(image_file, image_vec_s))
+    num_written += 1
+    
+print("{:d} images processed, COMPLETE".format(num_written))
+fvec.close()
+