File size: 7,670 Bytes
f4a41d8 |
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 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 |
import random
import modules.scripts as scripts
from modules import images
from modules.processing import process_images, Processed
from modules.processing import Processed
from modules.shared import opts, cmd_opts, state
import gradio as gr
import huggingface_hub
import onnxruntime as rt
import copy
import numpy as np
import cv2
from PIL import Image as im, ImageDraw
# Declare Execution Providers
providers = ['CUDAExecutionProvider', 'CPUExecutionProvider']
# Download and host the model
model_path = huggingface_hub.hf_hub_download(
"skytnt/anime-seg", "isnetis.onnx")
rmbg_model = rt.InferenceSession(model_path, providers=providers)
# Function to get mask
def get_mask(img, s=1024):
# Resize the img to a square shape with dimension s
# Convert img pixel values from integers 0-255 to float 0-1
img = (img / 255).astype(np.float32)
# get the amount of dimensions of img
dim = img.shape[2]
# Convert the input image to RGB format if it has an alpha channel
if dim == 4:
img = img[..., :3]
dim = 3
# Get height and width of the image
h, w = h0, w0 = img.shape[:-1]
# IF height is greater than width, set h as s and w as s*width/height
# ELSE, set w as s and h as s*height/width
h, w = (s, int(s * w / h)) if h > w else (int(s * h / w), s)
# Calculate padding for height and width
ph, pw = s - h, s - w
# Create a 1024x1024x3 array of 0's
img_input = np.zeros([s, s, dim], dtype=np.float32)
# Resize the original image to (w,h) and then pad with the calculated ph,pw
img_input[ph // 2:ph // 2 + h, pw //
2:pw // 2 + w] = cv2.resize(img, (w, h))
# Change the axes
img_input = np.transpose(img_input, (2, 0, 1))
# Add an extra axis (1,0)
img_input = img_input[np.newaxis, :]
# Run the model to get the mask
mask = rmbg_model.run(None, {'img': img_input})[0][0]
# Transpose axis
mask = np.transpose(mask, (1, 2, 0))
# Crop it to the images original dimensions (h0,w0)
mask = mask[ph // 2:ph // 2 + h, pw // 2:pw // 2 + w]
# Resize the mask to original image size (h0,w0)
mask = cv2.resize(mask, (w0, h0))[:, :, np.newaxis]
return mask
### Function to remove background
def rmbg_fn(img):
# Call get_mask() to get the mask
mask = get_mask(img)
# Multiply the image and the mask together to get the output image
img = (mask * img + 255 * (1 - mask)).astype(np.uint8)
# Convert mask value back to int 0-255
mask = (mask * 255).astype(np.uint8)
# Concatenate the output image and mask
img = np.concatenate([img, mask], axis=2, dtype=np.uint8)
# Stacking 3 identical copies of the mask for displaying
mask = mask.repeat(3, axis=2)
return mask, img
class Script(scripts.Script):
is_txt2img = False
# Function to set title
def title(self):
return "ABG Remover"
def ui(self, is_img2img):
with gr.Column():
only_save_background_free_pictures = gr.Checkbox(
label='Only save background free pictures')
do_not_auto_save = gr.Checkbox(label='Do not auto save')
with gr.Row():
custom_background = gr.Checkbox(label='Custom Background')
custom_background_color = gr.ColorPicker(
label='Background Color', default='#ff0000')
custom_background_random = gr.Checkbox(
label='Random Custom Background')
return [only_save_background_free_pictures, do_not_auto_save, custom_background, custom_background_color, custom_background_random]
# Function to show the script
def show(self, is_img2img):
return True
# Function to run the script
def run(self, p, only_save_background_free_pictures, do_not_auto_save, custom_background, custom_background_color, custom_background_random):
# If only_save_background_free_pictures is true, set do_not_save_samples to true
if only_save_background_free_pictures:
p.do_not_save_samples = True
# Create a process_images object
proc = process_images(p)
has_grid = False
unwanted_grid_because_of_img_count = len(
proc.images) < 2 and opts.grid_only_if_multiple
if (opts.return_grid or opts.grid_save) and not p.do_not_save_grid and not unwanted_grid_because_of_img_count:
has_grid = True
# Loop through all the images in proc
for i in range(len(proc.images)):
# Separate the background from the foreground
nmask, nimg = rmbg_fn(np.array(proc.images[i]))
# Check the number of channels in the nimg array, select only the first 3 or 4 channels
num_channels = nimg.shape[2]
if num_channels > 4:
nimg = nimg[:, :, :4]
# Ensure the data type is uint8 and convert the image back to a format that can be saved
nimg = nimg.astype(np.uint8)
img = im.fromarray(nimg)
# If only_save_background_free_pictures is true, check if the image has a background
if custom_background or custom_background_random:
# If custom_background_random is true, set the background color to a random color
if custom_background_random:
custom_background_color = (random.randint(0, 255), random.randint(0, 255), random.randint(0, 255))
# Create a new image with the same size as the original image
background = im.new('RGBA', img.size, custom_background_color)
# Draw a colored rectangle onto the new image
draw = ImageDraw.Draw(background)
draw.rectangle([(0, 0), img.size],
fill=custom_background_color)
# Merge the two images
img = im.alpha_composite(background, img)
# determine output path
outpath = p.outpath_grids if has_grid and i == 0 else p.outpath_samples
# If we are saving all images, save the mask and the image
if not only_save_background_free_pictures:
mask = im.fromarray(nmask)
# Dot not save the new images if checkbox is checked
if not do_not_auto_save:
# Save the new images
images.save_image(
mask, outpath, "mask_", proc.seed + i, proc.prompt, "png", info=proc.info, p=p)
images.save_image(
img, outpath, "img_", proc.seed + i, proc.prompt, "png", info=proc.info, p=p)
# Add the images to the proc object
proc.images.append(mask)
proc.images.append(img)
# If we are only saving background-free images, save the image and replace it in the proc object
else:
proc.images[i] = img
# Check if automatic saving is enabled
if not do_not_auto_save:
# Check if the image is the first one and has a grid
if has_grid and i == 0:
# Save the image
images.save_image(img, p.outpath_grids, "grid", p.all_seeds[0], p.all_prompts[0],
opts.grid_format, info=proc.info, short_filename=not opts.grid_extended_filename, p=p)
else:
# Save the image
images.save_image(img, outpath, "", proc.seed,
proc.prompt, "png", info=proc.info, p=p)
# Return the proc object
return proc
|