Change layout

app.py

@@ -9,6 +9,7 @@ import requests
 from PIL import Image
 import numpy as np
 import cv2
+import pathlib
 
 
 # create heatmap from mask on image
@@ -113,28 +114,62 @@ def _load_model(model_name: str):
 	model_finetuned.eval()
 	
 _load_model('ar_base.tar')
-#demo = gr.Interface(image_classifier, gr.inputs.Image(shape=(224,224)), [gr.outputs.Label(label="Our Classification", num_top_classes=3), gr.Image(label="Our Relevance",shape=(64,64)), gr.outputs.Label(label="Original Classification", num_top_classes=3), gr.Image(label="Original Relevance",shape=(64,64))],examples=["samples/augreg_base/tank.png", "samples/augreg_base/sundial.png", "samples/augreg_base/lizard.png", "samples/augreg_base/storck.png", "samples/augreg_base/hummingbird2.png", "samples/augreg_base/hummingbird.png"], capture_session=True)
-#demo.launch(debug=True)
 
-demo = gr.Blocks()
+def _set_example_image(example: list) -> dict:
+    return gr.Image.update(value=example[0])
+	
+def _clear_image():
+	return None
+
+demo = gr.Blocks(css='style.css')
 
 with demo:
-	gr.Markdown('Select an image and then click **submit** to see the output.')
 	
-	with gr.Row():
-		inp = gr.inputs.Image(shape=(224,224))
 	
-	outs = []
 	with gr.Row():
-		out1 = gr.outputs.Label(label="Our Classification", num_top_classes=3)
-		out2 = gr.Image(label="Our Relevance",shape=(224,224))
+		with gr.Column():
+			gr.Markdown('## [Optimizing Relevance Maps of Vision Transformers Improves Robustness](https://github.com/hila-chefer/RobustViT) - Official Demo')
+			# gr.Markdown('This is an official demo for [Optimizing Relevance Maps of Vision Transformers Improves Robustness](https://github.com/hila-chefer/RobustViT).')
+			gr.Markdown('Select or upload an image and then click **Submit** to see the output.')
+			with gr.Row():
+				input_image = gr.Image(shape=(224,224))
+			with gr.Row():	
+				btn = gr.Button("Submit", variant="primary")
+				clear_btn = gr.Button('Clear')
+		with gr.Column():
+			gr.Markdown('### Examples')
+			gr.Markdown('#### Corrected Prediction')
+			with gr.Row():
+				paths = sorted(pathlib.Path('samples/corrected').rglob('*.png'))
+				corrected_pred_examples = gr.Dataset(components=[input_image], headers=['header'],
+											samples=[[path.as_posix()] for path in paths])
+		
+			gr.Markdown('#### Improved Explainability')
+			with gr.Row():
+				paths = sorted(pathlib.Path('samples/better_expl').rglob('*.png'))
+				better_expl = gr.Dataset(components=[input_image], headers=['header'],
+											samples=[[path.as_posix()] for path in paths])
 	
+	
+	#gr.Markdown('### Results:')
+		
 	with gr.Row():
-		out3 = gr.outputs.Label(label="Original Classification", num_top_classes=3)
-		out4 = gr.Image(label="Original Relevance",shape=(224,224))
+		with gr.Column():
+			gr.Markdown('### Ours (finetuned model)')
+			out1 = gr.outputs.Label(label="Our Classification", num_top_classes=3)
+			out2 = gr.Image(label="Our Relevance",shape=(224,224), elem_id="expl1")
+			
+		with gr.Column():
+			gr.Markdown('### Original model')
+			out3 = gr.outputs.Label(label="Original Classification", num_top_classes=3)
+			out4 = gr.Image(label="Original Relevance",shape=(224,224),elem_id="expl2")
+		
+	
+	corrected_pred_examples.click(fn=_set_example_image, inputs=corrected_pred_examples, outputs=input_image)
+	better_expl.click(fn=_set_example_image, inputs=better_expl, outputs=input_image)
+	btn.click(fn=image_classifier, inputs=input_image, outputs=[out1, out2, out3, out4])
+	clear_btn.click(fn=_clear_image, inputs=[], outputs=[input_image])
 	
-	btn = gr.Button('Submit')
-	btn.click(fn=image_classifier, inputs=inp, outputs=[out1, out2, out3, out4])
 
 demo.launch()
 	

style.css

@@ -0,0 +1,9 @@
+div#expl1 {
+	max-width: 224px;
+	max-height: 224px;
+}
+
+div#expl2 {
+	max-width: 224px;
+	max-height: 224px;
+}