def test_transformation_consistency(scene, batch_size):
print('test_transformation_consistency')
res = render_scene(scene)
scene = make_torch_var(load_scene(scene))
pos_cc = res['pos'].reshape(-1, res['pos'].shape[-1])
normal_cc = res['normal'].reshape(-1, res['normal'].shape[-1])
surfels = cam_to_world(pos_cc, normal_cc, scene['camera'])
surfels_cc = world_to_cam(surfels['pos'], surfels['normal'],
scene['camera'])
np.testing.assert_array_almost_equal(get_data(pos_cc),
get_data(surfels_cc['pos'][:, :3]))
np.testing.assert_array_almost_equal(get_data(normal_cc),
get_data(surfels_cc['normal'][:, :3]))
def test_depth_to_world_consistency(scene, batch_size):
res = render_scene(scene)
scene = make_torch_var(load_scene(scene))
pos_wc1 = res['pos'].reshape(-1, res['pos'].shape[-1])
pos_cc1 = world_to_cam(pos_wc1, None, scene['camera'])['pos']
# First test the non-batched z_to_pcl_CC method:
# NOTE: z_to_pcl_CC takes as input the Z dimension in the camera coordinate
# and gets the full (X, Y, Z) in the camera coordinate.
pos_cc2 = z_to_pcl_CC(pos_cc1[:, 2], scene['camera'])
pos_wc2 = cam_to_world(pos_cc2, None, scene['camera'])['pos']
# Test Z -> (X, Y, Z)
np.testing.assert_array_almost_equal(get_data(pos_cc1[..., :3]),
get_data(pos_cc2[..., :3]))
# Test world -> camera -> Z -> (X, Y, Z) in camera -> world
np.testing.assert_array_almost_equal(get_data(pos_wc1[..., :3]),
get_data(pos_wc2[..., :3]))
# Then test the batched version:
camera = scene['camera']
camera['eye'] = camera['eye'].repeat(batch_size, 1)
camera['at'] = camera['at'].repeat(batch_size, 1)
camera['up'] = camera['up'].repeat(batch_size, 1)
pos_wc1 = pos_wc1.repeat(batch_size, 1, 1)
pos_cc1 = world_to_cam_batched(pos_wc1, None, scene['camera'])['pos']
pos_cc2 = z_to_pcl_CC_batched(pos_cc1[..., 2], camera) # NOTE: z = -depth
pos_wc2 = cam_to_world_batched(pos_cc2, None, camera)['pos']
# Test Z -> (X, Y, Z)
np.testing.assert_array_almost_equal(get_data(pos_cc1[..., :3]),
get_data(pos_cc2[..., :3]))
# Test world -> camera -> Z -> (X, Y, Z) in camera -> world
np.testing.assert_array_almost_equal(get_data(pos_wc1[..., :3]),
get_data(pos_wc2[..., :3]))
def test_sphere_splat_render_along_ray(out_dir, cam_pos, width, height, fovy, focal_length, use_quartic,
b_display=False):
"""
Create a sphere on a square as in render_sphere_world, and then convert to the camera's coordinate system
and then render using render_splats_along_ray.
"""
import copy
print('render sphere along ray')
sampling_time = []
rendering_time = []
num_samples = width * height
large_scene = copy.deepcopy(SCENE_TEST)
large_scene['camera']['viewport'] = [0, 0, width, height]
large_scene['camera']['eye'] = tch_var_f(cam_pos)
large_scene['camera']['fovy'] = np.deg2rad(fovy)
large_scene['camera']['focal_length'] = focal_length
large_scene['objects']['disk']['material_idx'] = tch_var_l(np.zeros(num_samples, dtype=int).tolist())
large_scene['materials']['albedo'] = tch_var_f([[0.6, 0.6, 0.6]])
large_scene['tonemap']['gamma'] = tch_var_f([1.0]) # Linear output
x, y = np.meshgrid(np.linspace(-1, 1, width), np.linspace(-1, 1, height))
#z = np.sqrt(1 - np.min(np.stack((x ** 2 + y ** 2, np.ones_like(x)), axis=-1), axis=-1))
unit_disk_mask = (x ** 2 + y ** 2) <= 1
z = np.sqrt(1 - unit_disk_mask * (x ** 2 + y ** 2))
# Make a hemi-sphere bulging out of the xy-plane scene
z[~unit_disk_mask] = 0
pos = np.stack((x.ravel(), y.ravel(), z.ravel() - 5, np.ones(num_samples)), axis=1)
# Normals outside the sphere should be [0, 0, 1]
x[~unit_disk_mask] = 0
y[~unit_disk_mask] = 0
z[~unit_disk_mask] = 1
normals = np_normalize(np.stack((x.ravel(), y.ravel(), z.ravel(), np.zeros(num_samples)), axis=1))
if b_display:
plt.ion()
plt.figure()
plt.subplot(131)
plt.imshow(pos[..., 0].reshape((height, width)))
plt.subplot(132)
plt.imshow(pos[..., 1].reshape((height, width)))
plt.subplot(133)
plt.imshow(pos[..., 2].reshape((height, width)))
plt.figure()
plt.imshow(normals[..., 2].reshape((height, width)))
## Convert to the camera's coordinate system
#Mcam = lookat(eye=large_scene['camera']['eye'], at=large_scene['camera']['at'], up=large_scene['camera']['up'])
pos_CC = tch_var_f(pos) #torch.matmul(tch_var_f(pos), Mcam.transpose(1, 0))
large_scene['objects']['disk']['pos'] = pos_CC
large_scene['objects']['disk']['normal'] = None # Estimate the normals tch_var_f(normals)
# large_scene['camera']['eye'] = tch_var_f([-10., 0., 10.])
# large_scene['camera']['eye'] = tch_var_f([2., 0., 10.])
large_scene['camera']['eye'] = tch_var_f([-5., 0., 0.])
# main render run
start_time = time()
res = render_splats_along_ray(large_scene, use_quartic=use_quartic)
rendering_time.append(time() - start_time)
# Test cam_to_world
res_world = cam_to_world(res['pos'].reshape(-1, 3), res['normal'].reshape(-1, 3), large_scene['camera'])
im = get_data(res['image'])
im = np.uint8(255. * im)
depth = get_data(res['depth'])
depth[depth >= large_scene['camera']['far']] = large_scene['camera']['far']
if b_display:
plt.figure()
plt.imshow(im, interpolation='none')
plt.title('Image')
plt.savefig(out_dir + '/fig_img_orig.png')
plt.figure()
plt.imshow(depth, interpolation='none')
plt.title('Depth Image')
#plt.savefig(out_dir + '/fig_depth_orig.png')
plt.figure()
pos_world = get_data(res_world['pos'])
posx_world = pos_world[:, 0].reshape((im.shape[0], im.shape[1]))
posy_world = pos_world[:, 1].reshape((im.shape[0], im.shape[1]))
posz_world = pos_world[:, 2].reshape((im.shape[0], im.shape[1]))
plt.subplot(131)
plt.imshow(posx_world)
plt.title('x_world')
plt.subplot(132)
plt.imshow(posy_world)
plt.title('y_world')
plt.subplot(133)
plt.imshow(posz_world)
plt.title('z_world')
fig = plt.figure()
ax = fig.add_subplot(111, projection='3d')
ax.scatter(pos_world[:, 0], pos_world[:, 1], pos_world[:, 2], s=1.3)
ax.set_xlabel('x')
ax.set_ylabel('y')
plt.figure()
pos_world = get_data(res['pos'].reshape(-1, 3))
posx_world = pos_world[:, 0].reshape((im.shape[0], im.shape[1]))
posy_world = pos_world[:, 1].reshape((im.shape[0], im.shape[1]))
posz_world = pos_world[:, 2].reshape((im.shape[0], im.shape[1]))
plt.subplot(131)
plt.imshow(posx_world)
plt.title('x_CC')
plt.subplot(132)
plt.imshow(posy_world)
plt.title('y_CC')
plt.subplot(133)
plt.imshow(posz_world)
plt.title('z_CC')
imsave(out_dir + '/img_orig.png', im)
#imsave(out_dir + '/depth_orig.png', im_depth)
# hold matplotlib figure
plt.ioff()
plt.show()
im = get_data(res['image'])
depth = get_data(res['depth'])
plt.figure()
plt.imshow(im)
plt.figure()
plt.imshow(depth)
pos = get_data(res['pos'])
fig = plt.figure()
ax = fig.add_subplot(111, projection='3d')
ax.scatter(pos[:, 0], pos[:, 1], pos[:, 2], s=1.3)
res_world = cam_to_world(pos=res['pos'], normal=res['normal'], camera=scene[idx]['camera'])
pos = get_data(res_world['pos'])
fig = plt.figure()
ax = fig.add_subplot(111, projection='3d')
ax.scatter(pos[:, 0], pos[:, 1], pos[:, 2], s=1.3)
plt.ioff()
plt.show()
# data = np.load('res_world_twogans.npy')
#
# # pos0 = get_data(data[0]['pos'])
# # fig = plt.figure()
# # ax = fig.add_subplot(111, projection='3d')
# # ax.scatter(pos0[:, 0], pos0[:, 1], pos0[:, 2], s=1.3)
#
# fig = plt.figure()
def render_batch(self, batch, batch_cond,light_pos):
"""Render a batch of splats."""
batch_size = batch.size()[0]
# Generate camera positions on a sphere
if batch_cond is None:
if self.opt.full_sphere_sampling:
cam_pos = uniform_sample_sphere(
radius=self.opt.cam_dist, num_samples=self.opt.batchSize,
axis=self.opt.axis, angle=np.deg2rad(self.opt.angle),
theta_range=self.opt.theta, phi_range=self.opt.phi)
# TODO: deg2grad!!
else:
cam_pos = uniform_sample_sphere(
radius=self.opt.cam_dist, num_samples=self.opt.batchSize,
axis=self.opt.axis, angle=self.opt.angle,
theta_range=np.deg2rad(self.opt.theta),
phi_range=np.deg2rad(self.opt.phi))
# TODO: deg2grad!!
rendered_data = []
rendered_data_depth = []
rendered_data_cond = []
scenes = []
inpath = self.opt.vis_images + '/'
inpath_xyz = self.opt.vis_xyz + '/'
z_min = self.scene['camera']['focal_length']
z_max = z_min + 3
# TODO (fmannan): Move this in init. This only needs to be done once!
# Set splats into rendering scene
if 'sphere' in self.scene['objects']:
del self.scene['objects']['sphere']
if 'triangle' in self.scene['objects']:
del self.scene['objects']['triangle']
if 'disk' not in self.scene['objects']:
self.scene['objects'] = {'disk': {'pos': None, 'normal': None,
'material_idx': None}}
lookat = self.opt.at if self.opt.at is not None else [0.0, 0.0, 0.0, 1.0]
self.scene['camera']['at'] = tch_var_f(lookat)
self.scene['objects']['disk']['material_idx'] = tch_var_l(
np.zeros(self.opt.splats_img_size * self.opt.splats_img_size))
loss = 0.0
loss_ = 0.0
z_loss_ = 0.0
z_norm_loss_ = 0.0
spatial_loss_ = 0.0
spatial_var_loss_ = 0.0
unit_normal_loss_ = 0.0
normal_away_from_cam_loss_ = 0.0
image_depth_consistency_loss_ = 0.0
for idx in range(batch_size):
# Get splats positions and normals
eps = 1e-3
if self.opt.rescaled:
z = F.relu(-batch[idx][:, 0]) + z_min
z = ((z - z.min()) / (z.max() - z.min() + eps) *
(z_max - z_min) + z_min)
pos = -z
else:
z = F.relu(-batch[idx][:, 0]) + z_min
pos = -F.relu(-batch[idx][:, 0]) - z_min
normals = batch[idx][:, 1:]
self.scene['objects']['disk']['pos'] = pos
# Normal estimation network and est_normals don't go together
self.scene['objects']['disk']['normal'] = normals if self.opt.est_normals is False else None
# Set camera position
if batch_cond is None:
if not self.opt.same_view:
self.scene['camera']['eye'] = tch_var_f(cam_pos[idx])
else:
self.scene['camera']['eye'] = tch_var_f(cam_pos[0])
else:
if not self.opt.same_view:
self.scene['camera']['eye'] = batch_cond[idx]
else:
self.scene['camera']['eye'] = batch_cond[0]
self.scene['lights']['pos'][0,:3]=tch_var_f(light_pos[idx])
#self.scene['lights']['pos'][1,:3]=tch_var_f(self.light_pos2[idx])
# Render scene
# res = render_splats_NDC(self.scene)
res = render_splats_along_ray(self.scene,
samples=self.opt.pixel_samples,
normal_estimation_method='plane')
world_tform = cam_to_world(res['pos'].view((-1, 3)),
res['normal'].view((-1, 3)),
self.scene['camera'])
# Get rendered output
res_pos = res['pos'].contiguous()
res_pos_2D = res_pos.view(res['image'].shape)
# The z_loss needs to be applied after supersampling
# TODO: Enable this (currently final loss becomes NaN!!)
# loss += torch.mean(
# (10 * F.relu(z_min - torch.abs(res_pos[..., 2]))) ** 2 +
# (10 * F.relu(torch.abs(res_pos[..., 2]) - z_max)) ** 2)
res_normal = res['normal']
# depth_grad_loss = spatial_3x3(res['depth'][..., np.newaxis])
# grad_img = grad_spatial2d(torch.mean(res['image'], dim=-1)[..., np.newaxis])
# grad_depth_img = grad_spatial2d(res['depth'][..., np.newaxis])
image_depth_consistency_loss = depth_rgb_gradient_consistency(
res['image'], res['depth'])
unit_normal_loss = unit_norm2_L2loss(res_normal, 10.0) # TODO: MN
normal_away_from_cam_loss = away_from_camera_penalty(
res_pos, res_normal)
z_pos = res_pos[..., 2]
z_loss = torch.mean((2 * F.relu(z_min - torch.abs(z_pos))) ** 2 +
(2 * F.relu(torch.abs(z_pos) - z_max)) ** 2)
z_norm_loss = normal_consistency_cost(
res_pos, res['normal'], norm=1)
spatial_loss = spatial_3x3(res_pos_2D)
spatial_var = torch.mean(res_pos[..., 0].var() +
res_pos[..., 1].var() +
res_pos[..., 2].var())
spatial_var_loss = (1 / (spatial_var + 1e-4))
loss = (self.opt.zloss * z_loss +
self.opt.unit_normalloss*unit_normal_loss +
self.opt.normal_consistency_loss_weight * z_norm_loss +
self.opt.spatial_var_loss_weight * spatial_var_loss +
self.opt.grad_img_depth_loss*image_depth_consistency_loss +
self.opt.spatial_loss_weight * spatial_loss)
pos_out_ = get_data(res['pos'])
loss_ += get_data(loss)
z_loss_ += get_data(z_loss)
z_norm_loss_ += get_data(z_norm_loss)
spatial_loss_ += get_data(spatial_loss)
spatial_var_loss_ += get_data(spatial_var_loss)
unit_normal_loss_ += get_data(unit_normal_loss)
normal_away_from_cam_loss_ += get_data(normal_away_from_cam_loss)
image_depth_consistency_loss_ += get_data(
image_depth_consistency_loss)
normals_ = get_data(res_normal)
if self.opt.render_img_nc == 1:
depth = res['depth']
im = depth.unsqueeze(0)
else:
depth = res['depth']
im_d = depth.unsqueeze(0)
im = res['image'].permute(2, 0, 1)
H, W = im.shape[1:]
target_normal_ = get_data(res['normal']).reshape((H, W, 3))
target_normalmap_img_ = get_normalmap_image(target_normal_)
target_worldnormal_ = get_data(world_tform['normal']).reshape(
(H, W, 3))
target_worldnormalmap_img_ = get_normalmap_image(
target_worldnormal_)
if self.iterationa_no % (self.opt.save_image_interval*5) == 0:
imsave((inpath + str(self.iterationa_no) +
'normalmap_{:05d}.png'.format(idx)),
target_normalmap_img_)
imsave((inpath + str(self.iterationa_no) +
'depthmap_{:05d}.png'.format(idx)),
get_data(res['depth']))
imsave((inpath + str(self.iterationa_no) +
'world_normalmap_{:05d}.png'.format(idx)),
target_worldnormalmap_img_)
if self.iterationa_no % 1000 == 0:
im2 = get_data(res['image'])
depth2 = get_data(res['depth'])
pos = get_data(res['pos'])
out_file2 = ("pos"+".npy")
np.save(inpath_xyz+out_file2, pos)
out_file2 = ("im"+".npy")
np.save(inpath_xyz+out_file2, im2)
out_file2 = ("depth"+".npy")
np.save(inpath_xyz+out_file2, depth2)
# Save xyz file
save_xyz((inpath_xyz + str(self.iterationa_no) +
'withnormal_{:05d}.xyz'.format(idx)),
pos=get_data(res['pos']),
normal=get_data(res['normal']))
# Save xyz file in world coordinates
save_xyz((inpath_xyz + str(self.iterationa_no) +
'withnormal_world_{:05d}.xyz'.format(idx)),
pos=get_data(world_tform['pos']),
normal=get_data(world_tform['normal']))
if self.opt.gz_gi_loss is not None and self.opt.gz_gi_loss > 0:
gradZ = grad_spatial2d(res_pos_2D[:, :, 2][:, :, np.newaxis])
gradImg = grad_spatial2d(torch.mean(im,
dim=0)[:, :, np.newaxis])
for (gZ, gI) in zip(gradZ, gradImg):
loss += (self.opt.gz_gi_loss * torch.mean(torch.abs(
torch.abs(gZ) - torch.abs(gI))))
# Store normalized depth into the data
rendered_data.append(im)
rendered_data_depth.append(im_d)
rendered_data_cond.append(self.scene['camera']['eye'])
scenes.append(self.scene)
rendered_data = torch.stack(rendered_data)
rendered_data_depth = torch.stack(rendered_data_depth)
return rendered_data, rendered_data_depth, loss/self.opt.batchSize
def optimize_splats_along_ray_shadow_with_normalest_test(
out_dir,
width,
height,
max_iter=100,
lr=1e-3,
scale=10,
shadow=True,
vis_only=False,
samples=1,
est_normals=False,
b_generate_normals=False,
print_interval=10,
imsave_interval=10,
xyz_save_interval=100):
"""A demo function to check if the differentiable renderer can optimize splats rendered along ray.
:param scene:
:param out_dir:
:return:
"""
import torch
import copy
from diffrend.torch.params import SCENE_SPHERE_HALFBOX_0
if not os.path.exists(out_dir):
os.mkdir(out_dir)
scene = SCENE_SPHERE_HALFBOX_0
scene['camera']['viewport'] = [0, 0, width, height]
scene['camera']['fovy'] = np.deg2rad(45)
scene['camera']['focal_length'] = 1
scene['camera']['eye'] = tch_var_f(
[2, 1, 2, 1]) # tch_var_f([1, 1, 1, 1]) # tch_var_f([2, 2, 2, 1]) #
scene['camera']['at'] = tch_var_f(
[0, 0.8, 0, 1]) # tch_var_f([0, 1, 0, 1]) # tch_var_f([2, 2, 0, 1]) #
scene['lights']['attenuation'] = tch_var_f([
[0., 0.0, 0.01],
[0., 0.0, 0.01],
[0., 0.0, 0.01],
])
scene['materials']['coeffs'] = tch_var_f([
[1.0, 0.0, 0.0],
[1.0, 0.0, 0.0],
[1.0, 0.0, 0.0],
[0.5, 0.2, 8.0],
[1.0, 0.0, 0.0],
[1.0, 0.0, 0.0],
])
target_res = render(scene, tiled=True, shadow=shadow)
target_im = normalize_maxmin(target_res['image'])
target_im.require_grad = False
target_im_ = get_data(target_im)
target_pos_ = get_data(target_res['pos'])
target_normal_ = get_data(target_res['normal'])
target_normalmap_img_ = get_normalmap_image(target_normal_)
target_depth_ = get_data(target_res['depth'])
print('[z_min, z_max] = [%f, %f]' %
(np.min(target_pos_[..., 2]), np.max(target_pos_[..., 2])))
print('[depth_min, depth_max] = [%f, %f]' %
(np.min(target_depth_), np.max(target_depth_)))
# world -> cam -> render_splats_along_ray
cc_tform = world_to_cam(target_res['pos'].view(
(-1, 3)), target_res['normal'].view((-1, 3)), scene['camera'])
wc_cc_tform = cam_to_world(cc_tform['pos'], cc_tform['normal'],
scene['camera'])
# Check normal estimation in camera space
pos_cc = cc_tform['pos'][:, :3].contiguous().view(target_im.shape)
normal_cc = cc_tform['normal'][:, :3].contiguous().view(target_im.shape)
plane_fit_est = estimate_surface_normals_plane_fit(pos_cc, None)
normal_cc_normalmap = get_normalmap_image(get_data(normal_cc))
plane_fit_est_normalmap = get_normalmap_image(get_data(plane_fit_est))
pos_diff = torch.abs(wc_cc_tform['pos'][:, :3] -
target_res['pos'].view((-1, 3)))
mean_pos_diff = torch.mean(pos_diff)
normal_diff = torch.abs(wc_cc_tform['normal'][:, :3] -
target_res['normal'].view(-1, 3))
mean_normal_diff = torch.mean(normal_diff)
print('mean_pos_diff', mean_pos_diff, 'mean_normal_diff', mean_normal_diff)
wc_cc_normal = wc_cc_tform['normal'].view(target_im_.shape)
wc_cc_normal_img = get_normalmap_image(get_data(wc_cc_normal))
material_idx = tch_var_l(np.ones(cc_tform['pos'].shape[0]) * 3)
input_scene = copy.deepcopy(scene)
del input_scene['objects']['sphere']
del input_scene['objects']['triangle']
light_vis = tch_var_f(
np.ones(
(input_scene['lights']['pos'].shape[0], cc_tform['pos'].shape[0])))
input_scene['objects'] = {
'disk': {
'pos': cc_tform['pos'],
'normal': cc_tform['normal'],
'material_idx': material_idx,
'light_vis': light_vis,
}
}
target_res_noshadow = render(scene, tiled=True, shadow=False)
res = render_splats_along_ray(input_scene)
test_img_ = get_data(normalize_maxmin(res['image']))
test_depth_ = get_data(res['depth'])
test_normal_ = get_data(res['normal']).reshape(test_img_.shape)
test_normalmap_ = get_normalmap_image(test_normal_)
im_diff = np.abs(test_img_ -
get_data(normalize_maxmin(target_res_noshadow['image'])))
print('mean image diff: {}'.format(np.mean(im_diff)))
#### PLOT
plt.ion()
plt.figure()
plt.imshow(test_img_, interpolation='none')
plt.title('Test Image')
plt.savefig(out_dir + '/test_img.png')
plt.figure()
plt.imshow(test_depth_, interpolation='none')
plt.title('Test Depth')
plt.savefig(out_dir + '/test_depth.png')
plt.figure()
plt.imshow(test_normalmap_, interpolation='none')
plt.title('Test Normals')
plt.savefig(out_dir + '/test_normal.png')
####
criterion = nn.L1Loss() #nn.MSELoss()
criterion = criterion.cuda()
plt.ion()
plt.figure()
plt.imshow(target_im_, interpolation='none')
plt.title('Target Image')
plt.savefig(out_dir + '/target.png')
plt.figure()
plt.imshow(target_normalmap_img_, interpolation='none')
plt.title('Normals')
plt.savefig(out_dir + '/normal.png')
plt.figure()
plt.imshow(wc_cc_normal_img, interpolation='none')
plt.title('WC_CC Normals')
plt.savefig(out_dir + '/wc_cc_normal.png')
plt.figure()
plt.imshow(normal_cc_normalmap, interpolation='none')
plt.title('Normal CC GT')
plt.savefig(out_dir + '/normal_cc.png')
plt.figure()
plt.imshow(plane_fit_est_normalmap, interpolation='none')
plt.title('Plane fit CC')
plt.savefig(out_dir + '/est_normal_cc.png')
plt.figure()
plt.subplot(121)
plt.imshow(normal_cc_normalmap, interpolation='none')
plt.title('Normal CC GT')
plt.subplot(122)
plt.imshow(plane_fit_est_normalmap, interpolation='none')
plt.title('Plane fit CC')
plt.savefig(out_dir + '/normal_and_estnormal_cc_comparison.png')
input_scene = copy.deepcopy(scene)
del input_scene['objects']['sphere']
del input_scene['objects']['triangle']
input_scene['camera']['viewport'] = [
0, 0, int(width / samples),
int(height / samples)
]
num_splats = int(width * height / (samples * samples))
#x, y = np.meshgrid(np.linspace(-1, 1, int(width / samples)), np.linspace(-1, 1, int(height / samples)))
z_min = scene['camera']['focal_length']
z_max = 3
z = -tch_var_f(
np.ones(num_splats) * (z_min + z_max) / 2
) # -torch.clamp(tch_var_f(2 * np.random.rand(num_splats)), z_min, z_max)
z.requires_grad = True
normal_angles = tch_var_f(np.random.rand(num_splats, 2))
normal_angles.requires_grad = True
material_idx = tch_var_l(np.ones(num_splats) * 3)
light_vis = tch_var_f(
np.ones((input_scene['lights']['pos'].shape[0], num_splats)))
light_vis.requires_grad = True
if vis_only:
assert shadow is True
opt_vars = [light_vis]
z = cc_tform['pos'][:, 2]
# FIXME: sph2cart
#normals = cc_tform['normal']
else:
opt_vars = [z, normal_angles]
if shadow:
opt_vars += [light_vis]
optimizer = optim.Adam(opt_vars, lr=lr)
lr_scheduler = StepLR(optimizer, step_size=10000, gamma=0.8)
h0 = plt.figure()
h1 = plt.figure()
h2 = plt.figure()
h3 = plt.figure()
h4 = plt.figure()
gs1 = gridspec.GridSpec(3, 3)
gs1.update(wspace=0.0025, hspace=0.02)
# Two options for z_norm_consistency
# 1. start after N iterations
# 2. start at the beginning and decay
# 3. start after N iterations and decay to 0
no_decay = lambda x: x
exp_decay = lambda x, scale: torch.exp(-x / scale)
linear_decay = lambda x, scale: scale / (x + 1e-6)
spatial_var_loss_weight = 10.0 #0.0
normal_away_from_cam_loss_weight = 0.0
grad_img_depth_loss_weight = 1.0
spatial_loss_weight = 2
z_norm_weight_init = 1 # 1e-5
z_norm_activate_iter = 0 # 1000
decay_fn = lambda x: linear_decay(x, 100)
loss_per_iter = []
if b_generate_normals:
est_normals = False
normal_est_network = NEstNetAffine(kernel_size=3, sph=False)
print(normal_est_network)
normal_est_network.cuda()
for iter in range(max_iter):
lr_scheduler.step()
zz = -F.relu(-z) - z_min # torch.clamp(z, -z_max, -z_min)
if b_generate_normals:
normals = generate_normals(zz, scene['camera'], normal_est_network)
#if iter > 100 and iter % 10 == 0:
# print(normals)
elif not est_normals:
phi = F.sigmoid(normal_angles[:, 0]) * 2 * np.pi
theta = F.sigmoid(
normal_angles[:, 1]
) * np.pi / 2 # F.tanh(normal_angles[:, 1]) * np.pi / 2
normals = sph2cart_unit(torch.stack((phi, theta), dim=1))
pos = zz # torch.stack((tch_var_f(x.ravel()), tch_var_f(y.ravel()), zz), dim=1)
input_scene['objects'] = {
'disk': {
'pos': pos,
'normal': normalize(normals) if not est_normals else None,
'material_idx': material_idx,
'light_vis': torch.sigmoid(light_vis),
}
}
res = render_splats_along_ray(input_scene,
samples=samples,
normal_estimation_method='plane')
res_pos = res['pos']
res_normal = res['normal']
spatial_loss = spatial_3x3(res_pos)
depth_grad_loss = spatial_3x3(res['depth'][..., np.newaxis])
grad_img = grad_spatial2d(
torch.mean(res['image'], dim=-1)[..., np.newaxis])
grad_depth_img = grad_spatial2d(res['depth'][..., np.newaxis])
image_depth_consistency_loss = depth_rgb_gradient_consistency(
res['image'], res['depth'])
unit_normal_loss = unit_norm2_L2loss(res_normal, 10.0)
normal_away_from_cam_loss = away_from_camera_penalty(
res_pos, res_normal)
z_pos = res_pos[..., 2]
z_loss = torch.mean((10 * F.relu(z_min - torch.abs(z_pos)))**2 +
(10 * F.relu(torch.abs(z_pos) - z_max))**2)
z_norm_loss = normal_consistency_cost(res_pos, res_normal, norm=1)
spatial_var = torch.mean(res_pos[..., 0].var() +
res_pos[..., 1].var() + res_pos[..., 2].var())
spatial_var_loss = (1 / (spatial_var + 1e-4))
im_out = normalize_maxmin(res['image'])
res_depth_ = get_data(res['depth'])
optimizer.zero_grad()
z_norm_weight = z_norm_weight_init * float(
iter > z_norm_activate_iter) * decay_fn(iter -
z_norm_activate_iter)
loss = criterion(scale * im_out, scale * target_im) + z_loss + unit_normal_loss + \
z_norm_weight * z_norm_loss + \
spatial_var_loss_weight * spatial_var_loss + \
grad_img_depth_loss_weight * image_depth_consistency_loss
#normal_away_from_cam_loss_weight * normal_away_from_cam_loss + \
#spatial_loss_weight * spatial_loss
im_out_ = get_data(im_out)
im_out_normal_ = get_data(res['normal'])
pos_out_ = get_data(res['pos'])
loss_ = get_data(loss)
z_loss_ = get_data(z_loss)
z_norm_loss_ = get_data(z_norm_loss)
spatial_loss_ = get_data(spatial_loss)
spatial_var_loss_ = get_data(spatial_var_loss)
unit_normal_loss_ = get_data(unit_normal_loss)
normal_away_from_cam_loss_ = get_data(normal_away_from_cam_loss)
normals_ = get_data(res_normal)
image_depth_consistency_loss_ = get_data(image_depth_consistency_loss)
loss_per_iter.append(loss_)
if iter == 0:
plt.figure(h0.number)
plt.imshow(im_out_)
plt.title('Initial')
if iter % print_interval == 0 or iter == max_iter - 1:
z_ = get_data(z)
z__ = pos_out_[..., 2]
print(
'%d. loss= %f nloss=%f z_loss=%f [%f, %f] [%f, %f], z_normal_loss: %f,'
' spatial_var_loss: %f, normal_away_loss: %f'
' nz_range: [%f, %f], spatial_loss: %f, imd_loss: %f' %
(iter, loss_, unit_normal_loss_, z_loss_, np.min(z_),
np.max(z_), np.min(z__), np.max(z__), z_norm_loss_,
spatial_var_loss_, normal_away_from_cam_loss_,
normals_[..., 2].min(), normals_[..., 2].max(), spatial_loss_,
image_depth_consistency_loss_))
if iter % xyz_save_interval == 0 or iter == max_iter - 1:
save_xyz(out_dir + '/res_{:05d}.xyz'.format(iter),
get_data(res_pos), get_data(res_normal))
if iter % imsave_interval == 0 or iter == max_iter - 1:
z_ = get_data(z)
plt.figure(h4.number)
plt.clf()
plt.suptitle('%d. loss= %f [%f, %f]' %
(iter, loss_, np.min(z_), np.max(z_)))
plt.subplot(121)
#plt.axis('off')
plt.imshow(im_out_, interpolation='none')
plt.title('Output')
plt.subplot(122)
#plt.axis('off')
plt.imshow(target_im_, interpolation='none')
plt.title('Ground truth')
# plt.subplot(223)
# plt.plot(loss_per_iter, linewidth=2)
# plt.xlabel('Iteration', fontsize=14)
# plt.title('Loss', fontsize=12)
# plt.grid(True)
plt.savefig(out_dir + '/fig_im_gt_loss_%05d.png' % iter)
plt.figure(h1.number, figsize=(4, 4))
plt.clf()
plt.suptitle('%d. loss= %f [%f, %f]' %
(iter, loss_, np.min(z_), np.max(z_)))
plt.subplot(gs1[0])
plt.axis('off')
plt.imshow(im_out_, interpolation='none')
plt.subplot(gs1[1])
plt.axis('off')
plt.imshow(get_normalmap_image(im_out_normal_),
interpolation='none')
ax = plt.subplot(gs1[2])
plt.axis('off')
im_tmp = ax.imshow(res_depth_, interpolation='none')
# create an axes on the right side of ax. The width of cax will be 5%
# of ax and the padding between cax and ax will be fixed at 0.05 inch.
divider = make_axes_locatable(ax)
cax = divider.append_axes("right", size="5%", pad=0.05)
plt.colorbar(im_tmp, cax=cax)
plt.subplot(gs1[3])
plt.axis('off')
plt.imshow(target_im_, interpolation='none')
plt.subplot(gs1[4])
plt.axis('off')
plt.imshow(test_normalmap_, interpolation='none')
ax = plt.subplot(gs1[5])
plt.axis('off')
im_tmp = ax.imshow(test_depth_, interpolation='none')
divider = make_axes_locatable(ax)
cax = divider.append_axes("right", size="5%", pad=0.05)
plt.colorbar(im_tmp, cax=cax)
W, H = input_scene['camera']['viewport'][2:]
light_vis_ = get_data(torch.sigmoid(light_vis))
plt.subplot(gs1[6])
plt.axis('off')
plt.imshow(light_vis_[0].reshape((H, W)), interpolation='none')
if (light_vis_.shape[0] > 1):
plt.subplot(gs1[7])
plt.axis('off')
plt.imshow(light_vis_[1].reshape((H, W)), interpolation='none')
if (light_vis_.shape[0] > 2):
plt.subplot(gs1[8])
plt.axis('off')
plt.imshow(light_vis_[2].reshape((H, W)), interpolation='none')
plt.savefig(out_dir + '/fig_%05d.png' % iter)
plt.figure(h2.number)
plt.clf()
plt.imshow(res_depth_)
plt.colorbar()
plt.savefig(out_dir + '/fig_depth_%05d.png' % iter)
plt.figure(h3.number)
plt.clf()
plt.imshow(z_.reshape(H, W))
plt.colorbar()
plt.savefig(out_dir + '/fig_z_%05d.png' % iter)
loss.backward()
optimizer.step()
plt.figure()
plt.plot(loss_per_iter, linewidth=2)
plt.xlabel('Iteration', fontsize=14)
plt.title('Loss', fontsize=12)
plt.grid(True)
plt.savefig(out_dir + '/loss.png')
plt.ioff()
plt.show()