pipeline_conf.yaml

device: cuda

# Valid robot types are franka, ur5, jaco, and sawyer
source_robot: "franka"
target_robot: "ur5"

# The tensorflow dataset name (found in dataset_info.json)
dataset: franka_tiger

# The path to the folder containing the dataset folder
data_dir: data/franka_tiger

# Where to write the full data containing the augmented images
target_dir: data/franka_tiger_aug

# The number of episodes converted & stored in memory before writing to disk
max_episodes_in_memory: 100

# Number of processes to handle the data loading and augmentation
# WARNING: Use 1 for now while the multiprocessing + single process GPU is being developed
n_workers: 1

# Ro-Aug related parameters
robot_mask:
  # Paths, should be good if running from rovi-aug root
  sam_checkpoint_path: "weights/mask/sam_vit_h_4b8939.pth"
  sam_package_path: deps/SAMed_h/SAMed_h

  # Make sure the LORA checkpoint points to the one for the corresponding robot!
  sam_lora_checkpoint_path: "weights/mask/${source_robot}.pth"

  batch_size: 256

  image_input_key: exterior_image_1_left
  mask_output_key: masks
  masked_images_output_key: masked_imgs

robot_to_robot:
  # Paths, should be good if running from rovi-aug root
  controlnet_module_path: deps/r2r/examples/controlnet
  controlnet_checkpoint_folder_path: "weights/r2r/${source_robot}_to_${target_robot}"

  target_robot: ${target_robot}

  batch_size: 200

  masked_images_input_key: masked_imgs
  image_output_key: robot_aug_imgs

video_inpaint:
  # Paths, should be good if running from rovi-aug root
  video_inpaint_checkpoint_path: weights/video-inpaint/E2FGVI-HQ-CVPR22.pth
  video_inpaint_module_path: deps/video-inpaint

  batch_size: 400

  image_input_key: image
  mask_input_key: masks
  image_output_key: inpainted_background

# For merging the Ro-Aug images into a single one
aug_merge:
  robot_aug_imgs_input_key: robot_aug_imgs
  inpainted_background_input_key: inpainted_background
  merged_output_key: merged_robot_aug

# Vi-Aug related parameters
view_augmentation:
  # Paths, should be good if running from rovi-aug root
  # For the sample image, the path just needs to point to any valid img
  sample_img_path: docs/splash_fig.png
  zeronvs_checkpoint_path: weights/zeronvs/zeronvs.ckpt
  zeronvs_config_path: deps/ZeroNVS/zeronvs_config.yaml
  zeronvs_module_path: deps/ZeroNVS

  batch_size: 80
  image_input_key: merged_robot_aug
  image_output_key: rovi_aug_imgs

  # Settings for the view sampling
  sampler:
    # Should the same random rotation + translation be used for all views in a trajectory?
    trajectory_consistent: False
    
    # TODO: Make this general to include specific translation / rotation values
    # Values can be lower, upper, ultra (goes from small translations / rotations to large ones)
    randomization_amount: "lower"