Update app.py

app.py

@@ -34,7 +34,12 @@ def predict(image):
         output = model(image_tensor)
     
     wsr, hsr = output.squeeze().tolist()
-    return f"Squish Ratio: (Width, Height)= ({wsr:.2f}, {hsr:.2f})"
+    if(wsr < hsr):
+      height = int(height * wsr)
+    else:
+      width = int(width * hsr)
+    reconstructed_image = TF.resize(image, (height, width))
+    return f"Squish Ratio: (Width, Height)= ({wsr:.2f}, {hsr:.2f})", reconstructed_image
 
 # Define the examples (provide paths to example images)
 examples = [
@@ -47,7 +52,7 @@ examples = [
 iface = gr.Interface(
     fn=predict,
     inputs=gr.Image(type="pil"),
-    outputs="text",
+    outputs=[gr.Textbox(label="Prediction"), gr.Image(type="pil", label="Reconstructed Aspect-Ratio")],
     examples=examples,
     title="Deep Image Squish Predictor",
     description="Upload an image to see the predicted squish ratios."