Update app.py

app.py

@@ -25,6 +25,7 @@ numeric_cols: List[str] = df.select_dtypes(include=["float64"]).columns.tolist()
 species: List[str] = df["Species"].unique().tolist()
 species.sort()
 
+### UI ###
 app_ui = ui.page_fillable(
     shinyswatch.theme.minty(),
     ui.layout_sidebar(
@@ -62,6 +63,7 @@ app_ui = ui.page_fillable(
     ),
 )
 
+### HELPER FUNCTIONS ###
 def tif_bytes_to_pil_image(tif_bytes):
   # Create a BytesIO object from the TIFF bytes
   bytes_io = io.BytesIO(tif_bytes)
@@ -89,6 +91,7 @@ def load_model():
 
     return model, processor, device
 
+### SERVER ###
 def server(input: Inputs, output: Outputs, session: Session):
 
     # set model, processor, device once
@@ -163,11 +166,14 @@ def server(input: Inputs, output: Outputs, session: Session):
         """Processes the uploaded image, loads the model, and evaluates to get predictions"""
 
         """ Get Image """
-        # Load the uploaded image
-        uploaded_image_bytes = input.tile_image()[0].read()
-
-        # Convert the uploaded TIFF bytes to a PIL Image object
-        uploaded_image = tif_bytes_to_pil_image(uploaded_image_bytes)
+        img_src = uploaded_image_path()
+        
+        # Read the image bytes from the file
+        with open(img_src, 'rb') as f:
+            image_bytes = f.read()
+        
+        # Convert the image bytes to a PIL Image
+        image = tif_bytes_to_pil_image(image_bytes)
 
         """ Prepare Inputs """
         # get input points prompt (grid of points)
@@ -176,10 +182,64 @@ def server(input: Inputs, output: Outputs, session: Session):
         # prepare image and prompt for the model
         inputs = processor(image, input_points=input_points, return_tensors="pt")
         
-        # remove batch dimension which the processor adds by default
-        inputs = {k:v.squeeze(0) for k,v in inputs.items()}
+        # # remove batch dimension which the processor adds by default
+        # inputs = {k:v.squeeze(0) for k,v in inputs.items()}
 
+        # Move the input tensor to the GPU if it's not already there
+        inputs = {k: v.to(device) for k, v in inputs.items()}
+        
         """ Get Predictions """
+        # forward pass
+        with torch.no_grad():
+        outputs = model(**inputs, multimask_output=False)
+        
+        # apply sigmoid
+        prob = torch.sigmoid(outputs.pred_masks.squeeze(1))
+        # convert soft mask to hard mask
+        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)