CVPR/DualStyleGAN

1#!/usr/bin/env python2 3from __future__ import annotations4 5import pathlib6 7import gradio as gr8 9from dualstylegan import Model10 11DESCRIPTION = """# Portrait Style Transfer with [DualStyleGAN](https://github.com/williamyang1991/DualStyleGAN)12 13<img id="overview" alt="overview" src="https://raw.githubusercontent.com/williamyang1991/DualStyleGAN/main/doc_images/overview.jpg" />14"""15 16 17def get_style_image_url(style_name: str) -> str:18    base_url = "https://raw.githubusercontent.com/williamyang1991/DualStyleGAN/main/doc_images"19    filenames = {20        "cartoon": "cartoon_overview.jpg",21        "caricature": "caricature_overview.jpg",22        "anime": "anime_overview.jpg",23        "arcane": "Reconstruction_arcane_overview.jpg",24        "comic": "Reconstruction_comic_overview.jpg",25        "pixar": "Reconstruction_pixar_overview.jpg",26        "slamdunk": "Reconstruction_slamdunk_overview.jpg",27    }28    return f"{base_url}/{filenames[style_name]}"29 30 31def get_style_image_markdown_text(style_name: str) -> str:32    url = get_style_image_url(style_name)33    return f'<img id="style-image" src="{url}" alt="style image">'34 35 36def update_slider(choice: str) -> dict:37    max_vals = {38        "cartoon": 316,39        "caricature": 198,40        "anime": 173,41        "arcane": 99,42        "comic": 100,43        "pixar": 121,44        "slamdunk": 119,45    }46    return gr.Slider(maximum=max_vals[choice])47 48 49def update_style_image(style_name: str) -> dict:50    text = get_style_image_markdown_text(style_name)51    return gr.Markdown(value=text)52 53 54model = Model()55 56with gr.Blocks(css="style.css") as demo:57    gr.Markdown(DESCRIPTION)58 59    with gr.Group():60        gr.Markdown(61            """## Step 1 (Preprocess Input Image)62 63- Drop an image containing a near-frontal face to the **Input Image**.64- If there are multiple faces in the image, hit the Edit button in the upper right corner and crop the input image beforehand.65- Hit the **Preprocess** button.66- Choose the encoder version. Default is Z+ encoder which has better stylization performance. W+ encoder better reconstructs the input image to preserve more details.67- The final result will be based on this **Reconstructed Face**. So, if the reconstructed image is not satisfactory, you may want to change the input image.68"""69        )70        with gr.Row():71            encoder_type = gr.Radio(72                label="Encoder Type",73                choices=["Z+ encoder (better stylization)", "W+ encoder (better reconstruction)"],74                value="Z+ encoder (better stylization)",75            )76        with gr.Row():77            with gr.Column():78                with gr.Row():79                    input_image = gr.Image(label="Input Image", type="filepath")80                with gr.Row():81                    preprocess_button = gr.Button("Preprocess")82            with gr.Column():83                with gr.Row():84                    aligned_face = gr.Image(label="Aligned Face", type="numpy", interactive=False)85            with gr.Column():86                reconstructed_face = gr.Image(label="Reconstructed Face", type="numpy")87                instyle = gr.State()88 89        with gr.Row():90            paths = sorted(pathlib.Path("images").glob("*.jpg"))91            gr.Examples(examples=[[path.as_posix()] for path in paths], inputs=input_image)92 93    with gr.Group():94        gr.Markdown(95            """## Step 2 (Select Style Image)96 97- Select **Style Type**.98- Select **Style Image Index** from the image table below.99"""100        )101        with gr.Row():102            with gr.Column():103                style_type = gr.Radio(label="Style Type", choices=model.style_types, value=model.style_types[0])104                text = get_style_image_markdown_text("cartoon")105                style_image = gr.Markdown(value=text, latex_delimiters=[])106                style_index = gr.Slider(label="Style Image Index", minimum=0, maximum=316, step=1, value=26)107 108        with gr.Row():109            gr.Examples(110                examples=[111                    ["cartoon", 26],112                    ["caricature", 65],113                    ["arcane", 63],114                    ["pixar", 80],115                ],116                inputs=[style_type, style_index],117            )118 119    with gr.Group():120        gr.Markdown(121            """## Step 3 (Generate Style Transferred Image)122 123- Adjust **Structure Weight** and **Color Weight**.124- These are weights for the style image, so the larger the value, the closer the resulting image will be to the style image.125- Tips: For W+ encoder, better way of (Structure Only) is to uncheck (Structure Only) and set Color weight to 0.126- Hit the **Generate** button.127"""128        )129        with gr.Row():130            with gr.Column():131                with gr.Row():132                    structure_weight = gr.Slider(label="Structure Weight", minimum=0, maximum=1, step=0.1, value=0.6)133                with gr.Row():134                    color_weight = gr.Slider(label="Color Weight", minimum=0, maximum=1, step=0.1, value=1)135                with gr.Row():136                    structure_only = gr.Checkbox(label="Structure Only", value=False)137                with gr.Row():138                    generate_button = gr.Button("Generate")139 140            with gr.Column():141                result = gr.Image(label="Result")142 143        with gr.Row():144            gr.Examples(145                examples=[146                    [0.6, 1.0],147                    [0.3, 1.0],148                    [0.0, 1.0],149                    [1.0, 0.0],150                ],151                inputs=[structure_weight, color_weight],152            )153 154    preprocess_button.click(155        fn=model.detect_and_align_face,156        inputs=[input_image],157        outputs=aligned_face,158    )159    aligned_face.change(160        fn=model.reconstruct_face,161        inputs=[aligned_face, encoder_type],162        outputs=[163            reconstructed_face,164            instyle,165        ],166    )167    style_type.change(168        fn=update_slider,169        inputs=style_type,170        outputs=style_index,171    )172    style_type.change(173        fn=update_style_image,174        inputs=style_type,175        outputs=style_image,176    )177    generate_button.click(178        fn=model.generate,179        inputs=[180            style_type,181            style_index,182            structure_weight,183            color_weight,184            structure_only,185            instyle,186        ],187        outputs=result,188    )189 190if __name__ == "__main__":191    demo.queue(max_size=20).launch()192