Success!

app.py

@@ -62,18 +62,17 @@ def server(input: Inputs, output: Outputs, session: Session):
 
             pred_prob, pred_prediction = run.pred(new_image)
             
-            fig, axes = plt.subplots(1, 3, figsize=(15, 5))
+            print("plotting...")
+            fig, axes = plt.subplots(1, 2, figsize=(15, 5))
 
-            axes[0].imshow(new_image, cmap='gray')
-            axes[0].set_title("Image")
+            
+            axes[0].imshow(pred_prediction, cmap='gray')
+            axes[0].set_title("Prediction")
 
             im = axes[1].imshow(pred_prob) 
             axes[1].set_title("Probability Map")
             cbar = fig.colorbar(im, ax=axes[1])
 
-            axes[2].imshow(pred_prediction, cmap='gray')
-            axes[2].set_title("Prediction")
-
             for ax in axes:
                 ax.set_xticks([])
                 ax.set_yticks([])

run.py

@@ -29,17 +29,20 @@ def pred(src):
     image = Image.open(src)
     rgbim = image.convert("RGB")
     new_image = np.array(rgbim)
-    print("Shape:",new_image.shape)
+    print()
+    print("image shape:",new_image.shape)
 
     inputs = processor(new_image, return_tensors="pt")
     model.eval()
 
     # forward pass
+    print("predicting...")
     with torch.no_grad():
         outputs = model(pixel_values=inputs["pixel_values"], 
                         multimask_output=False)
     
     # apply sigmoid
+    print("apply sigmoid...")
     pred_prob = torch.sigmoid(outputs.pred_masks.squeeze(1))
 
     # convert soft mask to hard mask
@@ -47,5 +50,4 @@ def pred(src):
     pred_prob = pred_prob.cpu().numpy().squeeze()
     pred_prediction = (pred_prob > PROBABILITY_THRES).astype(np.uint8)
     
-    x=1
     return pred_prob, pred_prediction