Update app.py

app.py

@@ -460,47 +460,52 @@ device = "cuda" if torch.cuda.is_available() else "cpu"
 model.to(device)
 
 def get_crop(image: Image.Image, result, idx: int, padding: int = 15):
-    img_np = np.array(image)
+    img_np = np.array(image)  # shape: (H_full, W_full, 3)
 
     if result.masks is not None:
-        # Original mask (float -> bool)
-        mask = result.masks.data[idx].cpu().numpy()  # shape: (h_orig, w_orig)
-        mask_bool = mask > 0.5
+        # Get the ORIGINAL bounding box (before any upsampling)
+        box = result.boxes.xyxy[idx].cpu().numpy().astype(int)  # [x1, y1, x2, y2]
+        x1, y1, x2, y2 = box
 
-        # Find bounding box of the mask
-        ys, xs = np.where(mask_bool)
-        if len(ys) == 0:
-            return None
+        # Get the mask – but make sure we use the mask at ORIGINAL size
+        # In many cases masks.data[idx] is already at input resolution → we crop it directly
+        mask = result.masks.data[idx].cpu().numpy()  # shape likely (H_full, W_full)
+        mask_bool = mask > 0.5
 
-        y_min, y_max = ys.min(), ys.max()
-        x_min, x_max = xs.min(), xs.max()
+        # Crop both image and mask using the **same box coordinates**
+        crop_img = img_np[y1:y2, x1:x2]          # shape ~ (h_box, w_box, 3)
+        crop_mask = mask_bool[y1:y2, x1:x2]      # shape ~ (h_box, w_box)
 
-        # Apply padding to coordinates
-        y_min_padded = max(0, y_min - padding)
-        y_max_padded = min(img_np.shape[0], y_max + padding + 1)
-        x_min_padded = max(0, x_min - padding)
-        x_max_padded = min(img_np.shape[1], x_max + padding + 1)
+        if crop_img.size == 0 or crop_mask.size == 0:
+            return None
 
-        # Extract padded RGB crop
-        crop = img_np[y_min_padded:y_max_padded, x_min_padded:x_max_padded]
+        # Now apply **padding** around the cropped region
+        h, w = crop_img.shape[:2]
+        pad_top    = min(padding, y1)
+        pad_bottom = min(padding, img_np.shape[0] - y2)
+        pad_left   = min(padding, x1)
+        pad_right  = min(padding, img_np.shape[1] - x2)
 
-        # Now pad the MASK to match the padded crop size
-        # Create full-size zero mask, place original mask in correct position
-        full_mask = np.zeros((img_np.shape[0], img_np.shape[1]), dtype=bool)
-        full_mask[y_min:y_max+1, x_min:x_max+1] = mask_bool
+        # Padded coordinates in full image
+        y_start = y1 - pad_top
+        y_end   = y2 + pad_bottom
+        x_start = x1 - pad_left
+        x_end   = x2 + pad_right
 
-        # Slice the padded region of the mask
-        mask_crop_padded = full_mask[y_min_padded:y_max_padded, x_min_padded:x_max_padded]
+        # Extract padded crops
+        padded_img  = img_np[y_start:y_end, x_start:x_end]
+        padded_mask = mask_bool[y_start:y_end, x_start:x_end]
 
-        # Now shapes match: crop.shape[:2] == mask_crop_padded.shape
-        crop[~mask_crop_padded] = 255
+        # Set background (outside mask) to white
+        padded_img[~padded_mask] = 255
 
-        return Image.fromarray(crop)
+        return Image.fromarray(padded_img)
 
     else:
-        # Bounding box fallback (unchanged)
+        # Bounding box fallback (no mask)
         xyxy = result.boxes.xyxy[idx].cpu().numpy().astype(int)
         x1, y1, x2, y2 = xyxy
+
         x1 = max(0, x1 - padding)
         y1 = max(0, y1 - padding)
         x2 = min(image.width, x2 + padding)
@@ -510,7 +515,6 @@ def get_crop(image: Image.Image, result, idx: int, padding: int = 15):
             return None
 
         return image.crop((x1, y1, x2, y2))
-
         
 def process_image(image: Image.Image):
     if image is None: