diff --git a/.gitignore b/.gitignore index 0ceb7b2..fa7a81b 100644 --- a/.gitignore +++ b/.gitignore @@ -164,4 +164,6 @@ cython_debug/ temp/ dataset/ wandb_login.py -custom_data/dijsktra.npy +custom_data/dijsktra* + +*.out diff --git a/custom_data/body2.ply b/custom_data/body2.ply new file mode 100644 index 0000000..ac5b721 Binary files /dev/null and b/custom_data/body2.ply differ diff --git a/custom_data/body3.ply b/custom_data/body3.ply new file mode 100644 index 0000000..6214641 Binary files /dev/null and b/custom_data/body3.ply differ diff --git a/custom_data/body_colors.ply b/custom_data/body_colors.ply new file mode 100644 index 0000000..e9f98e5 Binary files /dev/null and b/custom_data/body_colors.ply differ diff --git a/custom_data/body_colors_interpolated.npy b/custom_data/body_colors_interpolated.npy new file mode 100644 index 0000000..110cd07 Binary files /dev/null and b/custom_data/body_colors_interpolated.npy differ diff --git a/custom_data/body_colors_interpolated_unweighted.npy b/custom_data/body_colors_interpolated_unweighted.npy new file mode 100644 index 0000000..4f74948 Binary files /dev/null and b/custom_data/body_colors_interpolated_unweighted.npy differ diff --git a/data/pose_renderer.py b/data/pose_renderer.py index 0fd054b..2472c41 100644 --- a/data/pose_renderer.py +++ b/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() @@ -52,4 +54,30 @@ class PoseRenderer(Module): output = renderer(Meshes(verts[None,:,:], faces[None,:,:], textures=self.textures)) - return output \ No newline at end of file + 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) \ No newline at end of file diff --git a/utils/pose_coloring.py b/utils/pose_coloring.py index 6925cbc..eb8abf2 100644 --- a/utils/pose_coloring.py +++ b/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 \ No newline at end of file