112 lines
2.4 KiB
Python
112 lines
2.4 KiB
Python
import sys
|
|
|
|
import numpy as np
|
|
import pickle as pkl
|
|
from scipy import sparse
|
|
|
|
from time import time
|
|
|
|
from custom_utils.supersmpl_io import writeOBJ, readOBJ
|
|
|
|
|
|
def makeGraph(F):
|
|
G = {}
|
|
for f in F:
|
|
N = len(f)
|
|
for i,v in enumerate(f):
|
|
if not v in G: G[v] = {}
|
|
G[v][f[(i+1) % N]] = 1
|
|
return G
|
|
|
|
def extrapolateMissings(M, G):
|
|
for i in range(6890, M.shape[0]):
|
|
if np.isnan(M[i]).all():
|
|
M[i] = neighAvg(M, G, i)
|
|
elif np.isnan(M[i]).any():
|
|
print("HOW?")
|
|
print(M[i])
|
|
if np.isnan(M).any(): extrapolateMissings(M, G)
|
|
|
|
def neighAvg(M, G, i):
|
|
neigh = set(G[i].keys())
|
|
avg = 0
|
|
count = 0
|
|
for n in neigh:
|
|
if not np.isnan(M[n]).any():
|
|
avg += M[n]
|
|
count += 1
|
|
if count: return avg / count
|
|
return np.nan
|
|
|
|
|
|
if __name__ == '__main__':
|
|
|
|
g = 'f'
|
|
# Read Super topology
|
|
V, F, Vt, Ft = readOBJ('super_f_avg.obj')
|
|
N = V.shape[0]
|
|
G = makeGraph(F)
|
|
|
|
# Read PKL
|
|
print("Loading .pkl file")
|
|
with open('custom_data/model_' + g + '.pkl', 'rb') as f:
|
|
M = pkl.load(f, encoding='latin1')
|
|
|
|
# Fields
|
|
# print("Checking fields and shapes")
|
|
# for k in M:
|
|
# print(k, M[k].shape)
|
|
# Faces
|
|
print("New faces...")
|
|
M['f'] = F
|
|
# J_regressor (24, 6890)
|
|
# Sparse matrix, just change shape to add zeros
|
|
print("New regressor")
|
|
A = np.zeros((24, N), np.float32)
|
|
A[:,:6890] = M['J_regressor'].todense()
|
|
M['J_regressor'] = sparse.csr_matrix(A)
|
|
|
|
# Weights (6890, 24)
|
|
print("New weights")
|
|
start = time()
|
|
A = np.empty((N, 24), np.float32)
|
|
A[:] = np.nan
|
|
A[:6890] = M['weights'][:6890]
|
|
extrapolateMissings(A, G)
|
|
M['weights'] = A
|
|
print(time() - start)
|
|
|
|
# Posedirs (6890, 3, 207)
|
|
print("New posedirs")
|
|
start = time()
|
|
A = np.empty((N, 3, 207), np.float32)
|
|
A[:] = np.nan
|
|
A[:6890] = M['posedirs'][:6890]
|
|
extrapolateMissings(A, G)
|
|
M['posedirs'] = A
|
|
print(time() - start)
|
|
|
|
# v_tempalte (6890, 3)
|
|
print("New template")
|
|
start = time()
|
|
A = np.empty((N, 3), np.float32)
|
|
A[:] = np.nan
|
|
A[:6890] = M['v_template'][:6890]
|
|
extrapolateMissings(A, G)
|
|
M['v_template'] = A
|
|
print(time() - start)
|
|
|
|
# Shapedirs (6890, 3, 10)
|
|
print("New shapedirs")
|
|
start = time()
|
|
A = np.empty((N, 3, 10), np.float32)
|
|
A[:] = np.nan
|
|
A[:6890] = M['shapedirs'][:6890]
|
|
extrapolateMissings(A, G)
|
|
M['shapedirs'] = A
|
|
print(time() - start)
|
|
|
|
# Dump
|
|
print("Saving new model")
|
|
with open('super_model_' + g + '.pkl', 'wb') as f:
|
|
pkl.dump(M, f) |