Update inference_i2mv_sdxl.py

inference_i2mv_sdxl.py

@@ -151,28 +151,105 @@ def remove_bg(image: Image.Image, net, transform, device, mask: Image.Image = No
 #     return output_image
 
 
+def remove_bg(image: Image.Image, net, transform, device, mask: np.ndarray = None):
+    """
+    Applies a pre-existing mask to an image to make the background transparent.
+
+    Args:
+        image (PIL.Image.Image): The input image.
+        net: Pre-trained neural network (not used but kept for compatibility).
+        transform: Image transformation object (not used but kept for compatibility).
+        device: Device used for inference (not used but kept for compatibility).
+        mask (np.ndarray, optional): The mask to use. Should be the same size
+                                     as the input image, with values between 0 and 255.
+                                     If None, will return image with no changes.
 
+    Returns:
+        PIL.Image.Image: The modified image with transparent background.
+    """
+    if mask is None:
+        return image
 
+    # Ensure the mask is in the correct format
+    if mask.ndim == 2:  # If mask is 2D (H, W)
+        mask = mask.astype(np.uint8)  # Ensure mask is uint8
+        mask = np.expand_dims(mask, axis=-1)  # Add channel dimension
 
-def preprocess_image(image: Image.Image, height, width):
+    # Convert the mask to PIL Image
+    mask_pil = Image.fromarray(mask.squeeze(2) * 255)  # Convert to binary mask
+
+    # Resize the mask to match the original image size
+    mask_pil = mask_pil.resize(image.size, Image.LANCZOS)
+
+    # Create a new image with the same size and mode as the original
+    output_image = Image.new("RGBA", image.size)
+
+    # Apply the mask to the original image
+    image.putalpha(mask_pil)
+
+    # Composite the original image with the mask
+    output_image.paste(image, (0, 0), image)
+
+    return output_image
+
+
+# def preprocess_image(image: Image.Image, height, width):
     
-    alpha = image[..., 3] > 0
-    # alpha = image
-
-    #if image.mode in ("RGBA", "LA"):
-    #    image = np.array(image)
-    #    alpha = image[..., 3]  # Extract the alpha channel
-    #elif image.mode in ("RGB"):
-    #    image = np.array(image)
-        # Create default alpha for non-alpha images
-    #    alpha = np.ones(image[..., 0].shape, dtype=np.uint8) * 255 # Create
+#     alpha = image[..., 3] > 0
+#     # alpha = image
+
+#     #if image.mode in ("RGBA", "LA"):
+#     #    image = np.array(image)
+#     #    alpha = image[..., 3]  # Extract the alpha channel
+#     #elif image.mode in ("RGB"):
+#     #    image = np.array(image)
+#         # Create default alpha for non-alpha images
+#     #    alpha = np.ones(image[..., 0].shape, dtype=np.uint8) * 255 # Create
+#     H, W = alpha.shape
+#     # get the bounding box of alpha
+#     y, x = np.where(alpha)
+#     y0, y1 = max(y.min() - 1, 0), min(y.max() + 1, H)
+#     x0, x1 = max(x.min() - 1, 0), min(x.max() + 1, W)
+#     image_center = image[y0:y1, x0:x1]
+#     # resize the longer side to H * 0.9
+#     H, W, _ = image_center.shape
+#     if H > W:
+#         W = int(W * (height * 0.9) / H)
+#         H = int(height * 0.9)
+#     else:
+#         H = int(H * (width * 0.9) / W)
+#         W = int(width * 0.9)
+#     image_center = np.array(Image.fromarray(image_center).resize((W, H)))
+#     # pad to H, W
+#     start_h = (height - H) // 2
+#     start_w = (width - W) // 2
+#     image = np.zeros((height, width, 4), dtype=np.uint8)
+#     image[start_h : start_h + H, start_w : start_w + W] = image_center
+#     image = image.astype(np.float32) / 255.0
+#     image = image[:, :, :3] * image[:, :, 3:4] + (1 - image[:, :, 3:4]) * 0.5
+#     image = (image * 255).clip(0, 255).astype(np.uint8)
+#     image = Image.fromarray(image)
+
+#     return image
+
+def preprocess_image(image: Image.Image, height, width):
+    # Convert image to numpy array
+    image_np = np.array(image)
+
+    # Extract the alpha channel if present
+    if image_np.shape[-1] == 4:
+        alpha = image_np[..., 3] > 0  # Create a binary mask from the alpha channel
+    else:
+        alpha = np.ones(image_np[..., 0].shape, dtype=bool)  # Default to all true for RGB images
+
     H, W = alpha.shape
-    # get the bounding box of alpha
+    # Get the bounding box of the alpha
     y, x = np.where(alpha)
     y0, y1 = max(y.min() - 1, 0), min(y.max() + 1, H)
     x0, x1 = max(x.min() - 1, 0), min(x.max() + 1, W)
-    image_center = image[y0:y1, x0:x1]
-    # resize the longer side to H * 0.9
+    image_center = image_np[y0:y1, x0:x1]
+
+    # Resize the longer side to H * 0.9
     H, W, _ = image_center.shape
     if H > W:
         W = int(W * (height * 0.9) / H)
@@ -180,18 +257,17 @@ def preprocess_image(image: Image.Image, height, width):
     else:
         H = int(H * (width * 0.9) / W)
         W = int(width * 0.9)
+
     image_center = np.array(Image.fromarray(image_center).resize((W, H)))
-    # pad to H, W
+
+    # Pad to H, W
     start_h = (height - H) // 2
     start_w = (width - W) // 2
-    image = np.zeros((height, width, 4), dtype=np.uint8)
-    image[start_h : start_h + H, start_w : start_w + W] = image_center
-    image = image.astype(np.float32) / 255.0
-    image = image[:, :, :3] * image[:, :, 3:4] + (1 - image[:, :, 3:4]) * 0.5
-    image = (image * 255).clip(0, 255).astype(np.uint8)
-    image = Image.fromarray(image)
+    padded_image = np.zeros((height, width, 4), dtype=np.uint8)
+    padded_image[start_h:start_h + H, start_w:start_w + W] = image_center
 
-    return image
+    # Convert back to PIL Image
+    return Image.fromarray(padded_image)
 
 
 def run_pipeline(