Update app.py

app.py

@@ -371,9 +371,9 @@ def predict(inp):
     with torch.no_grad():
         classification, reconstruction = model(image_tensor.unsqueeze(0))
         periods_classification = periods_model(image_tensor.unsqueeze(0))
-        # recon_tensor = reconstruction[0].repeat(3, 1, 1)
-        # recon_tensor = invTrans(kornia.enhance.invert(recon_tensor))
-        # plot_recon = recon_tensor.permute(1, 2, 0).detach().numpy()
+        recon_tensor = reconstruction[0].repeat(3, 1, 1)
+        recon_tensor = invTrans(kornia.enhance.invert(recon_tensor))
+        plot_recon = recon_tensor.permute(1, 2, 0).detach().numpy()
         w, h = inp.size
         # plot_recon = resize(plot_recon, (h, w))
         m = nn.Sigmoid()
@@ -405,7 +405,7 @@ def predict(inp):
                 dists[image_name] = dist
         res = dict(sorted(dists.items(), key=itemgetter(1)))
     fig, names = create_retrieval_figure(res)
-    return confidences, periods_confidences, fig, names
+    return confidences, periods_confidences, plot_recon, fig, names
 
 
 gr.Interface(fn=predict,
@@ -415,4 +415,4 @@ gr.Interface(fn=predict,
                          "https://arxiv.org/abs/2211.09480.",
              examples=['anth.jpg', 'beetle_snakes.jpg', 'bird.jpg', 'cross.jpg', 'ibex.jpg',
                        'lion.jpg', 'lion2.jpg', 'sa.jpg'],
-             outputs=[gr.Label(num_top_classes=1), gr.Label(num_top_classes=1), 'plot', 'text'], ).launch(share=True, enable_queue=True)
+             outputs=[gr.Label(num_top_classes=1), gr.Label(num_top_classes=1), "image", 'plot', 'text'], ).launch(share=True, enable_queue=True)