class ToyNeuralGraphicsDataset(data.Dataset):
def __init__(self,
dir: str,
rasterization_settings: dict,
znear: float = 1.0,
zfar: float = 1000.0,
scale_min: float = 0.5,
scale_max: float = 2.0,
device: str = 'cuda'):
super(ToyNeuralGraphicsDataset, self).__init__()
device = torch.device(device)
self.device = device
self.scale_min = scale_min
self.scale_max = scale_max
self.scale_range = scale_max - scale_min
objs = [
os.path.join(dir, f) for f in os.listdir(dir) if f.endswith('.obj')
]
self.meshes = load_objs_as_meshes(objs, device=device)
R, T = look_at_view_transform(0, 0, 0)
self.cameras = FoVPerspectiveCameras(R=R,
T=T,
znear=znear,
zfar=zfar,
device=device)
self.renderer = MeshRenderer(rasterizer=MeshRasterizer(
cameras=self.cameras,
raster_settings=RasterizationSettings(**rasterization_settings),
),
shader=HardFlatShader(
device=device,
cameras=self.cameras,
))
def get_random_transform(self):
scale = (torch.rand(1).squeeze() * self.scale_range +
self.scale_min).item()
rot = random_rotation()
x, y, d = torch.rand(3)
x = x * 2.0 - 1.0
y = y * 2.0 - 1.0
trans = torch.Tensor([x, y, d])
trans = self.cameras.unproject_points(
trans.unsqueeze(0).to(self.device),
world_coordinates=False,
scaled_depth_input=True)[0].cpu()
return scale, rot, trans
def __getitem__(self, index):
index %= len(self.meshes)
scale, rot, trans = self.get_random_transform()
transform = Transform3d() \
.scale(scale) \
.compose(Rotate(rot)) \
.translate(*trans) \
.get_matrix() \
.squeeze()
mesh = self.meshes[index].scale_verts(scale)
pixels = self.renderer(mesh,
R=rot.unsqueeze(0).to(self.device),
T=trans.unsqueeze(0).to(self.device))
pixels = pixels[0, ..., :3].transpose(0, -1)
return (pixels, [transform.to(self.device)])
def __len__(self):
return len(self.meshes) * 1024
# With world coordinates +Y up, +X left and +Z in, the front of the cow is facing the -Z direction.
# So we move the camera by 180 in the azimuth direction so it is facing the front of the cow.
R, T = look_at_view_transform(0, 0, 0)
cameras = FoVPerspectiveCameras(device=device, R=R, T=T, zfar=zfar)
smin = 0.1
smax = 2.0
srange = smax - smin
scale = (torch.rand(1).squeeze() * srange + smin).item()
# Generate a random NDC coordinate https://pytorch3d.org/docs/cameras
x, y, d = torch.rand(3)
x = x * 2.0 - 1.0
y = y * 2.0 - 1.0
trans = torch.Tensor([x, y, d]).to(device)
trans = cameras.unproject_points(trans.unsqueeze(0),
world_coordinates=False,
scaled_depth_input=True)[0]
rot = random_rotations(1)[0].to(device)
transform = Transform3d() \
.scale(scale) \
.compose(Rotate(rot)) \
.translate(*trans)
# TODO: transform mesh
# Create a phong renderer by composing a rasterizer and a shader. The textured phong shader will
# interpolate the texture uv coordinates for each vertex, sample from a texture image and
# apply the Phong lighting model
renderer = MeshRenderer(rasterizer=MeshRasterizer(
cameras=cameras, raster_settings=raster_settings),
shader=SoftPhongShader(