| import argparse |
| import os |
| import torch |
| import torchvision |
| import numpy as np |
| from tqdm import tqdm |
| import sys |
| import concurrent.futures |
|
|
| sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..'))) |
| from core.registry import METHOD_REGISTRY |
| import methods |
| from ufd_evalkit.geometric import depth_to_normal |
|
|
| def parse_args(): |
| parser = argparse.ArgumentParser() |
| parser.add_argument("--method", type=str, required=True) |
| parser.add_argument("--source_path", type=str, required=True) |
| parser.add_argument("--model_path", type=str, required=True) |
| parser.add_argument("--iteration", type=int, default=30000) |
| parser.add_argument("--resolution", type=int, default=-1) |
| return parser.parse_args() |
|
|
| def save_np(data, path): np.save(path, data) |
|
|
| def main(): |
| args = parse_args() |
| |
| |
| model_class = methods.load_method(args.method) |
| |
| dataset_config = {"source_path": args.source_path, "model_path": args.model_path, "resolution": args.resolution} |
| model = model_class(dataset_config, hyperparams={}) |
| |
| print(f"📥 [SplatAtlas] 正在请求变体自行加载资产... (Iteration: {args.iteration})") |
| try: |
| |
| model.load(args.model_path, args.iteration) |
| except Exception as e: |
| print(f"❌ 资产加载失败，请检查 [{args.method}] 的 load() 方法实现: {e}") |
| return |
| |
| print(f"📷 [SplatAtlas] 启动离线资产渲染 (异步 I/O 开启)") |
| |
| executor = concurrent.futures.ThreadPoolExecutor(max_workers=16) |
| futures = [] |
| |
| test_render_dir = os.path.join(args.model_path, f"renders_test_{args.iteration}") |
| test_depth_dir = os.path.join(args.model_path, f"depths_test_{args.iteration}") |
| test_normal_dir = os.path.join(args.model_path, f"normals_test_{args.iteration}") |
| test_gt_dir = os.path.join(args.model_path, f"gt_test_{args.iteration}") |
| for d in [test_render_dir, test_depth_dir, test_normal_dir, test_gt_dir]: os.makedirs(d, exist_ok=True) |
| |
| train_render_dir = os.path.join(args.model_path, f"renders_train_{args.iteration}") |
| for d in [train_render_dir]: os.makedirs(d, exist_ok=True) |
|
|
| with torch.no_grad(): |
| test_cams = model.scene.getTestCameras() |
| if test_cams and isinstance(test_cams[0], list): |
| flat_cams = [c for sub in test_cams for c in sub] |
| else: |
| flat_cams = test_cams |
| for idx, view in enumerate(tqdm(flat_cams, desc="Rendering Test")): |
| render_pkg = model.render(view) |
| rendering = render_pkg.get("render", render_pkg.get("image")).detach().cpu() |
| depth = render_pkg.get("depth", None) |
| |
| futures.append(executor.submit(torchvision.utils.save_image, rendering, os.path.join(test_render_dir, f'{idx:05d}.png'))) |
| futures.append(executor.submit(torchvision.utils.save_image, view.original_image.cpu(), os.path.join(test_gt_dir, f'{idx:05d}.png'))) |
| |
| if depth is not None: |
| depth_cpu = depth.detach().cpu() |
| normal = render_pkg.get("normal", depth_to_normal(depth)).detach().cpu() |
| futures.append(executor.submit(save_np, depth_cpu.numpy(), os.path.join(test_depth_dir, f'{idx:05d}.npy'))) |
| futures.append(executor.submit(torchvision.utils.save_image, normal, os.path.join(test_normal_dir, f'{idx:05d}.png'))) |
| |
| for idx, view in enumerate(tqdm(model.scene.getTrainCameras(), desc="Rendering Train")): |
| render_pkg = model.render(view) |
| rendering = render_pkg.get("render", render_pkg.get("image")).detach().cpu() |
| futures.append(executor.submit(torchvision.utils.save_image, rendering, os.path.join(train_render_dir, f'{idx:05d}.png'))) |
|
|
| print("⏳ [Async I/O] GPU 渲染完毕！正在等待 CPU 将所有数据落盘...") |
| concurrent.futures.wait(futures) |
| executor.shutdown() |
| |
| flag_path = os.path.join(args.model_path, f"render_complete_{args.iteration}.flag") |
| with open(flag_path, "w") as f: f.write("Completed by Async Engine") |
| print("✅ 完美落盘，标记成功！") |
|
|
| if __name__ == "__main__": |
| main() |
|
|
