def on_epoch_end(self, epoch, logs=None):
""" Store the model loss and accuracy at the end of every epoch, and store a model prediction on data """
self.epoch_log.write(
json.dumps({
'epoch': epoch,
'loss': logs['loss']
}) + '\n')
if (epoch + 1) % self.period == 0 or epoch == 0:
# Predict on all of the given silhouettes
for data_type, data in self.pred_data.items():
if data is not None:
if not isinstance(data, list) or type(data) == np.array:
data = np.array(data)
data = data.reshape(
(1, data.shape[0], data.shape[1], data.shape[2]))
#for i, silhouette in enumerate(data):
# # Save silhouettes
# silhouette *= 255
# cv2.imwrite(os.path.join(self.pred_path, "{}_epoch.{:03d}.gt_silh_{:03d}.png".format(data_type, epoch + 1, i)), silhouette.astype("uint8"))
preds = self.model.predict(data)
#print("Predictions: " + str(preds))
for i, pred in enumerate(preds[1], 1):
#self.smpl.set_params(pred[:72].reshape((24, 3)), pred[72:82], pred[82:])
#self.smpl.save_to_obj(os.path.join(self.pred_path, "{}_pred_{:03d}.obj".format(data_type, i)))
#print_mesh(os.path.join(self.pred_path, "epoch.{:03d}.{}_gt_{:03d}.obj".format(epoch, data_type, i)), gt[i-1], smpl.faces)
print_mesh(
os.path.join(
self.pred_path,
"{}_epoch.{:03d}.pred_{:03d}.obj".format(
data_type, epoch + 1, i)), pred,
self.smpl.faces)
# Store predicted silhouette and the difference between it and the GT silhouette
gt_silhouette = (data[i - 1] *
255).astype("uint8").reshape(
data.shape[1], data.shape[2])
#cv2.imwrite(os.path.join(self.pred_path, "{}_epoch.{:03d}.gt_silh_{:03d}.png".format(data_type, epoch + 1, i)), gt_silhouette)
pred_silhouette = Mesh(
pointcloud=pred).render_silhouette(show=False)
#cv2.imwrite(os.path.join(self.pred_path, "{}_epoch.{:03d}.pred_silh_{:03d}.png".format(data_type, epoch + 1, i)), pred_silhouette)
diff_silh = abs(gt_silhouette - pred_silhouette)
#print(diff_silh.shape)
#cv2.imshow("Diff silh", diff_silh)
#cv2.imwrite(os.path.join(self.pred_path, "{}_epoch.{:03d}.diff_silh_{:03d}.png".format(data_type, epoch + 1, i)), diff_silh.astype("uint8"))
silh_comp = np.concatenate(
[gt_silhouette, pred_silhouette, diff_silh])
cv2.imwrite(
os.path.join(
self.pred_path,
"{}_epoch.{:03d}.silh_comp_{:03d}.png".format(
data_type, epoch + 1, i)),
silh_comp.astype("uint8"))
if self.gt_pc[data_type] is not None:
print_mesh(
os.path.join(
self.pred_path,
"{}_epoch.{:03d}.gt_pc_{:03d}.obj".format(
data_type, epoch + 1, i)),
self.gt_pc[data_type], self.smpl.faces)
print_point_clouds(
os.path.join(
self.pred_path,
"{}_epoch.{:03d}.comparison_{:03d}.obj".
format(data_type, epoch + 1, i)),
[pred, self.gt_pc[data_type]], [(255, 0, 0),
(0, 255, 0)])
if self.visualise:
# Show a random sample
rand_index = np.random.randint(low=0,
high=len(data)) + 1
mesh = Mesh(filepath=os.path.join(
self.pred_path, "{}_epoch.{:03d}.pred_{:03d}.obj".
format(data_type, epoch + 1, rand_index)))
# Show the true silhouette
true_silh = data[rand_index - 1]
true_silh = true_silh.reshape(true_silh.shape[:-1])
plt.imshow(true_silh, cmap='gray')
plt.title("True {} silhouette {:03d}".format(
data_type, rand_index))
plt.show()
# Show the predicted silhouette and mesh
mesh.render_silhouette(
title="Predicted {} silhouette {:03d}".format(
data_type, rand_index))
diff_silh = cv2.imread(
"{}_epoch.{:03d}.diff_silh_{:03d}.png".format(
data_type, epoch + 1, rand_index))
cv2.imshow(
"Predicted {} silhouette {:03d}".format(
data_type, rand_index), diff_silh)
try:
mesh.render3D()
except Exception:
pass
np.zeros((data_samples, 85)),
np.zeros((data_samples, 7))
]
x_test = X_data
y_test = Y_data
# Render silhouettes for the callback data
num_samples = 5
cb_indices = X_indices[:num_samples]
cb_params = X_params[:num_samples]
cb_pcs = X_pcs[:num_samples]
X_cb = [np.array(cb_indices), np.array(cb_params), np.array(cb_pcs)]
silh_cb = []
for pc in cb_pcs:
silh = Mesh(pointcloud=pc).render_silhouette(show=False)
silh_cb.append(silh)
# Initialise the embedding layer
def emb_init_weights_np(emb_params, distractor=np.pi):
def emb_init_wrapper(param, offset=False):
def emb_init(shape, dtype="float32"):
""" Initializer for the embedding layer """
emb_params_ = emb_params[:, param]
if offset:
k = distractor
#k = 1.0
offset_ = k * 2 * (np.random.rand(shape[0]) - 0.5)
emb_params_[:] += offset_
dilate = 1
if dilate == 1:
morph_mask = np.array([[0.34, 0.34, 0.34], [0.34, 1.00, 0.34],
[0.34, 0.34, 0.34]])
new_img = binary_closing(shifted_img != 0,
structure=morph_mask,
iterations=1).astype(np.uint8)
new_img *= 255
else:
new_img = shifted_img
return new_img
if __name__ == "__main__":
mesh_dir = "/data/cvfs/hjhb2/projects/deep_optimiser/example_meshes/"
obj_paths = os.listdir(mesh_dir)
for obj_path in obj_paths:
mesh = Mesh(os.path.join(mesh_dir, obj_path))
silh = mesh.render_silhouette(dim=[256, 256], show=True)
normalised_silh = normalise_img(silh, dim=(128, 128))
#plt.imshow(silh_cropped, cmap="gray")
plt.imshow(normalised_silh, cmap="gray")
plt.show()
augmented_silh = augment_image(normalised_silh)
plt.imshow(augmented_silh, cmap="gray")
plt.show()
np.zeros((data_samples, )),
np.zeros((data_samples, 85)),
np.zeros((data_samples, 85)),
np.zeros((data_samples, 85)),
np.zeros((data_samples, 85))
]
# Render silhouettes for the callback data
num_samples = 5
cb_indices = X_indices[:num_samples]
cb_params = X_params[:num_samples]
cb_pcs = X_pcs[:num_samples]
X_cb = [np.array(cb_indices), np.array(cb_params), np.array(cb_pcs)]
silh_cb = []
for pc in cb_pcs:
silh = Mesh(pointcloud=pc).render_silhouette(show=False)
silh_cb.append(silh)
""" Model set-up """
# Callback functions
# Create a model checkpoint after every few epochs
model_save_checkpoint = tf.keras.callbacks.ModelCheckpoint(
model_dir + "model.{epoch:02d}-{loss:.2f}.hdf5",
monitor='loss',
verbose=1,
save_best_only=False,
mode='auto',
period=50,
save_weights_only=True)
# Predict on sample params at the end of every few epochs
def get_pc_and_silh(path):
mesh = Mesh(filepath=path)
pc = mesh.render_silhouette_with_closing(show=False, closing=False)
silh = mesh.render_silhouette_with_closing(show=False, closing=True)
return pc, silh
def __getitem__(self, item):
""" Yield batches of data """
# Load the SMPL model
#smpl = SMPLModel('./keras_rotationnet_v2_demo_for_hidde/./basicModel_f_lbs_10_207_0_v1.0.0.pkl')
smpl = self.smpl
if self.debug:
if self.debug_X is None:
Y_batch_params = []
Y_batch_pc = []
X_batch = []
for i in range(self.batch_size):
# Generate artificial data
pose = 0.65 * (np.random.rand(smpl.pose_shape[0], smpl.pose_shape[1]) - 0.5)
beta = 0.2 * (np.random.rand(smpl.beta_shape[0]) - 0.5)
trans = np.zeros(smpl.trans_shape[0])
#trans = 0.1 * (np.random.rand(smpl.trans_shape[0]) - 0.5)
# Create the body mesh
pointcloud = smpl.set_params(beta=beta, pose=pose, trans=trans)
Y_batch_params.append(np.concatenate([pose.ravel(), beta, trans]))
Y_batch_pc.append(pointcloud)
# Render the silhouette
silhouette = Mesh(pointcloud=pointcloud).render_silhouette(dim=self.img_dim, show=False)
X_batch.append(np.array(silhouette))
# Preprocess the batches and yield them
Y_batch = [np.array(Y_batch_params), np.array(Y_batch_pc)]
X_batch = np.array(X_batch, dtype="float32")
X_batch /= 255
X_batch = X_batch.reshape((X_batch.shape[0], X_batch.shape[1], X_batch.shape[2], 1))
self.debug_X = X_batch
self.debug_Y = Y_batch
else:
X_batch = self.debug_X
Y_batch = self.debug_Y
else:
# Split of random and real data
num_artificial = int(np.round(self.frac_randomised * self.batch_size))
num_real = int(self.batch_size - num_artificial)
# Retrieve a random batch of parameters from the data directory
if num_real > 0:
data = np.array(os.listdir(self.data_dir))
Y_batch_ids = data[np.random.randint(low=0, high=data.shape[0], size=num_real)]
else:
Y_batch_ids = []
Y_batch_params = []
Y_batch_pc = []
X_batch = []
for Y_id in Y_batch_ids:
# Fetch the real data
Y = np.load(os.path.join(self.data_dir, Y_id))
# Add a small amount of noise to the data
Y += np.random.uniform(low=-self.noise, high=self.noise, size=Y.shape)
Y_batch_params.append(Y)
# Now generate the silhouette from the SMPL meshes
# Create the body mesh
pose = Y[:72]
beta = Y[72:82]
trans = Y[82:]
pointcloud = smpl.set_params(pose.reshape((24, 3)), beta, trans)
# Render the silhouette
silhouette = Mesh(pointcloud=pointcloud).render_silhouette(dim=self.img_dim, show=False)
X_batch.append(np.array(silhouette))
for i in range(num_artificial):
# Generate artificial data
pose = 0.65 * (np.random.rand(smpl.pose_shape[0], smpl.pose_shape[1]) - 0.5)
beta = 0.2 * (np.random.rand(smpl.beta_shape[0]) - 0.5)
trans = np.zeros(smpl.trans_shape[0])
#trans = 0.1 * (np.random.rand(smpl.trans_shape[0]) - 0.5)
# Create the body mesh
pointcloud = smpl.set_params(beta=beta, pose=pose, trans=trans)
Y_batch_params.append(np.concatenate([pose.ravel(), beta, trans]))
Y_batch_pc.append(pointcloud)
# Render the silhouette
silhouette = Mesh(pointcloud=pointcloud).render_silhouette(dim=self.img_dim, show=False)
X_batch.append(np.array(silhouette))
# Preprocess the batches and yield them
Y_batch = [np.array(Y_batch_params), np.array(Y_batch_pc)]
X_batch = np.array(X_batch, dtype="float32")
X_batch /= 255
X_batch = X_batch.reshape((X_batch.shape[0], X_batch.shape[1], X_batch.shape[2], 1))
#print("X_batch shape " + str(X_batch.shape))
#print("Y_batch shape " + str(Y_batch.shape))
#X_batch = list(X_batch)
return X_batch, Y_batch
#eval_string = ""
#for i in range(len(eval_log)):
# eval_string += str(optlearner_model.metrics_names[i]) + ": " + str(eval_log[i]) + " "
#print(eval_string)
# Predict the parameters from the silhouette and generate a mesh
for
prediction = optlearner_model.predict(test_sample_input)
for i, pred in enumerate(prediction, 1):
learned_pc = pred[2]
gt_silhoutte = test_sample_silh[i-1]
pred_silhouette = Mesh(pointcloud=learned_pc).render_silhouette(show=False)
diff_silh = abs(gt_silhouette - pred_silhouette)
#print(diff_silh.shape)
#cv2.imshow("Diff silh", diff_silh)
#cv2.imwrite(os.path.join(self.pred_path, "{}_epoch.{:03d}.diff_silh_{:03d}.png".format(data_type, epoch + 1, i)), diff_silh.astype("uint8"))
silh_comp = np.concatenate([gt_silhouette, pred_silhouette, diff_silh])
cv2.imwrite(os.path.join(test_vis_dir, "{}_epoch.{:05d}.silh_comp_{:03d}.png".format(data_type, epoch + 1, i)), silh_comp.astype("uint8"))
#pred_params = prediction[0]
#real_params = Y_test[0]
#
#for pred in pred_params:
# pointcloud = smpl.set_params(pred[10:82], pred[0:10], pred[82:85])
# pred_silhouette = Mesh(pointcloud=pointcloud).render_silhouette()
