fix sfm depth bug; TODO: numba dependency

handcrafted_solution.py

@@ -164,7 +164,20 @@ def get_smooth_uv_depth(vertices, depth, gest_seg_np, sfm_depth_np):
     vertex_depth = np.array(vertex_depth)
     return uv, vertex_depth
 
-# TODO: timeit
+from numba import njit, prange
+@njit(parallel=True)
+def fill_range(u, v, z, dilate_r, c, sfm_depth_np, sfm_color_np, H, W):
+  for i in prange(max(0, u - dilate_r), min(W, u + dilate_r)):
+      for j in prange(max(0, v - dilate_r), min(H, v + dilate_r)) :
+        #checked+=1
+        existing_z = sfm_depth_np[j, i]
+        if z > 0:
+          if (existing_z!=0 and z < existing_z) or (existing_z==0):
+            sfm_depth_np[j, i] = z
+            if DUMP_IMG: 
+              sfm_color_np[j, i] = c
+  return sfm_depth_np, sfm_color_np
+
 def get_SfM_depth(points3D, depth_np, gest_seg_np, K, R, t, dilate_r = 5):
   '''Project 3D sfm pointcloud to the image plane '''
   H, W = depth_np.shape[:2]
@@ -180,16 +193,31 @@ def get_SfM_depth(points3D, depth_np, gest_seg_np, K, R, t, dilate_r = 5):
   world_to_cam = K @ Rt
   xyz = world_to_cam @ XYZ1.transpose()
   xyz = np.transpose(xyz)  
-  us, vs, zs = xyz[0]/xyz[2], xyz[1]/xyz[2], xyz[2]
+  #print('shape of xyz: ', xyz.shape)
+  us, vs, zs = xyz[:,0]/xyz[:,2], xyz[:,1]/xyz[:,2], xyz[:,2]
+  #print('dim of us uv zs rgb:', len(us), len(vs), len(zs), len(rgb))
   us = us.astype(np.int32)    
   vs = vs.astype(np.int32)
+  #checked = 0  
+  #print('dim of us uv zs rgb:', len(us), len(vs), len(zs), len(rgb))
   for u,v,z,c in zip(us,vs,zs, rgb):
-    for i in range(max(0, u - dilate_r), min(W, u + dilate_r)):
-      for j in range(max(0, v - dilate_r), min(H, v + dilate_r)):        
-        if z > 0: 
-          if (sfm_depth_np[j, i]!=0 and z < sfm_depth_np[j, i]) or (sfm_depth_np[j, i]==0):
+    
+    sfm_depth_np, sfm_color_np = fill_range(u, v, z, dilate_r, c, sfm_depth_np, sfm_color_np, H, W)
+    '''
+    i_range = range(max(0, u - dilate_r), min(W, u + dilate_r))
+    j_range = range(max(0, v - dilate_r), min(H, v + dilate_r)) 
+    for i in i_range:
+      for j in j_range:
+        #checked+=1
+        existing_z = sfm_depth_np[j, i]
+        if z > 0:
+          if (existing_z!=0 and z < existing_z) or (existing_z==0):
             sfm_depth_np[j, i] = z
-            sfm_color_np[j, i] = c
+            if DUMP_IMG: 
+              sfm_color_np[j, i] = c
+    ''' 
+              
+  #print(f'checked {checked} pts')
 
   if DUMP_IMG:
     filename_sfm_depth = 'sfm_depth.png'
@@ -546,7 +574,7 @@ def predict(entry, visualize=False, prune_dist_thr=600, depth_scale=2.5, ) -> Tu
         elif i==2: # only visualize view 0,1
           continue 
         '''
-        
+
         ade_seg = ade.resize(depth.size)
         ade_seg_np = np.array(ade_seg).astype(np.uint8)
         gest_seg = gest.resize(depth.size)