generate_template_mesh.py

#!/usr/bin/env python3
# Copyright 2004-present Facebook. All Rights Reserved.

import argparse
import json
import numpy as np
import os
import torch

import deep_sdf
import deep_sdf.workspace as ws


def code_to_mesh(experiment_directory, checkpoint):

    specs_filename = os.path.join(experiment_directory, "specs.json")

    if not os.path.isfile(specs_filename):
        raise Exception(
            'The experiment directory does not include specifications file "specs.json"'
        )

    specs = json.load(open(specs_filename))

    arch = __import__("networks." + specs["NetworkArch"], fromlist=["Decoder"])

    latent_size = specs["CodeLength"]

    decoder = arch.Decoder(latent_size, **specs["NetworkSpecs"])
    decoder = torch.nn.DataParallel(decoder)

    saved_model_state = torch.load(
        os.path.join(experiment_directory, ws.model_params_subdir, checkpoint + ".pth")
    )
    saved_model_epoch = saved_model_state["epoch"]

    decoder.load_state_dict(saved_model_state["model_state_dict"])

    decoder = decoder.module.cuda()

    decoder.eval()

    mesh_dir = os.path.join(
        experiment_directory,
        ws.training_meshes_subdir,
        str(saved_model_epoch),
    )

    if not os.path.isdir(mesh_dir):
        os.makedirs(mesh_dir)

    mesh_filename = os.path.join(mesh_dir, 'template')

    print(mesh_filename)
    offset = None
    scale = None

    with torch.no_grad():
        deep_sdf.mesh.create_mesh(
            decoder.forward_template,
            None,
            mesh_filename,
            N=512,
            max_batch=int(2 ** 20),
            offset=offset,
            scale=scale,
            volume_size=2.0
        )


if __name__ == "__main__":

    arg_parser = argparse.ArgumentParser(
        description="Use a trained NDF decoder to generate a mesh given a latent code."
    )
    arg_parser.add_argument(
        "--experiment",
        "-e",
        dest="experiment_directory",
        required=True,
        help="The experiment directory which includes specifications and saved model "
        + "files to use for reconstruction",
    )
    arg_parser.add_argument(
        "--checkpoint",
        "-c",
        dest="checkpoint",
        default="latest",
        help="The checkpoint weights to use. This can be a number indicated an epoch "
        + "or 'latest' for the latest weights (this is the default)",
    )
    deep_sdf.add_common_args(arg_parser)

    args = arg_parser.parse_args()

    deep_sdf.configure_logging(args)

    code_to_mesh(args.experiment_directory, args.checkpoint)