fixed some bugs

app.py

@@ -50,16 +50,23 @@ def greet(image):
 
     model.eval()
     with torch.no_grad():
-        outputs = model(**inputs,multimask_outputs=False)
+        outputs = model(**inputs, multimask_outputs=False)
     
     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)
+    # Ensure the array is 2D (height, width) for grayscale image
+    if seg_prob.ndim > 2:
+        seg_prob = seg_prob.squeeze()  # Remove extra dimensions if any
+    elif seg_prob.ndim < 2:
+        raise ValueError("Output mask has less than 2 dimensions")
+
+    # Convert the processed mask back to a PIL image
+    seg_prob_image = Image.fromarray(seg_prob)
+
+    return seg_prob_image
 
-    return seg_prob
 
 
 iface = gr.Interface(fn= greet, inputs="image", outputs="image", title="Greeter")