Update app.py

app.py

@@ -226,10 +226,10 @@ def server(input: Inputs, output: Outputs, session: Session):
         fig, axes = plt.subplots(1, 4, figsize=(15, 5))
         
         # Extract the image data
-        image_data = image.cpu().detach().numpy()
+        #image_data = image.cpu().detach().numpy()
         
         # Plot the first image on the left
-        axes[0].imshow(image_data)  
+        axes[0].imshow(image)  
         axes[0].set_title("Image")
         
         # Plot the probability map on the right
@@ -241,7 +241,7 @@ def server(input: Inputs, output: Outputs, session: Session):
         axes[2].set_title("Prediction")
         
         # Plot the predicted mask on the right
-        axes[3].imshow(image_data) 
+        axes[3].imshow(image) 
         show_mask(prediction, axes[3])
         axes[3].set_title("Predicted Mask")