Update app.py

app.py

@@ -25,10 +25,6 @@ def process_image(image, prompt, threshold, alpha_value, draw_rectangles):
 
     pred = torch.sigmoid(preds)
     mat = pred.cpu().numpy()
-
-    # Squeeze to remove dimensions of size 1 (batch and channel dimensions)
-    mat = np.squeeze(mat, axis=(0, 1))  # Convert from [1, 1, H, W] to [H, W]
-
     mask = Image.fromarray(np.uint8(mat * 255), "L")
     mask = mask.convert("RGB")
     mask = mask.resize(image.size)
@@ -103,6 +99,6 @@ def mask_image_api():
     _, result_mask, result_output = process_image(image, prompt, threshold, alpha_value, draw_rectangles)
 
     return jsonify({'result_mask': result_mask, 'result_output': result_output})
-    
+
 if __name__ == '__main__':
     app.run(host='0.0.0.0', port=7860, debug=True)