Function export_model

spectrum_visualization/get_shape_spectra.py:26–72

(ckpt_path, model_name, N=512, model_type='bacon', hidden_layers=8,
                 hidden_size=256, output_layers=[1, 2, 4, 8], w0=30, pe=8,
                 filter_mesh=False, scaling=None, return_sdf=False)

def export_model(ckpt_path, model_name, N=512, model_type='bacon', hidden_layers=8,
                 hidden_size=256, output_layers=[1, 2, 4, 8], w0=30, pe=8,
                 filter_mesh=False, scaling=None, return_sdf=False):

    with HiddenPrints():
        # the network has 4 output levels of detail
        num_outputs = len(output_layers)
        max_frequency = 3*(32,)

        # load model
        if len(output_layers) > 1:
            model = modules.MultiscaleBACON(3, hidden_size, 1,
                                            hidden_layers=hidden_layers,
                                            bias=True,
                                            frequency=max_frequency,
                                            quantization_interval=np.pi,
                                            is_sdf=True,
                                            output_layers=output_layers,
                                            reuse_filters=True)

    print(model)
    ckpt = torch.load(ckpt_path, map_location=device)
    model.load_state_dict(ckpt)
    model = model.to(device)

    # write output
    x = torch.linspace(-0.5, 0.5, N)
    if return_sdf:
        x = torch.arange(-N//2, N//2) / N
        x = x.float()
    x, y, z = torch.meshgrid(x, x, x)
    render_coords = torch.stack((x.flatten(), y.flatten(), z.flatten()), dim=-1).to(device)
    sdf_values = [np.zeros((N**3, 1)) for i in range(num_outputs)]

    # render in a batched fashion to save memory
    bsize = int(128**2)
    for i in tqdm(range(int(N**3 / bsize))):
        coords = render_coords[i*bsize:(i+1)*bsize, :]
        out = model({'coords': coords})['model_out']

        if not isinstance(out, list):
            out = [out,]

        for idx, sdf in enumerate(out):
            sdf_values[idx][i*bsize:(i+1)*bsize] = sdf.detach().cpu().numpy()

    return [sdf.reshape(N, N, N) for sdf in sdf_values]


def normalize(coords, scaling=0.9):