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/bootstrap@5.3.2/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",
)