def forward(ctx, face_vertices, textures, image_size=256,
background_color=[0, 0, 0], near=1, far=100,
fill_back=True, eps=1e-3,
sigma_val=1e-5, dist_func='euclidean', dist_eps=1e-4,
gamma_val=1e-4, aggr_func_rgb='softmax', aggr_func_alpha='prod',
texture_type='surface'):
# face_vertices: [nb, nf, 9]
# textures: [nb, nf, 9]
func_dist_map = {'hard': 0, 'barycentric': 1, 'euclidean': 2}
func_rgb_map = {'hard': 0, 'softmax': 1}
func_alpha_map = {'hard': 0, 'sum': 1, 'prod': 2}
func_map_sample = {'surface': 0, 'vertex': 1}
ctx.image_size = image_size
ctx.background_color = background_color
ctx.near = near
ctx.far = far
ctx.eps = eps
ctx.sigma_val = sigma_val
ctx.gamma_val = gamma_val
ctx.func_dist_type = func_dist_map[dist_func]
ctx.dist_eps = np.log(1. / dist_eps - 1.)
ctx.func_rgb_type = func_rgb_map[aggr_func_rgb]
ctx.func_alpha_type = func_alpha_map[aggr_func_alpha]
ctx.texture_type = func_map_sample[texture_type]
ctx.fill_back = fill_back
face_vertices = face_vertices.clone()
textures = textures.clone()
ctx.device = face_vertices.device
ctx.batch_size, ctx.num_faces = face_vertices.shape[:2]
faces_info = torch.full((ctx.batch_size, ctx.num_faces, 9 * 3),
fill_value=0.0, device=ctx.device)
aggrs_info = torch.full((ctx.batch_size, 2, ctx.image_size, ctx.image_size),
fill_value=0.0, device=ctx.device)
soft_colors = torch.cat([torch.full((ctx.batch_size, 1, ctx.image_size, ctx.image_size),
fill_value=background_color[0], device=ctx.device),
torch.full((ctx.batch_size, 1, ctx.image_size, ctx.image_size),
fill_value=background_color[1], device=ctx.device),
torch.full((ctx.batch_size, 1, ctx.image_size, ctx.image_size),
fill_value=background_color[2], device=ctx.device),
torch.full((ctx.batch_size, 1, ctx.image_size, ctx.image_size),
fill_value=1.0, device=ctx.device)], dim=1)
faces_info, aggrs_info, soft_colors = \
soft_rasterize_cuda.forward_soft_rasterize(face_vertices, textures,
faces_info, aggrs_info,
soft_colors,
image_size, near, far, eps,
sigma_val, ctx.func_dist_type, ctx.dist_eps,
gamma_val, ctx.func_rgb_type, ctx.func_alpha_type,
ctx.texture_type, fill_back)
ctx.save_for_backward(face_vertices, textures, soft_colors, faces_info, aggrs_info)
return soft_colors # return d^2
def forward(ctx, face_vertices, textures, image_size=256,
background_color=[0, 0, 0], near=1, far=100,
fill_back=True, eps=1e-3,
sigma_val=1e-5, dist_func='euclidean', dist_eps=1e-4,
gamma_val=1e-4, aggr_func_rgb='softmax', aggr_func_alpha='prod',
texture_type='surface'):
# face_vertices: [nb, nf, 9]
# textures: [nb, nf, 9]
func_dist_map = {'hard': 0, 'barycentric': 1, 'euclidean': 2}
func_rgb_map = {'hard': 0, 'softmax': 1}
func_alpha_map = {'hard': 0, 'sum': 1, 'prod': 2}
func_map_sample = {'surface': 0, 'vertex': 1}
ctx.channel_size = textures.shape[-1]
ctx.image_size = image_size
ctx.background_color = background_color
ctx.near = near
ctx.far = far
ctx.eps = eps
ctx.sigma_val = sigma_val
ctx.gamma_val = gamma_val
ctx.func_dist_type = func_dist_map[dist_func]
ctx.dist_eps = np.log(1. / dist_eps - 1.)
ctx.func_rgb_type = func_rgb_map[aggr_func_rgb]
ctx.func_alpha_type = func_alpha_map[aggr_func_alpha]
ctx.texture_type = func_map_sample[texture_type]
ctx.fill_back = fill_back
face_vertices = face_vertices.clone()
textures = textures.clone()
ctx.device = face_vertices.device
ctx.batch_size, ctx.num_faces = face_vertices.shape[:2]
faces_info = torch.FloatTensor(ctx.batch_size, ctx.num_faces, 9*3).fill_(0.0).to(device=ctx.device) # [inv*9, sym*9, obt*3, 0*6]
aggrs_info = torch.FloatTensor(ctx.batch_size, 2, ctx.image_size, ctx.image_size).fill_(0.0).to(device=ctx.device)
soft_colors = torch.FloatTensor(ctx.batch_size, ctx.channel_size+1, ctx.image_size, ctx.image_size).fill_(1.0).to(device=ctx.device)
for k in range(min(ctx.channel_size, 3)):
soft_colors[:, k, :, :] *= background_color[k]
# soft_colors[:, 0, :, :] *= background_color[0]
# soft_colors[:, 1, :, :] *= background_color[1]
# soft_colors[:, 2, :, :] *= background_color[2]
faces_info, aggrs_info, soft_colors = \
soft_rasterize_cuda.forward_soft_rasterize(face_vertices, textures,
faces_info, aggrs_info,
soft_colors,
image_size, near, far, eps,
sigma_val, ctx.func_dist_type, ctx.dist_eps,
gamma_val, ctx.func_rgb_type, ctx.func_alpha_type,
ctx.texture_type, fill_back)
ctx.save_for_backward(face_vertices, textures, soft_colors, faces_info, aggrs_info)
return soft_colors
def execute(self, face_vertices, textures):
# face_vertices: [nb, nf, 9]
# textures: [nb, nf, 9]
func_dist_map = {'hard': 0, 'barycentric': 1, 'euclidean': 2}
func_rgb_map = {'hard': 0, 'softmax': 1}
func_alpha_map = {'hard': 0, 'sum': 1, 'prod': 2}
func_map_sample = {'surface': 0, 'vertex': 1}
image_size = self.image_size
background_color = self.background_color
near = self.near
far = self.far
eps = self.eps
sigma_val = self.sigma_val
gamma_val = self.gamma_val
dist_eps = self.dist_eps
fill_back = self.fill_back
dist_func = self.dist_func
aggr_func_rgb = self.aggr_func_rgb
aggr_func_alpha = self.aggr_func_alpha
texture_type = self.texture_type
self.func_dist_type = func_dist_map[dist_func]
self.func_rgb_type = func_rgb_map[aggr_func_rgb]
self.func_alpha_type = func_alpha_map[aggr_func_alpha]
self.texture_type = func_map_sample[texture_type]
face_vertices = face_vertices.clone()
textures = textures.clone()
self.batch_size, self.num_faces = face_vertices.shape[:2]
faces_info = jt.zeros((self.batch_size, self.num_faces,
9 * 3)) # [inv*9, sym*9, obt*3, 0*6]
aggrs_info = jt.zeros(
(self.batch_size, 2, self.image_size, self.image_size))
soft_colors = jt.ones(
(self.batch_size, 4, self.image_size, self.image_size))
soft_colors[:, 0, :, :] *= background_color[0]
soft_colors[:, 1, :, :] *= background_color[1]
soft_colors[:, 2, :, :] *= background_color[2]
# TODO: soft_colors do not init in cuda
faces_info, aggrs_info, soft_colors = \
soft_rasterize_cuda.forward_soft_rasterize(face_vertices, textures,
faces_info, aggrs_info,
soft_colors,
image_size, near, far, eps,
sigma_val, self.func_dist_type, self.dist_eps,
gamma_val, self.func_rgb_type, self.func_alpha_type,
self.texture_type, int(fill_back))
self.save_vars = face_vertices, textures, soft_colors, faces_info, aggrs_info
return soft_colors
def forward(ctx, face_vertices, textures, image_size=256,
background_color=[0, 0, 0], near=1, far=100,
fill_back=True, eps=1e-3,
sigma_val=1e-5, dist_func='euclidean', dist_eps=1e-4,
gamma_val=1e-4, aggr_func_rgb='softmax', aggr_func_alpha='prod',
texture_type='surface'):
# face_vertices: [nb, nf, 9]
# textures: [nb, nf, 9]
func_dist_map = {'hard': 0, 'barycentric': 1, 'euclidean': 2}
func_rgb_map = {'hard': 0, 'softmax': 1}
func_alpha_map = {'hard': 0, 'sum': 1, 'prod': 2}
func_map_sample = {'surface': 0, 'vertex': 1}
ctx.image_size = image_size
ctx.background_color = background_color
ctx.near = near
ctx.far = far
ctx.eps = eps
ctx.sigma_val = sigma_val
ctx.gamma_val = gamma_val
ctx.func_dist_type = func_dist_map[dist_func]
ctx.dist_eps = np.log(1. / dist_eps - 1.)
ctx.func_rgb_type = func_rgb_map[aggr_func_rgb]
ctx.func_alpha_type = func_alpha_map[aggr_func_alpha]
ctx.texture_type = func_map_sample[texture_type]
ctx.fill_back = fill_back
face_vertices = face_vertices.clone()
textures = textures.clone()
ctx.device = face_vertices.device
ctx.batch_size, ctx.num_faces = face_vertices.shape[:2]
faces_info = torch.FloatTensor(ctx.batch_size, ctx.num_faces, 9*3).fill_(0.0).to(device=ctx.device) # [inv*9, sym*9, obt*3, 0*6]
aggrs_info = torch.FloatTensor(ctx.batch_size, 2, ctx.image_size, ctx.image_size).fill_(0.0).to(device=ctx.device)
p2f_info = torch.FloatTensor(ctx.batch_size, ctx.num_faces, 2).fill_(0.0).to(device=ctx.device)
p2f_sum = torch.FloatTensor(ctx.batch_size, ctx.num_faces, 2).fill_(0.0).to(device=ctx.device)
soft_colors = torch.FloatTensor(ctx.batch_size, 4, ctx.image_size, ctx.image_size).fill_(1.0).to(device=ctx.device)
soft_colors[:, 0, :, :] *= background_color[0]
soft_colors[:, 1, :, :] *= background_color[1]
soft_colors[:, 2, :, :] *= background_color[2]
# a standard grid
theta = torch.tensor([[1, 0, 0], [0, 1, 0]], dtype=torch.float)
# grid: 1 x H x W x 2
grid = torch.nn.functional.affine_grid(theta.unsqueeze(0), (1, 1, image_size, image_size))
grid = grid.view(image_size, image_size, 2)
grid = grid.cuda()
faces_info, aggrs_info, p2f_info, p2f_sum, soft_colors = \
soft_rasterize_cuda.forward_soft_rasterize(face_vertices, textures,
faces_info, aggrs_info,
grid, p2f_info, p2f_sum, soft_colors,
image_size, near, far, eps,
sigma_val, ctx.func_dist_type, ctx.dist_eps,
gamma_val, ctx.func_rgb_type, ctx.func_alpha_type,
ctx.texture_type, fill_back)
p2f_info = p2f_info / p2f_sum.clamp_min(1e-12)
ctx.save_for_backward(face_vertices, textures, soft_colors, faces_info, aggrs_info)
return soft_colors, p2f_info, aggrs_info
def forward(ctx, face_vertices, textures, image_size=256,
background_color=[0, 0, 0], near=1, far=100,
fill_back=True, eps=1e-3,
sigma_val=1e-5, dist_func='euclidean', dist_eps=1e-4,
gamma_val=1e-4, aggr_func_rgb='softmax', aggr_func_alpha='prod',
texture_type='surface', shadow_map=None, shadow_vertices=None,
light_width=1, softmin_scale=10):
# face_vertices: [nb, nf, 9]
# textures: [nb, nf, 9]
func_dist_map = {'hard': 0, 'barycentric': 1, 'euclidean': 2}
func_rgb_map = {'hard': 0, 'softmax': 1, 'depth': 2}
func_alpha_map = {'hard': 0, 'sum': 1, 'prod': 2}
func_map_sample = {'surface': 0, 'vertex': 1}
ctx.image_size = image_size
ctx.background_color = background_color
ctx.near = near
ctx.far = far
ctx.eps = eps
ctx.sigma_val = sigma_val
ctx.gamma_val = gamma_val
ctx.func_dist_type = func_dist_map[dist_func]
ctx.dist_eps = np.log(1. / dist_eps - 1.)
ctx.func_rgb_type = func_rgb_map[aggr_func_rgb]
ctx.func_alpha_type = func_alpha_map[aggr_func_alpha]
ctx.texture_type = func_map_sample[texture_type]
ctx.fill_back = fill_back
ctx.light_width = light_width
ctx.softmin_scale = softmin_scale
face_vertices = face_vertices.clone()
textures = textures.clone()
ctx.device = face_vertices.device
ctx.batch_size, ctx.num_faces = face_vertices.shape[:2]
faces_info = torch.FloatTensor(ctx.batch_size, ctx.num_faces, 9*3).fill_(0.0).to(device=ctx.device) # [inv*9, sym*9, obt*3, 0*6]
aggrs_info = torch.FloatTensor(ctx.batch_size, 3, ctx.image_size, ctx.image_size).fill_(0.0).to(device=ctx.device)
soft_colors = torch.FloatTensor(ctx.batch_size, 4, ctx.image_size, ctx.image_size).fill_(1.0).to(device=ctx.device)
soft_colors[:, 0, :, :] *= background_color[0]
soft_colors[:, 1, :, :] *= background_color[1]
soft_colors[:, 2, :, :] *= background_color[2]
if shadow_map is None or shadow_vertices is None:
shadow_map = torch.tensor([]).to(device=ctx.device)
shadow_vertices = torch.tensor([]).to(device=ctx.device)
else:
shadow_map = shadow_map.to(device=ctx.device)
shadow_vertices = shadow_vertices.to(device=ctx.device)
faces_info, aggrs_info, soft_colors = \
soft_rasterize_cuda.forward_soft_rasterize(face_vertices, textures,
shadow_map, shadow_vertices,
faces_info, aggrs_info,
soft_colors,
image_size, near, far, eps,
sigma_val, ctx.func_dist_type, ctx.dist_eps,
gamma_val, ctx.func_rgb_type, ctx.func_alpha_type,
ctx.texture_type, fill_back,
light_width, softmin_scale)
ctx.save_for_backward(face_vertices, textures, soft_colors, faces_info, aggrs_info)
return soft_colors
