(
data,
model_name="u2net",
alpha_matting=False,
alpha_matting_foreground_threshold=240,
alpha_matting_background_threshold=10,
alpha_matting_erode_structure_size=10,
alpha_matting_base_size=1000,
only_mask=False,
background_color=None,
background_image=None,
mask_threshold=None,
)
| 255 | |
| 256 | |
| 257 | def remove( |
| 258 | data, |
| 259 | model_name="u2net", |
| 260 | alpha_matting=False, |
| 261 | alpha_matting_foreground_threshold=240, |
| 262 | alpha_matting_background_threshold=10, |
| 263 | alpha_matting_erode_structure_size=10, |
| 264 | alpha_matting_base_size=1000, |
| 265 | only_mask=False, |
| 266 | background_color=None, |
| 267 | background_image=None, |
| 268 | mask_threshold=None, |
| 269 | ): |
| 270 | model = get_model(model_name) |
| 271 | |
| 272 | if isinstance(data, np.ndarray): |
| 273 | img = Image.fromarray(data).convert("RGB") |
| 274 | else: |
| 275 | try: |
| 276 | img = Image.open(io.BytesIO(data)) |
| 277 | # Handle EXIF orientation to prevent rotated images (fixes #144) |
| 278 | img = ImageOps.exif_transpose(img) |
| 279 | img = img.convert("RGB") |
| 280 | except Exception as e: |
| 281 | raise ValueError(f"Invalid image input to `remove()`: {e}") |
| 282 | |
| 283 | mask = detect.predict(model, np.array(img)).convert("L") |
| 284 | |
| 285 | # Apply threshold for hard/sharp edges (fixes #122) |
| 286 | if mask_threshold is not None: |
| 287 | mask = mask.point(lambda p: 255 if p > mask_threshold else 0) |
| 288 | |
| 289 | # If only_mask is True, return just the mask |
| 290 | if only_mask: |
| 291 | bio = io.BytesIO() |
| 292 | mask.save(bio, "PNG") |
| 293 | return bio.getbuffer() |
| 294 | |
| 295 | if alpha_matting: |
| 296 | cutout = alpha_matting_cutout( |
| 297 | img, |
| 298 | mask, |
| 299 | alpha_matting_foreground_threshold, |
| 300 | alpha_matting_background_threshold, |
| 301 | alpha_matting_erode_structure_size, |
| 302 | alpha_matting_base_size, |
| 303 | ) |
| 304 | else: |
| 305 | cutout = naive_cutout(img, mask) |
| 306 | |
| 307 | # If background_image is specified, composite over that image |
| 308 | if background_image is not None: |
| 309 | if isinstance(background_image, np.ndarray): |
| 310 | bg = Image.fromarray(background_image).convert("RGB") |
| 311 | else: |
| 312 | try: |
| 313 | bg = Image.open(io.BytesIO(background_image)) |
| 314 | # Handle EXIF orientation for background image too |
no test coverage detected