96 lines
3.3 KiB
Python
96 lines
3.3 KiB
Python
|
#!/usr/bin/env python3
|
||
|
|
# -*- coding: utf-8 -*-
|
||
|
|
|
||
|
|
import os
|
||
|
|
import argparse
|
||
|
|
from PIL import Image
|
||
|
|
|
||
|
|
from trellis.pipelines import TrellisImageTo3DPipeline
|
||
|
|
from trellis.utils import render_utils, postprocessing_utils
|
||
|
|
|
||
|
|
def build_parser():
|
||
|
|
p = argparse.ArgumentParser("TRELLIS CLI: single image -> 3D")
|
||
|
|
|
||
|
|
p.add_argument("-i", "--image", required=True, help="Input image path")
|
||
|
|
p.add_argument("-o", "--out_dir", default="trellis_out", help="Output directory")
|
||
|
|
|
||
|
|
p.add_argument("--seed", type=int, default=1)
|
||
|
|
p.add_argument("--steps_sparse", type=int, default=12)
|
||
|
|
p.add_argument("--cfg_sparse", type=float, default=7.5)
|
||
|
|
p.add_argument("--steps_slat", type=int, default=12)
|
||
|
|
p.add_argument("--cfg_slat", type=float, default=3.0)
|
||
|
|
|
||
|
|
p.add_argument("--simplify", type=float, default=0.95)
|
||
|
|
p.add_argument("--texture_size", type=int, default=1024)
|
||
|
|
|
||
|
|
# Export GLB (default True)
|
||
|
|
p.add_argument("--export_glb", dest="export_glb", action="store_true", default=True)
|
||
|
|
p.add_argument("--no-export_glb", dest="export_glb", action="store_false")
|
||
|
|
|
||
|
|
# Save PLY (default True)
|
||
|
|
p.add_argument("--save_ply", dest="save_ply", action="store_true", default=True)
|
||
|
|
p.add_argument("--no-save_ply", dest="save_ply", action="store_false")
|
||
|
|
|
||
|
|
# Save videos (default False, plus explicit toggle)
|
||
|
|
p.add_argument("--save_video", dest="save_video", action="store_true", default=True)
|
||
|
|
p.add_argument("--no-save_video", dest="save_video", action="store_false")
|
||
|
|
|
||
|
|
p.add_argument("--fps", type=int, default=30)
|
||
|
|
p.add_argument("--video_gs_name", type=str, default="sample_gs.mp4")
|
||
|
|
p.add_argument("--video_rf_name", type=str, default="sample_rf.mp4")
|
||
|
|
p.add_argument("--video_mesh_name", type=str, default="sample_mesh.mp4")
|
||
|
|
|
||
|
|
return p
|
||
|
|
|
||
|
|
def main():
|
||
|
|
args = build_parser().parse_args()
|
||
|
|
os.makedirs(args.out_dir, exist_ok=True)
|
||
|
|
|
||
|
|
# Optional env
|
||
|
|
os.environ.setdefault("SPCONV_ALGO", "native")
|
||
|
|
|
||
|
|
pipeline = TrellisImageTo3DPipeline.from_pretrained("microsoft/TRELLIS-image-large")
|
||
|
|
pipeline.cuda()
|
||
|
|
|
||
|
|
image = Image.open(args.image)
|
||
|
|
|
||
|
|
outputs = pipeline.run(
|
||
|
|
image,
|
||
|
|
seed=args.seed,
|
||
|
|
sparse_structure_sampler_params={
|
||
|
|
"steps": args.steps_sparse,
|
||
|
|
"cfg_strength": args.cfg_sparse,
|
||
|
|
},
|
||
|
|
slat_sampler_params={
|
||
|
|
"steps": args.steps_slat,
|
||
|
|
"cfg_strength": args.cfg_slat,
|
||
|
|
},
|
||
|
|
)
|
||
|
|
|
||
|
|
if args.save_video:
|
||
|
|
import imageio
|
||
|
|
|
||
|
|
video = render_utils.render_video(outputs["gaussian"][0])["color"]
|
||
|
|
imageio.mimsave(os.path.join(args.out_dir, args.video_gs_name), video, fps=args.fps)
|
||
|
|
|
||
|
|
video = render_utils.render_video(outputs["radiance_field"][0])["color"]
|
||
|
|
imageio.mimsave(os.path.join(args.out_dir, args.video_rf_name), video, fps=args.fps)
|
||
|
|
|
||
|
|
video = render_utils.render_video(outputs["mesh"][0])["normal"]
|
||
|
|
imageio.mimsave(os.path.join(args.out_dir, args.video_mesh_name), video, fps=args.fps)
|
||
|
|
|
||
|
|
if args.export_glb:
|
||
|
|
glb = postprocessing_utils.to_glb(
|
||
|
|
outputs["gaussian"][0],
|
||
|
|
outputs["mesh"][0],
|
||
|
|
simplify=args.simplify,
|
||
|
|
texture_size=args.texture_size,
|
||
|
|
)
|
||
|
|
glb.export(os.path.join(args.out_dir, "sample.glb"))
|
||
|
|
|
||
|
|
if args.save_ply:
|
||
|
|
outputs["gaussian"][0].save_ply(os.path.join(args.out_dir, "sample.ply"))
|
||
|
|
|
||
|
|
if __name__ == "__main__":
|
||
|
|
main()
|