diff --git a/.gitignore b/.gitignore
index 5d381cc..9e0b2f3 100644
--- a/.gitignore
+++ b/.gitignore
@@ -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*
diff --git a/custom_data/uv_map/stretched_uv_map_3.npy b/custom_data/uv_map/stretched_uv_map_3.npy
new file mode 100644
index 0000000..5bd87e3
Binary files /dev/null and b/custom_data/uv_map/stretched_uv_map_3.npy differ
diff --git a/custom_data/uv_map/uv_data_plane_stretched_front_high.npz b/custom_data/uv_map/uv_data_plane_stretched_front_high.npz
new file mode 100644
index 0000000..9d25262
Binary files /dev/null and b/custom_data/uv_map/uv_data_plane_stretched_front_high.npz differ
diff --git a/custom_data/uv_map/uv_data_plane_stretched_front_low.npz b/custom_data/uv_map/uv_data_plane_stretched_front_low.npz
new file mode 100644
index 0000000..1a6b144
Binary files /dev/null and b/custom_data/uv_map/uv_data_plane_stretched_front_low.npz differ
diff --git a/custom_data/uv_map/uv_data_smpl_2.npz b/custom_data/uv_map/uv_data_smpl_2.npz
new file mode 100644
index 0000000..f200b6f
Binary files /dev/null and b/custom_data/uv_map/uv_data_smpl_2.npz differ
diff --git a/custom_data/uv_map/uv_data_smpl_2_less_dense.npz b/custom_data/uv_map/uv_data_smpl_2_less_dense.npz
new file mode 100644
index 0000000..7d1904f
Binary files /dev/null and b/custom_data/uv_map/uv_data_smpl_2_less_dense.npz differ
diff --git a/custom_utils/NN.mat b/custom_utils/NN.mat
new file mode 100644
index 0000000..f2e3ac2
Binary files /dev/null and b/custom_utils/NN.mat differ
diff --git a/custom_utils/create_uv_map.py b/custom_utils/create_uv_map.py
new file mode 100644
index 0000000..5958166
--- /dev/null
+++ b/custom_utils/create_uv_map.py
@@ -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
\ No newline at end of file
diff --git a/custom_utils/create_uv_map_2.py b/custom_utils/create_uv_map_2.py
new file mode 100644
index 0000000..fcb99db
--- /dev/null
+++ b/custom_utils/create_uv_map_2.py
@@ -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
\ No newline at end of file
diff --git a/custom_utils/custom_utils.py b/custom_utils/custom_utils.py
new file mode 100644
index 0000000..77508c1
--- /dev/null
+++ b/custom_utils/custom_utils.py
@@ -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 += "" + line + "
"
+ 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
\ No newline at end of file
diff --git a/custom_utils/fill_uv.py b/custom_utils/fill_uv.py
new file mode 100644
index 0000000..79ec71e
--- /dev/null
+++ b/custom_utils/fill_uv.py
@@ -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()
\ No newline at end of file
diff --git a/custom_utils/graph_utils.py b/custom_utils/graph_utils.py
new file mode 100644
index 0000000..1113727
--- /dev/null
+++ b/custom_utils/graph_utils.py
@@ -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()
\ No newline at end of file
diff --git a/custom_utils/keep_remove.mat b/custom_utils/keep_remove.mat
new file mode 100644
index 0000000..213b672
Binary files /dev/null and b/custom_utils/keep_remove.mat differ
diff --git a/custom_utils/m2uv.npy b/custom_utils/m2uv.npy
new file mode 100644
index 0000000..56708d7
Binary files /dev/null and b/custom_utils/m2uv.npy differ
diff --git a/custom_utils/m2uv_orig.npy b/custom_utils/m2uv_orig.npy
new file mode 100644
index 0000000..6abd9d7
Binary files /dev/null and b/custom_utils/m2uv_orig.npy differ
diff --git a/custom_utils/m2uv_skirt.npy b/custom_utils/m2uv_skirt.npy
new file mode 100644
index 0000000..6448f02
Binary files /dev/null and b/custom_utils/m2uv_skirt.npy differ
diff --git a/custom_utils/m2uv_skirt.py b/custom_utils/m2uv_skirt.py
new file mode 100644
index 0000000..1264b1c
--- /dev/null
+++ b/custom_utils/m2uv_skirt.py
@@ -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()
\ No newline at end of file
diff --git a/custom_utils/mask_256x512.npy b/custom_utils/mask_256x512.npy
new file mode 100644
index 0000000..770d714
Binary files /dev/null and b/custom_utils/mask_256x512.npy differ
diff --git a/custom_utils/mask_256x512_skirt.npy b/custom_utils/mask_256x512_skirt.npy
new file mode 100644
index 0000000..0230632
Binary files /dev/null and b/custom_utils/mask_256x512_skirt.npy differ
diff --git a/custom_utils/smpl.py b/custom_utils/smpl.py
new file mode 100644
index 0000000..77d09c6
--- /dev/null
+++ b/custom_utils/smpl.py
@@ -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
\ No newline at end of file
diff --git a/custom_utils/smpl_2.py b/custom_utils/smpl_2.py
new file mode 100644
index 0000000..bc12542
--- /dev/null
+++ b/custom_utils/smpl_2.py
@@ -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
\ No newline at end of file
diff --git a/custom_utils/smpl_faces.npy b/custom_utils/smpl_faces.npy
new file mode 100644
index 0000000..d1d74c9
Binary files /dev/null and b/custom_utils/smpl_faces.npy differ
diff --git a/custom_utils/smpl_faces_orig.npy b/custom_utils/smpl_faces_orig.npy
new file mode 100644
index 0000000..282dc29
Binary files /dev/null and b/custom_utils/smpl_faces_orig.npy differ
diff --git a/custom_utils/smpl_faces_skirt.npy b/custom_utils/smpl_faces_skirt.npy
new file mode 100644
index 0000000..d6aa40f
Binary files /dev/null and b/custom_utils/smpl_faces_skirt.npy differ
diff --git a/custom_utils/ssmpl0.mat b/custom_utils/ssmpl0.mat
new file mode 100644
index 0000000..e3cea72
Binary files /dev/null and b/custom_utils/ssmpl0.mat differ
diff --git a/custom_utils/ssmpl_faces.mat b/custom_utils/ssmpl_faces.mat
new file mode 100644
index 0000000..bbf7c78
Binary files /dev/null and b/custom_utils/ssmpl_faces.mat differ
diff --git a/custom_utils/ssmpl_skirt0.mat b/custom_utils/ssmpl_skirt0.mat
new file mode 100644
index 0000000..3755659
Binary files /dev/null and b/custom_utils/ssmpl_skirt0.mat differ
diff --git a/custom_utils/stretched_uv_map_3.npy b/custom_utils/stretched_uv_map_3.npy
new file mode 100644
index 0000000..5bd87e3
Binary files /dev/null and b/custom_utils/stretched_uv_map_3.npy differ
diff --git a/custom_utils/stretched_uv_map_3_orig.npy b/custom_utils/stretched_uv_map_3_orig.npy
new file mode 100644
index 0000000..296f66b
Binary files /dev/null and b/custom_utils/stretched_uv_map_3_orig.npy differ
diff --git a/custom_utils/stretched_uv_map_skirt.npy b/custom_utils/stretched_uv_map_skirt.npy
new file mode 100644
index 0000000..3b7bb91
Binary files /dev/null and b/custom_utils/stretched_uv_map_skirt.npy differ
diff --git a/custom_utils/super_model_f.pkl b/custom_utils/super_model_f.pkl
new file mode 100644
index 0000000..070c314
Binary files /dev/null and b/custom_utils/super_model_f.pkl differ
diff --git a/custom_utils/super_model_m.pkl b/custom_utils/super_model_m.pkl
new file mode 100644
index 0000000..bfb6bf0
Binary files /dev/null and b/custom_utils/super_model_m.pkl differ
diff --git a/custom_utils/supersmpl_io.py b/custom_utils/supersmpl_io.py
new file mode 100644
index 0000000..400d2bf
--- /dev/null
+++ b/custom_utils/supersmpl_io.py
@@ -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('= 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(' 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)
\ No newline at end of file
diff --git a/custom_utils/tps_custom.py b/custom_utils/tps_custom.py
new file mode 100644
index 0000000..651c2a3
--- /dev/null
+++ b/custom_utils/tps_custom.py
@@ -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__
diff --git a/custom_utils/unpose.py b/custom_utils/unpose.py
new file mode 100644
index 0000000..97bac12
--- /dev/null
+++ b/custom_utils/unpose.py
@@ -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()
+
+
+
\ No newline at end of file
diff --git a/custom_utils/uv_zero_body.npy b/custom_utils/uv_zero_body.npy
new file mode 100644
index 0000000..a231698
Binary files /dev/null and b/custom_utils/uv_zero_body.npy differ
diff --git a/data/dataset.py b/data/dataset.py
new file mode 100644
index 0000000..d2ee42d
--- /dev/null
+++ b/data/dataset.py
@@ -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
+
\ No newline at end of file
diff --git a/docker/Dockerfile b/docker/Dockerfile
new file mode 100644
index 0000000..2641a78
--- /dev/null
+++ b/docker/Dockerfile
@@ -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"]
\ No newline at end of file
diff --git a/docker/requierements.txt b/docker/requierements.txt
new file mode 100644
index 0000000..64e5458
--- /dev/null
+++ b/docker/requierements.txt
@@ -0,0 +1,12 @@
+numpy
+matplotlib
+opencv-python
+wandb
+OpenEXR
+tensorflow-graphics
+scipy
+plotly
+pandas
+tensorboardX
+opencv-python
+cmapy
\ No newline at end of file
diff --git a/models/base.py b/models/base.py
new file mode 100644
index 0000000..e9cfd67
--- /dev/null
+++ b/models/base.py
@@ -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
\ No newline at end of file
diff --git a/models/cond_encoder.py b/models/cond_encoder.py
new file mode 100644
index 0000000..c86f54c
--- /dev/null
+++ b/models/cond_encoder.py
@@ -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
\ No newline at end of file
diff --git a/models/convnext.py b/models/convnext.py
new file mode 100644
index 0000000..55bad56
--- /dev/null
+++ b/models/convnext.py
@@ -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]
+ ),
+)
+
diff --git a/models/cvae_curv_model.py b/models/cvae_curv_model.py
new file mode 100644
index 0000000..c7e8878
--- /dev/null
+++ b/models/cvae_curv_model.py
@@ -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
diff --git a/models/cvae_model.py b/models/cvae_model.py
new file mode 100644
index 0000000..f851741
--- /dev/null
+++ b/models/cvae_model.py
@@ -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
+
diff --git a/models/pyramid_level.py b/models/pyramid_level.py
new file mode 100644
index 0000000..9cf7322
--- /dev/null
+++ b/models/pyramid_level.py
@@ -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
\ No newline at end of file
diff --git a/models/pyramid_model_optimized.py b/models/pyramid_model_optimized.py
new file mode 100644
index 0000000..5301872
--- /dev/null
+++ b/models/pyramid_model_optimized.py
@@ -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
\ No newline at end of file
diff --git a/models/upscaler.py b/models/upscaler.py
new file mode 100644
index 0000000..c03e663
--- /dev/null
+++ b/models/upscaler.py
@@ -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
+
diff --git a/scripts/cvae.py b/scripts/cvae.py
new file mode 100644
index 0000000..b66874a
--- /dev/null
+++ b/scripts/cvae.py
@@ -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
+
\ No newline at end of file
diff --git a/scripts/cvae_pyramid.py b/scripts/cvae_pyramid.py
new file mode 100644
index 0000000..7594a89
--- /dev/null
+++ b/scripts/cvae_pyramid.py
@@ -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:]))
\ No newline at end of file