general save
This commit is contained in:
+132
-6
@@ -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"] = '0'
|
||||
os.environ["CUDA_VISIBLE_DEVICES"] = '3'
|
||||
# os.environ["CUDA_LAUNCH_BLOCKING"] = '1'
|
||||
|
||||
import torch
|
||||
@@ -58,6 +58,98 @@ class PoseRenderer(Module):
|
||||
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')
|
||||
@@ -75,20 +167,54 @@ if __name__ == '__main__':
|
||||
# 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()
|
||||
# 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[:,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()
|
||||
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
|
||||
|
||||
Reference in New Issue
Block a user