reaarrange files
This commit is contained in:
Binary file not shown.
@@ -0,0 +1,596 @@
|
||||
import os
|
||||
import sys
|
||||
import shutil
|
||||
import tqdm
|
||||
import random
|
||||
import math
|
||||
import collections
|
||||
import tps
|
||||
import numpy as np
|
||||
import scipy as sp
|
||||
import trimesh as tm
|
||||
import cmapy as cm
|
||||
import scipy.io as sio
|
||||
import matplotlib.pyplot as plt
|
||||
import plotly.express as ex
|
||||
|
||||
|
||||
|
||||
def m2uv_array(m2uv, is_smpl_orig=False):
|
||||
result = np.zeros((14475 if not is_smpl_orig else 3638, 2), dtype=np.int32)
|
||||
for k, v in m2uv.items():
|
||||
if len(v) == 2:
|
||||
result[k] = list(v)
|
||||
elif len(v) == 1:
|
||||
x = v.pop()
|
||||
result[k] = [x, x]
|
||||
else:
|
||||
assert False
|
||||
|
||||
return result
|
||||
|
||||
|
||||
def get_body():
|
||||
smpl = SMPL('FeatureLines/superSMPL')
|
||||
pose = np.zeros((1, 72), dtype=np.float32)
|
||||
beta = np.zeros((1, 11), dtype=np.float32)
|
||||
R = smpl.run_rodrigues(pose)
|
||||
T, _, _ = smpl.run(R, beta, True)
|
||||
return np.squeeze(T)
|
||||
|
||||
|
||||
def plot_3d(x, y=None):
|
||||
fig = plt.figure(figsize=(15, 15))
|
||||
ax = fig.add_subplot(111, projection='3d')
|
||||
ax.set_xlim3d(-1, 1)
|
||||
ax.set_ylim3d(-1, 1)
|
||||
ax.set_zlim3d(-1, 1)
|
||||
ax.set_xlabel("X")
|
||||
ax.set_ylabel("Y")
|
||||
ax.set_zlabel("Z")
|
||||
ax.scatter(x[:, 0], x[:, 1], x[:, 2], s=5)
|
||||
if y is not None:
|
||||
ax.scatter(y[:, 0], y[:, 1], y[:, 2], s=5)
|
||||
plt.show()
|
||||
|
||||
def plot_uv(x, y=None):
|
||||
fig = plt.figure(figsize=(15, 15))
|
||||
ax = fig.add_subplot(111)
|
||||
ax.set_xlim(0, 1)
|
||||
ax.set_ylim(0, 1)
|
||||
ax.scatter(x[:, 0], x[:, 1])
|
||||
if y is not None:
|
||||
ax.scatter(y[:, 0], y[:, 1])
|
||||
plt.show()
|
||||
|
||||
"""
|
||||
Mesh to UV map
|
||||
Computes correspondences between 3D mesh and UV map
|
||||
NOTE: 3D mesh vertices can have multiple correspondences with UV vertices
|
||||
"""
|
||||
def mesh2UV(F, Ft):
|
||||
m2uv = {v: set() for f in F for v in f}
|
||||
for f, ft in zip(F, Ft):
|
||||
for v, vt in zip(f, ft):
|
||||
m2uv[v].add(vt)
|
||||
# m2uv = {k:list(v) for k,v in m2uv.items()}
|
||||
return m2uv
|
||||
|
||||
|
||||
def filter_m2u(m2u, to_keep):
|
||||
for key in m2u:
|
||||
if len(m2u[key]) == 2:
|
||||
v = list(m2u[key])
|
||||
if v[0] in to_keep:
|
||||
m2u[key].remove(v[1])
|
||||
else:
|
||||
m2u[key].remove(v[0])
|
||||
return m2u
|
||||
|
||||
|
||||
def convert_to_numbers(m2u):
|
||||
for key in m2u:
|
||||
m2u[key] = m2u[key].pop()
|
||||
return m2u
|
||||
|
||||
|
||||
def get_masked_index(mask, index):
|
||||
m = np.zeros((mask.shape[0]), dtype=bool)
|
||||
m[index] = True
|
||||
m = m[mask]
|
||||
new_index = np.argmax(m)
|
||||
return new_index
|
||||
|
||||
|
||||
def get_closest_point(verts_1, verts_2):
|
||||
dist = sp.spatial.distance.cdist(verts_1, verts_2)
|
||||
closest = np.argmin(dist, axis=-1)
|
||||
return closest
|
||||
|
||||
def get_normal_edge(p1, p2):
|
||||
d = p2 - p1
|
||||
return np.array([-d[1], d[0]])
|
||||
|
||||
def get_normal_vertex(p1, p2, p3):
|
||||
n1 = get_normal_edge(p1, p2)
|
||||
n2 = get_normal_edge(p2, p3)
|
||||
n = (n1 + n2)
|
||||
n = n / np.linalg.norm(n) / 20
|
||||
return n
|
||||
|
||||
|
||||
def get_distance(p1, p2, points):
|
||||
d = np.array([np.linalg.norm(np.cross(p2 - p1, point - p1)) / np.linalg.norm(p2 - p1) for point in points])
|
||||
return d
|
||||
|
||||
|
||||
def get_normals(verts):
|
||||
v = verts.copy()
|
||||
v = np.concatenate([v[-1][np.newaxis, :], v, v[0][np.newaxis, :]])
|
||||
normals = np.zeros_like(verts)
|
||||
for i in range(1, v.shape[0] - 1):
|
||||
normals[i - 1] = get_normal_vertex(v[i - 1], v[i], v[i + 1])
|
||||
return normals
|
||||
|
||||
|
||||
def get_closest_points_normals(verts1, verts2):
|
||||
dist_2 = sp.spatial.distance.cdist(verts1, verts2)
|
||||
normals = get_normals(verts1)
|
||||
|
||||
# fig = plt.figure(figsize=(15, 15))
|
||||
# ax = fig.add_subplot(111)
|
||||
# ax.set_xlim(0, 1)
|
||||
# ax.set_ylim(0, 1)
|
||||
# ax.set_aspect('equal')
|
||||
# ax.scatter(verts1[:, 0], verts1[:, 1])
|
||||
# for p1, p2 in zip(verts1, verts1 + normals):
|
||||
# ax.plot([p1[0], p2[0]], [p1[1], p2[1]])
|
||||
# plt.show()
|
||||
|
||||
closest = []
|
||||
i = 0
|
||||
for p1, p2 in zip(verts1, verts1 + normals):
|
||||
dist = get_distance(p1, p2, verts2)
|
||||
closest.append(np.argmin(dist + dist_2[i, :]))
|
||||
i += 1
|
||||
|
||||
return np.array(closest)
|
||||
|
||||
|
||||
def back():
|
||||
# load data
|
||||
data_body = np.load("Blender/uv_data_1.npz", allow_pickle=True)
|
||||
data_back = np.load("Blender/uv_data_plane_1_2.npz", allow_pickle=True)
|
||||
|
||||
# translate back uv to match body
|
||||
uv_back = data_back.f.Vt
|
||||
# uv_back[:, 0] = uv_back[:, 0] + 0.1
|
||||
# uv_back[:, 1] = uv_back[:, 1] + 0.42
|
||||
|
||||
# find front and back indices in uv
|
||||
uv_verts = np.append(data_body.f.Vt, np.zeros((data_body.f.Vt.shape[0], 1)), axis=-1)
|
||||
m = tm.Trimesh(uv_verts, data_body.f.Ft, process=False, validate=False)
|
||||
front, back = tm.graph.connected_components(m.edges, min_len=1, nodes=None, engine=None)
|
||||
|
||||
# get mesh to uv mapping for the back vertices
|
||||
m2u_body = mesh2UV(data_body.f.F, data_body.f.Ft)
|
||||
m2u_body = filter_m2u(m2u_body, set(back))
|
||||
m2u_body = convert_to_numbers(m2u_body)
|
||||
|
||||
m2u_back = mesh2UV(data_back.f.F, data_back.f.Ft)
|
||||
m2u_back = convert_to_numbers(m2u_back)
|
||||
|
||||
# create mask for filtering vertices belonging to the back
|
||||
body_back_mask = np.zeros((data_body.f.Vt.shape[0]), dtype=bool)
|
||||
body_back_mask[back] = True
|
||||
|
||||
# define which points should match after the transform on the mesh
|
||||
matching_points = np.array([[3284, 3], [9557, 13], [9278, 15], [8095, 1], [10597, 10], [10602, 17], [7114, 9], [6989, 16], [6983, 11], [4219, 0], [5628, 14], [5908, 12], [1447, 2], [1086, 5], [21, 8], [33, 4], [1888, 7], [2931, 6]])
|
||||
|
||||
# map matching point to uv coordinates
|
||||
for i in range(matching_points.shape[0]):
|
||||
matching_points[i, :] = [get_masked_index(body_back_mask, m2u_body[matching_points[i,0]]), m2u_back[matching_points[i, 1]]]
|
||||
|
||||
x = data_body.f.Vt[body_back_mask, :]
|
||||
y = uv_back
|
||||
|
||||
fig = plt.figure(figsize=(15, 15))
|
||||
ax = fig.add_subplot(121)
|
||||
ax.set_xlim(0, 1)
|
||||
ax.set_ylim(0, 1)
|
||||
ax.set_aspect('equal')
|
||||
ax.scatter(x[:, 0], x[:, 1])
|
||||
if y is not None:
|
||||
ax.scatter(y[:, 0], y[:, 1])
|
||||
for pair in matching_points:
|
||||
ax.scatter([x[pair[0], 0], y[pair[1], 0]], [x[pair[0], 1], y[pair[1], 1]], s=50)
|
||||
|
||||
# get transformation from body to back contour
|
||||
source = x[matching_points[:, 0], :]
|
||||
destination = y[matching_points[:, 1], :]
|
||||
trans = tps.TPS(source, destination)
|
||||
transformed = trans(data_body.f.Vt[back, :])
|
||||
|
||||
ax = fig.add_subplot(122)
|
||||
ax.set_xlim(0, 1)
|
||||
ax.set_ylim(0, 1)
|
||||
ax.set_aspect('equal')
|
||||
ax.scatter(transformed[:, 0], transformed[:, 1])
|
||||
if y is not None:
|
||||
ax.scatter(y[:, 0], y[:, 1])
|
||||
for pair in matching_points:
|
||||
ax.scatter([transformed[pair[0], 0], y[pair[1], 0]], [transformed[pair[0], 1], y[pair[1], 1]], s=50)
|
||||
|
||||
plt.show()
|
||||
|
||||
data_back_high = np.load("Blender/uv_data_plane_1_high_res.npz", allow_pickle=True)
|
||||
y_high = data_back_high.f.Vt
|
||||
# y_high[:, 0] = y_high[:, 0] + 0.1
|
||||
# y_high[:, 1] = y_high[:, 1] + 0.42
|
||||
|
||||
closest_points = get_closest_points_normals(y_high, transformed)
|
||||
# closest_points = get_closest_point(y_high, transformed)
|
||||
unique_mask = np.zeros_like(closest_points, dtype=bool)
|
||||
unique_mask[np.unique(closest_points, return_index=True)[1]] = True
|
||||
matching_points = np.array([closest_points[unique_mask], np.array(range(y_high.shape[0]))[unique_mask]]).T
|
||||
# matching_points_2 = []
|
||||
# for i, p in enumerate(matching_points):
|
||||
# if i % 10 == 0:
|
||||
# matching_points_2.append(p)
|
||||
# matching_points = np.array(matching_points_2)
|
||||
|
||||
source_2 = transformed[matching_points[:, 0], :]
|
||||
destination_2 = y_high[matching_points[:, 1], :]
|
||||
trans_2 = tps.TPS(source_2, destination_2)
|
||||
transformed_2 = trans_2(transformed)
|
||||
|
||||
|
||||
|
||||
fig = plt.figure(figsize=(15, 15))
|
||||
ax = fig.add_subplot(121)
|
||||
ax.set_xlim(0, 1)
|
||||
ax.set_ylim(0, 1)
|
||||
ax.set_aspect('equal')
|
||||
ax.scatter(source_2[:, 0], source_2[:, 1])
|
||||
ax.scatter(destination_2[:, 0], destination_2[:, 1])
|
||||
for i, pair in enumerate(zip(source_2, destination_2)):
|
||||
if i % 5 == 0:
|
||||
s, d = pair
|
||||
ax.scatter([s[0], d[0]], [s[1], d[1]])
|
||||
|
||||
ax = fig.add_subplot(122)
|
||||
ax.set_xlim(0, 1)
|
||||
ax.set_ylim(0, 1)
|
||||
ax.set_aspect('equal')
|
||||
ax.scatter(transformed_2[:, 0], transformed_2[:, 1])
|
||||
ax.scatter(y_high[:, 0], y_high[:, 1])
|
||||
plt.show()
|
||||
|
||||
return transformed_2
|
||||
|
||||
|
||||
def match_back_to_front(front_uv):
|
||||
# load data
|
||||
data_body = np.load("Blender/uv_data_smpl_2.npz", allow_pickle=True)
|
||||
# data_back = np.load("Blender/uv_data_plane_1_2.npz", allow_pickle=True)
|
||||
|
||||
# translate back uv to match body
|
||||
# uv_back = data_back.f.Vt
|
||||
# uv_back[:, 0] = uv_back[:, 0] + 0.1
|
||||
# uv_back[:, 1] = uv_back[:, 1] + 0.42
|
||||
|
||||
# find front and back indices in uv
|
||||
uv_verts = np.append(data_body.f.Vt, np.zeros((data_body.f.Vt.shape[0], 1)), axis=-1)
|
||||
m = tm.Trimesh(uv_verts, data_body.f.Ft, process=False, validate=False)
|
||||
front, back = tm.graph.connected_components(m.edges, min_len=1, nodes=None, engine=None)
|
||||
|
||||
# get mesh to uv mapping for the back vertices
|
||||
m2u_body = mesh2UV(data_body.f.F, data_body.f.Ft)
|
||||
m2u_body = filter_m2u(m2u_body, set(back))
|
||||
m2u_body = convert_to_numbers(m2u_body)
|
||||
|
||||
# m2u_back = mesh2UV(data_back.f.F, data_back.f.Ft)
|
||||
# m2u_back = convert_to_numbers(m2u_back)
|
||||
|
||||
# create mask for filtering vertices belonging to the back
|
||||
body_back_mask = np.zeros((data_body.f.Vt.shape[0]), dtype=bool)
|
||||
body_back_mask[back] = True
|
||||
body_front_mask = np.zeros((data_body.f.Vt.shape[0]), dtype=bool)
|
||||
body_front_mask[front] = True
|
||||
|
||||
# define which points should match after the transform on the mesh
|
||||
# matching_points = np.array([[3284, 3], [9557, 13], [9278, 15], [8095, 1], [10597, 10], [10602, 17], [7114, 9], [6989, 16], [6983, 11], [4219, 0], [5628, 14], [5908, 12], [1447, 2], [1086, 5], [21, 8], [33, 4], [1888, 7], [2931, 6]])
|
||||
|
||||
# # map matching point to uv coordinates
|
||||
# for i in range(matching_points.shape[0]):
|
||||
# matching_points[i, :] = [get_masked_index(body_back_mask, m2u_body[matching_points[i,0]]), m2u_back[matching_points[i, 1]]]
|
||||
|
||||
m2u_body_all = mesh2UV(data_body.f.F, data_body.f.Ft)
|
||||
m2u_body_all = m2uv_array(m2u_body_all)
|
||||
border_pairs = m2u_body_all[:, 0] != m2u_body_all[:, 1]
|
||||
border_pairs = m2u_body_all[border_pairs, :]
|
||||
matching_points = border_pairs
|
||||
matching_points_sorted = []
|
||||
for p in matching_points:
|
||||
if p[0] in front:
|
||||
matching_points_sorted.append([p[0], p[1]])
|
||||
elif p[0] in back:
|
||||
matching_points_sorted.append([p[1], p[0]])
|
||||
else:
|
||||
assert False
|
||||
# add neck V shape control point
|
||||
matching_points_sorted.append([m2u_body_all[10228][0], m2u_body_all[33][0]])
|
||||
matching_points = np.array(matching_points_sorted)
|
||||
mm = matching_points.copy()
|
||||
|
||||
for i in range(matching_points.shape[0]):
|
||||
matching_points[i, :] = [get_masked_index(body_back_mask, matching_points[i,1]), get_masked_index(body_front_mask, matching_points[i,0])]
|
||||
|
||||
temp = np.zeros_like(data_body.f.Vt)
|
||||
temp[front, :] = front_uv
|
||||
|
||||
x = data_body.f.Vt[body_back_mask, :]
|
||||
y = temp[body_front_mask, :] # front_uv #
|
||||
|
||||
fig = plt.figure(figsize=(15, 15))
|
||||
ax = fig.add_subplot(121)
|
||||
ax.set_xlim(0, 1)
|
||||
ax.set_ylim(0, 1)
|
||||
ax.set_aspect('equal')
|
||||
ax.scatter(x[:, 0], x[:, 1])
|
||||
if y is not None:
|
||||
ax.scatter(y[:, 0], y[:, 1])
|
||||
for pair in matching_points:
|
||||
ax.scatter([x[pair[0], 0], y[pair[1], 0]], [x[pair[0], 1], y[pair[1], 1]], s=50)
|
||||
|
||||
# get transformation from body to back contour
|
||||
source = x[matching_points[:, 0], :]
|
||||
destination = y[matching_points[:, 1], :]
|
||||
trans = tps.TPS(source, destination)
|
||||
transformed = trans(data_body.f.Vt[back, :])
|
||||
|
||||
ax = fig.add_subplot(122)
|
||||
ax.set_xlim(0, 1)
|
||||
ax.set_ylim(0, 1)
|
||||
ax.set_aspect('equal')
|
||||
ax.scatter(transformed[:, 0], transformed[:, 1])
|
||||
if y is not None:
|
||||
ax.scatter(y[:, 0], y[:, 1])
|
||||
for pair in matching_points:
|
||||
ax.scatter([transformed[pair[0], 0], y[pair[1], 0]], [transformed[pair[0], 1], y[pair[1], 1]], s=50)
|
||||
|
||||
plt.show()
|
||||
|
||||
|
||||
return transformed
|
||||
|
||||
|
||||
def generate_mask():
|
||||
import cv2
|
||||
res = 1024
|
||||
front = np.load("Blender/uv_data_plane_0_high_res_2.npz", allow_pickle=True)
|
||||
back = np.load("Blender/uv_data_plane_1_high_res_2.npz", allow_pickle=True)
|
||||
points = np.round(np.concatenate((front.f.Vt, back.f.Vt)) * res).astype(np.int32)
|
||||
img = np.zeros((res, res, 3))
|
||||
img[points[:, 0], points[:, 1], :] = 255
|
||||
cv2.imwrite('MeshUtils/mask_1024.png', img)
|
||||
|
||||
# do paint now
|
||||
|
||||
img = cv2.imread('MeshUtils/mask_1024_2.png')
|
||||
img = img[:, :, 0] == 255
|
||||
np.save('MeshUtils/mask_1024.npy', img)
|
||||
|
||||
def generate_mask_2():
|
||||
import cv2
|
||||
res = 256
|
||||
front = np.load("Blender/uv_data_plane_stretched_front_high_extra.npz", allow_pickle=True)
|
||||
back = np.load("Blender/uv_data_plane_stretched_front_high_extra.npz", allow_pickle=True)
|
||||
points = np.floor(np.concatenate((front.f.Vt, back.f.Vt + [1, 0])) * res).astype(np.int32)
|
||||
img = np.zeros((res * 2, res, 3))
|
||||
img[points[:, 0], points[:, 1], :] = 255
|
||||
cv2.imwrite('MeshUtils/mask_256x512.png', img)
|
||||
|
||||
# do paint now
|
||||
|
||||
img = cv2.imread('MeshUtils/mask_256x512.png')
|
||||
img = img[:, :, 0] > 240
|
||||
np.save('MeshUtils/mask_256x512.npy', img)
|
||||
|
||||
pass
|
||||
|
||||
|
||||
def generate_uv_mask_border():
|
||||
import cv2
|
||||
mask = np.load('MeshUtils/mask_256x512.npy')
|
||||
uv_res = 256
|
||||
mask_border = (cv2.resize(mask * 255.0, (uv_res, uv_res * 2)) / 255).astype(bool)[:, :, np.newaxis]
|
||||
mask_border = cv2.erode(mask_border.astype(np.float32), cv2.getStructuringElement(cv2.MORPH_RECT, (2 * 2 + 1, 2 * 2 + 1), (2, 2)), iterations=1).astype(bool)
|
||||
|
||||
mask[mask_border == True] = 0
|
||||
# plt.imshow(mask)
|
||||
idx_front = np.transpose(np.vstack(np.nonzero(mask[:256, :])), [1,0])
|
||||
|
||||
test = np.zeros_like(mask)
|
||||
test[idx_front[:,0], idx_front[:,1]] = 1
|
||||
idx_back = idx_front.copy()
|
||||
idx_back[:, 0] += 256
|
||||
test[idx_back[:,0], idx_back[:,1]] = 1
|
||||
|
||||
|
||||
pass
|
||||
|
||||
|
||||
|
||||
def front():
|
||||
# load data
|
||||
data_body = np.load("custom_data/uv_map/uv_data_smpl_2.npz", allow_pickle=True)
|
||||
data_front = np.load("custom_data/uv_map/uv_data_plane_stretched_front_low.npz", allow_pickle=True)
|
||||
|
||||
# translate back uv to match body
|
||||
uv_front = data_front.f.Vt
|
||||
# uv_back[:, 0] = uv_back[:, 0] + 0.1
|
||||
# uv_back[:, 1] = uv_back[:, 1] + 0.42
|
||||
|
||||
# find front and back indices in uv
|
||||
uv_verts = np.append(data_body.f.Vt, np.zeros((data_body.f.Vt.shape[0], 1)), axis=-1)
|
||||
m = tm.Trimesh(uv_verts, data_body.f.Ft, process=False, validate=False)
|
||||
front, back = tm.graph.connected_components(m.edges, min_len=1, nodes=None, engine=None)
|
||||
|
||||
# get mesh to uv mapping for the back vertices
|
||||
m2u_body = mesh2UV(data_body.f.F, data_body.f.Ft)
|
||||
m2u_body = filter_m2u(m2u_body, set(front))
|
||||
m2u_body = convert_to_numbers(m2u_body)
|
||||
|
||||
m2u_front = mesh2UV(data_front.f.F, data_front.f.Ft)
|
||||
m2u_front = convert_to_numbers(m2u_front)
|
||||
|
||||
# create mask for filtering vertices belonging to the back
|
||||
body_front_mask = np.zeros((data_body.f.Vt.shape[0]), dtype=bool)
|
||||
body_front_mask[front] = True
|
||||
|
||||
# define which points should match after the transform on the mesh
|
||||
matching_points = np.array([ [9557, 12], [3284, 2], [3053, 5], [1888, 8], [10228, 4], [21, 7], [1209, 6], [1447, 3], [5908, 13], [4617, 15], [4219, 1], [6983, 10], [6989, 17], [7114, 9], [10602, 16], [10597, 11], [7884, 0], [8280, 14] ])
|
||||
# matching_points = np.array([ [9557, 12], [3284, 2], [3053, 5], [1888, 8], [10228, 4], [21, 7], [1209, 6], [1447, 3], [5908, 13], [4793, 15], [4219, 1], [6983, 10], [6989, 17], [7114, 9], [10602, 16], [10597, 11], [7884, 0], [8450, 14], [13566, 18], [11767, 19], [8668, 20], [5012, 21], [12716, 22], [10917, 23], [8280, 24], [4617, 25], [2760, 26], [911, 27] ])
|
||||
|
||||
# map matching point to uv coordinates
|
||||
for i in range(matching_points.shape[0]):
|
||||
matching_points[i, :] = [get_masked_index(body_front_mask, m2u_body[matching_points[i,0]]), m2u_front[matching_points[i, 1]]]
|
||||
|
||||
x = data_body.f.Vt[body_front_mask, :]
|
||||
y = uv_front
|
||||
|
||||
fig = plt.figure(figsize=(15, 15))
|
||||
ax = fig.add_subplot(121)
|
||||
# ax.set_xlim(0, 1)
|
||||
# ax.set_ylim(0, 1)
|
||||
ax.set_aspect('equal')
|
||||
ax.scatter(x[:, 0], x[:, 1])
|
||||
if y is not None:
|
||||
ax.scatter(y[:, 0], y[:, 1])
|
||||
for pair in matching_points:
|
||||
ax.scatter([x[pair[0], 0], y[pair[1], 0]], [x[pair[0], 1], y[pair[1], 1]], s=50)
|
||||
|
||||
# get transformation from body to front contour
|
||||
source = x[matching_points[:, 0], :]
|
||||
destination = y[matching_points[:, 1], :]
|
||||
trans = tps.TPS(source, destination)
|
||||
transformed = trans(data_body.f.Vt[front, :])
|
||||
|
||||
# transformed = (transformed - transformed.min()) / (transformed.max() - transformed.min())
|
||||
|
||||
ax = fig.add_subplot(122)
|
||||
# ax.set_xlim(0, 1)
|
||||
# ax.set_ylim(0, 1)
|
||||
ax.set_aspect('equal')
|
||||
ax.scatter(transformed[:, 0], transformed[:, 1])
|
||||
if y is not None:
|
||||
ax.scatter(y[:, 0], y[:, 1])
|
||||
for pair in matching_points:
|
||||
ax.scatter([transformed[pair[0], 0], y[pair[1], 0]], [transformed[pair[0], 1], y[pair[1], 1]], s=50)
|
||||
|
||||
plt.show()
|
||||
|
||||
data_front_high = np.load("custom_data/uv_map/uv_data_plane_stretched_front_high.npz", allow_pickle=True)
|
||||
y_high = data_front_high.f.Vt
|
||||
# y_high[:, 0] = y_high[:, 0] + 0.1
|
||||
# y_high[:, 1] = y_high[:, 1] + 0.42
|
||||
|
||||
closest_points = get_closest_points_normals(y_high, transformed)
|
||||
# closest_points = get_closest_point(y_high, transformed)
|
||||
unique_mask = np.zeros_like(closest_points, dtype=bool)
|
||||
unique_mask[np.unique(closest_points, return_index=True)[1]] = True
|
||||
matching_points = np.array([closest_points[unique_mask], np.array(range(y_high.shape[0]))[unique_mask]]).T
|
||||
# matching_points_2 = []
|
||||
# for i, p in enumerate(matching_points):
|
||||
# if i % 10 == 0:
|
||||
# matching_points_2.append(p)
|
||||
# matching_points = np.array(matching_points_2)
|
||||
|
||||
source_2 = transformed[matching_points[:, 0], :]
|
||||
destination_2 = y_high[matching_points[:, 1], :]
|
||||
trans_2 = tps.TPS(source_2, destination_2)
|
||||
transformed_2 = trans_2(transformed)
|
||||
|
||||
|
||||
|
||||
fig = plt.figure(figsize=(15, 15))
|
||||
ax = fig.add_subplot(121)
|
||||
ax.set_xlim(0, 1)
|
||||
ax.set_ylim(0, 1)
|
||||
ax.set_aspect('equal')
|
||||
ax.scatter(source_2[:, 0], source_2[:, 1])
|
||||
ax.scatter(destination_2[:, 0], destination_2[:, 1])
|
||||
for i, pair in enumerate(zip(source_2, destination_2)):
|
||||
if i % 5 == 0:
|
||||
s, d = pair
|
||||
ax.scatter([s[0], d[0]], [s[1], d[1]])
|
||||
|
||||
ax = fig.add_subplot(122)
|
||||
ax.set_xlim(0, 1)
|
||||
ax.set_ylim(0, 1)
|
||||
ax.set_aspect('equal')
|
||||
ax.scatter(transformed_2[:, 0], transformed_2[:, 1])
|
||||
ax.scatter(y_high[:, 0], y_high[:, 1])
|
||||
plt.show()
|
||||
|
||||
return transformed_2, mesh2UV(data_body.f.F, data_body.f.Ft), front, back
|
||||
|
||||
|
||||
def smoothing():
|
||||
data_body = np.load("Blender/uv_data_smpl_2.npz", allow_pickle=True)
|
||||
uv_map = np.load('MeshUtils\stretched_uv_map_3.npy')
|
||||
verts = np.concatenate((uv_map, np.ones((15068, 1))), axis=-1)
|
||||
m = tm.Trimesh(verts, data_body.f.Ft, process=False, validate=False)
|
||||
|
||||
|
||||
m2u_body_all = mesh2UV(data_body.f.F, data_body.f.Ft)
|
||||
m2u_body_all = m2uv_array(m2u_body_all)
|
||||
border_pairs = m2u_body_all[:, 0] != m2u_body_all[:, 1]
|
||||
border_pairs = m2u_body_all[border_pairs, :]
|
||||
pinned_vertices = border_pairs.reshape(-1)
|
||||
|
||||
lap = tm.smoothing.laplacian_calculation(m, equal_weight=False, pinned_vertices=pinned_vertices)
|
||||
# tm.smoothing.filter_laplacian(m, lamb=0.5, iterations=100, volume_constraint=False, laplacian_operator=lap)
|
||||
# tm.smoothing.filter_mut_dif_laplacian(m, lamb=0.5, iterations=100, volume_constraint=False, laplacian_operator=lap)
|
||||
# tm.smoothing.filter_humphrey(m, iterations=100, laplacian_operator=lap)
|
||||
tm.smoothing.filter_taubin(m, iterations=100, laplacian_operator=lap)
|
||||
|
||||
|
||||
x = uv_map
|
||||
y = m.vertices[:, :2]
|
||||
|
||||
fig = plt.figure(figsize=(15, 15))
|
||||
ax = fig.add_subplot(211)
|
||||
ax.set_xlim(0, 2)
|
||||
ax.set_ylim(0, 1)
|
||||
# ax.set_aspect('equal')
|
||||
ax.scatter(x[:, 0], x[:, 1])
|
||||
|
||||
ax = fig.add_subplot(212)
|
||||
ax.set_xlim(0, 2)
|
||||
ax.set_ylim(0, 1)
|
||||
# ax.set_aspect('equal')
|
||||
ax.scatter(y[:, 0], y[:, 1])
|
||||
plt.show()
|
||||
|
||||
|
||||
pass
|
||||
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
# generate_mask_2()
|
||||
# smoothing()
|
||||
# generate_uv_mask_border()
|
||||
|
||||
front_uv, m2uv, front_mask, back_mask = front()
|
||||
back_uv = match_back_to_front(front_uv)
|
||||
# back_uv = back()
|
||||
|
||||
final = np.zeros((15068, 2))
|
||||
final[front_mask] = front_uv
|
||||
final[back_mask] = back_uv + [1, 0]
|
||||
np.save('MeshUtils\stretched_uv_map_3.npy', final)
|
||||
|
||||
pass
|
||||
@@ -0,0 +1,846 @@
|
||||
import os
|
||||
import sys
|
||||
import shutil
|
||||
import tqdm
|
||||
import random
|
||||
import math
|
||||
import collections
|
||||
from tps_custom import TPS
|
||||
import numpy as np
|
||||
import scipy as sp
|
||||
import trimesh as tm
|
||||
import cmapy as cm
|
||||
import scipy.io as sio
|
||||
import matplotlib.pyplot as plt
|
||||
import plotly.express as ex
|
||||
import plotly.graph_objects as go
|
||||
from scipy.spatial.distance import cdist
|
||||
|
||||
|
||||
|
||||
def m2uv_array(m2uv, is_smpl_orig=False):
|
||||
result = np.zeros((14475 if not is_smpl_orig else 3638, 2), dtype=np.int32)
|
||||
for k, v in m2uv.items():
|
||||
if len(v) == 2:
|
||||
result[k] = list(v)
|
||||
elif len(v) == 1:
|
||||
x = v.pop()
|
||||
result[k] = [x, x]
|
||||
else:
|
||||
assert False
|
||||
|
||||
return result
|
||||
|
||||
|
||||
def get_body():
|
||||
smpl = SMPL('FeatureLines/superSMPL')
|
||||
pose = np.zeros((1, 72), dtype=np.float32)
|
||||
beta = np.zeros((1, 11), dtype=np.float32)
|
||||
R = smpl.run_rodrigues(pose)
|
||||
T, _, _ = smpl.run(R, beta, True)
|
||||
return np.squeeze(T)
|
||||
|
||||
|
||||
def plot_3d(x, y=None):
|
||||
fig = plt.figure(figsize=(15, 15))
|
||||
ax = fig.add_subplot(111, projection='3d')
|
||||
ax.set_xlim3d(-1, 1)
|
||||
ax.set_ylim3d(-1, 1)
|
||||
ax.set_zlim3d(-1, 1)
|
||||
ax.set_xlabel("X")
|
||||
ax.set_ylabel("Y")
|
||||
ax.set_zlabel("Z")
|
||||
ax.scatter(x[:, 0], x[:, 1], x[:, 2], s=5)
|
||||
if y is not None:
|
||||
ax.scatter(y[:, 0], y[:, 1], y[:, 2], s=5)
|
||||
plt.show()
|
||||
|
||||
def plot_uv(x, y=None):
|
||||
fig = plt.figure(figsize=(15, 15))
|
||||
ax = fig.add_subplot(111)
|
||||
ax.set_xlim(0, 1)
|
||||
ax.set_ylim(0, 1)
|
||||
ax.scatter(x[:, 0], x[:, 1])
|
||||
if y is not None:
|
||||
ax.scatter(y[:, 0], y[:, 1])
|
||||
plt.show()
|
||||
|
||||
"""
|
||||
Mesh to UV map
|
||||
Computes correspondences between 3D mesh and UV map
|
||||
NOTE: 3D mesh vertices can have multiple correspondences with UV vertices
|
||||
"""
|
||||
def mesh2UV(F, Ft):
|
||||
m2uv = {v: set() for f in F for v in f}
|
||||
for f, ft in zip(F, Ft):
|
||||
for v, vt in zip(f, ft):
|
||||
m2uv[v].add(vt)
|
||||
# m2uv = {k:list(v) for k,v in m2uv.items()}
|
||||
return m2uv
|
||||
|
||||
|
||||
def filter_m2u(m2u, to_keep):
|
||||
for key in m2u:
|
||||
if len(m2u[key]) == 2:
|
||||
v = list(m2u[key])
|
||||
if v[0] in to_keep:
|
||||
m2u[key].remove(v[1])
|
||||
else:
|
||||
m2u[key].remove(v[0])
|
||||
return m2u
|
||||
|
||||
|
||||
def convert_to_numbers(m2u):
|
||||
for key in m2u:
|
||||
m2u[key] = m2u[key].pop()
|
||||
return m2u
|
||||
|
||||
|
||||
def get_masked_index(mask, index):
|
||||
m = np.zeros((mask.shape[0]), dtype=bool)
|
||||
m[index] = True
|
||||
m = m[mask]
|
||||
new_index = np.argmax(m)
|
||||
return new_index
|
||||
|
||||
|
||||
def get_closest_point(verts_1, verts_2):
|
||||
dist = sp.spatial.distance.cdist(verts_1, verts_2)
|
||||
closest = np.argmin(dist, axis=-1)
|
||||
return closest
|
||||
|
||||
def get_normal_edge(p1, p2):
|
||||
d = p2 - p1
|
||||
return np.array([-d[1], d[0]])
|
||||
|
||||
def get_normal_vertex(p1, p2, p3):
|
||||
n1 = get_normal_edge(p1, p2)
|
||||
n2 = get_normal_edge(p2, p3)
|
||||
n = (n1 + n2)
|
||||
n = n / np.linalg.norm(n) / 20
|
||||
return n
|
||||
|
||||
|
||||
def get_distance(p1, p2, points):
|
||||
d = np.array([np.linalg.norm(np.cross(p2 - p1, point - p1)) / np.linalg.norm(p2 - p1) for point in points])
|
||||
return d
|
||||
|
||||
|
||||
def get_normals(verts):
|
||||
v = verts.copy()
|
||||
v = np.concatenate([v[-1][np.newaxis, :], v, v[0][np.newaxis, :]])
|
||||
normals = np.zeros_like(verts)
|
||||
for i in range(1, v.shape[0] - 1):
|
||||
normals[i - 1] = get_normal_vertex(v[i - 1], v[i], v[i + 1])
|
||||
return normals
|
||||
|
||||
|
||||
def get_closest_points_normals(verts1, verts2):
|
||||
dist_2 = sp.spatial.distance.cdist(verts1, verts2)
|
||||
normals = get_normals(verts1)
|
||||
|
||||
# fig = plt.figure(figsize=(15, 15))
|
||||
# ax = fig.add_subplot(111)
|
||||
# ax.set_xlim(0, 1)
|
||||
# ax.set_ylim(0, 1)
|
||||
# ax.set_aspect('equal')
|
||||
# ax.scatter(verts1[:, 0], verts1[:, 1])
|
||||
# for p1, p2 in zip(verts1, verts1 + normals):
|
||||
# ax.plot([p1[0], p2[0]], [p1[1], p2[1]])
|
||||
# plt.show()
|
||||
|
||||
closest = []
|
||||
i = 0
|
||||
for p1, p2 in zip(verts1, verts1 + normals):
|
||||
dist = get_distance(p1, p2, verts2)
|
||||
closest.append(np.argmin(dist + dist_2[i, :]))
|
||||
i += 1
|
||||
|
||||
return np.array(closest)
|
||||
|
||||
|
||||
def back():
|
||||
# load data
|
||||
data_body = np.load("Blender/uv_data_1.npz", allow_pickle=True)
|
||||
data_back = np.load("Blender/uv_data_plane_1_2.npz", allow_pickle=True)
|
||||
|
||||
# translate back uv to match body
|
||||
uv_back = data_back.f.Vt
|
||||
# uv_back[:, 0] = uv_back[:, 0] + 0.1
|
||||
# uv_back[:, 1] = uv_back[:, 1] + 0.42
|
||||
|
||||
# find front and back indices in uv
|
||||
uv_verts = np.append(data_body.f.Vt, np.zeros((data_body.f.Vt.shape[0], 1)), axis=-1)
|
||||
m = tm.Trimesh(uv_verts, data_body.f.Ft, process=False, validate=False)
|
||||
front, back = tm.graph.connected_components(m.edges, min_len=1, nodes=None, engine=None)
|
||||
|
||||
# get mesh to uv mapping for the back vertices
|
||||
m2u_body = mesh2UV(data_body.f.F, data_body.f.Ft)
|
||||
m2u_body = filter_m2u(m2u_body, set(back))
|
||||
m2u_body = convert_to_numbers(m2u_body)
|
||||
|
||||
m2u_back = mesh2UV(data_back.f.F, data_back.f.Ft)
|
||||
m2u_back = convert_to_numbers(m2u_back)
|
||||
|
||||
# create mask for filtering vertices belonging to the back
|
||||
body_back_mask = np.zeros((data_body.f.Vt.shape[0]), dtype=bool)
|
||||
body_back_mask[back] = True
|
||||
|
||||
# define which points should match after the transform on the mesh
|
||||
matching_points = np.array([[3284, 3], [9557, 13], [9278, 15], [8095, 1], [10597, 10], [10602, 17], [7114, 9], [6989, 16], [6983, 11], [4219, 0], [5628, 14], [5908, 12], [1447, 2], [1086, 5], [21, 8], [33, 4], [1888, 7], [2931, 6]])
|
||||
|
||||
# map matching point to uv coordinates
|
||||
for i in range(matching_points.shape[0]):
|
||||
matching_points[i, :] = [get_masked_index(body_back_mask, m2u_body[matching_points[i,0]]), m2u_back[matching_points[i, 1]]]
|
||||
|
||||
x = data_body.f.Vt[body_back_mask, :]
|
||||
y = uv_back
|
||||
|
||||
fig = plt.figure(figsize=(15, 15))
|
||||
ax = fig.add_subplot(121)
|
||||
ax.set_xlim(0, 1)
|
||||
ax.set_ylim(0, 1)
|
||||
ax.set_aspect('equal')
|
||||
ax.scatter(x[:, 0], x[:, 1])
|
||||
if y is not None:
|
||||
ax.scatter(y[:, 0], y[:, 1])
|
||||
for pair in matching_points:
|
||||
ax.scatter([x[pair[0], 0], y[pair[1], 0]], [x[pair[0], 1], y[pair[1], 1]], s=50)
|
||||
|
||||
# get transformation from body to back contour
|
||||
source = x[matching_points[:, 0], :]
|
||||
destination = y[matching_points[:, 1], :]
|
||||
trans = tps.TPS(source, destination)
|
||||
transformed = trans(data_body.f.Vt[back, :])
|
||||
|
||||
ax = fig.add_subplot(122)
|
||||
ax.set_xlim(0, 1)
|
||||
ax.set_ylim(0, 1)
|
||||
ax.set_aspect('equal')
|
||||
ax.scatter(transformed[:, 0], transformed[:, 1])
|
||||
if y is not None:
|
||||
ax.scatter(y[:, 0], y[:, 1])
|
||||
for pair in matching_points:
|
||||
ax.scatter([transformed[pair[0], 0], y[pair[1], 0]], [transformed[pair[0], 1], y[pair[1], 1]], s=50)
|
||||
|
||||
plt.show()
|
||||
|
||||
data_back_high = np.load("Blender/uv_data_plane_1_high_res.npz", allow_pickle=True)
|
||||
y_high = data_back_high.f.Vt
|
||||
# y_high[:, 0] = y_high[:, 0] + 0.1
|
||||
# y_high[:, 1] = y_high[:, 1] + 0.42
|
||||
|
||||
closest_points = get_closest_points_normals(y_high, transformed)
|
||||
# closest_points = get_closest_point(y_high, transformed)
|
||||
unique_mask = np.zeros_like(closest_points, dtype=bool)
|
||||
unique_mask[np.unique(closest_points, return_index=True)[1]] = True
|
||||
matching_points = np.array([closest_points[unique_mask], np.array(range(y_high.shape[0]))[unique_mask]]).T
|
||||
# matching_points_2 = []
|
||||
# for i, p in enumerate(matching_points):
|
||||
# if i % 10 == 0:
|
||||
# matching_points_2.append(p)
|
||||
# matching_points = np.array(matching_points_2)
|
||||
|
||||
source_2 = transformed[matching_points[:, 0], :]
|
||||
destination_2 = y_high[matching_points[:, 1], :]
|
||||
trans_2 = tps.TPS(source_2, destination_2)
|
||||
transformed_2 = trans_2(transformed)
|
||||
|
||||
|
||||
|
||||
fig = plt.figure(figsize=(15, 15))
|
||||
ax = fig.add_subplot(121)
|
||||
ax.set_xlim(0, 1)
|
||||
ax.set_ylim(0, 1)
|
||||
ax.set_aspect('equal')
|
||||
ax.scatter(source_2[:, 0], source_2[:, 1])
|
||||
ax.scatter(destination_2[:, 0], destination_2[:, 1])
|
||||
for i, pair in enumerate(zip(source_2, destination_2)):
|
||||
if i % 5 == 0:
|
||||
s, d = pair
|
||||
ax.scatter([s[0], d[0]], [s[1], d[1]])
|
||||
|
||||
ax = fig.add_subplot(122)
|
||||
ax.set_xlim(0, 1)
|
||||
ax.set_ylim(0, 1)
|
||||
ax.set_aspect('equal')
|
||||
ax.scatter(transformed_2[:, 0], transformed_2[:, 1])
|
||||
ax.scatter(y_high[:, 0], y_high[:, 1])
|
||||
plt.show()
|
||||
|
||||
return transformed_2
|
||||
|
||||
|
||||
def match_back_to_front(front_uv):
|
||||
# load data
|
||||
data_body = np.load("Blender/uv_data_smpl_2.npz", allow_pickle=True)
|
||||
# data_back = np.load("Blender/uv_data_plane_1_2.npz", allow_pickle=True)
|
||||
|
||||
# translate back uv to match body
|
||||
# uv_back = data_back.f.Vt
|
||||
# uv_back[:, 0] = uv_back[:, 0] + 0.1
|
||||
# uv_back[:, 1] = uv_back[:, 1] + 0.42
|
||||
|
||||
# find front and back indices in uv
|
||||
uv_verts = np.append(data_body.f.Vt, np.zeros((data_body.f.Vt.shape[0], 1)), axis=-1)
|
||||
m = tm.Trimesh(uv_verts, data_body.f.Ft, process=False, validate=False)
|
||||
front, back = tm.graph.connected_components(m.edges, min_len=1, nodes=None, engine=None)
|
||||
|
||||
# get mesh to uv mapping for the back vertices
|
||||
m2u_body = mesh2UV(data_body.f.F, data_body.f.Ft)
|
||||
m2u_body = filter_m2u(m2u_body, set(back))
|
||||
m2u_body = convert_to_numbers(m2u_body)
|
||||
|
||||
# m2u_back = mesh2UV(data_back.f.F, data_back.f.Ft)
|
||||
# m2u_back = convert_to_numbers(m2u_back)
|
||||
|
||||
# create mask for filtering vertices belonging to the back
|
||||
body_back_mask = np.zeros((data_body.f.Vt.shape[0]), dtype=bool)
|
||||
body_back_mask[back] = True
|
||||
body_front_mask = np.zeros((data_body.f.Vt.shape[0]), dtype=bool)
|
||||
body_front_mask[front] = True
|
||||
|
||||
# define which points should match after the transform on the mesh
|
||||
# matching_points = np.array([[3284, 3], [9557, 13], [9278, 15], [8095, 1], [10597, 10], [10602, 17], [7114, 9], [6989, 16], [6983, 11], [4219, 0], [5628, 14], [5908, 12], [1447, 2], [1086, 5], [21, 8], [33, 4], [1888, 7], [2931, 6]])
|
||||
|
||||
# # map matching point to uv coordinates
|
||||
# for i in range(matching_points.shape[0]):
|
||||
# matching_points[i, :] = [get_masked_index(body_back_mask, m2u_body[matching_points[i,0]]), m2u_back[matching_points[i, 1]]]
|
||||
|
||||
m2u_body_all = mesh2UV(data_body.f.F, data_body.f.Ft)
|
||||
m2u_body_all = m2uv_array(m2u_body_all)
|
||||
border_pairs = m2u_body_all[:, 0] != m2u_body_all[:, 1]
|
||||
border_pairs = m2u_body_all[border_pairs, :]
|
||||
matching_points = border_pairs
|
||||
matching_points_sorted = []
|
||||
for p in matching_points:
|
||||
if p[0] in front:
|
||||
matching_points_sorted.append([p[0], p[1]])
|
||||
elif p[0] in back:
|
||||
matching_points_sorted.append([p[1], p[0]])
|
||||
else:
|
||||
assert False
|
||||
# add neck V shape control point
|
||||
matching_points_sorted.append([m2u_body_all[10228][0], m2u_body_all[33][0]])
|
||||
matching_points = np.array(matching_points_sorted)
|
||||
mm = matching_points.copy()
|
||||
|
||||
for i in range(matching_points.shape[0]):
|
||||
matching_points[i, :] = [get_masked_index(body_back_mask, matching_points[i,1]), get_masked_index(body_front_mask, matching_points[i,0])]
|
||||
|
||||
temp = np.zeros_like(data_body.f.Vt)
|
||||
temp[front, :] = front_uv
|
||||
|
||||
x = data_body.f.Vt[body_back_mask, :]
|
||||
y = temp[body_front_mask, :] # front_uv #
|
||||
|
||||
fig = plt.figure(figsize=(15, 15))
|
||||
ax = fig.add_subplot(121)
|
||||
ax.set_xlim(0, 1)
|
||||
ax.set_ylim(0, 1)
|
||||
ax.set_aspect('equal')
|
||||
ax.scatter(x[:, 0], x[:, 1])
|
||||
if y is not None:
|
||||
ax.scatter(y[:, 0], y[:, 1])
|
||||
for pair in matching_points:
|
||||
ax.scatter([x[pair[0], 0], y[pair[1], 0]], [x[pair[0], 1], y[pair[1], 1]], s=50)
|
||||
|
||||
# get transformation from body to back contour
|
||||
source = x[matching_points[:, 0], :]
|
||||
destination = y[matching_points[:, 1], :]
|
||||
trans = tps.TPS(source, destination)
|
||||
transformed = trans(data_body.f.Vt[back, :])
|
||||
|
||||
ax = fig.add_subplot(122)
|
||||
ax.set_xlim(0, 1)
|
||||
ax.set_ylim(0, 1)
|
||||
ax.set_aspect('equal')
|
||||
ax.scatter(transformed[:, 0], transformed[:, 1])
|
||||
if y is not None:
|
||||
ax.scatter(y[:, 0], y[:, 1])
|
||||
for pair in matching_points:
|
||||
ax.scatter([transformed[pair[0], 0], y[pair[1], 0]], [transformed[pair[0], 1], y[pair[1], 1]], s=50)
|
||||
|
||||
plt.show()
|
||||
|
||||
|
||||
return transformed
|
||||
|
||||
|
||||
def generate_mask():
|
||||
import cv2
|
||||
res = 1024
|
||||
front = np.load("Blender/uv_data_plane_0_high_res_2.npz", allow_pickle=True)
|
||||
back = np.load("Blender/uv_data_plane_1_high_res_2.npz", allow_pickle=True)
|
||||
points = np.round(np.concatenate((front.f.Vt, back.f.Vt)) * res).astype(np.int32)
|
||||
img = np.zeros((res, res, 3))
|
||||
img[points[:, 0], points[:, 1], :] = 255
|
||||
cv2.imwrite('MeshUtils/mask_1024.png', img)
|
||||
|
||||
# do paint now
|
||||
|
||||
img = cv2.imread('MeshUtils/mask_1024_2.png')
|
||||
img = img[:, :, 0] == 255
|
||||
np.save('MeshUtils/mask_1024.npy', img)
|
||||
|
||||
def generate_mask_2():
|
||||
import cv2
|
||||
res = 256
|
||||
front = np.load("Blender/uv_data_plane_stretched_front_high_extra.npz", allow_pickle=True)
|
||||
back = np.load("Blender/uv_data_plane_stretched_front_high_extra.npz", allow_pickle=True)
|
||||
points = np.floor(np.concatenate((front.f.Vt, back.f.Vt + [1, 0])) * res).astype(np.int32)
|
||||
img = np.zeros((res * 2, res, 3))
|
||||
img[points[:, 0], points[:, 1], :] = 255
|
||||
cv2.imwrite('MeshUtils/mask_256x512.png', img)
|
||||
|
||||
# do paint now
|
||||
|
||||
img = cv2.imread('MeshUtils/mask_256x512.png')
|
||||
img = img[:, :, 0] > 240
|
||||
np.save('MeshUtils/mask_256x512.npy', img)
|
||||
|
||||
pass
|
||||
|
||||
|
||||
def generate_uv_mask_border():
|
||||
import cv2
|
||||
mask = np.load('MeshUtils/mask_256x512.npy')
|
||||
uv_res = 256
|
||||
mask_border = (cv2.resize(mask * 255.0, (uv_res, uv_res * 2)) / 255).astype(bool)[:, :, np.newaxis]
|
||||
mask_border = cv2.erode(mask_border.astype(np.float32), cv2.getStructuringElement(cv2.MORPH_RECT, (2 * 2 + 1, 2 * 2 + 1), (2, 2)), iterations=1).astype(bool)
|
||||
|
||||
mask[mask_border == True] = 0
|
||||
# plt.imshow(mask)
|
||||
idx_front = np.transpose(np.vstack(np.nonzero(mask[:256, :])), [1,0])
|
||||
|
||||
test = np.zeros_like(mask)
|
||||
test[idx_front[:,0], idx_front[:,1]] = 1
|
||||
idx_back = idx_front.copy()
|
||||
idx_back[:, 0] += 256
|
||||
test[idx_back[:,0], idx_back[:,1]] = 1
|
||||
|
||||
|
||||
pass
|
||||
|
||||
|
||||
|
||||
def front():
|
||||
# load data
|
||||
data_body = np.load("custom_data/uv_map/uv_data_smpl_2.npz", allow_pickle=True)
|
||||
data_front = np.load("custom_data/uv_map/uv_data_plane_stretched_front_low.npz", allow_pickle=True)
|
||||
|
||||
# translate back uv to match body
|
||||
uv_front = data_front.f.Vt
|
||||
# uv_back[:, 0] = uv_back[:, 0] + 0.1
|
||||
# uv_back[:, 1] = uv_back[:, 1] + 0.42
|
||||
|
||||
# find front and back indices in uv
|
||||
uv_verts = np.append(data_body.f.Vt, np.zeros((data_body.f.Vt.shape[0], 1)), axis=-1)
|
||||
m = tm.Trimesh(uv_verts, data_body.f.Ft, process=False, validate=False)
|
||||
front, back = tm.graph.connected_components(m.edges, min_len=1, nodes=None, engine=None)
|
||||
|
||||
# get mesh to uv mapping for the back vertices
|
||||
m2u_body = mesh2UV(data_body.f.F, data_body.f.Ft)
|
||||
m2u_body = filter_m2u(m2u_body, set(front))
|
||||
m2u_body = convert_to_numbers(m2u_body)
|
||||
|
||||
m2u_front = mesh2UV(data_front.f.F, data_front.f.Ft)
|
||||
m2u_front = convert_to_numbers(m2u_front)
|
||||
|
||||
# create mask for filtering vertices belonging to the back
|
||||
body_front_mask = np.zeros((data_body.f.Vt.shape[0]), dtype=bool)
|
||||
body_front_mask[front] = True
|
||||
|
||||
# define which points should match after the transform on the mesh
|
||||
matching_points = np.array([ [9557, 12], [3284, 2], [3053, 5], [1888, 8], [10228, 4], [21, 7], [1209, 6], [1447, 3], [5908, 13], [4617, 15], [4219, 1], [6983, 10], [6989, 17], [7114, 9], [10602, 16], [10597, 11], [7884, 0], [8280, 14] ])
|
||||
# matching_points = np.array([ [9557, 12], [3284, 2], [3053, 5], [1888, 8], [10228, 4], [21, 7], [1209, 6], [1447, 3], [5908, 13], [4793, 15], [4219, 1], [6983, 10], [6989, 17], [7114, 9], [10602, 16], [10597, 11], [7884, 0], [8450, 14], [13566, 18], [11767, 19], [8668, 20], [5012, 21], [12716, 22], [10917, 23], [8280, 24], [4617, 25], [2760, 26], [911, 27] ])
|
||||
|
||||
# map matching point to uv coordinates
|
||||
for i in range(matching_points.shape[0]):
|
||||
matching_points[i, :] = [get_masked_index(body_front_mask, m2u_body[matching_points[i,0]]), m2u_front[matching_points[i, 1]]]
|
||||
|
||||
x = data_body.f.Vt[body_front_mask, :]
|
||||
y = uv_front
|
||||
|
||||
fig = plt.figure(figsize=(15, 15))
|
||||
ax = fig.add_subplot(121)
|
||||
# ax.set_xlim(0, 1)
|
||||
# ax.set_ylim(0, 1)
|
||||
ax.set_aspect('equal')
|
||||
ax.scatter(x[:, 0], x[:, 1])
|
||||
if y is not None:
|
||||
ax.scatter(y[:, 0], y[:, 1])
|
||||
for pair in matching_points:
|
||||
ax.scatter([x[pair[0], 0], y[pair[1], 0]], [x[pair[0], 1], y[pair[1], 1]], s=50)
|
||||
|
||||
# get transformation from body to front contour
|
||||
source = x[matching_points[:, 0], :]
|
||||
destination = y[matching_points[:, 1], :]
|
||||
trans = tps.TPS(source, destination)
|
||||
transformed = trans(data_body.f.Vt[front, :])
|
||||
|
||||
# transformed = (transformed - transformed.min()) / (transformed.max() - transformed.min())
|
||||
|
||||
ax = fig.add_subplot(122)
|
||||
# ax.set_xlim(0, 1)
|
||||
# ax.set_ylim(0, 1)
|
||||
ax.set_aspect('equal')
|
||||
ax.scatter(transformed[:, 0], transformed[:, 1])
|
||||
if y is not None:
|
||||
ax.scatter(y[:, 0], y[:, 1])
|
||||
for pair in matching_points:
|
||||
ax.scatter([transformed[pair[0], 0], y[pair[1], 0]], [transformed[pair[0], 1], y[pair[1], 1]], s=50)
|
||||
|
||||
plt.show()
|
||||
|
||||
data_front_high = np.load("custom_data/uv_map/uv_data_plane_stretched_front_high.npz", allow_pickle=True)
|
||||
y_high = data_front_high.f.Vt
|
||||
# y_high[:, 0] = y_high[:, 0] + 0.1
|
||||
# y_high[:, 1] = y_high[:, 1] + 0.42
|
||||
|
||||
closest_points = get_closest_points_normals(y_high, transformed)
|
||||
# closest_points = get_closest_point(y_high, transformed)
|
||||
unique_mask = np.zeros_like(closest_points, dtype=bool)
|
||||
unique_mask[np.unique(closest_points, return_index=True)[1]] = True
|
||||
matching_points = np.array([closest_points[unique_mask], np.array(range(y_high.shape[0]))[unique_mask]]).T
|
||||
# matching_points_2 = []
|
||||
# for i, p in enumerate(matching_points):
|
||||
# if i % 10 == 0:
|
||||
# matching_points_2.append(p)
|
||||
# matching_points = np.array(matching_points_2)
|
||||
|
||||
source_2 = transformed[matching_points[:, 0], :]
|
||||
destination_2 = y_high[matching_points[:, 1], :]
|
||||
trans_2 = tps.TPS(source_2, destination_2)
|
||||
transformed_2 = trans_2(transformed)
|
||||
|
||||
|
||||
|
||||
fig = plt.figure(figsize=(15, 15))
|
||||
ax = fig.add_subplot(121)
|
||||
ax.set_xlim(0, 1)
|
||||
ax.set_ylim(0, 1)
|
||||
ax.set_aspect('equal')
|
||||
ax.scatter(source_2[:, 0], source_2[:, 1])
|
||||
ax.scatter(destination_2[:, 0], destination_2[:, 1])
|
||||
for i, pair in enumerate(zip(source_2, destination_2)):
|
||||
if i % 5 == 0:
|
||||
s, d = pair
|
||||
ax.scatter([s[0], d[0]], [s[1], d[1]])
|
||||
|
||||
ax = fig.add_subplot(122)
|
||||
ax.set_xlim(0, 1)
|
||||
ax.set_ylim(0, 1)
|
||||
ax.set_aspect('equal')
|
||||
ax.scatter(transformed_2[:, 0], transformed_2[:, 1])
|
||||
ax.scatter(y_high[:, 0], y_high[:, 1])
|
||||
plt.show()
|
||||
|
||||
return transformed_2, mesh2UV(data_body.f.F, data_body.f.Ft), front, back
|
||||
|
||||
|
||||
|
||||
def faces2graph(F, v=None):
|
||||
if v: v = set(v) # Only this verts
|
||||
G = {}
|
||||
for f in F:
|
||||
for i in f:
|
||||
if v and i not in v: continue
|
||||
if not i in G:
|
||||
G[i] = {}
|
||||
for j in f:
|
||||
if v and j not in v: continue
|
||||
if i == j: continue
|
||||
G[i][j] = 1
|
||||
return G
|
||||
|
||||
|
||||
def custom_smoothing(mesh, edges, pinned, F, iterations=10):
|
||||
# custom_scatter(mesh[pinned, :]).show()
|
||||
|
||||
mesh_orig = mesh.copy()
|
||||
G = faces2graph(F)
|
||||
for _ in tqdm.tqdm(range(iterations)):
|
||||
edge_lengths = np.sqrt(((mesh[edges[:,0], :] - mesh[edges[:,1], :]) ** 2 ).sum(-1))
|
||||
|
||||
edge_lengths_m = np.zeros((15068, 15068), np.float32)
|
||||
edge_lengths_m[edges[:,0], edges[:,1]] = edge_lengths
|
||||
edge_lengths_m[edges[:,1], edges[:,0]] = edge_lengths
|
||||
|
||||
# mean_edge = edge_lengths.mean()
|
||||
|
||||
# dist = cdist(mesh, mesh)
|
||||
# for i in range(dist.shape[0]):
|
||||
# dist[i,i] = 10000
|
||||
|
||||
# dist = sp.spatial.KDTree(mesh)
|
||||
|
||||
# closest_neighbour = edge_lengths_m.argmin(axis=0)
|
||||
# if edge_lengths_m[i, closest_neighbour[i]] > mean_edge:
|
||||
# continue
|
||||
|
||||
mesh_temp = mesh.copy()
|
||||
for i in range(15068):
|
||||
if i in pinned:
|
||||
continue
|
||||
|
||||
# closest_neighbour = edge_lengths_m[i, :].argmin()
|
||||
# if edge_lengths_m[i, closest_neighbour] > mean_edge:
|
||||
# continue
|
||||
|
||||
neighbours = list(G[i].keys())
|
||||
neighbours_dist = edge_lengths_m[i, neighbours]
|
||||
weights = neighbours_dist / neighbours_dist.sum()
|
||||
weights = np.ones_like(weights) / weights.shape[0]
|
||||
mesh_temp[i, :] = np.sum(mesh[neighbours, :] * weights[:, np.newaxis], axis=0)
|
||||
|
||||
# neighbours_dist, neighbours = dist.query(mesh[i,:], k=10, workers=10)
|
||||
# # if neighbours[0] == 15068:
|
||||
# # continue
|
||||
# neighbours_dist, neighbours = neighbours_dist[1:], neighbours[1:]
|
||||
# weights = neighbours_dist / neighbours_dist.sum()
|
||||
# mesh_temp[i, :] = np.sum(mesh[neighbours, :] * weights[:, np.newaxis], axis=0)
|
||||
|
||||
# mesh_temp[i, :] = np.mean(mesh[neighbours, :], axis=0)
|
||||
|
||||
# print()
|
||||
|
||||
# closest_neighbour = edge_lengths_m[i, :].argmin()
|
||||
# if edge_lengths_m[i, closest_neighbour] < mean_edge:
|
||||
# vector = mesh[i,:] - mesh[closest_neighbour,:]
|
||||
# vector = vector / np.sqrt(np.sum(vector[0] ** 2 + vector[1] ** 2 ))
|
||||
# to_move = vector / 1000
|
||||
# # to_move = vector / edge_lengths_m[i, closest_neighbour] * 100
|
||||
# mesh[i, :] -= to_move
|
||||
|
||||
mesh = mesh_temp
|
||||
|
||||
|
||||
custom_scatter(mesh_orig, mesh).show()
|
||||
print()
|
||||
|
||||
return mesh
|
||||
|
||||
def custom_scatter(x, y=None):
|
||||
data = [go.Scatter(
|
||||
x=x[:,0],
|
||||
y=x[:,1],
|
||||
mode='markers',
|
||||
marker=dict(
|
||||
size=3,
|
||||
line=dict(width=0)
|
||||
),
|
||||
),]
|
||||
|
||||
if y is not None:
|
||||
data.append(
|
||||
go.Scatter(
|
||||
x=y[:,0],
|
||||
y=y[:,1],
|
||||
mode='markers',
|
||||
marker=dict(
|
||||
size=3,
|
||||
line=dict(width=0)
|
||||
),
|
||||
),)
|
||||
|
||||
fig = go.Figure(data=data)
|
||||
|
||||
# fig.update_layout(scene_aspectmode='manual',
|
||||
# scene_aspectratio=dict(x=1, y=1, z=1))
|
||||
# fig.update_layout(
|
||||
# scene = dict(
|
||||
# xaxis = dict(nticks=4, range=[-1,1],),
|
||||
# yaxis = dict(nticks=4, range=[-1,1],),
|
||||
# zaxis = dict(nticks=4, range=[-1,1],),)
|
||||
# )
|
||||
|
||||
return fig
|
||||
|
||||
|
||||
def smoothing():
|
||||
data_body = np.load("custom_data/uv_map/uv_data_smpl_2.npz", allow_pickle=True)
|
||||
uv_map = np.load('custom_data/uv_map/stretched_uv_map_3.npy')
|
||||
F = np.load("custom_utils/smpl_faces.npy").astype(np.int32)
|
||||
verts = np.concatenate((uv_map, np.ones((15068, 1))), axis=-1)
|
||||
m = tm.Trimesh(verts, data_body.f.Ft, process=False, validate=False)
|
||||
|
||||
|
||||
m2u_body_all = mesh2UV(data_body.f.F, data_body.f.Ft)
|
||||
m2u_body_all = m2uv_array(m2u_body_all)
|
||||
border_pairs = m2u_body_all[:, 0] != m2u_body_all[:, 1]
|
||||
border_pairs = m2u_body_all[border_pairs, :]
|
||||
pinned_vertices = border_pairs.reshape(-1)
|
||||
right_legs_bottom = np.nonzero(((verts[:,0] > 1).astype(np.int8) * (verts[:,0] < 1.03).astype(np.int8)).astype(bool))[0]
|
||||
left_legs_bottom = np.nonzero(verts[:,0] < 0.025)[0]
|
||||
top_arms = np.nonzero(verts[:,1] > 0.995)[0]
|
||||
bottom_arms = np.nonzero(verts[:,1] < 0.005)[0]
|
||||
pinned_vertices = np.concatenate((pinned_vertices, right_legs_bottom, left_legs_bottom, top_arms, bottom_arms))
|
||||
|
||||
x = custom_smoothing(np.asarray(m.vertices)[:,:2], np.asarray(m.edges), pinned_vertices, data_body.f.Ft)
|
||||
# m = tm.Trimesh(x, data_body.f.Ft, process=False, validate=False)
|
||||
|
||||
lap = tm.smoothing.laplacian_calculation(m, equal_weight=False, pinned_vertices=pinned_vertices)
|
||||
# tm.smoothing.filter_laplacian(m, lamb=0.5, implicit_time_integration=False, iterations=1000, volume_constraint=False, laplacian_operator=lap)
|
||||
# tm.smoothing.filter_mut_dif_laplacian(m, lamb=0.9, iterations=100, volume_constraint=True, laplacian_operator=lap)
|
||||
tm.smoothing.filter_humphrey(m, iterations=1000, laplacian_operator=lap, alpha=0.01, beta=1.0)
|
||||
# tm.smoothing.filter_taubin(m, iterations=25, lamb=0.01, nu=0.0015, laplacian_operator=lap)
|
||||
|
||||
|
||||
x = uv_map
|
||||
y = m.vertices[:, :2]
|
||||
|
||||
|
||||
|
||||
fig = go.Figure(data=[
|
||||
# go.Scatter(
|
||||
# x=x[:,0],
|
||||
# y=x[:,1],
|
||||
# mode='markers',
|
||||
# marker=dict(
|
||||
# size=3,
|
||||
# line=dict(width=0)
|
||||
# ),
|
||||
# ),
|
||||
go.Scatter(
|
||||
x=y[:,0],
|
||||
y=y[:,1],
|
||||
mode='markers',
|
||||
marker=dict(
|
||||
size=3,
|
||||
line=dict(width=0)
|
||||
),
|
||||
),
|
||||
])
|
||||
|
||||
fig.update_layout(scene_aspectmode='manual',
|
||||
scene_aspectratio=dict(x=1, y=1, z=1))
|
||||
fig.update_layout(
|
||||
scene = dict(
|
||||
xaxis = dict(nticks=4, range=[-1,1],),
|
||||
yaxis = dict(nticks=4, range=[-1,1],),
|
||||
zaxis = dict(nticks=4, range=[-1,1],),)
|
||||
)
|
||||
|
||||
fig.show()
|
||||
|
||||
dist = cdist(y, y)
|
||||
for i in range(dist.shape[0]):
|
||||
dist[i,i] = 10000
|
||||
|
||||
res = 1 / dist.min()
|
||||
|
||||
print()
|
||||
|
||||
|
||||
pass
|
||||
|
||||
|
||||
def front_3dto2d():
|
||||
# load data
|
||||
data_body = np.load("custom_data/uv_map/uv_data_smpl_2.npz", allow_pickle=True)
|
||||
data_front = np.load("custom_data/uv_map/uv_data_plane_stretched_front_low.npz", allow_pickle=True)
|
||||
|
||||
laplacian_loss = sio.loadmat("custom_utils/ssmpl0.mat")['L'].tocoo()
|
||||
edges = np.stack(np.nonzero(laplacian_loss.todense())).transpose()
|
||||
|
||||
# translate back uv to match body
|
||||
uv_front = data_front.f.V
|
||||
# uv_back[:, 0] = uv_back[:, 0] + 0.1
|
||||
# uv_back[:, 1] = uv_back[:, 1] + 0.42
|
||||
|
||||
# find front and back indices in uv
|
||||
uv_verts = np.append(data_body.f.Vt, np.zeros((data_body.f.Vt.shape[0], 1)), axis=-1)
|
||||
m = tm.Trimesh(uv_verts, data_body.f.Ft, process=False, validate=False)
|
||||
front, back = tm.graph.connected_components(m.edges, min_len=1, nodes=None, engine=None)
|
||||
|
||||
# get mesh to uv mapping for the back vertices
|
||||
m2u_body = mesh2UV(data_body.f.F, data_body.f.Ft)
|
||||
m2u_body = filter_m2u(m2u_body, set(front))
|
||||
m2u_body = convert_to_numbers(m2u_body)
|
||||
|
||||
m2u_front = mesh2UV(data_front.f.F, data_front.f.Ft)
|
||||
m2u_front = convert_to_numbers(m2u_front)
|
||||
|
||||
# create mask for filtering vertices belonging to the back
|
||||
body_front_mask = np.zeros((data_body.f.Vt.shape[0]), dtype=bool)
|
||||
body_front_mask[front] = True
|
||||
|
||||
# define which points should match after the transform on the mesh
|
||||
matching_points = np.array([ [9557, 12], [3284, 2], [3053, 5], [1888, 8], [10228, 4], [21, 7], [1209, 6], [1447, 3], [5908, 13], [4617, 15], [4219, 1], [6983, 10], [6989, 17], [7114, 9], [10602, 16], [10597, 11], [7884, 0], [8280, 14] ])
|
||||
|
||||
# for i in range(matching_points.shape[0]):
|
||||
# matching_points[i, :] = [get_masked_index(body_front_mask, matching_points[i,0]),matching_points[i, 1]]
|
||||
|
||||
|
||||
x = data_body.f.V
|
||||
# x = data_body.f.V[body_front_mask, :]
|
||||
y = uv_front
|
||||
|
||||
# get transformation from body to front contour
|
||||
source = x[matching_points[:, 0], :]
|
||||
destination = y[matching_points[:, 1], :]
|
||||
trans = TPS(source, destination)
|
||||
transformed = trans(data_body.f.V)
|
||||
|
||||
z = transformed
|
||||
fig = go.Figure(data=[
|
||||
# go.Scatter3d(
|
||||
# x=x[:,0],
|
||||
# y=x[:,1],
|
||||
# z=x[:,2],
|
||||
# ),
|
||||
go.Scatter3d(
|
||||
x=y[:,0],
|
||||
y=y[:,1],
|
||||
z=y[:,2],
|
||||
mode='markers',
|
||||
marker=dict(
|
||||
size=10,
|
||||
line=dict(width=0)
|
||||
),
|
||||
),
|
||||
go.Scatter3d(
|
||||
x=z[:,0],
|
||||
y=z[:,1],
|
||||
z=z[:,2],
|
||||
mode='markers',
|
||||
marker=dict(
|
||||
size=10,
|
||||
line=dict(width=0)
|
||||
),
|
||||
|
||||
),
|
||||
])
|
||||
|
||||
fig.update_layout(scene_aspectmode='manual',
|
||||
scene_aspectratio=dict(x=1, y=1, z=1))
|
||||
fig.update_layout(
|
||||
scene = dict(
|
||||
xaxis = dict(nticks=4, range=[-1,1],),
|
||||
yaxis = dict(nticks=4, range=[-1,1],),
|
||||
zaxis = dict(nticks=4, range=[-1,1],),)
|
||||
)
|
||||
|
||||
fig.show()
|
||||
|
||||
print()
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
# generate_mask_2()
|
||||
smoothing()
|
||||
# generate_uv_mask_border()
|
||||
|
||||
# front_uv, m2uv, front_mask, back_mask = front_3dto2d()
|
||||
|
||||
|
||||
# back_uv = match_back_to_front(front_uv)
|
||||
# # back_uv = back()
|
||||
|
||||
# final = np.zeros((15068, 2))
|
||||
# final[front_mask] = front_uv
|
||||
# final[back_mask] = back_uv + [1, 0]
|
||||
# np.save('MeshUtils\stretched_uv_map_3.npy', final)
|
||||
|
||||
pass
|
||||
@@ -0,0 +1,405 @@
|
||||
import json
|
||||
import numpy as np
|
||||
import tensorflow as tf
|
||||
import scipy as sp
|
||||
import plotly.graph_objects as go
|
||||
from plotly.subplots import make_subplots
|
||||
|
||||
from tensorflow_graphics.geometry.convolution.graph_pooling import pool, unpool
|
||||
from tensorflow_graphics.geometry.convolution.graph_convolution import edge_convolution_template
|
||||
|
||||
|
||||
def m2uv_array(m2uv, is_smpl_orig=False):
|
||||
result = np.zeros((14475 if not is_smpl_orig else 3638, 2), dtype=np.int32)
|
||||
for k, v in m2uv.items():
|
||||
if len(v) == 2:
|
||||
result[k] = list(v)
|
||||
elif len(v) == 1:
|
||||
x = v.pop()
|
||||
result[k] = [x, x]
|
||||
else:
|
||||
assert False
|
||||
|
||||
return result
|
||||
|
||||
|
||||
class UV2Mesh(tf.keras.Model):
|
||||
def __init__(self, res=256, is_skirt=False, is_smpl_orig=False, is_adaptive=False):
|
||||
super(UV2Mesh, self).__init__()
|
||||
self.res = res
|
||||
uv_map = np.load('custom_utils/stretched_uv_map_' + ('skirt' if is_skirt else '3') + ('_orig' if is_smpl_orig else '') + '.npy')
|
||||
self.uv_map_discreet = np.floor(uv_map * self.res).astype(np.int32)
|
||||
m2uv = np.load('custom_utils/m2uv' + ('_skirt' if is_skirt else '') + ('_orig' if is_smpl_orig else '') + '.npy', allow_pickle=True).item()
|
||||
m2uv = m2uv_array(m2uv, is_smpl_orig=is_smpl_orig)
|
||||
self.uv_pixels = self.uv_map_discreet[m2uv]
|
||||
self.smpl_res = 14475 if not is_smpl_orig else 3638
|
||||
self.is_adaptive = is_adaptive
|
||||
|
||||
if self.is_adaptive:
|
||||
laplacian = sp.sparse.load_npz("custom_utils/laplacian/laplacian0" + ("_skirt" if is_skirt else "") + ".npz")
|
||||
self.gc = GraphConvolution((3, 3), L=laplacian, batchSize=4)
|
||||
|
||||
|
||||
@tf.function
|
||||
def uv_to_mesh(self, uv):
|
||||
if uv.shape[1] == 3:
|
||||
uv = tf.transpose(uv, [0,2,3,1])
|
||||
if uv.shape[1] != 512:
|
||||
uv = tf.image.resize(uv, (512, 256), tf.image.ResizeMethod.LANCZOS5)
|
||||
|
||||
batch_size = uv.shape[0]
|
||||
|
||||
x1_ind = tf.concat((tf.tile(self.uv_pixels[:, 0, :][tf.newaxis, :, :], [batch_size, 1, 1]), tf.zeros((batch_size, self.smpl_res, 1), dtype=tf.int32)), axis=-1)
|
||||
y1_ind = tf.concat((tf.tile(self.uv_pixels[:, 0, :][tf.newaxis, :, :], [batch_size, 1, 1]), tf.ones((batch_size, self.smpl_res, 1), dtype=tf.int32)), axis=-1)
|
||||
z1_ind = tf.concat((tf.tile(self.uv_pixels[:, 0, :][tf.newaxis, :, :], [batch_size, 1, 1]), tf.ones((batch_size, self.smpl_res, 1), dtype=tf.int32) * 2), axis=-1)
|
||||
x2_ind = tf.concat((tf.tile(self.uv_pixels[:, 1, :][tf.newaxis, :, :], [batch_size, 1, 1]), tf.zeros((batch_size, self.smpl_res, 1), dtype=tf.int32)), axis=-1)
|
||||
y2_ind = tf.concat((tf.tile(self.uv_pixels[:, 1, :][tf.newaxis, :, :], [batch_size, 1, 1]), tf.ones((batch_size, self.smpl_res, 1), dtype=tf.int32)), axis=-1)
|
||||
z2_ind = tf.concat((tf.tile(self.uv_pixels[:, 1, :][tf.newaxis, :, :], [batch_size, 1, 1]), tf.ones((batch_size, self.smpl_res, 1), dtype=tf.int32) * 2), axis=-1)
|
||||
|
||||
|
||||
x1 = tf.gather_nd(uv, x1_ind, batch_dims=1)
|
||||
y1 = tf.gather_nd(uv, y1_ind, batch_dims=1)
|
||||
z1 = tf.gather_nd(uv, z1_ind, batch_dims=1)
|
||||
x2 = tf.gather_nd(uv, x2_ind, batch_dims=1)
|
||||
y2 = tf.gather_nd(uv, y2_ind, batch_dims=1)
|
||||
z2 = tf.gather_nd(uv, z2_ind, batch_dims=1)
|
||||
|
||||
p1 = tf.stack((x1, y1, z1), axis=-1)
|
||||
p2 = tf.stack((x2, y2, z2), axis=-1)
|
||||
|
||||
mesh = (p1 + p2) / 2
|
||||
return mesh
|
||||
|
||||
|
||||
def __call__(self, uv):
|
||||
x = self.uv_to_mesh(uv)
|
||||
if self.is_adaptive:
|
||||
x = self.gc(x)
|
||||
return x
|
||||
|
||||
|
||||
|
||||
class UV2MeshLabel():
|
||||
def __init__(self, res=256, is_skirt=False, is_smpl_orig=False):
|
||||
super(UV2MeshLabel, self).__init__()
|
||||
self.res = res
|
||||
uv_map = np.load('custom_utils/stretched_uv_map_' + ('skirt' if is_skirt else '3') + ('_orig' if is_smpl_orig else '') + '.npy')
|
||||
self.uv_map_discreet = np.floor(uv_map * self.res).astype(np.int32)
|
||||
m2uv = np.load('custom_utils/m2uv' + ('_skirt' if is_skirt else '') + ('_orig' if is_smpl_orig else '') + '.npy', allow_pickle=True).item()
|
||||
m2uv = m2uv_array(m2uv, is_smpl_orig=is_smpl_orig)
|
||||
self.uv_pixels = self.uv_map_discreet[m2uv]
|
||||
|
||||
# input = channels_first
|
||||
# @tf.function
|
||||
def uv_to_mesh(self, uv):
|
||||
batch_size = uv.shape[0]
|
||||
|
||||
x1_ind = tf.tile(self.uv_pixels[:, 0, :][tf.newaxis, :, :], [batch_size, 1, 1])
|
||||
x2_ind = tf.tile(self.uv_pixels[:, 1, :][tf.newaxis, :, :], [batch_size, 1, 1])
|
||||
|
||||
x1 = tf.gather_nd(uv, x1_ind, batch_dims=1)
|
||||
x2 = tf.gather_nd(uv, x2_ind, batch_dims=1)
|
||||
|
||||
mesh = (x1 + x2) / 2
|
||||
return mesh
|
||||
|
||||
def __call__(self, uv):
|
||||
return self.uv_to_mesh(uv)
|
||||
|
||||
|
||||
def _maskFaces(V, F, mask):
|
||||
V_mask = mask
|
||||
F_mask = np.take(V_mask, F.reshape(-1), axis=0)
|
||||
F_mask = F_mask.reshape((-1, 3)).sum(axis=1) == 3
|
||||
|
||||
V_ind = np.zeros(V_mask.shape)
|
||||
V_ind[V_mask] = np.arange(V_mask.sum())
|
||||
F_new = np.take(V_ind, F[F_mask, :].reshape(-1), axis=0).reshape((-1, 3))
|
||||
V_new = V[V_mask, :]
|
||||
|
||||
return V_new, F_new
|
||||
|
||||
|
||||
# def _maskFaces_tf(V, F, mask):
|
||||
# V_mask = mask
|
||||
# F_mask = tf.gather(V_mask, F.reshape(-1), axis=0)
|
||||
# F_mask = F_mask.reshape((-1, 3)).sum(axis=1) == 3
|
||||
# F_mask = tf.reduce_sum(tf.reshape(F_mask, (-1, 3)), axis=1) == 3
|
||||
|
||||
# V_ind = tf.zeros(V_mask.shape)
|
||||
# V_ind[V_mask] = tf.range(V_mask.sum())
|
||||
# F_new = tf.gather(V_ind, F[F_mask, :].reshape(-1), axis=0).reshape((-1, 3))
|
||||
# V_new = V[V_mask, :]
|
||||
|
||||
# return V_new, F_new
|
||||
|
||||
|
||||
def maskFacesBatch(V, F, mask):
|
||||
Vs = []
|
||||
Fs = []
|
||||
for i in range(mask.shape[0]):
|
||||
v, f = _maskFaces(V[i], F, mask[i])
|
||||
Vs.append(v)
|
||||
Fs.append(f)
|
||||
return Vs, Fs
|
||||
|
||||
|
||||
def flip_axis(x, ind=None):
|
||||
y = x.copy()
|
||||
if ind is not None:
|
||||
y[ind:,1] = x[ind:,2]
|
||||
y[ind:,2] = x[ind:,1]
|
||||
else:
|
||||
y[:,1] = x[:,2]
|
||||
y[:,2] = x[:,1]
|
||||
return y
|
||||
|
||||
|
||||
def plot_3d(mesh_gt, mesh_pred_unposed, faces, collision, collision_faces, collision_color, flip_xy=False):
|
||||
fig = make_subplots(rows=1, cols=3, start_cell="top-left", specs = [ [{ 'type': 'scene' } for _ in range(3)] ])
|
||||
|
||||
if flip_xy:
|
||||
mesh_gt = flip_axis(mesh_gt)
|
||||
mesh_pred_unposed = flip_axis(mesh_pred_unposed)
|
||||
collision = flip_axis(collision, ind=14475)
|
||||
|
||||
fig.add_trace(go.Mesh3d(
|
||||
# 8 vertices of a cube
|
||||
x=mesh_gt[:,0],
|
||||
y=mesh_gt[:,2],
|
||||
z=mesh_gt[:,1],
|
||||
# i, j and k give the vertices of triangles
|
||||
i = faces[:,0],
|
||||
j = faces[:,1],
|
||||
k = faces[:,2]
|
||||
),
|
||||
row=1, col=1)
|
||||
|
||||
fig.add_trace(go.Mesh3d(
|
||||
# 8 vertices of a cube
|
||||
x=mesh_pred_unposed[:,0],
|
||||
y=mesh_pred_unposed[:,2],
|
||||
z=mesh_pred_unposed[:,1],
|
||||
# i, j and k give the vertices of triangles
|
||||
i = faces[:,0],
|
||||
j = faces[:,1],
|
||||
k = faces[:,2]
|
||||
),
|
||||
row=1, col=2)
|
||||
|
||||
fig.add_trace(go.Mesh3d(
|
||||
# 8 vertices of a cube
|
||||
x=collision[:,0],
|
||||
y=collision[:,2],
|
||||
z=collision[:,1],
|
||||
# i, j and k give the vertices of triangles
|
||||
i = collision_faces[:,0],
|
||||
j = collision_faces[:,1],
|
||||
k = collision_faces[:,2],
|
||||
vertexcolor=collision_color
|
||||
),
|
||||
row=1, col=3)
|
||||
|
||||
|
||||
fig.update_layout(scene_aspectmode='manual', scene_aspectratio=dict(x=1, y=1, z=1),
|
||||
scene2_aspectmode='manual', scene2_aspectratio=dict(x=1, y=1, z=1),
|
||||
scene3_aspectmode='manual', scene3_aspectratio=dict(x=1, y=1, z=1),
|
||||
)
|
||||
helper = dict(xaxis = dict(nticks=4, range=[-1,1],),
|
||||
yaxis = dict(nticks=4, range=[-1,1],),
|
||||
zaxis = dict(nticks=4, range=[-1,1],),)
|
||||
fig.update_layout(scene = helper,
|
||||
scene2 = helper,
|
||||
scene3 = helper,
|
||||
)
|
||||
|
||||
return fig
|
||||
|
||||
|
||||
def plot_3d_collision(gt, collision_1, collision_2, collision_faces, collision_color, flip_xy=False):
|
||||
fig = make_subplots(rows=1, cols=3, start_cell="top-left", specs = [ [{ 'type': 'scene' } for _ in range(3)] ])
|
||||
|
||||
if flip_xy:
|
||||
collision_1 = flip_axis(collision_1, ind=14475)
|
||||
collision_2 = flip_axis(collision_2, ind=14475)
|
||||
gt = flip_axis(gt, ind=14475)
|
||||
|
||||
fig.add_trace(go.Mesh3d(
|
||||
x=gt[:,0],
|
||||
y=gt[:,2],
|
||||
z=gt[:,1],
|
||||
i = collision_faces[:,0],
|
||||
j = collision_faces[:,1],
|
||||
k = collision_faces[:,2],
|
||||
vertexcolor=collision_color
|
||||
),
|
||||
row=1, col=1)
|
||||
|
||||
fig.add_trace(go.Mesh3d(
|
||||
x=collision_1[:,0],
|
||||
y=collision_1[:,2],
|
||||
z=collision_1[:,1],
|
||||
i = collision_faces[:,0],
|
||||
j = collision_faces[:,1],
|
||||
k = collision_faces[:,2],
|
||||
vertexcolor=collision_color
|
||||
),
|
||||
row=1, col=2)
|
||||
|
||||
fig.add_trace(go.Mesh3d(
|
||||
x=collision_2[:,0],
|
||||
y=collision_2[:,2],
|
||||
z=collision_2[:,1],
|
||||
i = collision_faces[:,0],
|
||||
j = collision_faces[:,1],
|
||||
k = collision_faces[:,2],
|
||||
vertexcolor=collision_color
|
||||
),
|
||||
row=1, col=3)
|
||||
|
||||
|
||||
fig.update_layout(scene_aspectmode='manual', scene_aspectratio=dict(x=1, y=1, z=1),
|
||||
scene2_aspectmode='manual', scene2_aspectratio=dict(x=1, y=1, z=1),
|
||||
scene3_aspectmode='manual', scene3_aspectratio=dict(x=1, y=1, z=1),
|
||||
)
|
||||
helper = dict(xaxis = dict(nticks=4, range=[-1,1],),
|
||||
yaxis = dict(nticks=4, range=[-1,1],),
|
||||
zaxis = dict(nticks=4, range=[-1,1],),)
|
||||
fig.update_layout(scene = helper,
|
||||
scene2 = helper,
|
||||
scene3 = helper,
|
||||
)
|
||||
|
||||
return fig
|
||||
|
||||
|
||||
def batch_to_dict(batch, signature):
|
||||
result = dict()
|
||||
for i, key in enumerate(signature.keys()):
|
||||
result[key] = batch[i]
|
||||
return result
|
||||
|
||||
conditions = ['pose_shape', 'uv_static', 'uv_mask', 'uv_shaped_body']
|
||||
def nest_conditions(batch):
|
||||
batch['conditions'] = {d: batch.pop(d) for d in conditions}
|
||||
return batch
|
||||
|
||||
|
||||
def sparse_to_tensor(L):
|
||||
idx = np.mat([L.row, L.col]).transpose().astype(np.float32)
|
||||
data = L.data.astype(dtype=np.float32)
|
||||
shape = L.shape
|
||||
return tf.SparseTensor(idx, data, shape)
|
||||
|
||||
|
||||
def load_model(model, checkpoint, disable_assert=False, assert_consumed_on=False, iteration=''):
|
||||
ckpt = tf.train.Checkpoint(model=model)
|
||||
manager = tf.train.CheckpointManager(ckpt, checkpoint, max_to_keep=1)
|
||||
if not disable_assert:
|
||||
assert manager.latest_checkpoint is not None, checkpoint
|
||||
assert manager.latest_checkpoint.split('/')[-1] != 'ckpt-0', checkpoint
|
||||
if assert_consumed_on:
|
||||
ckpt.restore(manager.latest_checkpoint).assert_consumed()
|
||||
else:
|
||||
if iteration != '':
|
||||
name_base = manager.latest_checkpoint.split('-')[0]
|
||||
ckpt_ind = [i for i in range(len(manager.checkpoints)) if manager.checkpoints[i].endswith(name_base + '-' + iteration)][0]
|
||||
ckpt.restore(manager.checkpoints[ckpt_ind]).assert_existing_objects_matched()#.assert_consumed()
|
||||
else:
|
||||
ckpt.restore(manager.latest_checkpoint).assert_existing_objects_matched()#.assert_consumed()
|
||||
|
||||
|
||||
def get_summary_str(model):
|
||||
lines = []
|
||||
model.summary(print_fn=lines.append, line_length=150,expand_nested=True,show_trainable=True)
|
||||
# Add initial spaces to avoid markdown formatting in TensorBoard
|
||||
return lines
|
||||
|
||||
def get_summary_html(model):
|
||||
lines = get_summary_str(model)
|
||||
html = ''
|
||||
for line in lines:
|
||||
html += "<pre>" + line + "</pre>"
|
||||
return html
|
||||
|
||||
|
||||
def repeat_laplacian(L, times):
|
||||
idx = np.repeat(np.mat([L.row, L.col]), repeats=times, axis=1)
|
||||
idx_offset = np.linspace(0, L.shape[0] * (times-1), times).astype(np.int32)
|
||||
idx_offset = np.repeat(idx_offset, L.row.shape[0], axis=0)
|
||||
idx += idx_offset[np.newaxis,:]
|
||||
data = np.repeat(L.data, repeats=times, axis=0)
|
||||
shape = (L.shape[0] * times, L.shape[1] * times)
|
||||
return sp.sparse.coo_matrix((data, (np.squeeze(np.asarray(idx[0,:])), np.squeeze(np.asarray(idx[1,:])))), shape)
|
||||
|
||||
|
||||
def merge_verts(Vs, batchDim):
|
||||
V = tf.concat(tf.unstack(Vs, num=batchDim), axis=0)
|
||||
return V
|
||||
|
||||
|
||||
class GraphConvolution(tf.keras.layers.Layer):
|
||||
def __init__(self, shape, activation=None, name='gcn', L=None, batchSize=None, normalization=None):
|
||||
super(GraphConvolution, self).__init__()
|
||||
self.shape = shape
|
||||
self.w = self.add_weight(name=name + '/weight', shape=(2, *shape), initializer=tf.initializers.glorot_normal, trainable=True)
|
||||
self.b = self.add_weight(name=name + '/bias', shape=shape[-1], initializer=tf.initializers.zeros, trainable=True)
|
||||
self.activation = activation or (lambda x: x)
|
||||
self.normalization = normalization
|
||||
self.batchSize = batchSize
|
||||
self.indices = tf.cast(tf.linspace(0, self.batchSize * L.shape[1], self.batchSize+1), dtype=tf.int32)
|
||||
self.L = repeat_laplacian(L, self.batchSize)
|
||||
|
||||
def gather(self):
|
||||
return [self.w, self.b]
|
||||
|
||||
def _single(self, X):
|
||||
X0 = tf.einsum('ab,bc->ac', X, self.w[0]) # Node
|
||||
X1 = tf.einsum('ab,bc->ac', X, self.w[1]) # Neigh
|
||||
# Neigh conv
|
||||
X1 = edge_convolution_template(
|
||||
X1, sparse_to_tensor(self.L), sizes=None,
|
||||
edge_function=lambda x, y: y,
|
||||
reduction='weighted',
|
||||
edge_function_kwargs={}
|
||||
)
|
||||
X = X0 + X1 + self.b
|
||||
return X
|
||||
|
||||
def call(self, X, training):
|
||||
X = merge_verts(X, self.batchSize)
|
||||
X = self._single(X)
|
||||
X = tf.stack([X[self.indices[i-1]:self.indices[i]] for i in range(1, len(self.indices))])
|
||||
X = self.activation(X)
|
||||
if self.normalization:
|
||||
X = self.normalization(X, training=training)
|
||||
return X
|
||||
|
||||
|
||||
def replace_arrays(string):
|
||||
out = string[:1]
|
||||
i = 0
|
||||
while i < len(string) and string[i+1:].find('[') != -1:
|
||||
start = string[i+1:].find('[')
|
||||
end = string[i+1+start:].find(']')
|
||||
out += string[i+1:i+1+start]
|
||||
array = string[i+1+start:i+1+start+end+1].replace('\n', '')
|
||||
out += array
|
||||
i = i + 1 + start + end
|
||||
|
||||
out += string[i+1:]
|
||||
|
||||
return out
|
||||
|
||||
|
||||
def group_inputs_for_levels(batch, config):
|
||||
levels_input = []
|
||||
for res in config['pyramid']['resolutions']:
|
||||
res_modifier = '_' + str(res) if res != 512 else ''
|
||||
if config['use_skinning']:
|
||||
uv_orig = batch['uv_unposed' + res_modifier]
|
||||
else:
|
||||
uv_orig = batch['uv_posed' + res_modifier]
|
||||
levels_input.append([uv_orig, batch['uv_static' + res_modifier], batch['uv_body_posed' + res_modifier], batch['uv_curvature' + res_modifier], batch['uv_mask' + res_modifier], batch['uv_mask']])
|
||||
conds = [batch['uv_static'], batch['uv_body_posed'], batch['uv_curvature']]
|
||||
return levels_input, conds
|
||||
@@ -0,0 +1,123 @@
|
||||
import os
|
||||
import tqdm
|
||||
import scipy.interpolate as spi
|
||||
import numpy as np
|
||||
import plotly.graph_objects as go
|
||||
import plotly.express as ex
|
||||
|
||||
os.environ["CUDA_VISIBLE_DEVICES"] = "1"
|
||||
import tensorflow as tf
|
||||
import tensorflow_addons as tfa
|
||||
def fill(img, mask):
|
||||
coords = np.transpose(np.stack(np.nonzero((1 - mask))), (1,0))
|
||||
img_idx = np.transpose(np.stack(np.nonzero((mask))), (1,0))
|
||||
img_data = img[img_idx[:,0], img_idx[:,1], :]
|
||||
interpolator = spi.NearestNDInterpolator(x=img_idx, y=img_data)
|
||||
img_2 = interpolator(coords)
|
||||
res = np.zeros((512,256,3), dtype=np.float32)
|
||||
res[coords[:,0], coords[:,1], :] = img_2
|
||||
res[img_idx[:,0], img_idx[:,1], :] = img_data
|
||||
|
||||
return res
|
||||
|
||||
# @tf.function
|
||||
def fill_tf(img, mask):
|
||||
img = img[tf.newaxis,:,:,:]
|
||||
coords_int = tf.transpose(tf.stack(tf.experimental.numpy.nonzero(1-mask)), (1,0))
|
||||
coords = tf.cast(coords_int, dtype=tf.float32)
|
||||
img_2 = tfa.image.interpolate_bilinear(img, coords[tf.newaxis,:,:])
|
||||
# coords_2 = list(zip(tf.ones(coords_int.shape[0]), coords_int[:,0], coords_int[:,1]))
|
||||
# coords = tf.transpose(tf.stack([coords_int[:,0], coords_int[:,1]]))
|
||||
out = tf.tensor_scatter_nd_update(img[0], coords_int, img_2[0])
|
||||
|
||||
return out
|
||||
|
||||
|
||||
|
||||
def fill_label(img, mask):
|
||||
coords = np.transpose(np.stack(np.nonzero((1 - mask))), (1,0))
|
||||
img_idx = np.transpose(np.stack(np.nonzero((mask))), (1,0))
|
||||
img_data = img[img_idx[:,0], img_idx[:,1]]
|
||||
|
||||
interpolator = spi.NearestNDInterpolator(x=img_idx, y=img_data)
|
||||
img_2 = interpolator(coords)
|
||||
res = np.zeros((512,256), dtype=np.float32)
|
||||
res[coords[:,0], coords[:,1]] = img_2
|
||||
res[img_idx[:,0], img_idx[:,1]] = img_data
|
||||
|
||||
return res
|
||||
|
||||
|
||||
def fill_uv_maps_posed_body(folder, type):
|
||||
os.makedirs(os.path.join(folder, type + '_filled'), exist_ok=True)
|
||||
mask = np.load('custom_utils/mask_256x512.npy')
|
||||
|
||||
frames = os.listdir(os.path.join(folder, type))
|
||||
|
||||
for frame in tqdm.tqdm(frames):
|
||||
id, _ = frame.split('_')
|
||||
# mask = np.load(os.path.join(folder, 'uv_mask', id + '.npz')); mask = mask['uv_mask']
|
||||
img = np.load(os.path.join(folder, type, frame)); img = img['uv_posed_body_no_orient']
|
||||
# filled = fill(img, mask)
|
||||
img = np.zeros((100, 100, 3), dtype=np.float32)
|
||||
mask = np.zeros((100, 100), dtype=np.int32)
|
||||
img[25:75, 25:75, :] = 1
|
||||
mask[25:75, 25:75] = 1
|
||||
|
||||
filled = fill_tf(img, mask)
|
||||
np.savez_compressed(os.path.join(folder, type + '_filled', frame), uv_posed_body_no_orient_filled = filled)
|
||||
|
||||
|
||||
def fill_uv_maps_curv(folder, type):
|
||||
os.makedirs(os.path.join(folder, type + '_filled'), exist_ok=True)
|
||||
mask = np.load('custom_utils/mask_256x512.npy')
|
||||
|
||||
frames = os.listdir(os.path.join(folder, type))
|
||||
|
||||
for frame in tqdm.tqdm(frames):
|
||||
id, _ = frame.split('_')
|
||||
mask = np.load(os.path.join(folder, 'uv_mask', id + '.npz')); mask = mask['uv_mask']
|
||||
img = np.load(os.path.join(folder, type, frame)); img = img['curvature_uv']
|
||||
filled = fill_label(img, mask)
|
||||
np.savez_compressed(os.path.join(folder, type + '_filled', frame), uv_curvature_filled = filled)
|
||||
|
||||
|
||||
|
||||
def fill_uv_maps_unposed(folder, type):
|
||||
os.makedirs(os.path.join(folder, type + '_filled'), exist_ok=True)
|
||||
mask = np.load('custom_utils/mask_256x512.npy')
|
||||
|
||||
frames = os.listdir(os.path.join(folder, type))
|
||||
|
||||
for frame in tqdm.tqdm(frames):
|
||||
id, _ = frame.split('_')
|
||||
mask = np.load(os.path.join(folder, 'uv_mask', id + '.npz')); mask = mask['uv_mask']
|
||||
img = np.load(os.path.join(folder, type, frame)); img = img['uv_unposed_no_shape']
|
||||
filled = fill(img, mask)
|
||||
np.savez_compressed(os.path.join(folder, type + '_filled', frame), uv_unposed_no_shape_filled = filled)
|
||||
|
||||
|
||||
|
||||
def fill_uv_maps_static(folder, type):
|
||||
os.makedirs(os.path.join(folder, type + '_filled'), exist_ok=True)
|
||||
mask = np.load('custom_utils/mask_256x512.npy')
|
||||
|
||||
frames = os.listdir(os.path.join(folder, type))
|
||||
|
||||
for frame in tqdm.tqdm(frames):
|
||||
id, _ = frame.split('.')
|
||||
mask = np.load(os.path.join(folder, 'uv_mask', id + '.npz')); mask = mask['uv_mask']
|
||||
img = np.load(os.path.join(folder, type, frame)); img = img['uv_static']
|
||||
filled = fill(img, mask)
|
||||
np.savez_compressed(os.path.join(folder, type + '_filled', frame), uv_static_filled = filled)
|
||||
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
# fill_uv_maps_static(folder='/home/workspace/local-data/train_100', type='uv_static')
|
||||
# fill_uv_maps_unposed(folder='/home/workspace/local-data/train_100', type='uv_unposed_no_shape')
|
||||
# fill_uv_maps_curv(folder='/home/workspace/local-data/train_100', type='uv_curvature')
|
||||
fill_uv_maps_posed_body(folder='/home/workspace/local-data/train_100', type='uv_posed_body_no_orient')
|
||||
|
||||
|
||||
print()
|
||||
@@ -0,0 +1,290 @@
|
||||
import os
|
||||
import sys
|
||||
import scipy as sp
|
||||
import numpy as np
|
||||
import scipy.io as sio
|
||||
import tensorflow as tf
|
||||
from tensorflow.keras.activations import tanh, relu
|
||||
from tensorflow.keras.layers import Dense, BatchNormalization, Add, Flatten
|
||||
|
||||
sys.path.append(os.path.dirname(os.path.abspath(__file__)) + '/..')
|
||||
# from custom_utils.custom_utils import UV2Mesh
|
||||
|
||||
from tensorflow_graphics.geometry.convolution.graph_pooling import pool, unpool
|
||||
from tensorflow_graphics.geometry.convolution.graph_convolution import edge_convolution_template
|
||||
|
||||
|
||||
def sparse_to_tensor(L):
|
||||
idx = np.mat([L.row, L.col]).transpose().astype(np.float32)
|
||||
data = L.data.astype(dtype=np.float32)
|
||||
shape = L.shape
|
||||
return tf.SparseTensor(idx, data, shape)
|
||||
|
||||
def merge_verts(Vs, batchDim):
|
||||
V = tf.concat(tf.unstack(Vs, num=batchDim), axis=0)
|
||||
return V
|
||||
|
||||
|
||||
def repeat_laplacian(L, times):
|
||||
idx = np.repeat(np.mat([L.row, L.col]), repeats=times, axis=1)
|
||||
idx_offset = np.linspace(0, L.shape[0] * (times-1), times).astype(np.int32)
|
||||
idx_offset = np.repeat(idx_offset, L.row.shape[0], axis=0)
|
||||
idx += idx_offset[np.newaxis,:]
|
||||
data = np.repeat(L.data, repeats=times, axis=0)
|
||||
shape = (L.shape[0] * times, L.shape[1] * times)
|
||||
return sp.sparse.coo_matrix((data, (np.squeeze(np.asarray(idx[0,:])), np.squeeze(np.asarray(idx[1,:])))), shape)
|
||||
|
||||
|
||||
def convert_pooling(pooling, batchSize):
|
||||
pool_map = np.zeros((np.max(pooling)+1, pooling.shape[0]), dtype=np.float32)
|
||||
for i in range(pooling.shape[0]):
|
||||
pool_map[pooling[i], i] = 1.0
|
||||
pool_map = np.tile(pool_map[np.newaxis, :, :], [batchSize, 1, 1])
|
||||
pool_map = tf.sparse.from_dense(pool_map)
|
||||
return pool_map
|
||||
|
||||
|
||||
class GraphConvolution(tf.keras.layers.Layer):
|
||||
def __init__(self, shape, activation=None, name='gcn', L=None, batchSize=None, normalization=None):
|
||||
super(GraphConvolution, self).__init__()
|
||||
self.shape = shape
|
||||
self.w = self.add_weight(name=name + '/weight', shape=(2, *shape), initializer=tf.initializers.glorot_normal, trainable=True)
|
||||
self.b = self.add_weight(name=name + '/bias', shape=shape[-1], initializer=tf.initializers.zeros, trainable=True)
|
||||
self.activation = activation or (lambda x: x)
|
||||
self.normalization = normalization
|
||||
self.batchSize = batchSize
|
||||
self.indices = tf.cast(tf.linspace(0, self.batchSize * L.shape[1], self.batchSize+1), dtype=tf.int32)
|
||||
self.L = repeat_laplacian(L, self.batchSize)
|
||||
|
||||
def gather(self):
|
||||
return [self.w, self.b]
|
||||
|
||||
def _single(self, X):
|
||||
X0 = tf.einsum('ab,bc->ac', X, self.w[0]) # Node
|
||||
X1 = tf.einsum('ab,bc->ac', X, self.w[1]) # Neigh
|
||||
# Neigh conv
|
||||
X1 = edge_convolution_template(
|
||||
X1, sparse_to_tensor(self.L), sizes=None,
|
||||
edge_function=lambda x, y: y,
|
||||
reduction='weighted',
|
||||
edge_function_kwargs={}
|
||||
)
|
||||
X = X0 + X1 + self.b
|
||||
return X
|
||||
|
||||
def call(self, X, training):
|
||||
X = merge_verts(X, self.batchSize)
|
||||
X = self._single(X)
|
||||
X = tf.stack([X[self.indices[i-1]:self.indices[i]] for i in range(1, len(self.indices))])
|
||||
X = self.activation(X)
|
||||
if self.normalization:
|
||||
X = self.normalization(X, training=training)
|
||||
return X
|
||||
|
||||
|
||||
class GraphUnPool(tf.keras.layers.Layer):
|
||||
def __init__(self, pool_map):
|
||||
super(GraphUnPool, self).__init__()
|
||||
self.pool_map = pool_map
|
||||
|
||||
def call(self, X):
|
||||
# print(X.dtype)
|
||||
# print(self.pool_map.dtype)
|
||||
return unpool(X, self.pool_map, None)
|
||||
# return tf.gather(X, self.pool_map, axis=1)
|
||||
|
||||
|
||||
|
||||
class GraphPool(tf.keras.layers.Layer):
|
||||
def __init__(self, pool_map, type, batch_size):
|
||||
super(GraphPool, self).__init__()
|
||||
self.type = type
|
||||
self.batch_size = batch_size
|
||||
self.indices = tf.cast(tf.linspace(0, batch_size * pool_map.shape[0], batch_size+1), dtype=tf.int32)
|
||||
# self.pool_map = repeat_laplacian(pool_map, self.batch_size)
|
||||
self.pool_map = np.squeeze(pool_map)
|
||||
|
||||
def _single(self, x):
|
||||
return pool(x, sparse_to_tensor(self.pool_map), algorithm=self.type, sizes=None)
|
||||
|
||||
def call(self, x, training=None):
|
||||
# x = merge_verts(x, self.batch_size)
|
||||
# x = self._single(x)
|
||||
# x = tf.stack([x[self.indices[i-1]:self.indices[i]] for i in range(1, len(self.indices))])
|
||||
|
||||
N, M = len(self.pool_map), len(set(self.pool_map)) # N = HQ verts, M = LQ verts
|
||||
# Assign correspondences as V0_lq = [V0_hq, V1_hq, ..., VN_hq]
|
||||
d = [[] for _ in range(M)] # list of lists
|
||||
for i in range(N): d[self.pool_map[i]] += [i]
|
||||
# max sublist len
|
||||
m = max([len(d[i]) for i in range(M)])
|
||||
# expand all sublist to max len
|
||||
for i in range(M):
|
||||
while len(d[i]) < m:
|
||||
# repeating elements will NOT change MIN/MAX output
|
||||
d[i] += [d[i][0]]
|
||||
|
||||
x = tf.reduce_max(tf.gather(x, d, axis=1), axis=2)
|
||||
return x
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
class GNNDiscriminator(tf.keras.Model):
|
||||
def __init__(self, batchSize, isSkirt=False, n_feature=256, pool_type="max", global_average=True):
|
||||
super(GNNDiscriminator, self).__init__()
|
||||
|
||||
self.uv2mesh = UV2Mesh(res=256)
|
||||
|
||||
laplacians = [sio.loadmat("custom_utils/laplacian/ssmpl" + ("_skirt" if isSkirt else "") + str(i) + ".mat")['L'].tocoo() for i in range(5)]
|
||||
poolings = [sio.loadmat("custom_utils/pooling/pool" + str(i) + str(i+1) + ".mat")['idx'] for i in range(5)]
|
||||
|
||||
self.conv1 = GraphConvolution((3, int(n_feature/2**2)), activation=relu, L=laplacians[0], batchSize=batchSize, normalization=BatchNormalization(), name='gcn_1')
|
||||
self.pool1 = GraphPool(pool_map=poolings[0], type=pool_type, batch_size=batchSize)
|
||||
|
||||
self.conv2 = GraphConvolution((int(n_feature/2**2), int(n_feature/2**2)), activation=relu, L=laplacians[1], batchSize=batchSize, normalization=BatchNormalization(), name='gcn_2')
|
||||
self.pool2 = GraphPool(pool_map=poolings[1], type=pool_type, batch_size=batchSize)
|
||||
|
||||
self.custom_layers = []
|
||||
for i in range(2,5):
|
||||
self.custom_layers.append(GraphConvolution((int(n_feature/2**(4- (i if i == 2 else i-1) )), int(n_feature/2**(4-i))), activation=relu, L=laplacians[i], batchSize=batchSize, normalization=BatchNormalization()))
|
||||
for j in range(2**(i-2)):
|
||||
conv1 = GraphConvolution((int(n_feature/2**(4-i)), int(n_feature/2**(4-i))), activation=relu, L=laplacians[i], batchSize=batchSize, normalization=BatchNormalization())
|
||||
conv2 = GraphConvolution((int(n_feature/2**(4-i)), int(n_feature/2**(4-i))), activation=relu, L=laplacians[i], batchSize=batchSize, normalization=BatchNormalization())
|
||||
self.custom_layers.append([conv1, conv2])
|
||||
self.custom_layers.append(GraphPool(pool_map=poolings[i], type=pool_type, batch_size=batchSize))
|
||||
|
||||
self.custom_layers.append(Dense(256, relu))
|
||||
self.custom_layers.append(BatchNormalization())
|
||||
self.custom_layers.append(Dense(int(n_feature/2**(4-i)), relu))
|
||||
self.custom_layers.append(BatchNormalization())
|
||||
|
||||
self.flatten = Flatten()
|
||||
self.last = Dense(1)
|
||||
|
||||
def call(self, inputs, training=None):
|
||||
inputs = self.uv2mesh(inputs)
|
||||
|
||||
x = self.conv1(inputs, training)
|
||||
x = self.pool1(x)
|
||||
|
||||
x = self.conv2(x, training)
|
||||
x = self.pool2(x)
|
||||
|
||||
for layer in self.custom_layers:
|
||||
if isinstance(layer, list):
|
||||
x1 = layer[0](x, training)
|
||||
x2 = layer[1](x1, training)
|
||||
x = Add()([x, x2])
|
||||
else:
|
||||
x = layer(x, training=training)
|
||||
|
||||
x = self.flatten(x)
|
||||
x = self.last(x)
|
||||
|
||||
return x
|
||||
|
||||
|
||||
class Reconstructor_V4_Decoder(tf.keras.Model):
|
||||
def __init__(self, batchSize, isSkirt=False, n_feature=64):
|
||||
super(Reconstructor_V4_Decoder, self).__init__()
|
||||
|
||||
|
||||
laplacians = [sio.loadmat("custom_utils/laplacian/ssmpl" + ("_skirt" if isSkirt else "") + str(i) + ".mat")['L'].tocoo() for i in range(5)]
|
||||
poolings = [sio.loadmat("custom_utils/pooling/pool" + str(i) + str(i+1) + ".mat")['idx'] for i in range(5)]
|
||||
poolings = [convert_pooling(poolings[i][0,:], batchSize) for i in range(len(poolings))]
|
||||
|
||||
self.custom_layers = []
|
||||
for i in range(4,1,-1):
|
||||
self.custom_layers.append(GraphUnPool(pool_map=poolings[i]))
|
||||
self.custom_layers.append(GraphConvolution((int(n_feature/2**(4-(i if i == 4 else i+1))), int(n_feature/2**(4-i)) ), activation=relu, L=laplacians[i], batchSize=batchSize, normalization=BatchNormalization()))
|
||||
for _ in range(2**(i-2)):
|
||||
conv1 = GraphConvolution((int(n_feature/2**(4-i)), int(n_feature/2**(4-i))), activation=relu, L=laplacians[i], batchSize=batchSize, normalization=BatchNormalization())
|
||||
conv2 = GraphConvolution((int(n_feature/2**(4-i)), int(n_feature/2**(4-i))), activation=relu, L=laplacians[i], batchSize=batchSize, normalization=BatchNormalization())
|
||||
self.custom_layers.append([conv1, conv2])
|
||||
|
||||
self.pool2 = GraphUnPool(pool_map=poolings[1])
|
||||
self.conv2 = GraphConvolution((int(n_feature/2**2), int(n_feature/2**2)), activation=relu, L=laplacians[1], batchSize=batchSize, normalization=BatchNormalization(), name='gcn_2')
|
||||
|
||||
self.pool1 = GraphUnPool(pool_map=poolings[0])
|
||||
self.conv1 = GraphConvolution((int(n_feature/2**2), int(n_feature/2**3)), activation=relu, L=laplacians[0], batchSize=batchSize, normalization=BatchNormalization(), name='gcn_1')
|
||||
|
||||
self.conv0 = GraphConvolution((int(n_feature/2**3), 3), L=laplacians[0], batchSize=batchSize, name='gcn_0')
|
||||
|
||||
def call(self, inputs, training=None):
|
||||
x = inputs
|
||||
|
||||
for layer in self.custom_layers:
|
||||
if isinstance(layer, list):
|
||||
x1 = layer[0](x, training)
|
||||
x2 = layer[1](x1, training)
|
||||
x = Add()([x, x2])
|
||||
else:
|
||||
x = layer(x, training=training)
|
||||
|
||||
x = self.pool2(x)
|
||||
x = self.conv2(x, training)
|
||||
|
||||
x = self.pool1(x)
|
||||
x = self.conv1(x, training)
|
||||
|
||||
x = self.conv0(x, training)
|
||||
|
||||
return x
|
||||
|
||||
|
||||
|
||||
class Reconstructor_v2(tf.keras.Model):
|
||||
def __init__(self, batchSize, isSkirt=False):
|
||||
super(Reconstructor_v2, self).__init__()
|
||||
|
||||
self.flatten = tf.keras.layers.Flatten()
|
||||
|
||||
self.fc2 = tf.keras.layers.Dense(14475*3, kernel_regularizer=tf.keras.regularizers.l2)
|
||||
|
||||
self.reshape2 = tf.keras.layers.Reshape((14475, 3))
|
||||
# self.drop = Dropout(0.2)
|
||||
# self.fc3 = Dense(512)
|
||||
self.fc1 = tf.keras.layers.Dense(3)
|
||||
|
||||
def call(self, inputs, training=None):
|
||||
x = self.flatten(inputs)
|
||||
x = self.fc2(x)
|
||||
# x = self.drop(x, training=training)
|
||||
x = self.reshape2(x)
|
||||
# x = self.fc3(x)
|
||||
x = self.fc1(x)
|
||||
return x
|
||||
|
||||
|
||||
class DatasetTest(tf.data.Dataset):
|
||||
|
||||
def _generator():
|
||||
for i in range(10):
|
||||
yield np.array([i], dtype=np.int32)
|
||||
|
||||
def __new__(cls):
|
||||
return tf.data.Dataset.from_generator(
|
||||
cls._generator,
|
||||
output_signature = tf.TensorSpec(shape = (1), dtype = tf.int32) ,
|
||||
args=()
|
||||
)
|
||||
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
import os
|
||||
os.environ["CUDA_VISIBLE_DEVICES"] = "2"
|
||||
|
||||
batch_size = 4
|
||||
# input = tf.ones((batch_size, 14475, 3))
|
||||
|
||||
# ge = GNNDiscriminator(batchSize=batch_size, n_feature=32)
|
||||
|
||||
# x = ge(input)
|
||||
|
||||
gd = Reconstructor_V4_Decoder(batch_size)
|
||||
x = gd(tf.ones((batch_size, 21, 64)))
|
||||
|
||||
print()
|
||||
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
@@ -0,0 +1,140 @@
|
||||
import os
|
||||
from socket import close
|
||||
import sys
|
||||
import tqdm
|
||||
import math
|
||||
import time
|
||||
import random
|
||||
import shutil
|
||||
import collections
|
||||
import argparse
|
||||
import threading
|
||||
import numpy as np
|
||||
import scipy as sp
|
||||
import trimesh as tm
|
||||
import cmapy as cm
|
||||
import cv2
|
||||
import plotly.express as ex
|
||||
import plotly.graph_objects as go
|
||||
|
||||
|
||||
def m2uv_array(m2uv):
|
||||
result = np.zeros((14475, 2), dtype=np.int32)
|
||||
for k, v in m2uv.items():
|
||||
if len(v) == 2:
|
||||
result[k] = list(v)
|
||||
elif len(v) == 1:
|
||||
x = v.pop()
|
||||
result[k] = [x, x]
|
||||
else:
|
||||
assert False
|
||||
|
||||
return result
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
F = np.load("custom_utils/smpl_faces.npy")
|
||||
uv_map = np.load('custom_utils/stretched_uv_map_3.npy')
|
||||
m2uv = np.load('custom_utils/m2uv.npy', allow_pickle=True).item()
|
||||
uv_mask = np.load('custom_utils/mask_256x512.npy')
|
||||
|
||||
left = [10602, 14090, 10601, 14299, 10623, 14102, 8739, 13152, 8738, 14292, 10737, 14244, 10738, 14288, 10736, 14243, 8728, 13149, 8730, 13982, 10419, 14241, 8721, 13147, 8723, 14239, 8717, 13145, 8716, 14237, 8711, 13143, 8710, 14235, 10734, 14283, 10729, 14226, 10381, 13948, 10383, 14225, 10725, 14222, 8620, 13107, 8622, 14305, 10761, 14293, 10732, 14233, 10720, 14214, 10721, 14217, 7114]
|
||||
right = list(reversed([12418, 7113, 12415, 7112, 12434, 7125, 12494, 7155, 12506, 4964, 11308, 4963, 12423, 7118, 12426, 6752, 12149, 6751, 12427, 7122, 12484, 7127, 12436, 5054, 11344, 5056, 12438, 5060, 11346, 5062, 12440, 5067, 11348, 5066, 12442, 6788, 12183, 5074, 11350, 5072, 12444, 7129, 12489, 7131, 12445, 7130, 12493, 5082, 11353, 5084, 12303, 7009, 12500, 6987, 12291, 6989]))
|
||||
corners = np.array([[0,0],[0,1],[2,0],[2,1]])
|
||||
|
||||
# left_orig = []
|
||||
# for l in left:
|
||||
# left_orig.append(list(m2uv[l])[0])
|
||||
# left_orig.append(list(m2uv[l])[1])
|
||||
# ex.scatter(x=np.concatenate((uv_map[left_orig, 0], corners[:,0]), axis=-1), y=np.concatenate((uv_map[left_orig, 1], corners[:,1]), axis=-1)).show()
|
||||
|
||||
|
||||
# left_1 = [[0,0],[0,1],[2,0],[2,1]]
|
||||
# right_1 = [[0,0],[0,1],[2,0],[2,1]]
|
||||
# left_2 = [[0,0],[0,1],[2,0],[2,1]]
|
||||
# right_2 = [[0,0],[0,1],[2,0],[2,1]]
|
||||
left_1 = []
|
||||
right_1 = []
|
||||
left_2 = []
|
||||
right_2 = []
|
||||
for l in left:
|
||||
print(list(m2uv[l])[0])
|
||||
if uv_map[list(m2uv[l])[0], 0] < 1:
|
||||
left_1.append(list(m2uv[l])[0])
|
||||
else:
|
||||
left_2.append(list(m2uv[l])[0])
|
||||
if uv_map[list(m2uv[l])[1], 0] < 1:
|
||||
left_1.append(list(m2uv[l])[1])
|
||||
else:
|
||||
left_2.append(list(m2uv[l])[1])
|
||||
left_1 = np.array(left_1)[:56]
|
||||
left_2 = np.array(left_2)[:56]
|
||||
|
||||
# ex.scatter(x=uv_map[:, 0], y=uv_map[:, 1]).show()
|
||||
|
||||
for r in right:
|
||||
print(list(m2uv[r])[0])
|
||||
if uv_map[list(m2uv[r])[0], 0] < 1:
|
||||
right_1.append(list(m2uv[r])[0])
|
||||
else:
|
||||
right_2.append(list(m2uv[r])[0])
|
||||
if uv_map[list(m2uv[r])[1], 0] < 1:
|
||||
right_1.append(list(m2uv[r])[1])
|
||||
else:
|
||||
right_2.append(list(m2uv[r])[1])
|
||||
right_1 = np.array(right_1)
|
||||
right_2 = np.array(right_2)
|
||||
|
||||
c = np.ones((60, 3), dtype=np.int32) * 50
|
||||
c[5, :] = 200
|
||||
# ex.scatter(x=np.concatenate((uv_map[left_1, 0], corners[:,0]), axis=-1), y=np.concatenate((uv_map[left_1, 1], corners[:,1]), axis=-1), color=c[:,0]).show()
|
||||
# ex.scatter(x=np.concatenate((uv_map[right_1, 0], corners[:,0]), axis=-1), y=np.concatenate((uv_map[right_1, 1], corners[:,1]), axis=-1), color=c[:,0]).show()
|
||||
# # ex.scatter(x=uv_map[right_1, 0], y=uv_map[right_1, 1], color=c[:,0]).show()
|
||||
|
||||
# ex.scatter(x=np.concatenate((uv_map[left_2, 0], corners[:,0]), axis=-1), y=np.concatenate((uv_map[left_2, 1], corners[:,1]), axis=-1), color=c[:,0]).show()
|
||||
# ex.scatter(x=np.concatenate((uv_map[right_2, 0], corners[:,0]), axis=-1), y=np.concatenate((uv_map[right_2, 1], corners[:,1]), axis=-1), color=c[:,0]).show()
|
||||
# # ex.scatter(x=uv_map[left_2, 0], y=uv_map[left_2, 1]).show()
|
||||
# # ex.scatter(x=uv_map[right_2, 0], y=uv_map[right_2, 1]).show()
|
||||
|
||||
# new_pairs_1 = np.concatenate((left_1, right_1), axis=-1)
|
||||
# new_pairs_2 = np.concatenate((left_2, right_2), axis=-1)
|
||||
|
||||
# ex.scatter(x=uv_map[new_pairs_1, 0], y=uv_map[new_pairs_1, 1]).show()
|
||||
# ex.scatter(x=uv_map[new_pairs_2, 0], y=uv_map[new_pairs_2, 1]).show()
|
||||
|
||||
|
||||
for i, l in enumerate(left[:56]):
|
||||
m2uv[l] = {left_2[i], right_2[i]}
|
||||
for i, r in enumerate(right):
|
||||
m2uv[r] = {left_1[i], right_1[i]}
|
||||
|
||||
m2uv[left[-1]] = {7131} #7133
|
||||
|
||||
p = uv_map[7133, :]
|
||||
closest = None
|
||||
min_dist = 10000
|
||||
for i in range(uv_map.shape[0]):
|
||||
dist = np.sqrt((p[0] - uv_map[i, 0])**2 + (p[1] - uv_map[i, 1]))
|
||||
if dist < min_dist and i != 7133:
|
||||
min_dist = dist
|
||||
closest = i
|
||||
|
||||
mesh_coord = None
|
||||
for key in m2uv.keys():
|
||||
if len(list(m2uv[key])) == 2:
|
||||
continue
|
||||
if list(m2uv[key])[0] == closest:
|
||||
mesh_coord = key
|
||||
|
||||
assert m2uv[mesh_coord] == { closest }
|
||||
|
||||
m2uv[mesh_coord] = { closest, 7133 }
|
||||
|
||||
np.save('custom_utils/m2uv_skirt.npy', m2uv)
|
||||
|
||||
# m2uv = m2uv_array(m2uv)
|
||||
# uv_counter = np.zeros((uv_map.shape[0]))
|
||||
# uv_counter[m2uv[:, 0]] = 1
|
||||
# uv_counter[m2uv[:, 1]] += 1
|
||||
|
||||
print()
|
||||
Binary file not shown.
Binary file not shown.
@@ -0,0 +1,476 @@
|
||||
import os
|
||||
import sys
|
||||
|
||||
import pickle
|
||||
import scipy.io as sio
|
||||
from math import *
|
||||
|
||||
import numpy as np
|
||||
import tensorflow as tf
|
||||
|
||||
gpus = tf.config.experimental.list_physical_devices('GPU')
|
||||
if gpus:
|
||||
try:
|
||||
# Currently, memory growth needs to be the same across GPUs
|
||||
for gpu in gpus:
|
||||
tf.config.experimental.set_memory_growth(gpu, True)
|
||||
logical_gpus = tf.config.experimental.list_logical_devices('GPU')
|
||||
print(len(gpus), "Physical GPUs,", len(logical_gpus), "Logical GPUs")
|
||||
except RuntimeError as e:
|
||||
# Memory growth must be set before GPUs have been initialized
|
||||
print(e)
|
||||
|
||||
# Left-Right
|
||||
eval_joints = {
|
||||
'All': [1, 2, 3, 4, 5, 7, 8, 9, 12, 15, 16, 17, 18, 19, 20, 21],
|
||||
'Head': [12, 15], # Neck, Head
|
||||
'Torso': [3, 6, 9], # Back joints
|
||||
'Shoulder': [13, 14, 16, 17], # Inner, Outer
|
||||
'Elbow': [18, 19],
|
||||
'Wrist': [20, 21],
|
||||
'Hip': [1, 2],
|
||||
'Knee': [4, 5],
|
||||
'Foot': [7, 8, 10, 11]} # Ankle, Foot
|
||||
|
||||
class SMPL:
|
||||
|
||||
def __init__(self, orig_smpl=False, use_one_hot=False, is_skirt=False):
|
||||
self.J_regressor = []
|
||||
self.v_template = []
|
||||
self.shapedirs = []
|
||||
self.posedirs = []
|
||||
self.weights = []
|
||||
self.use_one_hot = use_one_hot
|
||||
self.is_skirt = is_skirt
|
||||
# Load SMPL models
|
||||
# path = os.path.abspath(os.path.dirname(__file__))
|
||||
with open(os.path.join('custom_utils/super_model_f.pkl'), 'rb') as f:
|
||||
if sys.version_info[0]<3:
|
||||
female = pickle.load(f)
|
||||
else:
|
||||
female = pickle.load(f, encoding="latin1")
|
||||
self.J_regressor.append(tf.constant(np.array(female['J_regressor'].todense(), dtype=np.float32)))
|
||||
self.v_template.append(tf.constant(female['v_template'], dtype=np.float32))
|
||||
self.posedirs.append(tf.constant(female['posedirs'], dtype=np.float32))
|
||||
self.shapedirs.append(tf.constant(female['shapedirs'], dtype=np.float32))
|
||||
self.weights.append(tf.constant(female['weights'], dtype=np.float32))
|
||||
with open(os.path.join('custom_utils/super_model_m.pkl'), 'rb') as f:
|
||||
if sys.version_info[0]<3:
|
||||
male = pickle.load(f)
|
||||
else:
|
||||
male = pickle.load(f, encoding="latin1")
|
||||
self.J_regressor.append(tf.constant(np.array(male['J_regressor'].todense(), dtype=np.float32)))
|
||||
self.v_template.append(tf.constant(male['v_template'], dtype=np.float32))
|
||||
self.posedirs.append(tf.constant(male['posedirs'], dtype=np.float32))
|
||||
self.shapedirs.append(tf.constant(male['shapedirs'], dtype=np.float32))
|
||||
self.weights.append(tf.constant(male['weights'], dtype=np.float32))
|
||||
kintree_table = male['kintree_table']
|
||||
id_to_col = {kintree_table[1, i]: i for i in range(kintree_table.shape[1])}
|
||||
self.parent = {i: id_to_col[kintree_table[0, i]] for i in range(1, kintree_table.shape[1])}
|
||||
self.J_regressor = tf.stack(self.J_regressor, axis=0)
|
||||
self.v_template = tf.stack(self.v_template, axis=0)
|
||||
self.shapedirs = tf.stack(self.shapedirs, axis=0)
|
||||
self.posedirs = tf.stack(self.posedirs, axis=0)
|
||||
self.weights = tf.stack(self.weights, axis=0)
|
||||
self.weights_one_hot = tf.reduce_mean(self.weights, axis=0)
|
||||
self.weights_one_hot = tf.math.argmax(self.weights_one_hot, axis=-1)
|
||||
self.weights_one_hot = tf.one_hot(self.weights_one_hot, 24)
|
||||
|
||||
# Topology no head/hands/feet
|
||||
F = sio.loadmat('custom_utils/ssmpl_faces.mat')
|
||||
self.f = np.array(F['faces_default'])
|
||||
self.f_skirt = np.array(F['faces_skirt'])
|
||||
self.nn = np.array(sio.loadmat('custom_utils/NN.mat'))
|
||||
self.to_keep = sio.loadmat(os.path.join('custom_utils/keep_remove.mat'))['keep'].reshape(-1)
|
||||
|
||||
self.weights_all = self.weights
|
||||
if use_one_hot:
|
||||
self.weights = tf.gather(self.weights_one_hot, self.to_keep, axis=0)
|
||||
self.weights = tf.tile(self.weights[tf.newaxis,:,:], [2,1,1])
|
||||
else:
|
||||
self.weights = tf.gather(self.weights, self.to_keep, axis=1)
|
||||
|
||||
self.shape = (1,27578,3)
|
||||
|
||||
if orig_smpl:
|
||||
self.to_keep = self.to_keep[self.to_keep < 6890]
|
||||
self.J_regressor = self.J_regressor[:,:,:6890]
|
||||
self.v_template = self.v_template[:,:6890,:]
|
||||
self.shapedirs = self.shapedirs[:,:6890,:,:]
|
||||
self.posedirs = self.posedirs[:,:6890,:,:]
|
||||
self.weights_all = self.weights_all[:,:6890,:]
|
||||
self.weights = tf.gather(self.weights_all, self.to_keep, axis=1)
|
||||
self.shape = (1, 6890, 3)
|
||||
|
||||
def Q2R(self, w, x, y, z):
|
||||
def diag(a, b): # computes the diagonal entries, 1 - 2*a**2 - 2*b**2
|
||||
return 1 - 2 * tf.pow(a, 2) - 2 * tf.pow(b, 2)
|
||||
|
||||
def tr_add(a, b, c, d): # computes triangle entries with addition
|
||||
return 2 * a * b + 2 * c * d
|
||||
|
||||
def tr_sub(a, b, c, d): # computes triangle entries with subtraction
|
||||
return 2 * a * b - 2 * c * d
|
||||
|
||||
m = [[diag(y, z), tr_sub(x, y, z, w), tr_add(x, z, y, w)],
|
||||
[tr_add(x, y, z, w), diag(x, z), tr_sub(y, z, x, w)],
|
||||
[tr_sub(x, z, y, w), tr_add(y, z, x, w), diag(x, y)]]
|
||||
|
||||
R = tf.stack([tf.stack(m[i], axis=-1) for i in range(3)], axis=-2)
|
||||
return R
|
||||
|
||||
def rodrigues(self, r):
|
||||
theta = tf.norm(r, axis=(1, 2), keepdims=True)
|
||||
theta = tf.maximum(theta, np.finfo(np.float32).eps)
|
||||
r_hat = r / theta
|
||||
z_stick = tf.zeros(theta.get_shape().as_list()[0], dtype=tf.float32)
|
||||
m = tf.stack(
|
||||
(z_stick, -r_hat[:, 0, 2], r_hat[:, 0, 1], r_hat[:, 0, 2], z_stick,
|
||||
-r_hat[:, 0, 0], -r_hat[:, 0, 1], r_hat[:, 0, 0], z_stick), axis=1)
|
||||
m = tf.reshape(m, (-1, 3, 3))
|
||||
i_cube = tf.expand_dims(tf.eye(3, dtype=tf.float32), axis=0) + tf.zeros(
|
||||
(theta.get_shape().as_list()[0], 3, 3), dtype=tf.float32)
|
||||
# R = cos * i_cube + (1 - cos) * dot + tf.sin(theta) * m
|
||||
R = i_cube + m * tf.sin(theta) + tf.matmul(m, m) * (1 - tf.cos(theta))
|
||||
return R
|
||||
|
||||
def with_zeros(self, x):
|
||||
return tf.concat((x, tf.constant([[0.0, 0.0, 0.0, 1.0]], dtype=tf.float32)), axis=0)
|
||||
|
||||
def pack(self, x):
|
||||
return tf.concat((tf.zeros((x.get_shape().as_list()[0], 4, 3), dtype=tf.float32), x), axis=2)
|
||||
|
||||
def run_rodrigues(self, pose):
|
||||
batch_size = pose.shape[0]
|
||||
R_out = self.rodrigues(tf.reshape(pose, (-1, 1, 3)))
|
||||
R_out = tf.reshape(R_out, (batch_size, -1, 3, 3))
|
||||
#R_out = []
|
||||
#for i in range(batch_size):
|
||||
# R_out.append(self.rodrigues(tf.reshape(pose[i,:], (-1, 1, 3))))
|
||||
#R_out = tf.stack(R_out, axis=0)
|
||||
return R_out
|
||||
|
||||
def run(self, R, betas, remove=True):
|
||||
batch_size = betas.shape[0]
|
||||
# R = self.run_rodrigues(pose)
|
||||
|
||||
#J_regressor = tf.gather(self.J_regressor, gender, axis=0)
|
||||
#shapedirs = tf.gather(self.shapedirs, gender, axis=0)
|
||||
#v_template = tf.gather(self.v_template, gender, axis=0)
|
||||
|
||||
v_shaped = []
|
||||
v_shaped.append(tf.tensordot(betas[:,:10], self.shapedirs[0], axes=[[1], [2]]) + tf.reshape(self.v_template[0], shape=self.shape)) # female
|
||||
v_shaped.append(tf.tensordot(betas[:,:10], self.shapedirs[1], axes=[[1], [2]]) + tf.reshape(self.v_template[1], shape=self.shape)) # male
|
||||
v_shaped = tf.stack(v_shaped, axis=-1)
|
||||
|
||||
def single_run(x, J_regressor=self.J_regressor, posedirs=self.posedirs, weights=self.weights_all, parent=self.parent):
|
||||
R = x[0]
|
||||
v_shaped = x[1]
|
||||
g = x[2]
|
||||
|
||||
J = tf.matmul(J_regressor[g], v_shaped[:, :, g])
|
||||
R_cube_big = R
|
||||
R_cube = R_cube_big[1:]
|
||||
I_cube = tf.expand_dims(tf.eye(3, dtype=tf.float32), axis=0) + tf.zeros((R_cube.get_shape()[0], 3, 3), dtype=tf.float32)
|
||||
lrotmin = tf.squeeze(tf.reshape((R_cube - I_cube), (-1, 1)))
|
||||
v_posed = v_shaped[:, :, g] + tf.tensordot(posedirs[g], lrotmin, axes=[[2], [0]])
|
||||
results = []
|
||||
results.append(self.with_zeros(tf.concat((R_cube_big[0], tf.reshape(J[0, :], (3, 1))), axis=1)))
|
||||
for j in range(1, 24):
|
||||
results.append(tf.matmul(results[parent[j]], self.with_zeros(tf.concat((R_cube_big[j], tf.reshape(J[j, :] - J[parent[j], :], (3, 1))), axis=1))))
|
||||
stacked = tf.stack(results, axis=0)
|
||||
results = stacked - self.pack(tf.matmul(stacked, tf.reshape(tf.concat((J, tf.zeros((24, 1), dtype=tf.float32)), axis=1), (24, 4, 1))))
|
||||
|
||||
J = tf.matmul(results, tf.expand_dims(tf.concat([J, tf.ones([24, 1], tf.float32)], axis=1), axis=1), transpose_b=True)
|
||||
J = tf.gather(J, tf.constant([0, 1, 2], tf.int32), axis=1)
|
||||
|
||||
T = tf.tensordot(weights[g], results, axes=((1), (0)))
|
||||
rest_shape_h = tf.concat((v_posed, tf.ones((v_posed.get_shape().as_list()[0], 1), dtype=tf.float32)), axis=1)
|
||||
v = tf.matmul(T, tf.reshape(rest_shape_h, (-1, 4, 1)))
|
||||
v = tf.reshape(v, (-1, 4))[:, :3]
|
||||
|
||||
return v, tf.squeeze(J), v_posed
|
||||
|
||||
gender = tf.cast(tf.greater_equal(betas[:, 10] , 0.5), tf.int32)
|
||||
T_out, J_out, v_posed = tf.map_fn(single_run, (R, v_shaped, gender), (tf.float32, tf.float32, tf.float32))
|
||||
|
||||
J0 = tf.reshape(J_out[:,0,:], (-1, 1, 3))
|
||||
J_out = tf.subtract(J_out, J0)
|
||||
T_out = tf.subtract(T_out, J0)
|
||||
v_posed = tf.subtract(v_posed, J0)
|
||||
|
||||
if remove:
|
||||
T_out = tf.gather(T_out, self.to_keep, axis=1)
|
||||
v_posed = tf.gather(v_posed, self.to_keep, axis=1)
|
||||
|
||||
return T_out, J_out, v_posed
|
||||
|
||||
|
||||
def run_with_weights(self, R, betas, v_cloth, weights=None, filter=None):
|
||||
|
||||
v_shaped = []
|
||||
v_shaped.append(tf.tensordot(betas[:,:10], self.shapedirs[0], axes=[[1], [2]]) + tf.reshape(self.v_template[0], shape=self.shape)) # female
|
||||
v_shaped.append(tf.tensordot(betas[:,:10], self.shapedirs[1], axes=[[1], [2]]) + tf.reshape(self.v_template[1], shape=self.shape)) # male
|
||||
v_shaped = tf.stack(v_shaped, axis=-1)
|
||||
|
||||
if weights is None:
|
||||
w = self.weights_one_hot[np.newaxis, :, :]
|
||||
if filter is not None:
|
||||
w = tf.gather(tf.gather(w, self.to_keep, axis=1), filter, axis=1)
|
||||
if v_cloth.shape[1] == 14475:
|
||||
w = tf.gather(w, self.to_keep, axis=1)
|
||||
else:
|
||||
w = weights
|
||||
|
||||
def single_run(x, J_regressor=self.J_regressor, w=w, parent=self.parent):
|
||||
R = x[0]
|
||||
v_shaped = x[1]
|
||||
g = x[2]
|
||||
v_cloth = x[3]
|
||||
|
||||
J = tf.matmul(J_regressor[g], v_shaped[:, :, g])
|
||||
R_cube_big = R
|
||||
results = []
|
||||
results.append(self.with_zeros(tf.concat((R_cube_big[0], tf.reshape(J[0, :], (3, 1))), axis=1)))
|
||||
for j in range(1, 24):
|
||||
results.append(tf.matmul(results[parent[j]], self.with_zeros(tf.concat((R_cube_big[j], tf.reshape(J[j, :] - J[parent[j], :], (3, 1))), axis=1))))
|
||||
stacked = tf.stack(results, axis=0)
|
||||
results = stacked - self.pack(tf.matmul(stacked, tf.reshape(tf.concat((J, tf.zeros((24, 1), dtype=tf.float32)), axis=1), (24, 4, 1))))
|
||||
|
||||
J = tf.matmul(results, tf.expand_dims(tf.concat([J, tf.ones([24, 1], tf.float32)], axis=1), axis=1), transpose_b=True)
|
||||
J = tf.gather(J, tf.constant([0, 1, 2], tf.int32), axis=1)
|
||||
|
||||
T = tf.tensordot(w[0], results, axes=((1), (0)))
|
||||
rest_shape_h = tf.concat((v_cloth, tf.ones((v_cloth.shape[0], 1), dtype=tf.float32)), axis=1)
|
||||
v = tf.matmul(T, tf.reshape(rest_shape_h, (-1, 4, 1)))
|
||||
v = tf.reshape(v, (-1, 4))[:, :3]
|
||||
|
||||
return v, tf.squeeze(J)
|
||||
|
||||
gender = tf.cast(tf.greater_equal(betas[:, 10] , 0.5), tf.int32)
|
||||
T_out, J_out = tf.map_fn(single_run, (R, v_shaped, gender, v_cloth), (tf.float32, tf.float32))
|
||||
|
||||
J0 = tf.reshape(J_out[:,0,:], (-1, 1, 3))
|
||||
T_out = tf.subtract(T_out, J0)
|
||||
|
||||
return T_out
|
||||
|
||||
|
||||
def get_joints_template(self, betas, v_shaped=None):
|
||||
|
||||
if v_shaped is None:
|
||||
v_shaped = []
|
||||
v_shaped.append(tf.tensordot(betas[:,:10], self.shapedirs[0], axes=[[1], [2]]) + tf.reshape(self.v_template[0], shape=self.shape)) # female
|
||||
v_shaped.append(tf.tensordot(betas[:,:10], self.shapedirs[1], axes=[[1], [2]]) + tf.reshape(self.v_template[1], shape=self.shape)) # male
|
||||
v_shaped = tf.stack(v_shaped, axis=-1)
|
||||
|
||||
def single_run(x, J_regressor=self.J_regressor):
|
||||
v_shaped = x[0]
|
||||
g = x[1]
|
||||
J = tf.matmul(J_regressor[g], v_shaped[:, :, g])
|
||||
return J
|
||||
|
||||
gender = tf.cast(tf.greater(betas[:, 10] , 0.5), tf.int32)
|
||||
J = tf.map_fn(single_run, (v_shaped, gender), (tf.float32))
|
||||
J0 = tf.reshape(J[:,0,:], (-1, 1, 3))
|
||||
|
||||
return tf.subtract(J, J0)
|
||||
|
||||
def get_root_joint(self, betas, v_shaped=None):
|
||||
|
||||
if v_shaped is None:
|
||||
v_shaped = []
|
||||
v_shaped.append(tf.tensordot(betas[:,:10], self.shapedirs[0], axes=[[1], [2]]) + tf.reshape(self.v_template[0], shape=self.shape)) # female
|
||||
v_shaped.append(tf.tensordot(betas[:,:10], self.shapedirs[1], axes=[[1], [2]]) + tf.reshape(self.v_template[1], shape=self.shape)) # male
|
||||
v_shaped = tf.stack(v_shaped, axis=-1)
|
||||
|
||||
def single_run(x, J_regressor=self.J_regressor):
|
||||
v_shaped = x[0]
|
||||
g = x[1]
|
||||
J = tf.matmul(J_regressor[g], v_shaped[:, :, g])
|
||||
return J
|
||||
|
||||
gender = tf.cast(tf.greater(betas[:, 10] , 0.5), tf.int32)
|
||||
J = tf.map_fn(single_run, (v_shaped, gender), (tf.float32))
|
||||
J0 = tf.reshape(J[:,0,:], (-1, 1, 3))
|
||||
|
||||
return J[:,0,:]
|
||||
|
||||
# return J0
|
||||
|
||||
|
||||
def run_garment(self, pose, betas, v_cloth, v_shaped=None):
|
||||
J = self.get_joints_template(betas, v_shaped)
|
||||
|
||||
if self.use_one_hot and self.is_skirt:
|
||||
# p_avg = (pose[:, 3:6] + pose[:, 6:9]) / 2
|
||||
# pose_dress = tf.concat([pose[:, :3], p_avg, p_avg, pose[:, 9:]], axis=-1)
|
||||
# R = self.run_rodrigues(pose_dress).numpy()
|
||||
# for b in range(betas.shape[0]):
|
||||
# R[b][[4,5,7,8,10,11],:,:] = np.eye(3)
|
||||
|
||||
pose_dress = pose
|
||||
# p_avg = (pose_dress[:, 3:6] + pose_dress[:, 6:9]) / 2
|
||||
# pose_dress[:,3:9] = np.concatenate([p_avg, p_avg], axis=-1)
|
||||
R = self.run_rodrigues(pose_dress).numpy()
|
||||
for b in range(betas.shape[0]):
|
||||
R[b][[4,5,7,8,10,11],:,:] = np.eye(3)
|
||||
else:
|
||||
R = self.run_rodrigues(pose)
|
||||
|
||||
# if filter is not None:
|
||||
# weights = tf.boolean_mask(self.weights, filter)
|
||||
|
||||
def single_run(x, weights=self.weights, parent=self.parent):
|
||||
R = x[0]
|
||||
J = x[1]
|
||||
g = x[2]
|
||||
v_cloth = x[3]
|
||||
|
||||
R_cube_big = R
|
||||
|
||||
results = []
|
||||
results.append(self.with_zeros(tf.concat((R_cube_big[0], tf.reshape(J[0, :], (3, 1))), axis=1)))
|
||||
for j in range(1, 24):
|
||||
results.append(tf.matmul(results[parent[j]], self.with_zeros(tf.concat((R_cube_big[j], tf.reshape(J[j, :] - J[parent[j], :], (3, 1))), axis=1))))
|
||||
stacked = tf.stack(results, axis=0)
|
||||
results = stacked - self.pack(tf.matmul(stacked, tf.reshape(tf.concat((J, tf.zeros((24, 1), dtype=tf.float32)), axis=1), (24, 4, 1))))
|
||||
|
||||
J = tf.matmul(results, tf.expand_dims(tf.concat([J, tf.ones([24, 1], tf.float32)], axis=1), axis=1), transpose_b=True)
|
||||
J = tf.gather(J, tf.constant([0, 1, 2], tf.int32), axis=1)
|
||||
|
||||
T = tf.tensordot(weights[g, :, :], results, axes=((1), (0)))
|
||||
rest_shape_h = tf.concat((v_cloth, tf.ones((v_cloth.shape[0], 1), dtype=tf.float32)), axis=1)
|
||||
v = tf.matmul(T, tf.reshape(rest_shape_h, (-1, 4, 1)))
|
||||
v = tf.reshape(v, (-1, 4))[:, :3]
|
||||
|
||||
return v, tf.squeeze(J)
|
||||
|
||||
gender = tf.cast(tf.greater(betas[:, 10] , 0.5), tf.int32)
|
||||
T_out, J_out = tf.map_fn(single_run, (R, J, gender, v_cloth), (tf.float32, tf.float32))
|
||||
|
||||
J0 = tf.reshape(J_out[:,0,:], (-1, 1, 3))
|
||||
J_out = tf.subtract(J_out, J0)
|
||||
T_out = tf.subtract(T_out, J0)
|
||||
|
||||
|
||||
return T_out
|
||||
|
||||
|
||||
def run_garment_orig(self, R, betas, v_cloth, J=None, weights=None):
|
||||
batch_size = R.shape[0]
|
||||
|
||||
if J is None:
|
||||
J = self.get_joints_template(betas)
|
||||
|
||||
if weights is None:
|
||||
weights = tf.gather(self.weights_one_hot, self.to_keep, axis=0)
|
||||
# else:
|
||||
# weights = tf.gather(weights, self.to_keep, axis=0)
|
||||
|
||||
def single_run(x, J_regressor=self.J_regressor, posedirs=self.posedirs, weights=weights, parent=self.parent):
|
||||
R = x[0]
|
||||
J = x[1]
|
||||
g = x[2]
|
||||
v_cloth = x[3]
|
||||
|
||||
R_cube_big = R
|
||||
# R_cube = R_cube_big[1:]
|
||||
# I_cube = tf.expand_dims(tf.eye(3, dtype=tf.float32), axis=0) + tf.zeros((R_cube.get_shape()[0], 3, 3), dtype=tf.float32)
|
||||
# lrotmin = tf.squeeze(tf.reshape((R_cube - I_cube), (-1, 1)))
|
||||
|
||||
results = []
|
||||
results.append(self.with_zeros(tf.concat((R_cube_big[0], tf.reshape(J[0, :], (3, 1))), axis=1)))
|
||||
for j in range(1, 24):
|
||||
results.append(tf.matmul(results[parent[j]], self.with_zeros(tf.concat((R_cube_big[j], tf.reshape(J[j, :] - J[parent[j], :], (3, 1))), axis=1))))
|
||||
stacked = tf.stack(results, axis=0)
|
||||
results = stacked - self.pack(tf.matmul(stacked, tf.reshape(tf.concat((J, tf.zeros((24, 1), dtype=tf.float32)), axis=1), (24, 4, 1))))
|
||||
|
||||
J = tf.matmul(results, tf.expand_dims(tf.concat([J, tf.ones([24, 1], tf.float32)], axis=1), axis=1), transpose_b=True)
|
||||
J = tf.gather(J, tf.constant([0, 1, 2], tf.int32), axis=1)
|
||||
|
||||
T = tf.tensordot(weights, results, axes=((1), (0)))
|
||||
rest_shape_h = tf.concat((v_cloth, tf.ones((v_cloth.get_shape().as_list()[0], 1), dtype=tf.float32)), axis=1)
|
||||
v = tf.matmul(T, tf.reshape(rest_shape_h, (-1, 4, 1)))
|
||||
v = tf.reshape(v, (-1, 4))[:, :3]
|
||||
|
||||
return v, tf.squeeze(J)
|
||||
|
||||
gender = tf.cast(tf.greater_equal(betas[:, 10] , 0), tf.int32)
|
||||
T_out, J_out = tf.map_fn(single_run, (R, J, gender, v_cloth), (tf.float32, tf.float32))
|
||||
|
||||
J0 = tf.reshape(J_out[:,0,:], (-1, 1, 3))
|
||||
J_out = tf.subtract(J_out, J0)
|
||||
T_out = tf.subtract(T_out, J0)
|
||||
|
||||
return T_out, J_out
|
||||
|
||||
|
||||
def run_garment_with_weights(self, R, betas, v_cloth, weights, J=None):
|
||||
|
||||
if J is None:
|
||||
J = self.get_joints_template(betas)
|
||||
|
||||
def single_run(x, parent=self.parent):
|
||||
R = x[0]
|
||||
J = x[1]
|
||||
v_cloth = x[2]
|
||||
weights = x[3]
|
||||
|
||||
R_cube_big = R
|
||||
|
||||
results = []
|
||||
results.append(self.with_zeros(tf.concat((R_cube_big[0], tf.reshape(J[0, :], (3, 1))), axis=1)))
|
||||
for j in range(1, 24):
|
||||
results.append(tf.matmul(results[parent[j]], self.with_zeros(tf.concat((R_cube_big[j], tf.reshape(J[j, :] - J[parent[j], :], (3, 1))), axis=1))))
|
||||
stacked = tf.stack(results, axis=0)
|
||||
results = stacked - self.pack(tf.matmul(stacked, tf.reshape(tf.concat((J, tf.zeros((24, 1), dtype=tf.float32)), axis=1), (24, 4, 1))))
|
||||
|
||||
J = tf.matmul(results, tf.expand_dims(tf.concat([J, tf.ones([24, 1], tf.float32)], axis=1), axis=1), transpose_b=True)
|
||||
J = tf.gather(J, tf.constant([0, 1, 2], tf.int32), axis=1)
|
||||
|
||||
T = tf.tensordot(weights, results, axes=((1), (0)))
|
||||
rest_shape_h = tf.concat((v_cloth, tf.ones((v_cloth.get_shape().as_list()[0], 1), dtype=tf.float32)), axis=1)
|
||||
v = tf.matmul(T, tf.reshape(rest_shape_h, (-1, 4, 1)))
|
||||
v = tf.reshape(v, (-1, 4))[:, :3]
|
||||
|
||||
return v, tf.squeeze(J)
|
||||
|
||||
# gender = tf.cast(tf.greater_equal(betas[:, 10] , 0), tf.int32)
|
||||
T_out, J_out = tf.map_fn(single_run, (R, J, v_cloth, weights), (tf.float32, tf.float32))
|
||||
|
||||
J0 = tf.reshape(J_out[:,0,:], (-1, 1, 3))
|
||||
J_out = tf.subtract(J_out, J0)
|
||||
T_out = tf.subtract(T_out, J0)
|
||||
|
||||
return T_out
|
||||
|
||||
|
||||
def run_backward(self, R, J, v_posed, weights):
|
||||
batch_size = R.get_shape()[0]
|
||||
|
||||
if weights.shape[0] != 14475:
|
||||
weights = tf.gather(self.weights_one_hot, self.to_keep, axis=0)
|
||||
|
||||
|
||||
J_relative = J[:, 1:, :] - tf.gather(J, tf.constant(list(self.parent.values()), tf.int32), axis=1) # batch_size x 23 x 3
|
||||
|
||||
J_out = []
|
||||
J_out.append(tf.expand_dims(J[:, 0, :], axis=-1)) # batch_size x 3 x 1
|
||||
results = []
|
||||
results.append(tf.expand_dims(tf.eye(3, dtype=tf.float32), axis=0) + tf.zeros((batch_size, 3, 3), dtype=tf.float32))
|
||||
for j in range(1, 24):
|
||||
results.append(tf.matmul(R[:, self.parent[j], :, :], results[self.parent[j]]))
|
||||
J_out.append(tf.matmul(results[j], tf.expand_dims(J_relative[:, j-1, :], axis=-1)) + J_out[self.parent[j]]) # batch_size x 3 x 1
|
||||
results = tf.transpose(tf.stack(results, axis=0), [1, 0, 2, 3]) # (batch_size x 24 x 3 x 3)
|
||||
J_out = tf.transpose(tf.stack(J_out, axis=0), [1, 0, 2, 3]) # (batch_size x 24 x 3 x 1)
|
||||
|
||||
results = tf.matmul(R, results)
|
||||
T = tf.transpose(tf.tensordot(results, weights, axes=((1), (1))), [0, 3, 1, 2]) # (batch_size x 6890 x 3 x 3)
|
||||
v_posed = tf.expand_dims(v_posed, axis=-1) # (batch_size x 6890 x 3 x 1)
|
||||
J = tf.expand_dims(J, axis=-1) # (batch_size x 24 x 3 x 1)
|
||||
v = tf.matmul(T, v_posed) + tf.transpose(tf.tensordot(J_out - tf.matmul(results, J), weights, axes=((1), (1))), [0, 3, 1, 2]) # (batch_size x 6890 x 3 x 1)
|
||||
v = tf.squeeze(v, -1) # (batch_size x 6890 x 3)
|
||||
J_out = tf.squeeze(J_out, -1) # (batch_size x 24 x 3)
|
||||
|
||||
return v, J_out
|
||||
@@ -0,0 +1,346 @@
|
||||
import os
|
||||
import sys
|
||||
|
||||
import pickle
|
||||
import scipy.io as sio
|
||||
from math import *
|
||||
|
||||
import numpy as np
|
||||
import tensorflow as tf
|
||||
|
||||
# Left-Right
|
||||
eval_joints = {
|
||||
'All': [1, 2, 3, 4, 5, 7, 8, 9, 12, 15, 16, 17, 18, 19, 20, 21],
|
||||
'Head': [12, 15], # Neck, Head
|
||||
'Torso': [3, 6, 9], # Back joints
|
||||
'Shoulder': [13, 14, 16, 17], # Inner, Outer
|
||||
'Elbow': [18, 19],
|
||||
'Wrist': [20, 21],
|
||||
'Hip': [1, 2],
|
||||
'Knee': [4, 5],
|
||||
'Foot': [7, 8, 10, 11]} # Ankle, Foot
|
||||
|
||||
class SMPL:
|
||||
|
||||
def __init__(self, path):
|
||||
self.J_regressor = []
|
||||
self.v_template = []
|
||||
self.shapedirs = []
|
||||
self.posedirs = []
|
||||
self.weights = []
|
||||
# Load SMPL models
|
||||
# path = os.path.abspath(os.path.dirname(__file__))
|
||||
with open(os.path.join(path, 'super_model_f.pkl'), 'rb') as f:
|
||||
if sys.version_info[0]<3:
|
||||
female = pickle.load(f)
|
||||
else:
|
||||
female = pickle.load(f, encoding="latin1")
|
||||
self.J_regressor.append(tf.constant(np.array(female['J_regressor'].todense(), dtype=np.float32)))
|
||||
self.v_template.append(tf.constant(female['v_template'], dtype=np.float32))
|
||||
self.posedirs.append(tf.constant(female['posedirs'], dtype=np.float32))
|
||||
self.shapedirs.append(tf.constant(female['shapedirs'], dtype=np.float32))
|
||||
self.weights.append(tf.constant(female['weights'], dtype=np.float32))
|
||||
with open(os.path.join(path, 'super_model_m.pkl'), 'rb') as f:
|
||||
if sys.version_info[0]<3:
|
||||
male = pickle.load(f)
|
||||
else:
|
||||
male = pickle.load(f, encoding="latin1")
|
||||
self.J_regressor.append(tf.constant(np.array(male['J_regressor'].todense(), dtype=np.float32)))
|
||||
self.v_template.append(tf.constant(male['v_template'], dtype=np.float32))
|
||||
self.posedirs.append(tf.constant(male['posedirs'], dtype=np.float32))
|
||||
self.shapedirs.append(tf.constant(male['shapedirs'], dtype=np.float32))
|
||||
self.weights.append(tf.constant(male['weights'], dtype=np.float32))
|
||||
kintree_table = male['kintree_table']
|
||||
id_to_col = {kintree_table[1, i]: i for i in range(kintree_table.shape[1])}
|
||||
self.parent = {i: id_to_col[kintree_table[0, i]] for i in range(1, kintree_table.shape[1])}
|
||||
self.J_regressor = tf.stack(self.J_regressor, axis=0)
|
||||
self.v_template = tf.stack(self.v_template, axis=0)
|
||||
self.shapedirs = tf.stack(self.shapedirs, axis=0)
|
||||
self.posedirs = tf.stack(self.posedirs, axis=0)
|
||||
self.weights = tf.stack(self.weights, axis=0)
|
||||
self.weights_one_hot = tf.reduce_mean(self.weights, axis=0)
|
||||
self.weights_one_hot = tf.math.argmax(self.weights_one_hot, axis=-1)
|
||||
self.weights_one_hot = tf.one_hot(self.weights_one_hot, 24)
|
||||
|
||||
# Topology no head/hands/feet
|
||||
F = sio.loadmat(path + '/ssmpl_faces.mat')
|
||||
self.f = np.array(F['faces_default'])
|
||||
self.f_skirt = np.array(F['faces_skirt'])
|
||||
self.nn = np.array(sio.loadmat(path + '/NN.mat'))
|
||||
self.to_keep = sio.loadmat(os.path.join(path, 'keep_remove.mat'))['keep'].reshape(-1)
|
||||
|
||||
self.weights = tf.gather(self.weights, self.to_keep, axis=1)
|
||||
# self.shapedirs = tf.gather(self.shapedirs, self.to_keep, axis=1)
|
||||
# self.posedirs = tf.gather(self.posedirs, self.to_keep, axis=1)
|
||||
# self.J_regressor = tf.gather(self.J_regressor, self.to_keep, axis=2)
|
||||
# self.v_template = tf.gather(self.v_template, self.to_keep, axis=1)
|
||||
|
||||
|
||||
def Q2R(self, w, x, y, z):
|
||||
def diag(a, b): # computes the diagonal entries, 1 - 2*a**2 - 2*b**2
|
||||
return 1 - 2 * tf.pow(a, 2) - 2 * tf.pow(b, 2)
|
||||
|
||||
def tr_add(a, b, c, d): # computes triangle entries with addition
|
||||
return 2 * a * b + 2 * c * d
|
||||
|
||||
def tr_sub(a, b, c, d): # computes triangle entries with subtraction
|
||||
return 2 * a * b - 2 * c * d
|
||||
|
||||
m = [[diag(y, z), tr_sub(x, y, z, w), tr_add(x, z, y, w)],
|
||||
[tr_add(x, y, z, w), diag(x, z), tr_sub(y, z, x, w)],
|
||||
[tr_sub(x, z, y, w), tr_add(y, z, x, w), diag(x, y)]]
|
||||
|
||||
R = tf.stack([tf.stack(m[i], axis=-1) for i in range(3)], axis=-2)
|
||||
return R
|
||||
|
||||
def rodrigues(self, r):
|
||||
theta = tf.norm(r, axis=(1, 2), keepdims=True)
|
||||
theta = tf.maximum(theta, np.finfo(np.float32).eps)
|
||||
r_hat = r / theta
|
||||
z_stick = tf.zeros(theta.get_shape().as_list()[0], dtype=tf.float32)
|
||||
m = tf.stack(
|
||||
(z_stick, -r_hat[:, 0, 2], r_hat[:, 0, 1], r_hat[:, 0, 2], z_stick,
|
||||
-r_hat[:, 0, 0], -r_hat[:, 0, 1], r_hat[:, 0, 0], z_stick), axis=1)
|
||||
m = tf.reshape(m, (-1, 3, 3))
|
||||
i_cube = tf.expand_dims(tf.eye(3, dtype=tf.float32), axis=0) + tf.zeros(
|
||||
(theta.get_shape().as_list()[0], 3, 3), dtype=tf.float32)
|
||||
# R = cos * i_cube + (1 - cos) * dot + tf.sin(theta) * m
|
||||
R = i_cube + m * tf.sin(theta) + tf.matmul(m, m) * (1 - tf.cos(theta))
|
||||
return R
|
||||
|
||||
def with_zeros(self, x):
|
||||
return tf.concat((x, tf.ones((x.shape[0], 1, 4)) * tf.constant([[0.0, 0.0, 0.0, 1.0]], dtype=tf.float32)[tf.newaxis,:,:]), axis=1)
|
||||
|
||||
def pack(self, x):
|
||||
return tf.concat((tf.zeros((x.shape[0], 24, 4, 3), dtype=tf.float32), x), axis=3)
|
||||
|
||||
def with_zeros_old(self, x):
|
||||
return tf.concat((x, tf.constant([[0.0, 0.0, 0.0, 1.0]], dtype=tf.float32)), axis=0)
|
||||
|
||||
def pack_old(self, x):
|
||||
return tf.concat((tf.zeros((x.get_shape().as_list()[0], 4, 3), dtype=tf.float32), x), axis=2)
|
||||
|
||||
|
||||
def run_rodrigues(self, pose):
|
||||
batch_size = pose.shape[0]
|
||||
R_out = self.rodrigues(tf.reshape(pose, (-1, 1, 3)))
|
||||
R_out = tf.reshape(R_out, (batch_size, -1, 3, 3))
|
||||
#R_out = []
|
||||
#for i in range(batch_size):
|
||||
# R_out.append(self.rodrigues(tf.reshape(pose[i,:], (-1, 1, 3))))
|
||||
#R_out = tf.stack(R_out, axis=0)
|
||||
return R_out
|
||||
|
||||
def run(self, R, betas, remove=True):
|
||||
batch_size = betas.shape[0]
|
||||
|
||||
#J_regressor = tf.gather(self.J_regressor, gender, axis=0)
|
||||
#shapedirs = tf.gather(self.shapedirs, gender, axis=0)
|
||||
#v_template = tf.gather(self.v_template, gender, axis=0)
|
||||
|
||||
v_shaped = []
|
||||
v_shaped.append(tf.tensordot(betas[:,:10], self.shapedirs[0], axes=[[1], [2]]) + tf.reshape(self.v_template[0], shape=(1,27578,3))) # female
|
||||
v_shaped.append(tf.tensordot(betas[:,:10], self.shapedirs[1], axes=[[1], [2]]) + tf.reshape(self.v_template[1], shape=(1,27578,3))) # male
|
||||
v_shaped = tf.stack(v_shaped, axis=-1)
|
||||
|
||||
def single_run(x, J_regressor=self.J_regressor, posedirs=self.posedirs, weights=self.weights, parent=self.parent):
|
||||
R = x[0]
|
||||
v_shaped = x[1]
|
||||
g = x[2]
|
||||
|
||||
J = tf.matmul(J_regressor[g], v_shaped[:, :, g])
|
||||
R_cube_big = R
|
||||
R_cube = R_cube_big[1:]
|
||||
I_cube = tf.expand_dims(tf.eye(3, dtype=tf.float32), axis=0) + tf.zeros((R_cube.get_shape()[0], 3, 3), dtype=tf.float32)
|
||||
lrotmin = tf.squeeze(tf.reshape((R_cube - I_cube), (-1, 1)))
|
||||
v_posed = v_shaped[:, :, g] + tf.tensordot(posedirs[g], lrotmin, axes=[[2], [0]])
|
||||
results = []
|
||||
results.append(self.with_zeros(tf.concat((R_cube_big[0], tf.reshape(J[0, :], (3, 1))), axis=1)))
|
||||
for j in range(1, 24):
|
||||
results.append(tf.matmul(results[parent[j]], self.with_zeros(tf.concat((R_cube_big[j], tf.reshape(J[j, :] - J[parent[j], :], (3, 1))), axis=1))))
|
||||
stacked = tf.stack(results, axis=0)
|
||||
results = stacked - self.pack(tf.matmul(stacked, tf.reshape(tf.concat((J, tf.zeros((24, 1), dtype=tf.float32)), axis=1), (24, 4, 1))))
|
||||
|
||||
J = tf.matmul(results, tf.expand_dims(tf.concat([J, tf.ones([24, 1], tf.float32)], axis=1), axis=1), transpose_b=True)
|
||||
J = tf.gather(J, tf.constant([0, 1, 2], tf.int32), axis=1)
|
||||
|
||||
T = tf.tensordot(weights[g], results, axes=((1), (0)))
|
||||
rest_shape_h = tf.concat((v_posed, tf.ones((v_posed.get_shape().as_list()[0], 1), dtype=tf.float32)), axis=1)
|
||||
v = tf.matmul(T, tf.reshape(rest_shape_h, (-1, 4, 1)))
|
||||
v = tf.reshape(v, (-1, 4))[:, :3]
|
||||
|
||||
return v, tf.squeeze(J), v_posed
|
||||
|
||||
gender = tf.cast(tf.greater_equal(betas[:, 10] , 0), tf.int32)
|
||||
T_out, J_out, v_posed = tf.map_fn(single_run, (R, v_shaped, gender), (tf.float32, tf.float32, tf.float32))
|
||||
|
||||
J0 = tf.reshape(J_out[:,0,:], (-1, 1, 3))
|
||||
J_out = tf.subtract(J_out, J0)
|
||||
T_out = tf.subtract(T_out, J0)
|
||||
v_posed = tf.subtract(v_posed, J0)
|
||||
|
||||
if remove:
|
||||
T_out = tf.gather(T_out, self.to_keep, axis=1)
|
||||
v_posed = tf.gather(v_posed, self.to_keep, axis=1)
|
||||
|
||||
return T_out, J_out, v_posed
|
||||
|
||||
def run_with_weights(self, R, betas, v_cloth, weights=None, filter=None):
|
||||
|
||||
v_shaped = []
|
||||
v_shaped.append(tf.tensordot(betas[:,:10], self.shapedirs[0], axes=[[1], [2]]) + tf.reshape(self.v_template[0], shape=(1,27578,3))) # female
|
||||
v_shaped.append(tf.tensordot(betas[:,:10], self.shapedirs[1], axes=[[1], [2]]) + tf.reshape(self.v_template[1], shape=(1,27578,3))) # male
|
||||
v_shaped = tf.stack(v_shaped, axis=-1)
|
||||
|
||||
if weights is None:
|
||||
w = self.weights_one_hot[np.newaxis, :, :]
|
||||
if filter is not None:
|
||||
w = tf.gather(tf.gather(w, self.to_keep, axis=1), filter, axis=1)
|
||||
if v_cloth.shape[1] == 14475:
|
||||
w = tf.gather(w, self.to_keep, axis=1)
|
||||
else:
|
||||
w = weights
|
||||
|
||||
def single_run(x, J_regressor=self.J_regressor, w=w, parent=self.parent):
|
||||
R = x[0]
|
||||
v_shaped = x[1]
|
||||
g = x[2]
|
||||
v_cloth = x[3]
|
||||
|
||||
J = tf.matmul(J_regressor[g], v_shaped[:, :, g])
|
||||
R_cube_big = R
|
||||
results = []
|
||||
results.append(self.with_zeros(tf.concat((R_cube_big[0], tf.reshape(J[0, :], (3, 1))), axis=1)))
|
||||
for j in range(1, 24):
|
||||
results.append(tf.matmul(results[parent[j]], self.with_zeros(tf.concat((R_cube_big[j], tf.reshape(J[j, :] - J[parent[j], :], (3, 1))), axis=1))))
|
||||
stacked = tf.stack(results, axis=0)
|
||||
results = stacked - self.pack(tf.matmul(stacked, tf.reshape(tf.concat((J, tf.zeros((24, 1), dtype=tf.float32)), axis=1), (24, 4, 1))))
|
||||
|
||||
J = tf.matmul(results, tf.expand_dims(tf.concat([J, tf.ones([24, 1], tf.float32)], axis=1), axis=1), transpose_b=True)
|
||||
J = tf.gather(J, tf.constant([0, 1, 2], tf.int32), axis=1)
|
||||
|
||||
T = tf.tensordot(w[0], results, axes=((1), (0)))
|
||||
rest_shape_h = tf.concat((v_cloth, tf.ones((v_cloth.shape[0], 1), dtype=tf.float32)), axis=1)
|
||||
v = tf.matmul(T, tf.reshape(rest_shape_h, (-1, 4, 1)))
|
||||
v = tf.reshape(v, (-1, 4))[:, :3]
|
||||
|
||||
return v, tf.squeeze(J)
|
||||
|
||||
gender = tf.cast(tf.greater_equal(betas[:, 10] , 0), tf.int32)
|
||||
T_out, J_out = tf.map_fn(single_run, (R, v_shaped, gender, v_cloth), (tf.float32, tf.float32))
|
||||
|
||||
J0 = tf.reshape(J_out[:,0,:], (-1, 1, 3))
|
||||
T_out = tf.subtract(T_out, J0)
|
||||
|
||||
return T_out
|
||||
|
||||
|
||||
|
||||
def get_joints_template(self, betas, gender, batch_size):
|
||||
v_shaped = tf.tensordot(betas[tf.newaxis,:], self.shapedirs[gender], axes=[[1], [2]]) + tf.reshape(self.v_template[gender], shape=(1,27578,3))
|
||||
J = tf.matmul(self.J_regressor[gender], v_shaped)
|
||||
J0 = tf.reshape(J[:,0,:], (-1, 1, 3))
|
||||
J_single = tf.subtract(J, J0)#[0]
|
||||
J = tf.ones((batch_size, 24, 3)) * J_single
|
||||
|
||||
return J
|
||||
|
||||
|
||||
def run_garment(self, R, betas, gender, v_cloth, filter_idx=None):
|
||||
batch_size = v_cloth.shape[0]
|
||||
|
||||
J = self.get_joints_template(betas, gender, batch_size)
|
||||
|
||||
if filter_idx is not None:
|
||||
weights = tf.gather(self.weights, filter_idx, axis=1)
|
||||
|
||||
R_cube_big = R
|
||||
results = []
|
||||
results.append(self.with_zeros(tf.concat((R_cube_big[:,0,:,:], J[:, 0, :, tf.newaxis]), axis=2)))
|
||||
for j in range(1, 24):
|
||||
results.append(tf.matmul(results[self.parent[j]], self.with_zeros(tf.concat((R_cube_big[:,j,:,:], (J[:, j, :] - J[:, self.parent[j], :])[:,:,tf.newaxis]), axis=2))))
|
||||
stacked = tf.stack(results, axis=1)
|
||||
results = stacked - self.pack(tf.matmul(stacked, tf.reshape(tf.concat((J, tf.zeros((batch_size, 24, 1), dtype=tf.float32)), axis=2), (batch_size, 24, 4, 1))))
|
||||
|
||||
J = tf.matmul(results, tf.expand_dims(tf.concat([J, tf.ones([batch_size, 24, 1], tf.float32)], axis=2), axis=2), transpose_b=True)
|
||||
J = J[:, :, :3, :]
|
||||
|
||||
T = tf.tensordot(weights[gender, :, :], results, axes=((1), (1)))
|
||||
T = tf.transpose(T, [1,0,2,3])
|
||||
|
||||
rest_shape_h = tf.concat((v_cloth, tf.ones((batch_size, v_cloth.shape[1], 1), dtype=tf.float32)), axis=2)
|
||||
v = tf.matmul(T, tf.reshape(rest_shape_h, (batch_size, -1, 4, 1)))
|
||||
v = tf.reshape(v, (batch_size, -1, 4))[:, :, :3]
|
||||
|
||||
J_out = tf.squeeze(J)
|
||||
if J.shape[0] == 1:
|
||||
J_out = J_out[tf.newaxis, :, :]
|
||||
J0 = tf.reshape(J_out[:,0,:], (-1, 1, 3))
|
||||
J_out = tf.subtract(J_out, J0)
|
||||
T_out = tf.subtract(v, J0)
|
||||
|
||||
return T_out
|
||||
|
||||
def run_garment_with_weights(self, R, betas, v_cloth, weights, J=None):
|
||||
|
||||
if J is None:
|
||||
J = self.get_joints_template(betas)
|
||||
|
||||
def single_run(x, parent=self.parent):
|
||||
R = x[0]
|
||||
J = x[1]
|
||||
v_cloth = x[2]
|
||||
weights = x[3]
|
||||
|
||||
R_cube_big = R
|
||||
|
||||
results = []
|
||||
results.append(self.with_zeros(tf.concat((R_cube_big[0], tf.reshape(J[0, :], (3, 1))), axis=1)))
|
||||
for j in range(1, 24):
|
||||
results.append(tf.matmul(results[parent[j]], self.with_zeros(tf.concat((R_cube_big[j], tf.reshape(J[j, :] - J[parent[j], :], (3, 1))), axis=1))))
|
||||
stacked = tf.stack(results, axis=0)
|
||||
results = stacked - self.pack(tf.matmul(stacked, tf.reshape(tf.concat((J, tf.zeros((24, 1), dtype=tf.float32)), axis=1), (24, 4, 1))))
|
||||
|
||||
J = tf.matmul(results, tf.expand_dims(tf.concat([J, tf.ones([24, 1], tf.float32)], axis=1), axis=1), transpose_b=True)
|
||||
J = tf.gather(J, tf.constant([0, 1, 2], tf.int32), axis=1)
|
||||
|
||||
T = tf.tensordot(weights, results, axes=((1), (0)))
|
||||
rest_shape_h = tf.concat((v_cloth, tf.ones((v_cloth.get_shape().as_list()[0], 1), dtype=tf.float32)), axis=1)
|
||||
v = tf.matmul(T, tf.reshape(rest_shape_h, (-1, 4, 1)))
|
||||
v = tf.reshape(v, (-1, 4))[:, :3]
|
||||
|
||||
return v, tf.squeeze(J)
|
||||
|
||||
# gender = tf.cast(tf.greater_equal(betas[:, 10] , 0), tf.int32)
|
||||
T_out, J_out = tf.map_fn(single_run, (R, J, v_cloth, weights), (tf.float32, tf.float32))
|
||||
|
||||
J0 = tf.reshape(J_out[:,0,:], (-1, 1, 3))
|
||||
J_out = tf.subtract(J_out, J0)
|
||||
T_out = tf.subtract(T_out, J0)
|
||||
|
||||
return T_out
|
||||
|
||||
def run_backward(self, R, J, v_posed, weights):
|
||||
batch_size = R.get_shape()[0]
|
||||
|
||||
J_relative = J[:, 1:, :] - tf.gather(J, tf.constant(list(self.parent.values()), tf.int32), axis=1) # batch_size x 23 x 3
|
||||
|
||||
J_out = []
|
||||
J_out.append(tf.expand_dims(J[:, 0, :], axis=-1)) # batch_size x 3 x 1
|
||||
results = []
|
||||
results.append(tf.expand_dims(tf.eye(3, dtype=tf.float32), axis=0) + tf.zeros((batch_size, 3, 3), dtype=tf.float32))
|
||||
for j in range(1, 24):
|
||||
results.append(tf.matmul(R[:, self.parent[j], :, :], results[self.parent[j]]))
|
||||
J_out.append(tf.matmul(results[j], tf.expand_dims(J_relative[:, j-1, :], axis=-1)) + J_out[self.parent[j]]) # batch_size x 3 x 1
|
||||
results = tf.transpose(tf.stack(results, axis=0), [1, 0, 2, 3]) # (batch_size x 24 x 3 x 3)
|
||||
J_out = tf.transpose(tf.stack(J_out, axis=0), [1, 0, 2, 3]) # (batch_size x 24 x 3 x 1)
|
||||
|
||||
results = tf.matmul(R, results)
|
||||
T = tf.transpose(tf.tensordot(results, weights, axes=((1), (1))), [0, 3, 1, 2]) # (batch_size x 6890 x 3 x 3)
|
||||
v_posed = tf.expand_dims(v_posed, axis=-1) # (batch_size x 6890 x 3 x 1)
|
||||
J = tf.expand_dims(J, axis=-1) # (batch_size x 24 x 3 x 1)
|
||||
v = tf.matmul(T, v_posed) + tf.transpose(tf.tensordot(J_out - tf.matmul(results, J), weights, axes=((1), (1))), [0, 3, 1, 2]) # (batch_size x 6890 x 3 x 1)
|
||||
v = tf.squeeze(v, -1) # (batch_size x 6890 x 3)
|
||||
J_out = tf.squeeze(J_out, -1) # (batch_size x 24 x 3)
|
||||
|
||||
return v, J_out
|
||||
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
@@ -0,0 +1,175 @@
|
||||
import numpy as np
|
||||
from struct import pack, unpack
|
||||
|
||||
|
||||
"""
|
||||
Reads PC2 files, and proposed format PC16 files
|
||||
Inputs:
|
||||
- file: path to .pc2/.pc16 file
|
||||
- float16: False for PC2 files, True for PC16
|
||||
Output:
|
||||
- data: dictionary with .pc2/.pc16 file data
|
||||
NOTE: 16-bit floats lose precision with high values (positive or negative),
|
||||
we do not recommend using this format for data outside range [-2, 2]
|
||||
"""
|
||||
def readPC2(file, float16=False):
|
||||
assert file.endswith('.pc2') and not float16 or file.endswith('.pc16') and float16, 'File format not consistent with specified input format'
|
||||
data = {}
|
||||
bytes = 2 if float16 else 4
|
||||
dtype = np.float16 if float16 else np.float32
|
||||
with open(file, 'rb') as f:
|
||||
# Header
|
||||
data['sign'] = f.read(12)
|
||||
# data['version'] = int.from_bytes(f.read(4), 'little')
|
||||
data['version'] = unpack('<i', f.read(4))[0]
|
||||
# Num points
|
||||
# data['nPoints'] = int.from_bytes(f.read(4), 'little')
|
||||
data['nPoints'] = unpack('<i', f.read(4))[0]
|
||||
# Start frame
|
||||
data['startFrame'] = unpack('f', f.read(4))
|
||||
# Sample rate
|
||||
data['sampleRate'] = unpack('f', f.read(4))
|
||||
# Number of samples
|
||||
# data['nSamples'] = int.from_bytes(f.read(4), 'little')
|
||||
data['nSamples'] = unpack('<i', f.read(4))[0]
|
||||
# Animation data
|
||||
size = data['nPoints']*data['nSamples']*3*bytes
|
||||
data['V'] = np.frombuffer(f.read(size), dtype=dtype).astype(np.float32)
|
||||
data['V'] = data['V'].reshape(data['nSamples'], data['nPoints'], 3)
|
||||
|
||||
return data
|
||||
|
||||
"""
|
||||
Reads an specific frame of PC2/PC16 files
|
||||
Inputs:
|
||||
- file: path to .pc2/.pc16 file
|
||||
- frame: number of the frame to read
|
||||
- float16: False for PC2 files, True for PC16
|
||||
Output:
|
||||
- T: mesh vertex data at specified frame
|
||||
"""
|
||||
def readPC2Frame(file, frame, float16=False):
|
||||
assert file.endswith('.pc2') and not float16 or file.endswith('.pc16') and float16, 'File format not consistent with specified input format'
|
||||
assert frame >= 0 and isinstance(frame,int), 'Frame must be a positive integer'
|
||||
bytes = 2 if float16 else 4
|
||||
dtype = np.float16 if float16 else np.float32
|
||||
with open(file,'rb') as f:
|
||||
# Num points
|
||||
f.seek(16)
|
||||
# nPoints = int.from_bytes(f.read(4), 'little')
|
||||
nPoints = unpack('<i', f.read(4))[0]
|
||||
# Number of samples
|
||||
f.seek(28)
|
||||
# nSamples = int.from_bytes(f.read(4), 'little')
|
||||
nSamples = unpack('<i', f.read(4))[0]
|
||||
if frame > nSamples:
|
||||
print("Frame index outside size")
|
||||
print("\tN. frame: " + str(frame))
|
||||
print("\tN. samples: " + str(nSamples))
|
||||
return
|
||||
# Read frame
|
||||
size = nPoints * 3 * bytes
|
||||
f.seek(size * frame, 1) # offset from current '1'
|
||||
T = np.frombuffer(f.read(size), dtype=dtype).astype(np.float32)
|
||||
return T.reshape(nPoints, 3)
|
||||
|
||||
"""
|
||||
Writes PC2 and PC16 files
|
||||
Inputs:
|
||||
- file: path to file (overwrites if exists)
|
||||
- V: 3D animation data as a three dimensional array (N. Frames x N. Vertices x 3)
|
||||
- float16: False for writing as PC2 file, True for PC16
|
||||
This function assumes 'startFrame' to be 0 and 'sampleRate' to be 1
|
||||
NOTE: 16-bit floats lose precision with high values (positive or negative),
|
||||
we do not recommend using this format for data outside range [-2, 2]
|
||||
"""
|
||||
def writePC2(file, V, float16=False):
|
||||
assert file.endswith('.pc2') and not float16 or file.endswith('.pc16') and float16, 'File format not consistent with specified input format'
|
||||
if float16: V = V.astype(np.float16)
|
||||
else: V = V.astype(np.float32)
|
||||
with open(file, 'wb') as f:
|
||||
# Create the header
|
||||
headerFormat='<12siiffi'
|
||||
headerStr = pack(headerFormat, b'POINTCACHE2\0',
|
||||
1, V.shape[1], 0, 1, V.shape[0])
|
||||
f.write(headerStr)
|
||||
# Write vertices
|
||||
f.write(V.tobytes())
|
||||
|
||||
"""
|
||||
Reads OBJ files
|
||||
Only handles vertices, faces and UV maps
|
||||
Input:
|
||||
- file: path to .obj file
|
||||
Outputs:
|
||||
- V: 3D vertices
|
||||
- F: 3D faces
|
||||
- Vt: UV vertices
|
||||
- Ft: UV faces
|
||||
Correspondence between mesh and UV map is implicit in F to Ft correspondences
|
||||
If no UV map data in .obj file, it shall return Vt=None and Ft=None
|
||||
"""
|
||||
def readOBJ(file):
|
||||
V, Vt, F, Ft = [], [], [], []
|
||||
with open(file, 'r') as f:
|
||||
T = f.readlines()
|
||||
for t in T:
|
||||
# 3D vertex
|
||||
if t.startswith('v '):
|
||||
v = [float(n) for n in t.replace('v ','').split(' ')]
|
||||
V += [v]
|
||||
# UV vertex
|
||||
elif t.startswith('vt '):
|
||||
v = [float(n) for n in t.replace('vt ','').split(' ')]
|
||||
Vt += [v]
|
||||
# Face
|
||||
elif t.startswith('f '):
|
||||
idx = [n.split('/') for n in t.replace('f ','').split(' ')]
|
||||
f = [int(n[0]) - 1 for n in idx]
|
||||
F += [f]
|
||||
# UV face
|
||||
if '/' in t:
|
||||
f = [int(n[1]) - 1 for n in idx]
|
||||
Ft += [f]
|
||||
V = np.array(V, np.float32)
|
||||
Vt = np.array(Vt, np.float32)
|
||||
if Ft: assert len(F) == len(Ft), 'Inconsistent .obj file, mesh and UV map do not have the same number of faces'
|
||||
else: Vt, Ft = None, None
|
||||
return V, F, Vt, Ft
|
||||
|
||||
|
||||
|
||||
"""
|
||||
Writes OBJ files
|
||||
Only handles vertices, faces and UV maps
|
||||
Inputs:
|
||||
- file: path to .obj file (overwrites if exists)
|
||||
- V: 3D vertices
|
||||
- F: 3D faces
|
||||
- Vt: UV vertices
|
||||
- Ft: UV faces
|
||||
Correspondence between mesh and UV map is implicit in F to Ft correspondences
|
||||
If no UV map data as input, it will write only 3D data in .obj file
|
||||
"""
|
||||
def writeOBJ(file, V, F, Vt=None, Ft=None):
|
||||
if not Vt is None:
|
||||
assert len(F) == len(Ft), 'Inconsistent data, mesh and UV map do not have the same number of faces'
|
||||
|
||||
with open(file, 'w') as file:
|
||||
# Vertices
|
||||
for v in V:
|
||||
line = 'v ' + ' '.join([str(_) for _ in v]) + '\n'
|
||||
file.write(line)
|
||||
# UV verts
|
||||
if not Vt is None:
|
||||
for v in Vt:
|
||||
line = 'vt ' + ' '.join([str(_) for _ in v]) + '\n'
|
||||
file.write(line)
|
||||
# 3D Faces / UV faces
|
||||
if Ft:
|
||||
F = [[str(i+1)+'/'+str(j+1) for i,j in zip(f,ft)] for f,ft in zip(F,Ft)]
|
||||
else:
|
||||
F = [[str(i + 1) for i in f] for f in F]
|
||||
for f in F:
|
||||
line = 'f ' + ' '.join(f) + '\n'
|
||||
file.write(line)
|
||||
@@ -0,0 +1,210 @@
|
||||
import numpy
|
||||
import pygeodesic.geodesic as geodesic
|
||||
|
||||
__all__ = ['find_coefficients', 'transform']
|
||||
|
||||
|
||||
def cdist(K: numpy.ndarray, B: numpy.ndarray) -> numpy.ndarray:
|
||||
"""Calculate Euclidean distance between K[i, :] and B[j, :].
|
||||
|
||||
Arguments
|
||||
---------
|
||||
K : numpy.array
|
||||
B : numpy.array
|
||||
"""
|
||||
K = numpy.atleast_2d(K)
|
||||
B = numpy.atleast_2d(B)
|
||||
assert K.ndim == 2
|
||||
assert B.ndim == 2
|
||||
|
||||
K = numpy.expand_dims(K, 1)
|
||||
B = numpy.expand_dims(B, 0)
|
||||
D = K - B
|
||||
return numpy.linalg.norm(D, axis=2)
|
||||
|
||||
|
||||
# def cdist(K: numpy.ndarray, B: numpy.ndarray, faces: numpy.ndarray) -> numpy.ndarray:
|
||||
# geoalg = geodesic.PyGeodesicAlgorithmExact(points, faces)
|
||||
|
||||
|
||||
|
||||
def pairwise_radial_basis(K: numpy.ndarray, B: numpy.ndarray) -> numpy.ndarray:
|
||||
"""Compute the TPS radial basis function phi(r) between every row-pair of K
|
||||
and B where r is the Euclidean distance.
|
||||
|
||||
Arguments
|
||||
---------
|
||||
K : numpy.array
|
||||
n by d vector containing n d-dimensional points.
|
||||
B : numpy.array
|
||||
m by d vector containing m d-dimensional points.
|
||||
|
||||
Return
|
||||
------
|
||||
P : numpy.array
|
||||
n by m matrix where.
|
||||
P(i, j) = phi( norm( K(i,:) - B(j,:) ) ),
|
||||
where phi(r) = r^2*log(r), if r >= 1
|
||||
r*log(r^r), if r < 1
|
||||
"""
|
||||
# r_mat(i, j) is the Euclidean distance between K(i, :) and B(j, :).
|
||||
r_mat = cdist(K, B)
|
||||
|
||||
pwise_cond_ind1 = r_mat >= 1
|
||||
pwise_cond_ind2 = r_mat < 1
|
||||
r_mat_p1 = r_mat[pwise_cond_ind1]
|
||||
r_mat_p2 = r_mat[pwise_cond_ind2]
|
||||
|
||||
# P correcponds to the matrix K from [1].
|
||||
P = numpy.empty(r_mat.shape)
|
||||
P[pwise_cond_ind1] = (r_mat_p1**2) * numpy.log(r_mat_p1)
|
||||
P[pwise_cond_ind2] = r_mat_p2 * numpy.log(numpy.power(r_mat_p2, r_mat_p2))
|
||||
|
||||
return P
|
||||
|
||||
|
||||
def find_coefficients(control_points: numpy.ndarray,
|
||||
target_points: numpy.ndarray,
|
||||
lambda_: float = 0.,
|
||||
solver: str = 'exact') -> numpy.ndarray:
|
||||
"""Given a set of control points and their corresponding points, compute the
|
||||
coefficients of the TPS interpolant deforming surface.
|
||||
|
||||
Arguments
|
||||
---------
|
||||
control_points : numpy.array
|
||||
p by d vector of control points
|
||||
target_points : numpy.array
|
||||
p by d vector of corresponding target points on the deformed
|
||||
surface
|
||||
lambda_ : float
|
||||
regularization parameter
|
||||
solver : str
|
||||
the solver to get the coefficients. default is 'exact' for the exact
|
||||
solution. Or use 'lstsq' for the least square solution.
|
||||
|
||||
Return
|
||||
------
|
||||
coef : numpy.ndarray
|
||||
the coefficients
|
||||
|
||||
.. seealso::
|
||||
|
||||
http://cseweb.ucsd.edu/~sjb/pami_tps.pdf
|
||||
"""
|
||||
# ensure data type and shape
|
||||
control_points = numpy.atleast_2d(control_points)
|
||||
target_points = numpy.atleast_2d(target_points)
|
||||
if control_points.shape != target_points.shape:
|
||||
raise ValueError(
|
||||
'Shape of and control points {cp} and target points {tp} are not the same.'.
|
||||
format(cp=control_points.shape, tp=target_points.shape))
|
||||
|
||||
p, d = control_points.shape
|
||||
|
||||
# The matrix
|
||||
K = pairwise_radial_basis(control_points, control_points)
|
||||
P = numpy.hstack([numpy.ones((p, 1)), control_points])
|
||||
|
||||
# Relax the exact interpolation requirement by means of regularization.
|
||||
K = K + lambda_ * numpy.identity(p)
|
||||
|
||||
# Target points
|
||||
M = numpy.vstack([
|
||||
numpy.hstack([K, P]),
|
||||
numpy.hstack([P.T, numpy.zeros((d + 1, d + 1))])
|
||||
])
|
||||
Y = numpy.vstack([target_points, numpy.zeros((d + 1, d))])
|
||||
|
||||
# solve for M*X = Y.
|
||||
# At least d+1 control points should not be in a subspace; e.g. for d=2, at
|
||||
# least 3 points are not on a straight line. Otherwise M will be singular.
|
||||
solver = solver.lower()
|
||||
if solver == 'exact':
|
||||
X = numpy.linalg.solve(M, Y)
|
||||
elif solver == 'lstsq':
|
||||
X, _, _, _ = numpy.linalg.lstsq(M, Y, None)
|
||||
else:
|
||||
raise ValueError('Unknown solver: ' + solver)
|
||||
|
||||
return X
|
||||
|
||||
|
||||
def transform(source_points: numpy.ndarray, control_points: numpy.ndarray,
|
||||
coefficient: numpy.ndarray) -> numpy.ndarray:
|
||||
"""Transform the source points form the original surface to the destination
|
||||
(deformed) surface.
|
||||
|
||||
Arguments
|
||||
---------
|
||||
source_points : numpy.array
|
||||
n by d array of source points to be transformed
|
||||
control_points : numpy.array
|
||||
the control points used in the function `find_coefficients`
|
||||
coefficient : numpy.array
|
||||
the computed coefficients
|
||||
|
||||
Return
|
||||
------
|
||||
deformed_points : numpy.array
|
||||
n by d array of the transformed point on the target surface
|
||||
"""
|
||||
source_points = numpy.atleast_2d(source_points)
|
||||
control_points = numpy.atleast_2d(control_points)
|
||||
if source_points.shape[-1] != control_points.shape[-1]:
|
||||
raise ValueError(
|
||||
'Dimension of source points ({sd}D) and control points ({cd}D) are not the same.'.
|
||||
format(sd=source_points.shape[-1], cd=control_points.shape[-1]))
|
||||
|
||||
n = source_points.shape[0]
|
||||
|
||||
A = pairwise_radial_basis(source_points, control_points)
|
||||
K = numpy.hstack([A, numpy.ones((n, 1)), source_points])
|
||||
|
||||
deformed_points = numpy.dot(K, coefficient)
|
||||
return deformed_points
|
||||
|
||||
|
||||
|
||||
|
||||
class TPS:
|
||||
"""The thin plate spline deformation warpping.
|
||||
"""
|
||||
|
||||
def __init__(self,
|
||||
control_points: numpy.ndarray,
|
||||
target_points: numpy.ndarray,
|
||||
lambda_: float = 0.,
|
||||
solver: str = 'exact'):
|
||||
"""Create a instance that preserve the TPS coefficients.
|
||||
|
||||
Arguments
|
||||
---------
|
||||
control_points : numpy.array
|
||||
p by d vector of control points
|
||||
target_points : numpy.array
|
||||
p by d vector of corresponding target points on the deformed
|
||||
surface
|
||||
lambda_ : float
|
||||
regularization parameter
|
||||
solver : str
|
||||
the solver to get the coefficients. default is 'exact' for the
|
||||
exact solution. Or use 'lstsq' for the least square solution.
|
||||
"""
|
||||
self.control_points = control_points
|
||||
self.coefficient = find_coefficients(
|
||||
control_points, target_points, lambda_, solver)
|
||||
|
||||
def __call__(self, source_points):
|
||||
"""Transform the source points form the original surface to the
|
||||
destination (deformed) surface.
|
||||
|
||||
Arguments
|
||||
---------
|
||||
source_points : numpy.array
|
||||
n by d array of source points to be transformed
|
||||
"""
|
||||
return transform(source_points, self.control_points,
|
||||
self.coefficient)
|
||||
|
||||
transform = __call__
|
||||
@@ -0,0 +1,311 @@
|
||||
import os
|
||||
import sys
|
||||
import pickle
|
||||
import time
|
||||
import json
|
||||
import math
|
||||
import argparse
|
||||
|
||||
import numpy as np
|
||||
import tensorflow as tf
|
||||
import scipy.io as sio
|
||||
import scipy as sp
|
||||
import plotly.graph_objects as go
|
||||
from scipy.spatial.transform import Rotation
|
||||
|
||||
from tqdm import tqdm
|
||||
from datetime import datetime
|
||||
|
||||
|
||||
from custom_utils.supersmpl_io import readPC2, writePC2
|
||||
from custom_utils.smpl import SMPL
|
||||
from custom_utils.smpl_2 import SMPL as SMPL_2
|
||||
|
||||
|
||||
|
||||
os.environ["CUDA_VISIBLE_DEVICES"] = "1"
|
||||
gpus = tf.config.experimental.list_physical_devices('GPU')
|
||||
if gpus:
|
||||
try:
|
||||
for gpu in gpus:
|
||||
tf.config.experimental.set_memory_growth(gpu, True)
|
||||
|
||||
except RuntimeError as e:
|
||||
print(e)
|
||||
|
||||
|
||||
def read_garment(folder, id):
|
||||
garment_dir = os.path.join(folder, id)
|
||||
|
||||
garments = [f.replace('.pc16','') for f in os.listdir(garment_dir) if f.find('_') == -1 and f != 'info.mat' and not f.endswith('.obj')]
|
||||
assert len (garments) == 1
|
||||
|
||||
garment = garments[0]
|
||||
|
||||
posed = readPC2(os.path.join(folder, id, garment + '_registered.pc16'), True)['V']
|
||||
info = sio.loadmat(os.path.join(folder, id, 'info.mat'))
|
||||
pose, shape, gender = info['poses'].astype(np.float32), info['shape'], info['gender']
|
||||
shape = np.squeeze(np.concatenate((shape, gender), axis=1)).astype(np.float32)
|
||||
mask = np.load(os.path.join(folder, id, garment + '_mask.npy'))
|
||||
offsets = sio.loadmat(os.path.join(folder, id, garment + '_offsets_static.mat'))['label'][:,:3]
|
||||
template = sio.loadmat(os.path.join(folder, id, 'info_body.mat'))['B']
|
||||
mesh_static = (template + offsets) * mask[:, np.newaxis]
|
||||
|
||||
posed, mesh_static = posed[:, mask, :], mesh_static[mask, :]
|
||||
mask_idx = np.array(range(14475))[mask]
|
||||
|
||||
return posed, mask, pose, shape, mesh_static, mask_idx
|
||||
|
||||
def plot_3d(mesh_orig, mesh_pred, unposed, faces):
|
||||
_mesh_pred = mesh_pred.copy()
|
||||
_unposed = unposed.copy()
|
||||
_mesh_pred[:,1] += 1
|
||||
_unposed[:,1] += 3
|
||||
fig = go.Figure(data=[
|
||||
go.Mesh3d(
|
||||
# 8 vertices of a cube
|
||||
x=mesh_orig[:,0],
|
||||
y=mesh_orig[:,1],
|
||||
z=mesh_orig[:,2],
|
||||
# i, j and k give the vertices of triangles
|
||||
i = faces[:,0],
|
||||
j = faces[:,1],
|
||||
k = faces[:,2]
|
||||
),
|
||||
go.Mesh3d(
|
||||
# 8 vertices of a cube
|
||||
x=_mesh_pred[:,0],
|
||||
y=_mesh_pred[:,1],
|
||||
z=_mesh_pred[:,2],
|
||||
# i, j and k give the vertices of triangles
|
||||
i = faces[:,0],
|
||||
j = faces[:,1],
|
||||
k = faces[:,2],
|
||||
),
|
||||
go.Mesh3d(
|
||||
# 8 vertices of a cube
|
||||
x=_unposed[:,0],
|
||||
y=_unposed[:,1],
|
||||
z=_unposed[:,2],
|
||||
# i, j and k give the vertices of triangles
|
||||
i = faces[:,0],
|
||||
j = faces[:,1],
|
||||
k = faces[:,2],
|
||||
# vertexcolor=collision_color
|
||||
)
|
||||
])
|
||||
|
||||
fig.update_layout(scene_aspectmode='manual',
|
||||
scene_aspectratio=dict(x=1, y=2, z=1))
|
||||
fig.update_layout(
|
||||
scene = dict(
|
||||
xaxis = dict(nticks=4, range=[-1,1],),
|
||||
yaxis = dict(nticks=4, range=[-1,3],),
|
||||
zaxis = dict(nticks=4, range=[-1,1],),)
|
||||
)
|
||||
|
||||
return fig
|
||||
|
||||
|
||||
def unpose(mesh_static, posed, pose, shape, gender, mask, mask_idx, smpl_2, id):
|
||||
init = Rotation.from_euler('zyz', [90,90,90], degrees=True).apply(mesh_static.copy()).astype(np.float32)
|
||||
|
||||
# existing_frames = os.listdir(os.path.join(folder, id, 'mesh_unposed'))
|
||||
# if len(existing_frames) == posed.shape[0]:
|
||||
# return
|
||||
|
||||
R = smpl_2.run_rodrigues(np.transpose(pose, [1,0]))
|
||||
|
||||
frames = posed.shape[0]
|
||||
|
||||
lr_schedule = tf.keras.optimizers.schedules.ExponentialDecay(1000.0, decay_steps=40000, decay_rate=0.1, staircase=False)
|
||||
optimizer = tf.keras.optimizers.SGD(learning_rate=lr_schedule)
|
||||
unposed = tf.Variable(tf.tile(init[tf.newaxis,:,:], [frames,1,1]))
|
||||
|
||||
|
||||
@tf.function
|
||||
def skinning(verts):
|
||||
result = smpl_2.run_garment(R, shape[:10], gender, verts, filter_idx=mask_idx)
|
||||
return result
|
||||
|
||||
|
||||
@tf.function
|
||||
def error():
|
||||
pred = skinning(unposed)
|
||||
loss = tf.reduce_mean(tf.math.abs(pred - posed))
|
||||
|
||||
return loss
|
||||
|
||||
start = time.perf_counter()
|
||||
euclid_error = 10000
|
||||
|
||||
|
||||
while True:
|
||||
optimizer.minimize(error, [unposed])
|
||||
|
||||
i = optimizer.iterations.numpy()
|
||||
if i % 1000 == 0:
|
||||
pred = skinning(unposed)
|
||||
# loss = tf.reduce_mean(tf.reduce_sum(tf.math.abs(pred - posed) ** 2, axis=-1))
|
||||
loss = tf.reduce_mean(tf.math.abs(pred - posed))
|
||||
euclid_error = tf.reduce_mean(tf.math.sqrt(tf.reduce_sum((posed - pred) ** 2, axis=-1)))
|
||||
print('id: ' + id + ' iteration: ' + str(i) + ' loss: ' + str(loss.numpy()) + ' error: ' + str(euclid_error.numpy()) + ' time: ' + str(time.perf_counter() - start))
|
||||
start = time.perf_counter()
|
||||
|
||||
if euclid_error < 0.0001:
|
||||
break
|
||||
|
||||
if euclid_error < 0.0005 and i > 15000:
|
||||
break
|
||||
|
||||
if euclid_error < 0.001 and i > 30000:
|
||||
break
|
||||
|
||||
if i > 40000:
|
||||
break
|
||||
|
||||
# b = np.random.choice(np.array(range(frames)))
|
||||
# _, faces = _maskFaces(np.zeros((14475,3)), F, mask)
|
||||
# plot_3d(posed[b], pred.numpy()[b], unposed.numpy()[b], faces).show()
|
||||
|
||||
|
||||
result = np.zeros((frames, 14475, 3), dtype = np.float32)
|
||||
result[:, mask,:] = unposed.numpy()
|
||||
|
||||
return result
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("-f", "--folder", help = "end of range")
|
||||
parser.add_argument("-s", "--start", help = "start of range")
|
||||
parser.add_argument("-e", "--end", help = "end of range")
|
||||
parser.add_argument("-n", "--name", help = "end of range")
|
||||
args = parser.parse_args()
|
||||
|
||||
F = np.load('FeatureLines/superSMPL/smpl_faces.npy')
|
||||
folder = args.folder
|
||||
smpl = SMPL('FeatureLines/superSMPL')
|
||||
smpl_2 = SMPL_2('FeatureLines/superSMPL')
|
||||
iterations = 10000
|
||||
|
||||
ids = os.listdir(folder)[int(args.start):int(args.end)]
|
||||
|
||||
for id in tqdm(ids):
|
||||
if not os.path.isdir(os.path.join(folder, id)):
|
||||
continue
|
||||
|
||||
os.makedirs(os.path.join(folder, id, 'mesh_unposed'), exist_ok=True)
|
||||
|
||||
posed, mask, pose, shape, mesh_static, mask_idx = read_garment(folder, id)
|
||||
gender = int(shape[-1] > 0)
|
||||
|
||||
|
||||
unposed = unpose(mesh_static, posed, pose, shape, gender, mask, mask_idx, smpl_2)
|
||||
|
||||
for frame in range(unposed.shape[0]):
|
||||
np.savez_compressed(os.path.join(folder, id, 'mesh_unposed', id + '_' + str(frame) + '.npz'), mesh_unposed=unposed[frame])
|
||||
|
||||
|
||||
|
||||
continue
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
offset = 1
|
||||
|
||||
|
||||
for frame in range(posed.shape[0]):
|
||||
|
||||
lr_schedule = tf.keras.optimizers.schedules.ExponentialDecay(10.0, decay_steps=10000, decay_rate=0.1, staircase=False)
|
||||
optimizer = tf.keras.optimizers.SGD(learning_rate=lr_schedule)
|
||||
unposed = tf.Variable(np.random.rand(offset, posed.shape[1], 3).astype(np.float32) * 2 - 1)
|
||||
# unposed = tf.Variable(init)
|
||||
|
||||
@tf.function
|
||||
def skinning(verts):
|
||||
# verts = smpl.run(R[frame:frame+offset], tf.tile(shape[tf.newaxis,:], [offset,1]))[0].numpy()[:,mask,:]
|
||||
|
||||
|
||||
# start = time.perf_counter()
|
||||
result = smpl.run_garment(R[frame:frame+offset], tf.tile(shape[tf.newaxis,:], [offset,1]), verts, filter_idx=mask_idx)
|
||||
# print(time.perf_counter() - start)
|
||||
|
||||
# start = time.perf_counter()
|
||||
# result_2 = smpl_2.run_garment(R[frame:frame+offset], shape[:10], gender, verts, filter_idx=mask_idx)
|
||||
# print(time.perf_counter() - start)
|
||||
|
||||
# test = np.abs(result.numpy() - result_2.numpy()) > 1e-4
|
||||
# x = np.sum(test, axis=-1).sum(-1).sum(-1)
|
||||
# euclid_error = tf.reduce_mean(tf.math.sqrt(tf.reduce_sum((result - result_2) ** 2, axis=-1)))
|
||||
|
||||
# batch = 25
|
||||
# _, faces = _maskFaces(np.zeros((14475,3)), F, mask)
|
||||
# go.Figure(data=[
|
||||
# go.Mesh3d(
|
||||
# # 8 vertices of a cube
|
||||
# x=result[batch,:,0],
|
||||
# y=result[batch,:,1],
|
||||
# z=result[batch,:,2],
|
||||
# # i, j and k give the vertices of triangles
|
||||
# i = faces[:,0],
|
||||
# j = faces[:,1],
|
||||
# k = faces[:,2]
|
||||
# ),
|
||||
# go.Mesh3d(
|
||||
# # 8 vertices of a cube
|
||||
# x=result_2[batch,:,0],
|
||||
# y=result_2[batch,:,1],
|
||||
# z=result_2[batch,:,2],
|
||||
# # i, j and k give the vertices of triangles
|
||||
# i = faces[:,0],
|
||||
# j = faces[:,1],
|
||||
# k = faces[:,2],
|
||||
# )]).show()
|
||||
|
||||
return result
|
||||
|
||||
@tf.function
|
||||
def error():
|
||||
pred = skinning(unposed)
|
||||
loss = tf.reduce_mean(tf.math.abs(pred - posed[frame]))
|
||||
|
||||
return loss
|
||||
|
||||
start = time.perf_counter()
|
||||
euclid_error = 10000
|
||||
|
||||
for i in range(iterations):
|
||||
optimizer.minimize(error, [unposed])
|
||||
if i % 50 == 0:
|
||||
pred = skinning(unposed)
|
||||
# loss = tf.reduce_mean(tf.reduce_sum(tf.math.abs(pred - posed) ** 2, axis=-1))
|
||||
loss = tf.reduce_mean(tf.math.abs(pred - posed[frame]))
|
||||
euclid_error = tf.reduce_mean(tf.math.sqrt(tf.reduce_sum((posed[frame] - pred) ** 2, axis=-1)))
|
||||
print('id: ' + id + ' frame: ' + str(frame) + ' iteration: ' + str(i) + ' loss: ' + str(loss.numpy()) + ' error: ' + str(euclid_error.numpy()) + ' time: ' + str(time.perf_counter() - start))
|
||||
start = time.perf_counter()
|
||||
|
||||
if euclid_error < 0.001:
|
||||
break
|
||||
|
||||
result = np.zeros((14475, 3), dtype = np.float32)
|
||||
result[mask,:] = unposed.numpy()
|
||||
np.savez_compressed(os.path.join(folder, id, 'mesh_unposed', id + '_' + str(frame) + '.npz'), mesh_unposed=result)
|
||||
|
||||
init = unposed
|
||||
|
||||
# pred, unposed = pred.numpy(), unposed.numpy()
|
||||
# _, faces = _maskFaces(np.zeros((14475,3)), F, mask)
|
||||
# k = 0
|
||||
# plot_3d(posed[k], pred[k], unposed[k], faces).show()
|
||||
|
||||
# print()
|
||||
|
||||
|
||||
|
||||
Binary file not shown.
Reference in New Issue
Block a user