Update app.py

app.py

@@ -9,7 +9,7 @@ import torch.nn.functional as F
 from torchvision import transforms
 from torchvision.transforms import Compose
 import trimesh
-#from geometry import create_triangles
+from geometry import create_triangles
 import tempfile
 from functools import partial
 import spaces
@@ -239,6 +239,7 @@ def pano_depth_to_world_points(depth):
 
     pts3d = [[0,0,0]]
     uv = [[0,0]]
+    nl = [[0,0,0]]
     for i in range(0, 1): #(0,2)
         for j in range(0, 1): #(0,2)
             #rnd_lon = (np.random.rand(depth.shape[0]*depth.shape[1]) - 0.5) / 8
@@ -251,16 +252,22 @@ def pano_depth_to_world_points(depth):
             y = radius * np.cos(d_lat)
             z = radius * np.sin(d_lon) * np.sin(d_lat)
 
+            nx = -np.cos(d_lon) * np.sin(d_lat)
+            ny = -np.cos(d_lat)
+            nz = -np.sin(d_lon) * np.sin(d_lat)
+
             pts = np.stack([x, y, z], axis=1)
             uvs = np.stack([lon, lat], axis=1)
+            nls = np.stack([nx, ny, nz], axis=1)
             
             pts3d = np.concatenate((pts3d, pts), axis=0)
             uv = np.concatenate((uv, uvs), axis=0)
+            nl = np.concatenate((nl, nls), axis=0)
             #print(f'i: {i}, j: {j}')
             j = j+1
         i = i+1
         
-    return [pts3d, uv]
+    return [pts3d, uv, nl]
 
 def rgb2gray(rgb):
     return np.dot(rgb[...,:3], [0.333, 0.333, 0.333])
@@ -287,6 +294,7 @@ def get_mesh(image, depth, blur_data, loadall):
     points = pano_depth_to_world_points(gdepth)
     pts3d = points[0]
     uv = points[1]
+    nl = points[2]
     print('radius from depth - ok')
 
     # Create a trimesh mesh from the points
@@ -295,6 +303,7 @@ def get_mesh(image, depth, blur_data, loadall):
     uvs = uv.reshape(-1, 2)
     print(uvs)
     verts = pts3d.reshape(-1, 3)
+    normals = nl.reshape(-1, 3)
     rgba = cv2.cvtColor(blur_img, cv2.COLOR_RGB2RGBA)
     colors = rgba.reshape(-1, 4)
     clrs = [[128,128,128,0]]
@@ -312,18 +321,18 @@ def get_mesh(image, depth, blur_data, loadall):
             #if (clrs[j] == pcolors[i]).all():
                 #cverts[i].append(verts[j])
                 #j=j+1
-    mesh.append(trimesh.PointCloud(verts, colors=None))
-    mesh[len(mesh)-1].merge_vertices()
+    #mesh.append(trimesh.PointCloud(verts, colors=None))
+    #mesh[len(mesh)-1].merge_vertices()
         #i=i+1
     mesh_n.append(str(fnum))
 
-    #triangles = create_triangles(rgba.shape[0], rgba.shape[1])
-    #mesh.append(trimesh.Trimesh(vertices=verts, faces=triangles))
+    triangles = create_triangles(rgba.shape[0], rgba.shape[1])
+    mesh.append(trimesh.Trimesh(vertices=verts, faces=triangles, vertex_normals=normals))
     rgba_pil = Image.fromarray(rgba.astype(np.uint8))
     material = trimesh.visual.texture.SimpleMaterial(image=rgba_pil)
     visuals = trimesh.visual.TextureVisuals(uv=uvs, image=rgba_pil, material=material)
     mesh[len(mesh)-1].visual = visuals
-    #print('triangles - ok')
+    print('triangles - ok')
     
     #ln = []
     #cr = []