Update app.py

app.py

@@ -53,7 +53,7 @@ app_ui = ui.page_fillable(
             ui.input_switch("show_margins", "Show marginal plots", value=True),
         ),
         ui.output_image("uploaded_image"),  # display the uploaded TIFF sidewalk tile image
-        ui.output_text("processed_output"),
+        ui.output_plot("prediction_plots"),
         ui.output_ui("value_boxes"),
         ui.output_plot("scatter", fill=True),
         ui.help_text(
@@ -176,8 +176,6 @@ def server(input: Inputs, output: Outputs, session: Session):
         """ Prepare Inputs """
         # get input points prompt (grid of points)
         input_points = generate_input_points(image)
-        print("INPUT POINTS SHAPE", np.array(input_points).shape)
-        print("IMAGE DIMENSIONS", image.width, image.height)
         
         # prepare image and prompt for the model
         inputs = processor(image, input_points=input_points, return_tensors="pt")
@@ -199,66 +197,53 @@ def server(input: Inputs, output: Outputs, session: Session):
         prob = prob.cpu().numpy().squeeze()
         prediction = (prob > 0.5).astype(np.uint8)
         
-        # fig, axes = plt.subplots(1, 5, figsize=(15, 5))
-        
-        # # Extract the image data from the batch
-        # image_data = batch['image'].cpu().detach().numpy()[0]  # Assuming batch size is 1
-        
-        # # Plot the first image on the left
-        # axes[0].imshow(image_data)  
-        # axes[0].set_title("Image")
-        
-        # # Plot the second image on the right
-        # axes[1].imshow(prob)
-        # axes[1].set_title("Probability Map")
-        
-        # # Plot the prediction image on the right
-        # axes[2].imshow(prediction)
-        # axes[2].set_title("Prediction")
-        
-        # # Plot the predicted mask on the right
-        # axes[3].imshow(image_data) 
-        # show_mask(prediction, axes[3])
-        # axes[3].set_title("Predicted Mask")
-        
-        # # Extract the ground truth mask data from the batch
-        # ground_truth_mask_data = inputs['ground_truth_mask'].cpu().detach().numpy()[0]  # Assuming batch size is 1
-        
-        # # Plot the ground truth mask on the right
-        # axes[4].imshow(image_data) 
-        # axes[4].imshow(ground_truth_mask_data)
-        # #show_mask(inputs['ground_truth_mask'], axes[4])
-        # axes[4].set_title("Ground Truth Mask")
-        
-        # # Hide axis ticks and labels
-        # for ax in axes:
-        # ax.set_xticks([])
-        # ax.set_yticks([])
-        # ax.set_xticklabels([])
-        # ax.set_yticklabels([])
-        
-        # # Display the images side by side
-        # plt.show()
-
-        # Evaluate the image with the model
-        # Example: predictions = model.predict(image_array)
-        
-        # Return the processed result (replace 'result' with the actual processed result)
-        return "Processed result"
+        # Return the processed result 
+        return image, prob, prediction
 
     @reactive.Calc
-    def processed_result():
-        """Processes the image when uploaded"""
+    def get_predictions():
+        """Processes the image when uploaded to get predictions"""
         if input.tile_image() is not None:
             return process_image()
         else:
             return None
 
     @output
-    @render.text
-    def processed_output():
-        """Displays the predictions of the uploaded image"""
-        return processed_result()
+    @render.plot
+    def prediction_plots():
+        # get prediction results when image is uploaded
+        image, prob, prediction = get_predictions()
+        
+        fig, axes = plt.subplots(1, 4, figsize=(15, 5))
+        
+        # Extract the image data
+        image_data = image.cpu().detach().numpy()
+        
+        # Plot the first image on the left
+        axes[0].imshow(image_data)  
+        axes[0].set_title("Image")
+        
+        # Plot the probability map on the right
+        axes[1].imshow(prob)
+        axes[1].set_title("Probability Map")
+        
+        # Plot the prediction image on the right
+        axes[2].imshow(prediction)
+        axes[2].set_title("Prediction")
+        
+        # Plot the predicted mask on the right
+        axes[3].imshow(image_data) 
+        show_mask(prediction, axes[3])
+        axes[3].set_title("Predicted Mask")
+        
+        # Hide axis ticks and labels
+        for ax in axes:
+        ax.set_xticks([])
+        ax.set_yticks([])
+        ax.set_xticklabels([])
+        ax.set_yticklabels([])
+        
+        return fig