Spaces:
Runtime error
Runtime error
File size: 3,091 Bytes
81b1a0e 4adc5f9 6284dc0 e797135 6be00d8 e797135 81b1a0e 53ff575 81b1a0e 621c740 81b1a0e 0972107 81b1a0e 1592dab 81b1a0e 6284dc0 81b1a0e 0972107 d967d62 cb61e6f 0972107 fbe03e2 cb61e6f 0972107 cb61e6f 0972107 1acca69 0972107 741bf59 0972107 cb61e6f 1acca69 0972107 1acca69 0972107 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 |
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()
@spaces.GPU
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) |