update

app.py

@@ -21,9 +21,7 @@ app_ui = ui.page_fillable(
             ui.input_file("image_input", "Upload image: ", multiple=True),
         ),
         ui.output_image("image"),
-        ui.output_image("image_output"),
-        ui.output_image("single_patch_prediction"),
-        ui.output_image("single_patch_prob")
+        ui.output_plot("plot_output"),
     ),
 )
 
@@ -32,7 +30,6 @@ def server(input: Inputs, output: Outputs, session: Session):
     @output
     @render.image
     def image():
-        here = Path(__file__).parent
         if input.image_input():
             # print(input.image_input())
             src = input.image_input()[0]['datapath']
@@ -41,15 +38,26 @@ def server(input: Inputs, output: Outputs, session: Session):
         return None
 
     @output
-    @render.image
-    def image_output():
-        here = Path(__file__).parent
+    @render.plot
+    def plot_output():
         if input.image_input():
             src = input.image_input()[0]['datapath']
-            img = {"src": src, "width": "500px"} 
-            x = run.pred(src)
-            print(x)
-            return img
+            single_patch_prob, single_patch_prediction = run.pred(src)
+            fig, axes = plt.subplots(1, 2, figsize=(15, 5))
+
+            axes[0].imshow(single_patch_prob, cmap='gray')
+            axes[0].set_title("Probability Map")
+
+            im = axes[1].imshow(single_patch_prediction) 
+            axes[1].set_title("Prediction")
+            cbar = fig.colorbar(im, ax=axes[1])
+
+            for ax in axes:
+                ax.set_xticks([])
+                ax.set_yticks([])
+                ax.set_xticklabels([])
+                ax.set_yticklabels([])
+            return fig
         return None
 
 

run.py

@@ -5,22 +5,7 @@ import app
 import os
 from PIL import Image
 
-
-# def customized_patchify(large_image):
-#     print(len(large_image), len(large_image[0]))
-#     patch_size = 256
-#     all_img_patches = []
-#     patches_img = patchify(large_image, (patch_size, patch_size), step=256)  #Step=256 for 256 patches means no overlap
-#     for i in range(patches_img.shape[0]):
-#         for j in range(patches_img.shape[1]):
-#             single_patch_img = patches_img[i,j,:,:]
-#             all_img_patches.append(single_patch_img)
-#     # image = np.array(all_img_patches)
-#     return all_img_patches
-
-
 def pred(src):
-    # os.environ['HUGGINGFACE_HUB_HOME'] = './.cache'
     # Load the model configuration
     cache_dir = "/code/cache"
 
@@ -31,8 +16,8 @@ def pred(src):
     my_sam_model = SamModel(config=model_config)
     # #Update the model by loading the weights from saved file.
     my_sam_model.load_state_dict(torch.load("sam_model.pth", map_location=torch.device('cpu')))
-    print(src)
-    new_image = np.array(Image.open(src))
+
+    new_image = np.array(Image.open(src).convert("RGB"))
     inputs = processor(new_image, return_tensors="pt")
     my_sam_model.eval()
 
@@ -45,60 +30,5 @@ def pred(src):
     # # convert soft mask to hard mask
     single_patch_prob = single_patch_prob.cpu().numpy().squeeze()
     single_patch_prediction = (single_patch_prob > 0.5).astype(np.uint8)
-    # patches = customized_patchify(new_image)
-
-    # # Define the size of your array
-    # array_size = 256
-
-    # # Define the size of your grid
-    # grid_size = 10
-
-    # # Generate the grid points
-    # x = np.linspace(0, array_size-1, grid_size)
-    # y = np.linspace(0, array_size-1, grid_size)
-
-    # # Generate a grid of coordinates
-    # xv, yv = np.meshgrid(x, y)
-
-    # # Convert the numpy arrays to lists
-    # xv_list = xv.tolist()
-    # yv_list = yv.tolist()
-
-    # # Combine the x and y coordinates into a list of list of lists
-    # input_points = [[[int(x), int(y)] for x, y in zip(x_row, y_row)] for x_row, y_row in zip(xv_list, yv_list)]
-    # input_points = torch.tensor(input_points).view(1, 1, grid_size*grid_size, 2)
-
-    # i, j = 1, 2
-
-    # # Selectelected patch for segmentation
-    # random_array = patches[i, j]
 
-    # single_patch = Image.fromarray(random_array)
-    # inputs = processor(single_patch, input_points=input_points, return_tensors="pt")
-
-    # inputs = {k: v.to(device) for k, v in inputs.items()}
-
-    # my_mito_model.eval()
-
-    # # forward pass
-    # with torch.no_grad():
-    #     outputs = my_mito_model(**inputs, multimask_output=False)
-
-    # # apply sigmoid
-    # single_patch_prob = torch.sigmoid(outputs.pred_masks.squeeze(1))
-    # # convert soft mask to hard mask
-    # single_patch_prob = single_patch_prob.cpu().numpy().squeeze()
-    # single_patch_prediction = (single_patch_prob > 0.5).astype(np.uint8)
-
-    x = 1
-    # my_sam_model.eval()
-    # # forward pass
-    # with torch.no_grad():
-    #     outputs = my_sam_model(**inputs, multimask_output=False)
-
-    # # apply sigmoid
-    # single_patch_prob = torch.sigmoid(outputs.pred_masks.squeeze(1))
-    # # convert soft mask to hard mask
-    # single_patch_prob = single_patch_prob.cpu().numpy().squeeze()
-    # single_patch_prediction = (single_patch_prob > 0.5).astype(np.uint8)
-    return x
+    return single_patch_prob, single_patch_prediction