Revision 364211dce1b0ec238d9e40c7c3240912ff82f399 authored by Guilherme G. Schardong on 06 November 2023, 11:51:46 UTC, committed by Guilherme G. Schardong on 06 November 2023, 11:51:46 UTC
1 parent 2cac1e4
reconstruct.py
#!/usr/bin/env python
# coding: utf-8
"""
Simple script to run multiple SDF reconstructions given a base log directory
and a set of checkpoints.
"""
import argparse
import os
import os.path as osp
import torch
from i3d.meshing import create_mesh
from i3d.util import from_pth
if __name__ == "__main__":
parser = argparse.ArgumentParser(
description="Run marching cubes using a trained model."
)
parser.add_argument(
"model_path",
help="Path to the PyTorch weights file"
)
parser.add_argument(
"output_path",
help="Path to the output mesh file"
)
parser.add_argument(
"--w0", type=int, default=1,
help="Value for \\omega_0. Default is 1."
)
parser.add_argument(
"--resolution", "-r", default=128, type=int,
help="Resolution to use on marching cubes. Default is 128."
)
args = parser.parse_args()
out_dir = osp.split(args.output_path)[0]
if out_dir and not osp.exists(out_dir):
os.makedirs(out_dir)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = from_pth(args.model_path, w0=args.w0).eval().to(device)
print(model)
print(f"Running marching cubes running with resolution {args.resolution}")
create_mesh(
model,
args.output_path,
N=args.resolution,
device=device
)
print("Done")
Computing file changes ...