Update app.py

app.py

@@ -24,8 +24,13 @@ def process_image(image, prompt, threshold, alpha_value, draw_rectangles):
         preds = outputs.logits
 
     pred = torch.sigmoid(preds)
-    mat = pred.cpu().numpy()
-    mask = Image.fromarray(np.uint8(mat[0, 0] * 255), "L")  # Access the first channel of the output
+    
+    if len(pred.shape) == 4:  # Check if the shape is (batch_size, channels, height, width)
+        mat = pred[0, 0].cpu().numpy()  # Access the first channel of the first batch
+    else:
+        mat = pred[0].cpu().numpy()  # If the shape is (channels, height, width)
+
+    mask = Image.fromarray(np.uint8(mat * 255), "L")  # Convert to PIL Image
 
     # normalize the mask
     mask_min = mask.min()
@@ -42,6 +47,7 @@ def process_image(image, prompt, threshold, alpha_value, draw_rectangles):
     return bmask
 
 
+
 @app.route('/')
 def index():
     return "Hello, World! clipseg2"