option to subtract mean

app.py

@@ -15,7 +15,9 @@ def plot_to_png(fig):
     return img
 
 
-def generate_relevancy(img: np.array, labels: str, prompt: str, saliency_config: str):
+def generate_relevancy(
+    img: np.array, labels: str, prompt: str, saliency_config: str, subtract_mean: bool
+):
     labels = labels.split(",")
     prompts = [prompt]
     assert img.dtype == np.uint8
@@ -27,11 +29,15 @@ def generate_relevancy(img: np.array, labels: str, prompt: str, saliency_config:
         prompts=prompts,
         **saliency_configs[saliency_config](h),
     )[0]
-    grads -= grads.mean(axis=0)
+    if subtract_mean:
+        grads -= grads.mean(axis=0)
     grads = grads.cpu().numpy()
     num_axes = int(np.ceil(np.sqrt(len(labels))))
     fig, axes = plt.subplots(num_axes, num_axes)
-    axes = axes.flatten()
+    if num_axes == 1:
+        axes = [axes]
+    else:
+        axes = axes.flatten()
     vmin = 0.002
     cmap = plt.get_cmap("jet")
     vmax = 0.008
@@ -53,10 +59,15 @@ def generate_relevancy(img: np.array, labels: str, prompt: str, saliency_config:
 iface = gr.Interface(
     fn=generate_relevancy,
     inputs=[
-        gr.Image(type="numpy"),
-        gr.Textbox(),
-        gr.Textbox(),
-        gr.Dropdown(value="ours", choices=["ours", "chefer_et_al"]),
+        gr.Image(type="numpy", label="Image"),
+        gr.Textbox(label="Labels (comma separated)"),
+        gr.Textbox(label="Prompt"),
+        gr.Dropdown(
+            value="ours",
+            choices=["ours", "chefer_et_al"],
+            label="Relevancy Configuration",
+        ),
+        gr.Checkbox(value=True, label="subtract mean"),
     ],
     outputs=gr.Image(type="numpy"),
     examples=[