def call(self, dataset_element, training=False, **kwargs):
known = stack_known(dataset_element)
mean, logvar = self.encode(known)
sampled_features = sample_gaussian(mean, logvar)
if self.hparams['use_flow_during_inference']:
sampled_features = self.apply_flow_to_latent_box(sampled_features)
predicted_occ = self.decode(sampled_features, apply_sigmoid=True)
output = {
'predicted_occ': predicted_occ,
'predicted_free': 1 - predicted_occ,
'latent_mean': mean,
'latent_logvar': logvar,
'sampled_latent': sampled_features
}
return output
def train_step(self, train_element):
bb = tf.keras.layers.Flatten()(tf.cast(train_element['bounding_box'],
tf.float32))
gt_latent_box = self.flow.bijector.inverse(bb)
gt_latent_box = tf.stop_gradient(gt_latent_box)
# train_element['gt_latent_box'] = gt_latent_box
with tf.GradientTape() as tape:
# train_outputs = self.call(train_element, training=True)
# train_losses = self.compute_loss(gt_latent_box, train_outputs)
known = stack_known(train_element)
mean, logvar = self.encode(known)
sampled_latent = sample_gaussian(mean, logvar)
corrected_latent = self.replace_true_box(sampled_latent,
gt_latent_box)
if self.hparams['use_flow_during_inference']:
corrected_latent = self.apply_flow_to_latent_box(
corrected_latent)
logits = self.decode(corrected_latent)
mean_f, mean_box = self.split_box(mean)
logvar_f, logvar_box = self.split_box(logvar)
sampled_f, _ = self.split_box(sampled_latent)
vae_loss = compute_vae_loss(sampled_f, mean_f, logvar_f, logits,
train_element['gt_occ'])
box_loss = compute_box_loss(gt_latent_box, mean_box, logvar_box)
train_losses = {
"loss/vae_loss": vae_loss,
"loss/box_loss": box_loss,
"loss": vae_loss + box_loss
}
train_outputs = None
gradient_metrics = self.apply_gradients(tape, train_element,
train_outputs, train_losses)
other_metrics = self.calculate_metrics(train_element, train_outputs)
metrics = {}
metrics.update(train_losses)
metrics.update(gradient_metrics)
metrics.update(other_metrics)
return train_outputs, metrics
def complete_shape(req: CompleteShapeRequest):
print("Complete Shape Request Received")
pssnet = model_runner.model # type: PSSNet
elem = get_elem(test_records, 90)
# Sample random latent tables
# for i in range(100):
# # latent = tf.Variable(tf.random.normal(shape=[1, pssnet.hparams["num_latent_layers"]]))
# latent = tf.Variable(pssnet.sample_latent(elem))
# vg = pssnet.decode(latent, apply_sigmoid=True)
# VG_PUB.publish('predicted_occ', vg)
# # PT_PUB.publish('predicted_occ', vg)
# rospy.sleep(0.1)
# latent = tf.Variable(tf.random.normal(shape=[1, pssnet.hparams["num_latent_layers"]]))
latent = tf.Variable(pssnet.sample_latent(elem))
mean, logvar = pssnet.encode(stack_known(elem))
latent = tf.Variable(pssnet.apply_flow_to_latent_box(sample_gaussian(mean, logvar)))
log_normal_pdf(latent, mean, logvar)
p = lambda v: np.exp(log_normal_pdf(v, mean, logvar).numpy()[0])
known_contact = tf.Variable(tf.zeros((1, 64, 64, 64, 1)))
# known_free = tf.Variable(tf.zeros((1, 64, 64, 64, 1)))
transformed_free = PT_PUB.transformer.transform_to_frame(req.known_free, "object")
known_free = pointcloud2_msg_to_vox(transformed_free, scale=0.05)
# VG_PUB.publish('known_free', known_free)
# time.sleep(3)
for i in range(10):
vg = pssnet.decode(latent, apply_sigmoid=True)
VG_PUB.publish('predicted_occ', vg)
VG_PUB.publish('known_free', known_free)
if i == 0:
rospy.sleep(3)
# PT_PUB.publish('predicted_occ', vg)
pssnet.grad_step_towards_output(latent, known_contact, known_free)
rospy.sleep(0.1)
return CompleteShapeResponse()
def sample(mean, var):
u = np.array([[mean]], dtype=np.float32)
sigma = np.array([[var]], dtype=np.float32)
return sample_gaussian(u, np.log(sigma)).numpy()[0, 0]
def sample_latent_from_mean_and_logvar(self, mean, logvar):
sampled_features = sample_gaussian(mean, logvar)
if self.hparams['use_flow_during_inference']:
sampled_features = self.apply_flow_to_latent_box(sampled_features)
return sampled_features