Spaces:
Running
Running
File size: 4,235 Bytes
97c0d3c 9ff4f09 97c0d3c 9ff4f09 97c0d3c 9ff4f09 97c0d3c 9ff4f09 5585941 9ff4f09 97c0d3c 5585941 97c0d3c 9ff4f09 03dffaa 9ff4f09 97c0d3c 9ff4f09 5585941 9ff4f09 97c0d3c 9ff4f09 |
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 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 |
############################################################################################################
#
# Source from
# https://github.com/eugenesiow/practical-ml/blob/master/notebooks/Remove_Image_Background_DeepLabV3.ipynb
#
############################################################################################################
import os
import cv2
import torch
import PIL.Image
import numpy as np
import gradio as gr
import torchvision.transforms as transforms
os.system("pip freeze")
model = torch.hub.load('pytorch/vision:v0.6.0', 'deeplabv3_resnet101', weights='DEFAULT')
model.eval()
def image_to_tensor(image):
return transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)),
])(image)
def custom_background(background, foreground):
x = (background.size[0] - foreground.size[0]) / 2
y = (background.size[1] - foreground.size[1]) / 2
box = (x, y, foreground.size[0] + x, foreground.size[1] + y)
crop = background.crop(box)
final_image = crop.copy()
# put the foreground in the centre of the background
paste_box = (0, final_image.size[1] - foreground.size[1], final_image.size[0], final_image.size[1])
final_image.paste(foreground, paste_box, mask=foreground)
return final_image
def make_transparent_foreground(image, mask):
# split the image into channels
b, g, r = cv2.split(np.array(image).astype('uint8'))
# add an alpha channel with and fill all with transparent pixels (max 255)
a = np.ones(mask.shape, dtype='uint8') * 255
# merge the alpha channel back
alpha_im = cv2.merge([b, g, r, a], 4)
# create a transparent background
bg = np.zeros(alpha_im.shape)
# set up the new mask
new_mask = np.stack([mask, mask, mask, mask], axis=2)
# copy only the foreground color pixels from the original image where mask is set
return np.where(new_mask, alpha_im, bg).astype(np.uint8)
def makeMask(image):
input_tensor = image_to_tensor(image)
input_batch = input_tensor.unsqueeze(0) # create a mini-batch as expected by the model
# move the input and model to GPU for speed if available
if torch.cuda.is_available():
input_batch = input_batch.to('cuda')
model.to('cuda')
with torch.no_grad():
output = model(input_batch)['out'][0]
output_predictions = output.argmax(0)
# create a binary (black and white) mask of the profile foreground
mask = output_predictions.byte().cpu().numpy()
background = np.zeros(mask.shape)
return np.where(mask, 255, background).astype(np.uint8)
def predict(image, new_background=None):
mask = makeMask(image)
foreground = make_transparent_foreground(image, mask)
if new_background is not None:
foreground = PIL.Image.fromarray(foreground)
return custom_background(new_background, foreground)
return foreground
title = "Zero Background"
description = r"""
## Remove image background
This is another implementation of <a href='https://github.com/eugenesiow/practical-ml/blob/master/notebooks/Remove_Image_Background_DeepLabV3.ipynb' target='_blank'>eugenesiow</a>.
It has no any particular purpose than start research on AI models.
"""
article = r"""
Questions, doubts, comments, please email 📧 `[email protected]`
This demo is running on a CPU, if you like this project please make us a donation to run on a GPU or just give us a <a href='https://github.com/leonelhs/face-shine' target='_blank'>Github ⭐</a>
<a href="https://www.buymeacoffee.com/leonelhs">
<img src="https://img.buymeacoffee.com/button-api/?text=Buy me a coffee&emoji=&slug=leonelhs&button_colour=FFDD00&font_colour=000000&font_family=Cookie&outline_colour=000000&coffee_colour=ffffff" />
</a>
<center><img src='https://visitor-badge.glitch.me/badge?page_id=deoldify.visitor-badge' alt='visitor badge'></center>
"""
demo = gr.Interface(
predict, [
gr.Image(type="pil", label="Image"),
gr.Image(type="pil", label="Optionally: Set a new background")
], [
gr.Image(type="pil", label="Image alpha background")
],
title=title,
description=description,
article=article)
demo.queue().launch()
|