227 lines
8.4 KiB
Python
227 lines
8.4 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"] = '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 |