app.py

import os
import sys

import cv2
import gradio as gr

from ldivider.ld_convertor import cv2pil, pil2cv
from ldivider.ld_processor_fast import get_base, get_composite_layer, get_normal_layer
from ldivider.ld_utils import save_psd
from pytoshop.enums import BlendMode

path = os.getcwd()
output_dir = f"{path}/output"
input_dir = f"{path}/input"
model_dir = f"{path}/segment_model"


class webui:
    def __init__(self):
        self.demo = gr.Blocks()

    def color_base_divide(
        self, input_image, loops, init_cluster, ciede_threshold, blur_size, layer_mode
    ):
        image = pil2cv(input_image)
        self.input_image = cv2.cvtColor(image, cv2.COLOR_BGR2RGBA)

        base_image, label = get_base(
            self.input_image, loops, init_cluster, ciede_threshold, blur_size
        )

        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(self.input_image, base_image, label)
            filename = save_psd(
                self.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(
                self.input_image, base_image, label
            )
            filename = save_psd(
                self.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 launch(self, share):
        with self.demo:
            with gr.Row():
                with gr.Column():
                    input_image = gr.Image(type="pil")
                    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,
                        )

                    submit = gr.Button(value="Create PSD")
                with gr.Row():
                    with gr.Column():
                        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(
                self.color_base_divide,
                inputs=[
                    input_image,
                    loops,
                    init_cluster,
                    ciede_threshold,
                    blur_size,
                    layer_mode,
                ],
                outputs=[output_0, output_1, output_2, output_3, output_file],
            )

        self.demo.queue()
        self.demo.launch(share=share)


if __name__ == "__main__":
    ui = webui()
    if len(sys.argv) > 1:
        if sys.argv[1] == "share":
            ui.launch(share=True)
        else:
            ui.launch(share=False)
    else:
        ui.launch(share=False)