Spaces:
Running
on
Zero
Running
on
Zero
Update app.py
#6
by
linoyts
HF Staff
- opened
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 |
-
#
|
175 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
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
|