[ { "path": ".gitignore", "content": "scripts/__pycache__" }, { "path": "README.md", "content": "

\nABG extension\n

\n\n

\n\t\n\t\n\t\n

\n\n## Installation\n\n 1. Install extension by going to Extensions tab -> Install from URL -> Paste github URL and click Install.\n 2. After it's installed, go back to the Installed tab in Extensions and press Apply and restart UI.\n 3. Installation finished.\n 4. If the script does not show up or work, please restart the WebUI.\n\n## Usage\n\n### txt2img\n\n 1. In the bottom of the WebUI in Script, select **ABG Remover**.\n 2. Select the desired options: **Only save background free pictures** or **Do not auto save**.\n 3. Generate an image and you will see the result in the output area.\n\n### img2img\n\n 1. In the bottom of the WebUI in Script, select **ABG Remover**.\n 2. Select the desired options: **Only save background free pictures** or **Do not auto save**.\n 3. **IMPORTANT**: Set **Denoising strength** to a low value, like **0.01**\n\nBased on https://huggingface.co/spaces/skytnt/anime-remove-background\n" }, { "path": "install.py", "content": "import launch\n\nfor dep in ['onnx', 'onnxruntime', 'numpy']:\n if not launch.is_installed(dep):\n launch.run_pip(f\"install {dep}\", f\"{dep} for ABG_extension\")\n\nif not launch.is_installed(\"cv2\"):\n launch.run_pip(\"install opencv-python\", \"opencv-python\")\n\nif not launch.is_installed(\"PIL\"):\n launch.run_pip(\"install Pillow\", \"Pillow\")\n" }, { "path": "scripts/app.py", "content": "import random\nimport modules.scripts as scripts\nfrom modules import images\nfrom modules.processing import process_images, Processed\nfrom modules.processing import Processed\nfrom modules.shared import opts, cmd_opts, state\n\nimport gradio as gr\nimport huggingface_hub\nimport onnxruntime as rt\nimport copy\nimport numpy as np\nimport cv2\nfrom PIL import Image as im, ImageDraw\n\n\n# Declare Execution Providers\nproviders = ['CUDAExecutionProvider', 'CPUExecutionProvider']\n\n# Download and host the model\nmodel_path = huggingface_hub.hf_hub_download(\n \"skytnt/anime-seg\", \"isnetis.onnx\")\nrmbg_model = rt.InferenceSession(model_path, providers=providers)\n\n# Function to get mask\ndef get_mask(img, s=1024):\n # Resize the img to a square shape with dimension s\n # Convert img pixel values from integers 0-255 to float 0-1\n img = (img / 255).astype(np.float32)\n # get the amount of dimensions of img\n dim = img.shape[2]\n # Convert the input image to RGB format if it has an alpha channel\n if dim == 4:\n img = img[..., :3]\n dim = 3\n # Get height and width of the image\n h, w = h0, w0 = img.shape[:-1]\n # IF height is greater than width, set h as s and w as s*width/height\n # ELSE, set w as s and h as s*height/width\n h, w = (s, int(s * w / h)) if h > w else (int(s * h / w), s)\n # Calculate padding for height and width\n ph, pw = s - h, s - w\n # Create a 1024x1024x3 array of 0's\n img_input = np.zeros([s, s, dim], dtype=np.float32)\n # Resize the original image to (w,h) and then pad with the calculated ph,pw\n img_input[ph // 2:ph // 2 + h, pw //\n 2:pw // 2 + w] = cv2.resize(img, (w, h))\n # Change the axes\n img_input = np.transpose(img_input, (2, 0, 1))\n # Add an extra axis (1,0)\n img_input = img_input[np.newaxis, :]\n # Run the model to get the mask\n mask = rmbg_model.run(None, {'img': img_input})[0][0]\n # Transpose axis\n mask = np.transpose(mask, (1, 2, 0))\n # Crop it to the images original dimensions (h0,w0)\n mask = mask[ph // 2:ph // 2 + h, pw // 2:pw // 2 + w]\n # Resize the mask to original image size (h0,w0)\n mask = cv2.resize(mask, (w0, h0))[:, :, np.newaxis]\n return mask\n\n### Function to remove background\ndef rmbg_fn(img):\n # Call get_mask() to get the mask\n mask = get_mask(img)\n # Multiply the image and the mask together to get the output image\n img = (mask * img + 255 * (1 - mask)).astype(np.uint8)\n # Convert mask value back to int 0-255\n mask = (mask * 255).astype(np.uint8)\n # Concatenate the output image and mask\n img = np.concatenate([img, mask], axis=2, dtype=np.uint8)\n # Stacking 3 identical copies of the mask for displaying\n mask = mask.repeat(3, axis=2)\n return mask, img\n\n\nclass Script(scripts.Script):\n is_txt2img = False\n\n # Function to set title\n def title(self):\n return \"ABG Remover\"\n\n def ui(self, is_img2img):\n\n with gr.Column():\n only_save_background_free_pictures = gr.Checkbox(\n label='Only save background free pictures')\n do_not_auto_save = gr.Checkbox(label='Do not auto save')\n with gr.Row():\n custom_background = gr.Checkbox(label='Custom Background')\n custom_background_color = gr.ColorPicker(\n label='Background Color', default='#ff0000')\n custom_background_random = gr.Checkbox(\n label='Random Custom Background')\n\n return [only_save_background_free_pictures, do_not_auto_save, custom_background, custom_background_color, custom_background_random]\n\n # Function to show the script\n def show(self, is_img2img):\n return True\n\n # Function to run the script\n def run(self, p, only_save_background_free_pictures, do_not_auto_save, custom_background, custom_background_color, custom_background_random):\n # If only_save_background_free_pictures is true, set do_not_save_samples to true\n if only_save_background_free_pictures:\n p.do_not_save_samples = True\n\n # Create a process_images object\n proc = process_images(p)\n\n has_grid = False\n\n unwanted_grid_because_of_img_count = len(\n proc.images) < 2 and opts.grid_only_if_multiple\n if (opts.return_grid or opts.grid_save) and not p.do_not_save_grid and not unwanted_grid_because_of_img_count:\n has_grid = True\n\n # Loop through all the images in proc\n for i in range(len(proc.images)):\n # Separate the background from the foreground\n nmask, nimg = rmbg_fn(np.array(proc.images[i]))\n\n # Check the number of channels in the nimg array, select only the first 3 or 4 channels\n num_channels = nimg.shape[2]\n if num_channels > 4:\n nimg = nimg[:, :, :4]\n \n # Ensure the data type is uint8 and convert the image back to a format that can be saved\n nimg = nimg.astype(np.uint8)\n img = im.fromarray(nimg)\n\n # If only_save_background_free_pictures is true, check if the image has a background\n if custom_background or custom_background_random:\n # If custom_background_random is true, set the background color to a random color\n if custom_background_random:\n custom_background_color = (random.randint(0, 255), random.randint(0, 255), random.randint(0, 255))\n\n # Create a new image with the same size as the original image\n background = im.new('RGBA', img.size, custom_background_color)\n\n # Draw a colored rectangle onto the new image\n draw = ImageDraw.Draw(background)\n draw.rectangle([(0, 0), img.size],\n fill=custom_background_color)\n\n # Merge the two images\n img = im.alpha_composite(background, img)\n\n # determine output path\n outpath = p.outpath_grids if has_grid and i == 0 else p.outpath_samples\n\n # If we are saving all images, save the mask and the image\n if not only_save_background_free_pictures:\n mask = im.fromarray(nmask)\n # Dot not save the new images if checkbox is checked\n if not do_not_auto_save:\n # Save the new images\n images.save_image(\n mask, outpath, \"mask_\", proc.seed + i, proc.prompt, \"png\", info=proc.info, p=p)\n images.save_image(\n img, outpath, \"img_\", proc.seed + i, proc.prompt, \"png\", info=proc.info, p=p)\n # Add the images to the proc object\n proc.images.append(mask)\n proc.images.append(img)\n # If we are only saving background-free images, save the image and replace it in the proc object\n else:\n proc.images[i] = img\n\n # Check if automatic saving is enabled\n if not do_not_auto_save:\n # Check if the image is the first one and has a grid\n if has_grid and i == 0:\n # Save the image\n images.save_image(img, p.outpath_grids, \"grid\", p.all_seeds[0], p.all_prompts[0],\n opts.grid_format, info=proc.info, short_filename=not opts.grid_extended_filename, p=p)\n else:\n # Save the image\n images.save_image(img, outpath, \"\", proc.seed,\n proc.prompt, \"png\", info=proc.info, p=p)\n\n # Return the proc object\n return proc\n" } ]