Add original model results

app.py

@@ -74,6 +74,7 @@ def generate_visualization(model, original_image, class_index=None):
     return vis
 	
 model_finetuned = None
+model = None
 
 normalize = transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])
 transform_224 = transforms.Compose([
@@ -93,10 +94,14 @@ def image_classifier(inp):
 		prediction = torch.nn.functional.softmax(model_finetuned(image.unsqueeze(0))[0], dim=0)
 		confidences = {labels[i]: float(prediction[i]) for i in range(1000)}    
 		heatmap = generate_visualization(model_finetuned, image)
-	return confidences, heatmap
+		
+		prediction_orig = torch.nn.functional.softmax(model(image.unsqueeze(0))[0], dim=0)
+		confidences_orig = {labels[i]: float(prediction_orig[i]) for i in range(1000)}    
+		heatmap_orig = generate_visualization(model, image)
+	return confidences, heatmap, confidences_orig, heatmap_orig
 
 def _load_model(model_name: str):
-	global model_finetuned
+	global model_finetuned, model
 	path = hf_hub_download('Hila/RobustViT',
 						   f'{model_name}')
 						   
@@ -108,5 +113,5 @@ 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=(224,224))],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 = 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=(224,224)), gr.outputs.Label(label="Original Classification", num_top_classes=3), gr.Image(label="Original Relevance",shape=(224,224))],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)