def main():
set_cuda_visible_devices([get_freer_gpu()])
#
# model_file, dataset_name = get_DD_model_file()
dataset_name = "box_push_1000_top_down"
dataset_name = "box_push_1000_angled"
dataset_paths = exp_utils.get_dataset_paths(dataset_name)
dataset_root = dataset_paths['dataset_root']
model_file = dataset_paths['dense_descriptor_model_chkpt']
multi_episode_dict = DCDrakeSimEpisodeReader.load_dataset(
dataset_root, max_num_episodes=None)
camera_list = [dataset_paths['main_camera_name']]
target_camera_names = [dataset_paths['main_camera_name']]
# camera_list = ['camera_1_top_down']
# target_camera_names = ['camera_1_top_down']
model_config_file = os.path.join(os.path.dirname(model_file),
'config.yaml')
model_config = getDictFromYamlFilename(model_config_file)
model_config['dataset']['data_augmentation'] = False
dataset = DynamicDrakeSimDataset(
model_config, multi_episode_dict,
phase="valid") # could also use train data
model = torch.load(model_file)
model = model.cuda()
model = model.eval()
heatmap_vis = HeatmapVisualization(model_config,
dataset,
model,
visualize_3D=False,
camera_list=camera_list,
target_camera_names=target_camera_names,
verbose=True,
sample_same_episode=False)
heatmap_vis.run()
from key_dynam.utils.utils import get_project_root, load_yaml, get_current_YYYY_MM_DD_hh_mm_ss_ms, get_data_root, load_pickle
from key_dynam.utils.torch_utils import get_freer_gpu
from key_dynam.dynamics.utils import set_seed
from dense_correspondence_manipulation.utils.utils import set_cuda_visible_devices
from key_dynam.dense_correspondence.dc_drake_sim_episode_reader import DCDrakeSimEpisodeReader
from key_dynam.dense_correspondence.precompute_descriptors import precompute_descriptor_keypoints
from key_dynam.dense_correspondence.keypoint_selection import score_and_select_spatially_separated_keypoints
from key_dynam.experiments.exp_18_box_on_side import utils as exp_18_utils
from key_dynam.experiments.exp_20_mugs import utils as exp_20_utils
# CUDA_VISIBLE_DEVICES = [0]
# set_cuda_visible_devices(CUDA_VISIBLE_DEVICES)
set_cuda_visible_devices([get_freer_gpu()])
from dense_correspondence_manipulation.utils.visualization import draw_reticles
def get_DD_model_file(dataset_name):
model_name = None
model_file = None
if dataset_name == "dps_box_on_side_600":
model_name = ""
model_file = "/home/manuelli/data/key_dynam/dev/experiments/drake_pusher_slider_box_on_side/dataset_dps_box_on_side_600/trained_models/perception/dense_descriptor/3D_loss_camera_angled_2020-05-13-23-39-35-818188/net_best_dy_model.pth"
elif dataset_name == "correlle_mug-small_single_color_600":
model_name = ""
model_file = "/home/manuelli/data/key_dynam/dev/experiments/20/dataset_correlle_mug-small_single_color_600/trained_models/perception/dense_descriptors/data_aug_2020-06-03-00-27-50-738970/net_best_model.pth"
elif dataset_name == "correlle_mug-small_many_colors_600":