reaarrange files
This commit is contained in:
+11
@@ -160,3 +160,14 @@ cython_debug/
|
||||
# option (not recommended) you can uncomment the following to ignore the entire idea folder.
|
||||
#.idea/
|
||||
|
||||
wandb/
|
||||
output/*
|
||||
*.zip
|
||||
*.out
|
||||
temp.jpg
|
||||
corrupt/
|
||||
checkpoints/
|
||||
temp/
|
||||
ids/
|
||||
examples/
|
||||
custom_data/dijkstra*
|
||||
|
||||
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,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.
+610
@@ -0,0 +1,610 @@
|
||||
import os
|
||||
import sys
|
||||
import time
|
||||
import cv2
|
||||
import random
|
||||
import numpy as np
|
||||
import tensorflow as tf
|
||||
|
||||
from custom_utils.custom_utils import nest_conditions, batch_to_dict
|
||||
|
||||
|
||||
|
||||
# class DatasetSimple(tf.data.Dataset):
|
||||
|
||||
# def _generator(folder):
|
||||
# # Opening the file
|
||||
# folder = folder.decode("utf-8")
|
||||
# uv_folder = os.path.join(folder, 'uv_posed')
|
||||
# pose_hape_folder = os.path.join(folder, 'poses_shapes')
|
||||
# frames = os.listdir(uv_folder)
|
||||
# random.shuffle(frames)
|
||||
|
||||
# for frame in frames:
|
||||
# id, frame_nr = frame.split('.')[0].split('_')
|
||||
# uv_posed = np.load(os.path.join(folder, 'uv_posed', id + '_' + frame_nr + '.npz')); uv_posed = uv_posed['uv_gt']; uv_posed = np.transpose(uv_posed, [2, 0, 1])
|
||||
# uv_body_posed = np.load(os.path.join(folder, 'uv_body_posed', id + '_' + frame_nr + '.npz')); uv_body_posed = uv_body_posed['uv_body']; uv_body_posed = np.transpose(uv_body_posed, [2, 0, 1])
|
||||
# pose_shape = np.load(os.path.join(folder, 'pose_shape', id + '_' + frame_nr + '.npy'))
|
||||
# uv_static = np.load(os.path.join(folder, 'uv_static', id + '.npz')); uv_static = uv_static['uv_static']; uv_static = np.transpose(uv_static, [2, 0, 1])
|
||||
# mesh_static = np.load(os.path.join(folder, 'mesh_static', id + '.npz')); mesh_static = mesh_static['mesh_static']
|
||||
# mesh_body_posed = np.load(os.path.join(folder, 'mesh_body_posed', id + '_' + frame_nr + '.npy'))
|
||||
# mesh_mask = np.load(os.path.join(folder, 'mesh_mask', id + '.npy')); mesh_mask = mesh_mask.astype(np.float32)
|
||||
# uv_mask = np.load(os.path.join(folder, 'uv_mask', id + '.npz')); uv_mask=uv_mask['uv_mask']; uv_mask = uv_mask.astype(np.float32)
|
||||
# mesh_posed = np.load(os.path.join(folder, 'mesh_posed', id + '_' + frame_nr + '.npz')); mesh_posed = mesh_posed['mesh_posed']
|
||||
# mesh_unposed = np.load(os.path.join(folder, 'mesh_unposed', id + '_' + frame_nr + '.npz')); mesh_unposed = mesh_unposed['mesh_unposed']
|
||||
# uv_unposed = np.load(os.path.join(folder, 'uv_unposed', id + '_' + frame_nr + '.npz')); uv_unposed = uv_unposed['uv_unposed']; uv_unposed = np.transpose(uv_unposed, [2, 0, 1])
|
||||
|
||||
# # uv_posed, (pose, shape, uv_static), mesh_static, mesh_body_posed, mask
|
||||
# yield uv_posed, (pose_shape[:72], pose_shape[72:], uv_static, uv_mask, uv_body_posed), mesh_static, mesh_body_posed, mesh_mask, mesh_posed, mesh_unposed, uv_unposed
|
||||
|
||||
# def __new__(cls, folder, res=256):
|
||||
# return tf.data.Dataset.from_generator(
|
||||
# cls._generator,
|
||||
# output_signature = (tf.TensorSpec(shape = (3, res * 2, res,), dtype = tf.float32),
|
||||
# (tf.TensorSpec(shape = (72), dtype = tf.float32), tf.TensorSpec(shape = (11), dtype = tf.float32), tf.TensorSpec(shape = (3, res * 2, res,), dtype = tf.float32), tf.TensorSpec(shape = (res * 2, res,), dtype = tf.float32), tf.TensorSpec(shape = (3, res * 2, res,), dtype = tf.float32) ),
|
||||
# tf.TensorSpec(shape = (14475, 3,), dtype = tf.float32),
|
||||
# tf.TensorSpec(shape = (14475, 3,), dtype = tf.float32),
|
||||
# tf.TensorSpec(shape = (14475,), dtype = tf.float32),
|
||||
# tf.TensorSpec(shape = (14475, 3,), dtype = tf.float32),
|
||||
# tf.TensorSpec(shape = (14475, 3,), dtype = tf.float32),
|
||||
# tf.TensorSpec(shape = (3, res * 2, res,), dtype = tf.float32), ),
|
||||
# args=(folder, )
|
||||
# )
|
||||
|
||||
|
||||
|
||||
class StaticDataset():
|
||||
def __init__(self, folder, signature, use_filled=False, resolutions=[512], in_memory=True, is_cape=False) -> None:
|
||||
self.store = dict()
|
||||
self.signature = signature
|
||||
self.folder = folder
|
||||
self.in_memory = in_memory
|
||||
assert set(resolutions).issubset(set([32, 64, 128, 256, 512]))
|
||||
self.possible_keys = ['uv_static', 'uv_static_32', 'uv_static_64', 'uv_static_128', 'uv_static_256',
|
||||
'uv_mask', 'uv_mask_32', 'uv_mask_64', 'uv_mask_128', 'uv_mask_256',
|
||||
'mesh_static', 'mesh_mask', 'uv_shaped_body']
|
||||
self.keys = [value for value in signature.keys() if value in self.possible_keys]
|
||||
if 'uv_static_128' in signature.keys() and 'uv_static_128' not in self.keys:
|
||||
self.keys.append('uv_static_128')
|
||||
self.resolutions = resolutions
|
||||
self.filled = ''
|
||||
self.is_cape = is_cape
|
||||
if use_filled:
|
||||
self.filled = '_filled'
|
||||
|
||||
if is_cape:
|
||||
self.uv_masks = dict()
|
||||
uv_mask = np.load('custom_data/garment_uv_mask_full.npy').astype(np.int32)
|
||||
self.uv_masks[512] = uv_mask.astype(np.float32)
|
||||
for res in [32, 64, 128, 256]:
|
||||
down = tf.image.resize(uv_mask[:,:,tf.newaxis], size=(res, int(res/2)), method=tf.image.ResizeMethod.NEAREST_NEIGHBOR)
|
||||
down = np.squeeze(down.numpy().astype(np.float32))
|
||||
self.uv_masks[res] = down
|
||||
|
||||
|
||||
|
||||
def read(self, id):
|
||||
data = dict()
|
||||
|
||||
for res in self.resolutions:
|
||||
res_string = ''
|
||||
if res != 512:
|
||||
res_string = '_' + str(res)
|
||||
|
||||
if 'uv_static' + res_string in self.signature.keys():
|
||||
uv_static = np.load(os.path.join(self.folder, 'uv_static' + self.filled + res_string, id + '.npz')); uv_static = uv_static['uv_static' + self.filled]; data['uv_static' + res_string] = uv_static
|
||||
if 'uv_mask' + res_string in self.signature.keys():
|
||||
if self.is_cape:
|
||||
data['uv_mask' + res_string] = self.uv_masks[res]
|
||||
else:
|
||||
uv_mask = np.load(os.path.join(self.folder, 'uv_mask' + res_string, id + '.npz')); uv_mask=uv_mask['uv_mask']; uv_mask = uv_mask.astype(np.float32); data['uv_mask' + res_string] = uv_mask
|
||||
|
||||
|
||||
if 'mesh_static' in self.signature.keys():
|
||||
mesh_static = np.load(os.path.join(self.folder, 'mesh_static', id + '.npz')); mesh_static = mesh_static['mesh_static']; data['mesh_static'] = mesh_static.astype(np.float32)
|
||||
if 'mesh_mask' in self.signature.keys():
|
||||
if self.is_cape:
|
||||
data['mesh_mask'] = np.ones((3638), np.float32)
|
||||
else:
|
||||
mesh_mask = np.load(os.path.join(self.folder, 'mesh_mask', id + '.npy')); mesh_mask = mesh_mask.astype(np.float32); data['mesh_mask'] = mesh_mask
|
||||
if 'uv_shaped_body' in self.signature.keys():
|
||||
uv_shaped_body = np.load(os.path.join(self.folder, 'uv_shaped_body', id + '.npz')); uv_shaped_body = uv_shaped_body['uv_shaped_body']; data['uv_shaped_body'] = uv_shaped_body
|
||||
|
||||
return data
|
||||
|
||||
def __getitem__(self, id) -> dict:
|
||||
if id not in self.store:
|
||||
item = self.read(id)
|
||||
if not self.in_memory:
|
||||
return item
|
||||
else:
|
||||
self.store[id] = item
|
||||
|
||||
return self.store[id]
|
||||
|
||||
def expand_batch(self, batch):
|
||||
for key in self.keys:
|
||||
batch[key] = tf.zeros(((0,) + self.signature[key]))
|
||||
for i in range(batch['id'].shape[0]):
|
||||
item = self.__getitem__(batch['id'][i].numpy().decode('UTF-8'))
|
||||
for key in self.keys:
|
||||
batch[key] = tf.concat((batch[key], tf.expand_dims(item[key], axis=0)), axis=0)
|
||||
return batch
|
||||
|
||||
|
||||
|
||||
class DatasetIdx(tf.data.Dataset):
|
||||
|
||||
def __new__(cls, folder, shuffle=True, is_cape=False, testing=False, frames_file=None):
|
||||
if frames_file is None:
|
||||
uv_static_folder = os.path.join(folder, 'pose_shape')
|
||||
frames = sorted(os.listdir(uv_static_folder))
|
||||
|
||||
if is_cape:
|
||||
frames = ['.'.join(frame.split('.')[:-1]) for frame in frames] # .split('_')[0]
|
||||
else:
|
||||
frames = [frame.split('.')[0] for frame in frames] # .split('_')[0]
|
||||
|
||||
if testing:
|
||||
frames = sorted(os.listdir(os.path.join(folder, 'mesh_static') ))
|
||||
frames = [frame.split('.')[0] + '_4' for frame in frames]
|
||||
|
||||
if frames_file is not None:
|
||||
frames = []
|
||||
with open(os.path.join(folder, frames_file), 'r') as f:
|
||||
lines = f.readlines()
|
||||
lines = [line.strip() for line in lines]
|
||||
frames = lines
|
||||
|
||||
if shuffle:
|
||||
random.shuffle(frames)
|
||||
|
||||
return tf.data.Dataset.from_tensor_slices(frames)
|
||||
|
||||
|
||||
# class DatasetFactory(tf.data.Dataset):
|
||||
class DatasetFactory():
|
||||
|
||||
def __init__(self):
|
||||
super(DatasetFactory, self).__init__()
|
||||
|
||||
def _build_tf_signature(self, signature):
|
||||
tf_signature = []
|
||||
for key in signature.keys():
|
||||
if key == 'id':
|
||||
tf_signature.append(tf.TensorSpec(shape = (), dtype = tf.string))
|
||||
elif key == 'frame_nr':
|
||||
tf_signature.append(tf.TensorSpec(shape = (), dtype = tf.string))
|
||||
elif key == 'mesh_curvature':
|
||||
tf_signature.append(tf.TensorSpec(shape = signature[key], dtype = tf.int32))
|
||||
else:
|
||||
tf_signature.append(tf.TensorSpec(shape = signature[key], dtype = tf.float32))
|
||||
return tf_signature
|
||||
|
||||
|
||||
def get_dataset(self,
|
||||
folder,
|
||||
batch_size,
|
||||
signature,
|
||||
shuffle=True,
|
||||
resolutions=[512],
|
||||
in_memory=False,
|
||||
do_repeat=True,
|
||||
is_cape=False,
|
||||
use_one_hot=False,
|
||||
use_no_orient=True,
|
||||
testing=False,
|
||||
frames_file=None):
|
||||
assert set(resolutions).issubset(set([32, 64, 128, 256, 512]))
|
||||
|
||||
dataset_idx = DatasetIdx(folder, shuffle, is_cape=is_cape, testing=testing, frames_file=frames_file)
|
||||
|
||||
separator = '_'
|
||||
if is_cape:
|
||||
separator = '.'
|
||||
|
||||
one_hot = ''
|
||||
one_hot_2 = ''
|
||||
if use_one_hot:
|
||||
one_hot = '_one_hot'
|
||||
one_hot_2 = '_one_hot_2'
|
||||
|
||||
uv_posed_body_no_orient_folder_name = 'uv_posed_body_no_orient_2'
|
||||
if is_cape:
|
||||
uv_posed_body_no_orient_folder_name = 'uv_posed_body_no_orient'
|
||||
|
||||
if use_no_orient == False:
|
||||
assert resolutions == [512], 'Using global orientation is only possible with non-pyramid model!'
|
||||
|
||||
|
||||
def read_item(frame):
|
||||
frame_split = tf.strings.split(frame, sep=separator)
|
||||
id, frame_nr = frame_split[0].numpy().decode("utf-8"), frame_split[1].numpy().decode("utf-8")
|
||||
# id, frame_nr = frame_split[0].numpy(), frame_split[1].numpy()
|
||||
# static_info = static_data[id]
|
||||
item = dict()
|
||||
|
||||
if 'id' in signature.keys():
|
||||
item['id'] = id
|
||||
if 'frame_nr' in signature.keys():
|
||||
item['frame_nr'] = frame_nr
|
||||
if 'pose_shape' in signature.keys():
|
||||
pose_shape = np.load(os.path.join(folder, 'pose_shape', id + separator + frame_nr + '.npy'))
|
||||
if use_no_orient:
|
||||
pose_shape[:3] = 0
|
||||
item['pose_shape'] = pose_shape
|
||||
if 'mesh_body_posed' in signature.keys():
|
||||
if use_no_orient:
|
||||
mesh_posed_body_no_orient = np.load(os.path.join(folder, 'mesh_posed_body_no_orient', id + separator + frame_nr + '.npz')); mesh_posed_body_no_orient = mesh_posed_body_no_orient['mesh_posed_body_no_orient']
|
||||
item['mesh_body_posed'] = mesh_posed_body_no_orient
|
||||
else:
|
||||
mesh_body_posed = np.load(os.path.join(folder, 'mesh_body_posed', id + '_' + frame_nr + '.npy'))
|
||||
item['mesh_body_posed'] = mesh_body_posed
|
||||
if 'mesh_posed' in signature.keys():
|
||||
if use_no_orient:
|
||||
mesh_posed = np.load(os.path.join(folder, 'mesh_posed_no_orient', id + separator + frame_nr + '.npz')); mesh_posed = mesh_posed['mesh_posed_no_orient'].astype(np.float32)
|
||||
item['mesh_posed'] = mesh_posed
|
||||
else:
|
||||
mesh_posed = np.load(os.path.join(folder, 'mesh_posed', id + separator + frame_nr + '.npz')); mesh_posed = mesh_posed['mesh_posed'].astype(np.float32)
|
||||
item['mesh_posed'] = mesh_posed
|
||||
# if 'mesh_unposed' in signature.keys():
|
||||
# mesh_unposed = np.load(os.path.join(folder, 'mesh_unposed', id + '_' + frame_nr + '.npz')); mesh_unposed = mesh_unposed['mesh_unposed']
|
||||
# item['mesh_unposed'] = mesh_unposed
|
||||
# if 'uv_unposed' in signature.keys():
|
||||
# uv_unposed = np.load(os.path.join(folder, 'uv_unposed', id + '_' + frame_nr + '.npz')); uv_unposed = uv_unposed['uv_unposed']; uv_unposed = np.transpose(uv_unposed, [2, 0, 1])
|
||||
# item['uv_unposed'] = uv_unposed
|
||||
if 'mesh_unposed' in signature.keys():
|
||||
# print('one_hot')
|
||||
mesh_unposed_no_shape = np.load(os.path.join(folder, 'mesh_unposed_no_shape' + one_hot, id + separator + frame_nr + '.npz')); mesh_unposed_no_shape = mesh_unposed_no_shape['mesh_unposed_no_shape']
|
||||
item['mesh_unposed'] = mesh_unposed_no_shape
|
||||
|
||||
|
||||
if 'mesh_curvature' in signature.keys():
|
||||
mesh_curvature = np.load(os.path.join(folder, 'mesh_curvature', id + separator + frame_nr + '.npz')); mesh_curvature = mesh_curvature['mesh_curvature']
|
||||
item['mesh_curvature'] = mesh_curvature.astype(np.int32)
|
||||
|
||||
|
||||
for res in resolutions:
|
||||
res_string = ''
|
||||
if res != 512:
|
||||
res_string = '_' + str(res)
|
||||
|
||||
if 'uv_curvature' + res_string in signature.keys():
|
||||
uv_curvature = np.load(os.path.join(folder, 'uv_curvature' + res_string, id + separator + frame_nr + '.npz'));
|
||||
if res_string != '':
|
||||
uv_curvature = uv_curvature['uv_curvature']
|
||||
else:
|
||||
uv_curvature = uv_curvature['curvature_uv']
|
||||
uv_curvature = tf.one_hot(uv_curvature, 7)
|
||||
item['uv_curvature' + res_string] = uv_curvature
|
||||
|
||||
if 'uv_unposed' + res_string in signature.keys():
|
||||
uv_unposed_no_shape = np.load(os.path.join(folder, 'uv_unposed_no_shape' + one_hot_2 + res_string, id + separator + frame_nr + '.npz')); uv_unposed_no_shape = uv_unposed_no_shape['uv_unposed_no_shape']
|
||||
item['uv_unposed' + res_string] = uv_unposed_no_shape
|
||||
|
||||
if 'uv_body_posed' + res_string in signature.keys():
|
||||
if use_no_orient:
|
||||
uv_posed_body_no_orient = np.load(os.path.join(folder, uv_posed_body_no_orient_folder_name + res_string, id + separator + frame_nr + '.npz'))
|
||||
uv_posed_body_no_orient = uv_posed_body_no_orient['uv_posed_body_no_orient']
|
||||
item['uv_body_posed' + res_string] = uv_posed_body_no_orient
|
||||
else:
|
||||
uv_body_posed = np.load(os.path.join(folder, 'uv_body_posed', id + '_' + frame_nr + '.npz')); uv_body_posed = uv_body_posed['uv_body']
|
||||
item['uv_body_posed' + res_string] = uv_body_posed
|
||||
|
||||
if 'uv_posed' + res_string in signature.keys():
|
||||
if use_no_orient:
|
||||
uv_posed = np.load(os.path.join(folder, 'uv_posed_no_orient' + res_string, id + separator + frame_nr + '.npz')); uv_posed = uv_posed['uv_posed_no_orient']
|
||||
item['uv_posed' + res_string] = uv_posed
|
||||
else:
|
||||
uv_posed = np.load(os.path.join(folder, 'uv_posed' + res_string, id + separator + frame_nr + '.npz')); uv_posed = uv_posed['uv_gt']
|
||||
item['uv_posed' + res_string] = uv_posed
|
||||
|
||||
# print('here')
|
||||
|
||||
# item = { **static_info, **item }
|
||||
as_list = []
|
||||
for key in signature.keys():
|
||||
if key in item:
|
||||
as_list.append(item[key])
|
||||
|
||||
return as_list
|
||||
|
||||
# dataset = dataset_idx.map(lambda i: tf.py_function(func=read_item,
|
||||
# inp=[i],
|
||||
# Tout=self._build_tf_signature(signature)
|
||||
# ), num_parallel_calls=tf.data.AUTOTUNE).cache().batch(batch_size=batch_size, drop_remainder=True).repeat()
|
||||
|
||||
if in_memory:
|
||||
dataset = dataset_idx.map(lambda i: tf.py_function(func=read_item,
|
||||
inp=[i],
|
||||
Tout=self._build_tf_signature(signature)
|
||||
), num_parallel_calls=tf.data.AUTOTUNE).cache().batch(batch_size=batch_size, drop_remainder=True)
|
||||
else:
|
||||
dataset = dataset_idx.map(lambda i: tf.py_function(func=read_item,
|
||||
inp=[i],
|
||||
Tout=self._build_tf_signature(signature)
|
||||
), num_parallel_calls=tf.data.AUTOTUNE).batch(batch_size=batch_size, drop_remainder=True)
|
||||
|
||||
if do_repeat:
|
||||
dataset = dataset.repeat()
|
||||
|
||||
# dataset = dataset_idx.map(read_item, num_parallel_calls=tf.data.AUTOTUNE).cache().batch(batch_size=batch_size, drop_remainder=True).repeat()
|
||||
|
||||
return dataset
|
||||
|
||||
# def __new__(cls, folder):
|
||||
# return tf.data.Dataset.from_generator(
|
||||
# cls._generator,
|
||||
# # output_signature = (tf.TensorSpec(shape = (), dtype = tf.string), ),
|
||||
# output_types=(tf.string),
|
||||
# args=(folder, )
|
||||
# )
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
os.environ["CUDA_VISIBLE_DEVICES"] = "6"
|
||||
from custom_utils.custom_utils import batch_to_dict, UV2Mesh, _maskFaces
|
||||
import plotly.graph_objects as go
|
||||
import plotly.express as ex
|
||||
# import cmapy as cm
|
||||
uv2mesh = UV2Mesh(is_skirt=True)
|
||||
|
||||
F = np.load("custom_utils/smpl_faces.npy")
|
||||
|
||||
|
||||
# dataset = DatasetIdx('/home/workspace/data/dataset/stylegan_new/train')
|
||||
# for d in dataset:
|
||||
# print(d)
|
||||
|
||||
factory = DatasetFactory()
|
||||
signature_dynamic = { 'id': (), 'frame_nr': (),
|
||||
# 'pose_shape': (72 + 11,),
|
||||
# 'uv_posed': (3, 512, 256,),
|
||||
# 'uv_body_posed': (512, 256, 3),
|
||||
# 'uv_body_posed_32': (32, 16, 3),
|
||||
# 'mesh_body_posed': (14475, 3,),
|
||||
# 'mesh_posed': (14475, 3,),
|
||||
# 'mesh_unposed': (14475, 3,),
|
||||
'uv_unposed': (512, 256, 3),
|
||||
# 'uv_unposed_32': (3, 32, 16,),
|
||||
'uv_curvature': (512, 256, 7),
|
||||
# 'mesh_curvature': (14475)
|
||||
}
|
||||
signature_static = { #'uv_static': (3, 512, 256,),
|
||||
# 'uv_static_32': (3, 32, 16,),
|
||||
# 'mesh_static': (14475, 3,),
|
||||
# 'mesh_mask': (14475,),
|
||||
'uv_mask': (512, 256,),
|
||||
# 'uv_shaped_body': (3, 512, 256,),
|
||||
}
|
||||
|
||||
dataset = factory.get_dataset(folder='/home/workspace/local-data/train_skirt',
|
||||
batch_size=1,
|
||||
signature=signature_dynamic,
|
||||
resolutions=[512],
|
||||
do_repeat=False,
|
||||
use_one_hot=False,
|
||||
is_cape=False,
|
||||
frames_file='frames_dress_30k.txt')
|
||||
static_dataset = StaticDataset(folder='/home/workspace/local-data/train_skirt',
|
||||
signature=signature_static,
|
||||
resolutions=[512],
|
||||
is_cape=False)
|
||||
|
||||
|
||||
# for i, batch in enumerate(dataset):
|
||||
# print(i)
|
||||
|
||||
|
||||
sum_uv = np.zeros((512, 256, 3), dtype=np.float32)
|
||||
sum_mask = np.zeros((512, 256), dtype=np.float32)
|
||||
for i, batch in enumerate(dataset):
|
||||
print(i)
|
||||
batch = static_dataset.expand_batch(batch_to_dict(batch, signature_dynamic))
|
||||
ex.imshow(batch['uv_curvature'][0].numpy().argmax(-1) / 6.0).show()
|
||||
sum_uv += batch['uv_unposed'].numpy().sum(axis=0)
|
||||
sum_mask += batch['uv_mask'].numpy().sum(axis=0)
|
||||
|
||||
sum_mask[sum_mask == 0] = 1
|
||||
avg_garment = sum_uv / sum_mask[:, :, np.newaxis]
|
||||
|
||||
# x = np.transpose(avg_garment, [1,2,0])
|
||||
x = uv2mesh(avg_garment[tf.newaxis,:,:,:])[0]
|
||||
go.Figure(data=[
|
||||
go.Mesh3d(
|
||||
# 8 vertices of a cube
|
||||
x=x[:,0],
|
||||
y=x[:,1],
|
||||
z=x[:,2],
|
||||
# i, j and k give the vertices of triangles
|
||||
i = F[:,0],
|
||||
j = F[:,1],
|
||||
k = F[:,2]
|
||||
)]).show()
|
||||
|
||||
np.save('custom_data/train_dress_30k_no_shape_avg_uv.npy', avg_garment)
|
||||
|
||||
print()
|
||||
|
||||
# pass
|
||||
|
||||
|
||||
# uv_res = 256
|
||||
|
||||
# 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
|
||||
|
||||
# def unique_uv_map_mask(uv_map_discreet):
|
||||
# mask = np.zeros((uv_map_discreet.shape[0]), dtype=bool)
|
||||
# temp = np.zeros((uv_res * 2, uv_res))
|
||||
# for p in uv_map_discreet:
|
||||
# temp[p[0], p[1]] += 1
|
||||
# # temp = temp == 1
|
||||
# for i, p in enumerate(uv_map_discreet):
|
||||
# if temp[p[0], p[1]] == 1:
|
||||
# mask[i] = True
|
||||
# elif temp[p[0], p[1]] == 2:
|
||||
# mask[i] = True
|
||||
# temp[p[0], p[1]] = 0
|
||||
# return mask
|
||||
|
||||
# def uv_to_mesh(uv, m2uv):
|
||||
# res = uv.shape[0]
|
||||
# mesh = np.zeros((14475, 3))
|
||||
# uv_pixels = uv_map_discreet[m2uv]
|
||||
# # counter = np.ones((14475))
|
||||
# # counter[m2uv[:, 0] == m2uv[:, 1]] = 2
|
||||
# mesh = np.mean(uv[uv_pixels[:, :, 0], uv_pixels[:, :, 1], :], axis=1) # / counter[:, np.newaxis]
|
||||
# return mesh
|
||||
|
||||
|
||||
# uv_map = np.load('custom_utils/stretched_uv_map_3.npy')
|
||||
# m2uv = np.load('custom_utils/m2uv.npy', allow_pickle=True).item()
|
||||
# m2uv = m2uv_array(m2uv)
|
||||
# uv_mask = np.load('custom_utils/mask_256x512.npy')
|
||||
# curvature_label_type = '6_25'
|
||||
|
||||
# uv_res = 256
|
||||
# uv_mask = (cv2.resize(uv_mask * 255.0, (uv_res, uv_res * 2)) / 255).astype(bool)[:, :, np.newaxis]
|
||||
# uv_mask = cv2.dilate(uv_mask.astype(np.float32), cv2.getStructuringElement(cv2.MORPH_RECT, (2 * 2 + 1, 2 * 2 + 1), (2, 2)), iterations=2).astype(bool)[:, :, np.newaxis]
|
||||
|
||||
# uv_map_discreet = np.floor(uv_map * uv_res).astype(np.int32)
|
||||
# uv_map_mask = unique_uv_map_mask(uv_map_discreet)
|
||||
# uv_counter = np.zeros((uv_map.shape[0]))
|
||||
# uv_counter[m2uv[:, 0]] = 1
|
||||
# uv_counter[m2uv[:, 1]] += 1
|
||||
# uv_counter = uv_counter[uv_map_mask]
|
||||
|
||||
|
||||
|
||||
# for i, batch in enumerate(dataset):
|
||||
# print(i)
|
||||
# batch = static_dataset.expand_batch(batch_to_dict(batch, signature_dynamic))
|
||||
|
||||
# uv_unposed = np.squeeze(batch['uv_unposed'].numpy())
|
||||
# mesh_unposed = np.squeeze(batch['mesh_unposed'].numpy())
|
||||
# mask = np.squeeze(batch['mesh_mask'].numpy()).astype(bool)
|
||||
|
||||
|
||||
# values = mesh_unposed
|
||||
# uv_values = np.zeros((uv_map.shape[0], 3))
|
||||
# uv_values[m2uv[:, 0]] = values
|
||||
# uv_values[m2uv[:, 1]] += values
|
||||
# uv_counter = np.zeros((uv_map.shape[0]))
|
||||
# uv_counter[m2uv[:, 0]] = 1
|
||||
# uv_counter[m2uv[:, 1]] += 1
|
||||
# uv_values /= uv_counter[:, np.newaxis]
|
||||
# x = uv_values
|
||||
# # verts = np.concatenate((verts, np.zeros((15068, 1))), axis=-1)
|
||||
# im = ((x - x.min()) / (x.max() - x.min())) * 255
|
||||
|
||||
|
||||
# x = np.transpose(uv_unposed, [1,2,0])
|
||||
# x = uv2mesh(uv_unposed[tf.newaxis,:,:,:])[0].numpy()[mask, :]
|
||||
# _, faces = _maskFaces(np.zeros((14475, 3), dtype=np.float32), F, mask)
|
||||
# go.Figure(data=[
|
||||
# go.Mesh3d(
|
||||
# # 8 vertices of a cube
|
||||
# x=x[:,0],
|
||||
# y=x[:,1],
|
||||
# z=x[:,2],
|
||||
# # i, j and k give the vertices of triangles
|
||||
# i = faces[:,0],
|
||||
# j = faces[:,1],
|
||||
# k = faces[:,2]
|
||||
# )]).show()
|
||||
|
||||
|
||||
# print()
|
||||
|
||||
|
||||
|
||||
|
||||
# # @tf.function
|
||||
# def euclid_error(mask, gt, pred):
|
||||
# # return tf.reduce_sum(tf.math.sqrt(tf.reduce_sum(mask[:,:,tf.newaxis] * (gt - pred) ** 2, axis=-1)), axis=-1) / tf.reduce_sum(mask, axis=-1)
|
||||
# return tf.math.sqrt(tf.reduce_sum(mask[:,:,tf.newaxis] * (gt - pred) ** 2, axis=-1))
|
||||
|
||||
# all_error = []
|
||||
# for batch in dataset:
|
||||
# batch = static_dataset.expand_batch(batch_to_dict(batch, signature_dynamic))
|
||||
# mesh_reconstruct = uv2mesh(batch['uv_posed'])
|
||||
# err = euclid_error(batch['mesh_mask'], batch['mesh_posed'], uv2mesh(batch['uv_posed']))
|
||||
|
||||
# i = 296
|
||||
# err_normed = (err[i].numpy() - err[i].numpy().min()) / (err[i].numpy().max() - err[i].numpy().min())
|
||||
# colors = np.array([cm.color('jet', 1 - i) for i in err_normed])
|
||||
|
||||
# _, faces = _maskFaces(batch['mesh_posed'][i].numpy(), F, batch['mesh_mask'][i].numpy().astype(bool))
|
||||
# x = batch['mesh_posed'][i].numpy()[batch['mesh_mask'][i].numpy().astype(bool), :]
|
||||
# y = mesh_reconstruct[i].numpy()[batch['mesh_mask'][i].numpy().astype(bool), :]
|
||||
# colors = colors[batch['mesh_mask'][i].numpy().astype(bool), :]
|
||||
# fig = go.Figure(data=[
|
||||
# go.Mesh3d(
|
||||
# # 8 vertices of a cube
|
||||
# x=x[:,0],
|
||||
# y=x[:,1],
|
||||
# z=x[:,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=y[:,0],
|
||||
# y=y[:,1],
|
||||
# z=y[:,2],
|
||||
# # i, j and k give the vertices of triangles
|
||||
# i = faces[:,0],
|
||||
# j = faces[:,1],
|
||||
# k = faces[:,2],
|
||||
# )])
|
||||
# fig.show()
|
||||
|
||||
# fig = ex.scatter_3d(x=x[:,0], y=x[:,1], z=x[:,2])
|
||||
# fig.update_traces(marker={'color': colors, 'size': 5})
|
||||
# fig.show()
|
||||
|
||||
# all_error.append(err.numpy().sum())
|
||||
# # b = nest_conditions(static_dataset.expand_batch(batch_to_dict(d, signature_dynamic)))
|
||||
# print(len(all_error))
|
||||
|
||||
|
||||
# dataset = DatasetSimple('/home/workspace/data/dataset/stylegan_new/train').batch(2)
|
||||
|
||||
# from losses import collision_loss as collision_penalty
|
||||
# from losses import collision_penalty as collision_penalty_2
|
||||
|
||||
|
||||
# for d in dataset:
|
||||
# d
|
||||
|
||||
# v = uv2mesh(d[0])
|
||||
# b = d[3]
|
||||
# m = d[4]
|
||||
|
||||
# b2 = tf.concat((b[1][tf.newaxis,:,:], b[0][tf.newaxis,:,:]), axis=0)
|
||||
|
||||
# c = collision_penalty(v, b, m)
|
||||
# c2 = collision_penalty(v, b2, m)
|
||||
|
||||
# c_ = collision_penalty_2(v, b, m)
|
||||
# c2_ = collision_penalty_2(v, b2, m)
|
||||
|
||||
# print()
|
||||
# pass
|
||||
|
||||
@@ -0,0 +1,42 @@
|
||||
FROM tensorflow/tensorflow:2.9.0-gpu
|
||||
|
||||
RUN apt-get update
|
||||
RUN apt-get upgrade -y
|
||||
|
||||
# Add the current user to the image
|
||||
ARG USER_NAME=lhunor
|
||||
ARG USER_ID=11017
|
||||
ARG GROUP_NAME=guest
|
||||
ARG GROUP_ID=11000
|
||||
|
||||
# Update the system and install required packages
|
||||
RUN apt-get update && apt-get upgrade -y && apt-get install -y sudo git
|
||||
|
||||
# Create a new group with the same GID as the host
|
||||
RUN groupadd --gid $GROUP_ID $GROUP_NAME
|
||||
|
||||
# Create a new user with the same UID and GID as the host
|
||||
RUN useradd --uid $USER_ID --gid $GROUP_ID --shell /bin/bash --create-home $USER_NAME
|
||||
|
||||
# don't require password with sudo, for convenience
|
||||
# not the safest thing to do, but hopefully okay inside the container
|
||||
RUN echo '%sudo ALL=(ALL) NOPASSWD:ALL' >> /etc/sudoers
|
||||
|
||||
# Add the new user to the sudo group
|
||||
RUN usermod -aG sudo $USER_NAME
|
||||
|
||||
# Set the default user to the new user
|
||||
USER $USER_NAME
|
||||
|
||||
WORKDIR //home-net
|
||||
|
||||
RUN apt-get install git ffmpeg libsm6 libxext6 libopenexr-dev zip unzip htop nano -y
|
||||
RUN pip install numpy matplotlib opencv-python wandb OpenEXR tensorflow-graphics scipy plotly pandas tensorboardX opencv-python cmapy
|
||||
|
||||
# defining the volumes that will be mounted
|
||||
VOLUME [ "/data-net" ]
|
||||
VOLUME [ "/data-local" ]
|
||||
VOLUME [ "/home-net" ]
|
||||
|
||||
# Start the container with a bash shell
|
||||
CMD ["/bin/bash"]
|
||||
@@ -0,0 +1,12 @@
|
||||
numpy
|
||||
matplotlib
|
||||
opencv-python
|
||||
wandb
|
||||
OpenEXR
|
||||
tensorflow-graphics
|
||||
scipy
|
||||
plotly
|
||||
pandas
|
||||
tensorboardX
|
||||
opencv-python
|
||||
cmapy
|
||||
+207
@@ -0,0 +1,207 @@
|
||||
import tensorflow as tf
|
||||
from convnext.models.convnext_tf import ConvNeXt, ConvNeXtReverse, Block, ConvNeXtReverseLabels
|
||||
|
||||
class Bottleneck(tf.keras.Model):
|
||||
|
||||
def __init__(self, dim, latent_dim, kernel_size=7):
|
||||
super(Bottleneck, self).__init__()
|
||||
|
||||
self.bottleneck = tf.keras.Sequential(
|
||||
[Block(dim=dim, drop_path=0, kernel_size=kernel_size,
|
||||
layer_scale_init_value=1e-6, prefix=f'bottleneck{j}') for j in range(5)]
|
||||
)
|
||||
|
||||
self.flatten = tf.keras.layers.Flatten()
|
||||
self.dense = tf.keras.layers.Dense(latent_dim)
|
||||
|
||||
def call(self, inputs, training):
|
||||
x = inputs
|
||||
|
||||
x = self.bottleneck(x, training=training)
|
||||
x = self.flatten(x)
|
||||
x = self.dense(x, training=training)
|
||||
|
||||
return x
|
||||
|
||||
|
||||
class VariationalEncoder(tf.keras.Model):
|
||||
|
||||
def __init__(self, latent_dim, depths, dims, kernel_size=7):
|
||||
super(VariationalEncoder, self).__init__()
|
||||
|
||||
self.encoder = ConvNeXt(depths=depths, dims=dims, include_top=False, kernel_size=kernel_size)
|
||||
|
||||
self.flatten = tf.keras.layers.Flatten()
|
||||
self.mean = tf.keras.layers.Dense(latent_dim)
|
||||
self.log_var = tf.keras.layers.Dense(latent_dim)
|
||||
|
||||
def call(self, inputs, training):
|
||||
x = inputs
|
||||
|
||||
x = self.encoder(x, training=training)
|
||||
x = self.flatten(x)
|
||||
mean = self.mean(x)
|
||||
log_var = self.log_var(x)
|
||||
|
||||
return mean, log_var
|
||||
|
||||
|
||||
class AEEncoder(tf.keras.Model):
|
||||
|
||||
def __init__(self, latent_dim, depths, dims, bottleneck_size=32, kernel_size=7):
|
||||
super(AEEncoder, self).__init__()
|
||||
|
||||
self.encoder = ConvNeXt(depths=depths, dims=dims, include_top=False, kernel_size=kernel_size)
|
||||
self.bottleneck = Bottleneck(bottleneck_size, latent_dim, kernel_size)
|
||||
|
||||
|
||||
def call(self, inputs, training):
|
||||
if len(inputs) == 2:
|
||||
x, _ = inputs
|
||||
else:
|
||||
x = inputs
|
||||
|
||||
x = self.encoder(x, training=training)
|
||||
x = self.bottleneck(x, training=training)
|
||||
|
||||
if len(inputs) == 2:
|
||||
return x, inputs[1]
|
||||
else:
|
||||
return x
|
||||
|
||||
|
||||
class Decoder(tf.keras.Model):
|
||||
|
||||
def __init__(self, input_size, depths, dims, kernel_size=7, is_label=False, upscale=False):
|
||||
super(Decoder, self).__init__()
|
||||
|
||||
self.dense = tf.keras.layers.Dense(tf.reduce_prod(input_size))
|
||||
self.reshape = tf.keras.layers.Reshape(input_size)
|
||||
self.upscale = upscale
|
||||
|
||||
if not is_label:
|
||||
self.decoder = ConvNeXtReverse(depths=list(reversed(depths)), dims=list(reversed(dims)), kernel_size=kernel_size)
|
||||
else:
|
||||
self.decoder = ConvNeXtReverseLabels(depths=list(reversed(depths)), dims=list(reversed(dims)), nr_of_classes=7)
|
||||
|
||||
|
||||
def call(self, inputs, training):
|
||||
if len(inputs) == 2:
|
||||
x, mask = inputs
|
||||
else:
|
||||
x = inputs
|
||||
|
||||
x = self.dense(x)
|
||||
x = self.reshape(x)
|
||||
x = self.decoder(x, training=training)
|
||||
|
||||
if self.upscale:
|
||||
x = tf.image.resize(x, size=(512, 256), method=tf.image.ResizeMethod.LANCZOS5)
|
||||
|
||||
|
||||
if len(inputs) == 2:
|
||||
if len(mask.shape) == 2:
|
||||
x = x * mask[tf.newaxis,:,:,tf.newaxis]
|
||||
else:
|
||||
x = x * mask[:,:,:,tf.newaxis]
|
||||
|
||||
return x
|
||||
|
||||
|
||||
|
||||
class VariationalEncoderNew(tf.keras.Model):
|
||||
|
||||
def __init__(self, depths, dims, kernel_size=7):
|
||||
super(VariationalEncoderNew, self).__init__()
|
||||
|
||||
self.encoder = ConvNeXt(depths=depths, dims=dims, include_top=False, kernel_size=kernel_size)
|
||||
|
||||
self.mean_branch = tf.keras.Sequential(
|
||||
[Block(dim=64, drop_path=0, kernel_size=kernel_size,
|
||||
layer_scale_init_value=1e-6, prefix=f'bottleneck{j}') for j in range(5)]
|
||||
)
|
||||
|
||||
self.var_branch = tf.keras.Sequential(
|
||||
[Block(dim=64, drop_path=0, kernel_size=kernel_size,
|
||||
layer_scale_init_value=1e-6, prefix=f'bottleneck{j}') for j in range(5)]
|
||||
)
|
||||
|
||||
self.flatten_mean = tf.keras.layers.Flatten()
|
||||
self.flatten_var = tf.keras.layers.Flatten()
|
||||
|
||||
|
||||
|
||||
def call(self, inputs, training):
|
||||
x = inputs
|
||||
|
||||
x = self.encoder(x, training=training)
|
||||
mean = self.mean_branch(x)
|
||||
log_var = self.var_branch(x)
|
||||
|
||||
mean = self.flatten_mean(mean)
|
||||
log_var = self.flatten_var(log_var)
|
||||
|
||||
return mean, log_var
|
||||
|
||||
|
||||
|
||||
class AEEncoderNew(tf.keras.Model):
|
||||
|
||||
def __init__(self, depths, dims, kernel_size=7, activation='tanh'):
|
||||
super(AEEncoderNew, self).__init__()
|
||||
|
||||
self.encoder = ConvNeXt(depths=depths, dims=dims, include_top=False, kernel_size=kernel_size)
|
||||
self.act = tf.keras.layers.Activation(activation)
|
||||
|
||||
|
||||
def call(self, inputs, training):
|
||||
if len(inputs) == 2:
|
||||
x, _ = inputs
|
||||
else:
|
||||
x = inputs
|
||||
|
||||
x = self.encoder(x, training=training)
|
||||
x = self.act(x)
|
||||
|
||||
if len(inputs) == 2:
|
||||
return x, inputs[1]
|
||||
else:
|
||||
return x
|
||||
|
||||
|
||||
class DecoderNew(tf.keras.Model):
|
||||
|
||||
def __init__(self, depths, dims, kernel_size=7, is_label=False):
|
||||
super(DecoderNew, self).__init__()
|
||||
|
||||
self.is_label = is_label
|
||||
|
||||
if not is_label:
|
||||
self.decoder = ConvNeXtReverse(depths=depths, dims=dims, kernel_size=kernel_size)
|
||||
else:
|
||||
self.decoder = ConvNeXtReverseLabels(depths=depths, dims=dims, nr_of_classes=7)
|
||||
|
||||
|
||||
def call(self, inputs, training):
|
||||
if len(inputs) == 2:
|
||||
x, mask = inputs
|
||||
else:
|
||||
x = inputs
|
||||
|
||||
x = self.decoder(x, training=training)
|
||||
|
||||
if self.is_label:
|
||||
x = x * mask[:,:,:,tf.newaxis]
|
||||
|
||||
background = tf.ones((mask.shape[0],512,256,1)) * (1 - mask[:,:,:,tf.newaxis])
|
||||
x = tf.concat((background, x), axis=-1)
|
||||
|
||||
x = tf.math.softmax(x)
|
||||
|
||||
if len(inputs) == 2:
|
||||
if len(mask.shape) == 2:
|
||||
x = x * mask[tf.newaxis,:,:,tf.newaxis]
|
||||
else:
|
||||
x = x * mask[:,:,:,tf.newaxis]
|
||||
|
||||
return x
|
||||
@@ -0,0 +1,38 @@
|
||||
import tensorflow as tf
|
||||
from models.base import AEEncoder, Decoder, AEEncoderNew, DecoderNew
|
||||
|
||||
|
||||
class AutoEncoder(tf.keras.Model):
|
||||
|
||||
def __init__(self, config):
|
||||
super(AutoEncoder, self).__init__()
|
||||
|
||||
self.encoder : AEEncoder = AEEncoder(config['latent_dim'], config['encoder_depths'], config['encoder_dims'], config['bottleneck_size'], config['kernel_size'])
|
||||
self.decoder : Decoder = Decoder(config['decoder_input'], config['decoder_depths'], config['decoder_dims'], kernel_size=config['kernel_size'], is_label=config['is_label'], upscale=config['upscale_output'])
|
||||
|
||||
|
||||
def call(self, inputs, training):
|
||||
x = inputs
|
||||
|
||||
x = self.encoder(x, training=training)
|
||||
x = self.decoder(x, training=training)
|
||||
|
||||
return x
|
||||
|
||||
|
||||
class AutoEncoderNew(tf.keras.Model):
|
||||
|
||||
def __init__(self, config):
|
||||
super(AutoEncoderNew, self).__init__()
|
||||
|
||||
self.encoder : AEEncoderNew = AEEncoderNew(config['encoder_depths'], config['encoder_dims'], config['kernel_size'], config['activation'])
|
||||
self.decoder : DecoderNew = DecoderNew(config['decoder_depths'], config['decoder_dims'], kernel_size=config['kernel_size'], is_label=config['is_label'])
|
||||
|
||||
|
||||
def call(self, inputs, training):
|
||||
x = inputs
|
||||
|
||||
x = self.encoder(x, training=training)
|
||||
x = self.decoder(x, training=training)
|
||||
|
||||
return x
|
||||
@@ -0,0 +1,306 @@
|
||||
from unittest import skip
|
||||
import numpy as np
|
||||
import tensorflow as tf
|
||||
from tensorflow.keras.layers import (Conv2D, Conv2DTranspose, Dense, DepthwiseConv2D,
|
||||
GlobalAveragePooling2D, Layer,
|
||||
LayerNormalization, Flatten, Reshape, UpSampling2D)
|
||||
|
||||
def drop_path(inputs, drop_prob, is_training):
|
||||
# borrowed from https://github.com/rishigami/Swin-Transformer-TF/blob/main/swintransformer/model.py
|
||||
if (not is_training) or (drop_prob == 0.):
|
||||
return inputs
|
||||
|
||||
# Compute keep_prob
|
||||
keep_prob = 1.0 - drop_prob
|
||||
|
||||
# Compute drop_connect tensor
|
||||
random_tensor = keep_prob
|
||||
shape = (tf.shape(inputs)[0],) + (1,) * \
|
||||
(len(tf.shape(inputs)) - 1)
|
||||
random_tensor += tf.random.uniform(shape, dtype=inputs.dtype)
|
||||
binary_tensor = tf.floor(random_tensor)
|
||||
output = tf.math.divide(inputs, keep_prob) * binary_tensor
|
||||
return output
|
||||
|
||||
|
||||
class DropPath(tf.keras.layers.Layer):
|
||||
# borrowed from https://github.com/rishigami/Swin-Transformer-TF/blob/main/swintransformer/model.py
|
||||
def __init__(self, drop_prob=None):
|
||||
super().__init__()
|
||||
self.drop_prob = drop_prob
|
||||
|
||||
def call(self, x, training):
|
||||
return drop_path(x, self.drop_prob, training)
|
||||
|
||||
|
||||
class Block(Layer):
|
||||
r""" ConvNeXt Block. There are two equivalent implementations:
|
||||
(1) DwConv -> LayerNorm (channels_first) -> 1x1 Conv -> GELU -> 1x1 Conv; all in (N, C, H, W)
|
||||
(2) DwConv -> Permute to (N, H, W, C); LayerNorm (channels_last) -> Linear -> GELU -> Linear; Permute back
|
||||
We use (2) as we find it slightly faster in PyTorch
|
||||
|
||||
Args:
|
||||
dim (int): Number of input channels.
|
||||
drop_path (float): Stochastic depth rate. Default: 0.0
|
||||
layer_scale_init_value (float): Init value for Layer Scale. Default: 1e-6.
|
||||
"""
|
||||
|
||||
def __init__(self, dim, kernel_size=7, drop_path=0., layer_scale_init_value=1e-6, prefix=''):
|
||||
super().__init__()
|
||||
self.dwconv = DepthwiseConv2D(
|
||||
kernel_size=kernel_size, padding='same') # depthwise conv
|
||||
self.norm = LayerNormalization(epsilon=1e-6)
|
||||
# pointwise/1x1 convs, implemented with linear layers
|
||||
self.pwconv1 = Dense(4 * dim)
|
||||
self.act = tf.keras.activations.gelu
|
||||
self.pwconv2 = Dense(dim)
|
||||
self.drop_path = DropPath(drop_path)
|
||||
self.dim = dim
|
||||
self.layer_scale_init_value = layer_scale_init_value
|
||||
self.prefix = prefix
|
||||
|
||||
def build(self, input_shape):
|
||||
self.gamma = tf.Variable(
|
||||
initial_value=self.layer_scale_init_value * tf.ones((self.dim)),
|
||||
trainable=True,
|
||||
name=f'{self.prefix}/gamma')
|
||||
self.built = True
|
||||
|
||||
def call(self, x, training):
|
||||
input = x
|
||||
x = self.dwconv(x, training=training)
|
||||
# x = x.permute(0, 2, 3, 1) # (N, C, H, W) -> (N, H, W, C)
|
||||
x = self.norm(x, training=training)
|
||||
x = self.pwconv1(x, training=training)
|
||||
x = self.act(x)
|
||||
x = self.pwconv2(x, training=training)
|
||||
if self.gamma is not None:
|
||||
x = self.gamma * x
|
||||
# x = x.permute(0, 3, 1, 2) # (N, H, W, C) -> (N, C, H, W)
|
||||
|
||||
x = input + self.drop_path(x, training=training)
|
||||
return x
|
||||
|
||||
|
||||
class ConvNeXt(tf.keras.Model):
|
||||
r""" ConvNeXt
|
||||
A Tensorflow keras impl of : `A ConvNet for the 2020s` -
|
||||
https://arxiv.org/pdf/2201.03545.pdf
|
||||
|
||||
Args:
|
||||
num_classes (int): Number of classes for classification head. Default: 1000
|
||||
depths (tuple(int)): Number of blocks at each stage. Default: [3, 3, 9, 3]
|
||||
dims (int): Feature dimension at each stage. Default: [96, 192, 384, 768]
|
||||
include_top (bool): whether to add head or just use it as feature extractor. Default: True
|
||||
drop_path_rate (float): Stochastic depth rate. Default: 0.
|
||||
layer_scale_init_value (float): Init value for Layer Scale. Default: 1e-6.
|
||||
head_init_scale (float): Init scaling value for classifier weights and biases. Default: 1.
|
||||
"""
|
||||
|
||||
def __init__(self, num_classes=1000, kernel_size=7,
|
||||
depths=[3, 3, 9, 3], dims=[96, 192, 384, 768], include_top=True,
|
||||
drop_path_rate=0., layer_scale_init_value=1e-6, head_init_scale=1.,
|
||||
):
|
||||
super().__init__()
|
||||
self.include_top = include_top
|
||||
self.downsample_layers = [] # stem and 3 intermediate downsampling conv layers
|
||||
# stem = tf.keras.Sequential([
|
||||
# Conv2D(dims[0], kernel_size=4, strides=4, padding='same'),
|
||||
# LayerNormalization(epsilon=1e-6)]
|
||||
# )
|
||||
stem = [Conv2D(dims[0], kernel_size=4, strides=4, padding='same'), LayerNormalization(epsilon=1e-6)]
|
||||
self.downsample_layers.append(stem)
|
||||
for i in range(len(depths) - 1):
|
||||
# downsample_layer = tf.keras.Sequential([
|
||||
# LayerNormalization(epsilon=1e-6),
|
||||
# Conv2D(dims[i+1], kernel_size=2, strides=2, padding='same')]
|
||||
# )
|
||||
downsample_layer = [LayerNormalization(epsilon=1e-6), Conv2D(dims[i+1], kernel_size=2, strides=2, padding='same')]
|
||||
self.downsample_layers.append(downsample_layer)
|
||||
|
||||
self.stages = [] # 4 feature resolution stages, each consisting of multiple residual blocks
|
||||
dp_rates = [x for x in np.linspace(0, drop_path_rate, sum(depths))]
|
||||
cur = 0
|
||||
for i in range(len(depths)):
|
||||
stage = tf.keras.Sequential(
|
||||
[Block(dim=dims[i], drop_path=dp_rates[cur + j], kernel_size=kernel_size,
|
||||
layer_scale_init_value=layer_scale_init_value, prefix=f'block{i}') for j in range(depths[i])]
|
||||
)
|
||||
self.stages.append(stage)
|
||||
cur += depths[i]
|
||||
|
||||
if self.include_top:
|
||||
self.avg = GlobalAveragePooling2D()
|
||||
self.norm = LayerNormalization(epsilon=1e-6) # final norm layer
|
||||
self.head = Dense(num_classes)
|
||||
else:
|
||||
self.avg = None
|
||||
self.norm = None
|
||||
self.head = None
|
||||
|
||||
def forward_features(self, x, training):
|
||||
for i in range(len(self.downsample_layers)):
|
||||
# x = self.downsample_layers[i](x)
|
||||
x = self.downsample_layers[i][0](x, training=training)
|
||||
x = self.downsample_layers[i][1](x, training=training)
|
||||
x = self.stages[i](x, training=training)
|
||||
return x
|
||||
|
||||
def call(self, x, training):
|
||||
x = self.forward_features(x, training=training)
|
||||
return x
|
||||
|
||||
|
||||
class ConvNeXtReverse(tf.keras.Model):
|
||||
r""" ConvNeXt
|
||||
A Tensorflow keras impl of : `A ConvNet for the 2020s` -
|
||||
https://arxiv.org/pdf/2201.03545.pdf
|
||||
|
||||
Args:
|
||||
num_classes (int): Number of classes for classification head. Default: 1000
|
||||
depths (tuple(int)): Number of blocks at each stage. Default: [3, 3, 9, 3]
|
||||
dims (int): Feature dimension at each stage. Default: [96, 192, 384, 768]
|
||||
include_top (bool): whether to add head or just use it as feature extractor. Default: True
|
||||
drop_path_rate (float): Stochastic depth rate. Default: 0.
|
||||
layer_scale_init_value (float): Init value for Layer Scale. Default: 1e-6.
|
||||
head_init_scale (float): Init scaling value for classifier weights and biases. Default: 1.
|
||||
"""
|
||||
|
||||
def __init__(self,
|
||||
depths=[3, 3, 9, 3], dims=[96, 192, 384, 768], kernel_size=7,
|
||||
drop_path_rate=0., layer_scale_init_value=1e-6, head_init_scale=1.,
|
||||
):
|
||||
super().__init__()
|
||||
self.upsample_layers = [] # stem and 3 intermediate downsampling conv layers
|
||||
stem = tf.keras.Sequential([
|
||||
Conv2DTranspose(dims[0], kernel_size=4, strides=4, padding='same'),
|
||||
# UpSampling2D((4, 4), interpolation='nearest'),
|
||||
# Conv2D(dims[0], kernel_size=4, strides=1, padding='same'),
|
||||
# Dense(dims[0]),
|
||||
LayerNormalization(epsilon=1e-6)]
|
||||
)
|
||||
self.upsample_layers.append(stem)
|
||||
for i in range(len(depths) - 1):
|
||||
downsample_layer = tf.keras.Sequential([
|
||||
LayerNormalization(epsilon=1e-6),
|
||||
Conv2DTranspose(dims[i+1], kernel_size=2, strides=2, padding='same')
|
||||
# UpSampling2D((2,2), interpolation='nearest'),
|
||||
# Conv2D(dims[i+1], kernel_size=2, strides=1, padding='same'),
|
||||
# Dense(dims[i+1]),
|
||||
])
|
||||
self.upsample_layers.append(downsample_layer)
|
||||
|
||||
self.stages = [] # 4 feature resolution stages, each consisting of multiple residual blocks
|
||||
dp_rates = [x for x in np.linspace(0, drop_path_rate, sum(depths))]
|
||||
cur = 0
|
||||
for i in range(len(depths)):
|
||||
stage = tf.keras.Sequential(
|
||||
[Block(dim=dims[i], drop_path=dp_rates[cur + j], kernel_size=kernel_size,
|
||||
layer_scale_init_value=layer_scale_init_value, prefix=f'block{i}') for j in range(depths[i])]
|
||||
)
|
||||
self.stages.append(stage)
|
||||
cur += depths[i]
|
||||
|
||||
self.last = Dense(3, activation='tanh')
|
||||
|
||||
def call(self, x, training):
|
||||
for i in reversed(range(len(self.upsample_layers))):
|
||||
x = self.upsample_layers[i](x, training=training)
|
||||
x = self.stages[i](x, training=training)
|
||||
x = self.last(x, training=training)
|
||||
return x
|
||||
|
||||
|
||||
class ConvNeXtReverseLabels(tf.keras.Model):
|
||||
r""" ConvNeXt
|
||||
A Tensorflow keras impl of : `A ConvNet for the 2020s` -
|
||||
https://arxiv.org/pdf/2201.03545.pdf
|
||||
|
||||
Args:
|
||||
num_classes (int): Number of classes for classification head. Default: 1000
|
||||
depths (tuple(int)): Number of blocks at each stage. Default: [3, 3, 9, 3]
|
||||
dims (int): Feature dimension at each stage. Default: [96, 192, 384, 768]
|
||||
include_top (bool): whether to add head or just use it as feature extractor. Default: True
|
||||
drop_path_rate (float): Stochastic depth rate. Default: 0.
|
||||
layer_scale_init_value (float): Init value for Layer Scale. Default: 1e-6.
|
||||
head_init_scale (float): Init scaling value for classifier weights and biases. Default: 1.
|
||||
"""
|
||||
|
||||
def __init__(self,
|
||||
depths=[3, 3, 9, 3], dims=[96, 192, 384, 768],
|
||||
drop_path_rate=0., layer_scale_init_value=1e-6, head_init_scale=1.,
|
||||
nr_of_classes=7):
|
||||
super().__init__()
|
||||
self.downsample_layers = [] # stem and 3 intermediate downsampling conv layers
|
||||
stem = tf.keras.Sequential([
|
||||
Conv2DTranspose(dims[0], kernel_size=4, strides=4, padding='same'),
|
||||
LayerNormalization(epsilon=1e-6)]
|
||||
)
|
||||
self.downsample_layers.append(stem)
|
||||
for i in range(len(depths) - 1):
|
||||
downsample_layer = tf.keras.Sequential([
|
||||
LayerNormalization(epsilon=1e-6),
|
||||
Conv2DTranspose(dims[i+1], kernel_size=2, strides=2, padding='same')]
|
||||
)
|
||||
self.downsample_layers.append(downsample_layer)
|
||||
|
||||
self.stages = [] # 4 feature resolution stages, each consisting of multiple residual blocks
|
||||
dp_rates = [x for x in np.linspace(0, drop_path_rate, sum(depths))]
|
||||
cur = 0
|
||||
for i in range(len(depths)):
|
||||
stage = tf.keras.Sequential(
|
||||
[Block(dim=dims[i], drop_path=dp_rates[cur + j],
|
||||
layer_scale_init_value=layer_scale_init_value, prefix=f'block{i}') for j in range(depths[i])]
|
||||
)
|
||||
self.stages.append(stage)
|
||||
cur += depths[i]
|
||||
|
||||
# self.last = Dense(nr_of_classes, activation='softmax')
|
||||
self.last = Dense(nr_of_classes - 1)
|
||||
|
||||
def call(self, x, training):
|
||||
for i in reversed(range(len(self.downsample_layers))):
|
||||
x = self.downsample_layers[i](x, training=training)
|
||||
x = self.stages[i](x, training=training)
|
||||
x = self.last(x, training=training)
|
||||
return x
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
model_urls = {
|
||||
"convnext_tiny_224": "https://github.com/bamps53/convnext-tf/releases/download/v0.1/convnext_tiny_1k_224_ema.h5",
|
||||
"convnext_small_224": "https://github.com/bamps53/convnext-tf/releases/download/v0.1/convnext_small_1k_224_ema.h5",
|
||||
"convnext_base_224": "https://github.com/bamps53/convnext-tf/releases/download/v0.1/convnext_base_22k_1k_224.h5",
|
||||
"convnext_base_384": "https://github.com/bamps53/convnext-tf/releases/download/v0.1/convnext_base_22k_1k_384.h5",
|
||||
"convnext_large_224": "https://github.com/bamps53/convnext-tf/releases/download/v0.1/convnext_large_22k_1k_224.h5",
|
||||
"convnext_large_384": "https://github.com/bamps53/convnext-tf/releases/download/v0.1/convnext_large_22k_1k_384.h5",
|
||||
"convnext_xlarge_224": "https://github.com/bamps53/convnext-tf/releases/download/v0.1/convnext_xlarge_22k_1k_224_ema.h5",
|
||||
"convnext_xlarge_384": "https://github.com/bamps53/convnext-tf/releases/download/v0.1/convnext_xlarge_22k_1k_384_ema.h5",
|
||||
}
|
||||
|
||||
|
||||
model_configs = dict(
|
||||
convnext_tiny=dict(
|
||||
depths=[3, 3, 9, 3],
|
||||
dims=[8, 16, 32, 512]
|
||||
),
|
||||
convnext_small=dict(
|
||||
depths=[3, 3, 27, 3],
|
||||
dims=[96, 192, 384, 768]
|
||||
),
|
||||
convnext_base=dict(
|
||||
depths=[3, 3, 27, 3],
|
||||
dims=[128, 256, 512, 1024]
|
||||
),
|
||||
convnext_large=dict(
|
||||
depths=[3, 3, 27, 3],
|
||||
dims=[192, 384, 768, 1536]
|
||||
),
|
||||
convnext_xlarge=dict(
|
||||
depths=[3, 3, 27, 3],
|
||||
dims=[256, 512, 1024, 2048]
|
||||
),
|
||||
)
|
||||
|
||||
@@ -0,0 +1,140 @@
|
||||
import os
|
||||
import time
|
||||
import random
|
||||
import tensorflow as tf
|
||||
import numpy as np
|
||||
import scipy as sp
|
||||
# import plotly.graph_objects as go
|
||||
# import plotly.express as ex
|
||||
|
||||
from custom_utils.custom_utils import load_model
|
||||
from models.base import AEEncoder, VariationalEncoder, Decoder
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
class CVAE(tf.keras.Model):
|
||||
|
||||
def __init__(self, config):
|
||||
super(CVAE, self).__init__()
|
||||
self.latent_dim = config['latent_dim']
|
||||
self.latent_dim = config['cvae']['latent_dim']
|
||||
self.size = config['resolution']
|
||||
|
||||
self.static_enc : AEEncoder = AEEncoder(config['cond_ae']['latent_dim'], config['cond_ae']['encoder_depths'], config['cond_ae']['encoder_dims'], config['cond_ae']['bottleneck_size'], config['cond_ae']['kernel_size'])
|
||||
if config['static_model'] != '':
|
||||
load_model(self.static_enc, config['static_model'] + '_enc')
|
||||
|
||||
self.body_enc : AEEncoder = AEEncoder(config['cond_ae_body']['latent_dim'], config['cond_ae_body']['encoder_depths'], config['cond_ae_body']['encoder_dims'], config['cond_ae_body']['bottleneck_size'], config['cond_ae_body']['kernel_size'])
|
||||
if config['body_model'] != '':
|
||||
load_model(self.body_enc, config['body_model'] + '_enc')
|
||||
|
||||
if config['freeze_encoders']:
|
||||
self.static_enc.trainable = False
|
||||
self.body_enc.trainable = False
|
||||
|
||||
self.cond_encoder = tf.keras.Sequential([tf.keras.layers.Dense(512, activation='relu'),
|
||||
tf.keras.layers.Dense(256, activation='relu'),
|
||||
tf.keras.layers.Dense(128, activation='relu'),
|
||||
tf.keras.layers.Dense(self.latent_dim, activation='tanh'),])
|
||||
|
||||
self.encoder = VariationalEncoder(config['cvae']['latent_dim'], config['cvae']['encoder_depths'], config['cvae']['encoder_dims'], kernel_size=config['cvae']['kernel_size'])
|
||||
self.decoder = Decoder(config['cvae']['decoder_input'], config['cvae']['decoder_depths'], config['cvae']['decoder_dims'], config['cvae']['kernel_size'], is_label=config['cvae']['is_label'])
|
||||
|
||||
self.add_noise = config['cvae']['add_noise']
|
||||
self.noise_amount = config['cvae']['noise_amount']
|
||||
|
||||
# @tf.function
|
||||
def sample_single(self, uv_template, uv_body_posed, uv_mask, eps=None, nr_of_samples=4):
|
||||
if eps is None:
|
||||
eps = tf.random.normal(shape=(nr_of_samples, self.latent_dim))
|
||||
|
||||
uv_template = tf.tile(tf.transpose(uv_template, [1,2,0])[tf.newaxis,:,:,:], [nr_of_samples,1,1,1])
|
||||
uv_body_posed = tf.tile(tf.transpose(uv_body_posed, [1,2,0])[tf.newaxis,:,:,:], [nr_of_samples,1,1,1])
|
||||
uv_mask = tf.tile(uv_mask[tf.newaxis,:,:], [nr_of_samples,1,1])
|
||||
|
||||
# uv_template, uv_body_posed = tf.transpose(uv_template, [0,2,3,1]), tf.transpose(uv_body_posed, [0,2,3,1])
|
||||
z_static = self.static_enc(uv_template)
|
||||
z_body = self.body_enc(uv_body_posed)
|
||||
|
||||
z_cond = tf.concat([z_static, z_body], axis=-1)
|
||||
z_cond = self.cond_encoder(z_cond)
|
||||
z = tf.concat([eps, z_cond], axis=-1)
|
||||
|
||||
x = self.decoder(z)
|
||||
x = x * uv_mask[:,:,:,tf.newaxis]
|
||||
|
||||
return x
|
||||
|
||||
|
||||
def encode_conds(self, uv_template, uv_body_posed, training):
|
||||
z_static = self.static_enc(uv_template, training=training)
|
||||
z_body = self.body_enc(uv_body_posed, training=training)
|
||||
|
||||
z_cond = tf.concat([z_static, z_body], axis=-1)
|
||||
z_cond = self.cond_encoder(z_cond, training=training)
|
||||
|
||||
return z_cond
|
||||
|
||||
def sample(self, uv_template, uv_body_posed, uv_mask, noise=None):
|
||||
if noise is None:
|
||||
eps = tf.random.normal(shape=(uv_template.shape[0], self.latent_dim))
|
||||
else:
|
||||
eps = noise
|
||||
|
||||
z_cond = self.encode_conds(uv_template, uv_body_posed, training=False)
|
||||
z = tf.concat([eps, z_cond], axis=-1)
|
||||
|
||||
x = self.decoder(z, training=False)
|
||||
|
||||
x = x * uv_mask[:,:,:,tf.newaxis]
|
||||
|
||||
background = tf.ones((uv_template.shape[0],512,256,1)) * (1 - uv_mask[:,:,:,tf.newaxis])
|
||||
x = tf.concat((background, x), axis=-1)
|
||||
|
||||
x = tf.math.softmax(x)
|
||||
|
||||
return x
|
||||
|
||||
|
||||
# encoders have to be the same for curv_vae and general CVAE
|
||||
def sample_without_encoding(self, z_static, z_body, uv_mask, eps=None):
|
||||
eps = tf.random.normal(shape=(z_static.shape[0], self.latent_dim))
|
||||
|
||||
z_cond = tf.concat([z_static, z_body], axis=-1)
|
||||
z = tf.concat([eps, z_cond], axis=-1)
|
||||
|
||||
x = self.decoder(z)
|
||||
x = x * uv_mask[:,:,:,tf.newaxis]
|
||||
|
||||
return x
|
||||
|
||||
|
||||
def reparameterize(self, mean, logvar):
|
||||
eps = tf.random.normal(shape=mean.shape)
|
||||
return eps * tf.exp(logvar * .5) + mean
|
||||
|
||||
|
||||
def call(self, inputs, training):
|
||||
uv_template, uv_body_posed, uv_curv, uv_mask = inputs
|
||||
|
||||
|
||||
x = tf.concat([uv_template, uv_body_posed, uv_curv], axis=-1)
|
||||
mean, logvar = self.encoder(x)
|
||||
z = self.reparameterize(mean, logvar)
|
||||
|
||||
z_cond = self.encode_conds(uv_template, uv_body_posed, training=training)
|
||||
|
||||
z = tf.concat([z, z_cond], axis=-1)
|
||||
x = self.decoder(z, training=training)
|
||||
|
||||
x = x * uv_mask[:,:,:,tf.newaxis]
|
||||
|
||||
background = tf.ones((4,512,256,1)) * (1 - uv_mask[:,:,:,tf.newaxis])
|
||||
x = tf.concat((background, x), axis=-1)
|
||||
|
||||
x = tf.math.softmax(x)
|
||||
|
||||
return x, mean, logvar
|
||||
@@ -0,0 +1,285 @@
|
||||
import os
|
||||
import time
|
||||
import random
|
||||
import tensorflow as tf
|
||||
import numpy as np
|
||||
import scipy as sp
|
||||
# import plotly.graph_objects as go
|
||||
# import plotly.express as ex
|
||||
|
||||
from custom_utils.custom_utils import load_model
|
||||
from models.base import AEEncoder, VariationalEncoder, Decoder
|
||||
from models.base import AEEncoderNew, VariationalEncoderNew, DecoderNew
|
||||
|
||||
|
||||
class CVAE(tf.keras.Model):
|
||||
|
||||
def __init__(self, config):
|
||||
super(CVAE, self).__init__()
|
||||
self.latent_dim = config['cvae']['latent_dim']
|
||||
self.size = config['resolution']
|
||||
|
||||
if not('only_curvatures' in config.keys() and config['only_curvatures'] == True):
|
||||
self.static_enc : AEEncoder = AEEncoder(config['cond_ae']['latent_dim'], config['cond_ae']['encoder_depths'], config['cond_ae']['encoder_dims'], config['cond_ae']['bottleneck_size'], config['cond_ae']['kernel_size'])
|
||||
if config['static_model'] != '':
|
||||
load_model(self.static_enc, config['static_model'] + '_enc')
|
||||
|
||||
self.body_enc : AEEncoder = AEEncoder(config['cond_ae_body']['latent_dim'], config['cond_ae_body']['encoder_depths'], config['cond_ae_body']['encoder_dims'], config['cond_ae_body']['bottleneck_size'], config['cond_ae_body']['kernel_size'])
|
||||
if config['body_model'] != '':
|
||||
load_model(self.body_enc, config['body_model'] + '_enc')
|
||||
|
||||
if config['use_curvature']:
|
||||
self.curv_enc : AEEncoder = AEEncoder(config['cond_ae']['latent_dim'], config['cond_ae']['encoder_depths'], config['cond_ae']['encoder_dims'], config['cond_ae']['bottleneck_size'], config['cond_ae']['kernel_size'])
|
||||
if config['curvature_model'] != '':
|
||||
load_model(self.curv_enc, config['curvature_model'] + '_enc')
|
||||
|
||||
if config['freeze_encoders']:
|
||||
if not('only_curvatures' in config.keys() and config['only_curvatures'] == True):
|
||||
self.static_enc.trainable = False
|
||||
self.body_enc.trainable = False
|
||||
if config['use_curvature']:
|
||||
self.curv_enc.trainable = False
|
||||
|
||||
self.cond_encoder = tf.keras.Sequential([tf.keras.layers.Dense(512, activation='relu'),
|
||||
tf.keras.layers.Dense(256, activation='relu'),
|
||||
tf.keras.layers.Dense(128, activation='relu'),
|
||||
# tf.keras.layers.Lambda(lambda x: x / 1000),
|
||||
tf.keras.layers.Dense(config['latent_dim'], activation='tanh')])
|
||||
|
||||
self.encoder = VariationalEncoder(config['cvae']['latent_dim'], config['cvae']['encoder_depths'], config['cvae']['encoder_dims'], kernel_size=config['cvae']['kernel_size'])
|
||||
self.decoder = Decoder(config['cvae']['decoder_input'], config['cvae']['decoder_depths'], config['cvae']['decoder_dims'], config['cvae']['kernel_size'])
|
||||
|
||||
if config['avg_garment_path'] != '':
|
||||
self.avg_garment = np.load(config['avg_garment_path'])
|
||||
else:
|
||||
self.avg_garment = None
|
||||
|
||||
self.add_noise = config['cvae']['add_noise']
|
||||
self.noise_amount = config['cvae']['noise_amount']
|
||||
self.use_curvature = config['use_curvature']
|
||||
self.only_curvatures = config['only_curvatures'] if 'only_curvatures' in config.keys() and config['only_curvatures'] == True else None
|
||||
self.use_offsets_over_body = not config['use_skinning']
|
||||
|
||||
if self.use_offsets_over_body:
|
||||
assert self.avg_garment is None, 'When we are not using skinning, we can\'t use avg_garment!'
|
||||
if self.avg_garment is not None:
|
||||
assert not self.use_offsets_over_body, 'When avg_garment is given, we must use skinning!'
|
||||
|
||||
# @tf.function
|
||||
def sample(self, uv_template, uv_body_posed, uv_curv, uv_mask, noise=None, add_noise=False):
|
||||
if noise is None:
|
||||
eps = tf.random.normal(shape=(uv_template.shape[0], self.latent_dim))
|
||||
else:
|
||||
eps = noise
|
||||
|
||||
z_cond = self.encode_conds(uv_template, uv_body_posed, uv_curv, training=False)
|
||||
if add_noise:
|
||||
z_cond = z_cond + (tf.random.uniform(z_cond.shape, maxval=self.noise_amount * 2) - self.noise_amount)
|
||||
|
||||
z = tf.concat([eps, z_cond], axis=-1)
|
||||
|
||||
x = self.decoder(z, training=False)
|
||||
|
||||
if self.avg_garment is not None:
|
||||
x = x + self.avg_garment[tf.newaxis,:,:,:]
|
||||
|
||||
if self.use_offsets_over_body:
|
||||
x = x + uv_body_posed
|
||||
|
||||
x = x * uv_mask[:,:,:,tf.newaxis]
|
||||
|
||||
return x
|
||||
|
||||
|
||||
def reparameterize(self, mean, logvar, training, fix_noise=False):
|
||||
if training == True:
|
||||
if not fix_noise:
|
||||
eps = tf.random.normal(shape=mean.shape)
|
||||
else:
|
||||
eps = tf.random.normal(shape=(1, mean.shape[1]))
|
||||
eps = tf.tile(eps, [mean.shape[0],1])
|
||||
out = eps * tf.exp(logvar * .5) + mean
|
||||
else:
|
||||
out = tf.exp(logvar * .5) + mean
|
||||
return out
|
||||
|
||||
|
||||
def encode_conds(self, uv_template, uv_body_posed, uv_curv, training):
|
||||
if self.only_curvatures:
|
||||
z_cond = self.curv_enc(uv_curv, training=training)
|
||||
else:
|
||||
z_body = self.body_enc(uv_body_posed, training=training)
|
||||
z_static = self.static_enc(uv_template, training=training)
|
||||
if self.use_curvature:
|
||||
z_curv = self.curv_enc(uv_curv, training=training)
|
||||
z_cond = tf.concat([z_static, z_body, z_curv], axis=-1)
|
||||
else:
|
||||
z_cond = tf.concat([z_static, z_body], axis=-1)
|
||||
|
||||
z_cond = self.cond_encoder(z_cond, training=training)
|
||||
|
||||
return z_cond
|
||||
|
||||
|
||||
def call(self, inputs, training):
|
||||
uv_unposed, uv_template, uv_body_posed, uv_curv, uv_mask = inputs
|
||||
if self.only_curvatures:
|
||||
x = tf.concat([uv_unposed, uv_curv], axis=-1)
|
||||
else:
|
||||
x = tf.concat([uv_unposed, uv_template, uv_body_posed, uv_curv], axis=-1)
|
||||
|
||||
mean, logvar = self.encoder(x, training=training)
|
||||
z = self.reparameterize(mean, logvar, training)
|
||||
z_cond = self.encode_conds(uv_template, uv_body_posed, uv_curv, training=training)
|
||||
|
||||
if self.add_noise and training == True:
|
||||
z_cond = z_cond + (tf.random.uniform(z_cond.shape, maxval=self.noise_amount * 2) - self.noise_amount)
|
||||
|
||||
|
||||
z = tf.concat([z, z_cond], axis=-1)
|
||||
x = self.decoder(z, training=training)
|
||||
|
||||
if self.avg_garment is not None:
|
||||
x = x + self.avg_garment[tf.newaxis,:,:,:]
|
||||
|
||||
if self.use_offsets_over_body:
|
||||
x = x + uv_body_posed
|
||||
|
||||
x = x * uv_mask[:,:,:,tf.newaxis]
|
||||
|
||||
return x, mean, logvar
|
||||
|
||||
|
||||
|
||||
|
||||
class CVAENew(tf.keras.Model):
|
||||
|
||||
def __init__(self, config):
|
||||
super(CVAENew, self).__init__()
|
||||
self.size = config['resolution']
|
||||
|
||||
if not('only_curvatures' in config.keys() and config['only_curvatures'] == True):
|
||||
self.static_enc : AEEncoderNew = AEEncoderNew(config['cond_ae_static']['encoder_depths'], config['cond_ae_static']['encoder_dims'], config['cond_ae_static']['kernel_size'], config['cond_ae_static']['activation'])
|
||||
if config['static_model'] != '':
|
||||
load_model(self.static_enc, config['static_model'] + '_enc')
|
||||
|
||||
self.body_enc : AEEncoderNew = AEEncoderNew(config['cond_ae_body']['encoder_depths'], config['cond_ae_body']['encoder_dims'], config['cond_ae_body']['kernel_size'])
|
||||
if config['body_model'] != '':
|
||||
load_model(self.body_enc, config['body_model'] + '_enc')
|
||||
|
||||
if config['use_curvature']:
|
||||
self.curv_enc : AEEncoderNew = AEEncoderNew(config['cond_ae_curv']['encoder_depths'], config['cond_ae_curv']['encoder_dims'], config['cond_ae_curv']['kernel_size'], config['cond_ae_curv']['activation'])
|
||||
if config['curvature_model'] != '':
|
||||
load_model(self.curv_enc, config['curvature_model'] + '_enc')
|
||||
|
||||
if config['freeze_encoders']:
|
||||
if not('only_curvatures' in config.keys() and config['only_curvatures'] == True):
|
||||
self.static_enc.trainable = False
|
||||
self.body_enc.trainable = False
|
||||
if config['use_curvature']:
|
||||
self.curv_enc.trainable = False
|
||||
|
||||
self.encoder = VariationalEncoderNew(config['cvae']['encoder_depths'], config['cvae']['encoder_dims'], kernel_size=config['cvae']['kernel_size'])
|
||||
self.decoder = DecoderNew(config['cvae']['decoder_depths'], config['cvae']['decoder_dims'], config['cvae']['kernel_size'])
|
||||
|
||||
self.reshape = tf.keras.layers.Reshape((8, 4, 64))
|
||||
|
||||
if config['avg_garment_path'] != '':
|
||||
self.avg_garment = np.load(config['avg_garment_path'])
|
||||
else:
|
||||
self.avg_garment = None
|
||||
|
||||
self.add_noise = config['cvae']['add_noise']
|
||||
self.noise_amount = config['cvae']['noise_amount']
|
||||
self.use_curvature = config['use_curvature']
|
||||
self.only_curvatures = config['only_curvatures'] if 'only_curvatures' in config.keys() and config['only_curvatures'] == True else None
|
||||
self.use_offsets_over_body = not config['use_skinning']
|
||||
|
||||
if self.use_offsets_over_body:
|
||||
assert self.avg_garment is None, 'When we are not using skinning, we can\'t use avg_garment!'
|
||||
if self.avg_garment is not None:
|
||||
assert not self.use_offsets_over_body, 'When avg_garment is given, we must use skinning!'
|
||||
|
||||
# @tf.function
|
||||
def sample(self, uv_template, uv_body_posed, uv_curv, uv_mask, noise=None, add_noise=False):
|
||||
if noise is None:
|
||||
eps = tf.random.normal(shape=(uv_template.shape[0], self.latent_dim))
|
||||
else:
|
||||
eps = noise
|
||||
|
||||
z_cond = self.encode_conds(uv_template, uv_body_posed, uv_curv, training=False)
|
||||
if add_noise:
|
||||
z_cond = z_cond + (tf.random.uniform(z_cond.shape, maxval=self.noise_amount * 2) - self.noise_amount)
|
||||
|
||||
z = tf.concat([eps, z_cond], axis=-1)
|
||||
|
||||
x = self.decoder(z, training=False)
|
||||
|
||||
if self.avg_garment is not None:
|
||||
x = x + self.avg_garment[tf.newaxis,:,:,:]
|
||||
|
||||
if self.use_offsets_over_body:
|
||||
x = x + uv_body_posed
|
||||
|
||||
x = x * uv_mask[:,:,:,tf.newaxis]
|
||||
|
||||
return x
|
||||
|
||||
|
||||
def reparameterize(self, mean, logvar, training, fix_noise=False):
|
||||
if training == True:
|
||||
if not fix_noise:
|
||||
eps = tf.random.normal(shape=mean.shape)
|
||||
else:
|
||||
eps = tf.random.normal(shape=(1, mean.shape[1]))
|
||||
eps = tf.tile(eps, [mean.shape[0],1])
|
||||
out = eps * tf.exp(logvar * .5) + mean
|
||||
else:
|
||||
out = tf.exp(logvar * .5) + mean
|
||||
return out
|
||||
|
||||
|
||||
def encode_conds(self, uv_template, uv_body_posed, uv_curv, training):
|
||||
if self.only_curvatures:
|
||||
z_cond = self.curv_enc(uv_curv, training=training)
|
||||
else:
|
||||
z_body = self.body_enc(uv_body_posed, training=training)
|
||||
z_static = self.static_enc(uv_template, training=training)
|
||||
if self.use_curvature:
|
||||
z_curv = self.curv_enc(uv_curv, training=training)
|
||||
z_cond = tf.concat([z_static, z_body, z_curv], axis=-1)
|
||||
else:
|
||||
z_cond = tf.concat([z_static, z_body], axis=-1)
|
||||
|
||||
return z_cond
|
||||
|
||||
|
||||
def call(self, inputs, training):
|
||||
uv_unposed, uv_template, uv_body_posed, uv_curv, uv_mask = inputs
|
||||
if self.only_curvatures:
|
||||
x = tf.concat([uv_unposed, uv_curv], axis=-1)
|
||||
else:
|
||||
x = tf.concat([uv_unposed, uv_template, uv_body_posed, uv_curv], axis=-1)
|
||||
|
||||
mean, logvar = self.encoder(x, training=training)
|
||||
z = self.reparameterize(mean, logvar, training)
|
||||
z = self.reshape(z)
|
||||
z_cond = self.encode_conds(uv_template, uv_body_posed, uv_curv, training=training)
|
||||
|
||||
if self.add_noise and training == True:
|
||||
z_cond = z_cond + (tf.random.uniform(z_cond.shape, maxval=self.noise_amount * 2) - self.noise_amount)
|
||||
|
||||
|
||||
z = tf.concat([z, z_cond], axis=-1)
|
||||
x = self.decoder(z, training=training)
|
||||
|
||||
if self.avg_garment is not None:
|
||||
x = x + self.avg_garment[tf.newaxis,:,:,:]
|
||||
|
||||
if self.use_offsets_over_body:
|
||||
x = x + uv_body_posed
|
||||
|
||||
x = x * uv_mask[:,:,:,tf.newaxis]
|
||||
|
||||
return x, mean, logvar
|
||||
|
||||
@@ -0,0 +1,181 @@
|
||||
import os
|
||||
import time
|
||||
import random
|
||||
import tensorflow as tf
|
||||
import numpy as np
|
||||
import scipy as sp
|
||||
# import plotly.graph_objects as go
|
||||
# import plotly.express as ex
|
||||
|
||||
from custom_utils.custom_utils import load_model
|
||||
from models.base import AEEncoder, VariationalEncoder, Decoder, VariationalEncoderNew, DecoderNew
|
||||
from models.upscaler import UpscalerBasic, UpscalerBasicConv, UpscalerBasicNew
|
||||
|
||||
|
||||
class PyramidLevel(tf.keras.Model):
|
||||
|
||||
def __init__(self, config):
|
||||
super(PyramidLevel, self).__init__()
|
||||
self.latent_dim = config['latent_dim']
|
||||
self.size = config['resolution']
|
||||
|
||||
self.encoder = VariationalEncoder(config['latent_dim'], config['encoder_depths'], config['encoder_dims'], kernel_size=config['kernel_size'])
|
||||
self.decoder = Decoder(config['decoder_input'], config['decoder_depths'], config['decoder_dims'], config['kernel_size'])
|
||||
|
||||
self.upscale_enhancer = None
|
||||
if config['upscale_enhancer'] != '':
|
||||
k_size = config['upscale_enhancer_kernel_size'] if 'upscale_enhancer_kernel_size' in config.keys() else ((3 if config['resolution'] <= 64 else 5))
|
||||
self.upscale_enhancer = UpscalerBasicConv(config['resolution'], kernel_size=k_size)
|
||||
load_model(self.upscale_enhancer, config['upscale_enhancer'])
|
||||
self.upscale_enhancer.trainable = not config['upscale_enhancer_freeze']
|
||||
|
||||
if config['avg_garment_path'] != '':
|
||||
self.avg_garment = np.load(config['avg_garment_path'])
|
||||
else:
|
||||
self.avg_garment = None
|
||||
|
||||
|
||||
def enhance(self, low, mask, training=None):
|
||||
x = self.upscale_enhancer([low, mask], training=training)
|
||||
return x
|
||||
|
||||
|
||||
def sample(self, z_cond, uv_mask, noise=None):
|
||||
if noise is None:
|
||||
eps = tf.random.normal(shape=(uv_mask.shape[0], self.latent_dim))
|
||||
else:
|
||||
eps = noise
|
||||
|
||||
z = tf.concat([eps, z_cond], axis=-1)
|
||||
|
||||
x = self.decoder(z)
|
||||
|
||||
if self.avg_garment is not None:
|
||||
x = x + self.avg_garment[tf.newaxis,:,:,:]
|
||||
|
||||
x = x * uv_mask[:,:,:,tf.newaxis]
|
||||
|
||||
return x
|
||||
|
||||
|
||||
def reparameterize(self, mean, logvar, training, fix_noise=False):
|
||||
if training == True:
|
||||
if not fix_noise:
|
||||
eps = tf.random.normal(shape=mean.shape)
|
||||
else:
|
||||
eps = tf.random.normal(shape=(1, mean.shape[1]))
|
||||
eps = tf.tile(eps, [mean.shape[0],1])
|
||||
out = eps * tf.exp(logvar * .5) + mean
|
||||
else:
|
||||
out = tf.exp(logvar * .5) + mean
|
||||
return out
|
||||
|
||||
|
||||
def call(self, inputs, z_cond, training=None):
|
||||
uv_unposed, uv_template, uv_body_posed, uv_curv, uv_mask, uv_mask_high = inputs
|
||||
|
||||
x = tf.concat([uv_unposed, uv_template, uv_body_posed, uv_curv], axis=-1)
|
||||
mean, logvar = self.encoder(x, training=training)
|
||||
z = self.reparameterize(mean, logvar, training)
|
||||
|
||||
z = tf.concat([z, z_cond], axis=-1)
|
||||
x = self.decoder(z, training=training)
|
||||
|
||||
if self.avg_garment is not None:
|
||||
x = x + self.avg_garment[tf.newaxis,:,:,:]
|
||||
|
||||
x = x * uv_mask[:,:,:,tf.newaxis]
|
||||
|
||||
# if self.upscale_enhancer is not None:
|
||||
# pred_low = self.enhance(x, uv_mask, training=training)
|
||||
# x = tf.image.resize(pred_low, (512, 256), tf.image.ResizeMethod.LANCZOS5)
|
||||
|
||||
# x = x * uv_mask_high[:,:,:,tf.newaxis]
|
||||
|
||||
return x, mean, logvar
|
||||
|
||||
|
||||
|
||||
class PyramidLevelNew(tf.keras.Model):
|
||||
|
||||
def __init__(self, config):
|
||||
super(PyramidLevelNew, self).__init__()
|
||||
self.size = config['resolution']
|
||||
|
||||
self.encoder = VariationalEncoderNew(config['encoder_depths'], config['encoder_dims'], kernel_size=config['kernel_size'])
|
||||
self.decoder = DecoderNew(config['decoder_depths'], config['decoder_dims'], config['kernel_size'])
|
||||
|
||||
self.upscale_enhancer = None
|
||||
if config['upscale_enhancer'] != '':
|
||||
k_size = config['upscale_enhancer_kernel_size'] if 'upscale_enhancer_kernel_size' in config.keys() else ((3 if config['resolution'] <= 64 else 5))
|
||||
self.upscale_enhancer = UpscalerBasicConv(config['resolution'], kernel_size=k_size)
|
||||
load_model(self.upscale_enhancer, config['upscale_enhancer'])
|
||||
self.upscale_enhancer.trainable = not config['upscale_enhancer_freeze']
|
||||
|
||||
if config['avg_garment_path'] != '':
|
||||
self.avg_garment = np.load(config['avg_garment_path'])
|
||||
else:
|
||||
self.avg_garment = None
|
||||
|
||||
self.reshape = tf.keras.layers.Reshape((8, 4, 64))
|
||||
|
||||
|
||||
def enhance(self, low, mask, training=None):
|
||||
x = self.upscale_enhancer([low, mask], training=training)
|
||||
return x
|
||||
|
||||
|
||||
def sample(self, z_cond, uv_mask, noise=None):
|
||||
if noise is None:
|
||||
eps = tf.random.normal(shape=(uv_mask.shape[0], self.latent_dim))
|
||||
else:
|
||||
eps = noise
|
||||
|
||||
z = tf.concat([eps, z_cond], axis=-1)
|
||||
|
||||
x = self.decoder(z)
|
||||
|
||||
if self.avg_garment is not None:
|
||||
x = x + self.avg_garment[tf.newaxis,:,:,:]
|
||||
|
||||
x = x * uv_mask[:,:,:,tf.newaxis]
|
||||
|
||||
return x
|
||||
|
||||
|
||||
def reparameterize(self, mean, logvar, training, fix_noise=False):
|
||||
if training == True:
|
||||
if not fix_noise:
|
||||
eps = tf.random.normal(shape=mean.shape)
|
||||
else:
|
||||
eps = tf.random.normal(shape=(1, mean.shape[1]))
|
||||
eps = tf.tile(eps, [mean.shape[0],1])
|
||||
out = eps * tf.exp(logvar * .5) + mean
|
||||
else:
|
||||
out = tf.exp(logvar * .5) + mean
|
||||
return out
|
||||
|
||||
|
||||
def call(self, inputs, z_cond, training=None):
|
||||
uv_unposed, uv_template, uv_body_posed, uv_curv, uv_mask, uv_mask_high = inputs
|
||||
|
||||
x = tf.concat([uv_unposed, uv_template, uv_body_posed, uv_curv], axis=-1)
|
||||
mean, logvar = self.encoder(x, training=training)
|
||||
z = self.reparameterize(mean, logvar, training)
|
||||
z = self.reshape(z)
|
||||
|
||||
z = tf.concat([z, z_cond], axis=-1)
|
||||
x = self.decoder(z, training=training)
|
||||
|
||||
if self.avg_garment is not None:
|
||||
x = x + self.avg_garment[tf.newaxis,:,:,:]
|
||||
|
||||
x = x * uv_mask[:,:,:,tf.newaxis]
|
||||
|
||||
# if self.upscale_enhancer is not None:
|
||||
# pred_low = self.enhance(x, uv_mask, training=training)
|
||||
# x = tf.image.resize(pred_low, (512, 256), tf.image.ResizeMethod.LANCZOS5)
|
||||
|
||||
# x = x * uv_mask_high[:,:,:,tf.newaxis]
|
||||
|
||||
return x, mean, logvar
|
||||
@@ -0,0 +1,365 @@
|
||||
import os
|
||||
import time
|
||||
import random
|
||||
import tensorflow as tf
|
||||
import numpy as np
|
||||
# import plotly.graph_objects as go
|
||||
# import plotly.express as ex
|
||||
|
||||
from custom_utils.custom_utils import load_model
|
||||
from models.base import AEEncoder, AEEncoderNew
|
||||
from models.pyramid_level import PyramidLevel, PyramidLevelNew
|
||||
from models.upscaler import UpscalerBasic, UpscalerBasicConv, UpscalerBasicNew
|
||||
|
||||
|
||||
class Pyramid(tf.keras.Model):
|
||||
|
||||
def __init__(self, config):
|
||||
super(Pyramid, self).__init__()
|
||||
|
||||
self.static_enc : AEEncoder = AEEncoder(config['cond_ae']['latent_dim'], config['cond_ae']['encoder_depths'], config['cond_ae']['encoder_dims'], config['cond_ae']['bottleneck_size'], config['cond_ae']['kernel_size'])
|
||||
if config['pyramid']['static_model'] != '':
|
||||
load_model(self.static_enc, config['pyramid']['static_model'] + '_enc')
|
||||
|
||||
self.body_enc : AEEncoder = AEEncoder(config['cond_ae_body']['latent_dim'], config['cond_ae_body']['encoder_depths'], config['cond_ae_body']['encoder_dims'], config['cond_ae_body']['bottleneck_size'], config['cond_ae_body']['kernel_size'])
|
||||
if config['pyramid']['body_model'] != '':
|
||||
load_model(self.body_enc, config['pyramid']['body_model'] + '_enc')
|
||||
|
||||
self.curv_enc : AEEncoder = AEEncoder(config['cond_ae']['latent_dim'], config['cond_ae']['encoder_depths'], config['cond_ae']['encoder_dims'], config['cond_ae']['bottleneck_size'], config['cond_ae']['kernel_size'])
|
||||
if config['pyramid']['curvature_model'] != '':
|
||||
load_model(self.curv_enc, config['pyramid']['curvature_model'] + '_enc')
|
||||
|
||||
if config['pyramid']['freeze_encoders']:
|
||||
self.static_enc.trainable = False
|
||||
self.body_enc.trainable = False
|
||||
self.curv_enc.trainable = False
|
||||
|
||||
self.cond_encoder = tf.keras.Sequential([tf.keras.layers.Dense(512, activation='relu'),
|
||||
tf.keras.layers.Dense(256, activation='relu'),
|
||||
tf.keras.layers.Dense(128, activation='relu'),
|
||||
tf.keras.layers.Dense(config['pyramid']['latent_dim'], activation='tanh'),])
|
||||
|
||||
self.levels = [PyramidLevel({ **config['pyramid']['levels'][i], **{ 'cond_ae': config['cond_ae']} }) for i in range(len(config['pyramid']['levels']))]
|
||||
|
||||
# if config['is_cape']:
|
||||
# if len(config['pyramid']['resolutions']) == 2:
|
||||
# self.levels[-1].decoder.decoder.last = tf.keras.Sequential([tf.keras.layers.Dense(3), tf.keras.layers.Lambda(lambda x: x / 100)])
|
||||
# if len(config['pyramid']['resolutions']) == 3:
|
||||
# self.levels[-1].decoder.decoder.last = tf.keras.Sequential([tf.keras.layers.Dense(3), tf.keras.layers.Lambda(lambda x: x / 100)])
|
||||
# self.levels[-2].decoder.decoder.last = tf.keras.Sequential([tf.keras.layers.Dense(3), tf.keras.layers.Lambda(lambda x: x / 100)])
|
||||
# if len(config['pyramid']['resolutions']) == 4:
|
||||
# self.levels[-1].decoder.decoder.last = tf.keras.Sequential([tf.keras.layers.Dense(3), tf.keras.layers.Lambda(lambda x: x / 100)])
|
||||
# self.levels[-2].decoder.decoder.last = tf.keras.Sequential([tf.keras.layers.Dense(3), tf.keras.layers.Lambda(lambda x: x / 100)])
|
||||
# self.levels[-3].decoder.decoder.last = tf.keras.Sequential([tf.keras.layers.Dense(3), tf.keras.layers.Lambda(lambda x: x / 100)])
|
||||
# else:
|
||||
if len(config['pyramid']['resolutions']) == 2:
|
||||
self.levels[-1].decoder.decoder.last = tf.keras.Sequential([tf.keras.layers.Dense(3), tf.keras.layers.Lambda(lambda x: x / 100)])
|
||||
if len(config['pyramid']['resolutions']) == 3:
|
||||
self.levels[-1].decoder.decoder.last = tf.keras.Sequential([tf.keras.layers.Dense(3), tf.keras.layers.Lambda(lambda x: x / 1000)])
|
||||
self.levels[-2].decoder.decoder.last = tf.keras.Sequential([tf.keras.layers.Dense(3), tf.keras.layers.Lambda(lambda x: x / 100)])
|
||||
if len(config['pyramid']['resolutions']) == 4:
|
||||
self.levels[-1].decoder.decoder.last = tf.keras.Sequential([tf.keras.layers.Dense(3), tf.keras.layers.Lambda(lambda x: x / 10000)])
|
||||
self.levels[-2].decoder.decoder.last = tf.keras.Sequential([tf.keras.layers.Dense(3), tf.keras.layers.Lambda(lambda x: x / 1000)])
|
||||
self.levels[-3].decoder.decoder.last = tf.keras.Sequential([tf.keras.layers.Dense(3), tf.keras.layers.Lambda(lambda x: x / 100)])
|
||||
if len(config['pyramid']['resolutions']) == 5:
|
||||
self.levels[-1].decoder.decoder.last = tf.keras.Sequential([tf.keras.layers.Dense(3), tf.keras.layers.Lambda(lambda x: x / 100000)])
|
||||
self.levels[-2].decoder.decoder.last = tf.keras.Sequential([tf.keras.layers.Dense(3), tf.keras.layers.Lambda(lambda x: x / 10000)])
|
||||
self.levels[-3].decoder.decoder.last = tf.keras.Sequential([tf.keras.layers.Dense(3), tf.keras.layers.Lambda(lambda x: x / 1000)])
|
||||
self.levels[-4].decoder.decoder.last = tf.keras.Sequential([tf.keras.layers.Dense(3), tf.keras.layers.Lambda(lambda x: x / 100)])
|
||||
|
||||
self.add_noise = config['pyramid']['add_noise']
|
||||
self.noise_amount = config['pyramid']['noise_amount']
|
||||
self.use_offsets = config['use_offsets']
|
||||
|
||||
self.upscale_enhancer = None
|
||||
if config['pyramid']['final_upscale_enhancer'] != '':
|
||||
k_size = config['pyramid']['final_upscale_enhancer_kernel_size'] if 'final_upscale_enhancer_kernel_size' in config['pyramid'].keys() else ((3 if config['pyramid']['resolutions'][-1] <= 64 else 5))
|
||||
self.upscale_enhancer = UpscalerBasicConv(low_res=config['pyramid']['resolutions'][-1], kernel_size=k_size)
|
||||
load_model(self.upscale_enhancer, config['pyramid']['final_upscale_enhancer'])
|
||||
self.upscale_enhancer.trainable = False
|
||||
|
||||
def enhance(self, low, mask, training):
|
||||
x = self.upscale_enhancer([low, mask], training=training)
|
||||
return x
|
||||
|
||||
def encode_conds(self, uv_template, uv_body_posed, uv_curv, training):
|
||||
z_static = self.static_enc(uv_template, training)
|
||||
z_body = self.body_enc(uv_body_posed, training)
|
||||
z_curv = self.curv_enc(uv_curv, training)
|
||||
|
||||
z_cond = tf.concat([z_static, z_body, z_curv], axis=-1)
|
||||
z_cond = self.cond_encoder(z_cond, training=training)
|
||||
|
||||
return z_cond
|
||||
|
||||
def call(self, inputs, training):
|
||||
levels_input, conds = inputs
|
||||
uv_static, uv_body, uv_curv = conds[0], conds[1], conds[2]
|
||||
|
||||
z_cond = self.encode_conds(uv_static, uv_body, uv_curv, training)
|
||||
|
||||
if self.add_noise == True and training == True:
|
||||
z_cond = z_cond + (tf.random.uniform(z_cond.shape, maxval=self.noise_amount * 2) - self.noise_amount)
|
||||
|
||||
out_levels = []
|
||||
pred, mean, logvar = self.levels[0](levels_input[0], z_cond, training=training)
|
||||
if self.use_offsets:
|
||||
pred = pred + levels_input[0][2] # uv_body_posed_32
|
||||
if 1 != len(self.levels):
|
||||
pred = self.levels[0].enhance(pred, levels_input[0][-2], training=training) # levels_input[i][-2]=uv_mask_res
|
||||
pred = tf.image.resize(pred, (self.levels[1].size, self.levels[1].size//2), tf.image.ResizeMethod.LANCZOS5)
|
||||
out_levels.append([pred, mean, logvar])
|
||||
pred = pred * levels_input[1][-2][:,:,:,tf.newaxis]
|
||||
else:
|
||||
pred = self.enhance(pred, levels_input[0][-2], training=training)
|
||||
pred = tf.image.resize(pred, (512, 256), tf.image.ResizeMethod.LANCZOS5)
|
||||
pred = pred * levels_input[-1][-1][:,:,:,tf.newaxis] # full res uv mask
|
||||
out_levels.append([pred, mean, logvar])
|
||||
|
||||
for i in range(1, len(self.levels)):
|
||||
pred_level, mean, logvar = self.levels[i](levels_input[i], z_cond, training=training)
|
||||
pred = pred + pred_level
|
||||
if i != len(self.levels) - 1:
|
||||
pred = self.levels[i].enhance(pred, levels_input[i][-2], training=training)
|
||||
out_levels.append([pred, mean, logvar])
|
||||
pred = tf.image.resize(pred, (self.levels[i+1].size, self.levels[i+1].size//2), tf.image.ResizeMethod.LANCZOS5)
|
||||
pred = pred * levels_input[i+1][-2][:,:,:,tf.newaxis]
|
||||
else:
|
||||
if self.upscale_enhancer is not None:
|
||||
pred = self.enhance(pred, levels_input[-1][-2], training=training)
|
||||
pred = tf.image.resize(pred, (512, 256), tf.image.ResizeMethod.LANCZOS5)
|
||||
pred = pred * levels_input[-1][-1][:,:,:,tf.newaxis] # full res uv mask
|
||||
out_levels.append([pred, mean, logvar])
|
||||
|
||||
return out_levels
|
||||
|
||||
# # 32
|
||||
# pred_0, mean_0, logvar_0 = self.levels[0](levels_input[0], z_cond, training=training)
|
||||
# pred_0 = self.enhance(pred_0, levels_input[0][-2], training=training)
|
||||
# pred_0 = tf.image.resize(pred_0, (512, 256), tf.image.ResizeMethod.LANCZOS5)
|
||||
# pred_0 = pred_0 * levels_input[-1][-1][:,:,:,tf.newaxis] # full res uv mask
|
||||
# return [ [pred_0, mean_0, logvar_0] ]
|
||||
|
||||
# # 32 64
|
||||
# pred_0, mean_0, logvar_0 = self.levels[0](levels_input[0], z_cond, training=training)
|
||||
# pred_0 = self.levels[0].enhance(pred_0, levels_input[0][-2], training=training) # levels_input[i][-2]=uv_mask_res
|
||||
# pred_0 = tf.image.resize(pred_0, (self.levels[1].size, self.levels[1].size//2), tf.image.ResizeMethod.LANCZOS5)
|
||||
# pred = pred_0 * levels_input[1][-2][:,:,:,tf.newaxis]
|
||||
|
||||
# pred_1, mean_1, logvar_1 = self.levels[1](levels_input[1], z_cond, training=training)
|
||||
# pred_1 = pred + pred_1
|
||||
# pred_1 = self.enhance(pred_1, levels_input[-1][-2], training=training)
|
||||
# pred_1 = tf.image.resize(pred_1, (512, 256), tf.image.ResizeMethod.LANCZOS5)
|
||||
# pred_1 = pred_1 * levels_input[-1][-1][:,:,:,tf.newaxis] # full res uv mask
|
||||
|
||||
# return [ [pred_0, mean_0, logvar_0], [pred_1, mean_1, logvar_1] ]
|
||||
|
||||
|
||||
# class Pyramid(tf.keras.Model):
|
||||
|
||||
# def __init__(self, config):
|
||||
# super(Pyramid, self).__init__()
|
||||
|
||||
# self.static_enc : AEEncoder = AEEncoder(config['cond_ae']['latent_dim'], config['cond_ae']['encoder_depths'], config['cond_ae']['encoder_dims'], config['cond_ae']['bottleneck_size'], config['cond_ae']['kernel_size'])
|
||||
# if config['pyramid']['static_model'] != '':
|
||||
# load_model(self.static_enc, config['pyramid']['static_model'] + '_enc')
|
||||
|
||||
# self.body_enc : AEEncoder = AEEncoder(config['cond_ae']['latent_dim'], config['cond_ae']['encoder_depths'], config['cond_ae']['encoder_dims'], config['cond_ae']['bottleneck_size'], config['cond_ae']['kernel_size'])
|
||||
# if config['pyramid']['body_model'] != '':
|
||||
# load_model(self.body_enc, config['pyramid']['body_model'] + '_enc')
|
||||
|
||||
# self.curv_enc : AEEncoder = AEEncoder(config['cond_ae']['latent_dim'], config['cond_ae']['encoder_depths'], config['cond_ae']['encoder_dims'], config['cond_ae']['bottleneck_size'], config['cond_ae']['kernel_size'])
|
||||
# if config['pyramid']['curvature_model'] != '':
|
||||
# load_model(self.curv_enc, config['pyramid']['curvature_model'] + '_enc')
|
||||
|
||||
# if config['pyramid']['freeze_encoders']:
|
||||
# self.static_enc.trainable = False
|
||||
# self.body_enc.trainable = False
|
||||
# self.curv_enc.trainable = False
|
||||
|
||||
# self.cond_encoder = tf.keras.Sequential([tf.keras.layers.Dense(512, activation='relu'),
|
||||
# tf.keras.layers.Dense(256, activation='relu'),
|
||||
# tf.keras.layers.Dense(128, activation='relu'),
|
||||
# tf.keras.layers.Dense(config['pyramid']['latent_dim']),])
|
||||
|
||||
# self.levels = [PyramidLevel({ **config['pyramid']['levels'][i], **{ 'cond_ae': config['cond_ae']} }) for i in range(len(config['pyramid']['levels']))]
|
||||
|
||||
|
||||
# self.upscale_enhancer = None
|
||||
# if config['pyramid']['final_upscale_enhancer'] != '':
|
||||
# self.upscale_enhancer = UpscalerBasicNew(low_res=config['pyramid']['resolutions'][-1])
|
||||
# load_model(self.upscale_enhancer, config['pyramid']['final_upscale_enhancer'])
|
||||
# self.upscale_enhancer.trainable = False
|
||||
|
||||
# def enhance(self, low, mask, training):
|
||||
# x = self.upscale_enhancer([low, mask], training=training)
|
||||
# return x
|
||||
|
||||
# def encode_conds(self, uv_template, uv_body_posed, uv_curv):
|
||||
# z_static = self.static_enc(uv_template)
|
||||
# z_body = self.body_enc(uv_body_posed)
|
||||
# z_curv = self.curv_enc(uv_curv)
|
||||
|
||||
# z_cond = tf.concat([z_static, z_body, z_curv], axis=-1)
|
||||
# z_cond = self.cond_encoder(z_cond)
|
||||
|
||||
# return z_cond
|
||||
|
||||
# def call(self, inputs, training):
|
||||
# levels_input, conds = inputs
|
||||
# uv_static, uv_body, uv_curv = conds[0], conds[1], conds[2]
|
||||
# z_cond = self.encode_conds(uv_static, uv_body, uv_curv)
|
||||
|
||||
# out_levels = []
|
||||
# pred_0, mean_0, logvar_0 = self.levels[0](levels_input[0], z_cond, training=training)
|
||||
# pred_0 = self.levels[0].enhance(pred_0, levels_input[0][-2], training=training) # levels_input[i][-2]=uv_mask_res
|
||||
# # pred_0 = self.enhance(pred_0, levels_input[0][-2], training=training)
|
||||
|
||||
# ##################
|
||||
# # if more than 1 level
|
||||
# # x = tf.image.resize(pred_0, (self.levels[1].size, self.levels[1].size//2), tf.image.ResizeMethod.LANCZOS5)
|
||||
# # x = x * levels_input[1][-2][:,:,:,tf.newaxis]
|
||||
|
||||
# # if only 1 level
|
||||
# pred_0 = tf.image.resize(pred_0, (512, 256), tf.image.ResizeMethod.LANCZOS5)
|
||||
# ##################
|
||||
|
||||
|
||||
# # pred_1, mean_1, logvar_1 = self.levels[1](levels_input[1], z_cond, training=training)
|
||||
# # pred_1 = pred_0 + pred_1
|
||||
# # pred_1 = self.levels[1].enhance(pred, levels_input[1][-2], training=training)
|
||||
|
||||
# # for i in range(1, len(self.levels)):
|
||||
# # pred_level, mean, logvar = self.levels[i](levels_input[i], z_cond, training=training)
|
||||
# # pred = pred + pred_level
|
||||
# # if i != len(self.levels) - 1:
|
||||
# # pred = self.levels[i].enhance(pred, levels_input[i][-2], training=training)
|
||||
# # out_levels.append([pred, mean, logvar])
|
||||
# # pred = tf.image.resize(pred, (self.levels[i+1].size, self.levels[i+1].size//2), tf.image.ResizeMethod.LANCZOS5)
|
||||
# # pred = pred * levels_input[i+1][-2][:,:,:,tf.newaxis]
|
||||
# # else:
|
||||
# # if self.upscale_enhancer is not None:
|
||||
# # pred = self.enhance(pred, levels_input[-1][-2], training=training)
|
||||
# # pred = tf.image.resize(pred, (512, 256), tf.image.ResizeMethod.LANCZOS5)
|
||||
# # out_levels.append([pred, mean, logvar])
|
||||
# # pred = pred * levels_input[-1][-1][:,:,:,tf.newaxis] # full res uv mask
|
||||
|
||||
# out_levels = [
|
||||
# [pred_0, mean_0, logvar_0]
|
||||
# ]
|
||||
|
||||
# return out_levels
|
||||
|
||||
|
||||
class PyramidNew(tf.keras.Model):
|
||||
|
||||
def __init__(self, config):
|
||||
super(PyramidNew, self).__init__()
|
||||
|
||||
self.static_enc : AEEncoderNew = AEEncoderNew(config['cond_ae_static']['encoder_depths'], config['cond_ae_static']['encoder_dims'], config['cond_ae_static']['kernel_size'], config['cond_ae_static']['activation'])
|
||||
if config['pyramid']['static_model'] != '':
|
||||
load_model(self.static_enc, config['pyramid']['static_model'] + '_enc')
|
||||
|
||||
self.body_enc : AEEncoderNew = AEEncoderNew(config['cond_ae_body']['encoder_depths'], config['cond_ae_body']['encoder_dims'], config['cond_ae_body']['kernel_size'])
|
||||
if config['pyramid']['body_model'] != '':
|
||||
load_model(self.body_enc, config['pyramid']['body_model'] + '_enc')
|
||||
|
||||
self.curv_enc : AEEncoderNew = AEEncoderNew(config['cond_ae_curv']['encoder_depths'], config['cond_ae_curv']['encoder_dims'], config['cond_ae_curv']['kernel_size'], config['cond_ae_curv']['activation'])
|
||||
if config['pyramid']['curvature_model'] != '':
|
||||
load_model(self.curv_enc, config['pyramid']['curvature_model'] + '_enc')
|
||||
|
||||
if config['pyramid']['freeze_encoders']:
|
||||
self.static_enc.trainable = False
|
||||
self.body_enc.trainable = False
|
||||
self.curv_enc.trainable = False
|
||||
|
||||
self.levels = [PyramidLevelNew({ **config['pyramid']['levels'][i] }) for i in range(len(config['pyramid']['levels']))]
|
||||
|
||||
|
||||
# if len(config['pyramid']['resolutions']) == 2:
|
||||
# self.levels[-1].decoder.decoder.last = tf.keras.Sequential([tf.keras.layers.Dense(3), tf.keras.layers.Lambda(lambda x: x / 100)])
|
||||
# if len(config['pyramid']['resolutions']) == 3:
|
||||
# self.levels[-1].decoder.decoder.last = tf.keras.Sequential([tf.keras.layers.Dense(3), tf.keras.layers.Lambda(lambda x: x / 1000)])
|
||||
# self.levels[-2].decoder.decoder.last = tf.keras.Sequential([tf.keras.layers.Dense(3), tf.keras.layers.Lambda(lambda x: x / 100)])
|
||||
# if len(config['pyramid']['resolutions']) == 4:
|
||||
# self.levels[-1].decoder.decoder.last = tf.keras.Sequential([tf.keras.layers.Dense(3), tf.keras.layers.Lambda(lambda x: x / 10000)])
|
||||
# self.levels[-2].decoder.decoder.last = tf.keras.Sequential([tf.keras.layers.Dense(3), tf.keras.layers.Lambda(lambda x: x / 1000)])
|
||||
# self.levels[-3].decoder.decoder.last = tf.keras.Sequential([tf.keras.layers.Dense(3), tf.keras.layers.Lambda(lambda x: x / 100)])
|
||||
# if len(config['pyramid']['resolutions']) == 5:
|
||||
# self.levels[-1].decoder.decoder.last = tf.keras.Sequential([tf.keras.layers.Dense(3), tf.keras.layers.Lambda(lambda x: x / 100000)])
|
||||
# self.levels[-2].decoder.decoder.last = tf.keras.Sequential([tf.keras.layers.Dense(3), tf.keras.layers.Lambda(lambda x: x / 10000)])
|
||||
# self.levels[-3].decoder.decoder.last = tf.keras.Sequential([tf.keras.layers.Dense(3), tf.keras.layers.Lambda(lambda x: x / 1000)])
|
||||
# self.levels[-4].decoder.decoder.last = tf.keras.Sequential([tf.keras.layers.Dense(3), tf.keras.layers.Lambda(lambda x: x / 100)])
|
||||
|
||||
self.add_noise = config['pyramid']['add_noise']
|
||||
self.noise_amount = config['pyramid']['noise_amount']
|
||||
self.use_offsets = config['use_offsets']
|
||||
|
||||
self.upscale_enhancer = None
|
||||
if config['pyramid']['final_upscale_enhancer'] != '':
|
||||
k_size = config['pyramid']['final_upscale_enhancer_kernel_size'] if 'final_upscale_enhancer_kernel_size' in config['pyramid'].keys() else ((3 if config['pyramid']['resolutions'][-1] <= 64 else 5))
|
||||
self.upscale_enhancer = UpscalerBasicConv(low_res=config['pyramid']['resolutions'][-1], kernel_size=k_size)
|
||||
load_model(self.upscale_enhancer, config['pyramid']['final_upscale_enhancer'])
|
||||
self.upscale_enhancer.trainable = False
|
||||
|
||||
def enhance(self, low, mask, training):
|
||||
x = self.upscale_enhancer([low, mask], training=training)
|
||||
return x
|
||||
|
||||
def encode_conds(self, uv_template, uv_body_posed, uv_curv, training):
|
||||
z_static = self.static_enc(uv_template, training)
|
||||
z_body = self.body_enc(uv_body_posed, training)
|
||||
z_curv = self.curv_enc(uv_curv, training)
|
||||
|
||||
z_cond = tf.concat([z_static, z_body, z_curv], axis=-1)
|
||||
|
||||
return z_cond
|
||||
|
||||
def call(self, inputs, training):
|
||||
levels_input, conds = inputs
|
||||
uv_static, uv_body, uv_curv = conds[0], conds[1], conds[2]
|
||||
|
||||
z_cond = self.encode_conds(uv_static, uv_body, uv_curv, training)
|
||||
|
||||
if self.add_noise == True and training == True:
|
||||
z_cond = z_cond + (tf.random.uniform(z_cond.shape, maxval=self.noise_amount * 2) - self.noise_amount)
|
||||
|
||||
out_levels = []
|
||||
pred, mean, logvar = self.levels[0](levels_input[0], z_cond, training=training)
|
||||
if self.use_offsets:
|
||||
pred = pred + levels_input[0][2] # uv_body_posed_32
|
||||
if 1 != len(self.levels):
|
||||
pred = self.levels[0].enhance(pred, levels_input[0][-2], training=training) # levels_input[i][-2]=uv_mask_res
|
||||
pred = tf.image.resize(pred, (self.levels[1].size, self.levels[1].size//2), tf.image.ResizeMethod.BILINEAR)
|
||||
out_levels.append([pred, mean, logvar])
|
||||
pred = pred * levels_input[1][-2][:,:,:,tf.newaxis]
|
||||
else:
|
||||
# pred = self.enhance(pred, levels_input[0][-2], training=training)
|
||||
pred = self.levels[0].enhance(pred, levels_input[0][-2], training=training) # levels_input[i][-2]=uv_mask_res
|
||||
pred = tf.image.resize(pred, (512, 256), tf.image.ResizeMethod.BILINEAR)
|
||||
pred = pred * levels_input[-1][-1][:,:,:,tf.newaxis] # full res uv mask
|
||||
out_levels.append([pred, mean, logvar])
|
||||
|
||||
for i in range(1, len(self.levels)):
|
||||
input = levels_input[i]
|
||||
input[0] = input[0] - pred
|
||||
input_min, input_max = tf.reduce_min(input[0], axis=[1,2,3]), tf.reduce_max(input[0], axis=[1,2,3])
|
||||
input[0] = (input[0] - input_min[:,None,None,None]) / (input_max - input_min)[:,None,None,None]
|
||||
input[0] = input[0] * input[-2][:,:,:,None] # uv_mask_res
|
||||
pred_level, mean, logvar = self.levels[i](input, z_cond, training=training)
|
||||
pred_level = ( (pred_level * (input_max - input_min)[:,None,None,None]) + input_min[:,None,None,None] ) * input[-2][:,:,:,None] # uv_mask_res
|
||||
pred = pred + pred_level
|
||||
if i != len(self.levels) - 1:
|
||||
pred = self.levels[i].enhance(pred, levels_input[i][-2], training=training)
|
||||
out_levels.append([pred, mean, logvar])
|
||||
pred = tf.image.resize(pred, (self.levels[i+1].size, self.levels[i+1].size//2), tf.image.ResizeMethod.BILINEAR)
|
||||
pred = pred * levels_input[i+1][-2][:,:,:,tf.newaxis]
|
||||
else:
|
||||
# if self.upscale_enhancer is not None:
|
||||
# pred = self.enhance(pred, levels_input[-1][-2], training=training)
|
||||
pred = self.levels[-1].enhance(pred, levels_input[-1][-2], training=training)
|
||||
pred = tf.image.resize(pred, (512, 256), tf.image.ResizeMethod.BILINEAR)
|
||||
pred = pred * levels_input[-1][-1][:,:,:,tf.newaxis] # full res uv mask
|
||||
out_levels.append([pred, mean, logvar])
|
||||
|
||||
return out_levels
|
||||
@@ -0,0 +1,172 @@
|
||||
import tensorflow as tf
|
||||
from convnext.models.convnext_tf import ConvNeXt, ConvNeXtReverse, Block, ConvNeXtReverseWithSkip, ConvNeXtWithSkip
|
||||
|
||||
|
||||
class Bottleneck(tf.keras.Model):
|
||||
|
||||
def __init__(self, dim, kernel_size=7):
|
||||
super(Bottleneck, self).__init__()
|
||||
|
||||
self.bottleneck = tf.keras.Sequential(
|
||||
[Block(dim=dim, drop_path=0, kernel_size=kernel_size,
|
||||
layer_scale_init_value=1e-6, prefix=f'bottleneck{j}') for j in range(5)]
|
||||
)
|
||||
|
||||
def call(self, inputs, training):
|
||||
x = inputs
|
||||
|
||||
x = self.bottleneck(x, training=training)
|
||||
|
||||
return x
|
||||
|
||||
|
||||
|
||||
class Upscaler(tf.keras.Model):
|
||||
|
||||
def __init__(self, bottleneck_size):
|
||||
super(Upscaler, self).__init__()
|
||||
|
||||
depths = [27, 27]
|
||||
dims = [512, 512]
|
||||
reverse_dims = [16] + dims[:-1]
|
||||
self.encoder = ConvNeXtWithSkip(depths=depths, dims=dims, kernel_size=3)
|
||||
|
||||
self.bottleneck = Bottleneck(dims[-1], kernel_size=3)
|
||||
|
||||
self.decoder = ConvNeXtReverseWithSkip(depths=depths, dims=reverse_dims, kernel_size=3)
|
||||
|
||||
def call(self, inputs, training):
|
||||
img, mask = inputs
|
||||
|
||||
img = tf.transpose(img, [0,2,3,1])
|
||||
|
||||
x, skip_features = self.encoder(img, training=training)
|
||||
|
||||
x = self.bottleneck(x, training=training)
|
||||
|
||||
x = self.decoder([x, skip_features], training=training)
|
||||
|
||||
x = x * (1 - mask[:,:,:,tf.newaxis])
|
||||
x = x + img
|
||||
x = tf.image.resize(x, size=(512, 256), method=tf.image.ResizeMethod.LANCZOS5)
|
||||
|
||||
return x
|
||||
|
||||
|
||||
class UpscalerBasic(tf.keras.Model):
|
||||
|
||||
def __init__(self, bottleneck_size, low_res=32, high_res=64):
|
||||
super(UpscalerBasic, self).__init__()
|
||||
|
||||
self.low_res = low_res
|
||||
self.high_res = high_res
|
||||
self.flatten = tf.keras.layers.Flatten()
|
||||
|
||||
self.seq = tf.keras.Sequential([
|
||||
tf.keras.layers.Dense(1024, activation='relu'),
|
||||
tf.keras.layers.BatchNormalization(),
|
||||
tf.keras.layers.Dense(1024, activation='relu'),
|
||||
tf.keras.layers.BatchNormalization(),
|
||||
tf.keras.layers.Dense(1024, activation='relu'),
|
||||
tf.keras.layers.BatchNormalization(),
|
||||
tf.keras.layers.Dense(low_res * (low_res//2) * 3),
|
||||
])
|
||||
|
||||
self.reshape = tf.keras.layers.Reshape((low_res, (low_res//2), 3))
|
||||
|
||||
def call(self, inputs, training):
|
||||
img, mask = inputs
|
||||
|
||||
x = self.flatten(img)
|
||||
x = self.seq(x, training=training)
|
||||
x = self.reshape(x)
|
||||
|
||||
x = x * (1 - mask[:,:,:,tf.newaxis])
|
||||
x = x + img
|
||||
x = tf.image.resize(x, size=(self.high_res, self.high_res//2), method=tf.image.ResizeMethod.LANCZOS5)
|
||||
|
||||
return x
|
||||
|
||||
|
||||
|
||||
|
||||
class UpscalerBasicNew(tf.keras.Model):
|
||||
|
||||
def __init__(self, low_res=32):
|
||||
super(UpscalerBasicNew, self).__init__()
|
||||
|
||||
self.low_res = low_res
|
||||
self.flatten = tf.keras.layers.Flatten(name='upscale_enhancer_flattent')
|
||||
|
||||
self.seq = tf.keras.Sequential([
|
||||
tf.keras.layers.Dense(4096, activation='relu', name='upscale_enhancer_dense_1'),
|
||||
tf.keras.layers.BatchNormalization(name='upscale_enhancer_batch_norm_1'),
|
||||
tf.keras.layers.Dense(2048, activation='relu', name='upscale_enhancer_dense_2'),
|
||||
tf.keras.layers.BatchNormalization(name='upscale_enhancer_batch_norm_2'),
|
||||
tf.keras.layers.Dense(2048, activation='relu', name='upscale_enhancer_dense_3'),
|
||||
tf.keras.layers.BatchNormalization(name='upscale_enhancer_batch_norm_3'),
|
||||
tf.keras.layers.Dense(low_res * (low_res//2) * 3, activation='tanh', name='upscale_enhancer_dense_final'),
|
||||
])
|
||||
|
||||
self.reshape = tf.keras.layers.Reshape((low_res, (low_res//2), 3), name='upscale_enhancer_reshape')
|
||||
|
||||
def call(self, inputs, training):
|
||||
img, mask = inputs
|
||||
|
||||
x = self.flatten(img)
|
||||
x = self.seq(x, training=training)
|
||||
x = self.reshape(x)
|
||||
|
||||
x = x * (1 - mask[:,:,:,tf.newaxis])
|
||||
x = x + img
|
||||
|
||||
return x
|
||||
|
||||
|
||||
class UpscalerBasicConv(tf.keras.Model):
|
||||
|
||||
def __init__(self, low_res=32, kernel_size=3):
|
||||
super(UpscalerBasicConv, self).__init__()
|
||||
|
||||
self.low_res = low_res
|
||||
|
||||
self.seq = tf.keras.Sequential([
|
||||
tf.keras.layers.Conv2D(256, kernel_size=kernel_size, padding='same', activation='relu', name='upscale_enhancer_conv_1'),
|
||||
tf.keras.layers.BatchNormalization(name='upscale_enhancer_batch_norm_1'),
|
||||
tf.keras.layers.Conv2D(128, kernel_size=kernel_size, padding='same', activation='relu', name='upscale_enhancer_conv_2'),
|
||||
tf.keras.layers.BatchNormalization(name='upscale_enhancer_batch_norm_2'),
|
||||
tf.keras.layers.Conv2D(64, kernel_size=kernel_size, padding='same', activation='relu', name='upscale_enhancer_conv_3'),
|
||||
tf.keras.layers.BatchNormalization(name='upscale_enhancer_batch_norm_3'),
|
||||
tf.keras.layers.Conv2D(32, kernel_size=kernel_size, padding='same', activation='relu', name='upscale_enhancer_conv_4'),
|
||||
tf.keras.layers.BatchNormalization(name='upscale_enhancer_batch_norm_4'),
|
||||
tf.keras.layers.Dense(3, activation='tanh', name='upscale_enhancer_dense_final'),
|
||||
])
|
||||
|
||||
|
||||
def call(self, inputs, training):
|
||||
img, mask = inputs
|
||||
|
||||
x = self.seq(img, training=training)
|
||||
|
||||
x = x * (1 - mask[:,:,:,tf.newaxis])
|
||||
x = x + img
|
||||
|
||||
return x
|
||||
|
||||
# class Upscaler(tf.keras.Model):
|
||||
|
||||
# def __init__(self, bottleneck_size):
|
||||
# super(Upscaler, self).__init__()
|
||||
|
||||
# self.decoder = ConvNeXtReverse(depths=list(reversed([9, 9, 9])), dims=list(reversed([bottleneck_size, 128, 32, 16])))
|
||||
|
||||
# def call(self, inputs, training):
|
||||
# x, mask = inputs
|
||||
|
||||
# x = tf.transpose(x, [0,2,3,1])
|
||||
# x = self.decoder(x)
|
||||
|
||||
# x = x * mask[:,:,:,tf.newaxis]
|
||||
|
||||
# return x
|
||||
|
||||
+573
@@ -0,0 +1,573 @@
|
||||
import os
|
||||
import sys
|
||||
from configs.config_abl_all_post import config
|
||||
|
||||
os.environ["CUDA_VISIBLE_DEVICES"] = config['gpu']
|
||||
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)
|
||||
|
||||
import time
|
||||
import argparse
|
||||
import wandb
|
||||
import tqdm
|
||||
import random
|
||||
import numpy as np
|
||||
import plotly.graph_objects as go
|
||||
import plotly.express as ex
|
||||
from PIL import Image
|
||||
from postprocessing import fix_collision
|
||||
|
||||
if config['is_skirt']:
|
||||
from losses import edge_loss_skirt as edge_loss
|
||||
else:
|
||||
from losses import edge_loss
|
||||
|
||||
|
||||
|
||||
from losses import collision_ratio, kl_loss, normal_loss #, concept_loss
|
||||
# from losses import collision_penalty as collision_penalty
|
||||
from losses import collision_loss as collision_penalty
|
||||
from custom_utils.custom_utils import batch_to_dict, load_model, plot_3d, UV2Mesh, get_summary_html, maskFacesBatch, _maskFaces, plot_3d_collision
|
||||
from dataset import DatasetFactory, StaticDataset
|
||||
from custom_utils.smpl import SMPL
|
||||
from models.cvae_model import CVAENew as CVAE
|
||||
|
||||
|
||||
def init_wandb():
|
||||
run = wandb.init(project='ablation', name=config['name'], notes=config['notes'], dir='/wandb', tags=config['tags'])
|
||||
artifact = wandb.Artifact('code', type='code')
|
||||
artifact.add_file('cvae.py')
|
||||
artifact.add_file('models/cvae_model.py')
|
||||
artifact.add_file('losses.py')
|
||||
# artifact.add_file('configs/config_base.py')
|
||||
# artifact.add_file('configs/config_cvae.py')
|
||||
artifact.add_file('models/base.py')
|
||||
artifact.add_file('models/cond_encoder.py')
|
||||
artifact.add_file('dataset.py')
|
||||
artifact.add_file('convnext/models/convnext_tf.py')
|
||||
config_file_names = [key for key in sys.modules.keys() if key.startswith('configs.')]
|
||||
for config_file_name in config_file_names:
|
||||
file_name = '/'.join(config_file_name.split('.')) + '.py'
|
||||
if os.path.exists(file_name):
|
||||
artifact.add_file(file_name)
|
||||
run.log_artifact(artifact)
|
||||
|
||||
uv2mesh = UV2Mesh(res=256, is_skirt=config['is_skirt'], is_adaptive=False, is_smpl_orig=config['is_cape'])
|
||||
|
||||
|
||||
def kl_loss(mean, logvar):
|
||||
kl_loss = 1 + logvar - tf.square(mean) - tf.exp(logvar)
|
||||
kl_loss = tf.reduce_mean(kl_loss)
|
||||
kl_loss *= -0.5
|
||||
return kl_loss
|
||||
|
||||
|
||||
# 3d loss on posed !!!!!!!!!!!!
|
||||
|
||||
# @tf.function
|
||||
def train_step(batch):
|
||||
|
||||
with tf.GradientTape() as tape:
|
||||
if config['use_skinning']:
|
||||
mesh_orig, uv_orig = batch['mesh_unposed'], batch['uv_unposed']
|
||||
else:
|
||||
mesh_orig, uv_orig = batch['mesh_posed'], batch['uv_posed']
|
||||
|
||||
pred, mean, logvar = model([uv_orig, batch['uv_static'], batch['uv_body_posed'], batch['uv_curvature'], batch['uv_mask']], training=True)
|
||||
|
||||
mesh_pred = uv2mesh(pred)
|
||||
|
||||
if config['use_skinning']:
|
||||
shape = batch['pose_shape'][:, 72:]
|
||||
gender = tf.cast(tf.greater_equal(shape[:, 10] , 0.5), tf.int32)
|
||||
shape_offsets = tf.stack([tf.gather(tf.tensordot(shape[i,:10][tf.newaxis, :], smpl.shapedirs[gender[i]], axes=[[1], [2]])[0,:,:], smpl.to_keep, axis=0) for i in range(shape.shape[0])])
|
||||
|
||||
mesh_pred_shaped = mesh_pred + shape_offsets * batch['mesh_mask'][:,:,tf.newaxis]
|
||||
mesh_pred_posed = smpl.run_garment(batch['pose_shape'][:, :72], batch['pose_shape'][:, 72:], mesh_pred_shaped)
|
||||
else:
|
||||
mesh_pred_posed = mesh_pred
|
||||
|
||||
# is the output already masked?
|
||||
|
||||
e_loss = tf.reduce_mean(edge_loss(mesh_pred, mesh_orig, batch['mesh_mask']))
|
||||
c_loss, c_ratio = collision_penalty(mesh_pred_posed, batch['mesh_body_posed'], batch['mesh_mask'])
|
||||
c_loss, c_ratio = tf.reduce_mean(c_loss), tf.reduce_mean(c_ratio)
|
||||
l2_loss = tf.reduce_mean(1e3 * tf.reduce_mean(tf.reduce_sum(((uv_orig - pred) ** 2), axis=-1), axis=[1, 2]))
|
||||
l2_mesh_loss = tf.reduce_mean(1e3 * tf.reduce_mean(tf.reduce_sum(batch['mesh_mask'][:,:,tf.newaxis] * (batch['mesh_posed'] - mesh_pred_posed) ** 2, axis=-1), axis=-1))
|
||||
n_loss = normal_loss(mesh_pred_posed, batch['mesh_posed'], batch['mesh_mask'], is_skirt=config['is_skirt'], is_smpl_orig=config['is_cape'])
|
||||
# conc_loss = concept_loss(pred, batch['uv_unposed'])
|
||||
kl = kl_loss(mean, logvar)
|
||||
|
||||
reconstr_loss = l2_loss
|
||||
if config['use_extra_losses']:
|
||||
if 'block_collision_loss' in config.keys() and config['block_collision_loss'] == True:
|
||||
reconstr_loss = reconstr_loss + l2_mesh_loss + (e_loss) * config['reg_loss_coeff']
|
||||
else:
|
||||
reconstr_loss = reconstr_loss + l2_mesh_loss + (e_loss + c_loss) * config['reg_loss_coeff']
|
||||
if config['use_norm_loss'] and step > 2000:
|
||||
reconstr_loss = reconstr_loss + n_loss * config['reg_loss_coeff']
|
||||
|
||||
loss = reconstr_loss + g_balancing_coeff * kl
|
||||
|
||||
mean = tf.reduce_mean(mean)
|
||||
logvar = tf.reduce_mean(logvar)
|
||||
|
||||
# ms_ssim = tf.reduce_mean(tf.image.ssim_multiscale(pred + 5, batch['uv_unposed'] + 5, max_val=10, power_factors=(0.5, 0.2, 0.1, 0.1, 0.1)))
|
||||
metric = tf.reduce_sum(tf.math.sqrt(tf.reduce_sum(batch['mesh_mask'][:,:,tf.newaxis] * (mesh_orig - mesh_pred) ** 2, axis=-1))) / tf.reduce_sum(batch['mesh_mask'])
|
||||
gradients = tape.gradient(loss, model.trainable_variables)
|
||||
optimizer.apply_gradients(zip(gradients, model.trainable_variables))
|
||||
return { 'train/loss': loss, 'train/edge_loss': e_loss, 'train/collision_loss': c_loss, 'train/euclid_error': metric,
|
||||
'train/reconstr_loss': reconstr_loss, 'train/collision_ratio': c_ratio,
|
||||
'train/mean': mean, 'train/logvar': logvar, 'train/kl_loss': kl, 'train/var': np.e ** logvar,
|
||||
'train/l2_loss': l2_loss, 'train/l2_mesh_loss': l2_mesh_loss, 'train/normal_loss': n_loss,
|
||||
# 'train/ms_ssim': ms_ssim,
|
||||
}
|
||||
|
||||
|
||||
@tf.function
|
||||
def valid_step(batch):
|
||||
if config['use_skinning']:
|
||||
mesh_orig, uv_orig = batch['mesh_unposed'], batch['uv_unposed']
|
||||
else:
|
||||
mesh_orig, uv_orig = batch['mesh_posed'], batch['uv_posed']
|
||||
|
||||
pred, mean, logvar = model([uv_orig, batch['uv_static'], batch['uv_body_posed'], batch['uv_curvature'], batch['uv_mask']], training=False)
|
||||
|
||||
mesh_pred = uv2mesh(pred)
|
||||
|
||||
if config['use_skinning']:
|
||||
shape = batch['pose_shape'][:, 72:]
|
||||
gender = tf.cast(tf.greater_equal(shape[:, 10] , 0.5), tf.int32)
|
||||
shape_offsets = tf.stack([tf.gather(tf.tensordot(shape[i,:10][tf.newaxis, :], smpl.shapedirs[gender[i]], axes=[[1], [2]])[0,:,:], smpl.to_keep, axis=0) for i in range(shape.shape[0])])
|
||||
|
||||
mesh_pred_shaped = mesh_pred + shape_offsets * batch['mesh_mask'][:,:,tf.newaxis]
|
||||
mesh_pred_posed = smpl.run_garment(batch['pose_shape'][:, :72], batch['pose_shape'][:, 72:], mesh_pred_shaped)
|
||||
else:
|
||||
mesh_pred_posed = mesh_pred
|
||||
|
||||
# is the output already masked?
|
||||
|
||||
e_loss = tf.reduce_mean(edge_loss(mesh_pred, mesh_orig, batch['mesh_mask']))
|
||||
c_loss, c_ratio = collision_penalty(mesh_pred_posed, batch['mesh_body_posed'], batch['mesh_mask'])
|
||||
c_loss, c_ratio = tf.reduce_mean(c_loss), tf.reduce_mean(c_ratio)
|
||||
l2_loss = tf.reduce_mean(1e3 * tf.reduce_mean(tf.reduce_sum(((uv_orig - pred) ** 2), axis=-1), axis=[1, 2]))
|
||||
l2_mesh_loss = tf.reduce_mean(1e3 * tf.reduce_mean(tf.reduce_sum(batch['mesh_mask'][:,:,tf.newaxis] * (batch['mesh_posed'] - mesh_pred_posed) ** 2, axis=-1), axis=-1))
|
||||
n_loss = normal_loss(mesh_pred_posed, batch['mesh_posed'], batch['mesh_mask'], is_skirt=config['is_skirt'], is_smpl_orig=config['is_cape'])
|
||||
# conc_loss = concept_loss(pred, batch['uv_unposed'])
|
||||
kl = kl_loss(mean, logvar)
|
||||
|
||||
reconstr_loss = l2_loss
|
||||
if config['use_extra_losses']:
|
||||
reconstr_loss = reconstr_loss + l2_mesh_loss + (e_loss + c_loss) * config['reg_loss_coeff']
|
||||
if config['use_norm_loss']:
|
||||
reconstr_loss = reconstr_loss + n_loss * config['reg_loss_coeff']
|
||||
|
||||
loss = reconstr_loss + g_balancing_coeff * kl
|
||||
|
||||
mean = tf.reduce_mean(mean)
|
||||
logvar = tf.reduce_mean(logvar)
|
||||
|
||||
# ms_ssim = tf.reduce_mean(tf.image.ssim_multiscale(pred + 5, batch['uv_unposed'] + 5, max_val=10, power_factors=(0.5, 0.2, 0.1, 0.1, 0.1)))
|
||||
metric = tf.reduce_sum(tf.math.sqrt(tf.reduce_sum(batch['mesh_mask'][:,:,tf.newaxis] * (mesh_orig - mesh_pred) ** 2, axis=-1))) / tf.reduce_sum(batch['mesh_mask'])
|
||||
# c_ratio = tf.reduce_mean(collision_ratio(mesh_pred_posed, batch['mesh_body_posed'], batch['mesh_mask']))
|
||||
|
||||
return { 'valid/loss': loss, 'valid/edge_loss': e_loss, 'valid/collision_loss': c_loss, 'valid/euclid_error': metric,
|
||||
'valid/reconstr_loss': reconstr_loss, 'valid/l2_mesh_loss': l2_mesh_loss, 'valid/collision_ratio': c_ratio,
|
||||
'valid/mean': mean, 'valid/logvar': logvar, 'valid/kl_loss': kl, 'valid/var': np.e ** logvar,
|
||||
'valid/normal_loss': n_loss,# 'valid/ms_ssim': ms_ssim,
|
||||
}
|
||||
|
||||
|
||||
@tf.function
|
||||
def test_step(batch):
|
||||
if config['use_skinning']:
|
||||
mesh_orig, uv_orig = batch['mesh_unposed'], batch['uv_unposed']
|
||||
else:
|
||||
mesh_orig, uv_orig = batch['mesh_posed'], batch['uv_posed']
|
||||
|
||||
pred, mean, logvar = model([uv_orig, batch['uv_static'], batch['uv_body_posed'], batch['uv_curvature'], batch['uv_mask']], training=False)
|
||||
|
||||
mesh_pred = uv2mesh(pred)
|
||||
|
||||
if config['use_skinning']:
|
||||
shape = batch['pose_shape'][:, 72:]
|
||||
gender = tf.cast(tf.greater_equal(shape[:, 10] , 0.5), tf.int32)
|
||||
shape_offsets = tf.stack([tf.gather(tf.tensordot(shape[i,:10][tf.newaxis, :], smpl.shapedirs[gender[i]], axes=[[1], [2]])[0,:,:], smpl.to_keep, axis=0) for i in range(shape.shape[0])])
|
||||
|
||||
mesh_pred_shaped = mesh_pred + shape_offsets * batch['mesh_mask'][:,:,tf.newaxis]
|
||||
mesh_pred_posed = smpl.run_garment(batch['pose_shape'][:, :72], batch['pose_shape'][:, 72:], mesh_pred_shaped)
|
||||
mesh_orig_shaped = mesh_orig + shape_offsets * batch['mesh_mask'][:,:,tf.newaxis]
|
||||
mesh_orig_posed = smpl.run_garment(batch['pose_shape'][:, :72], batch['pose_shape'][:, 72:], mesh_orig_shaped)
|
||||
else:
|
||||
mesh_pred_posed = mesh_pred
|
||||
mesh_orig_posed = mesh_orig
|
||||
|
||||
mesh_pred_posed_fixed = fix_collision(mesh_pred_posed, batch['mesh_body_posed'], batch['mesh_mask'])
|
||||
|
||||
|
||||
# is the output already masked?
|
||||
|
||||
e_loss = tf.reduce_mean(edge_loss(mesh_pred, mesh_orig, batch['mesh_mask']))
|
||||
c_loss, c_ratio = collision_penalty(mesh_pred_posed, batch['mesh_body_posed'], batch['mesh_mask'])
|
||||
c_loss, c_ratio = tf.reduce_mean(c_loss), tf.reduce_mean(c_ratio)
|
||||
l2_loss = tf.reduce_mean(1e3 * tf.reduce_mean(tf.reduce_sum(((uv_orig - pred) ** 2), axis=-1), axis=[1, 2]))
|
||||
l2_mesh_loss = tf.reduce_mean(1e3 * tf.reduce_mean(tf.reduce_sum(batch['mesh_mask'][:,:,tf.newaxis] * (mesh_orig - mesh_pred) ** 2, axis=-1), axis=-1))
|
||||
# conc_loss = concept_loss(pred, batch['uv_unposed'])
|
||||
kl = kl_loss(mean, logvar)
|
||||
|
||||
reconstr_loss = l2_loss
|
||||
if config['use_extra_losses']:
|
||||
reconstr_loss = reconstr_loss + l2_mesh_loss + (e_loss + c_loss) * config['reg_loss_coeff']
|
||||
|
||||
loss = reconstr_loss + g_balancing_coeff * kl
|
||||
|
||||
mean = tf.reduce_mean(mean)
|
||||
logvar = tf.reduce_mean(logvar)
|
||||
|
||||
ms_ssim = tf.reduce_mean(tf.image.ssim_multiscale(pred + 5, uv_orig + 5, max_val=10, power_factors=(0.5, 0.2, 0.1, 0.1, 0.1)))
|
||||
metric = tf.reduce_sum(tf.math.sqrt(tf.reduce_sum(batch['mesh_mask'][:,:,tf.newaxis] * (mesh_orig - mesh_pred) ** 2, axis=-1))) / tf.reduce_sum(batch['mesh_mask'])
|
||||
euclid_posed = tf.reduce_sum(tf.math.sqrt(tf.reduce_sum(batch['mesh_mask'][:,:,tf.newaxis] * (mesh_orig_posed - mesh_pred_posed) ** 2, axis=-1))) / tf.reduce_sum(batch['mesh_mask'])
|
||||
euclid_posed_fixed = tf.reduce_sum(tf.math.sqrt(tf.reduce_sum(batch['mesh_mask'][:,:,tf.newaxis] * (mesh_orig_posed - mesh_pred_posed_fixed) ** 2, axis=-1))) / tf.reduce_sum(batch['mesh_mask'])
|
||||
|
||||
return { 'test/loss': loss, 'test/edge_loss': e_loss, 'test/collision_loss': c_loss, 'test/euclid_error': metric, 'test/ms_ssim': ms_ssim,
|
||||
'test/reconstr_loss': reconstr_loss, 'test/collision_ratio': c_ratio, 'test/l2_mesh_loss': l2_mesh_loss,
|
||||
'test/mean': mean, 'test/logvar': logvar, 'test/kl_loss': kl, 'test/var': np.e ** logvar,
|
||||
'test/euclid_posed': euclid_posed, 'test/euclid_posed_fixed': euclid_posed_fixed }
|
||||
|
||||
|
||||
def plot_everything(batch: dict, type: str):
|
||||
if config['use_skinning']:
|
||||
mesh_orig, uv_orig = batch['mesh_unposed'], batch['uv_unposed']
|
||||
else:
|
||||
mesh_orig, uv_orig = batch['mesh_posed'], batch['uv_posed']
|
||||
|
||||
pred, mean, logvar = model([uv_orig, batch['uv_static'], batch['uv_body_posed'], batch['uv_curvature'], batch['uv_mask']], training=False)
|
||||
|
||||
mesh_pred = uv2mesh(pred).numpy()
|
||||
|
||||
orig_img = uv_orig
|
||||
uv_static_img = batch['uv_static']
|
||||
uv_body_posed_img = batch['uv_body_posed']
|
||||
pred_img = pred
|
||||
img = tf.concat((uv_static_img, uv_body_posed_img, orig_img, pred_img), axis=2)
|
||||
imgs = [wandb.Image((img[i] + 1) / 2.0) for i in range(4)]
|
||||
|
||||
F = np.load("custom_utils/smpl_faces" + skirt_modifier + ("_orig" if config['is_cape'] else "") + ".npy")
|
||||
mask = batch['mesh_mask'].numpy().astype(bool)
|
||||
|
||||
if config['use_skinning']:
|
||||
shape = batch['pose_shape'][:, 72:]
|
||||
gender = tf.cast(tf.greater_equal(shape[:, 10] , 0.5), tf.int32)
|
||||
shape_offsets = tf.stack([tf.gather(tf.tensordot(shape[i,:10][tf.newaxis, :], smpl.shapedirs[gender[i]], axes=[[1], [2]])[0,:,:], smpl.to_keep, axis=0) for i in range(shape.shape[0])])
|
||||
|
||||
garment_unposed, faces = maskFacesBatch(mesh_pred, F, mask)
|
||||
garment_unposed = [garment_unposed[i] + shape_offsets.numpy()[i, mask[i], :] for i in range(mask.shape[0])]
|
||||
gt = [mesh_orig[i].numpy()[mask[i], :] + shape_offsets.numpy()[i, mask[i], :] for i in range(pred.shape[0])]
|
||||
|
||||
for i in range(mesh_orig.shape[0]):
|
||||
mesh_pred[i, mask[i], :] += shape_offsets.numpy()[i, mask[i], :]
|
||||
mesh_posed = smpl.run_garment(batch['pose_shape'][:, :72], batch['pose_shape'][:, 72:], mesh_pred).numpy()
|
||||
else:
|
||||
mesh_posed = mesh_pred
|
||||
gt = [mesh_orig[i].numpy()[mask[i], :] for i in range(pred.shape[0])]
|
||||
garment_unposed, faces = maskFacesBatch(mesh_pred, F, mask)
|
||||
|
||||
gt_posed = [batch['mesh_posed'].numpy()[i, mask[i], :] for i in range(mask.shape[0])]
|
||||
|
||||
|
||||
body = batch['mesh_body_posed'].numpy()
|
||||
mesh_posed_fixed = fix_collision(mesh_posed, body, mask).numpy()
|
||||
|
||||
garment_posed = [mesh_posed[i, mask[i], :] for i in range(body.shape[0])]
|
||||
collision = [np.concatenate((body[i], garment_posed[i]), axis=0) for i in range(len(garment_posed))]
|
||||
collision_faces = [np.concatenate((F, faces[i] + config['smpl_res']), axis=0) for i in range(len(garment_posed))]
|
||||
collision_color = [np.concatenate((np.ones_like(body[i], dtype=np.int32)*100, np.ones_like(garment_posed[i], dtype=np.int32)*230), axis=0) for i in range(len(garment_posed))]
|
||||
|
||||
garment_posed_fixed = [mesh_posed_fixed[i, mask[i], :] for i in range(body.shape[0])]
|
||||
collision_fixed = [np.concatenate((body[i], garment_posed_fixed[i]), axis=0) for i in range(len(garment_posed))]
|
||||
collision_gt = [np.concatenate((body[i], gt_posed[i]), axis=0) for i in range(len(garment_posed))]
|
||||
|
||||
figs = { type + "/garment_" + str(i): plot_3d(gt[i], garment_unposed[i], faces[i], collision[i], collision_faces[i], collision_color[i], flip_xy=(not config['use_skinning'] and config['remove_global_orientation'])) for i in range(4) }
|
||||
figs_coll = { type + "/garment_" + str(i) + '_collision': plot_3d_collision(collision_gt[i], collision[i], collision_fixed[i], collision_faces[i], collision_color[i], flip_xy=(not config['use_skinning'] and config['remove_global_orientation'])) for i in range(4) }
|
||||
|
||||
mean_hist = wandb.Histogram(sequence=mean.numpy().mean(axis=0))
|
||||
logvar_hist = wandb.Histogram(sequence=logvar.numpy().mean(axis=0))
|
||||
|
||||
return { **figs, **figs_coll, **{ type + '/images': imgs, type + '/mean_hist': mean_hist, type + '/logvar_hist': logvar_hist } }
|
||||
|
||||
|
||||
def sample(folder):
|
||||
# run = wandb.init(project='vae_plots', name=folder.split('/')[-1])
|
||||
load_model(model,folder)
|
||||
count = 0
|
||||
for batch in dataset_train:
|
||||
batch = static_dataset_train.expand_batch(batch_to_dict(batch, config['signature_dynamic']))
|
||||
uv_curvature = tf.one_hot(batch['uv_curvature'], 7)
|
||||
nr_of_samples=2
|
||||
|
||||
pred = model.sample_single(batch['uv_static'][0], batch['uv_body_posed'][0], uv_curvature[0], batch['uv_mask'][0], nr_of_samples=nr_of_samples)
|
||||
pose, shape = batch['pose_shape'][0, :72], batch['pose_shape'][0, 72:]
|
||||
pose, shape = tf.tile(pose[tf.newaxis,:], [nr_of_samples,1]), tf.tile(shape[tf.newaxis,:], [nr_of_samples,1])
|
||||
|
||||
mesh_pred = uv2mesh(pred)
|
||||
mesh_posed = smpl.run_garment(pose, shape, mesh_pred).numpy()
|
||||
mesh_orig = batch['mesh_posed'][0].numpy()
|
||||
mesh_orig_reposed = smpl.run_garment(pose[:1], shape[:1], batch['mesh_unposed'][:1]).numpy()[0]
|
||||
|
||||
body = batch['mesh_body_posed'][0].numpy()
|
||||
body_color = np.ones_like(body, dtype=np.int32)*100
|
||||
|
||||
F = np.load("custom_utils/smpl_faces.npy")
|
||||
mask = batch['mesh_mask'].numpy().astype(bool)[0]
|
||||
_, faces = _maskFaces(np.ones((config['smpl_res'], 3)), F, mask)
|
||||
|
||||
pred_zero = model.sample_single(batch['uv_static'][0], batch['uv_body_posed'][0], uv_curvature[0], batch['uv_mask'][0], nr_of_samples=nr_of_samples, zero_latent=True)
|
||||
mesh_pred_zero = uv2mesh(pred_zero)
|
||||
mesh_posed_zero = smpl.run_garment(pose, shape, mesh_pred_zero).numpy()
|
||||
|
||||
|
||||
meshes = []
|
||||
meshes.append(go.Mesh3d(
|
||||
# 8 vertices of a cube
|
||||
x=mesh_orig_reposed[mask,0],
|
||||
y=mesh_orig_reposed[mask,2],
|
||||
z=mesh_orig_reposed[mask,1],
|
||||
# i, j and k give the vertices of triangles
|
||||
i = faces[:,0],
|
||||
j = faces[:,1],
|
||||
k = faces[:,2]
|
||||
))
|
||||
|
||||
for i in range(nr_of_samples):
|
||||
x = mesh_posed[i][mask, :].copy()
|
||||
# x[:, 2] += 1 * (i + 1)
|
||||
x[:, 2] += 1
|
||||
meshes.append(go.Mesh3d(
|
||||
# 8 vertices of a cube
|
||||
x=x[:,0],
|
||||
y=x[:,2],
|
||||
z=x[:,1],
|
||||
# i, j and k give the vertices of triangles
|
||||
i = faces[:,0],
|
||||
j = faces[:,1],
|
||||
k = faces[:,2]
|
||||
))
|
||||
|
||||
|
||||
for i in range(nr_of_samples):
|
||||
x = mesh_posed_zero[i][mask, :].copy()
|
||||
# x[:, 2] += 1 * (i + 1)
|
||||
x[:, 2] += 2
|
||||
meshes.append(go.Mesh3d(
|
||||
# 8 vertices of a cube
|
||||
x=x[:,0],
|
||||
y=x[:,2],
|
||||
z=x[:,1],
|
||||
# i, j and k give the vertices of triangles
|
||||
i = faces[:,0],
|
||||
j = faces[:,1],
|
||||
k = faces[:,2]
|
||||
))
|
||||
|
||||
# for i in range(5):
|
||||
# x = body.copy()
|
||||
# x[:, 2] += 1 * i
|
||||
# meshes.append(go.Mesh3d(
|
||||
# # 8 vertices of a cube
|
||||
# x=x[:,0],
|
||||
# y=x[:,2],
|
||||
# z=x[:,1],
|
||||
# # i, j and k give the vertices of triangles
|
||||
# i = F[:,0],
|
||||
# j = F[:,1],
|
||||
# k = F[:,2],
|
||||
# vertexcolor=body_color
|
||||
# ))
|
||||
|
||||
|
||||
fig = go.Figure(data=meshes)
|
||||
|
||||
fig.update_layout(scene_aspectmode='manual', scene_aspectratio=dict(x=1, y=2, z=1))
|
||||
# fig.update_layout(scene_aspectmode='manual', scene_aspectratio=dict(x=1, y=nr_of_samples-1, z=1))
|
||||
fig.update_layout(
|
||||
scene = dict(
|
||||
xaxis = dict(nticks=4, range=[-1,1],),
|
||||
# yaxis = dict(nticks=4, range=[-1,nr_of_samples+1],),
|
||||
yaxis = dict(nticks=4, range=[-1,3],),
|
||||
zaxis = dict(nticks=4, range=[-1,1],),)
|
||||
)
|
||||
|
||||
# wandb.log({'plot_' + str(count): fig})
|
||||
fig.show()
|
||||
|
||||
pred = pred.numpy()
|
||||
frames = [Image.fromarray(((pred[i] + 1) * 128).astype(np.int8), mode='RGB') for i in range(pred.shape[0])]
|
||||
frames[0].save("temp/garment_1.gif", format="GIF", append_images=frames[1:],
|
||||
save_all=True, duration=500, loop=0)
|
||||
|
||||
|
||||
if count > 10:
|
||||
break
|
||||
|
||||
print(count)
|
||||
count += 1
|
||||
|
||||
print()
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
def test():
|
||||
str_split = config['name'].split('_')
|
||||
test_idx = str_split.index('test')
|
||||
checkpoint_name = '_'.join(str_split[:test_idx])
|
||||
load_model(model, 'checkpoints/' + checkpoint_name)
|
||||
|
||||
dataset_test = factory.get_dataset(folder=config['test_dir'], batch_size=config['batch_size'], signature=config['signature_dynamic'], use_filled=config['use_filled'], resolutions=[config['resolution']], in_memory=False, do_repeat=False)
|
||||
static_dataset_test = StaticDataset(folder=config['test_dir'], signature=config['signature_static'], use_filled=config['use_filled'], resolutions=[config['resolution']], in_memory=False)
|
||||
|
||||
step = 0
|
||||
losses = []
|
||||
|
||||
for batch in tqdm.tqdm(dataset_test):
|
||||
batch = static_dataset_test.expand_batch(batch_to_dict(batch, config['signature_dynamic']))
|
||||
|
||||
test_loss = test_step(batch)
|
||||
losses.append(test_loss)
|
||||
|
||||
to_log = { **{ 'current_step': step }, **test_loss }
|
||||
if step % 1000 == 0:
|
||||
figs = plot_everything(batch, 'test')
|
||||
to_log = { **to_log, **figs }
|
||||
|
||||
wandb.log(to_log, step = step)
|
||||
|
||||
# print(step)
|
||||
step += 1
|
||||
|
||||
# averages = dict()
|
||||
path_to_plotly_html = "temp/plotly_figure.html"
|
||||
|
||||
histograms = dict()
|
||||
for key in test_loss.keys():
|
||||
accumulated = [losses[i][key].numpy() for i in range(len(losses))]
|
||||
fig = ex.histogram(accumulated, title='Average: ' + str(np.asarray(accumulated).mean()))
|
||||
fig.write_html(path_to_plotly_html, auto_play = False)
|
||||
|
||||
histograms['all/' + key.split('/')[-1] + '_hist'] = wandb.Html(path_to_plotly_html)
|
||||
# averages['all/' + key.split('/')[-1] + '_avg'] = np.asarray(accumulated).mean()
|
||||
|
||||
|
||||
|
||||
wandb.log({ **histograms }, step = step)
|
||||
wandb.finish()
|
||||
pass
|
||||
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
# init_wandb()
|
||||
parser = argparse.ArgumentParser(description='')
|
||||
parser.add_argument('--name', type=str)
|
||||
args = vars(parser.parse_args())
|
||||
|
||||
smpl = SMPL(orig_smpl=config['is_cape'])
|
||||
skirt_modifier = '_skirt' if config['is_skirt'] else ''
|
||||
|
||||
factory = DatasetFactory()
|
||||
|
||||
dataset_train = factory.get_dataset(folder=config['train_dir'],
|
||||
batch_size=config['batch_size'],
|
||||
signature=config['signature_dynamic'],
|
||||
resolutions=[config['resolution']],
|
||||
in_memory=config['dataset_in_memory'],
|
||||
use_one_hot=config['use_one_hot'],
|
||||
use_no_orient=config['remove_global_orientation'],
|
||||
frames_file=config['frames_file'] if 'frames_file' in config.keys() else None,
|
||||
is_cape=config['is_cape'])
|
||||
static_dataset_train = StaticDataset(folder=config['train_dir'],
|
||||
signature=config['signature_static'],
|
||||
resolutions=[config['resolution']],
|
||||
in_memory=True,
|
||||
is_cape=config['is_cape'])
|
||||
|
||||
dataset_valid = factory.get_dataset(folder=config['valid_dir'],
|
||||
batch_size=config['batch_size'],
|
||||
signature=config['signature_dynamic'],
|
||||
resolutions=[config['resolution']],
|
||||
in_memory=config['dataset_in_memory'],
|
||||
use_one_hot=config['use_one_hot'],
|
||||
use_no_orient=config['remove_global_orientation'],
|
||||
frames_file=config['frames_file_valid'] if 'frames_file_valid' in config.keys() else None,
|
||||
is_cape=config['is_cape'])
|
||||
static_dataset_valid = StaticDataset(folder=config['valid_dir'],
|
||||
signature=config['signature_static'],
|
||||
resolutions=[config['resolution']],
|
||||
in_memory=True,
|
||||
is_cape=config['is_cape'])
|
||||
|
||||
model = CVAE(config)
|
||||
|
||||
optimizer = tf.optimizers.Adam(learning_rate=config['lr'])
|
||||
balancing_activated_step = 0
|
||||
g_balancing_coeff = config['starting_balancing_coeff']
|
||||
|
||||
# sample('checkpoints/cvae_13')
|
||||
# test()
|
||||
|
||||
ckpt = tf.train.Checkpoint(model=model)
|
||||
manager = tf.train.CheckpointManager(ckpt, 'checkpoints/post/' + config['name'], max_to_keep=20)
|
||||
|
||||
|
||||
if config['pretrained_model'] != '':
|
||||
load_model(model, config['pretrained_model'])
|
||||
|
||||
# load_model(model, 'checkpoints/' + config['name'])
|
||||
|
||||
valid_iterator = iter(dataset_valid)
|
||||
step = 0
|
||||
for batch in dataset_train:
|
||||
batch = static_dataset_train.expand_batch(batch_to_dict(batch, config['signature_dynamic']))
|
||||
|
||||
train_loss = train_step(batch)
|
||||
|
||||
data_valid = valid_iterator.get_next()
|
||||
data_valid = static_dataset_valid.expand_batch(batch_to_dict(data_valid, config['signature_dynamic']))
|
||||
valid_loss = valid_step(data_valid)
|
||||
|
||||
to_log = { **train_loss, **valid_loss, **{ 'current_step': step} }
|
||||
if step % 5000 == 0 or (step < 5000 and step % 500 == 0):
|
||||
train_figs = plot_everything(batch, 'train')
|
||||
valid_figs = plot_everything(data_valid, 'valid')
|
||||
to_log = { **to_log, **train_figs, **valid_figs }
|
||||
|
||||
if step % 5000 == 0:
|
||||
manager.save(checkpoint_number=step)
|
||||
if step == 0:
|
||||
to_log = { **to_log, **{ 'model/summary': wandb.Html(get_summary_html(model)) } }
|
||||
|
||||
if train_loss['train/euclid_error'] < (config['balancing_activation_limit'] if 'balancing_activation_limit' in config.keys() else 0.017):
|
||||
if balancing_activated_step == 0:
|
||||
balancing_activated_step = step
|
||||
g_balancing_coeff = g_balancing_coeff + 1 / 10000
|
||||
|
||||
to_log = { **to_log, **{ 'train/balancing_coeff': g_balancing_coeff } }
|
||||
|
||||
wandb.log(to_log, step = step)
|
||||
|
||||
print(step)
|
||||
|
||||
step += 1
|
||||
|
||||
@@ -0,0 +1,881 @@
|
||||
import os
|
||||
import sys
|
||||
from configs.config_abl_pyr_32_skirt import config
|
||||
|
||||
os.environ["CUDA_VISIBLE_DEVICES"] = config['gpu']
|
||||
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)
|
||||
|
||||
import time
|
||||
import tqdm
|
||||
import argparse
|
||||
import wandb
|
||||
import random
|
||||
import numpy as np
|
||||
import scipy as sp
|
||||
from custom_utils.custom_utils import UV2Mesh, get_summary_html, maskFacesBatch, _maskFaces, get_summary_str, plot_3d_collision, group_inputs_for_levels
|
||||
import plotly.graph_objects as go
|
||||
import plotly.express as ex
|
||||
from plotly.subplots import make_subplots
|
||||
from PIL import Image
|
||||
|
||||
if config['is_skirt']:
|
||||
from losses import edge_loss_skirt as edge_loss
|
||||
else:
|
||||
from losses import edge_loss
|
||||
|
||||
from losses import uv_mask_loss, bend_loss, kl_loss, collision_ratio, normal_loss
|
||||
# from losses import collision_penalty as collision_penalty
|
||||
from losses import collision_loss as collision_penalty
|
||||
from custom_utils.custom_utils import batch_to_dict, load_model
|
||||
from dataset import DatasetFactory, StaticDataset
|
||||
from custom_utils.smpl import SMPL
|
||||
from models.pyramid_model_optimized import Pyramid
|
||||
from postprocessing import fix_collision
|
||||
|
||||
|
||||
def init_wandb():
|
||||
run = wandb.init(project='ablation', name=config['name'], notes=config['notes'], dir='/wandb', tags=config['tags'])
|
||||
|
||||
artifact = wandb.Artifact('code', type='code')
|
||||
artifact.add_file('cvae_pyramid.py')
|
||||
artifact.add_file('models/cvae_model.py')
|
||||
artifact.add_file('models/pyramid_model_optimized.py')
|
||||
artifact.add_file('models/pyramid_level.py')
|
||||
artifact.add_file('models/upscaler.py')
|
||||
artifact.add_file('losses.py')
|
||||
artifact.add_file('models/base.py')
|
||||
artifact.add_file('models/cond_encoder.py')
|
||||
artifact.add_file('dataset.py')
|
||||
artifact.add_file('convnext/models/convnext_tf.py')
|
||||
config_file_names = [key for key in sys.modules.keys() if key.startswith('configs')]
|
||||
for config_file_name in config_file_names:
|
||||
file_name = '/'.join(config_file_name.split('.')) + '.py'
|
||||
if os.path.exists(file_name):
|
||||
artifact.add_file(file_name)
|
||||
run.log_artifact(artifact)
|
||||
|
||||
uv2mesh = UV2Mesh(res=256, is_skirt=config['is_skirt'], is_smpl_orig=config['is_cape'])
|
||||
|
||||
def plot_3d(mesh_gt, mesh_pred_unposed_levels, faces, mesh_posed, collision, collision_faces, collision_color):
|
||||
fig = make_subplots(rows=1, cols=1+len(mesh_pred_unposed_levels)+1+1, start_cell="top-left", specs = [ [{ 'type': 'scene' } for _ in range(1+len(mesh_pred_unposed_levels)+1+1)] ])
|
||||
|
||||
fig.add_trace(go.Mesh3d(
|
||||
x=mesh_gt[:,0],
|
||||
y=mesh_gt[:,2],
|
||||
z=mesh_gt[:,1],
|
||||
i = faces[:,0],
|
||||
j = faces[:,1],
|
||||
k = faces[:,2],
|
||||
text="Ground Truth"
|
||||
),
|
||||
row=1, col=1)
|
||||
|
||||
for i, mesh in enumerate(mesh_pred_unposed_levels):
|
||||
fig.add_trace(go.Mesh3d(
|
||||
x=mesh[:,0],
|
||||
y=mesh[:,2],
|
||||
z=mesh[:,1],
|
||||
i = faces[:,0],
|
||||
j = faces[:,1],
|
||||
k = faces[:,2],
|
||||
text='Unposed level_' + str(i)
|
||||
),
|
||||
row=1, col=i+2)
|
||||
|
||||
fig.add_trace(go.Mesh3d(
|
||||
x=mesh_posed[:,0],
|
||||
y=mesh_posed[:,1],
|
||||
z=mesh_posed[:,2],
|
||||
i = faces[:,0],
|
||||
j = faces[:,1],
|
||||
k = faces[:,2],
|
||||
text='Posed GT'
|
||||
),
|
||||
row=1, col=1+len(mesh_pred_unposed_levels)+1)
|
||||
|
||||
fig.add_trace(go.Mesh3d(
|
||||
x=collision[:,0],
|
||||
y=collision[:,2],
|
||||
z=collision[:,1],
|
||||
i = collision_faces[:,0],
|
||||
j = collision_faces[:,1],
|
||||
k = collision_faces[:,2],
|
||||
vertexcolor=collision_color,
|
||||
text='Posed output on body'
|
||||
),
|
||||
row=1, col=1+len(mesh_pred_unposed_levels)+1+1)
|
||||
|
||||
if len(mesh_pred_unposed_levels) == 1:
|
||||
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),
|
||||
scene4_aspectmode='manual', scene4_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,
|
||||
scene4 = helper,
|
||||
)
|
||||
if len(mesh_pred_unposed_levels) == 2:
|
||||
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),
|
||||
scene4_aspectmode='manual', scene4_aspectratio=dict(x=1, y=1, z=1),
|
||||
scene5_aspectmode='manual', scene5_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,
|
||||
scene4 = helper,
|
||||
scene5 = helper,
|
||||
)
|
||||
if len(mesh_pred_unposed_levels) == 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),
|
||||
scene4_aspectmode='manual', scene4_aspectratio=dict(x=1, y=1, z=1),
|
||||
scene5_aspectmode='manual', scene5_aspectratio=dict(x=1, y=1, z=1),
|
||||
scene6_aspectmode='manual', scene6_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,
|
||||
scene4 = helper,
|
||||
scene5 = helper,
|
||||
scene6 = helper,
|
||||
)
|
||||
if len(mesh_pred_unposed_levels) == 4:
|
||||
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),
|
||||
scene4_aspectmode='manual', scene4_aspectratio=dict(x=1, y=1, z=1),
|
||||
scene5_aspectmode='manual', scene5_aspectratio=dict(x=1, y=1, z=1),
|
||||
scene6_aspectmode='manual', scene6_aspectratio=dict(x=1, y=1, z=1,),
|
||||
scene7_aspectmode='manual', scene7_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,
|
||||
scene4 = helper,
|
||||
scene5 = helper,
|
||||
scene6 = helper,
|
||||
scene7 = helper,
|
||||
)
|
||||
if len(mesh_pred_unposed_levels) == 5:
|
||||
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),
|
||||
scene4_aspectmode='manual', scene4_aspectratio=dict(x=1, y=1, z=1),
|
||||
scene5_aspectmode='manual', scene5_aspectratio=dict(x=1, y=1, z=1),
|
||||
scene6_aspectmode='manual', scene6_aspectratio=dict(x=1, y=1, z=1,),
|
||||
scene7_aspectmode='manual', scene7_aspectratio=dict(x=1, y=1, z=1,),
|
||||
scene8_aspectmode='manual', scene8_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,
|
||||
scene4 = helper,
|
||||
scene5 = helper,
|
||||
scene6 = helper,
|
||||
scene7 = helper,
|
||||
scene8 = helper,
|
||||
)
|
||||
return fig
|
||||
|
||||
|
||||
# @tf.function
|
||||
def train_step(batch):
|
||||
levels_input, conds = group_inputs_for_levels(batch, config)
|
||||
|
||||
with tf.GradientTape() as tape:
|
||||
out_levels = model([levels_input, conds], training=True)
|
||||
|
||||
if config['use_skinning']:
|
||||
mesh_orig, uv_orig = batch['mesh_unposed'], batch['uv_unposed']
|
||||
else:
|
||||
mesh_orig, uv_orig = batch['mesh_posed'], batch['uv_posed']
|
||||
|
||||
output = dict()
|
||||
reconstr_loss_total = 0
|
||||
kl_total = 0
|
||||
|
||||
#for i, res in enumerate(config['pyramid']['resolutions']):
|
||||
i = -1
|
||||
res = config['pyramid']['resolutions'][-1]
|
||||
res_modifier = '_' + str(res) if res != 512 else ''
|
||||
pred, mean, logvar = out_levels[i]
|
||||
|
||||
mesh_pred = uv2mesh(pred)
|
||||
|
||||
if config['use_skinning']:
|
||||
shape = batch['pose_shape'][:, 72:]
|
||||
gender = tf.cast(tf.greater_equal(shape[:, 10] , 0.5), tf.int32)
|
||||
shape_offsets = tf.stack([tf.gather(tf.tensordot(shape[i,:10][tf.newaxis, :], smpl.shapedirs[gender[i]], axes=[[1], [2]])[0,:,:], smpl.to_keep, axis=0) for i in range(shape.shape[0])])
|
||||
|
||||
mesh_pred_shaped = mesh_pred + shape_offsets * batch['mesh_mask'][:,:,tf.newaxis]
|
||||
mesh_pred_posed = smpl.run_garment(batch['pose_shape'][:, :72], batch['pose_shape'][:, 72:], mesh_pred_shaped)
|
||||
else:
|
||||
mesh_pred_posed = mesh_pred
|
||||
|
||||
# is the output already masked?
|
||||
|
||||
e_loss = tf.reduce_mean(edge_loss(mesh_pred, mesh_orig, batch['mesh_mask']))
|
||||
c_loss, c_ratio = collision_penalty(mesh_pred_posed, batch['mesh_body_posed'], batch['mesh_mask'])
|
||||
c_loss, c_ratio = tf.reduce_mean(c_loss), tf.reduce_mean(c_ratio)
|
||||
l2_loss = tf.reduce_mean(1e3 * tf.reduce_mean(tf.reduce_sum(((uv_orig - pred) ** 2), axis=-1), axis=[1, 2]))
|
||||
# ms_ssim = tf.reduce_mean(tf.image.ssim_multiscale(pred + 5, batch['uv_unposed'] + 5, max_val=10, power_factors=(0.5, 0.2, 0.1, 0.1, 0.1)))
|
||||
# ms_ssim = tf.reduce_mean(tf.image.ssim(pred + 1, batch['uv_unposed'] + 1, max_val=3))
|
||||
l2_mesh_loss = tf.reduce_mean(1e3 * tf.reduce_mean(tf.reduce_sum(batch['mesh_mask'][:,:,tf.newaxis] * (batch['mesh_posed'] - mesh_pred_posed) ** 2, axis=-1), axis=-1))
|
||||
n_loss = normal_loss(mesh_pred_posed, batch['mesh_posed'], batch['mesh_mask'], is_skirt=config['is_skirt'])
|
||||
kl = kl_loss(mean, logvar)
|
||||
|
||||
if config['use_extra_losses']:
|
||||
reconstr_loss_total = reconstr_loss_total + config['pyramid']['loss_weights'][i] * (l2_loss + l2_mesh_loss)
|
||||
if not config['is_cape']:
|
||||
reconstr_loss_total = reconstr_loss_total + config['pyramid']['loss_weights'][i] * (e_loss + c_loss) * config['reg_loss_coeff']
|
||||
else:
|
||||
reconstr_loss_total = reconstr_loss_total + config['pyramid']['loss_weights'][i] * (l2_loss)
|
||||
|
||||
# if config['use_norm_loss'] and step > 2000:
|
||||
# reconstr_loss_total = reconstr_loss_total + n_loss * config['reg_loss_coeff']
|
||||
|
||||
kl_total += kl
|
||||
|
||||
mean = tf.reduce_mean(mean)
|
||||
logvar = tf.reduce_mean(logvar)
|
||||
|
||||
metric = tf.reduce_sum(tf.math.sqrt(tf.reduce_sum(batch['mesh_mask'][:,:,tf.newaxis] * (mesh_orig - mesh_pred) ** 2, axis=-1))) / tf.reduce_sum(batch['mesh_mask'])
|
||||
|
||||
output = { **output, **{ 'train/edge_loss' + res_modifier: e_loss, 'train/collision_loss' + res_modifier: c_loss,'train/euclid_error' + res_modifier: metric,
|
||||
'train/mean' + res_modifier: mean, 'train/logvar' + res_modifier: logvar, 'train/kl_loss' + res_modifier: kl,
|
||||
'train/var' + res_modifier: np.e ** logvar, 'train/l2_loss' + res_modifier: l2_loss, 'train/normal_loss': n_loss,
|
||||
'train/euclid_error': metric, 'train/l2_mesh_loss': l2_mesh_loss } } # 'train/ms_ssim' + res_modifier: ms_ssim,
|
||||
|
||||
loss = reconstr_loss_total + balancing_coeff * kl_total
|
||||
|
||||
gradients = tape.gradient(loss, model.trainable_variables)
|
||||
optimizer.apply_gradients(zip(gradients, model.trainable_variables))
|
||||
return { **{ 'train/loss': loss, 'train/reconstr_loss_total': reconstr_loss_total, 'train/kl_total': kl_total }, **output }
|
||||
|
||||
|
||||
@tf.function
|
||||
def valid_step(batch):
|
||||
levels_input, conds = group_inputs_for_levels(batch, config)
|
||||
out_levels = model([levels_input, conds], training=False)
|
||||
|
||||
|
||||
if config['use_skinning']:
|
||||
mesh_orig, uv_orig = batch['mesh_unposed'], batch['uv_unposed']
|
||||
else:
|
||||
mesh_orig, uv_orig = batch['mesh_posed'], batch['uv_posed']
|
||||
|
||||
output = dict()
|
||||
reconstr_loss_total = 0
|
||||
kl_total = 0
|
||||
|
||||
#for i, res in enumerate(config['pyramid']['resolutions']):
|
||||
i = -1
|
||||
res = config['pyramid']['resolutions'][-1]
|
||||
res_modifier = '_' + str(res) if res != 512 else ''
|
||||
pred, mean, logvar = out_levels[i]
|
||||
|
||||
mesh_pred = uv2mesh(pred)
|
||||
|
||||
if config['use_skinning']:
|
||||
shape = batch['pose_shape'][:, 72:]
|
||||
gender = tf.cast(tf.greater_equal(shape[:, 10] , 0.5), tf.int32)
|
||||
shape_offsets = tf.stack([tf.gather(tf.tensordot(shape[i,:10][tf.newaxis, :], smpl.shapedirs[gender[i]], axes=[[1], [2]])[0,:,:], smpl.to_keep, axis=0) for i in range(shape.shape[0])])
|
||||
|
||||
mesh_pred_shaped = mesh_pred + shape_offsets * batch['mesh_mask'][:,:,tf.newaxis]
|
||||
mesh_pred_posed = smpl.run_garment(batch['pose_shape'][:, :72], batch['pose_shape'][:, 72:], mesh_pred_shaped)
|
||||
else:
|
||||
mesh_pred_posed = mesh_pred
|
||||
# is the output already masked?
|
||||
|
||||
e_loss = tf.reduce_mean(edge_loss(mesh_pred, mesh_orig, batch['mesh_mask']))
|
||||
c_loss, c_ratio = collision_penalty(mesh_pred_posed, batch['mesh_body_posed'], batch['mesh_mask'])
|
||||
c_loss, c_ratio = tf.reduce_mean(c_loss), tf.reduce_mean(c_ratio)
|
||||
l2_loss = tf.reduce_mean(1e3 * tf.reduce_mean(tf.reduce_sum(((uv_orig - pred) ** 2), axis=-1), axis=[1, 2]))
|
||||
# ms_ssim = tf.reduce_mean(tf.image.ssim_multiscale(pred + 5, batch['uv_unposed'] + 5, max_val=10, power_factors=(0.5, 0.2, 0.1, 0.1, 0.1)))
|
||||
# ms_ssim = tf.reduce_mean(tf.image.ssim(pred + 1, batch['uv_unposed'] + 1, max_val=3))
|
||||
l2_mesh_loss = tf.reduce_mean(1e3 * tf.reduce_mean(tf.reduce_sum(batch['mesh_mask'][:,:,tf.newaxis] * (batch['mesh_posed'] - mesh_pred_posed) ** 2, axis=-1), axis=-1))
|
||||
n_loss = normal_loss(mesh_pred_posed, batch['mesh_posed'], batch['mesh_mask'], is_skirt=config['is_skirt'])
|
||||
kl = kl_loss(mean, logvar)
|
||||
|
||||
if config['use_extra_losses']:
|
||||
reconstr_loss_total = reconstr_loss_total + config['pyramid']['loss_weights'][i] * (l2_loss + l2_mesh_loss)
|
||||
if not config['is_cape']:
|
||||
reconstr_loss_total = reconstr_loss_total + config['pyramid']['loss_weights'][i] * (e_loss + c_loss) * config['reg_loss_coeff']
|
||||
else:
|
||||
reconstr_loss_total = reconstr_loss_total + config['pyramid']['loss_weights'][i] * (l2_loss)
|
||||
|
||||
# if config['use_norm_loss'] and step > 2000:
|
||||
# reconstr_loss_total = reconstr_loss_total + n_loss * config['reg_loss_coeff']
|
||||
|
||||
kl_total += kl
|
||||
|
||||
mean = tf.reduce_mean(mean)
|
||||
logvar = tf.reduce_mean(logvar)
|
||||
|
||||
metric = tf.reduce_sum(tf.math.sqrt(tf.reduce_sum(batch['mesh_mask'][:,:,tf.newaxis] * (mesh_orig - mesh_pred) ** 2, axis=-1))) / tf.reduce_sum(batch['mesh_mask'])
|
||||
|
||||
output = { **output, **{ 'valid/edge_loss' + res_modifier: e_loss, 'valid/collision_loss' + res_modifier: c_loss,'valid/euclid_error' + res_modifier: metric,
|
||||
'valid/mean' + res_modifier: mean, 'valid/logvar' + res_modifier: logvar, 'valid/kl_loss' + res_modifier: kl,
|
||||
'valid/var' + res_modifier: np.e ** logvar, 'valid/l2_loss' + res_modifier: l2_loss, 'valid/normal_loss': n_loss,
|
||||
'valid/euclid_error': metric, 'valid/l2_mesh_loss': l2_mesh_loss } } # 'valid/ms_ssim' + res_modifier: ms_ssim,
|
||||
|
||||
loss = reconstr_loss_total + balancing_coeff * kl_total
|
||||
|
||||
return { **{ 'valid/loss': loss, 'valid/reconstr_loss_total': reconstr_loss_total, 'valid/kl_total': kl_total }, **output }
|
||||
|
||||
|
||||
@tf.function
|
||||
def test_step(batch):
|
||||
levels_input, conds = group_inputs_for_levels(batch, config)
|
||||
out_levels = model([levels_input, conds], training=False)
|
||||
|
||||
shape = batch['pose_shape'][:, 72:]
|
||||
gender = tf.cast(tf.greater_equal(shape[:, 10] , 0.5), tf.int32)
|
||||
shape_offsets = tf.stack([tf.gather(tf.tensordot(shape[i,:10][tf.newaxis, :], smpl.shapedirs[gender[i]], axes=[[1], [2]])[0,:,:], smpl.to_keep, axis=0) for i in range(shape.shape[0])])
|
||||
mesh_orig = batch['mesh_unposed']
|
||||
|
||||
output = dict()
|
||||
reconstr_loss_total = 0
|
||||
kl_total = 0
|
||||
|
||||
#for i, res in enumerate(config['pyramid']['resolutions']):
|
||||
i = -1
|
||||
res = config['pyramid']['resolutions'][-1]
|
||||
res_modifier = '_' + str(res) if res != 512 else ''
|
||||
pred, mean, logvar = out_levels[i]
|
||||
|
||||
mesh_pred = uv2mesh(pred)
|
||||
mesh_pred_shaped = mesh_pred + shape_offsets * batch['mesh_mask'][:,:,tf.newaxis]
|
||||
mesh_pred_posed = smpl.run_garment(batch['pose_shape'][:, :72], batch['pose_shape'][:, 72:], mesh_pred_shaped)
|
||||
|
||||
# is the output already masked?
|
||||
|
||||
e_loss = tf.reduce_mean(edge_loss(mesh_pred, mesh_orig, batch['mesh_mask']))
|
||||
c_loss, c_ratio = collision_penalty(mesh_pred_posed, batch['mesh_body_posed'], batch['mesh_mask'])
|
||||
c_loss, c_ratio = tf.reduce_mean(c_loss), tf.reduce_mean(c_ratio)
|
||||
l2_loss = tf.reduce_mean(1e3 * tf.reduce_mean(tf.reduce_sum(((batch['uv_unposed'] - pred) ** 2), axis=-1), axis=[1, 2]))
|
||||
l2_mesh_loss = tf.reduce_mean(1e3 * tf.reduce_mean(tf.reduce_sum(batch['mesh_mask'][:,:,tf.newaxis] * (mesh_orig - mesh_pred) ** 2, axis=-1), axis=-1))
|
||||
# ms_ssim = tf.reduce_mean(tf.image.ssim_multiscale(pred + 5, batch['uv_unposed'] + 5, max_val=10, power_factors=(0.5, 0.2, 0.1, 0.1, 0.1)))
|
||||
# ms_ssim = tf.reduce_mean(tf.image.ssim(pred + 1, batch['uv_unposed'] + 1, max_val=3))
|
||||
kl = kl_loss(mean, logvar)
|
||||
|
||||
reconstr_loss_total = reconstr_loss_total + config['pyramid']['loss_weights'][i] * (l2_loss + l2_mesh_loss)
|
||||
if not config['is_cape']:
|
||||
reconstr_loss_total = reconstr_loss_total + config['pyramid']['loss_weights'][i] * (e_loss + c_loss) * config['reg_loss_coeff']
|
||||
# reconstr_loss_total += config['pyramid']['loss_weights'][i] * (l2_loss + (e_loss + c_loss) * config['reg_loss_coeff'])
|
||||
kl_total += kl
|
||||
|
||||
mean = tf.reduce_mean(mean)
|
||||
logvar = tf.reduce_mean(logvar)
|
||||
|
||||
metric = tf.reduce_sum(tf.math.sqrt(tf.reduce_sum(batch['mesh_mask'][:,:,tf.newaxis] * (mesh_orig - mesh_pred) ** 2, axis=-1))) / tf.reduce_sum(batch['mesh_mask'])
|
||||
euclid_error_per_garment = tf.reduce_sum(tf.math.sqrt(tf.reduce_sum(batch['mesh_mask'][:,:,tf.newaxis] * (mesh_orig - mesh_pred) ** 2, axis=-1)), axis=1) / tf.reduce_sum(batch['mesh_mask'], axis=1)
|
||||
|
||||
output = { **output, **{ 'test/edge_loss' + res_modifier: e_loss, 'test/collision_loss' + res_modifier: c_loss,'test/euclid_error' + res_modifier: metric,
|
||||
'test/mean' + res_modifier: mean, 'test/logvar' + res_modifier: logvar, 'test/kl_loss' + res_modifier: kl,
|
||||
'test/var' + res_modifier: np.e ** logvar, 'test/l2_loss' + res_modifier: l2_loss,
|
||||
'test/euclid_error': metric, 'test/collision_ratio' + res_modifier: c_ratio } }
|
||||
|
||||
loss = reconstr_loss_total + balancing_coeff * kl_total
|
||||
|
||||
return { **{ 'test/loss': loss, 'test/reconstr_loss_total': reconstr_loss_total, 'test/kl_total': kl_total }, **output }, euclid_error_per_garment
|
||||
|
||||
|
||||
|
||||
def plot_everything(batch: dict, type: str):
|
||||
if config['use_skinning']:
|
||||
mesh_orig, uv_orig = batch['mesh_unposed'].numpy(), batch['uv_unposed'].numpy()
|
||||
else:
|
||||
mesh_orig, uv_orig = batch['mesh_posed'].numpy(), batch['uv_posed'].numpy()
|
||||
|
||||
levels_input = group_inputs_for_levels(batch, config)
|
||||
output_levels = model(levels_input, training=False)
|
||||
|
||||
mesh_pred_levels = [uv2mesh(output_levels[i][0]).numpy() for i in range(len(output_levels))]
|
||||
mesh_orig_posed = batch['mesh_posed'].numpy()
|
||||
body = batch['mesh_body_posed'].numpy()
|
||||
|
||||
orig_img = (uv_orig + 1) / 2
|
||||
uv_static_img = (batch['uv_static'].numpy() + 1) / 2
|
||||
uv_body_posed_img = (batch['uv_body_posed'].numpy() + 1) / 2
|
||||
uv_curvature = (np.tile(batch['uv_curvature'].numpy().argmax(-1)[:,:,:,np.newaxis], [1,1,1,3]) / 6)
|
||||
pred_img = np.concatenate([tf.image.resize(output_levels[i][0], (512, 256), tf.image.ResizeMethod.LANCZOS5) for i in range(len(output_levels))], axis=2)
|
||||
pred_img = (pred_img - pred_img.min(axis=(1,2,3))[:,np.newaxis,np.newaxis,np.newaxis]) / (pred_img.max(axis=(1,2,3))[:,np.newaxis,np.newaxis,np.newaxis] - pred_img.min(axis=(1,2,3))[:,np.newaxis,np.newaxis,np.newaxis])
|
||||
img = tf.concat((uv_static_img, uv_body_posed_img, uv_curvature, orig_img, pred_img), axis=2)
|
||||
imgs = [wandb.Image(img[i]) for i in range(4)]
|
||||
|
||||
F = np.load("custom_utils/smpl_faces" + skirt_modifier + ("_orig" if config['is_cape'] else "") + ".npy")
|
||||
mask = batch['mesh_mask'].numpy().astype(bool)
|
||||
_, faces = maskFacesBatch(np.ones_like(mesh_orig), F, mask)
|
||||
|
||||
if config['use_skinning']:
|
||||
shape = batch['pose_shape'][:, 72:]
|
||||
gender = tf.cast(tf.greater_equal(shape[:, 10] , 0.5), tf.int32)
|
||||
shape_offsets = tf.stack([tf.gather(tf.tensordot(shape[i,:10][tf.newaxis, :], smpl.shapedirs[gender[i]], axes=[[1], [2]])[0,:,:], smpl.to_keep, axis=0) for i in range(shape.shape[0])]).numpy()
|
||||
|
||||
|
||||
gt = [mesh_orig[i][mask[i], :] + shape_offsets[i, mask[i], :] for i in range(mask.shape[0])]
|
||||
|
||||
mesh_unposed_shaped_levels = [[] for _ in range(mask.shape[0])]
|
||||
for i in range(len(output_levels)):
|
||||
mesh_unposed_shaped = [mesh_pred_levels[i][j, mask[j], :] + shape_offsets[j, mask[j]] for j in range(mask.shape[0])]
|
||||
for k in range(mask.shape[0]):
|
||||
mesh_unposed_shaped_levels[k].append(mesh_unposed_shaped[k])
|
||||
|
||||
# full res version
|
||||
mesh_unposed_shaped_full = np.zeros((mask.shape[0], config['smpl_res'], 3), dtype=np.float32)
|
||||
for i in range(mask.shape[0]):
|
||||
mesh_unposed_shaped_full[i, mask[i], :] = mesh_unposed_shaped[i]
|
||||
|
||||
mesh_posed = smpl.run_garment(batch['pose_shape'][:, :72], batch['pose_shape'][:, 72:], mesh_unposed_shaped_full).numpy()
|
||||
else:
|
||||
mesh_posed = mesh_pred_levels[-1]
|
||||
gt = [mesh_orig[i][mask[i], :] for i in range(mesh_posed.shape[0])]
|
||||
mesh_unposed_shaped_levels = [[mesh_pred_levels[i][j, mask[j], :] for i in range(len(output_levels))] for j in range(mask.shape[0])]
|
||||
|
||||
|
||||
gt_posed = [mesh_orig_posed[i, mask[i], :] for i in range(mask.shape[0])]
|
||||
mesh_posed_fixed = fix_collision(mesh_posed, body, mask).numpy()
|
||||
|
||||
mesh_posed = [mesh_posed[j, mask[j], :] for j in range(mask.shape[0])]
|
||||
|
||||
if config['is_cape']:
|
||||
collision = mesh_posed
|
||||
collision_faces = faces
|
||||
collision_color = [np.ones_like(mesh_posed[i], dtype=np.int32)*230 for i in range(len(mesh_posed))]
|
||||
else:
|
||||
collision = [np.concatenate((body[i], mesh_posed[i]), axis=0) for i in range(len(mesh_posed))]
|
||||
collision_faces = [np.concatenate((F, faces[i] + config['smpl_res']), axis=0) for i in range(len(mesh_posed))]
|
||||
collision_color = [np.concatenate((np.ones_like(body[i], dtype=np.int32)*100, np.ones_like(mesh_posed[i], dtype=np.int32)*230), axis=0) for i in range(len(mesh_posed))]
|
||||
|
||||
figs = { type + "/garment_" + str(i): plot_3d(gt[i], mesh_unposed_shaped_levels[i], faces[i], gt_posed[i], collision[i], collision_faces[i], collision_color[i]) for i in range(4) }
|
||||
|
||||
if config['is_cape']:
|
||||
figs_coll = dict()
|
||||
else:
|
||||
garment_posed_fixed = [mesh_posed_fixed[i, mask[i], :] for i in range(body.shape[0])]
|
||||
collision_fixed = [np.concatenate((body[i], garment_posed_fixed[i]), axis=0) for i in range(len(mesh_posed))]
|
||||
gt_posed = [batch['mesh_posed'].numpy()[i, mask[i], :] for i in range(mask.shape[0])]
|
||||
collision_gt = [np.concatenate((body[i], gt_posed[i]), axis=0) for i in range(len(gt_posed))]
|
||||
figs_coll = { type + "/garment_" + str(i) + '_collision': plot_3d_collision(collision_gt[i], collision[i], collision_fixed[i], collision_faces[i], collision_color[i]) for i in range(4) }
|
||||
|
||||
return { **figs, **figs_coll, **{ type + '/images': imgs, } }
|
||||
|
||||
|
||||
def sample(folder):
|
||||
# run = wandb.init(project='vae_plots', name=folder.split('/')[-1])
|
||||
load_model(model,folder)
|
||||
count = 0
|
||||
for batch in dataset_train:
|
||||
batch = static_dataset_train.expand_batch(batch_to_dict(batch, config['signature_dynamic']))
|
||||
uv_curvature = tf.one_hot(batch['uv_curvature'], 7)
|
||||
nr_of_samples=2
|
||||
|
||||
pred = model.sample(batch['uv_static'][0], batch['uv_body_posed'][0], uv_curvature[0], batch['uv_mask'][0], nr_of_samples=nr_of_samples)
|
||||
pose, shape = batch['pose_shape'][0, :72], batch['pose_shape'][0, 72:]
|
||||
pose, shape = tf.tile(pose[tf.newaxis,:], [nr_of_samples,1]), tf.tile(shape[tf.newaxis,:], [nr_of_samples,1])
|
||||
|
||||
mesh_pred = uv2mesh(tf.transpose(pred, [0,3,1,2]))
|
||||
mesh_posed = smpl.run_garment(pose, shape, mesh_pred).numpy()
|
||||
mesh_orig = batch['mesh_posed'][0].numpy()
|
||||
mesh_orig_reposed = smpl.run_garment(pose[:1], shape[:1], batch['mesh_unposed'][:1]).numpy()[0]
|
||||
|
||||
body = batch['mesh_body_posed'][0].numpy()
|
||||
body_color = np.ones_like(body, dtype=np.int32)*100
|
||||
|
||||
F = np.load("custom_utils/smpl_faces.npy")
|
||||
mask = batch['mesh_mask'].numpy().astype(bool)[0]
|
||||
_, faces = _maskFaces(np.ones((14475, 3)), F, mask)
|
||||
|
||||
|
||||
pred_zero = model.sample(batch['uv_static'][0], batch['uv_body_posed'][0], uv_curvature[0], batch['uv_mask'][0], nr_of_samples=nr_of_samples, zero_latent=True)
|
||||
mesh_pred_zero = uv2mesh(tf.transpose(pred_zero, [0,3,1,2]))
|
||||
mesh_posed_zero = smpl.run_garment(pose, shape, mesh_pred_zero).numpy()
|
||||
|
||||
|
||||
|
||||
|
||||
meshes = []
|
||||
meshes.append(go.Mesh3d(
|
||||
# 8 vertices of a cube
|
||||
x=mesh_orig_reposed[mask,0],
|
||||
y=mesh_orig_reposed[mask,2],
|
||||
z=mesh_orig_reposed[mask,1],
|
||||
# i, j and k give the vertices of triangles
|
||||
i = faces[:,0],
|
||||
j = faces[:,1],
|
||||
k = faces[:,2]
|
||||
))
|
||||
|
||||
for i in range(nr_of_samples):
|
||||
x = mesh_posed[i][mask, :].copy()
|
||||
# x[:, 2] += 1 * (i + 1)
|
||||
x[:, 2] += 1
|
||||
meshes.append(go.Mesh3d(
|
||||
# 8 vertices of a cube
|
||||
x=x[:,0],
|
||||
y=x[:,2],
|
||||
z=x[:,1],
|
||||
# i, j and k give the vertices of triangles
|
||||
i = faces[:,0],
|
||||
j = faces[:,1],
|
||||
k = faces[:,2]
|
||||
))
|
||||
|
||||
|
||||
for i in range(nr_of_samples):
|
||||
x = mesh_posed_zero[i][mask, :].copy()
|
||||
# x[:, 2] += 1 * (i + 1)
|
||||
x[:, 2] += 2
|
||||
meshes.append(go.Mesh3d(
|
||||
# 8 vertices of a cube
|
||||
x=x[:,0],
|
||||
y=x[:,2],
|
||||
z=x[:,1],
|
||||
# i, j and k give the vertices of triangles
|
||||
i = faces[:,0],
|
||||
j = faces[:,1],
|
||||
k = faces[:,2]
|
||||
))
|
||||
|
||||
# for i in range(5):
|
||||
# x = body.copy()
|
||||
# x[:, 2] += 1 * i
|
||||
# meshes.append(go.Mesh3d(
|
||||
# # 8 vertices of a cube
|
||||
# x=x[:,0],
|
||||
# y=x[:,2],
|
||||
# z=x[:,1],
|
||||
# # i, j and k give the vertices of triangles
|
||||
# i = F[:,0],
|
||||
# j = F[:,1],
|
||||
# k = F[:,2],
|
||||
# vertexcolor=body_color
|
||||
# ))
|
||||
|
||||
|
||||
fig = go.Figure(data=meshes)
|
||||
|
||||
fig.update_layout(scene_aspectmode='manual', scene_aspectratio=dict(x=1, y=2, z=1))
|
||||
# fig.update_layout(scene_aspectmode='manual', scene_aspectratio=dict(x=1, y=nr_of_samples-1, z=1))
|
||||
fig.update_layout(
|
||||
scene = dict(
|
||||
xaxis = dict(nticks=4, range=[-1,1],),
|
||||
# yaxis = dict(nticks=4, range=[-1,nr_of_samples+1],),
|
||||
yaxis = dict(nticks=4, range=[-1,3],),
|
||||
zaxis = dict(nticks=4, range=[-1,1],),)
|
||||
)
|
||||
|
||||
# wandb.log({'plot_' + str(count): fig})
|
||||
fig.show()
|
||||
|
||||
pred = pred.numpy()
|
||||
frames = [Image.fromarray(((pred[i] + 1) * 128).astype(np.int8), mode='RGB') for i in range(pred.shape[0])]
|
||||
frames[0].save("temp/garment_1.gif", format="GIF", append_images=frames[1:],
|
||||
save_all=True, duration=500, loop=0)
|
||||
|
||||
|
||||
if count > 10:
|
||||
break
|
||||
|
||||
print(count)
|
||||
count += 1
|
||||
|
||||
print()
|
||||
|
||||
|
||||
def test():
|
||||
garment_type_table = np.load('custom_data/garment_type_table_test.npy', allow_pickle=True).item()
|
||||
str_split = config['name'].split('_')
|
||||
test_idx = str_split.index('test')
|
||||
checkpoint_name = '_'.join(str_split[:test_idx])
|
||||
load_model(model, 'checkpoints/' + checkpoint_name)
|
||||
|
||||
dataset_test = factory.get_dataset(folder=config['test_dir'], batch_size=config['batch_size'], signature=config['signature_dynamic'], use_filled=config['use_filled'], resolutions=dataset_resolutions, in_memory=False, do_repeat=False)
|
||||
static_dataset_test = StaticDataset(folder=config['test_dir'], signature=config['signature_static'], use_filled=config['use_filled'], resolutions=dataset_resolutions, in_memory=False)
|
||||
|
||||
step = 0
|
||||
losses = []
|
||||
|
||||
if config['is_skirt']:
|
||||
per_garment_errors = { 'Tshirt': [], 'Dress': [], 'Skirt': [], 'Top': [] }
|
||||
else:
|
||||
per_garment_errors = { 'Tshirt': [], 'Jumpsuit': [], 'Trousers': [], 'Top': [] }
|
||||
|
||||
for batch in tqdm.tqdm(dataset_test):
|
||||
batch = static_dataset_test.expand_batch(batch_to_dict(batch, config['signature_dynamic']))
|
||||
|
||||
test_loss, euclid_error_per_garment = test_step(batch)
|
||||
losses.append(test_loss)
|
||||
|
||||
for i, id in enumerate(batch['id']):
|
||||
id_str = id.numpy().decode("utf-8")
|
||||
per_garment_errors[garment_type_table[id_str]].append(euclid_error_per_garment[i].numpy())
|
||||
|
||||
to_log = { **{ 'current_step': step }, **test_loss }
|
||||
if step % 1000 == 0:
|
||||
figs = plot_everything(batch, 'test')
|
||||
to_log = { **to_log, **figs }
|
||||
|
||||
wandb.log(to_log, step = step)
|
||||
|
||||
# print(step)
|
||||
step += 1
|
||||
|
||||
# averages = dict()
|
||||
path_to_plotly_html = "temp/plotly_figure.html"
|
||||
|
||||
histograms = dict()
|
||||
for key in test_loss.keys():
|
||||
accumulated = [losses[i][key].numpy() for i in range(len(losses))]
|
||||
fig = ex.histogram(accumulated, title='Average: ' + str(np.asarray(accumulated).mean()))
|
||||
fig.write_html(path_to_plotly_html, auto_play = False)
|
||||
|
||||
histograms['all/' + key.split('/')[-1] + '_hist'] = wandb.Html(path_to_plotly_html)
|
||||
# averages['all/' + key.split('/')[-1] + '_avg'] = np.asarray(accumulated).mean()
|
||||
|
||||
for key in per_garment_errors.keys():
|
||||
fig = ex.histogram(per_garment_errors[key], title='Average: ' + str(np.asarray(per_garment_errors[key]).mean()))
|
||||
fig.write_html(path_to_plotly_html, auto_play = False)
|
||||
histograms['all/' + key + 'euclid_error_hist'] = wandb.Html(path_to_plotly_html)
|
||||
|
||||
|
||||
wandb.log({ **histograms }, step = step)
|
||||
wandb.finish()
|
||||
pass
|
||||
|
||||
|
||||
def generate_results():
|
||||
load_model(model, 'checkpoints/' + config['name'])
|
||||
|
||||
dataset_test = factory.get_dataset(folder=config['test_dir'],
|
||||
batch_size=config['batch_size'],
|
||||
# batch_size=1,
|
||||
signature=config['signature_dynamic'],
|
||||
resolutions=dataset_resolutions,
|
||||
in_memory=config['dataset_in_memory'],
|
||||
is_cape=config['is_cape'],
|
||||
frames_file=config['frames_file'] if 'frames_file' in config.keys() else None,
|
||||
do_repeat=False,
|
||||
shuffle=False)
|
||||
static_dataset_test = StaticDataset(folder=config['test_dir'],
|
||||
signature=config['signature_static'],
|
||||
resolutions=dataset_resolutions,
|
||||
in_memory=True,
|
||||
is_cape=config['is_cape'])
|
||||
folder = os.path.join('/home/workspace/data/results/cape', config['name'])
|
||||
os.makedirs(folder, exist_ok=True)
|
||||
|
||||
for batch in tqdm.tqdm(dataset_test):
|
||||
batch = static_dataset_test.expand_batch(batch_to_dict(batch, config['signature_dynamic']))
|
||||
|
||||
levels_input, conds = group_inputs_for_levels(batch, config)
|
||||
out_levels = model([levels_input, conds], training=False)
|
||||
|
||||
shape = batch['pose_shape'][:, 72:]
|
||||
gender = tf.cast(tf.greater_equal(shape[:, 10] , 0.5), tf.int32)
|
||||
shape_offsets = tf.stack([tf.gather(tf.tensordot(shape[i,:10][tf.newaxis, :], smpl.shapedirs[gender[i]], axes=[[1], [2]])[0,:,:], smpl.to_keep, axis=0) for i in range(shape.shape[0])])
|
||||
|
||||
i = -1
|
||||
pred = out_levels[i][0]
|
||||
|
||||
mesh_pred = uv2mesh(pred)
|
||||
mesh_pred_shaped = mesh_pred + shape_offsets * batch['mesh_mask'][:,:,tf.newaxis]
|
||||
mesh_pred_posed = smpl.run_garment(batch['pose_shape'][:, :72], batch['pose_shape'][:, 72:], mesh_pred_shaped)
|
||||
|
||||
for i in range(batch['id'].shape[0]):
|
||||
mask = batch['mesh_mask'][i].numpy().astype(bool)
|
||||
id, frame_nr = batch['id'][i].numpy().decode("utf-8"), batch['frame_nr'][i].numpy().decode("utf-8")
|
||||
unposed = mesh_pred[i].numpy()[mask, :]
|
||||
posed = mesh_pred_posed[i].numpy()[mask, :]
|
||||
np.save(os.path.join(folder, id + '_' + frame_nr + '_posed.npy'), posed)
|
||||
np.save(os.path.join(folder, id + '_' + frame_nr + '_unposed.npy'), unposed)
|
||||
|
||||
pass
|
||||
|
||||
@tf.function
|
||||
def speed_step(levels_input, conds):
|
||||
|
||||
return model.sample([levels_input, conds], training=False)
|
||||
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
# init_wandb()
|
||||
parser = argparse.ArgumentParser(description='')
|
||||
parser.add_argument('--name', type=str, default='lhunor')
|
||||
args = vars(parser.parse_args())
|
||||
|
||||
# config_file_names = [key for key in sys.modules.keys() if key.startswith('configs.')]
|
||||
# assert len(config_file_names) == 1
|
||||
# config_file_name = config_file_names[0]
|
||||
# # assert config_file_name.split('.')[1].endswith(args.name)
|
||||
# assert config_file_name.split('.')[1].endswith(config['name'])
|
||||
|
||||
smpl = SMPL(orig_smpl=config['is_cape'])
|
||||
skirt_modifier = '_skirt' if config['is_skirt'] else ''
|
||||
dataset_resolutions = config['pyramid']['resolutions'].copy()
|
||||
if 512 not in dataset_resolutions:
|
||||
dataset_resolutions.append(512)
|
||||
|
||||
factory = DatasetFactory()
|
||||
dataset_train = factory.get_dataset(folder=config['train_dir'],
|
||||
batch_size=config['batch_size'],
|
||||
# batch_size=1,
|
||||
signature=config['signature_dynamic'],
|
||||
resolutions=dataset_resolutions,
|
||||
in_memory=config['dataset_in_memory'],
|
||||
is_cape=config['is_cape'],
|
||||
frames_file=config['frames_file'] if 'frames_file' in config.keys() else None)
|
||||
static_dataset_train = StaticDataset(folder=config['train_dir'],
|
||||
signature=config['signature_static'],
|
||||
resolutions=dataset_resolutions,
|
||||
in_memory=True,
|
||||
is_cape=config['is_cape'])
|
||||
dataset_valid = factory.get_dataset(folder=config['valid_dir'],
|
||||
batch_size=config['batch_size'],
|
||||
signature=config['signature_dynamic'],
|
||||
resolutions=dataset_resolutions,
|
||||
in_memory=config['dataset_in_memory'],
|
||||
is_cape=config['is_cape'],
|
||||
frames_file=config['frames_file_valid'] if 'frames_file_valid' in config.keys() else None)
|
||||
static_dataset_valid = StaticDataset(folder=config['valid_dir'],
|
||||
signature=config['signature_static'],
|
||||
resolutions=dataset_resolutions,
|
||||
in_memory=True,
|
||||
is_cape=config['is_cape'])
|
||||
|
||||
model: Pyramid = Pyramid(config)
|
||||
|
||||
optimizer = tf.optimizers.Adam(learning_rate=config['lr'])
|
||||
balancing_coeff = config['starting_balancing_coeff']
|
||||
reg_coeff = 1
|
||||
valid_iterator = iter(dataset_valid)
|
||||
gamma = 1
|
||||
balancing_activated_step = 0
|
||||
|
||||
# sample('checkpoints/cvae_13')
|
||||
# test()
|
||||
|
||||
# generate_results()
|
||||
|
||||
ckpt = tf.train.Checkpoint(model=model)
|
||||
manager = tf.train.CheckpointManager(ckpt, 'checkpoints/' + config['name'], max_to_keep=20)
|
||||
|
||||
# load_model(model, 'checkpoints/pyr_abl_32_skirt', iteration='180000')
|
||||
|
||||
if config['pretrained_model'] != '':
|
||||
# assert len(config['pyramid']['resolutions']) == 1
|
||||
load_model(model, config['pretrained_model'], iteration=config['pretrained_model_iteration'])
|
||||
|
||||
if len(config['pyramid']['resolutions']) == 1:
|
||||
model.static_enc.trainable = False
|
||||
model.body_enc.trainable = False
|
||||
model.curv_enc.trainable = False
|
||||
else:
|
||||
load_model(model, config['pyramid']['pretrained_pyramid'], iteration=(config['pyramid']['pretrained_pyramid_iteration'] if 'pretrained_pyramid_iteration' in config.keys() else ''))
|
||||
model.static_enc.trainable = False
|
||||
model.body_enc.trainable = False
|
||||
model.curv_enc.trainable = False
|
||||
model.cond_encoder.trainable = False
|
||||
|
||||
# load_model(model.levels[0].upscale_enhancer, config['pyramid']['levels'][0]['upscale_enhancer'])
|
||||
# model.levels[0].upscale_enhancer.trainable = False
|
||||
|
||||
|
||||
for i in range(len(model.levels)-1):
|
||||
load_model(model.levels[i].upscale_enhancer, config['pyramid']['levels'][i]['upscale_enhancer'])
|
||||
model.levels[i].upscale_enhancer.trainable = False
|
||||
if config['pyramid']['resolutions'][-1] != 512:
|
||||
load_model(model.upscale_enhancer, config['pyramid']['final_upscale_enhancer'])
|
||||
model.upscale_enhancer.trainable = False
|
||||
|
||||
for i in range(len(model.levels)-1):
|
||||
model.levels[i].trainable = False
|
||||
|
||||
# load_model(model.cond_encoder, config['pyramid']['pretrained_pyramid'])
|
||||
# load_model(model.levels[0], config['pyramid']['pretrained_pyramid'])
|
||||
# load_model(model.levels[-1], 'checkpoints/cvae_19')
|
||||
|
||||
|
||||
step = 0
|
||||
times = []
|
||||
for batch in dataset_train:
|
||||
batch = static_dataset_train.expand_batch(batch_to_dict(batch, config['signature_dynamic']))
|
||||
|
||||
# levels_input, conds = group_inputs_for_levels(batch, config)
|
||||
|
||||
# start = time.perf_counter()
|
||||
# out = speed_step(levels_input, conds)
|
||||
# t = time.perf_counter() - start
|
||||
# print(t)
|
||||
# times.append(t)
|
||||
|
||||
# if step > 200:
|
||||
# break
|
||||
|
||||
|
||||
train_loss = train_step(batch)
|
||||
|
||||
data_valid = valid_iterator.get_next()
|
||||
data_valid = static_dataset_valid.expand_batch(batch_to_dict(data_valid, config['signature_dynamic']))
|
||||
valid_loss = valid_step(data_valid)
|
||||
|
||||
to_log = { **train_loss, **valid_loss, **{ 'current_step': step} }
|
||||
if step % 5000 == 0 or (step < 5000 and step % 500 == 0):
|
||||
train_figs = plot_everything(batch, 'train')
|
||||
valid_figs = plot_everything(data_valid, 'valid')
|
||||
to_log = { **to_log, **train_figs, **valid_figs }
|
||||
|
||||
if step % 5000 == 0:
|
||||
manager.save(checkpoint_number=step)
|
||||
if step == 0:
|
||||
to_log = { **to_log, **{ 'model/summary': wandb.Html(get_summary_html(model)) } }
|
||||
|
||||
if train_loss['train/euclid_error'] < 0.004:
|
||||
if balancing_activated_step == 0:
|
||||
balancing_activated_step = step
|
||||
balancing_coeff = balancing_coeff + 1 / 20000
|
||||
|
||||
|
||||
to_log = { **to_log, **{ 'train/balancing_coeff': balancing_coeff, 'train/gamma': gamma } }
|
||||
|
||||
wandb.log(to_log, step = step)
|
||||
|
||||
print(step)
|
||||
|
||||
step += 1
|
||||
# print('mean:')
|
||||
# print(np.mean(times[1:]))
|
||||
Reference in New Issue
Block a user