Code examples
Code examples on how to use the folding_demonstrations API can be found here.
Installation instructions are available in the readme of the git project.
Scenario 1 - query RGB images from one perspective
from folding_demonstrations.dataset import FoldingDemonstrationDataSet
# Set to the directory where the folding demonstrations dataset is stored
home_dir = '/media/data/folding_data_output'
# Load the data
dataset = FoldingDemonstrationDataSet(home_dir, perspectives: Tuple = ('left'),
rgb=True, depth=False, pose=False, subtask=False, reward=False)
# Iterate over data and query available information
for demonstration in dataset:
for frame in demonstration:
rgb = frame['left']['rgb']
# Do something with the RGB image
Scenario 2 - query wrist pose estimation from folding sub-task only
from folding_demonstrations.dataset import FoldingDemonstrationDataSet
# Set to the directory where the folding demonstrations dataset is stored
home_dir = '/media/data/folding_data_output'
# Load the data
dataset = FoldingDemonstrationDataSet(home_dir, perspectives: Tuple = ('right'),
rgb=False, depth=False, pose=True, subtask=True, reward=False)
minimum_score = 0.5
# Iterate over data and query available information
for demonstration in dataset:
for frame in demonstration:
if frame['subtask'] == 'folding':
pose = frame['right']['pose']
if pose.confidence >= minimum_score:
LWrist_xy = pose['LWrist'][0:2]
RWrist_xy = pose['RWrist'][0:2]
# Do something with the pose data