| """Oracle upper bound: feed GT stimulus frames to TRELLIS.2 and export shape meshes.""" |
| import os |
| os.environ.setdefault("OPENCV_IO_ENABLE_OPENEXR", "1") |
| os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True") |
| import sys |
| import json |
| import time |
| import argparse |
|
|
| import numpy as np |
| import torch |
| import trimesh |
| import imageio.v3 as iio |
| from PIL import Image |
|
|
| sys.path.insert(0, "/home/hubin/trellis_work/TRELLIS.2") |
| from trellis2.pipelines import Trellis2ImageTo3DPipeline |
|
|
|
|
| def load_frame(video_path, frame_idx): |
| frames = iio.imread(video_path, index=frame_idx) |
| return Image.fromarray(frames) |
|
|
|
|
| def main(): |
| parser = argparse.ArgumentParser() |
| parser.add_argument("--weights", default="/home/hubin/trellis_work/weights/TRELLIS.2-4B-merged") |
| parser.add_argument("--config_file", default="pipeline_local.json") |
| parser.add_argument("--test_list", default="/home/hubin/data/fMRI-Shape/annotations/core_test_list.txt") |
| parser.add_argument("--stimuli_dir", default="/home/hubin/data/fMRI-Shape/stimuli_test/stimuli") |
| parser.add_argument("--frame_idx", type=int, default=24) |
| parser.add_argument("--pipeline_type", default="512", choices=["512", "1024"]) |
| parser.add_argument("--out_dir", default="/home/hubin/trellis_work/outputs/upper_bound_512_f24") |
| parser.add_argument("--seed", type=int, default=42) |
| parser.add_argument("--limit", type=int, default=0) |
| parser.add_argument("--shard", type=int, default=0) |
| parser.add_argument("--num_shards", type=int, default=1) |
| args = parser.parse_args() |
|
|
| ids = [l.strip() for l in open(args.test_list) if l.strip()] |
| ids = ids[args.shard::args.num_shards] |
| if args.limit: |
| ids = ids[:args.limit] |
| os.makedirs(os.path.join(args.out_dir, "meshes"), exist_ok=True) |
| os.makedirs(os.path.join(args.out_dir, "inputs"), exist_ok=True) |
|
|
| pipeline = Trellis2ImageTo3DPipeline.from_pretrained(args.weights, config_file=args.config_file) |
| pipeline.low_vram = False |
| pipeline.cuda() |
|
|
| res = int(args.pipeline_type) |
| ss_res = {512: 32, 1024: 64}[res] |
| flow_key = f"shape_slat_flow_model_{res}" |
| stats = [] |
| for i, obj in enumerate(ids): |
| name = obj.replace("/", "_") |
| mesh_path = os.path.join(args.out_dir, "meshes", f"{name}.ply") |
| if os.path.exists(mesh_path): |
| continue |
| t0 = time.time() |
| image = load_frame(os.path.join(args.stimuli_dir, f"{obj}.mp4"), args.frame_idx) |
| image = pipeline.preprocess_image(image) |
| image.save(os.path.join(args.out_dir, "inputs", f"{name}.png")) |
|
|
| torch.manual_seed(args.seed) |
| with torch.no_grad(): |
| cond = pipeline.get_cond([image], res) |
| coords = pipeline.sample_sparse_structure(cond, ss_res, 1) |
| shape_slat = pipeline.sample_shape_slat(cond, pipeline.models[flow_key], coords) |
| meshes, _ = pipeline.decode_shape_slat(shape_slat, res) |
| mesh = meshes[0] |
| mesh.fill_holes() |
| trimesh.Trimesh( |
| vertices=mesh.vertices.detach().cpu().numpy(), |
| faces=mesh.faces.detach().cpu().numpy(), |
| process=False, |
| ).export(mesh_path) |
| dt = time.time() - t0 |
| stats.append({"id": obj, "time": dt, "n_voxels": int(coords.shape[0]), |
| "n_verts": int(mesh.vertices.shape[0])}) |
| print(f"[{i + 1}/{len(ids)}] {obj} {dt:.1f}s voxels={coords.shape[0]} verts={mesh.vertices.shape[0]}", flush=True) |
| torch.cuda.empty_cache() |
|
|
| with open(os.path.join(args.out_dir, f"stats_shard{args.shard}.json"), "w") as f: |
| json.dump(stats, f, indent=1) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
