(
image: PIL.Image.Image,
rembg: Any = None,
force: bool = False,
**rembg_kwargs,
)
| 42 | |
| 43 | @torch.inference_mode() |
| 44 | def remove_background( |
| 45 | image: PIL.Image.Image, |
| 46 | rembg: Any = None, |
| 47 | force: bool = False, |
| 48 | **rembg_kwargs, |
| 49 | ) -> PIL.Image.Image: |
| 50 | do_remove = True |
| 51 | if image.mode == "RGBA" and image.getextrema()[3][0] < 255: |
| 52 | do_remove = False |
| 53 | do_remove = do_remove or force |
| 54 | if do_remove: |
| 55 | W, H = image.size |
| 56 | k = (256.0 / float(H * W)) ** 0.5 |
| 57 | feed = resize_without_crop(image, int(64 * round(W * k)), int(64 * round(H * k))) |
| 58 | feed = numpy2pytorch([feed]).to(device=rembg.device, dtype=torch.float32) |
| 59 | alpha = rembg(feed)[0][0] |
| 60 | alpha = torch.nn.functional.interpolate(alpha, size=(H, W), mode="bilinear") |
| 61 | alpha = alpha.squeeze().clamp(0, 1) |
| 62 | alpha = (alpha * 255).cpu().data.numpy().astype(np.uint8) |
| 63 | alpha = Image.fromarray(alpha) |
| 64 | |
| 65 | no_bg_image = Image.new("RGBA", alpha.size, (0, 0, 0, 0)) |
| 66 | no_bg_image.paste(image, mask=alpha) |
| 67 | image = no_bg_image |
| 68 | return image |
| 69 | |
| 70 | |
| 71 | @torch.inference_mode() |
no test coverage detected