app.py

import cv2, os
import gradio as gr
import numpy as np
from demo.generation import call_generation, call_generation_t2v
from demo.instructions import INSTRUCTIONS_VECTORIZE_SIMPLIFY, INSTRUCTIONS_T2V

VERSION = 'v0.1'

GALLERY_LIST = [os.path.join('demo/gallery',path) for path in os.listdir('demo/gallery')]

PROMPT_EXAMPLE=[['a cute dog'],
    ['a cute cat'],
    ['a beautiful woman'],
    ['a bus'],]

STYLE_MAPPING = {"0": "商业扁平插画",
                 "1": "矢量风扁平插画",
                 "2": "高级时尚系列插画",
                 "3": "s矢量插画",
                 "4": "彩色插画",
                 "5": "扁平海报插画",
                 "6": "扁平职人风插画",
                 "7": "矢量风插画",
                 "8": "轮廓插画",
                 "9": "F矢量插画",
                 "10": "Material图标",
                 "11": "图标元素",
                 "12": "矢量风插画0",
                 "13": "Logo风格图标",
                 "14": "sticker风图标",
                 "15": "纹理插画",
                 "16": "矢量风插画2",
                 "17": "矢量风插画M"}

STYLE_MAPPING_EN = {
    "0": "Commercial flat illustration",
    "1": "Vector-style flat illustration",
    "2": "High-end fashion series illustration",
    "3": "s vector illustration",
    "4": "Color illustration",
    "5": "Flat poster illustration",
    "6": "Flat artisan-style illustration",
    "7": "Vector-style illustration",
    "8": "Outline illustration",
    "9": "F vector illustration",
    "10": "Material icon",
    "11": "Icon elements",
    "12": "Vector-style illustration 0",
    "13": "Logo-style icon",
    "14": "Sticker-style icon",
    "15": "Texture illustration",
    "16": "Vector-style illustration 2",
    "17": "Vector-style illustration M"
}


STYLE_MAPPING_REVERT = {v: k for k, v in STYLE_MAPPING_EN.items()}


STYLE_IMG_DICT = {"demo/style_img/商业扁平插画.webp": "0", 
                  "demo/style_img/矢量风扁平插画.webp":"1",
                  "demo/style_img/高级时尚系列插画.webp":"2",
                  "demo/style_img/s矢量插画.webp": "3", 
                  "demo/style_img/彩色插画.webp": "4", 
                  "demo/style_img/扁平海报插画.webp": "5",
                  "demo/style_img/扁平职人.webp": "6", 
                  "demo/style_img/vectorart.webp": "7", 
                  "demo/style_img/outline.webp": "8", 
                  "demo/style_img/fllust.webp": "9", 
                  "demo/style_img/icon.webp": "11", 
                  "demo/style_img/icon_material.webp": "10",
                  "demo/style_img/last.webp": "12", 
                  "demo/style_img/logo.webp": "13", 
                  "demo/style_img/sdmai.webp": "14", 
                  "demo/style_img/texture.webp": "15", 
                  "demo/style_img/vector_illustration.webp":"16",
                  "demo/style_img/vectorizeM.webp": "17", 
                  }

STYLE_IMG_DICT_REVERT = {v: k for k, v in STYLE_IMG_DICT.items()}

details_slider_dict = {"Minimal 简约":"minimal",
                       "Medium 中等":"medium",
                       "Rich 丰富": "rich"
                    }

def resize_image(image, size):
    # find the minimal size of the image, resize it to size
    # H, W, C = image.shape
    return cv2.resize(image, (size[0], size[1]), interpolation=cv2.INTER_LINEAR)

def HWC3(x):
    assert x.dtype == np.uint8
    if x.ndim == 2:
        x = x[:, :, None]
    assert x.ndim == 3
    H, W, C = x.shape
    assert C == 1 or C == 3 or C == 4
    if C == 3:
        return x
    if C == 1:
        return np.concatenate([x, x, x], axis=2)
    if C == 4:
        color = x[:, :, 0:3].astype(np.float32)
        alpha = x[:, :, 3:4].astype(np.float32) / 255.0
        y = color * alpha + 255.0 * (1.0 - alpha)
        y = y.clip(0, 255).astype(np.uint8)
        return y

    


def process_vector(input_image, upsample_method, svg_simplify, svg_optimize, trace_mode, subsample_ratio, speckle_removal,sorting_method, sorting_order, use_gpu):
    print("Processing vector:",upsample_method, svg_simplify, svg_optimize, trace_mode)
    if input_image is not None:
        ## save input_image to a temp file
        
        ## process the image
        file_list = call_generation(input_image, 
                        preprocess=upsample_method, 
                        simplify=svg_simplify, 
                        optimize=svg_optimize, 
                        mode=trace_mode, 
                        subsample_ratio=subsample_ratio, 
                        speckle_removal=speckle_removal, 
                        sorting_method=sorting_method,
                        sorting_order=sorting_order,
                        use_gpu=use_gpu)
        
        return file_list


def process_t2v(prompt, num_imgs, image_resolution_h, image_resolution_w, details_slider, style_slider, vectorize, upsample_method, svg_simplify, svg_optimize, trace_mode, subsample_ratio, speckle_removal,sorting_method, sorting_order, use_gpu):
    print("Processing t2v:",upsample_method, svg_simplify, svg_optimize, trace_mode)
    if prompt is not None:
        ## save input_image to a temp file
        details = details_slider_dict[details_slider]
        style = STYLE_MAPPING_REVERT[style_slider]
        ## process the image
        file_list = call_generation_t2v(prompt,
                        num_imgs,
                        image_resolution_h,
                        image_resolution_w,
                        details,
                        style,
                        vectorize,
                        preprocess=upsample_method, 
                        simplify=svg_simplify, 
                        optimize=svg_optimize, 
                        mode=trace_mode, 
                        subsample_ratio=subsample_ratio, 
                        speckle_removal=speckle_removal, 
                        sorting_method=sorting_method,
                        sorting_order=sorting_order,
                        use_gpu=use_gpu)
        
        return file_list



block = gr.Blocks(
        title = "VectorizeAnything",
        theme=gr.themes.Soft(
            radius_size=gr.themes.sizes.radius_none,
            text_size=gr.themes.sizes.text_md
        ),
        css="css/style.css",
        ).queue()

with block:
    state = gr.State(value={
        'gallery_selected_img_path': None,  # 当前选中的图片路径
        'gallery_selected_img_path_idx': 0,  # 当前选中的图片路径索引
    })
    with gr.Row():
        gr.HTML(f"""
                    </br>
                    <div>
                        <h1 style="font-size:3rem; "><center>Vectorize Anything: {VERSION} </center></h1>
                    </div>
                    </br>
            """)
    # tab_0 = gr.Tab(label="Gallery (画廊)")
    # with tab_0:
    #     with gr.Row():
    #         gr.Gallery(label='图像生成结果', value=GALLERY_LIST,show_label=False, elem_id="Gallery", columns=5, height=1000)

    tab_1 = gr.Tab(label="TEXT to Vector (文生矢量图)")
    
    with tab_1:
        with gr.Accordion('🕹Usage (操作说明)', open=True,):
            with gr.Tabs():
                gr.HTML(INSTRUCTIONS_T2V)
        with gr.Row():
            with gr.Column():
                with gr.Row():
                    lora_radio_t2v = gr.Radio(choices=STYLE_MAPPING_REVERT.keys(), type="value", value='Commercial flat illustration', label="Style Selection (风格选择)")
                    style_image_t2v = gr.Image(value=STYLE_IMG_DICT_REVERT["0"],type="numpy", image_mode="RGBA", label="Style Image (风格图片)")
                prompt_t2v = gr.Textbox(label="Prompt (提示词)", value="a cute dog")
                details_slider_t2v = gr.Radio(label="图形细节 Details", choices=["Minimal 简约","Medium 中等","Rich 丰富" ], type="value", value="Medium 中等")
                num_samples_t2v = gr.Slider(label="Images (图像数量)", minimum=1, maximum=2, value=1, step=1, visible=False)
                vectorize = gr.Checkbox(label='generate svg file', value=True, visible=True)
                run_button_t2v = gr.Button(value="Generation Vector (生成矢量)",elem_id="btnSEG")

                with gr.Accordion("Image options", open=False):
                    image_resolution_h = gr.Slider(label="Image height (高)", minimum=256, maximum=1024, value=512, step=64, visible=True)
                    image_resolution_w = gr.Slider(label="Image width (宽)", minimum=256, maximum=1024, value=512, step=64, visible=True)
                     

                with gr.Accordion("Vector options", open=False):
                    upsample_method_t2v = gr.Dropdown(choices=["None", "x4", "x2"], type="value", value="None", label="Upsample Method (超分倍数)[无, 4倍, 2倍], 默认无")
                    sorting_method_t2v = gr.Dropdown(choices=["brightness","area"], type="value", value="brightness", label="Sorting Method (排序方法)[亮度,面积], 默认亮度")
                    sorting_order_t2v = gr.Dropdown(choices=["ascend","descend"], type="value", value="descend", label="Sorting Order (排序顺序)[递增,递减], 默认递减")
                    trace_mode_t2v = gr.Radio(choices=["overlap", "cutout"], type="value", value="overlap", label="Trace Mode (追踪模式)[堆叠,分开]")
                    use_gpu_t2v = gr.Checkbox(label='use GPU (是否使用GPU加速[针对复杂图形较为有效]), 默认关闭', value=False, visible=True)
                    svg_simplify_t2v = gr.Checkbox(label='Simplify SVG (简化矢量图, 默认关闭)', value=False, visible=True)
                    svg_optimize_t2v = gr.Checkbox(label='Optimize SVG (优化矢量图, 默认关闭)', value=False, visible=True)
                    speckle_removal_t2v = gr.Checkbox(label='Remove small speckle[是否移除面积过小的图形]', value=False)
                    subsample_ratio_t2v = gr.Slider(label="Subsample Ratio", minimum=1, maximum=10000, value=12, step=1, visible=False)

                with gr.Tab("Prompt Examples"):
                    t2v_examples = gr.Examples(
                        PROMPT_EXAMPLE,
                        [prompt_t2v],
                        examples_per_page=5,
                        label=''
                    )

            with gr.Column():
                result_gallery_t2v = gr.Gallery(label='图像生成结果 Generation Results', show_label=False, elem_id="Gallery")

            def update_lora_image(lora_radio: gr.Radio):
                return gr.Image(value=STYLE_IMG_DICT_REVERT[STYLE_MAPPING_REVERT[lora_radio]],)
            
            lora_radio_t2v.change(fn=update_lora_image, inputs=[lora_radio_t2v], outputs=[style_image_t2v])



            t2v_ips = [prompt_t2v, num_samples_t2v, image_resolution_h, image_resolution_w, details_slider_t2v, lora_radio_t2v, vectorize, upsample_method_t2v, svg_simplify_t2v, svg_optimize_t2v, trace_mode_t2v, subsample_ratio_t2v, speckle_removal_t2v,sorting_method_t2v, sorting_order_t2v, use_gpu_t2v]


            run_button_t2v.click(fn=process_t2v, inputs=t2v_ips, outputs=result_gallery_t2v)

    tab_3 = gr.Tab(label="IMG to SVG")
    with tab_3:   
        with gr.Accordion('🕹Usage', open=True,):
            with gr.Tabs():
                gr.HTML(INSTRUCTIONS_VECTORIZE_SIMPLIFY)
        with gr.Row():
            with gr.Column():
                input_image = gr.Image(type="numpy", image_mode="RGBA")
                run_vectorize = gr.Button(value="Vectorize",elem_id="btnVEC")

                with gr.Accordion("Vector options", open=True):
                    upsample_method = gr.Dropdown(choices=["None", "x4", "x2"], type="value", value="None", label="Upsample Method")
                    sorting_method = gr.Dropdown(choices=["brightness","area"], type="value", value="brightness", label="Sorting Method")
                    sorting_order = gr.Dropdown(choices=["ascend","descend"], type="value", value="descend", label="Sorting Order")
                    trace_mode = gr.Radio(choices=["overlap", "cutout"], type="value", value="overlap", label="Trace Mode")
                    use_gpu = gr.Checkbox(label='use GPU', value=False, visible=True)
                    svg_simplify = gr.Checkbox(label='Simplify SVG', value=False, visible=True)
                    svg_optimize = gr.Checkbox(label='Optimize SVG', value=False, visible=True)
                    speckle_removal = gr.Checkbox(label='Remove small speckle', value=False)
                    subsample_ratio = gr.Slider(label="Subsample Ratio", minimum=1, maximum=10000, value=12, step=1, visible=False)


            def exp_gen_click():
                return [gr.Slider(value=512), gr.Slider(value=512)]  # all examples are 512x512, refresh draw_img

            with gr.Column():
                result_vector_gallery = gr.Gallery(label='Output', show_label=False, elem_id="Gallery_vector")
            

            
                with gr.Tab("Image Examples"):
                    exp_gen_en = gr.Examples(
                    [
                        ["test_imgs/demo1.png"],
                        ["test_imgs/demo2.jpg"],
                        ["test_imgs/demo3.png"],
                        ["test_imgs/demo4.png"],
                        ["test_imgs/demo5.png"],
                        ["test_imgs/demo6.png"],
                        ["test_imgs/demo7.png"],
                        ["test_imgs/demo8.png"],
                        ["test_imgs/demo9.png"],
                        ["test_imgs/demo10.png"],
                        ["test_imgs/demo11.png"],
                        ["test_imgs/demo12.png"],
                    ],
                    [input_image],
                    examples_per_page=20,
                    label=''
                    )
                    exp_gen_en.dataset.click(exp_gen_click, None)

    
    vector_ips = [input_image, upsample_method, svg_simplify, svg_optimize, trace_mode, subsample_ratio, speckle_removal,sorting_method, sorting_order, use_gpu]
    run_vectorize.click(fn=process_vector, inputs=vector_ips, outputs=result_vector_gallery)

block.launch(server_name='0.0.0.0', share=False,debug=True, root_path=f"/{os.getenv('GRADIO_PROXY_PATH')}" if os.getenv('GRADIO_PROXY_PATH') else "")