Update app.py

app.py

@@ -1,7 +1,8 @@
 import gradio as gr
 from PIL import Image
 from torchvision import transforms
-from visualization import generate_visualization
+
+from explain import do_explain
 
 normalize = transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])
 TRANSFORM = transforms.Compose(
@@ -21,10 +22,10 @@ def generate_viz(image, class_index=None, use_threshold=False):
     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)
     
-    return Image.open(viz).convert("RGB")
+    viz = do_explain(TRANSFORM, image, class_index=class_index, use_threshold=use_threshold)
+    viz.savefig("visualization.png")
+    return Image.open("visualization.png").convert("RGB")
 
 title = "Explain ViT 😊"