body coloring trials

.gitignore

@@ -164,4 +164,6 @@ cython_debug/
 temp/
 dataset/
 wandb_login.py
-custom_data/dijsktra.npy
+custom_data/dijsktra*
+
+*.out

data/pose_renderer.py

@@ -30,7 +30,9 @@ from pytorch3d.structures import Meshes
 class PoseRenderer(Module):
     def __init__(self, device='cuda') -> None:
         super().__init__()
-        colors = np.load('custom_data/pose_colors.npy')
+        # colors = np.load('custom_data/pose_colors.npy')
+        colors = np.load('custom_data/body_colors_interpolated_unweighted.npy')
+        
         self.textures = TexturesVertex(verts_features=torch.from_numpy(colors).to(device)[None, :, :])
         self.cam = torch.tensor([3.0, 0.0, 0.0], device=device).cpu().numpy()
 
@@ -53,3 +55,29 @@ class PoseRenderer(Module):
         output = renderer(Meshes(verts[None,:,:], faces[None,:,:], textures=self.textures))
 
         return output
+    
+if __name__ == '__main__':
+    os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
+    os.environ["CUDA_VISIBLE_DEVICES"] = '0'
+    from superSMPL.smpl_torch import SMPL
+    import plotly.express as ex
+    renderer = PoseRenderer()
+    smpl = SMPL(path='superSMPL', device='cuda')
+    pose = torch.rand(72, dtype=torch.float32, device='cuda')
+    pose[:3] = 0
+    shape = torch.rand(11, dtype=torch.float32, device='cuda')
+    # pose_out[1] = 0.5 # np.random.uniform(-np.pi, np.pi)
+    # generate random number between -np.pi, np.pi
+    
+    # start = time()
+
+    body = smpl(pose.unsqueeze(0), shape.unsqueeze(0), remove=False)
+    body_to_render = body[0,:,[0,2,1]]
+    # and rotate it 90 degress along the z axis
+    body_to_render = torch.matmul(torch.tensor([[0,1,0],[-1,0,0],[0,0,1]], dtype=torch.float32, device=body_to_render.device), body_to_render.T).T
+
+    out = renderer(body_to_render, torch.from_numpy(smpl.f_all).to(body_to_render.get_device()))
+
+    ex.imshow(out[0, ..., :3].cpu().numpy()).show()
+
+    print(out)

utils/pose_coloring.py

@@ -3,7 +3,7 @@ import sys
 sys.path.append(os.path.dirname(os.path.abspath(__file__)) + '/..')
 
 os.environ["CUDA_DEVICE_ORDER"]="PCI_BUS_ID"
-os.environ["CUDA_VISIBLE_DEVICES"] = '3'
+os.environ["CUDA_VISIBLE_DEVICES"] = '0'
 # os.environ["CUDA_LAUNCH_BLOCKING"] = '1'
 
 import torch
@@ -11,7 +11,7 @@ from torch.nn import Module
 import numpy as np
 from superSMPL.smpl_torch import SMPL
 import trimesh
-
+import time
 import plotly.express as ex
 import plotly.graph_objects as go
 
@@ -78,11 +78,13 @@ def get_distances(vertices, edges):
     from scipy.sparse.csgraph import dijkstra
     num_vertices = vertices.shape[0]
 
-    edge_lengths = np.linalg.norm(vertices[edges[:, 0]] - vertices[edges[:,1]], axis=-1)
+    # edge_lengths = np.linalg.norm(vertices[edges[:, 0]] - vertices[edges[:,1]], axis=-1)
+    edge_lengths = np.ones(edges.shape[0])
     adjacency_matrix = csr_matrix((edge_lengths.flatten(), (edges[:, 0], edges[:, 1])), shape=(num_vertices, num_vertices))
 
     # Compute the distances between all vertices using Dijkstra's algorithm
-    distances = dijkstra(adjacency_matrix, directed=False)
+    # distances = dijkstra(adjacency_matrix, directed=False)
+    distances = dijkstra(adjacency_matrix, directed=False, unweighted=True)
 
     return distances
 
@@ -108,14 +110,147 @@ def get_distances2(vertices, faces):
     return distances
 
 
+import networkx as nx
+from scipy.spatial.distance import cdist
+ 
+# def order_mesh_vertices(vertices, distances, edges):
+#     # Create a graph with vertices as nodes and distances as edge weights
+#     graph = nx.Graph()
+#     graph.add_nodes_from(range(len(vertices)))
+#     # calcualate the edge lengths
+#     edge_lengths = np.array([np.linalg.norm(vertices[edges[i, 0]] - vertices[edges[i, 1]]) for i in range(edges.shape[0])])
+
+#     # Add edges to the graph based on the mesh edges
+#     for edge, edge_length in zip(edges, edge_lengths):
+#         graph.add_edge(edge[0], edge[1], weight=edge_length)
+ 
+#     ordered_vertices = []
+#     ordered_indices = []
+ 
+#     start_vertex = 0
+#     # Perform breadth-first search starting from the selected vertex
+#     visited = set()
+#     queue = [(1221, 1221)]  # Tuple of (vertex, index)
+
+#     while queue:
+#         vertex, index = queue.pop(0)
+#         if vertex not in visited:
+#             visited.add(vertex)
+#             ordered_vertices.append(vertices[vertex])
+#             ordered_indices.append(vertex)
+#             # neighbors = sorted(graph.neighbors(vertex), key=lambda n: graph.edges[vertex, n]['weight'])
+#             neighbors = graph.neighbors(vertex)
+#             queue.extend([(n, n) for n in neighbors])
+ 
+#     return np.array(ordered_vertices), np.array(ordered_indices)
+
+from collections import deque
+def order_mesh_vertices(vertices, distances, edges):
+    # Create a graph with vertices as nodes and distances as edge weights
+    graph = nx.Graph()
+    graph.add_nodes_from(range(len(vertices)))
+    
+    # Calculate the edge lengths
+    edge_lengths = np.linalg.norm(vertices[edges[:, 0]] - vertices[edges[:, 1]], axis=1)
+    
+    # Add edges to the graph based on the mesh edges
+    for edge, edge_length in zip(edges, edge_lengths):
+        graph.add_edge(edge[0], edge[1], weight=edge_length)
+ 
+    ordered_indices = []
+ 
+    start_vertex = 13833 
+    # Perform breadth-first search starting from the selected vertex
+    visited = set()
+    queue = deque([start_vertex])
+
+    while queue:
+        vertex = queue.popleft()
+        if vertex not in visited:
+            visited.add(vertex)
+            ordered_indices.append(vertex)
+            neighbors = sorted(graph.neighbors(vertex), key=lambda n: graph.edges[vertex, n]['weight'])
+            # neighbors = graph.neighbors(vertex)
+            queue.extend(neighbors)
+ 
+    return np.array(ordered_indices)
+
+
 def interpolate_colors(distance_matrix, colored_indices, colored_colors):
-    from scipy.spatial import distance
+    from scipy.spatial import distance, cKDTree
     all_colors = np.array(np.meshgrid(np.arange(0,256), np.arange(0,256), np.arange(0,256))).T.reshape(-1,3)
     color_map = dict()
     for idx, color in enumerate(all_colors):
         color_map[tuple(color)] = color
+    color_map_array = np.array(list(color_map.values()))
 
     num_vertices = distance_matrix.shape[0]
+    kd_tree = cKDTree(color_map_array)
+
+    # Create an array to store the interpolated colors
+    interpolated_colors = np.zeros((num_vertices, 3), dtype=np.uint8)
+
+    # Assign the given colors to the corresponding vertices
+    for index, color in zip(colored_indices, colored_colors):
+        interpolated_colors[index] = color
+
+    used_colors = set([tuple(color) for color in colored_colors])
+
+    # Interpolate colors for the remaining vertices
+    for vertex in range(num_vertices):
+        print(vertex)
+        if vertex not in colored_indices:
+            # start = time.time()
+            # Find the distances from the current vertex to the colored vertices
+            distances = distance_matrix[vertex, colored_indices]
+            # take the three closest one only
+            closest_indices = np.argsort(distances)[:3]
+            distances = distances[closest_indices]
+            # print(time.time() - start)
+
+            # Calculate the weights for color interpolation based on inverse distances
+            weights = 1.0 / distances
+            # Normalize the weights to sum up to 1
+            weights /= np.sum(weights)
+            # Interpolate the colors based on the weights
+            color = np.dot(weights, np.array(colored_colors)[closest_indices]).astype(np.uint8)
+
+            # get the closest color to this from color_map if this color is not available
+            if tuple(color) in used_colors:
+                start = time.time()
+                color_idx = kd_tree.query(color, workers=48)[1]
+                color = color_map_array[color_idx]
+                # if color is already in used_colors then get the next closest color until find one that is not in used_colors
+                k = 2
+                while tuple(color) in used_colors:
+                    neighbour_dist, color_idx = kd_tree.query(color, k=k, workers=48)
+                    color = color_map_array[color_idx[-1]]
+                    k += 1
+                # assert neighbour_dist[-1] < 5
+                # color = color_map[min(color_map.keys(), key=lambda x: distance.euclidean(x, color))]
+                print(time.time() - start)
+
+            # start = time.time()
+            interpolated_colors[vertex] = color
+            used_colors.add(tuple(color))
+            # color_map.pop(tuple(color))
+            # color_map_array = np.delete(color_map_array, np.where(np.all(color_map_array == color, axis=1))[0], axis=0)
+            # np.delete(color_map_array, np.where , axis=0)
+            # print(time.time() - start)
+
+    return interpolated_colors
+
+
+def interpolate_colors_2(distance_matrix, colored_indices, colored_colors):
+    from scipy.spatial import distance, cKDTree
+    all_colors = np.array(np.meshgrid(np.arange(0,256), np.arange(0,256), np.arange(0,256))).T.reshape(-1,3)
+    color_map = dict()
+    for idx, color in enumerate(all_colors):
+        color_map[tuple(color)] = color
+    color_map_array = np.array(list(color_map.values()))
+
+    num_vertices = distance_matrix.shape[0]
+    kd_tree = cKDTree(color_map_array)
 
     # Create an array to store the interpolated colors
     interpolated_colors = np.zeros((num_vertices, 3), dtype=np.uint8)
@@ -126,6 +261,7 @@ def interpolate_colors(distance_matrix, colored_indices, colored_colors):
 
     # Interpolate colors for the remaining vertices
     for vertex in range(num_vertices):
+        print(vertex)
         if vertex not in colored_indices:
             # Find the distances from the current vertex to the colored vertices
             distances = distance_matrix[vertex, colored_indices]
@@ -141,14 +277,83 @@ def interpolate_colors(distance_matrix, colored_indices, colored_colors):
 
             # get the closest color to this from color_map if this color is not available
             if tuple(color) not in color_map:
-                color = color_map[min(color_map.keys(), key=lambda x: distance.euclidean(x, color))]
+                kd_tree = cKDTree(color_map_array)
+                color_idx = kd_tree.query(color, workers=48)[1]
+                color = color_map_array[color_idx]
+                # color = color_map[min(color_map.keys(), key=lambda x: distance.euclidean(x, color))]
 
             interpolated_colors[vertex] = color
+            color_map.pop(tuple(color))
+            color_map_array = np.delete(color_map_array, np.where(np.all(color_map_array == color, axis=1))[0], axis=0)
+            # np.delete(color_map_array, np.where , axis=0)
 
     return interpolated_colors
 
 
+def get_edges(faces):
+    edges = []
+    unique_edges = set()
 
+    # Iterate over each face
+    for face in faces:
+        # Get the vertices of the current face
+        v0, v1, v2 = face
+
+        # Define the three edges of the current face
+        # and add them to the unique_edges set
+        edge0 = tuple(sorted((v0, v1)))
+        edge1 = tuple(sorted((v1, v2)))
+        edge2 = tuple(sorted((v2, v0)))
+
+        unique_edges.add(edge0)
+        unique_edges.add(edge1)
+        unique_edges.add(edge2)
+
+    # Convert the unique_edges set to a list
+    edges = list(unique_edges)
+
+    # Convert the list of edges to a numpy array
+    edges = np.array(edges)
+
+    return edges
+
+
+def get_data():
+    from plyfile import PlyData
+    # Read the .ply file
+    plydata = PlyData.read('custom_data/body_color_init.ply')
+
+    # Get the vertex coordinates
+    vertices = plydata['vertex']
+    coordinates = np.vstack((vertices['x'], vertices['y'], vertices['z'])).T
+
+    # Get the vertex colors
+    colors = np.vstack((vertices['red'], vertices['green'], vertices['blue'])).T
+
+    return coordinates, colors
+
+
+
+def find_corresponding_vertices(mesh1_vertices, mesh2_vertices, indices_mesh1):
+    corresponding_vertices_mesh2 = []
+    # mesh1_vertices[:, [1, 2]] = mesh1_vertices[:, [2, 1]]
+    for index_mesh1 in indices_mesh1:
+        vertex_mesh1 = mesh1_vertices[index_mesh1]
+        for i, vertex_mesh2 in enumerate(mesh2_vertices):
+            if np.allclose(vertex_mesh1, vertex_mesh2):
+                corresponding_vertices_mesh2.append(i)
+                break
+    return np.array(corresponding_vertices_mesh2)
+
+def transform_selected_indices(selected_indices, new_order):
+    mapping = {old_idx: new_idx for new_idx, old_idx in enumerate(new_order)}
+    transformed_indices = np.array([mapping[idx] for idx in selected_indices])
+    return transformed_indices
+
+def transform_faces(faces, new_order):
+    mapping = {old_idx: new_idx for new_idx, old_idx in enumerate(new_order)}
+    transformed_faces = np.array([[mapping[idx] for idx in face] for face in faces])
+    return transformed_faces
 
 if __name__ == '__main__':
 
@@ -185,21 +390,69 @@ if __name__ == '__main__':
             continue
         colored_vertices[idx] = color
 
-    # dist = get_distances(body, m.edges)
-    # # dist = compute_distances(body, faces)
-    # np.save('custom_data/dijsktra.npy', dist)
+    # create a 3d plotly scatter plot fo body vertices with vertex 26157 highlited and limits between -1 and 1
+    # fig = go.Figure(data=[go.Scatter3d(x=body[:,0], y=body[:,1], z=body[:,2], mode='markers',  marker=dict(color='blue', size=1))])
+    # fig.add_trace(go.Scatter3d(x=[body[26157,0]], y=[body[26157,1]], z=[body[26157,2]], mode='markers', marker=dict(color='red', size=5)))
+    # fig.update_layout(scene=dict(xaxis=dict(range=[-1,1]), yaxis=dict(range=[-1,1]), zaxis=dict(range=[-1,1])))
+    # fig.show()
 
-    dist = np.load('custom_data/dijsktra.npy')
-    colors = interpolate_colors(dist, list(colored_vertices.keys()), list(colored_vertices.values()))
 
-    weights = dist[0,:]
-    # define a gradient rgb colors for body based on weights that is not just gray
-    colors = np.zeros((body.shape[0], 3))
-    colors[:,0] = weights
-    colors[:,1] = weights
-    colors[:,2] = weights
-    colors = colors * 255
-    colors = colors.astype(np.uint8)
+
+
+    # dist = get_distances(body, get_edges(faces))
+    # # # dist = compute_distances(body, faces)
+    # np.save('custom_data/dijsktra_int.npy', dist)
+
+    body_ply, colors_ply = get_data()
+    # rotate body_ply -90 degress along x axis
+    body_ply = body_ply[:,[0,2,1]]
+    # rotate body_ply 180 degress along y axis
+    body_ply[:,2] = -body_ply[:,2]
+
+    # colored_vertices.pop(13833)
+    # colored_vertices[1221] = np.array([0, 128, 0]).astype(np.uint8)
+    c = find_corresponding_vertices(body_ply, body, list(colored_vertices.keys()))
+    # c = np.delete(c, np.where(c == 180))
+
+    dist = np.load('custom_data/dijsktra2.npy')
+    ordered_indices = order_mesh_vertices(body, dist, get_edges(faces))
+    # reverse ordered_indices
+    # ordered_indices = ordered_indices[::-1]
+
+    c = transform_selected_indices(c, ordered_indices)
+    body = body[ordered_indices]
+    dist = dist[ordered_indices]
+    dist = dist[:,ordered_indices]
+    faces = transform_faces(faces, ordered_indices)
+    colors = interpolate_colors(dist, list(c), list(colored_vertices.values()))
+    np.save('custom_data/body_colors_interpolated_unweighted.npy', colors)
+ 
+    #c = [14110, 3192, 25113, 10822, 27513, 25872]
+    # generate a gradient grayscale colors for the vertice based ont heir indices
+    # colors = np.zeros((body.shape[0], 3))
+    # colors[:,0] = np.arange(body.shape[0]) / body.shape[0]
+    # colors[:,1] = np.arange(body.shape[0]) / body.shape[0]
+    # colors[:,2] = np.arange(body.shape[0]) / body.shape[0]
+    # colors = colors * 255
+    # colors = colors.astype(np.uint8)
+
+
+    # weights = dist[0,:]
+    # # define a gradient rgb colors for body based on weights that is not just gray
+    # colors = np.zeros((body.shape[0], 3))
+    # colors[:,0] = weights
+    # colors[:,1] = weights
+    # colors[:,2] = weights
+    # colors = colors * 255
+    # colors = colors.astype(np.uint8)
+
+    # # define a gradient rgb colors for body based on their distance from vertex 1221
+    # colors = np.zeros((body.shape[0], 3))
+    # colors[:,0] = dist[1221,:]
+    # colors[:,1] = dist[1221,:]
+    # colors[:,2] = dist[1221,:]
+    # norm = colors / np.max(colors)
+    # colors = colors * 255
 
     # get the distances between all vertices using dijsktra algorithm without get_dijsktra
     # from scipy.sparse import csr_matrix
@@ -212,7 +465,7 @@ if __name__ == '__main__':
     # colors[:,0] = (body[:,0] - np.min(body[:,0])) / (np.max(body[:,0]) - np.min(body[:,0]))
     # colors[:,1] = (body[:,1] - np.min(body[:,1])) / (np.max(body[:,1]) - np.min(body[:,1]))
     # colors[:,2] = (body[:,2] - np.min(body[:,2])) / (np.max(body[:,2]) - np.min(body[:,2]))
-    fig = go.Figure(data=[go.Mesh3d(x=body[:,0], y=body[:,2], z=body[:,1], i=faces[:,0], j=faces[:,2], k=faces[:,1], vertexcolor=colors)])
+    fig = go.Figure(data=[go.Mesh3d(x=body[:,0], y=body[:,2], z=body[:,1], i=faces[:,0], j=faces[:,1], k=faces[:,2], vertexcolor=colors)])
     fig.update_layout(scene=dict(xaxis=dict(range=[-1,1]), yaxis=dict(range=[-1,1]), zaxis=dict(range=[-1,1])))
     fig.show()  
     # np.save('custom_data/pose_colors.npy', colors.astype(np.float32))
@@ -222,6 +475,10 @@ if __name__ == '__main__':
     # fig.update_layout(scene=dict(xaxis=dict(range=[-1,1]), yaxis=dict(range=[-1,1]), zaxis=dict(range=[-1,1])))
     # fig.show()
 
+    # save body with colors into a ply file
+    m = trimesh.Trimesh(vertices=body, faces=faces, vertex_colors=colors)
+    m.export('custom_data/body_colors.ply')
+
 
     
     pass