ndrplz / differentiable-renderer Goto Github PK

import glob

import matplotlib.pyplot as plt
import numpy as np
import torch

from rastering import Rasterer
from rastering import RotoTranslation
from rastering import Vector


if __name__ == '__main__':
    device = 'cpu'
    if device not in ['cpu', 'cuda']:
        raise ValueError('Unknown device.')

    camera_pose = RotoTranslation(rotation=Vector(x=0., y=0., z=0.),
                                  translation=Vector(x=0., y=0., z=0.),
                                  angle_unit='degrees')
    print(camera_pose, '\n')

    # Toy dataset containing 5 random car meshes from Shapenetcore.
    mesh_dataset = [np.load(mesh_path) for mesh_path in glob.glob('data/*/*.npy')]

    print(mesh_dataset[0][:5])
    print([m.max() for m in mesh_dataset])
   

    mesh = np.array([[[0, 0, 0], [0, 1, 0], [1, 1, 0]],
             [[0, 0, 0], [1, 1, 0], [1, 0, 0]]]).astype(np.float32)

    print(mesh.shape, 'MESH')

    camera_intrinsics = {'resolution_px': (5, 5), 'diagonal_mm': 35, 'focal_len_mm': 30}
    renderer = Rasterer(meshes=[mesh], max_triangles=100, **camera_intrinsics)
    renderer.to(device)

    # Sample a bunch of models from the mesh dataset. There will be rendered in the same batch.
    n_mesh_to_render = 1
    #model_idxs = np.random.choice(np.arange(len(mesh_dataset)), size=n_mesh_to_render)
    model_idxs = np.array([0])

    # Load camera pose. In this case it is the same for all renderings.
    camera_matrix = camera_pose.matrix
    camera_matrices = np.tile(camera_matrix[None, ...], [n_mesh_to_render, 1, 1])

    render_output = renderer(torch.from_numpy(model_idxs).to(device),
                             torch.from_numpy(camera_matrices).to(device))

    for b in range(n_mesh_to_render):
        print(render_output)
        print(render_output[0].shape, 'OUTPUT SHAPE')
        plt.imshow(render_output[b].to('cpu').numpy(), cmap='gray')
        plt.waitforbuttonpress()