Update app.py

app.py

@@ -14,16 +14,14 @@ TRANSFORM = transforms.Compose(
     ]
 )
 
-use_threshold = [True, False]
 
-def generate_viz(image, class_index=None, use_threshold=False):
+def generate_viz(image, class_index=None):
     if class_index is not None:
         class_index = int(class_index)
     print(f"Image: {image.size}")
     print(f"Class: {class_index}")
-    print(f"use_threshold: {use_threshold}")
     
-    viz = do_explain(TRANSFORM, image, class_index=class_index, use_threshold=use_threshold)
+    viz = do_explain(TRANSFORM, image, class_index=class_index)
     viz.savefig("visualization.png")
     return Image.open("visualization.png").convert("RGB")
 
@@ -31,10 +29,7 @@ title = "Explain ViT 😊"
 
 iface = gr.Interface(fn=generate_viz, inputs=[
         gr.Image(type="pil", label="Input Image"),
-        gr.Dropdown(
-            list(use_threshold),
-            label="use_threshold",
-        ),
+       
         gr.Number(label="Class Index", info="Class index to explain"),
     ],
     outputs=gr.Image(),
@@ -46,7 +41,6 @@ iface = gr.Interface(fn=generate_viz, inputs=[
         ["ViT_DeiT/samples/catdog.png", 243],
         ["ViT_DeiT/samples/el2.png", None],
         ["ViT_DeiT/samples/el2.png", 340],
-        ["ViT_DeiT/samples/dogbird.png", 161],
     ],
 )