Skip to content

Commit 74c3c52

Browse files
HochCCHochCC
authored
feat(texture): Refine texture back-projection. (#58)
1 parent fb637f9 commit 74c3c52

17 files changed

Lines changed: 727 additions & 108 deletions

‎README.md‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -331,4 +331,4 @@ EmbodiedGen builds upon the following amazing projects and models:
331331

332332
## ⚖️ License
333333

334-
This project is licensed under the [Apache License 2.0](LICENSE). See the `LICENSE` file for details.
334+
This project is licensed under the [Apache License 2.0](docs/LICENSE). See the `LICENSE` file for details.

‎apps/app_style.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -4,7 +4,7 @@
44
lighting_css = """
55
<style>
66
#lighter_mesh canvas {
7-
filter: brightness(1.9) !important;
7+
filter: brightness(2.0) !important;
88
}
99
</style>
1010
"""

‎apps/common.py‎

Lines changed: 62 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -32,8 +32,9 @@
3232
from easydict import EasyDict as edict
3333
from PIL import Image
3434
from embodied_gen.data.backproject_v2 import entrypoint as backproject_api
35+
from embodied_gen.data.backproject_v3 import entrypoint as backproject_api_v3
3536
from embodied_gen.data.differentiable_render import entrypoint as render_api
36-
from embodied_gen.data.utils import trellis_preprocess, zip_files
37+
from embodied_gen.data.utils import resize_pil, trellis_preprocess, zip_files
3738
from embodied_gen.models.delight_model import DelightingModel
3839
from embodied_gen.models.gs_model import GaussianOperator
3940
from embodied_gen.models.segment_model import (
@@ -131,8 +132,8 @@ def build_covariance_from_scaling_rotation(
131132
Gaussian.setup_functions = patched_setup_functions
132133

133134

134-
DELIGHT = DelightingModel()
135-
IMAGESR_MODEL = ImageRealESRGAN(outscale=4)
135+
# DELIGHT = DelightingModel()
136+
# IMAGESR_MODEL = ImageRealESRGAN(outscale=4)
136137
# IMAGESR_MODEL = ImageStableSR()
137138
if os.getenv("GRADIO_APP") == "imageto3d":
138139
RBG_REMOVER = RembgRemover()
@@ -169,6 +170,8 @@ def build_covariance_from_scaling_rotation(
169170
)
170171
os.makedirs(TMP_DIR, exist_ok=True)
171172
elif os.getenv("GRADIO_APP") == "texture_edit":
173+
DELIGHT = DelightingModel()
174+
IMAGESR_MODEL = ImageRealESRGAN(outscale=4)
172175
PIPELINE_IP = build_texture_gen_pipe(
173176
base_ckpt_dir="./weights",
174177
ip_adapt_scale=0.7,
@@ -205,7 +208,7 @@ def preprocess_image_fn(
205208
elif isinstance(image, np.ndarray):
206209
image = Image.fromarray(image)
207210

208-
image_cache = image.copy().resize((512, 512))
211+
image_cache = resize_pil(image.copy(), 1024)
209212

210213
bg_remover = RBG_REMOVER if rmbg_tag == "rembg" else RBG14_REMOVER
211214
image = bg_remover(image)
@@ -221,7 +224,7 @@ def preprocess_sam_image_fn(
221224
image = Image.fromarray(image)
222225

223226
sam_image = SAM_PREDICTOR.preprocess_image(image)
224-
image_cache = Image.fromarray(sam_image).resize((512, 512))
227+
image_cache = sam_image.copy()
225228
SAM_PREDICTOR.predictor.set_image(sam_image)
226229

227230
return sam_image, image_cache
@@ -512,6 +515,60 @@ def extract_3d_representations_v2(
512515
return mesh_glb_path, gs_path, mesh_obj_path, aligned_gs_path
513516

514517

518+
def extract_3d_representations_v3(
519+
state: dict,
520+
enable_delight: bool,
521+
texture_size: int,
522+
req: gr.Request,
523+
):
524+
output_root = TMP_DIR
525+
user_dir = os.path.join(output_root, str(req.session_hash))
526+
gs_model, mesh_model = unpack_state(state, device="cpu")
527+
528+
filename = "sample"
529+
gs_path = os.path.join(user_dir, f"{filename}_gs.ply")
530+
gs_model.save_ply(gs_path)
531+
532+
# Rotate mesh and GS by 90 degrees around Z-axis.
533+
rot_matrix = [[0, 0, -1], [0, 1, 0], [1, 0, 0]]
534+
gs_add_rot = [[1, 0, 0], [0, -1, 0], [0, 0, -1]]
535+
mesh_add_rot = [[1, 0, 0], [0, 0, -1], [0, 1, 0]]
536+
537+
# Addtional rotation for GS to align mesh.
538+
gs_rot = np.array(gs_add_rot) @ np.array(rot_matrix)
539+
pose = GaussianOperator.trans_to_quatpose(gs_rot)
540+
aligned_gs_path = gs_path.replace(".ply", "_aligned.ply")
541+
GaussianOperator.resave_ply(
542+
in_ply=gs_path,
543+
out_ply=aligned_gs_path,
544+
instance_pose=pose,
545+
device="cpu",
546+
)
547+
548+
mesh = trimesh.Trimesh(
549+
vertices=mesh_model.vertices.cpu().numpy(),
550+
faces=mesh_model.faces.cpu().numpy(),
551+
)
552+
mesh.vertices = mesh.vertices @ np.array(mesh_add_rot)
553+
mesh.vertices = mesh.vertices @ np.array(rot_matrix)
554+
555+
mesh_obj_path = os.path.join(user_dir, f"{filename}.obj")
556+
mesh.export(mesh_obj_path)
557+
558+
mesh = backproject_api_v3(
559+
gs_path=aligned_gs_path,
560+
mesh_path=mesh_obj_path,
561+
output_path=mesh_obj_path,
562+
skip_fix_mesh=False,
563+
texture_size=texture_size,
564+
)
565+
566+
mesh_glb_path = os.path.join(user_dir, f"{filename}.glb")
567+
mesh.export(mesh_glb_path)
568+
569+
return mesh_glb_path, gs_path, mesh_obj_path, aligned_gs_path
570+
571+
515572
def extract_urdf(
516573
gs_path: str,
517574
mesh_obj_path: str,

‎apps/image_to_3d.py‎

Lines changed: 16 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -27,7 +27,7 @@
2727
VERSION,
2828
active_btn_by_content,
2929
end_session,
30-
extract_3d_representations_v2,
30+
extract_3d_representations_v3,
3131
extract_urdf,
3232
get_seed,
3333
image_to_3d,
@@ -179,17 +179,17 @@
179179
)
180180

181181
generate_btn = gr.Button(
182-
"🚀 1. Generate(~0.5 mins)",
182+
"🚀 1. Generate(~2 mins)",
183183
variant="primary",
184184
interactive=False,
185185
)
186186
model_output_obj = gr.Textbox(label="raw mesh .obj", visible=False)
187-
with gr.Row():
188-
extract_rep3d_btn = gr.Button(
189-
"🔍 2. Extract 3D Representation(~2 mins)",
190-
variant="primary",
191-
interactive=False,
192-
)
187+
# with gr.Row():
188+
# extract_rep3d_btn = gr.Button(
189+
# "🔍 2. Extract 3D Representation(~2 mins)",
190+
# variant="primary",
191+
# interactive=False,
192+
# )
193193
with gr.Accordion(
194194
label="Enter Asset Attributes(optional)", open=False
195195
):
@@ -207,7 +207,7 @@
207207
)
208208
with gr.Row():
209209
extract_urdf_btn = gr.Button(
210-
"🧩 3. Extract URDF with physics(~1 mins)",
210+
"🧩 2. Extract URDF with physics(~1 mins)",
211211
variant="primary",
212212
interactive=False,
213213
)
@@ -230,7 +230,7 @@
230230
)
231231
with gr.Row():
232232
download_urdf = gr.DownloadButton(
233-
label="⬇️ 4. Download URDF",
233+
label="⬇️ 3. Download URDF",
234234
variant="primary",
235235
interactive=False,
236236
)
@@ -326,7 +326,7 @@
326326
image_prompt.change(
327327
lambda: tuple(
328328
[
329-
gr.Button(interactive=False),
329+
# gr.Button(interactive=False),
330330
gr.Button(interactive=False),
331331
gr.Button(interactive=False),
332332
None,
@@ -344,7 +344,7 @@
344344
]
345345
),
346346
outputs=[
347-
extract_rep3d_btn,
347+
# extract_rep3d_btn,
348348
extract_urdf_btn,
349349
download_urdf,
350350
model_output_gs,
@@ -375,7 +375,7 @@
375375
image_prompt_sam.change(
376376
lambda: tuple(
377377
[
378-
gr.Button(interactive=False),
378+
# gr.Button(interactive=False),
379379
gr.Button(interactive=False),
380380
gr.Button(interactive=False),
381381
None,
@@ -394,7 +394,7 @@
394394
]
395395
),
396396
outputs=[
397-
extract_rep3d_btn,
397+
# extract_rep3d_btn,
398398
extract_urdf_btn,
399399
download_urdf,
400400
model_output_gs,
@@ -447,12 +447,7 @@
447447
],
448448
outputs=[output_buf, video_output],
449449
).success(
450-
lambda: gr.Button(interactive=True),
451-
outputs=[extract_rep3d_btn],
452-
)
453-
454-
extract_rep3d_btn.click(
455-
extract_3d_representations_v2,
450+
extract_3d_representations_v3,
456451
inputs=[
457452
output_buf,
458453
project_delight,
@@ -495,4 +490,4 @@
495490

496491

497492
if __name__ == "__main__":
498-
demo.launch()
493+
demo.launch(server_port=8081)

‎apps/text_to_3d.py‎

Lines changed: 14 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -27,7 +27,7 @@
2727
VERSION,
2828
active_btn_by_text_content,
2929
end_session,
30-
extract_3d_representations_v2,
30+
extract_3d_representations_v3,
3131
extract_urdf,
3232
get_cached_image,
3333
get_seed,
@@ -178,17 +178,17 @@
178178
)
179179

180180
generate_btn = gr.Button(
181-
"🚀 2. Generate 3D(~0.5 mins)",
181+
"🚀 2. Generate 3D(~2 mins)",
182182
variant="primary",
183183
interactive=False,
184184
)
185185
model_output_obj = gr.Textbox(label="raw mesh .obj", visible=False)
186-
with gr.Row():
187-
extract_rep3d_btn = gr.Button(
188-
"🔍 3. Extract 3D Representation(~1 mins)",
189-
variant="primary",
190-
interactive=False,
191-
)
186+
# with gr.Row():
187+
# extract_rep3d_btn = gr.Button(
188+
# "🔍 3. Extract 3D Representation(~1 mins)",
189+
# variant="primary",
190+
# interactive=False,
191+
# )
192192
with gr.Accordion(
193193
label="Enter Asset Attributes(optional)", open=False
194194
):
@@ -206,13 +206,13 @@
206206
)
207207
with gr.Row():
208208
extract_urdf_btn = gr.Button(
209-
"🧩 4. Extract URDF with physics(~1 mins)",
209+
"🧩 3. Extract URDF with physics(~1 mins)",
210210
variant="primary",
211211
interactive=False,
212212
)
213213
with gr.Row():
214214
download_urdf = gr.DownloadButton(
215-
label="⬇️ 5. Download URDF",
215+
label="⬇️ 4. Download URDF",
216216
variant="primary",
217217
interactive=False,
218218
)
@@ -336,7 +336,7 @@
336336
generate_img_btn.click(
337337
lambda: tuple(
338338
[
339-
gr.Button(interactive=False),
339+
# gr.Button(interactive=False),
340340
gr.Button(interactive=False),
341341
gr.Button(interactive=False),
342342
gr.Button(interactive=False),
@@ -358,7 +358,7 @@
358358
]
359359
),
360360
outputs=[
361-
extract_rep3d_btn,
361+
# extract_rep3d_btn,
362362
extract_urdf_btn,
363363
download_urdf,
364364
generate_btn,
@@ -428,12 +428,7 @@
428428
],
429429
outputs=[output_buf, video_output],
430430
).success(
431-
lambda: gr.Button(interactive=True),
432-
outputs=[extract_rep3d_btn],
433-
)
434-
435-
extract_rep3d_btn.click(
436-
extract_3d_representations_v2,
431+
extract_3d_representations_v3,
437432
inputs=[
438433
output_buf,
439434
project_delight,
@@ -476,4 +471,4 @@
476471

477472

478473
if __name__ == "__main__":
479-
demo.launch()
474+
demo.launch(server_port=8082)

‎apps/texture_edit.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -381,4 +381,4 @@ def active_btn_by_content(mesh_content: gr.Model3D, text_content: gr.Textbox):
381381

382382

383383
if __name__ == "__main__":
384-
demo.launch()
384+
demo.launch(server_port=8083)

‎apps/visualize_asset.py‎

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -727,7 +727,6 @@ def update_assets(p, s, c):
727727

728728
if __name__ == "__main__":
729729
demo.launch(
730-
server_name="10.34.8.77",
731730
server_port=8088,
732731
allowed_paths=[
733732
"/horizon-bucket/robot_lab/datasets/embodiedgen/assets"

LICENSE renamed to docs/LICENSE

File renamed without changes.

‎docs/install.md‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,8 @@ conda activate embodiedgen
1414
bash install.sh basic
1515
```
1616

17+
Please `huggingface-cli login` to ensure that the ckpts can be downloaded automatically afterwards.
18+
1719
## ✅ Starting from Docker
1820

1921
We provide a pre-built Docker image on [Docker Hub](https://hub.docker.com/repository/docker/wangxinjie/embodiedgen) with a configured environment for your convenience. For more details, please refer to [Docker documentation](https://github.com/HorizonRobotics/EmbodiedGen/tree/master/docker).

‎embodied_gen/data/asset_converter.py‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -589,6 +589,8 @@ def convert(self, urdf_path: str, output_file: str):
589589
stage = Usd.Stage.Open(usd_path)
590590
layer = stage.GetRootLayer()
591591
with Usd.EditContext(stage, layer):
592+
base_prim = stage.GetPseudoRoot().GetChildren()[0]
593+
base_prim.SetMetadata("kind", "component")
592594
for prim in stage.Traverse():
593595
# Change texture path to relative path.
594596
if prim.GetName() == "material_0":

0 commit comments

Comments
 (0)