add description and fast option

CLIP/clip/__init__.py

@@ -30,6 +30,18 @@ saliency_configs = {
             {"tile_size": img_dim // 4, "stride": (img_dim // 4) // 4},
         ],
     },
+    "ours_fast": lambda img_dim: {
+        "distractor_labels": {},
+        "horizontal_flipping": True,
+        "augmentations": 2,
+        "imagenet_prompt_ensemble": False,
+        "positive_attn_only": True,
+        "cropping_augmentations": [
+            {"tile_size": img_dim, "stride": img_dim // 4},
+            {"tile_size": int(img_dim * 2 / 3), "stride": int(img_dim * 2 / 3) // 4},
+            {"tile_size": img_dim // 2, "stride": (img_dim // 2) // 4},
+        ],
+    },
     "chefer_et_al": lambda img_dim: {
         "distractor_labels": {},
         "horizontal_flipping": False,

app.py

@@ -4,7 +4,7 @@ from CLIP.clip import ClipWrapper, saliency_configs
 from time import time
 from matplotlib import pyplot as plt
 import io
-from PIL import Image
+from PIL import Image, ImageDraw, ImageFont
 
 
 def plot_to_png(fig):
@@ -15,12 +15,32 @@ def plot_to_png(fig):
     return img
 
 
+def add_text_to_image(
+    image: np.ndarray,
+    text,
+    position,
+    color="rgb(255, 255, 255)",
+    fontsize=60,
+):
+    image = Image.fromarray(image)
+    draw = ImageDraw.Draw(image)
+    draw.text(
+        position,
+        text,
+        fill=color,
+        font=ImageFont.truetype(
+            "/usr/share/fonts/truetype/lato/Lato-Medium.ttf", fontsize
+        ),
+    )
+    return np.array(image)
+
+
 def generate_relevancy(
     img: np.array, labels: str, prompt: str, saliency_config: str, subtract_mean: bool
 ):
     labels = labels.split(",")
     prompts = [prompt]
-    img = np.asarray(Image.fromarray(img).resize((244 * 2, 244 * 2)))
+    img = np.asarray(Image.fromarray(img).resize((244 * 4, 244 * 4)))
     assert img.dtype == np.uint8
     h, w, c = img.shape
     grads = ClipWrapper.get_clip_saliency(
@@ -38,14 +58,17 @@ def generate_relevancy(
 
     returns = []
     for label_grad, label in zip(grads, labels):
-        fig, ax = plt.subplots(1, 1)
+        fig, ax = plt.subplots(1, 1, figsize=(4, 4))
         ax.axis("off")
         ax.imshow(img)
-        ax.set_title(label, fontsize=12)
         grad = np.clip((label_grad - vmin) / (vmax - vmin), a_min=0.0, a_max=1.0)
         colored_grad = cmap(grad)
         grad = 1 - grad
         colored_grad[..., -1] = grad * 0.7
+        colored_grad = add_text_to_image(
+            (colored_grad * 255).astype(np.uint8), text=label, position=(0, 0)
+        )
+        colored_grad = colored_grad.astype(float) / 255
         ax.imshow(colored_grad)
         plt.tight_layout(pad=0)
         returns.append(plot_to_png(fig))
@@ -54,6 +77,10 @@ def generate_relevancy(
 
 
 iface = gr.Interface(
+    title="Semantic Abstraction Multi-scale Relevancy Extractor",
+    description="""A CPU-only demo of [Semantic Abstraction](https://semantic-abstraction.cs.columbia.edu/)'s Multi-Scale Relevancy Extractor. To run GPU inference locally, use the [official codebase release](https://github.com/columbia-ai-robotics/semantic-abstraction).
+
+This relevancy extractor builds heavily on [Chefer et al.'s codebase](https://github.com/hila-chefer/Transformer-MM-Explainability) and [CLIP on Wheels' codebase](https://cow.cs.columbia.edu/).""",
     fn=generate_relevancy,
     inputs=[
         gr.Image(type="numpy", label="Image"),
@@ -61,7 +88,7 @@ iface = gr.Interface(
         gr.Textbox(label="Prompt"),
         gr.Dropdown(
             value="ours",
-            choices=["ours", "chefer_et_al"],
+            choices=["ours", "ours_fast", "chefer_et_al"],
             label="Relevancy Configuration",
         ),
         gr.Checkbox(value=True, label="subtract mean"),
@@ -69,12 +96,33 @@ iface = gr.Interface(
     outputs=gr.Gallery(label="Relevancy Maps", type="numpy"),
     examples=[
         [
-            "https://semantic-abstraction.cs.columbia.edu/downloads/matterport.png",
+            "https://semantic-abstraction.cs.columbia.edu/downloads/gameroom.png",
             "basketball jersey,nintendo switch,television,ping pong table,vase,fireplace,abstract painting of a vespa,carpet,wall",
             "a photograph of a {} in a home.",
             "ours",
             True,
-        ]
+        ],
+        [
+            "https://semantic-abstraction.cs.columbia.edu/downloads/livingroom.png",
+            "monopoly boardgame set,door knob,sofa,coffee table,plant,carpet,wall",
+            "a photograph of a {} in a home.",
+            "ours",
+            True,
+        ],
+        [
+            "https://semantic-abstraction.cs.columbia.edu/downloads/fireplace.png",
+            "fireplace,beige armchair,candle,large indoor plant in a pot,forest painting,cheetah-patterned pillow,floor,carpet,wall",
+            "a photograph of a {} in a home.",
+            "ours",
+            True,
+        ],
+        [
+            "https://semantic-abstraction.cs.columbia.edu/downloads/walle.png",
+            "WALL-E,a fire extinguisher",
+            "a 3D render of {}.",
+            "ours",
+            True,
+        ],
     ],
 )
 # iface.launch(share=True)