1

2026-03-17 11:29:31 +08:00
parent 24e4c120be
commit a6d9bac6d0
11 changed files with 2159 additions and 0 deletions
--- a/single_image_to_3D.py
+++ b/single_image_to_3D.py
@@ -0,0 +1,95 @@
+#!/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()