Update app.py

app.py

@@ -1,7 +1,52 @@
 import gradio as gr
+from PIL import Image
+from torchvision import transforms
+from visualization import generate_visualization
 
-def greet(name):
-    return "Hello " + name + "!!"
+normalize = transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])
+TRANSFORM = transforms.Compose(
+    [
+        transforms.Resize(256),
+        transforms.CenterCrop(224),
+        transforms.ToTensor(),
+        normalize,
+    ]
+)
 
-iface = gr.Interface(fn=greet, inputs="text", outputs="text")
-iface.launch()
+use_threshold = [True, False]
+
+def generate_viz(image, class_index=None, use_threshold=False):
+    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}")
+    image_trans = TRANSFORM(image)
+    viz = generate_visualization(image_trans, class_index=class_index, use_threshold=use_threshold)
+    viz.savefig("visualization.png")
+    return Image.open("visualization.png").convert("RGB")
+
+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(),
+    title=title,
+    allow_flagging="never",
+    cache_examples=True,
+    examples=[
+        ["ViT_DeiT/samples/catdog.png",None],
+        ["ViT_DeiT/samples/catdog.png", 243],
+        ["ViT_DeiT/samples/el2.png", None],
+        ["ViT_DeiT/samples/el2.png", 340],
+        ["ViT_DeiT/samples/dogbird.png", 161],
+    ],
+)
+
+iface.launch(debug=True)