fixed some bugs

app.py

@@ -52,14 +52,13 @@ def greet(image):
     with torch.no_grad():
         outputs = model(**inputs,multimask_outputs=False)
     
-    # outputs = outputs.logits[0].cpu().numpy()
-    # outputs = np.argmax(outputs, axis=0)
-    # outputs = Image.fromarray(outputs)
-    # return outputs
     seg_prob = torch.sigmoid(outputs.pred_masks.squeeze(0))
     seg_prob = seg_prob.cpu().numpy().squeeze()
     seg_prob = (seg_prob > 0.5).astype(np.uint8)
 
+    # Add an extra dimension to seg_prob
+    seg_prob = np.expand_dims(seg_prob, axis=-1)
+
     return seg_prob