196 lines
9.1 KiB
Python
196 lines
9.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 ldivider.ld_convertor import pil2cv, cv2pil, df2bgra
|
|
from ldivider.ld_processor import get_base, get_normal_layer, get_composite_layer, get_seg_base
|
|
from ldivider.ld_utils import save_psd, load_masks, divide_folder, load_seg_model
|
|
from ldivider.ld_segment import get_mask_generator, get_masks, show_anns
|
|
|
|
from modules.paths_internal import extensions_dir
|
|
from collections import OrderedDict
|
|
|
|
from pytoshop.enums import BlendMode
|
|
|
|
|
|
|
|
|
|
model_cache = OrderedDict()
|
|
base_dir = "layerdivider"
|
|
output_dir = os.path.join(extensions_dir, base_dir, "output")
|
|
input_dir = os.path.join(extensions_dir, base_dir, "input")
|
|
model_dir = os.path.join(extensions_dir, base_dir, "segment_model")
|
|
|
|
class Script(scripts.Script):
|
|
def __init__(self) -> None:
|
|
super().__init__()
|
|
|
|
def title(self):
|
|
return "LayerDivider"
|
|
|
|
def show(self, is_img2img):
|
|
return scripts.AlwaysVisible
|
|
|
|
def ui(self, is_img2img):
|
|
return ()
|
|
|
|
def segment_image(input_image, pred_iou_thresh, stability_score_thresh, min_mask_region_area):
|
|
load_seg_model(model_dir)
|
|
mask_generator = get_mask_generator(pred_iou_thresh, stability_score_thresh,min_mask_region_area, model_dir, "extension")
|
|
masks = get_masks(pil2cv(input_image), mask_generator)
|
|
input_image.putalpha(255)
|
|
masked_image = show_anns(input_image, masks, output_dir)
|
|
return masked_image
|
|
|
|
def divide_layer(divide_mode, input_image, loops, init_cluster, ciede_threshold, blur_size, layer_mode, h_split, v_split, n_cluster, alpha, th_rate, split_bg, area_th):
|
|
if divide_mode == "segment_mode":
|
|
return segment_divide(input_image, loops, init_cluster, ciede_threshold, blur_size, layer_mode, h_split, v_split, n_cluster, alpha, th_rate, split_bg, area_th)
|
|
elif divide_mode == "color_base_mode":
|
|
return color_base_divide(input_image, loops, init_cluster, ciede_threshold, blur_size, layer_mode, h_split, v_split, n_cluster, alpha, th_rate, split_bg)
|
|
|
|
|
|
def segment_divide(input_image, loops, init_cluster, ciede_threshold, blur_size, layer_mode, h_split, v_split, n_cluster, alpha, th_rate, split_bg, area_th):
|
|
image = pil2cv(input_image)
|
|
input_image = cv2.cvtColor(image, cv2.COLOR_BGR2RGBA)
|
|
masks = load_masks(output_dir)
|
|
|
|
df = get_seg_base(input_image, masks, area_th)
|
|
|
|
|
|
base_image = cv2pil(df2bgra(df))
|
|
image = cv2pil(image)
|
|
if layer_mode == "composite":
|
|
base_layer_list, shadow_layer_list, bright_layer_list, addition_layer_list, subtract_layer_list = get_composite_layer(input_image, df)
|
|
filename = save_psd(
|
|
input_image,
|
|
[base_layer_list, bright_layer_list, shadow_layer_list, subtract_layer_list, addition_layer_list],
|
|
["base", "screen", "multiply", "subtract", "addition"],
|
|
[BlendMode.normal, BlendMode.screen, BlendMode.multiply, BlendMode.subtract, BlendMode.linear_dodge],
|
|
output_dir,
|
|
layer_mode
|
|
|
|
)
|
|
base_layer_list = [cv2pil(layer) for layer in base_layer_list]
|
|
divide_folder(filename, input_dir)
|
|
return [image, base_image], base_layer_list, bright_layer_list, shadow_layer_list, filename
|
|
elif layer_mode == "normal":
|
|
base_layer_list, bright_layer_list, shadow_layer_list = get_normal_layer(input_image, df)
|
|
filename = save_psd(
|
|
input_image,
|
|
[base_layer_list, bright_layer_list, shadow_layer_list],
|
|
["base", "bright", "shadow"],
|
|
[BlendMode.normal, BlendMode.normal, BlendMode.normal],
|
|
output_dir,
|
|
layer_mode
|
|
)
|
|
|
|
divide_folder(filename, input_dir)
|
|
return [image, base_image], base_layer_list, bright_layer_list, shadow_layer_list, filename
|
|
else:
|
|
return None
|
|
|
|
|
|
|
|
def color_base_divide(input_image, loops, init_cluster, ciede_threshold, blur_size, layer_mode, h_split, v_split, n_cluster, alpha, th_rate, split_bg):
|
|
image = pil2cv(input_image)
|
|
input_image = cv2.cvtColor(image, cv2.COLOR_BGR2RGBA)
|
|
|
|
df = get_base(input_image, loops, init_cluster, ciede_threshold, blur_size, h_split, v_split, n_cluster, alpha, th_rate, split_bg, False)
|
|
|
|
base_image = cv2pil(df2bgra(df))
|
|
image = cv2pil(image)
|
|
if layer_mode == "composite":
|
|
base_layer_list, shadow_layer_list, bright_layer_list, addition_layer_list, subtract_layer_list = get_composite_layer(input_image, df)
|
|
filename = save_psd(
|
|
input_image,
|
|
[base_layer_list, bright_layer_list, shadow_layer_list, subtract_layer_list, addition_layer_list],
|
|
["base", "screen", "multiply", "subtract", "addition"],
|
|
[BlendMode.normal, BlendMode.screen, BlendMode.multiply, BlendMode.subtract, BlendMode.linear_dodge],
|
|
output_dir,
|
|
layer_mode,
|
|
)
|
|
base_layer_list = [cv2pil(layer) for layer in base_layer_list]
|
|
return [image, base_image], base_layer_list, bright_layer_list, shadow_layer_list, filename
|
|
elif layer_mode == "normal":
|
|
base_layer_list, bright_layer_list, shadow_layer_list = get_normal_layer(input_image, df)
|
|
filename = save_psd(
|
|
input_image,
|
|
[base_layer_list, bright_layer_list, shadow_layer_list],
|
|
["base", "bright", "shadow"],
|
|
[BlendMode.normal, BlendMode.normal, BlendMode.normal],
|
|
output_dir,
|
|
layer_mode,
|
|
)
|
|
return [image, base_image], base_layer_list, bright_layer_list, shadow_layer_list, filename
|
|
else:
|
|
return None
|
|
|
|
def on_ui_tabs():
|
|
with gr.Blocks(analytics_enabled=False) as LayerDivider:
|
|
with gr.Row():
|
|
with gr.Column():
|
|
input_image = gr.Image(type="pil")
|
|
divide_mode = gr.Dropdown(["segment_mode", "color_base_mode"], value = "segment_mode", label="divide_mode", show_label=True)
|
|
|
|
with gr.Accordion("Segment Settings", open=True):
|
|
area_th = gr.Slider(1, 100000, value=20000, step=100, label="area_threshold", show_label=True)
|
|
|
|
with gr.Accordion("ColorBase Settings", open=True):
|
|
loops = gr.Slider(1, 20, value=1, step=1, label="loops", show_label=True)
|
|
init_cluster = gr.Slider(1, 50, value=10, step=1, label="init_cluster", show_label=True)
|
|
ciede_threshold = gr.Slider(1, 50, value=5, step=1, label="ciede_threshold", show_label=True)
|
|
blur_size = gr.Slider(1, 20, value=5, label="blur_size", show_label=True)
|
|
layer_mode = gr.Dropdown(["normal", "composite"], value = "normal", label="output_layer_mode", show_label=True)
|
|
|
|
with gr.Accordion("BG Settings", open=True):
|
|
split_bg = gr.Checkbox(label="split bg", 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)
|
|
|
|
submit = gr.Button(value="Create PSD")
|
|
with gr.Row():
|
|
with gr.Column():
|
|
SAM_output = gr.Image(type="pil")
|
|
pred_iou_thresh = gr.Slider(0, 1, value=0.8, step=0.01, label="pred_iou_thresh", show_label=True)
|
|
stability_score_thresh = gr.Slider(0, 1, value=0.8, step=0.01, label="stability_score_thresh", show_label=True)
|
|
#crop_n_layers = gr.Slider(0, 10, value=0, step=1, label="crop_n_layers", show_label=True)
|
|
#crop_n_points_downscale_factor = gr.Slider(1, 10, value=1, step=1, label="crop_n_points_downscale_factor", show_label=True)
|
|
min_mask_region_area = gr.Slider(1, 1000, value=100, step=1, label="min_mask_region_area", show_label=True)
|
|
segment = gr.Button(value="Segment")
|
|
with gr.Tab("output"):
|
|
output_0 = gr.Gallery()
|
|
with gr.Tab("base"):
|
|
output_1 = gr.Gallery()
|
|
with gr.Tab("bright"):
|
|
output_2 = gr.Gallery()
|
|
with gr.Tab("shadow"):
|
|
output_3 = gr.Gallery()
|
|
|
|
output_file = gr.File()
|
|
|
|
submit.click(
|
|
divide_layer,
|
|
inputs=[divide_mode, input_image, loops, init_cluster, ciede_threshold, blur_size, layer_mode, h_split, v_split, n_cluster, alpha, th_rate, split_bg, area_th],
|
|
outputs=[output_0, output_1, output_2, output_3, output_file]
|
|
)
|
|
segment.click(
|
|
segment_image,
|
|
inputs=[input_image, pred_iou_thresh, stability_score_thresh, min_mask_region_area],
|
|
outputs=[SAM_output]
|
|
)
|
|
|
|
return [(LayerDivider, "LayerDivider", "layerdivider")]
|
|
|
|
script_callbacks.on_ui_tabs(on_ui_tabs)
|