85 lines
3.1 KiB
Python
85 lines
3.1 KiB
Python
import os
|
|
import io
|
|
import json
|
|
import numpy as np
|
|
import cv2
|
|
|
|
import gradio as gr
|
|
|
|
import modules.scripts as scripts
|
|
from modules import script_callbacks
|
|
|
|
from td_abg import get_foreground
|
|
from convertor import pil2cv
|
|
|
|
|
|
"""
|
|
body_estimation = None
|
|
presets_file = os.path.join(scripts.basedir(), "presets.json")
|
|
presets = {}
|
|
|
|
try:
|
|
with open(presets_file) as file:
|
|
presets = json.load(file)
|
|
except FileNotFoundError:
|
|
pass
|
|
"""
|
|
|
|
def processing(self, input_image, td_abg_enabled, h_split, v_split, n_cluster, alpha, th_rate, cascadePSP_enabled, fast, psp_L):
|
|
image = pil2cv(input_image)
|
|
image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
|
|
mask, image = get_foreground(image, td_abg_enabled, h_split, v_split, n_cluster, alpha, th_rate, cascadePSP_enabled, fast, psp_L)
|
|
return image, mask
|
|
|
|
class Script(scripts.Script):
|
|
def __init__(self) -> None:
|
|
super().__init__()
|
|
|
|
def title(self):
|
|
return "PBRemTools"
|
|
|
|
def show(self, is_img2img):
|
|
return scripts.AlwaysVisible
|
|
|
|
def ui(self, is_img2img):
|
|
return ()
|
|
|
|
def on_ui_tabs():
|
|
with gr.Blocks(analytics_enabled=False) as PBRemTools:
|
|
with gr.Row():
|
|
with gr.Column():
|
|
input_image = gr.Image(type="pil")
|
|
with gr.Accordion("tile division ABG", open=True):
|
|
with gr.Box():
|
|
td_abg_enabled = gr.Checkbox(label="enabled", show_label=True)
|
|
h_split = gr.Slider(1, 2048, value=256, step=4, label="horizontal split num", show_label=True)
|
|
v_split = gr.Slider(1, 2048, value=256, step=4, label="vertical split num", show_label=True)
|
|
|
|
n_cluster = gr.Slider(1, 1000, value=500, step=10, label="cluster num", show_label=True)
|
|
alpha = gr.Slider(1, 255, value=100, step=1, label="alpha threshold", show_label=True)
|
|
th_rate = gr.Slider(0, 1, value=0.1, step=0.01, label="mask content ratio", show_label=True)
|
|
|
|
with gr.Accordion("cascadePSP", open=True):
|
|
with gr.Box():
|
|
cascadePSP_enabled = gr.Checkbox(label="enabled", show_label=True)
|
|
fast = gr.Checkbox(label="fast", show_label=True)
|
|
psp_L = gr.Slider(1, 2048, value=900, step=1, label="Memory usage", show_label=True)
|
|
|
|
submit = gr.Button(value="Submit")
|
|
with gr.Row():
|
|
with gr.Column():
|
|
with gr.Tab("output"):
|
|
output_img = gr.Image()
|
|
with gr.Tab("mask"):
|
|
output_mask = gr.Image()
|
|
#dummy_component = gr.Label(visible=False)
|
|
#preset = gr.Text(visible=False)
|
|
submit.click(
|
|
processing,
|
|
inputs=[input_image, td_abg_enabled, h_split, v_split, n_cluster, alpha, th_rate, cascadePSP_enabled, fast, psp_L],
|
|
outputs=[output_img, output_mask]
|
|
)
|
|
|
|
return [(PBRemTools, "PBRemTools", "pbremtools")]
|
|
|
|
script_callbacks.on_ui_tabs(on_ui_tabs) |