Update app.py

app.py

@@ -24,34 +24,25 @@ def process_image(image, prompt, threshold, alpha_value, draw_rectangles):
         preds = outputs.logits
 
     pred = torch.sigmoid(preds)
-    
-    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
-
-    # Convert the mask to a NumPy array for calculation
-    mask_array = np.array(mask)
+    mat = pred.cpu().numpy()
+    mask = Image.fromarray(np.uint8(mat * 255), "L")
+    mask = mask.resize(image.size)
+    mask = np.array(mask)[:, :, 0]
 
     # normalize the mask
-    mask_min = mask_array.min()
-    mask_max = mask_array.max()
-    mask_array = (mask_array - mask_min) / (mask_max - mask_min)
+    mask_min = mask.min()
+    mask_max = mask.max()
+    mask = (mask - mask_min) / (mask_max - mask_min)
 
     # threshold the mask
-    bmask = mask_array > threshold
+    bmask = mask > threshold
     # zero out values below the threshold
-    mask_array[mask_array < threshold] = 0
+    mask[mask < threshold] = 0
 
-    bmask = Image.fromarray((bmask * 255).astype(np.uint8), "L")
+    bmask = Image.fromarray(bmask.astype(np.uint8) * 255, "L")
 
     return bmask
 
-
-
-
 @app.route('/')
 def index():
     return "Hello, World! clipseg2"