def __init__(self, parent=None):
super(MainWindow, self).__init__(parent=parent)
# GUI
self.ui = Ui_SmplViewer()
self.ui.setupUi(self)
# SMPL object
self.smpl = SMPLModel('./smpl_data/male_model.pkl')
# initialize the sliders
self.poseSliders = []
self.shapeSliders = []
for item in self.ui.__dict__:
if 'horizontalSlider' in item:
self.poseSliders.append(self.ui.__dict__[item])
if 'verticalSlider' in item:
self.shapeSliders.append(self.ui.__dict__[item])
# set events
for i in range(len(self.poseSliders)):
self.poseSliders[i].setMinimum(-8)
self.poseSliders[i].setMaximum(8)
self.poseSliders[i].valueChanged[int].connect(self.changevalue)
for i in range(len(self.shapeSliders)):
self.shapeSliders[i].setMinimum(0)
self.shapeSliders[i].setMaximum(10)
self.shapeSliders[i].valueChanged[int].connect(self.changevalue)
self.ui.pushButton.clicked.connect(self.pushButton_Click)
def generate_objects(hmmr_out):
folder_path = model_folder
model = SMPLModel("assets/basicmodel_m_lbs_10_207_0_v1.0.0.pkl")
counter = 0
for i in hmmr_out["verts"]:
model.verts = i
model.save_to_obj(folder_path + "/" + str(counter) + ".obj")
counter += 1
def setup_test_data(trainable_params, DISTRACTOR, PARAMS_TO_OFFSET, POSE_OFFSET, ARCHITECTURE, JOINT_LEVELS, data_samples, num_test_samples=5, num_cb_samples=5, MODE="RODRIGUES", LOAD_DATA_DIR=None):
# Gather trainable params
param_ids = ["param_{:02d}".format(i) for i in range(85)]
trainable_params = format_offsetable_params(trainable_params)
param_trainable = { param: (param in trainable_params) for param in param_ids }
DISTRACTOR = format_distractor_dict(DISTRACTOR, trainable_params)
# Gather offsetable params
if PARAMS_TO_OFFSET == "trainable_params":
PARAMS_TO_OFFSET = trainable_params
PARAMS_TO_OFFSET = format_offsetable_params(PARAMS_TO_OFFSET)
POSE_OFFSET = format_distractor_dict(POSE_OFFSET, PARAMS_TO_OFFSET)
# Define the kinematic tree
kin_tree = format_joint_levels(JOINT_LEVELS)
# Generate the data from the SMPL parameters
print("loading SMPL...")
smpl = SMPLModel('/data/cvfs/hjhb2/projects/deep_optimiser/code/keras_rotationnet_v2_demo_for_hidde/basicModel_f_lbs_10_207_0_v1.0.0.pkl')
# Generate and format the data
print("Gather input data...")
X_test, Y_test = gather_input_data(data_samples, smpl, PARAMS_TO_OFFSET, POSE_OFFSET, ARCHITECTURE, param_trainable, num_test_samples=num_test_samples, MODE=MODE, LOAD_DATA_DIR=LOAD_DATA_DIR)
# Render silhouettes for the callback data
X_cb, Y_cb, silh_cb = gather_cb_data(X_test, Y_test, data_samples, num_cb_samples, where="front")
return X_test, Y_test, X_cb, Y_cb, silh_cb, smpl, kin_tree, trainable_params, param_trainable, DISTRACTOR
def video_example():
target_folder = './video'
try:
os.makedirs(target_folder)
except:
pass
lv0 = './data'
lv1s = os.listdir(lv0)
for lv1 in tqdm(lv1s, ncols=120):
lv2s = os.listdir('/'.join([lv0, lv1]))
asf_path = os.path.join(lv0, lv1, lv1 + '.asf')
joints = reader.parse_asf(asf_path)
im = Imitator(joints, SMPLModel('./model.pkl'))
random.shuffle(lv2s)
lv2 = None
for lv2 in lv2s:
if lv2.split('.')[-1] == 'amc':
break
amc_path = os.path.join(lv0, lv1, lv2)
video_path = os.path.join(target_folder, '%s.avi' % lv2.split('.')[0])
motions = reader.parse_amc(amc_path)
viewer = MeshViewer(im, motions)
viewer.run(video_path=video_path,
render_fps=-1,
auto_run=True,
auto_rerun=False,
close_after_run=True)
def test_mesh():
subject = '13'
im = Imitator(reader.parse_asf('./data/%s/%s.asf' % (subject, subject)),
SMPLModel('./model.pkl'))
sequence = '01'
motions = reader.parse_amc('./data/%s/%s_%s.amc' %
(subject, subject, sequence))
viewer = MeshViewer(im, motions)
viewer.run()
def __init__(self,
num_samples,
reset_period,
POSE_OFFSET,
PARAMS_TO_OFFSET,
batch_size=32,
smpl=None,
shuffle=True,
save_path="./cb_samples.npz"):
self.num_samples = num_samples
self.reset_period = reset_period
if isinstance(POSE_OFFSET, (int, float)):
k = {param: POSE_OFFSET for param in PARAMS_TO_OFFSET}
else:
# k must be a dict with an entry for each variable parameter
k = POSE_OFFSET
self.k = k
self.PARAMS_TO_OFFSET = PARAMS_TO_OFFSET
self.batch_size = batch_size
if smpl is None:
smpl = SMPLModel(
'./keras_rotationnet_v2_demo_for_hidde/basicModel_f_lbs_10_207_0_v1.0.0.pkl'
)
self.smpl = smpl
self.shuffle = shuffle
self.params, self.pcs = gen_data(POSE_OFFSET,
PARAMS_TO_OFFSET,
self.smpl,
data_samples=num_samples,
save_dir=None,
render_silhouette=False)
self.indices = np.array([i for i in range(num_samples)])
self.save_path = save_path
#print("generator params shape at init: " + str(self.params.shape))
#print("generator pcs shape at init: " + str(self.pcs.shape))
print("generator params first entry: " + str(self.params[0]))
print("generator params final entry: " + str(self.params[-1]))
# Hard-coded parameters
self.epoch = 0
self.num_examples = 5
self.cb_samples = np.linspace(0,
self.num_samples,
num=self.num_examples,
dtype=int)
self.cb_samples[-1] -= 1
print("samples: " + str(self.cb_samples))
# Store initial data
if self.save_path is not None:
with open(self.save_path, 'w') as f:
np.savez(f,
indices=self.indices[self.cb_samples],
params=self.params[self.cb_samples],
pcs=self.pcs[self.cb_samples])
def save_generated_smpl(self, epoch, examples=3, render=True):
""" Save generated SMPL parameters """
if self.ex_data is None:
self.ex_data = np.random.randn(examples, self._gen_dim)
generated_params = self.generator.predict(self.ex_data)
for i, params in enumerate(generated_params):
np.save(
"../gen_pose_params/gan_example.E{:03d}:{:02d}.npy".format(
epoch, i), params)
if render:
smpl = SMPLModel('../model.pkl')
beta = np.zeros(shape=(10, ))
trans = np.zeros(shape=(3, ))
for i, pose_params in enumerate(generated_params):
smpl.set_params(beta=beta, pose=pose_params, trans=trans)
smpl.save_to_obj('../meshes/smpl_np.E{:03d}:{:02d}.obj'.format(
epoch, i))
# Format the distractor and offset dictionaries
DISTRACTOR = format_distractor_dict(DISTRACTOR, trainable_params)
if PARAMS_TO_OFFSET == "trainable_params":
PARAMS_TO_OFFSET = trainable_params
elif PARAMS_TO_OFFSET == "all_pose_and_global_rotation":
not_trainable = [0, 2]
#not_trainable = []
#trainable_params_indices = [index for index in range(85) if index not in not_trainable and index < 72]
trainable_params_indices = [index for index in range(85) if index not in not_trainable and index < 66]
PARAMS_TO_OFFSET = [param_ids[index] for index in trainable_params_indices]
POSE_OFFSET = format_distractor_dict(POSE_OFFSET, PARAMS_TO_OFFSET)
# Generate the data from the SMPL parameters
print("loading SMPL")
smpl = SMPLModel('./keras_rotationnet_v2_demo_for_hidde/basicModel_f_lbs_10_207_0_v1.0.0.pkl')
zero_params = np.zeros(shape=(85,))
zero_pc = smpl.set_params(beta=zero_params[72:82], pose=zero_params[0:72].reshape((24,3)), trans=zero_params[82:85])
#print("zero_pc: " + str(zero_pc))
# Generate and format the data
X_indices = np.array([i for i in range(data_samples)])
X_params = np.array([zero_params for i in range(data_samples)], dtype="float32")
if LOAD_DATA_DIR is not None:
X_params, X_pcs = load_data(LOAD_DATA_DIR, num_samples=data_samples, load_silhouettes=False)
else:
if not all(value == 0.0 for value in POSE_OFFSET.values()):
num_test_samples = 5
print("Offsetting parameters")
X_params = offset_params(X_params, PARAMS_TO_OFFSET, POSE_OFFSET)
X_params = X_params[:num_test_samples]
from smpl_np import SMPLModel
from render_mesh import Mesh
parser = ArgumentParser()
parser.add_argument("--model", help="Path to the model weights")
parser.add_argument("--data", help="Path to the test image")
args = parser.parse_args()
einfo = {
'model_weights_path': '../models/model.01-0.07 2019-11-01_13:30:30.hdf5'
}
# # Ensure that TF2.0 is not used
tf.disable_v2_behavior()
tf.enable_eager_execution()
smpl = SMPLModel('../SMPL/model.pkl')
def get_random_sample(smpl):
pose = 0.65 * (np.random.rand(*smpl.pose_shape) - 0.5)
beta = 0.06 * (np.random.rand(*smpl.beta_shape) - 0.5)
trans = np.zeros(smpl.trans_shape)
parameters = np.concatenate([pose.ravel(), beta, trans])
# Create the body mesh
pointcloud = smpl.set_params(beta=beta, pose=pose, trans=trans)
# Render the silhouette
mesh = Mesh(pointcloud=pointcloud)
mesh.faces = smpl.faces
self.asf_to_smpl_joints(translate)
self.smpl_joints_to_mesh()
def imitate(self, motion, translate=False):
"""
A warpper for `set_asf_motion` with a cool name.
Prameter
--------
translate: Whether add global translate to the mesh. Not recommended since
the mesh is likely to move beyond the screen.
"""
self.set_asf_motion(motion, translate)
if __name__ == '__main__':
import reader
import pickle
subject = '01'
im = Imitator(reader.parse_asf('./data/%s/%s.asf' % (subject, subject)),
SMPLModel('./model.pkl'))
sequence = '09'
frame_idx = 0
motions = reader.parse_amc('./data/%s/%s_%s.amc' %
(subject, subject, sequence))
theta = im.motion2theta(motions[frame_idx])
np.save('./pose.npy', theta)
class MainWindow(QMainWindow):
def __init__(self, parent=None):
super(MainWindow, self).__init__(parent=parent)
# GUI
self.ui = Ui_SmplViewer()
self.ui.setupUi(self)
# SMPL object
self.smpl = SMPLModel('./smpl_data/male_model.pkl')
# initialize the sliders
self.poseSliders = []
self.shapeSliders = []
for item in self.ui.__dict__:
if 'horizontalSlider' in item:
self.poseSliders.append(self.ui.__dict__[item])
if 'verticalSlider' in item:
self.shapeSliders.append(self.ui.__dict__[item])
# set events
for i in range(len(self.poseSliders)):
self.poseSliders[i].setMinimum(-8)
self.poseSliders[i].setMaximum(8)
self.poseSliders[i].valueChanged[int].connect(self.changevalue)
for i in range(len(self.shapeSliders)):
self.shapeSliders[i].setMinimum(0)
self.shapeSliders[i].setMaximum(10)
self.shapeSliders[i].valueChanged[int].connect(self.changevalue)
self.ui.pushButton.clicked.connect(self.pushButton_Click)
def loadScene(self):
glClear(GL_COLOR_BUFFER_BIT | GL_DEPTH_BUFFER_BIT)
glMatrixMode(GL_PROJECTION)
glLoadIdentity()
x, y, width, height = glGetDoublev(GL_VIEWPORT)
gluPerspective(
45, # field of view in degrees
width / float(height or 1), # aspect ratio
.25, # near clipping plane
200, # far clipping plane
)
glMatrixMode(GL_MODELVIEW)
glLoadIdentity()
gluLookAt(1.8, 1.8, 1.8, 0, 0, 0, 0, 1, 0)
def paintGL(self):
self.loadScene()
for face in self.smpl.faces:
glColor3f(1, 1, 0.4)
glBegin(GL_TRIANGLES)
x_0, y_0, z_0 = self.smpl.verts[face[0]]
x_1, y_1, z_1 = self.smpl.verts[face[1]]
x_2, y_2, z_2 = self.smpl.verts[face[2]]
glVertex3f(x_0, y_0, z_0)
glVertex3f(x_1, y_1, z_1)
glVertex3f(x_2, y_2, z_2)
glEnd()
self.setupRC()
def setupViewer(self):
self.ui.openGLWidget.initializeGL()
self.ui.openGLWidget.paintGL = self.paintGL
timer = QTimer(self)
timer.timeout.connect(self.ui.openGLWidget.update)
timer.start(5)
def setupRC(self):
# Light values and coordinates
ambientLight = [0.4, 0.4, 0.4, 1.0]
diffuseLight = [0.7, 0.7, 0.7, 1.0]
specular = [0.9, 0.9, 0.9, 1.0]
lightPos = [-0.03, 0.15, 0.1, 1.0]
specref = [0.6, 0.6, 0.6, 1.0]
glEnable(GL_DEPTH_TEST) # Hidden surface removal
glEnable(GL_CULL_FACE) # Do not calculate inside of solid object
glFrontFace(GL_CCW)
glEnable(GL_LIGHTING)
# Setup light 0
glLightModelfv(GL_LIGHT_MODEL_AMBIENT, ambientLight)
glLightfv(GL_LIGHT0, GL_AMBIENT, ambientLight)
glLightfv(GL_LIGHT0, GL_DIFFUSE, diffuseLight)
glLightfv(GL_LIGHT0, GL_SPECULAR, specular)
# Position and turn on the light
glLightfv(GL_LIGHT0, GL_POSITION, lightPos)
glEnable(GL_LIGHT0)
# Enable color tracking
glEnable(GL_COLOR_MATERIAL)
# Set Material properties to follow glColor values
glColorMaterial(GL_FRONT, GL_AMBIENT_AND_DIFFUSE)
# All materials hereafter have full specular reflectivity with a moderate shine
glMaterialfv(GL_FRONT, GL_SPECULAR, specref)
glMateriali(GL_FRONT, GL_SHININESS, 64)
def changevalue(self, value):
pose_vals = []
shape_vals = []
for i in range(24):
vector = np.zeros(3)
vector[0] = 3 * self.poseSliders[3 * i + 0].value() / 8
vector[1] = 3 * self.poseSliders[3 * i + 1].value() / 8
vector[2] = 3 * self.poseSliders[3 * i + 2].value() / 8
pose_vals.append(vector)
for item in self.shapeSliders:
shape_vals.append(item.value())
pose_vals = np.array(pose_vals)
shape_vals = np.array(shape_vals)
trans_vals = np.zeros(3)
self.smpl.set_params(pose=pose_vals, beta=shape_vals, trans=trans_vals)
# save the model
def pushButton_Click(self):
mesh = trimesh.Trimesh(vertices=self.smpl.verts,
faces=self.smpl.faces,
process=False)
output_path = self.ui.lineEdit.text()
mesh.export(output_path)
print('Saved to %s' % output_path)
sys.path.append(os.path.dirname(os.path.abspath(__file__)) + '/..')
from IO import *
from values import *
"""
Precomputes a blend weights prior for each sample based on proximity to SMPL in canonical pose.
"""
# CLOTH3D rest pose (star)
rest_pose = np.zeros((24, 3))
rest_pose[0, 0] = np.pi / 2
rest_pose[1, 2] = .15
rest_pose[2, 2] = -.15
smpl_path = os.path.dirname(os.path.abspath(__file__)) + '/../Model/smpl/'
SMPL = {
0: SMPLModel(smpl_path + 'model_f.pkl', rest_pose),
1: SMPLModel(smpl_path + 'model_m.pkl', rest_pose)
}
# get sample list
samples = os.listdir(SRC)
N = len(samples)
# for each sample
for i, sample in enumerate(samples):
print(str(i + 1) + " / " + str(N))
src = SRC + sample + '/'
src_pre = SRC_PREPROCESS + sample + '/'
dst = src_pre + 'weights.npy'
if os.path.isfile(dst): continue
# load info
info = loadInfo(src + 'info.mat')
import matplotlib
matplotlib.use("TkAgg")
import matplotlib.pyplot as plt
import numpy as np
from scipy.ndimage.morphology import binary_closing
from smpl_np import SMPLModel
from render_mesh import Mesh
# Generate the silhouettes from the SMPL parameters
smpl = SMPLModel(
'./keras_rotationnet_v2_demo_for_hidde/./basicModel_f_lbs_10_207_0_v1.0.0.pkl'
)
# Artificial sample data
sample_pose = 0.0 * (np.random.rand(smpl.pose_shape[0], smpl.pose_shape[1]) -
0.5)
sample_beta = 0.0 * (np.random.rand(smpl.beta_shape[0]) - 0.5)
sample_trans = np.zeros(smpl.trans_shape[0])
#sample_trans = 0.1 * (np.random.rand(smpl.trans_shape[0]) - 0.5)
sample_y = np.array([sample_pose.ravel(), sample_beta, sample_trans])
sample_pc = smpl.set_params(sample_pose, sample_beta, sample_trans)
#pc_subset = sample_pc[5020:5150]
#pc_subset = sample_pc[5400:5600]
#pc_subset = sample_pc[5020:5600]
#pc_subset = sample_pc[1850:2100]
pc_subset = sample_pc[[1850, 1600, 2050, 5350, 5050, 5500]]
def __init__(self, num_samples, reset_period, POSE_OFFSET, PARAMS_TO_OFFSET, ARCHITECTURE, batch_size=32, smpl=None, shuffle=True, save_path="./", num_trainable=24, trainable_params_mask=np.zeros((85,)), kin_tree=None, train_period=5, dist="uniform"):
self.num_samples = num_samples
self.reset_period = reset_period
if isinstance(POSE_OFFSET, (int, float)):
k = {param: POSE_OFFSET for param in PARAMS_TO_OFFSET}
else:
# k must be a dict with an entry for each variable parameter
k = POSE_OFFSET
self.k = k
self.PARAMS_TO_OFFSET = PARAMS_TO_OFFSET
recognised_dists = ["uniform", "gaussian", "normal"]
assert dist in recognised_dists, "distribution not implemented"
self.dist = dist
self.ARCHITECTURE = ARCHITECTURE
self.batch_size = batch_size
if smpl is None:
smpl = SMPLModel('./keras_rotationnet_v2_demo_for_hidde/basicModel_f_lbs_10_207_0_v1.0.0.pkl')
self.smpl = smpl
self.shuffle = shuffle
self.params, self.pcs = gen_data(POSE_OFFSET, PARAMS_TO_OFFSET, self.smpl, data_samples=num_samples, save_dir=None, render_silhouette=False, dist=dist)
self.indices = np.array([i for i in range(num_samples)])
self.save_path = save_path
self.num_trainable = num_trainable
self.trainable_params_mask = trainable_params_mask
self.kin_tree = []
for level in kin_tree:
level_params = []
for param in level:
level_params.append(int(param.replace("param_", "")))
self.kin_tree.append(level_params)
self.curr_kin_level = 0
self.params_ordered = [i for i in range(85)]
self.train_period = train_period
if self.train_period > 0:
self.params_to_train = [param in self.kin_tree[self.curr_kin_level] for param in self.params_ordered]
else:
self.params_to_train = [1 for _ in range(85)]
#print("generator params shape at init: " + str(self.params.shape))
#print("generator pcs shape at init: " + str(self.pcs.shape))
print("generator params first entry: " + str(self.params[0]))
print("generator params final entry: " + str(self.params[-1]))
# Hard-coded parameters
self.epoch = 0
self.num_examples = 5
self.cb_samples = np.linspace(0, self.num_samples, num=self.num_examples, dtype=int)
self.cb_samples[-1] -= 1
print("samples: " + str(self.cb_samples))
# Store initial data
if self.save_path is not None:
sample_path = self.save_path + "cb_samples_E{}.npz".format(self.epoch - 1)
with open(sample_path, 'w') as f:
if self.ARCHITECTURE == "NewDeepConv1DOptLearnerArchitecture":
mask_to_save = np.tile(self.trainable_params_mask, (len(self.cb_samples), 1))
np.savez(f, indices=self.cb_samples, params=self.params[self.cb_samples], pcs=self.pcs[self.cb_samples],\
trainable_params=mask_to_save)
elif self.ARCHITECTURE == "PeriodicOptLearnerArchitecture":
save_params_to_train = np.tile(self.params_to_train, (len(self.cb_samples), 1))
np.savez(f, indices=self.cb_samples, params=self.params[self.cb_samples], pcs=self.pcs[self.cb_samples],\
params_to_train=save_params_to_train)
else:
np.savez(f, indices=self.cb_samples, params=self.params[self.cb_samples], pcs=self.pcs[self.cb_samples])
from imitator import Imitator
from smpl_np import SMPLModel
from tqdm import tqdm
import reader
import os
import numpy as np
if __name__ == '__main__':
target_folder = './pose'
lv0 = './data'
lv1s = os.listdir(lv0)
for lv1 in tqdm(lv1s, ncols=120):
lv2s = os.listdir('/'.join([lv0, lv1]))
asf_path = '%s/%s/%s.asf' % (lv0, lv1, lv1)
joints = reader.parse_asf(asf_path)
im = Imitator(joints, SMPLModel('./model.pkl'))
for lv2 in tqdm(lv2s):
pose = []
if lv2.split('.')[-1] != 'amc':
continue
amc_path = '%s/%s/%s' % (lv0, lv1, lv2)
save_path = '%s/%s' % (target_folder, lv2.split('.')[0] + '.npy')
motions = reader.parse_amc(amc_path)
for idx, motion in enumerate(motions):
pose.append(im.motion2theta(motions[idx]))
np.save(save_path, np.array(pose))
im.imitate(motions[0])
im.smpl.output_mesh('.'.join(save_path.split('.')[:-1]) + '.obj')
train_dir = "../data/train/"
val_dir = "../data/val/"
test_dir = "../data/test/"
# Store the batch size and number of epochs
batch_size = params["GENERATOR"]["BATCH_SIZE"]
epochs = params["MODEL"]["EPOCHS"]
steps_per_epoch = params["MODEL"]["STEPS_PER_EPOCH"]
validation_steps = params["MODEL"]["VALIDATION_STEPS"]
save_period = params["MODEL"]["SAVE_PERIOD"]
pred_period = params["MODEL"]["PRED_PERIOD"]
""" Data generation """
# Generate the silhouettes from the SMPL parameters
smpl = SMPLModel('./keras_rotationnet_v2_demo_for_hidde/./basicModel_f_lbs_10_207_0_v1.0.0.pkl')
# Load the SMPL data
train_gen = OptLearnerDataGenerator(train_dir, batch_size=batch_size, debug=False)
val_gen = OptLearnerDataGenerator(train_dir, batch_size=batch_size) # use same data as for training, but vary the position of the arm by a different amount
test_gen = OptLearnerDataGenerator(test_dir, batch_size=batch_size)
# Get a sample train input for the prediction callback
train_ids = os.listdir(train_dir)[:2]
train_sample_indices = []
train_sample_params = []
train_sample_pcs = []
train_sample_silh = []
for i, sample in enumerate(train_ids):
with open(os.path.join(train_dir, sample), 'r') as handle:
data_dict = pickle.load(handle)
if PARAMS_TO_OFFSET == "trainable_params":
PARAMS_TO_OFFSET = trainable_params
elif PARAMS_TO_OFFSET == "all_pose_and_global_rotation":
not_trainable = [0, 2]
#not_trainable = []
#trainable_params_indices = [index for index in range(85) if index not in not_trainable and index < 72]
trainable_params_indices = [
index for index in range(85)
if index not in not_trainable and index < 66
]
PARAMS_TO_OFFSET = [param_ids[index] for index in trainable_params_indices]
POSE_OFFSET = format_distractor_dict(POSE_OFFSET, PARAMS_TO_OFFSET)
# Generate the data from the SMPL parameters
smpl = SMPLModel(
'./keras_rotationnet_v2_demo_for_hidde/basicModel_f_lbs_10_207_0_v1.0.0.pkl'
)
zero_params = np.zeros(shape=(85, ))
zero_pc = smpl.set_params(beta=zero_params[72:82],
pose=zero_params[0:72].reshape((24, 3)),
trans=zero_params[82:85])
#print("zero_pc: " + str(zero_pc))
if USE_GENERATOR:
update_generator = OptLearnerUpdateGenerator(data_samples,
RESET_PERIOD,
POSE_OFFSET,
PARAMS_TO_OFFSET,
batch_size=BATCH_SIZE,
smpl=smpl,
shuffle=True)
def setup_train_data(trainable_params,
DISTRACTOR,
PARAMS_TO_OFFSET,
POSE_OFFSET,
ARCHITECTURE,
data_samples,
BATCH_SIZE,
train_vis_dir,
num_cb_samples=5,
MODE="RODRIGUES",
LOAD_DATA_DIR=None,
use_generator=False,
RESET_PERIOD=None,
groups=None,
TRAIN_PERIOD=5):
# Gather trainable params
param_ids = ["param_{:02d}".format(i) for i in range(85)]
trainable_params = format_offsetable_params(trainable_params)
param_trainable = {
param: (param in trainable_params)
for param in param_ids
}
DISTRACTOR = format_distractor_dict(DISTRACTOR, trainable_params)
# Gather offsetable params
if PARAMS_TO_OFFSET == "trainable_params":
PARAMS_TO_OFFSET = trainable_params
PARAMS_TO_OFFSET = format_offsetable_params(PARAMS_TO_OFFSET)
POSE_OFFSET = format_distractor_dict(POSE_OFFSET, PARAMS_TO_OFFSET)
# Define the kinematic tree
kin_tree = format_joint_levels(groups)
#print(kin_tree)
efficient_kin_tree = [[
param for param in level if param in trainable_params
] for level in kin_tree]
efficient_kin_tree = [entry for entry in efficient_kin_tree
if entry] # filter empty entries
#print(efficient_kin_tree)
# Generate the data from the SMPL parameters
print("loading SMPL...")
smpl = SMPLModel(
'/data/cvfs/hjhb2/projects/deep_optimiser/code/keras_rotationnet_v2_demo_for_hidde/basicModel_f_lbs_10_207_0_v1.0.0.pkl'
)
# Generate and format the data
print("Gather input data...")
if use_generator:
trainable_params_mask = [
int(param_trainable[key])
for key in sorted(param_trainable.keys(),
key=lambda x: int(x[6:8]))
]
update_generator = OptLearnerUpdateGenerator(
data_samples,
RESET_PERIOD,
POSE_OFFSET,
PARAMS_TO_OFFSET,
ARCHITECTURE,
batch_size=BATCH_SIZE,
smpl=smpl,
shuffle=True,
save_path=train_vis_dir,
trainable_params_mask=trainable_params_mask,
kin_tree=efficient_kin_tree,
train_period=TRAIN_PERIOD)
X_train, Y_train = update_generator.yield_data()
print("Y_train shapes: " + str([datum.shape for datum in Y_train]))
else:
update_generator = None
X_train, Y_train = gather_input_data(data_samples,
smpl,
PARAMS_TO_OFFSET,
POSE_OFFSET,
ARCHITECTURE,
param_trainable,
num_test_samples=-1,
MODE=MODE,
LOAD_DATA_DIR=LOAD_DATA_DIR,
use_generator=use_generator,
kin_tree=efficient_kin_tree)
# Render silhouettes for the callback data
X_cb, Y_cb, silh_cb = gather_cb_data(X_train,
Y_train,
data_samples,
num_cb_samples,
where="spread")
return X_train, Y_train, X_cb, Y_cb, silh_cb, smpl, trainable_params, param_trainable, DISTRACTOR, POSE_OFFSET, update_generator
