import os 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 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()) os.system(f"pip install {code_dir}/gradio_model3dcolor-0.0.1-py3-none-any.whl") os.system(f"pip install {code_dir}/gradio_model3dnormal-0.0.1-py3-none-any.whl") from gradio_model3dcolor import Model3DColor from gradio_model3dnormal import Model3DNormal with open(f'{code_dir}/api.json', 'r') as file: api_dict = json.load(file) SEGM_i_CALL = api_dict["SEGM_i_CALL"] SEGM_CALL = api_dict["SEGM_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 = """
Check out Hillbot (sudoAI) for more details and advanced features.
""" STYLE = """ """ # info (info-circle-fill), cursor (hand-index-thumb), wait (hourglass-split), done (check-circle) ICONS = { "info": """ """, "cursor": """ """, "wait": """ """, "done": """ """, } icons2alert = { "info": "primary", # blue "cursor": "info", # light blue "wait": "secondary", # gray "done": "success", # green } def message(text, icon_type="info"): return f"""{STYLE} """ 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, idx=None): if isinstance(input_img, list) and idx is None: for i, img_tuple in enumerate(input_img): Image.open(img_tuple[0]).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))] if idx is not None: print("image idx:", int(idx)) input_img = Image.open(input_img[int(idx)][0]) input_img.save(f"{tmp_dir}/input.png") os.system(SEGM_CALL.replace("{tmp_dir}", tmp_dir)) processed_img = Image.open(f"{tmp_dir}/seg.png") return processed_img.resize((320, 320), Image.Resampling.LANCZOS) def ply_to_glb(ply_path): script_path = "../utils/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, ) 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, ) 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 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 # gr.Error(visible=False) doesn't work, trick for not showing error message # 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=False, ).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 Generate to generate mesh.", "cursor"), outputs=[guide_text], 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=lambda: None, outputs=[mesh_output], queue=False, ).success( fn=lambda: None, outputs=[mesh_output_normal], 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=lambda: None, outputs=[mesh_output], queue=False, ).success( fn=lambda: None, outputs=[mesh_output_normal], queue=False, ).success( fn=create_tmp_dir, outputs=tmp_dir_unposed, queue=False, ).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 Generate to generate mesh.", "cursor"), outputs=[guide_text], 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=mesh_gen, inputs=[tmp_dir_unposed, bg_removed_checkbox], outputs=[mesh_output, mesh_output_normal], queue=True, ).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=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", )