def _init_viz_action(scenario: ExperimentScenario, example: Dict):
action = {k: example[k][0] for k in fwd_model.action_keys}
pred_0 = index_time_with_metadata(metadata,
example,
fwd_model.state_keys,
t=0)
scenario.plot_action_rviz(pred_0, action)
def debug_rviz(self, input_dict: Dict, debug_info_seq: List[Tuple]):
import pickle
from merrrt_visualization.rviz_animation_controller import RvizSimpleStepper
from link_bot_pycommon.bbox_visualization import grid_to_bbox
from moonshine.moonshine_utils import numpify
from moonshine.indexing import index_dict_of_batched_tensors_tf, index_time_with_metadata
from link_bot_pycommon.grid_utils import environment_to_occupancy_msg, send_occupancy_tf
import numpy as np
from time import time
debug_filename = f'debug_{int(time())}.pkl'
print(f"Saving debug info to {debug_filename}")
with open(debug_filename, "wb") as debug_file:
pickle.dump({'input_dict': input_dict, 'debug_info': debug_info_seq}, debug_file)
stepper = RvizSimpleStepper()
batch_size = input_dict.pop("batch_size").numpy().astype(np.int32)
input_dict.pop("time")
for b in range(batch_size):
example = index_dict_of_batched_tensors_tf(input_dict, b)
for t, debug_info_t in enumerate(debug_info_seq):
state_t, local_env_origin_t, local_env_t, local_voxel_grid_t = debug_info_t
for i, state_component_k_voxel_grid in enumerate(tf.transpose(local_voxel_grid_t, [4, 0, 1, 2, 3])):
raster_dict = {
'env': tf.clip_by_value(state_component_k_voxel_grid[b], 0, 1),
'origin': local_env_origin_t[b].numpy(),
'res': input_dict['res'][b].numpy(),
}
raster_msg = environment_to_occupancy_msg(raster_dict, frame='local_occupancy')
self.raster_debug_pubs[i].publish(raster_msg)
local_env_dict = {
'env': local_env_t[b],
'origin': local_env_origin_t[b].numpy(),
'res': input_dict['res'][b].numpy(),
}
send_occupancy_tf(self.scenario.tf.tf_broadcaster, local_env_dict, frame='local_occupancy')
pred_t = index_time_with_metadata({}, example, self.pred_state_keys, t)
action_t = {k: example[k][t] for k in self.action_keys}
self.scenario.plot_state_rviz(numpify(pred_t), label='predicted', color='#0000ffff')
if action_t is not None:
self.scenario.plot_action_rviz(numpify(pred_t), numpify(action_t), label='action',
color='#0000ffff')
# # Ground-Truth
# true_t = index_time_with_metadata({}, example, self.true_state_keys, t)
# self.scenario.plot_state_rviz(numpify(true_t), label='actual', color='#ff0000ff', scale=1.1)
# label_t = example['is_close'][1]
# self.scenario.plot_is_close(label_t)
bbox_msg = grid_to_bbox(rows=self.local_env_h_rows,
cols=self.local_env_w_cols,
channels=self.local_env_c_channels,
resolution=example['res'].numpy())
bbox_msg.header.frame_id = 'local_occupancy'
self.local_env_bbox_pub.publish(bbox_msg)
stepper.step()
def _classifier_transition_viz_t(scenario: ExperimentScenario,
example: Dict, t: int):
pred_t = index_time_with_metadata(metadata,
example,
predicted_state_keys,
t=t)
scenario.plot_state_rviz(pred_t, label='predicted', color='#0000ffff')
label_t = example['is_close'][t]
scenario.plot_is_close(label_t)
if true_state_keys is not None:
true_t = index_time_with_metadata(metadata,
example,
true_state_keys,
t=t)
scenario.plot_state_rviz(true_t,
label='actual',
color='#ff0000ff',
scale=1.1)
def index_pred_state_time(self, example: Dict, t: int):
return index_time_with_metadata(self.scenario_metadata, example,
self.predicted_state_keys, t)
def _recovery_transition_viz_t(scenario: ExperimentScenario, example: Dict,
t: int):
e_t = index_time_with_metadata(metadata, example, state_keys, t=t)
scenario.plot_state_rviz(e_t, label='', color='#ff0000ff', scale=1.1)