Spaces:
Runtime error
Runtime error
| import os | |
| import cv2 | |
| import numpy as np | |
| import torch | |
| import gradio as gr | |
| import spaces | |
| from PIL import Image | |
| from transformers import AutoModelForImageSegmentation | |
| from torchvision import transforms | |
| torch.set_float32_matmul_precision('high') | |
| torch.jit.script = lambda f: f | |
| device = "cuda" if torch.cuda.is_available() else "cpu" | |
| def refine_foreground(image, mask, r=90): | |
| if mask.size != image.size: | |
| mask = mask.resize(image.size) | |
| image = np.array(image) / 255.0 | |
| mask = np.array(mask) / 255.0 | |
| estimated_foreground = FB_blur_fusion_foreground_estimator_2(image, mask, r=r) | |
| image_masked = Image.fromarray((estimated_foreground * 255.0).astype(np.uint8)) | |
| return image_masked | |
| def FB_blur_fusion_foreground_estimator_2(image, alpha, r=90): | |
| alpha = alpha[:, :, None] | |
| F, blur_B = FB_blur_fusion_foreground_estimator( | |
| image, image, image, alpha, r) | |
| return FB_blur_fusion_foreground_estimator(image, F, blur_B, alpha, r=6)[0] | |
| def FB_blur_fusion_foreground_estimator(image, F, B, alpha, r=90): | |
| if isinstance(image, Image.Image): | |
| image = np.array(image) / 255.0 | |
| blurred_alpha = cv2.blur(alpha, (r, r))[:, :, None] | |
| blurred_FA = cv2.blur(F * alpha, (r, r)) | |
| blurred_F = blurred_FA / (blurred_alpha + 1e-5) | |
| blurred_B1A = cv2.blur(B * (1 - alpha), (r, r)) | |
| blurred_B = blurred_B1A / ((1 - blurred_alpha) + 1e-5) | |
| F = blurred_F + alpha * \ | |
| (image - alpha * blurred_F - (1 - alpha) * blurred_B) | |
| F = np.clip(F, 0, 1) | |
| return F, blurred_B | |
| class ImagePreprocessor(): | |
| def __init__(self, resolution=(1024, 1024)) -> None: | |
| self.transform_image = transforms.Compose([ | |
| transforms.Resize(resolution), | |
| transforms.ToTensor(), | |
| transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]), | |
| ]) | |
| def proc(self, image: Image.Image) -> torch.Tensor: | |
| image = self.transform_image(image) | |
| return image | |
| birefnet = AutoModelForImageSegmentation.from_pretrained('zhengpeng7/BiRefNet-matting', trust_remote_code=True) | |
| birefnet.to(device) | |
| birefnet.eval() | |
| def predict(image): | |
| if image is None: | |
| raise gr.Error("Please upload an image.") | |
| image_ori = Image.fromarray(image).convert('RGB') | |
| # Preprocess the image | |
| image_preprocessor = ImagePreprocessor(resolution=(1024, 1024)) | |
| image_proc = image_preprocessor.proc(image_ori) | |
| image_proc = image_proc.unsqueeze(0) | |
| # Prediction | |
| with torch.no_grad(): | |
| preds = birefnet(image_proc.to(device))[-1].sigmoid().cpu() | |
| pred = preds[0].squeeze() | |
| # Show Results | |
| pred_pil = transforms.ToPILImage()(pred) | |
| image_masked = refine_foreground(image_ori, pred_pil) | |
| image_masked.putalpha(pred_pil.resize(image_ori.size)) | |
| torch.cuda.empty_cache() | |
| # Save as PNG | |
| output_path = "output.png" | |
| image_masked.save(output_path) | |
| return output_path | |
| iface = gr.Interface( | |
| fn=predict, | |
| inputs=gr.Image(type="numpy"), | |
| outputs=gr.Image(type="filepath"), | |
| ) | |
| if __name__ == "__main__": | |
| iface.launch(debug=True) |