Spaces:
Runtime error
Runtime error
File size: 4,642 Bytes
881fc76 17d4ba5 881fc76 33400d1 881fc76 33400d1 881fc76 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 |
import gradio as gr
import numpy as np
import spaces
import torch
import rembg
from PIL import Image
from functools import partial
import logging
import os
import shlex
import subprocess
import tempfile
import time
subprocess.run(shlex.split('pip install wheel/torchmcubes-0.1.0-cp310-cp310-linux_x86_64.whl'))
from tsr.system import TSR
from tsr.utils import remove_background, resize_foreground, to_gradio_3d_orientation
STEP1_HEADER = """
# Step 1: Generate the 3D Mesh
For this step, we use TripoSR, an open-source model for **fast** feedforward 3D reconstruction from a single image, developed in collaboration between [Tripo AI](https://www.tripo3d.ai/) and [Stability AI](https://stability.ai/).
During this step, you need to upload an image of what you want to generate a 3D Model from.
## 💡 Tips
- If there's a background, ✅ Remove background.
- If you find the result is unsatisfied, please try to change the foreground ratio. It might improve the results.
"""
# These part of the code (check_input_image and preprocess were taken from https://huggingface.co/spaces/stabilityai/TripoSR/blob/main/app.py)
if torch.cuda.is_available():
device = "cuda:0"
else:
device = "cpu"
model = TSR.from_pretrained(
"stabilityai/TripoSR",
config_name="config.yaml",
weight_name="model.ckpt",
)
model.renderer.set_chunk_size(131072)
model.to(device)
rembg_session = rembg.new_session()
def check_input_image(input_image):
if input_image is None:
raise gr.Error("No image uploaded!")
def preprocess(input_image, do_remove_background, foreground_ratio):
def fill_background(image):
image = np.array(image).astype(np.float32) / 255.0
image = image[:, :, :3] * image[:, :, 3:4] + (1 - image[:, :, 3:4]) * 0.5
image = Image.fromarray((image * 255.0).astype(np.uint8))
return image
if do_remove_background:
image = input_image.convert("RGB")
image = remove_background(image, rembg_session)
image = resize_foreground(image, foreground_ratio)
image = fill_background(image)
else:
image = input_image
if image.mode == "RGBA":
image = fill_background(image)
return image
@spaces.GPU
def generate(image, mc_resolution, formats=["obj", "glb"]):
scene_codes = model(image, device=device)
mesh = model.extract_mesh(scene_codes, resolution=mc_resolution)[0]
mesh = to_gradio_3d_orientation(mesh)
mesh_path_glb = tempfile.NamedTemporaryFile(suffix=f".glb", delete=False)
mesh.export(mesh_path_glb.name)
mesh_path_obj = tempfile.NamedTemporaryFile(suffix=f".obj", delete=False)
mesh.apply_scale([-1, 1, 1]) # Otherwise the visualized .obj will be flipped
mesh.export(mesh_path_obj.name)
return mesh_path_obj.name, mesh_path_glb.name
with gr.Blocks() as demo:
gr.Markdown(STEP1_HEADER)
with gr.Row(variant = "panel"):
with gr.Column():
with gr.Row():
input_image = gr.Image(
label = "Input Image",
image_mode = "RGBA",
sources = "upload",
type="pil",
elem_id="content_image")
processed_image = gr.Image(label="Processed Image", interactive=False)
with gr.Row():
with gr.Group():
do_remove_background = gr.Checkbox(
label="Remove Background",
value=True)
foreground_ratio = gr.Slider(
label="Foreground Ratio",
minimum=0.5,
maximum=1.0,
value=0.85,
step=0.05,
)
mc_resolution = gr.Slider(
label="Marching Cubes Resolution",
minimum=32,
maximum=320,
value=256,
step=32
)
with gr.Row():
step1_submit = gr.Button("Generate", elem_id="generate", variant="primary")
with gr.Column():
with gr.Tab("OBJ"):
output_model_obj = gr.Model3D(
label = "Output Model (OBJ Format)",
interative = False,
)
gr.Markdown("Note: Downloaded object will be flipped in case of .obj export. Export .glb instead or manually flip it before usage.")
with gr.Tab("GLB"):
output_model_glb = gr.Model3D(
label="Output Model (GLB Format)",
interactive=False,
)
gr.Markdown("Note: The model shown here has a darker appearance. Download to get correct results.")
step1_submit.click(fn=check_input_image, inputs=[input_image]).success(
fn=preprocess,
inputs=[input_image, do_remove_background, foreground_ratio],
outputs=[processed_image],
).success(
fn=generate,
inputs=[processed_image, mc_resolution],
outputs=[output_model_obj, output_model_glb],
)
demo.queue(max_size=10)
demo.launch()
|