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