484 lines
19 KiB
Python
484 lines
19 KiB
Python
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"] = '0'
|
|
# 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 time
|
|
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)
|
|
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, unweighted=True)
|
|
|
|
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
|
|
|
|
|
|
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, 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)
|
|
|
|
# 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):
|
|
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]
|
|
|
|
# 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:
|
|
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__':
|
|
|
|
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
|
|
|
|
# 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 = 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
|
|
# 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()
|
|
# 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()
|
|
|
|
# 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 |