launch.py

import os
import sys
import json
import uuid
import numpy as np
import gradio as gr
import trimesh
import zipfile
import subprocess
from datetime import datetime
from functools import partial
from PIL import Image, ImageChops
from huggingface_hub import snapshot_download

from gradio_model3dcolor import Model3DColor
from gradio_model3dnormal import Model3DNormal

is_local_run = os.path.exists("../SpaRP_API")
code_dir = (
    snapshot_download("sudo-ai/SpaRP_API", token=os.environ["HF_TOKEN"])
    if not is_local_run
    else "../SpaRP_API"
)

if not is_local_run:
    zip_file_path = f"{code_dir}/examples.zip"
    # Unzipping the file into the current directory
    with zipfile.ZipFile(zip_file_path, "r") as zip_ref:
        zip_ref.extractall(os.getcwd())

with open(f"{code_dir}/api.json", "r") as file:
    api_dict = json.load(file)
    SEGM_i_CALL = api_dict["SEGM_i_CALL"]
    UNPOSED_CALL = api_dict["UNPOSED_CALL"]
    MESH_CALL = api_dict["MESH_CALL"]


_TITLE = (
    """SpaRP: Fast 3D Object Reconstruction and Pose Estimation from Sparse Views"""
)
_DESCRIPTION = (
    """Try SpaRP to reconstruct 3D textured mesh from one or a few unposed images!"""
)
_PR = """
<div>
<b><em>Project page: <a href="https://chaoxu.xyz/sparp/">https://chaoxu.xyz/sparp/</a></em></b>
</div>
<div>
<b><em>Check out <a href="https://www.sudo.ai/3dgen">Hillbot (sudoAI)</a> for more details and advanced features.</em></b>
</div>
"""


STYLE = """
    <link href="https://cdn.jsdelivr.net/npm/[email protected]/dist/css/bootstrap.min.css" rel="stylesheet" integrity="sha384-T3c6CoIi6uLrA9TneNEoa7RxnatzjcDSCmG1MXxSR1GAsXEV/Dwwykc2MPK8M2HN" crossorigin="anonymous">
    <style>
        .alert, .alert div, .alert b {
            color: black !important;
        }
    </style>
"""
# info (info-circle-fill), cursor (hand-index-thumb), wait (hourglass-split), done (check-circle)
ICONS = {
    "info": """<svg xmlns="http://www.w3.org/2000/svg" width="16" height="16" fill="#0d6efd" class="bi bi-info-circle-fill flex-shrink-0 me-2" viewBox="0 0 16 16">
    <path d="M8 16A8 8 0 1 0 8 0a8 8 0 0 0 0 16zm.93-9.412-1 4.705c-.07.34.029.533.304.533.194 0 .487-.07.686-.246l-.088.416c-.287.346-.92.598-1.465.598-.703 0-1.002-.422-.808-1.319l.738-3.468c.064-.293.006-.399-.287-.47l-.451-.081.082-.381 2.29-.287zM8 5.5a1 1 0 1 1 0-2 1 1 0 0 1 0 2z"/>
    </svg>""",
    "cursor": """<svg xmlns="http://www.w3.org/2000/svg" width="16" height="16" fill="#0dcaf0" class="bi bi-hand-index-thumb-fill flex-shrink-0 me-2" viewBox="0 0 16 16">
    <path d="M8.5 1.75v2.716l.047-.002c.312-.012.742-.016 1.051.046.28.056.543.18.738.288.273.152.456.385.56.642l.132-.012c.312-.024.794-.038 1.158.108.37.148.689.487.88.716.075.09.141.175.195.248h.582a2 2 0 0 1 1.99 2.199l-.272 2.715a3.5 3.5 0 0 1-.444 1.389l-1.395 2.441A1.5 1.5 0 0 1 12.42 16H6.118a1.5 1.5 0 0 1-1.342-.83l-1.215-2.43L1.07 8.589a1.517 1.517 0 0 1 2.373-1.852L5 8.293V1.75a1.75 1.75 0 0 1 3.5 0z"/>
    </svg>""",
    "wait": """<svg xmlns="http://www.w3.org/2000/svg" width="16" height="16" fill="#6c757d" class="bi bi-hourglass-split flex-shrink-0 me-2" viewBox="0 0 16 16">
    <path d="M2.5 15a.5.5 0 1 1 0-1h1v-1a4.5 4.5 0 0 1 2.557-4.06c.29-.139.443-.377.443-.59v-.7c0-.213-.154-.451-.443-.59A4.5 4.5 0 0 1 3.5 3V2h-1a.5.5 0 0 1 0-1h11a.5.5 0 0 1 0 1h-1v1a4.5 4.5 0 0 1-2.557 4.06c-.29.139-.443.377-.443.59v.7c0 .213.154.451.443.59A4.5 4.5 0 0 1 12.5 13v1h1a.5.5 0 0 1 0 1h-11zm2-13v1c0 .537.12 1.045.337 1.5h6.326c.216-.455.337-.963.337-1.5V2h-7zm3 6.35c0 .701-.478 1.236-1.011 1.492A3.5 3.5 0 0 0 4.5 13s.866-1.299 3-1.48V8.35zm1 0v3.17c2.134.181 3 1.48 3 1.48a3.5 3.5 0 0 0-1.989-3.158C8.978 9.586 8.5 9.052 8.5 8.351z"/>
    </svg>""",
    "done": """<svg xmlns="http://www.w3.org/2000/svg" width="16" height="16" fill="#198754" class="bi bi-check-circle-fill flex-shrink-0 me-2" viewBox="0 0 16 16">
    <path d="M16 8A8 8 0 1 1 0 8a8 8 0 0 1 16 0zm-3.97-3.03a.75.75 0 0 0-1.08.022L7.477 9.417 5.384 7.323a.75.75 0 0 0-1.06 1.06L6.97 11.03a.75.75 0 0 0 1.079-.02l3.992-4.99a.75.75 0 0 0-.01-1.05z"/>
    </svg>""",
}

icons2alert = {
    "info": "primary",  # blue
    "cursor": "info",  # light blue
    "wait": "secondary",  # gray
    "done": "success",  # green
}


def message(text, icon_type="info"):
    return f"""{STYLE}  <div class="alert alert-{icons2alert[icon_type]} d-flex align-items-center" role="alert"> {ICONS[icon_type]}
                            <div> 
                                {text} 
                            </div>
                        </div>"""


def create_tmp_dir():
    tmp_dir = (
        "../demo_exp/"
        + datetime.now().strftime("%Y-%m-%d_%H-%M-%S")
        + "_"
        + str(uuid.uuid4())[:4]
    )
    os.makedirs(tmp_dir, exist_ok=True)
    print("create tmp_exp_dir", tmp_dir)
    return tmp_dir


def preprocess_imgs(tmp_dir, input_img):
    for i, img_tuple in enumerate(input_img):
        img = Image.open(img_tuple[0])
        img.thumbnail([2048, 2048], Image.Resampling.LANCZOS)
        img.save(f"{tmp_dir}/input_{i}.png")
        os.system(SEGM_i_CALL.replace("{tmp_dir}", tmp_dir).replace("{i}", str(i)))
    return [Image.open(f"{tmp_dir}/seg_{i}.png") for i in range(len(input_img))]


def ply_to_glb(ply_path):
    script_path = f"{code_dir}/ply2glb.py"
    result = subprocess.run(
        ["python", script_path, "--", ply_path],
        capture_output=True,
        text=True,
    )

    print("Output of blender script:")
    print(result.stdout)

    glb_path = ply_path.replace(".ply", ".glb")
    return glb_path


def mesh_gen(tmp_dir, use_seg):
    os.system(
        UNPOSED_CALL.replace("{tmp_dir}", tmp_dir).replace("{use_seg}", str(use_seg))
    )
    os.system(MESH_CALL.replace("{tmp_dir}", tmp_dir))

    mesh = trimesh.load_mesh(f"{tmp_dir}/mesh.ply")
    vertex_normals = mesh.vertex_normals
    colors = (-vertex_normals + 1) / 2.0
    colors = (colors * 255).astype(np.uint8)  # Convert to 8-bit color
    mesh.visual.vertex_colors = colors
    mesh.export(f"{tmp_dir}/mesh_normal.ply", file_type="ply")

    color_path = ply_to_glb(f"{tmp_dir}/mesh.ply")
    normal_path = ply_to_glb(f"{tmp_dir}/mesh_normal.ply")

    return color_path, normal_path


def feed_example_to_gallery(img):
    for display_img in display_imgs:
        display_img = display_img[0]
        diff = ImageChops.difference(img, display_img)
        if not diff.getbbox():  # two images are the same
            img_id = display_img.filename
            data_dir = os.path.join(data_folder, str(img_id))
            data_fns = os.listdir(data_dir)
            data_fns.sort()
            data_imgs = []
            for data_fn in data_fns:
                file_path = os.path.join(data_dir, data_fn)
                img = Image.open(file_path)
                data_imgs.append(img)
            return data_imgs
    return [img]


custom_theme = gr.themes.Soft(primary_hue="blue").set(
    button_secondary_background_fill="*neutral_100",
    button_secondary_background_fill_hover="*neutral_200",
)

# Gradio blocks
with gr.Blocks(title=_TITLE, css="style.css", theme=custom_theme) as demo:
    tmp_dir_unposed = gr.State("./demo_exp/placeholder")
    display_folder = os.path.join(os.path.dirname(__file__), "examples_display")
    display_fns = os.listdir(display_folder)
    display_fns.sort()
    display_imgs = []
    for i, display_fn in enumerate(display_fns):
        file_path = os.path.join(display_folder, display_fn)
        img = Image.open(file_path)
        img.filename = i
        display_imgs.append([img])
    data_folder = os.path.join(os.path.dirname(__file__), "examples_data")

    # UI
    with gr.Row():
        gr.Markdown("# " + _TITLE)
    with gr.Row():
        gr.Markdown("### " + _DESCRIPTION)
    with gr.Row():
        gr.Markdown(_PR)
    with gr.Row():
        guide_text = gr.HTML(
            message("Input image(s) of object that you want to generate mesh with.")
        )
    with gr.Row(variant="panel"):
        with gr.Column():
            with gr.Row():
                with gr.Column(scale=5):
                    input_gallery = gr.Gallery(
                        label="Input Images",
                        show_label=False,
                        columns=[3],
                        rows=[2],
                        object_fit="contain",
                        height=400,
                        show_share_button=False,
                    )
                    input_image = gr.Image(
                        type="pil",
                        image_mode="RGBA",
                        visible=False,
                    )
                with gr.Column(scale=5):
                    processed_gallery = gr.Gallery(
                        label="Background Removal",
                        columns=[3],
                        rows=[2],
                        object_fit="contain",
                        height=400,
                        interactive=False,
                        show_share_button=False,
                    )
            with gr.Row():
                with gr.Column(scale=5):
                    example = gr.Examples(
                        examples=display_imgs,
                        inputs=[input_image],
                        outputs=[input_gallery],
                        fn=feed_example_to_gallery,
                        label="Image Examples (Click one of the images below to start)",
                        examples_per_page=10,
                        run_on_click=True,
                    )
                with gr.Column(scale=5):
                    with gr.Row():
                        bg_removed_checkbox = gr.Checkbox(
                            value=True,
                            label="Use background removed images (uncheck to use original)",
                            interactive=True,
                        )
                    with gr.Row():
                        run_btn = gr.Button(
                            "Generate",
                            variant="primary",
                            interactive=False,
                        )
            with gr.Row():
                with gr.Column(scale=5):
                    mesh_output = Model3DColor(
                        label="Generated Mesh (color)",
                        elem_id="mesh-out",
                        height=400,
                    )
                with gr.Column(scale=5):
                    mesh_output_normal = Model3DNormal(
                        label="Generated Mesh (normal)",
                        elem_id="mesh-normal-out",
                        height=400,
                    )

    # Callbacks
    generating_mesh = gr.State(False)
    disable_button = lambda: gr.Button(interactive=False)
    enable_button = lambda: gr.Button(interactive=True)
    update_guide = lambda GUIDE_TEXT, icon_type="info": gr.HTML(
        value=message(GUIDE_TEXT, icon_type)
    )

    def is_cleared(content):
        if content:
            raise ValueError  # gr.Error(visible=False) doesn't work, trick for not showing error message

    def not_cleared(content):
        if not content:
            raise ValueError

    def toggle_mesh_generation_status(generating_mesh):
        generating_mesh = not generating_mesh
        return generating_mesh

    def is_generating_mesh(generating_mesh):
        if generating_mesh:
            raise ValueError

    # Upload event listener for input gallery
    input_gallery.upload(
        fn=disable_button,
        outputs=[run_btn],
        queue=False,
    ).success(
        fn=create_tmp_dir,
        outputs=[tmp_dir_unposed],
        queue=True,
    ).success(
        fn=partial(
            update_guide, "Removing background of the input image(s)...", "wait"
        ),
        outputs=[guide_text],
        queue=False,
    ).success(
        fn=preprocess_imgs,
        inputs=[tmp_dir_unposed, input_gallery],
        outputs=[processed_gallery],
        queue=True,
    ).success(
        fn=partial(update_guide, "Click <b>Generate</b> to generate mesh.", "cursor"),
        outputs=[guide_text],
        queue=False,
    ).success(
        fn=is_generating_mesh,
        inputs=[generating_mesh],
        queue=False,
    ).success(
        fn=enable_button,
        outputs=[run_btn],
        queue=False,
    )

    # Clear event listener for input gallery
    input_gallery.change(
        fn=is_cleared,
        inputs=[input_gallery],
        queue=False,
    ).success(
        fn=disable_button,
        outputs=[run_btn],
        queue=False,
    ).success(
        fn=lambda: None,
        outputs=[input_image],
        queue=False,
    ).success(
        fn=lambda: None,
        outputs=[processed_gallery],
        queue=False,
    ).success(
        fn=partial(
            update_guide,
            "Input image(s) of object that you want to generate mesh with.",
            "info",
        ),
        outputs=[guide_text],
        queue=False,
    )

    # Change event listener for input image
    input_image.change(
        fn=not_cleared,
        inputs=[input_image],
        queue=False,
    ).success(
        fn=disable_button,
        outputs=run_btn,
        queue=False,
    ).success(
        fn=create_tmp_dir,
        outputs=tmp_dir_unposed,
        queue=True,
    ).success(
        fn=partial(
            update_guide, "Removing background of the input image(s)...", "wait"
        ),
        outputs=[guide_text],
        queue=False,
    ).success(
        fn=preprocess_imgs,
        inputs=[tmp_dir_unposed, input_gallery],
        outputs=[processed_gallery],
        queue=True,
    ).success(
        fn=partial(update_guide, "Click <b>Generate</b> to generate mesh.", "cursor"),
        outputs=[guide_text],
        queue=False,
    ).success(
        fn=is_generating_mesh,
        inputs=[generating_mesh],
        queue=False,
    ).success(
        fn=enable_button,
        outputs=run_btn,
        queue=False,
    )

    # Click event listener for run button
    run_btn.click(
        fn=disable_button,
        outputs=[run_btn],
        queue=False,
    ).success(
        fn=lambda: None,
        outputs=[mesh_output],
        queue=False,
    ).success(
        fn=lambda: None,
        outputs=[mesh_output_normal],
        queue=False,
    ).success(
        fn=partial(update_guide, "Generating the mesh...", "wait"),
        outputs=[guide_text],
        queue=False,
    ).success(
        fn=toggle_mesh_generation_status,
        inputs=[generating_mesh],
        outputs=[generating_mesh],
        queue=False,
    ).success(
        fn=mesh_gen,
        inputs=[tmp_dir_unposed, bg_removed_checkbox],
        outputs=[mesh_output, mesh_output_normal],
        queue=True,
    ).success(
        fn=toggle_mesh_generation_status,
        inputs=[generating_mesh],
        outputs=[generating_mesh],
        queue=False,
    ).success(
        fn=partial(
            update_guide,
            "Successfully generated the mesh. (It might take a few seconds to load the mesh)",
            "done",
        ),
        outputs=[guide_text],
        queue=False,
    ).success(
        fn=not_cleared,
        inputs=[input_gallery],
        queue=False,
    ).success(
        fn=enable_button,
        outputs=[run_btn],
        queue=False,
    )

demo.queue().launch(
    debug=False,
    share=False,
    inline=False,
    show_api=False,
    server_name="0.0.0.0",
)