import os 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_LAUNCH_BLOCKING"] = '1' import torch from torch.nn import Module import numpy as np from superSMPL.smpl_torch import SMPL import trimesh import plotly.express as ex import plotly.graph_objects as go from data_utils import intrinsic, extrinsic, get_camera from pytorch3d.renderer import ( RasterizationSettings, MeshRenderer, MeshRasterizer, SoftPhongShader, TexturesVertex, HardPhongShader, AmbientLights ) from pytorch3d.structures import Meshes class PoseRenderer(Module): def __init__(self, device='cuda') -> None: super().__init__() colors = np.load('custom_data/pose_colors.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() def forward(self, verts: torch.Tensor, faces: torch.Tensor): cameras, K = get_camera(self.cam, [512, 512], verts.get_device(), res_x=512, res_y=512) raster_settings = RasterizationSettings( image_size=(512, 512), blur_radius=0.0, faces_per_pixel=1, bin_size=0 ) renderer = MeshRenderer( rasterizer=MeshRasterizer(cameras=cameras, raster_settings=raster_settings), shader=HardPhongShader(device=verts.device, cameras=cameras, lights=AmbientLights(device=verts.device)) ) output = renderer(Meshes(verts[None,:,:], faces[None,:,:], textures=self.textures)) return output def get_dijsktra(m: trimesh.Trimesh) -> np.ndarray: edge_lengths = m.edges_unique_length # create and adjacency matrix based on the edge_lengths adjacency_matrix = np.zeros((m.vertices.shape[0], m.vertices.shape[0]), dtype=np.float32) adjacency_matrix[m.edges_unique[:,0], m.edges_unique[:,1]] = edge_lengths adjacency_matrix[m.edges_unique[:,1], m.edges_unique[:,0]] = edge_lengths # calculate the distance between all vertices using dijsktra algorithm from scipy.sparse import csr_matrix from scipy.sparse.csgraph import dijkstra dists = csr_matrix(adjacency_matrix) dists = dijkstra(dists, directed=False) return dists def get_distances(vertices, edges): from scipy.sparse import csr_matrix from scipy.sparse.csgraph import dijkstra num_vertices = vertices.shape[0] edge_lengths = np.linalg.norm(vertices[edges[:, 0]] - vertices[edges[:,1]], axis=-1) 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) return distances def get_distances2(vertices, faces): from scipy.sparse import coo_matrix, csgraph num_vertices = vertices.shape[0] num_faces = faces.shape[0] # Calculate the edge lengths edge_lengths = np.linalg.norm(vertices[faces[:, np.newaxis]] - vertices[faces], axis=2) # Create the adjacency matrix as a COO matrix row_indices = np.repeat(faces.flatten(), 3) col_indices = np.tile(faces.flatten(), 3) data = np.tile(edge_lengths.flatten(), 3) adjacency_matrix = coo_matrix((data, (row_indices, col_indices)), shape=(num_vertices, num_vertices)) # Compute the distances between all vertices using Dijkstra's algorithm distances = csgraph.dijkstra(adjacency_matrix, directed=False) return distances def interpolate_colors(distance_matrix, colored_indices, colored_colors): from scipy.spatial import distance 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 num_vertices = distance_matrix.shape[0] # 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 # Interpolate colors for the remaining vertices for vertex in range(num_vertices): if vertex not in colored_indices: # Find the distances from the current vertex to the colored vertices distances = distance_matrix[vertex, colored_indices] # 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, colored_colors).astype(np.uint8) # 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))] interpolated_colors[vertex] = color return interpolated_colors if __name__ == '__main__': smpl = SMPL(path='superSMPL', device='cuda') pose = torch.zeros((1,72), device='cuda') # generate random pose # pose[:,3:] = torch.randn((1,69), device='cuda') * 0.1 body, _ = smpl(pose, torch.zeros((1,11), device='cuda'), remove=False) colors = np.zeros((body.shape[1], 3)) renderer = PoseRenderer() body_to_render = body[0] # swap the y,z axes of the body_to_render body_to_render = body_to_render[:,[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='cuda'), body_to_render.T).T # out = renderer(body_to_render, torch.from_numpy(smpl.f_all).to(body.get_device())) # ex.imshow(out[0,:,:,:3].cpu().numpy()).show() body = body.cpu().numpy()[0] faces = smpl.f_all # m = trimesh.Trimesh(vertices=body, faces=faces, process=False, validate=False) # # export the mesh to an obj file # m.export('custom_data/body.obj') # imoprt the mesh from the obj file m = trimesh.load('custom_data/body_color_init.ply', process=False, validate=False) color_init = np.array(m.visual.vertex_colors)[:,:3] colored_vertices = {} for idx, color in enumerate(color_init): if np.all(color == 255): continue colored_vertices[idx] = color # dist = get_distances(body, m.edges) # # dist = compute_distances(body, faces) # np.save('custom_data/dijsktra.npy', dist) 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) # get the distances between all vertices using dijsktra algorithm without get_dijsktra # from scipy.sparse import csr_matrix # from scipy.sparse.csgraph import dijkstra # define a gradinet color cube for the body that goes from green to red based on the coordinates of the vertices and not their indices # colors = np.zeros((body.shape[0], 3)) # 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.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)) # same as a scatter plot # fig = go.Figure(data=[go.Scatter3d(x=body[:,0], y=body[:,2], z=body[:,1], mode='markers', marker=dict(size=2, color=colors, opacity=0.8))]) # fig.update_layout(scene=dict(xaxis=dict(range=[-1,1]), yaxis=dict(range=[-1,1]), zaxis=dict(range=[-1,1]))) # fig.show() pass