Files changed (1) hide show
  1. app.py +10 -2
app.py CHANGED
@@ -171,8 +171,16 @@ def remove_background_from_image(image: Image.Image) -> Image.Image:
171
  if result_image.mode != 'RGBA':
172
  result_image = result_image.convert('RGBA')
173
 
174
- # Apply alpha mask
175
- alpha_pil = Image.fromarray((alpha_mask * 255).astype(np.uint8), mode='L')
 
 
 
 
 
 
 
 
176
  result_image.putalpha(alpha_pil)
177
 
178
  return result_image
 
171
  if result_image.mode != 'RGBA':
172
  result_image = result_image.convert('RGBA')
173
 
174
+ # Handle alpha mask dimensions and convert to PIL
175
+ # The alpha_mask might have extra dimensions, so squeeze and ensure 2D
176
+ alpha_mask_2d = np.squeeze(alpha_mask)
177
+ if alpha_mask_2d.ndim > 2:
178
+ # If still more than 2D, take the first channel
179
+ alpha_mask_2d = alpha_mask_2d[:, :, 0] if alpha_mask_2d.shape[-1] == 1 else alpha_mask_2d[:, :, 0]
180
+
181
+ # Convert to uint8 and create PIL Image without deprecated mode parameter
182
+ alpha_array = (alpha_mask_2d * 255).astype(np.uint8)
183
+ alpha_pil = Image.fromarray(alpha_array, 'L')
184
  result_image.putalpha(alpha_pil)
185
 
186
  return result_image