Fixed Entropy regularization error

Author: yashbhalgat

camera_pose_visualizer.py

@@ -0,0 +1,73 @@
+import json
+import os
+import numpy as np
+import matplotlib as mpl
+import matplotlib.pyplot as plt
+from matplotlib.patches import Patch
+from mpl_toolkits.mplot3d.art3d import Poly3DCollection
+import pdb
+
+class CameraPoseVisualizer:
+    def __init__(self, xlim, ylim, zlim):
+        self.fig = plt.figure(figsize=(18, 7))
+        self.ax = self.fig.gca(projection='3d')
+        self.ax.set_aspect("auto")
+        self.ax.set_xlim(xlim)
+        self.ax.set_ylim(ylim)
+        self.ax.set_zlim(zlim)
+        self.ax.set_xlabel('x')
+        self.ax.set_ylabel('y')
+        self.ax.set_zlabel('z')
+        print('initialize camera pose visualizer')
+
+    def extrinsic2pyramid(self, extrinsic, color='r', focal_len_scaled=5, aspect_ratio=0.3):
+        focal_len_scaled = -1*focal_len_scaled
+        vertex_std = np.array([[0, 0, 0, 1],
+                               [focal_len_scaled * aspect_ratio, -focal_len_scaled * aspect_ratio, focal_len_scaled, 1],
+                               [focal_len_scaled * aspect_ratio, focal_len_scaled * aspect_ratio, focal_len_scaled, 1],
+                               [-focal_len_scaled * aspect_ratio, focal_len_scaled * aspect_ratio, focal_len_scaled, 1],
+                               [-focal_len_scaled * aspect_ratio, -focal_len_scaled * aspect_ratio, focal_len_scaled, 1]])
+        vertex_transformed = vertex_std @ extrinsic.T
+        meshes = [[vertex_transformed[0, :-1], vertex_transformed[1][:-1], vertex_transformed[2, :-1]],
+                            [vertex_transformed[0, :-1], vertex_transformed[2, :-1], vertex_transformed[3, :-1]],
+                            [vertex_transformed[0, :-1], vertex_transformed[3, :-1], vertex_transformed[4, :-1]],
+                            [vertex_transformed[0, :-1], vertex_transformed[4, :-1], vertex_transformed[1, :-1]],
+                            [vertex_transformed[1, :-1], vertex_transformed[2, :-1], vertex_transformed[3, :-1], vertex_transformed[4, :-1]]]
+        self.ax.add_collection3d(
+            Poly3DCollection(meshes, facecolors=color, linewidths=0.3, edgecolors=color, alpha=0.35))
+
+    def customize_legend(self, list_label):
+        list_handle = []
+        for idx, label in enumerate(list_label):
+            color = plt.cm.rainbow(idx / len(list_label))
+            patch = Patch(color=color, label=label)
+            list_handle.append(patch)
+        plt.legend(loc='right', bbox_to_anchor=(1.8, 0.5), handles=list_handle)
+
+    def colorbar(self, max_frame_length):
+        cmap = mpl.cm.rainbow
+        norm = mpl.colors.Normalize(vmin=0, vmax=max_frame_length)
+        self.fig.colorbar(mpl.cm.ScalarMappable(norm=norm, cmap=cmap), orientation='vertical', label='Frame Number')
+
+    def show(self):
+        plt.title('Extrinsic Parameters')
+        plt.show()
+
+if __name__ == '__main__':
+    poses = []
+    with open(os.path.join('data/nerf_synthetic/chair/', 'transforms_train.json'), 'r') as fp:
+        meta = json.load(fp)
+        for frame in meta['frames']:
+            poses.append(np.array(frame['transform_matrix']))
+    t_arr = np.array([pose[:3,-1] for pose in poses])
+    maxes = t_arr.max(axis=0)
+    mins = t_arr.min(axis=0)
+    
+    # argument : the minimum/maximum value of x, y, z
+    visualizer = CameraPoseVisualizer([mins[0]-1, maxes[0]+1], [mins[1]-1, maxes[1]+1], [mins[2]-1, maxes[2]+1])
+
+    # argument : extrinsic matrix, color, scaled focal length(z-axis length of frame body of camera
+    for pose in poses:
+        visualizer.extrinsic2pyramid(pose, 'c', 1)
+
+    visualizer.show() 

run_nerf.py

@@ -360,9 +360,11 @@ def raw2outputs(raw, z_vals, rays_d, raw_noise_std=0, white_bkgd=False, pytest=F
         rgb_map = rgb_map + (1.-acc_map[...,None])
 
     # Calculate weights sparsity loss
-    mask = weights.sum(-1) > 0.5
-    entropy = Categorical(probs = weights+1e-5).entropy()
-    sparsity_loss = entropy * mask
+    try:
+        entropy = Categorical(probs = torch.cat([weights, 1.0-weights.sum(-1, keepdim=True)+1e-6], dim=-1)).entropy()
+    except:
+        pdb.set_trace()
+    sparsity_loss = entropy
 
     return rgb_map, disp_map, acc_map, weights, depth_map, sparsity_loss
 
@@ -645,7 +647,6 @@ def train():
     elif args.dataset_type == 'blender':
         images, poses, render_poses, hwf, i_split, bounding_box = load_blender_data(args.datadir, args.half_res, args.testskip)
         args.bounding_box = bounding_box
-        pdb.set_trace()
         print('Loaded blender', images.shape, render_poses.shape, hwf, args.datadir)
         i_train, i_val, i_test = i_split