Full Code of KutsuyaYuki/ABG_extension for AI

Repository: KutsuyaYuki/ABG_extension
Branch: main
Commit: eb3ff7610abe
Files: 4
Total size: 9.3 KB

Directory structure:
gitextract_6c8i53hp/

├── .gitignore
├── README.md
├── install.py
└── scripts/
    └── app.py

================================================
FILE CONTENTS
================================================

================================================
FILE: .gitignore
================================================
scripts/__pycache__

================================================
FILE: README.md
================================================
<h3 align="center">
<b>ABG extension</b>
</h3>

<p align="center">
	<a href="https://github.com/KutsuyaYuki/ABG_extension/stargazers"><img src="https://img.shields.io/github/stars/KutsuyaYuki/ABG_extension?style=for-the-badge"></a>
	<a href="https://github.com/KutsuyaYuki/ABG_extension/issues"><img src="https://img.shields.io/github/issues/KutsuyaYuki/ABG_extension?style=for-the-badge"></a>
	<a href="https://github.com/KutsuyaYuki/ABG_extension/commits/main"><img src="https://img.shields.io/github/last-commit/KutsuyaYuki/ABG_extension?style=for-the-badge"></a>
</p>

## Installation

 1. Install extension by going to Extensions tab -> Install from URL -> Paste github URL and click Install.
 2. After it's installed, go back to the Installed tab in Extensions and press Apply and restart UI.
 3. Installation finished.
 4. If the script does not show up or work, please restart the WebUI.

## Usage

### txt2img

 1. In the bottom of the WebUI in Script, select **ABG Remover**.
 2. Select the desired options: **Only save background free pictures** or **Do not auto save**.
 3. Generate an image and you will see the result in the output area.

### img2img

 1. In the bottom of the WebUI in Script, select **ABG Remover**.
 2. Select the desired options: **Only save background free pictures** or **Do not auto save**.
 3. **IMPORTANT**: Set **Denoising strength** to a low value, like **0.01**

Based on https://huggingface.co/spaces/skytnt/anime-remove-background


================================================
FILE: install.py
================================================
import launch

for dep in ['onnx', 'onnxruntime', 'numpy']:
    if not launch.is_installed(dep):
        launch.run_pip(f"install {dep}", f"{dep} for ABG_extension")

if not launch.is_installed("cv2"):
    launch.run_pip("install opencv-python", "opencv-python")

if not launch.is_installed("PIL"):
    launch.run_pip("install Pillow", "Pillow")


================================================
FILE: scripts/app.py
================================================
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