Update

app_lib/main.py

@@ -29,7 +29,7 @@ def main(device=torch.device("cuda" if torch.cuda.is_available() else "cpu")):
         image_col, concepts_col = st.columns(2)
 
         with image_col:
-            image = get_image()
+            image_name, image = get_image()
             st.image(image, use_column_width=True)
 
             change_image_button = st.button(
@@ -42,8 +42,8 @@ def main(device=torch.device("cuda" if torch.cuda.is_available() else "cpu")):
                 st.experimental_rerun()
         with concepts_col:
             model_name = get_model_name()
-            class_name, class_ready, class_error = get_class_name()
-            concepts, concepts_ready, concepts_error = get_concepts()
+            class_name, class_ready, class_error = get_class_name(image_name)
+            concepts, concepts_ready, concepts_error = get_concepts(image_name)
 
         ready = class_ready and concepts_ready
 

app_lib/user_input.py

@@ -1,9 +1,17 @@
+import json
+import os
+
 import streamlit as st
 from PIL import Image
 from streamlit_image_select import image_select
 
 from app_lib.utils import SUPPORTED_DATASETS, SUPPORTED_MODELS
 
+IMAGE_DIR = os.path.join("assets", "images")
+IMAGE_NAMES = list(sorted(filter(lambda x: x.endswith(".jpg"), os.listdir(IMAGE_DIR))))
+IMAGE_PATHS = list(map(lambda x: os.path.join(IMAGE_DIR, x), IMAGE_NAMES))
+IMAGE_PRESETS = json.load(open("assets/image_presets.json"))
+
 
 def _validate_class_name(class_name):
     if class_name is None:
@@ -125,37 +133,49 @@ def get_model_name():
 def get_image():
     with st.sidebar:
         uploaded_file = st.file_uploader("Upload an image", type=["jpg", "png", "jpeg"])
-        image = uploaded_file or image_select(
-            label="or select one",
-            images=[
-                "assets/ace.jpg",
-                "assets/ace.jpg",
-                "assets/ace.jpg",
-                "assets/ace.jpg",
-            ],
-        )
-    return Image.open(image)
-
-
-def get_class_name():
+        if uploaded_file is not None:
+            return (None, Image.open(uploaded_file))
+        else:
+            DEFAULT = IMAGE_NAMES.index("ace.jpg")
+            image_idx = image_select(
+                label="or select one",
+                images=IMAGE_PATHS,
+                index=DEFAULT,
+                return_value="index",
+            )
+            image_name, image_path = IMAGE_NAMES[image_idx], IMAGE_PATHS[image_idx]
+            return (image_name, Image.open(image_path))
+
+
+def get_class_name(image_name=None):
+    DEFAULT = (
+        IMAGE_PRESETS[image_name.split(".")[0]]["class_name"] if image_name else ""
+    )
     class_name = st.text_input(
         "Class to test",
         help="Name of the class to build the zero-shot CLIP classifier with.",
-        value="cat",
+        value=DEFAULT,
         disabled=st.session_state.disabled,
+        placeholder="Type class name here",
     )
 
     class_ready, class_error = _validate_class_name(class_name)
     return class_name, class_ready, class_error
 
 
-def get_concepts():
+def get_concepts(image_name=None):
+    DEFAULT = (
+        "\n".join(IMAGE_PRESETS[image_name.split(".")[0]]["concepts"])
+        if image_name
+        else ""
+    )
     concepts = st.text_area(
         "Concepts to test",
         help="List of concepts to test the predictions of the model with. Write one concept per line. Maximum 10 concepts allowed.",
         height=160,
-        value="piano\ncute\nwhiskers\nmusic\nwild",
+        value=DEFAULT,
         disabled=st.session_state.disabled,
+        placeholder="Type one concept\nper line",
     )
     concepts = concepts.split("\n")
     concepts = [concept.strip() for concept in concepts]

app_lib/viz.py

@@ -41,23 +41,32 @@ def _viz_rank(results):
             y=rank_df["concept"],
             orientation="h",
             marker=dict(color="#a6cee3"),
-            name="Normalized rejection time",
+            name="Rejection time",
         )
     )
-    fig.add_shape(
-        type="line",
-        yref="paper",
-        line=dict(color="black", dash="dash"),
-        x0=significance_level,
-        x1=significance_level,
-        y0=0,
-        y1=1,
-        name="significance level",
-        showlegend=True,
+    fig.add_trace(
+        go.Scatter(
+            x=[significance_level, significance_level],
+            y=[sorted_concepts[0], sorted_concepts[0]],
+            mode="lines",
+            line=dict(color="black", dash="dash"),
+            name="significance level",
+        )
     )
+    fig.add_vline(significance_level, line_dash="dash", line_color="black")
+
     fig.update_layout(yaxis_title="Rank of importance", xaxis_title="")
+    if rank_df["tau"].min() <= 0.3:
+        fig.update_layout(
+            legend=dict(
+                x=0.3,
+                y=1.0,
+                bordercolor="black",
+                borderwidth=1,
+            ),
+        )
 
-    _, centercol, _ = st.columns([1, 4, 1])
+    _, centercol, _ = st.columns([1, 3, 1])
     with centercol:
         st.plotly_chart(fig, use_container_width=True)
 
@@ -86,7 +95,6 @@ def _viz_wealth(results):
         annotation_position="bottom right",
     )
     fig.update_yaxes(range=[0, 1.5 * 1 / significance_level])
-    # fig.update_layout(legend=dict(orientation="h", x=0, y=1.2))
     st.plotly_chart(fig, use_container_width=True)
 
 

assets/image_presets.json

@@ -0,0 +1,52 @@
+{
+    "ace": {
+        "class_name": "cat",
+        "concepts": [
+            "piano",
+            "cute",
+            "whiskers",
+            "music",
+            "wild"
+        ]
+    },
+    "english_springer_1": {
+        "class_name": "English springer",
+        "concepts": [
+            "spaniel",
+            "sibling",
+            "fluffy",
+            "patch",
+            "portrait"
+        ]
+    },
+    "english_springer_2": {
+        "class_name": "English springer",
+        "concepts": [
+            "spaniel",
+            "fetch",
+            "fishing",
+            "trumpet",
+            "cathedral"
+        ]
+    },
+    "french_horn": {
+        "class_name": "French horn",
+        "concepts": [
+            "trumpet",
+            "band",
+            "instrument",
+            "major",
+            "naval"
+        ]
+    },
+    "parachute": {
+        "class_name": "parachute",
+        "concepts": [
+            "flew",
+            "descending",
+            "tandem",
+            "instrument",
+            "band"
+        ]
+    }
+}