├── data └── .gitignore ├── assets ├── panda_pc.npy ├── panda_tubes.obj └── panda_gripper.obj ├── docs ├── evaluate_grasp_1.jpg ├── evaluate_grasp_2.jpg ├── evaluate_grasp_3.png ├── evaluate_grasp_4.png ├── evaluate_grasp_5.png ├── evaluate_grasp_6.png ├── view_sequence_1.gif ├── view_sequence_2.gif ├── visualize_pose_1.jpg ├── visualize_pose_2.jpg ├── plot_grasp_curve_1.png ├── render_sequence_1_1.jpg ├── render_sequence_1_2.jpg ├── render_sequence_1_3.jpg ├── render_sequence_1_4.jpg ├── render_sequence_1_5.jpg ├── render_sequence_1_6.jpg ├── render_sequence_2_1.jpg ├── render_sequence_2_2.jpg ├── render_sequence_2_3.jpg ├── render_sequence_2_4.jpg ├── render_sequence_2_5.jpg ├── render_sequence_2_6.jpg ├── render_sequence_3_1.jpg ├── render_sequence_3_2.jpg ├── render_sequence_3_3.jpg ├── render_sequence_3_4.jpg ├── render_sequence_3_5.jpg └── render_sequence_3_6.jpg ├── .gitmodules ├── .gitignore ├── setup.py ├── examples ├── create_dataset.py ├── evaluate_coco.py ├── evaluate_hpe.py ├── evaluate_bop.py ├── view_sequence.py ├── evaluate_grasp.py ├── visualize_grasps.py ├── visualize_pose.py ├── plot_grasp_curve.py ├── render_sequence.py └── all_cvpr2021_results_eval_scripts.sh ├── dex_ycb_toolkit ├── factory.py ├── logging.py ├── layers │ ├── mano_layer.py │ ├── mano_group_layer.py │ ├── ycb_group_layer.py │ └── ycb_layer.py ├── obj.py ├── hpe_eval.py ├── coco_eval.py ├── dex_ycb.py ├── bop_eval.py ├── sequence_loader.py ├── window.py └── grasp_eval.py └── results ├── fetch_cvpr2021_results.sh └── fetch_example_results.sh /data/.gitignore: -------------------------------------------------------------------------------- 1 | * 2 | -------------------------------------------------------------------------------- /assets/panda_pc.npy: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NVlabs/dex-ycb-toolkit/HEAD/assets/panda_pc.npy -------------------------------------------------------------------------------- /docs/evaluate_grasp_1.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NVlabs/dex-ycb-toolkit/HEAD/docs/evaluate_grasp_1.jpg -------------------------------------------------------------------------------- /docs/evaluate_grasp_2.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NVlabs/dex-ycb-toolkit/HEAD/docs/evaluate_grasp_2.jpg -------------------------------------------------------------------------------- /docs/evaluate_grasp_3.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NVlabs/dex-ycb-toolkit/HEAD/docs/evaluate_grasp_3.png -------------------------------------------------------------------------------- /docs/evaluate_grasp_4.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NVlabs/dex-ycb-toolkit/HEAD/docs/evaluate_grasp_4.png -------------------------------------------------------------------------------- /docs/evaluate_grasp_5.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NVlabs/dex-ycb-toolkit/HEAD/docs/evaluate_grasp_5.png -------------------------------------------------------------------------------- /docs/evaluate_grasp_6.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NVlabs/dex-ycb-toolkit/HEAD/docs/evaluate_grasp_6.png -------------------------------------------------------------------------------- /docs/view_sequence_1.gif: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NVlabs/dex-ycb-toolkit/HEAD/docs/view_sequence_1.gif -------------------------------------------------------------------------------- /docs/view_sequence_2.gif: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NVlabs/dex-ycb-toolkit/HEAD/docs/view_sequence_2.gif -------------------------------------------------------------------------------- /docs/visualize_pose_1.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NVlabs/dex-ycb-toolkit/HEAD/docs/visualize_pose_1.jpg -------------------------------------------------------------------------------- /docs/visualize_pose_2.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NVlabs/dex-ycb-toolkit/HEAD/docs/visualize_pose_2.jpg -------------------------------------------------------------------------------- /docs/plot_grasp_curve_1.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NVlabs/dex-ycb-toolkit/HEAD/docs/plot_grasp_curve_1.png -------------------------------------------------------------------------------- /docs/render_sequence_1_1.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NVlabs/dex-ycb-toolkit/HEAD/docs/render_sequence_1_1.jpg -------------------------------------------------------------------------------- /docs/render_sequence_1_2.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NVlabs/dex-ycb-toolkit/HEAD/docs/render_sequence_1_2.jpg -------------------------------------------------------------------------------- /docs/render_sequence_1_3.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NVlabs/dex-ycb-toolkit/HEAD/docs/render_sequence_1_3.jpg -------------------------------------------------------------------------------- /docs/render_sequence_1_4.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NVlabs/dex-ycb-toolkit/HEAD/docs/render_sequence_1_4.jpg -------------------------------------------------------------------------------- /docs/render_sequence_1_5.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NVlabs/dex-ycb-toolkit/HEAD/docs/render_sequence_1_5.jpg -------------------------------------------------------------------------------- /docs/render_sequence_1_6.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NVlabs/dex-ycb-toolkit/HEAD/docs/render_sequence_1_6.jpg -------------------------------------------------------------------------------- /docs/render_sequence_2_1.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NVlabs/dex-ycb-toolkit/HEAD/docs/render_sequence_2_1.jpg -------------------------------------------------------------------------------- /docs/render_sequence_2_2.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NVlabs/dex-ycb-toolkit/HEAD/docs/render_sequence_2_2.jpg -------------------------------------------------------------------------------- /docs/render_sequence_2_3.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NVlabs/dex-ycb-toolkit/HEAD/docs/render_sequence_2_3.jpg -------------------------------------------------------------------------------- /docs/render_sequence_2_4.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NVlabs/dex-ycb-toolkit/HEAD/docs/render_sequence_2_4.jpg -------------------------------------------------------------------------------- /docs/render_sequence_2_5.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NVlabs/dex-ycb-toolkit/HEAD/docs/render_sequence_2_5.jpg -------------------------------------------------------------------------------- /docs/render_sequence_2_6.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NVlabs/dex-ycb-toolkit/HEAD/docs/render_sequence_2_6.jpg -------------------------------------------------------------------------------- /docs/render_sequence_3_1.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NVlabs/dex-ycb-toolkit/HEAD/docs/render_sequence_3_1.jpg -------------------------------------------------------------------------------- /docs/render_sequence_3_2.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NVlabs/dex-ycb-toolkit/HEAD/docs/render_sequence_3_2.jpg -------------------------------------------------------------------------------- /docs/render_sequence_3_3.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NVlabs/dex-ycb-toolkit/HEAD/docs/render_sequence_3_3.jpg -------------------------------------------------------------------------------- /docs/render_sequence_3_4.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NVlabs/dex-ycb-toolkit/HEAD/docs/render_sequence_3_4.jpg -------------------------------------------------------------------------------- /docs/render_sequence_3_5.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NVlabs/dex-ycb-toolkit/HEAD/docs/render_sequence_3_5.jpg -------------------------------------------------------------------------------- /docs/render_sequence_3_6.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NVlabs/dex-ycb-toolkit/HEAD/docs/render_sequence_3_6.jpg -------------------------------------------------------------------------------- /.gitmodules: -------------------------------------------------------------------------------- 1 | [submodule "bop_toolkit"] 2 | path = bop_toolkit 3 | url = git@github.com:ywchao/bop_toolkit.git 4 | branch = dex-ycb-toolkit 5 | [submodule "freihand"] 6 | path = freihand 7 | url = git@github.com:ywchao/freihand.git 8 | branch = dex-ycb-toolkit 9 | [submodule "manopth"] 10 | path = manopth 11 | url = git@github.com:hassony2/manopth.git 12 | branch = master 13 | -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | *.pyc 2 | 3 | dex_ycb_toolkit.egg-info 4 | 5 | results/example_results* 6 | results/anno_coco_*.json 7 | results/anno_hpe_*.pkl 8 | results/anno_grasp_*.pkl 9 | results/coco_eval_*.log 10 | results/bop_eval_*.log 11 | results/hpe_eval_*.log 12 | results/grasp_eval_*.log 13 | results/bop-*_*-test* 14 | results/hpe_curve_* 15 | results/grasp_res_*.json 16 | results/grasp_vis_* 17 | results/grasp_precision_coverage_*.pdf 18 | results/cvpr2021_results* 19 | -------------------------------------------------------------------------------- /setup.py: -------------------------------------------------------------------------------- 1 | # DexYCB Toolkit 2 | # Copyright (C) 2021 NVIDIA Corporation 3 | # Licensed under the GNU General Public License v3.0 [see LICENSE for details] 4 | 5 | from setuptools import find_packages, setup 6 | 7 | setup( 8 | name='dex-ycb-toolkit', 9 | version='1.0', 10 | packages=find_packages(), 11 | install_requires=[ 12 | 'chumpy', 13 | 'numpy', 14 | 'matplotlib', 15 | 'open3d', 16 | 'opencv-python', 17 | 'pycocotools', 18 | 'pyglet', 19 | 'pyrender', 20 | 'python-fcl', 21 | 'pyyaml', 22 | 'scikit-image', 23 | 'tabulate', 24 | 'torch', 25 | ], 26 | ) 27 | -------------------------------------------------------------------------------- /examples/create_dataset.py: -------------------------------------------------------------------------------- 1 | # DexYCB Toolkit 2 | # Copyright (C) 2021 NVIDIA Corporation 3 | # Licensed under the GNU General Public License v3.0 [see LICENSE for details] 4 | 5 | """Example of creating DexYCB datasets.""" 6 | 7 | import json 8 | 9 | from dex_ycb_toolkit.factory import get_dataset 10 | 11 | 12 | def main(): 13 | for setup in ('s0', 's1', 's2', 's3'): 14 | for split in ('train', 'val', 'test'): 15 | name = '{}_{}'.format(setup, split) 16 | print('Dataset name: {}'.format(name)) 17 | 18 | dataset = get_dataset(name) 19 | 20 | print('Dataset size: {}'.format(len(dataset))) 21 | 22 | sample = dataset[999] 23 | print('1000th sample:') 24 | print(json.dumps(sample, indent=4)) 25 | 26 | 27 | if __name__ == '__main__': 28 | main() 29 | -------------------------------------------------------------------------------- /dex_ycb_toolkit/factory.py: -------------------------------------------------------------------------------- 1 | # DexYCB Toolkit 2 | # Copyright (C) 2021 NVIDIA Corporation 3 | # Licensed under the GNU General Public License v3.0 [see LICENSE for details] 4 | 5 | """Factory method for easily getting datasets by name.""" 6 | 7 | from .dex_ycb import DexYCBDataset 8 | 9 | _sets = {} 10 | 11 | for setup in ('s0', 's1', 's2', 's3'): 12 | for split in ('train', 'val', 'test'): 13 | name = '{}_{}'.format(setup, split) 14 | _sets[name] = (lambda setup=setup, split=split: DexYCBDataset(setup, split)) 15 | 16 | 17 | def get_dataset(name): 18 | """Gets a dataset by name. 19 | 20 | Args: 21 | name: Dataset name. E.g., 's0_test'. 22 | 23 | Returns: 24 | A dataset. 25 | 26 | Raises: 27 | KeyError: If name is not supported. 28 | """ 29 | if name not in _sets: 30 | raise KeyError('Unknown dataset name: {}'.format(name)) 31 | return _sets[name]() 32 | -------------------------------------------------------------------------------- /results/fetch_cvpr2021_results.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | 3 | DIR="$( dirname "${BASH_SOURCE[0]}" )" 4 | cd $DIR 5 | 6 | FILE=cvpr2021_results.zip 7 | ID=1r8e5DS4e6rz-v0OQffgK_JoJBXCEBZq3 8 | CHECKSUM=6b1e883d2c134ffe6228fe727d90650a 9 | 10 | if [ -f $FILE ]; then 11 | echo "File already exists. Checking md5..." 12 | os=`uname -s` 13 | if [ "$os" = "Linux" ]; then 14 | checksum=`md5sum $FILE | awk '{ print $1 }'` 15 | elif [ "$os" = "Darwin" ]; then 16 | checksum=`cat $FILE | md5` 17 | fi 18 | if [ "$checksum" = "$CHECKSUM" ]; then 19 | echo "Checksum is correct. No need to download." 20 | exit 0 21 | else 22 | echo "Checksum is incorrect. Need to download again." 23 | fi 24 | fi 25 | 26 | echo "Downloading example results (2.5G)..." 27 | 28 | wget --no-check-certificate -r "https://drive.google.com/uc?export=download&id=$ID" -O $FILE 29 | 30 | echo "Unzipping..." 31 | 32 | unzip $FILE 33 | 34 | echo "Done." 35 | -------------------------------------------------------------------------------- /results/fetch_example_results.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | 3 | DIR="$( dirname "${BASH_SOURCE[0]}" )" 4 | cd $DIR 5 | 6 | FILE=example_results.tar.gz 7 | ID=1b_rS0EHoji9a41c17v2hMQKMnMoM1oC4 8 | CHECKSUM=93853ec3a1cddce2f459ababd5c42a6a 9 | 10 | if [ -f $FILE ]; then 11 | echo "File already exists. Checking md5..." 12 | os=`uname -s` 13 | if [ "$os" = "Linux" ]; then 14 | checksum=`md5sum $FILE | awk '{ print $1 }'` 15 | elif [ "$os" = "Darwin" ]; then 16 | checksum=`cat $FILE | md5` 17 | fi 18 | if [ "$checksum" = "$CHECKSUM" ]; then 19 | echo "Checksum is correct. No need to download." 20 | exit 0 21 | else 22 | echo "Checksum is incorrect. Need to download again." 23 | fi 24 | fi 25 | 26 | echo "Downloading example results (3M)..." 27 | 28 | wget --no-check-certificate -r "https://drive.google.com/uc?export=download&id=$ID" -O $FILE 29 | 30 | echo "Unzipping..." 31 | 32 | tar zxvf $FILE 33 | 34 | echo "Done." 35 | -------------------------------------------------------------------------------- /dex_ycb_toolkit/logging.py: -------------------------------------------------------------------------------- 1 | # DexYCB Toolkit 2 | # Copyright (C) 2021 NVIDIA Corporation 3 | # Licensed under the GNU General Public License v3.0 [see LICENSE for details] 4 | 5 | """Functions for logging.""" 6 | 7 | import logging 8 | import sys 9 | 10 | 11 | def get_logger(log_file): 12 | """Gets a logger given the path to the log file. 13 | 14 | Args: 15 | name: Path to the log file. 16 | 17 | Returns: 18 | A logger. 19 | """ 20 | logger = logging.getLogger() 21 | logger.setLevel(logging.INFO) 22 | 23 | formatter = logging.Formatter('%(asctime)s: %(message)s', '%Y-%m-%d %H:%M:%S') 24 | 25 | stdout_handler = logging.StreamHandler(sys.stdout) 26 | stdout_handler.setLevel(logging.INFO) 27 | stdout_handler.setFormatter(formatter) 28 | logger.addHandler(stdout_handler) 29 | 30 | file_handler = logging.FileHandler(log_file, mode='w') 31 | file_handler.setLevel(logging.INFO) 32 | file_handler.setFormatter(formatter) 33 | logger.addHandler(file_handler) 34 | 35 | return logger 36 | -------------------------------------------------------------------------------- /examples/evaluate_coco.py: -------------------------------------------------------------------------------- 1 | # DexYCB Toolkit 2 | # Copyright (C) 2021 NVIDIA Corporation 3 | # Licensed under the GNU General Public License v3.0 [see LICENSE for details] 4 | 5 | """Example of running COCO evaluation.""" 6 | 7 | import argparse 8 | import os 9 | 10 | from dex_ycb_toolkit.coco_eval import COCOEvaluator 11 | 12 | 13 | def parse_args(): 14 | parser = argparse.ArgumentParser(description='Run COCO evaluation.') 15 | parser.add_argument('--name', help='Dataset name', default=None, type=str) 16 | parser.add_argument('--res_file', 17 | help='Path to result file', 18 | default=None, 19 | type=str) 20 | parser.add_argument('--out_dir', 21 | help='Directory to save eval output', 22 | default=None, 23 | type=str) 24 | args = parser.parse_args() 25 | return args 26 | 27 | 28 | def main(): 29 | args = parse_args() 30 | 31 | if args.name is None and args.res_file is None: 32 | args.name = 's0_test' 33 | args.res_file = os.path.join( 34 | os.path.dirname(__file__), "..", "results", 35 | "example_results_coco_{}.json".format(args.name)) 36 | 37 | coco_eval = COCOEvaluator(args.name) 38 | coco_eval.evaluate(args.res_file, out_dir=args.out_dir) 39 | 40 | 41 | if __name__ == '__main__': 42 | main() 43 | -------------------------------------------------------------------------------- /examples/evaluate_hpe.py: -------------------------------------------------------------------------------- 1 | # DexYCB Toolkit 2 | # Copyright (C) 2021 NVIDIA Corporation 3 | # Licensed under the GNU General Public License v3.0 [see LICENSE for details] 4 | 5 | """Example of running HPE evaluation.""" 6 | 7 | import argparse 8 | import os 9 | 10 | from dex_ycb_toolkit.hpe_eval import HPEEvaluator 11 | 12 | 13 | def parse_args(): 14 | parser = argparse.ArgumentParser(description='Run HPE evaluation.') 15 | parser.add_argument('--name', help='Dataset name', default=None, type=str) 16 | parser.add_argument('--res_file', 17 | help='Path to result file', 18 | default=None, 19 | type=str) 20 | parser.add_argument('--out_dir', 21 | help='Directory to save eval output', 22 | default=None, 23 | type=str) 24 | args = parser.parse_args() 25 | return args 26 | 27 | 28 | def main(): 29 | args = parse_args() 30 | 31 | if args.name is None and args.res_file is None: 32 | args.name = 's0_test' 33 | args.res_file = os.path.join(os.path.dirname(__file__), "..", "results", 34 | "example_results_hpe_{}.txt".format(args.name)) 35 | 36 | hpe_eval = HPEEvaluator(args.name) 37 | hpe_eval.evaluate(args.res_file, out_dir=args.out_dir) 38 | 39 | 40 | if __name__ == '__main__': 41 | main() 42 | -------------------------------------------------------------------------------- /examples/evaluate_bop.py: -------------------------------------------------------------------------------- 1 | # DexYCB Toolkit 2 | # Copyright (C) 2021 NVIDIA Corporation 3 | # Licensed under the GNU General Public License v3.0 [see LICENSE for details] 4 | 5 | """Example of running BOP evaluation.""" 6 | 7 | import argparse 8 | import os 9 | 10 | from dex_ycb_toolkit.bop_eval import BOPEvaluator 11 | 12 | 13 | def parse_args(): 14 | parser = argparse.ArgumentParser(description='Run BOP evaluation.') 15 | parser.add_argument('--name', help='Dataset name', default=None, type=str) 16 | parser.add_argument('--res_file', 17 | help='Path to result file', 18 | default=None, 19 | type=str) 20 | parser.add_argument('--out_dir', 21 | help='Directory to save eval output', 22 | default=None, 23 | type=str) 24 | args = parser.parse_args() 25 | return args 26 | 27 | 28 | def main(): 29 | args = parse_args() 30 | 31 | if args.name is None and args.res_file is None: 32 | args.name = 's0_test' 33 | args.res_file = os.path.join(os.path.dirname(__file__), "..", "results", 34 | "example_results_bop_{}.csv".format(args.name)) 35 | 36 | bop_eval = BOPEvaluator(args.name) 37 | bop_eval.evaluate(args.res_file, out_dir=args.out_dir, renderer_type='python') 38 | 39 | 40 | if __name__ == '__main__': 41 | main() 42 | -------------------------------------------------------------------------------- /examples/view_sequence.py: -------------------------------------------------------------------------------- 1 | # DexYCB Toolkit 2 | # Copyright (C) 2021 NVIDIA Corporation 3 | # Licensed under the GNU General Public License v3.0 [see LICENSE for details] 4 | 5 | """Example of viewing a sequence.""" 6 | 7 | import argparse 8 | import pyglet 9 | 10 | from dex_ycb_toolkit.sequence_loader import SequenceLoader 11 | from dex_ycb_toolkit.window import Window 12 | 13 | 14 | def parse_args(): 15 | parser = argparse.ArgumentParser( 16 | description='View point cloud and ground-truth hand & object poses in 3D.' 17 | ) 18 | parser.add_argument('--name', 19 | help='Name of the sequence', 20 | default=None, 21 | type=str) 22 | parser.add_argument('--device', 23 | help='Device for data loader computation', 24 | default='cuda:0', 25 | type=str) 26 | parser.add_argument('--no-preload', action='store_true', default=False) 27 | args = parser.parse_args() 28 | return args 29 | 30 | 31 | if __name__ == '__main__': 32 | args = parse_args() 33 | 34 | loader = SequenceLoader(args.name, 35 | device=args.device, 36 | preload=(not args.no_preload), 37 | app='viewer') 38 | w = Window(loader) 39 | 40 | def run(dt): 41 | w.update() 42 | 43 | pyglet.clock.schedule(run) 44 | pyglet.app.run() 45 | -------------------------------------------------------------------------------- /dex_ycb_toolkit/layers/mano_layer.py: -------------------------------------------------------------------------------- 1 | # DexYCB Toolkit 2 | # Copyright (C) 2021 NVIDIA Corporation 3 | # Licensed under the GNU General Public License v3.0 [see LICENSE for details] 4 | 5 | """Wrapper layer for manopth ManoLayer.""" 6 | 7 | import torch 8 | 9 | from torch.nn import Module 10 | from manopth.manolayer import ManoLayer 11 | 12 | 13 | class MANOLayer(Module): 14 | """Wrapper layer for manopth ManoLayer.""" 15 | 16 | def __init__(self, side, betas): 17 | """Constructor. 18 | 19 | Args: 20 | side: MANO hand type. 'right' or 'left'. 21 | betas: A numpy array of shape [10] containing the betas. 22 | """ 23 | super(MANOLayer, self).__init__() 24 | 25 | self._side = side 26 | self._betas = betas 27 | self._mano_layer = ManoLayer(flat_hand_mean=False, 28 | ncomps=45, 29 | side=self._side, 30 | mano_root='manopth/mano/models', 31 | use_pca=True) 32 | 33 | b = torch.from_numpy(self._betas).unsqueeze(0) 34 | f = self._mano_layer.th_faces 35 | self.register_buffer('b', b) 36 | self.register_buffer('f', f) 37 | 38 | v = torch.matmul(self._mano_layer.th_shapedirs, self.b.transpose( 39 | 0, 1)).permute(2, 0, 1) + self._mano_layer.th_v_template 40 | r = torch.matmul(self._mano_layer.th_J_regressor[0], v) 41 | self.register_buffer('root_trans', r) 42 | 43 | def forward(self, p, t): 44 | """Forward function. 45 | 46 | Args: 47 | p: A tensor of shape [B, 48] containing the pose. 48 | t: A tensor of shape [B, 3] containing the trans. 49 | 50 | Returns: 51 | v: A tensor of shape [B, 778, 3] containing the vertices. 52 | j: A tensor of shape [B, 21, 3] containing the joints. 53 | """ 54 | v, j = self._mano_layer(p, self.b.expand(p.size(0), -1), t) 55 | v /= 1000 56 | j /= 1000 57 | return v, j 58 | -------------------------------------------------------------------------------- /examples/evaluate_grasp.py: -------------------------------------------------------------------------------- 1 | # DexYCB Toolkit 2 | # Copyright (C) 2021 NVIDIA Corporation 3 | # Licensed under the GNU General Public License v3.0 [see LICENSE for details] 4 | 5 | """Example of running Grasp evaluation.""" 6 | 7 | import argparse 8 | import os 9 | 10 | from dex_ycb_toolkit.grasp_eval import GraspEvaluator 11 | 12 | 13 | def parse_args(): 14 | parser = argparse.ArgumentParser(description='Run grasp evaluation.') 15 | parser.add_argument('--name', help='Dataset name', default=None, type=str) 16 | parser.add_argument('--bop_res_file', 17 | help='Path to BOP result file', 18 | default=None, 19 | type=str) 20 | parser.add_argument('--coco_res_file', 21 | help='Path to COCO result file', 22 | default=None, 23 | type=str) 24 | parser.add_argument('--out_dir', 25 | help='Directory to save eval output', 26 | default=None, 27 | type=str) 28 | parser.add_argument('--visualize', action='store_true', default=False) 29 | args = parser.parse_args() 30 | return args 31 | 32 | 33 | def main(): 34 | args = parse_args() 35 | 36 | if args.name is None and args.coco_res_file is None and args.bop_res_file is None: 37 | args.name = 's0_test' 38 | args.bop_res_file = os.path.join( 39 | os.path.dirname(__file__), "..", "results", 40 | "example_results_bop_{}.csv".format(args.name)) 41 | args.coco_res_file = os.path.join( 42 | os.path.dirname(__file__), "..", "results", 43 | "example_results_coco_{}.json".format(args.name)) 44 | 45 | grasp_eval = GraspEvaluator(args.name) 46 | grasp_eval.evaluate(args.bop_res_file, 47 | args.coco_res_file, 48 | out_dir=args.out_dir, 49 | visualize=args.visualize) 50 | 51 | 52 | if __name__ == '__main__': 53 | main() 54 | -------------------------------------------------------------------------------- /examples/visualize_grasps.py: -------------------------------------------------------------------------------- 1 | # DexYCB Toolkit 2 | # Copyright (C) 2021 NVIDIA Corporation 3 | # Licensed under the GNU General Public License v3.0 [see LICENSE for details] 4 | 5 | """Example of visualizing pre-generating grasps for each YCB object.""" 6 | 7 | import os 8 | import json 9 | import trimesh 10 | import pyrender 11 | import numpy as np 12 | import copy 13 | 14 | from dex_ycb_toolkit.factory import get_dataset 15 | 16 | 17 | def main(): 18 | dataset = get_dataset('s0_train') 19 | 20 | # Load pre-generated grasps for YCB objects. 21 | ycb_grasp_file = os.path.join(os.path.dirname(__file__), "..", "assets", 22 | "ycb_farthest_100_grasps.json") 23 | with open(ycb_grasp_file, 'r') as f: 24 | ycb_grasps = json.load(f) 25 | 26 | # Load simplified panda gripper mesh. 27 | gripper_mesh_file = os.path.join(os.path.dirname(__file__), "..", "assets", 28 | "panda_tubes.obj") 29 | gripper_mesh = trimesh.load(gripper_mesh_file) 30 | 31 | gripper_material = pyrender.MetallicRoughnessMaterial( 32 | alphaMode="BLEND", 33 | doubleSided=True, 34 | baseColorFactor=(0.00, 1.00, 0.04, 1.00), 35 | metallicFactor=0.0) 36 | 37 | # Visualize pre-generated grasps for each YCB object. 38 | for ycb_id, name in dataset.ycb_classes.items(): 39 | if name not in ycb_grasps.keys(): 40 | print('{} does not have pre-generated grasps: skip.'.format(name)) 41 | continue 42 | 43 | scene = pyrender.Scene(ambient_light=np.array([0.5, 0.5, 0.5, 1.0])) 44 | 45 | obj_mesh = trimesh.load(dataset.obj_file[ycb_id]) 46 | scene.add(pyrender.Mesh.from_trimesh(obj_mesh)) 47 | 48 | for grasp in ycb_grasps[name]: 49 | grasp_mesh = copy.deepcopy(gripper_mesh).apply_transform(grasp), 50 | scene.add(pyrender.Mesh.from_trimesh(grasp_mesh, gripper_material)) 51 | 52 | print('Visualizing pre-generated grasps of {}'.format(name)) 53 | print('Close the window to visualize next object.') 54 | 55 | pyrender.Viewer(scene, use_raymond_lighting=True) 56 | 57 | 58 | if __name__ == '__main__': 59 | main() 60 | -------------------------------------------------------------------------------- /dex_ycb_toolkit/layers/mano_group_layer.py: -------------------------------------------------------------------------------- 1 | # DexYCB Toolkit 2 | # Copyright (C) 2021 NVIDIA Corporation 3 | # Licensed under the GNU General Public License v3.0 [see LICENSE for details] 4 | 5 | """Wrapper layer to hold a group of MANOLayers.""" 6 | 7 | import torch 8 | 9 | from torch.nn import Module, ModuleList 10 | 11 | from .mano_layer import MANOLayer 12 | 13 | 14 | class MANOGroupLayer(Module): 15 | """Wrapper layer to hold a group of MANOLayers.""" 16 | 17 | def __init__(self, sides, betas): 18 | """Constructor. 19 | 20 | Args: 21 | sides: A list of MANO sides. 'right' or 'left'. 22 | betas: A list of numpy arrays of shape [10] containing the betas. 23 | """ 24 | super(MANOGroupLayer, self).__init__() 25 | 26 | self._sides = sides 27 | self._betas = betas 28 | self._layers = ModuleList( 29 | [MANOLayer(s, b) for s, b in zip(self._sides, self._betas)]) 30 | self._num_obj = len(self._sides) 31 | 32 | f = [] 33 | for i in range(self._num_obj): 34 | f.append(self._layers[i].f + 778 * i) 35 | f = torch.cat(f) 36 | self.register_buffer('f', f) 37 | 38 | r = torch.cat([l.root_trans for l in self._layers]) 39 | self.register_buffer('root_trans', r) 40 | 41 | @property 42 | def num_obj(self): 43 | return self._num_obj 44 | 45 | def forward(self, p, inds=None): 46 | """Forward function. 47 | 48 | Args: 49 | p: A tensor of shape [B, D] containing the pose vectors. 50 | inds: A list of sub-layer indices. 51 | 52 | Returns: 53 | v: A tensor of shape [B, N, 3] containing the vertices. 54 | j: A tensor of shape [B, J, 3] containing the joints. 55 | """ 56 | if inds is None: 57 | inds = range(self._num_obj) 58 | v = [ 59 | torch.zeros((p.size(0), 0, 3), 60 | dtype=torch.float32, 61 | device=self.f.device) 62 | ] 63 | j = [ 64 | torch.zeros((p.size(0), 0, 3), 65 | dtype=torch.float32, 66 | device=self.f.device) 67 | ] 68 | p, t = self._pose2pt(p) 69 | for i in inds: 70 | y = self._layers[i](p[:, i], t[:, i]) 71 | v.append(y[0]) 72 | j.append(y[1]) 73 | v = torch.cat(v, dim=1) 74 | j = torch.cat(j, dim=1) 75 | return v, j 76 | 77 | def _pose2pt(self, pose): 78 | """Extracts pose and trans from pose vectors. 79 | 80 | Args: 81 | pose: A tensor of shape [B, D] containing the pose vectors. 82 | 83 | Returns: 84 | p: A tensor of shape [B, O, 48] containing the pose. 85 | t: A tensor of shape [B, O, 3] containing the trans. 86 | """ 87 | p = torch.stack( 88 | [pose[:, 51 * i + 0:51 * i + 48] for i in range(self._num_obj)], dim=1) 89 | t = torch.stack( 90 | [pose[:, 51 * i + 48:51 * i + 51] for i in range(self._num_obj)], dim=1) 91 | return p, t 92 | -------------------------------------------------------------------------------- /dex_ycb_toolkit/layers/ycb_group_layer.py: -------------------------------------------------------------------------------- 1 | # DexYCB Toolkit 2 | # Copyright (C) 2021 NVIDIA Corporation 3 | # Licensed under the GNU General Public License v3.0 [see LICENSE for details] 4 | 5 | """Wrapper layer to hold a group of YCBLayers.""" 6 | 7 | import torch 8 | 9 | from torch.nn import Module, ModuleList 10 | 11 | from .ycb_layer import YCBLayer 12 | 13 | 14 | class YCBGroupLayer(Module): 15 | """Wrapper layer to hold a group of YCBLayers.""" 16 | 17 | def __init__(self, ids): 18 | """Constructor. 19 | 20 | Args: 21 | ids: A list of YCB object ids. 22 | """ 23 | super(YCBGroupLayer, self).__init__() 24 | 25 | self._ids = ids 26 | self._layers = ModuleList([YCBLayer(i) for i in self._ids]) 27 | self._num_obj = len(self._ids) 28 | 29 | f = [] 30 | offset = 0 31 | for i in range(self._num_obj): 32 | if i > 0: 33 | offset += self._layers[i - 1].v.size(1) 34 | f.append(self._layers[i].f + offset) 35 | f = torch.cat(f) 36 | self.register_buffer('f', f) 37 | 38 | @property 39 | def num_obj(self): 40 | return self._num_obj 41 | 42 | @property 43 | def obj_file(self): 44 | return [l.obj_file for l in self._layers] 45 | 46 | @property 47 | def count(self): 48 | return [l.f.numel() for l in self._layers] 49 | 50 | @property 51 | def material(self): 52 | return [l.material for l in self._layers] 53 | 54 | @property 55 | def tex_coords(self): 56 | return [l.tex_coords for l in self._layers] 57 | 58 | def forward(self, p, inds=None): 59 | """Forward function. 60 | 61 | Args: 62 | p: A tensor of shape [B, D] containing the pose vectors. 63 | inds: A list of sub-layer indices. 64 | 65 | Returns: 66 | v: A tensor of shape [B, N, 3] containing the transformed vertices. 67 | n: A tensor of shape [B, N, 3] containing the transformed normals. 68 | """ 69 | if inds is None: 70 | inds = range(self._num_obj) 71 | v = [ 72 | torch.zeros((p.size(0), 0, 3), 73 | dtype=torch.float32, 74 | device=self.f.device) 75 | ] 76 | n = [ 77 | torch.zeros((p.size(0), 0, 3), 78 | dtype=torch.float32, 79 | device=self.f.device) 80 | ] 81 | r, t = self._pose2rt(p) 82 | for i in inds: 83 | y = self._layers[i](r[:, i], t[:, i]) 84 | v.append(y[0]) 85 | n.append(y[1]) 86 | v = torch.cat(v, dim=1) 87 | n = torch.cat(n, dim=1) 88 | return v, n 89 | 90 | def _pose2rt(self, pose): 91 | """Extracts rotations and translations from pose vectors. 92 | 93 | Args: 94 | pose: A tensor of shape [B, D] containing the pose vectors. 95 | 96 | Returns: 97 | r: A tensor of shape [B, O, 3] containing the rotation vectors. 98 | t: A tensor of shape [B, O, 3] containing the translations. 99 | """ 100 | r = torch.stack( 101 | [pose[:, 6 * i + 0:6 * i + 3] for i in range(self._num_obj)], dim=1) 102 | t = torch.stack( 103 | [pose[:, 6 * i + 3:6 * i + 6] for i in range(self._num_obj)], dim=1) 104 | return r, t 105 | -------------------------------------------------------------------------------- /examples/visualize_pose.py: -------------------------------------------------------------------------------- 1 | # DexYCB Toolkit 2 | # Copyright (C) 2021 NVIDIA Corporation 3 | # Licensed under the GNU General Public License v3.0 [see LICENSE for details] 4 | 5 | """Example of visualizing object and hand pose of one image sample.""" 6 | 7 | import numpy as np 8 | import pyrender 9 | import trimesh 10 | import torch 11 | import cv2 12 | import matplotlib.pyplot as plt 13 | 14 | from manopth.manolayer import ManoLayer 15 | 16 | from dex_ycb_toolkit.factory import get_dataset 17 | 18 | 19 | def create_scene(sample, obj_file): 20 | """Creates the pyrender scene of an image sample. 21 | 22 | Args: 23 | sample: A dictionary holding an image sample. 24 | obj_file: A dictionary holding the paths to YCB OBJ files. 25 | 26 | Returns: 27 | A pyrender scene object. 28 | """ 29 | # Create pyrender scene. 30 | scene = pyrender.Scene(bg_color=np.array([0.0, 0.0, 0.0, 0.0]), 31 | ambient_light=np.array([1.0, 1.0, 1.0])) 32 | 33 | # Add camera. 34 | fx = sample['intrinsics']['fx'] 35 | fy = sample['intrinsics']['fy'] 36 | cx = sample['intrinsics']['ppx'] 37 | cy = sample['intrinsics']['ppy'] 38 | cam = pyrender.IntrinsicsCamera(fx, fy, cx, cy) 39 | scene.add(cam, pose=np.eye(4)) 40 | 41 | # Load poses. 42 | label = np.load(sample['label_file']) 43 | pose_y = label['pose_y'] 44 | pose_m = label['pose_m'] 45 | 46 | # Load YCB meshes. 47 | mesh_y = [] 48 | for i in sample['ycb_ids']: 49 | mesh = trimesh.load(obj_file[i]) 50 | mesh = pyrender.Mesh.from_trimesh(mesh) 51 | mesh_y.append(mesh) 52 | 53 | # Add YCB meshes. 54 | for o in range(len(pose_y)): 55 | if np.all(pose_y[o] == 0.0): 56 | continue 57 | pose = np.vstack((pose_y[o], np.array([[0, 0, 0, 1]], dtype=np.float32))) 58 | pose[1] *= -1 59 | pose[2] *= -1 60 | node = scene.add(mesh_y[o], pose=pose) 61 | 62 | # Load MANO layer. 63 | mano_layer = ManoLayer(flat_hand_mean=False, 64 | ncomps=45, 65 | side=sample['mano_side'], 66 | mano_root='manopth/mano/models', 67 | use_pca=True) 68 | faces = mano_layer.th_faces.numpy() 69 | betas = torch.tensor(sample['mano_betas'], dtype=torch.float32).unsqueeze(0) 70 | 71 | # Add MANO meshes. 72 | if not np.all(pose_m == 0.0): 73 | pose = torch.from_numpy(pose_m) 74 | vert, _ = mano_layer(pose[:, 0:48], betas, pose[:, 48:51]) 75 | vert /= 1000 76 | vert = vert.view(778, 3) 77 | vert = vert.numpy() 78 | vert[:, 1] *= -1 79 | vert[:, 2] *= -1 80 | mesh = trimesh.Trimesh(vertices=vert, faces=faces) 81 | mesh1 = pyrender.Mesh.from_trimesh(mesh) 82 | mesh1.primitives[0].material.baseColorFactor = [0.7, 0.7, 0.7, 1.0] 83 | mesh2 = pyrender.Mesh.from_trimesh(mesh, wireframe=True) 84 | mesh2.primitives[0].material.baseColorFactor = [0.0, 0.0, 0.0, 1.0] 85 | node1 = scene.add(mesh1) 86 | node2 = scene.add(mesh2) 87 | 88 | return scene 89 | 90 | 91 | def main(): 92 | name = 's0_train' 93 | dataset = get_dataset(name) 94 | 95 | idx = 70 96 | 97 | sample = dataset[idx] 98 | 99 | scene_r = create_scene(sample, dataset.obj_file) 100 | scene_v = create_scene(sample, dataset.obj_file) 101 | 102 | print('Visualizing pose in camera view using pyrender renderer') 103 | 104 | r = pyrender.OffscreenRenderer(viewport_width=dataset.w, 105 | viewport_height=dataset.h) 106 | 107 | im_render, _ = r.render(scene_r) 108 | 109 | im_real = cv2.imread(sample['color_file']) 110 | im_real = im_real[:, :, ::-1] 111 | 112 | im = 0.33 * im_real.astype(np.float32) + 0.67 * im_render.astype(np.float32) 113 | im = im.astype(np.uint8) 114 | 115 | print('Close the window to continue.') 116 | 117 | plt.imshow(im) 118 | plt.tight_layout() 119 | plt.show() 120 | 121 | print('Visualizing pose using pyrender 3D viewer') 122 | 123 | pyrender.Viewer(scene_v) 124 | 125 | 126 | if __name__ == '__main__': 127 | main() 128 | -------------------------------------------------------------------------------- /assets/panda_tubes.obj: -------------------------------------------------------------------------------- 1 | # https://github.com/mikedh/trimesh 2 | v 0.04002507 0.00022252 0.05840000 3 | v 0.04002507 0.00022252 0.11217000 4 | v 0.04100000 -0.00000000 0.05840000 5 | v 0.04100000 -0.00000000 0.11217000 6 | v 0.04100000 -0.00100000 0.05840000 7 | v 0.04100000 -0.00100000 0.11217000 8 | v 0.04197493 0.00022252 0.05840000 9 | v 0.04197493 0.00022252 0.11217000 10 | v 0.04021817 -0.00062349 0.05840000 11 | v 0.04021817 -0.00062349 0.11217000 12 | v 0.04056612 0.00090097 0.05840000 13 | v 0.04056612 0.00090097 0.11217000 14 | v 0.04143388 0.00090097 0.05840000 15 | v 0.04143388 0.00090097 0.11217000 16 | v 0.04178183 -0.00062349 0.05840000 17 | v 0.04178183 -0.00062349 0.11217000 18 | v -0.04002507 -0.00022252 0.05840000 19 | v -0.04002507 -0.00022252 0.11217000 20 | v -0.04100000 0.00000000 0.05840000 21 | v -0.04100000 0.00000000 0.11217000 22 | v -0.04100000 0.00100000 0.05840000 23 | v -0.04100000 0.00100000 0.11217000 24 | v -0.04197493 -0.00022252 0.05840000 25 | v -0.04197493 -0.00022252 0.11217000 26 | v -0.04021817 0.00062349 0.05840000 27 | v -0.04021817 0.00062349 0.11217000 28 | v -0.04056612 -0.00090097 0.05840000 29 | v -0.04056612 -0.00090097 0.11217000 30 | v -0.04143388 -0.00090097 0.05840000 31 | v -0.04143388 -0.00090097 0.11217000 32 | v -0.04178183 0.00062349 0.05840000 33 | v -0.04178183 0.00062349 0.11217000 34 | v 0.00000000 0.00000000 0.00000000 35 | v 0.00000000 0.00000000 0.06600000 36 | v -0.00000000 0.00100000 0.00000000 37 | v -0.00000000 0.00100000 0.06600000 38 | v -0.00097493 -0.00022252 0.00000000 39 | v -0.00097493 -0.00022252 0.06600000 40 | v 0.00078183 0.00062349 0.00000000 41 | v 0.00078183 0.00062349 0.06600000 42 | v -0.04250000 0.00000000 0.05840000 43 | v -0.04250000 -0.00090097 0.05796612 44 | v -0.04250000 -0.00090097 0.05883388 45 | v -0.04250000 -0.00022252 0.05937493 46 | v -0.04250000 -0.00022252 0.05742507 47 | v -0.04250000 0.00062349 0.05761817 48 | v -0.04250000 0.00062349 0.05918183 49 | v -0.04250000 0.00100000 0.05840000 50 | v 0.00043388 -0.00090097 0.00000000 51 | v 0.00043388 -0.00090097 0.06600000 52 | v -0.00043388 -0.00090097 0.00000000 53 | v -0.00043388 -0.00090097 0.06600000 54 | v 0.04250000 0.00000000 0.05840000 55 | v 0.04250000 -0.00090097 0.05796612 56 | v 0.04250000 -0.00090097 0.05883388 57 | v 0.04250000 -0.00022252 0.05937493 58 | v 0.04250000 -0.00022252 0.05742507 59 | v 0.04250000 0.00062349 0.05761817 60 | v 0.04250000 0.00062349 0.05918183 61 | v 0.04250000 0.00100000 0.05840000 62 | v -0.00078183 0.00062349 0.00000000 63 | v -0.00078183 0.00062349 0.06600000 64 | v 0.00097493 -0.00022252 0.00000000 65 | v 0.00097493 -0.00022252 0.06600000 66 | f 3 9 1 67 | f 3 1 11 68 | f 3 11 13 69 | f 3 13 7 70 | f 3 7 15 71 | f 3 15 5 72 | f 3 5 9 73 | f 2 10 4 74 | f 12 2 4 75 | f 14 12 4 76 | f 8 14 4 77 | f 16 8 4 78 | f 6 16 4 79 | f 10 6 4 80 | f 10 2 9 81 | f 9 2 1 82 | f 2 12 1 83 | f 1 12 11 84 | f 12 14 11 85 | f 11 14 13 86 | f 14 8 13 87 | f 13 8 7 88 | f 8 16 7 89 | f 7 16 15 90 | f 6 10 5 91 | f 5 10 9 92 | f 16 6 15 93 | f 15 6 5 94 | f 19 25 17 95 | f 19 17 27 96 | f 19 27 29 97 | f 19 29 23 98 | f 19 23 31 99 | f 19 31 21 100 | f 19 21 25 101 | f 18 26 20 102 | f 28 18 20 103 | f 30 28 20 104 | f 24 30 20 105 | f 32 24 20 106 | f 22 32 20 107 | f 26 22 20 108 | f 26 18 25 109 | f 25 18 17 110 | f 18 28 17 111 | f 17 28 27 112 | f 28 30 27 113 | f 27 30 29 114 | f 30 24 29 115 | f 29 24 23 116 | f 24 32 23 117 | f 23 32 31 118 | f 22 26 21 119 | f 21 26 25 120 | f 32 22 31 121 | f 31 22 21 122 | f 33 39 63 123 | f 33 63 49 124 | f 33 49 51 125 | f 33 51 37 126 | f 33 37 61 127 | f 33 61 35 128 | f 33 35 39 129 | f 64 40 34 130 | f 50 64 34 131 | f 52 50 34 132 | f 38 52 34 133 | f 62 38 34 134 | f 36 62 34 135 | f 40 36 34 136 | f 40 64 39 137 | f 39 64 63 138 | f 64 50 63 139 | f 63 50 49 140 | f 50 52 49 141 | f 49 52 51 142 | f 52 38 51 143 | f 51 38 37 144 | f 38 62 37 145 | f 37 62 61 146 | f 36 40 35 147 | f 35 40 39 148 | f 62 36 61 149 | f 61 36 35 150 | f 41 46 45 151 | f 41 45 42 152 | f 41 42 43 153 | f 41 43 44 154 | f 41 44 47 155 | f 41 47 48 156 | f 41 48 46 157 | f 57 58 53 158 | f 54 57 53 159 | f 55 54 53 160 | f 56 55 53 161 | f 59 56 53 162 | f 60 59 53 163 | f 58 60 53 164 | f 58 57 46 165 | f 46 57 45 166 | f 57 54 45 167 | f 45 54 42 168 | f 54 55 42 169 | f 42 55 43 170 | f 55 56 43 171 | f 43 56 44 172 | f 56 59 44 173 | f 44 59 47 174 | f 60 58 48 175 | f 48 58 46 176 | f 59 60 47 177 | f 47 60 48 -------------------------------------------------------------------------------- /examples/plot_grasp_curve.py: -------------------------------------------------------------------------------- 1 | # DexYCB Toolkit 2 | # Copyright (C) 2021 NVIDIA Corporation 3 | # Licensed under the GNU General Public License v3.0 [see LICENSE for details] 4 | 5 | """Example of plotting grasp precision-coverage curve.""" 6 | 7 | import os 8 | import argparse 9 | import json 10 | import numpy as np 11 | import matplotlib.pyplot as plt 12 | 13 | from tabulate import tabulate 14 | 15 | from dex_ycb_toolkit.grasp_eval import GraspEvaluator 16 | 17 | res_dir = os.path.join(os.path.dirname(__file__), "..", "results") 18 | 19 | methods = [ 20 | { 21 | 'name': 'PoseCNN RGB', 22 | 'path': os.path.join(res_dir, "cvpr2021_results", "grasp_res_{name}_bop_posecnn_{name}_coco_maskrcnn_{name}.json"), 23 | 'linestyle': '-', 24 | 'marker': '.', 25 | }, 26 | { 27 | 'name': 'PoseCNN + Depth', 28 | 'path': os.path.join(res_dir, "cvpr2021_results", "grasp_res_{name}_bop_posecnn_{name}_refined_coco_maskrcnn_{name}.json"), 29 | 'linestyle': '-', 30 | 'marker': 'o', 31 | }, 32 | { 33 | 'name': 'DeepIM RGB', 34 | 'path': os.path.join(res_dir, "cvpr2021_results", "grasp_res_{name}_bop_deepim_{name}_COLOR_coco_maskrcnn_{name}.json"), 35 | 'linestyle': '-', 36 | 'marker': 'v', 37 | }, 38 | { 39 | 'name': 'DeepIM RGB-D', 40 | 'path': os.path.join(res_dir, "cvpr2021_results", "grasp_res_{name}_bop_deepim_{name}_RGBD_coco_maskrcnn_{name}.json"), 41 | 'linestyle': '-', 42 | 'marker': '^', 43 | }, 44 | { 45 | 'name': 'PoseRBPF RGB', 46 | 'path': os.path.join(res_dir, "cvpr2021_results", "grasp_res_{name}_bop_poserbpf_{name}_rgb_coco_maskrcnn_{name}.json"), 47 | 'linestyle': '-', 48 | 'marker': '>', 49 | }, 50 | { 51 | 'name': 'PoseRBPF RGB-D', 52 | 'path': os.path.join(res_dir, "cvpr2021_results", "grasp_res_{name}_bop_poserbpf_{name}_rgbd_coco_maskrcnn_{name}.json"), 53 | 'linestyle': '-', 54 | 'marker': '*', 55 | }, 56 | { 57 | 'name': 'CosyPose RGB', 58 | 'path': os.path.join(res_dir, "cvpr2021_results", "grasp_res_{name}_bop_cosypose_{name}_coco_maskrcnn_{name}.json"), 59 | 'linestyle': '-', 60 | 'marker': 'D', 61 | }, 62 | ] 63 | 64 | 65 | def parse_args(): 66 | parser = argparse.ArgumentParser(description='Plot grasp precision-coverage curve.') 67 | parser.add_argument('--name', help='Dataset name', default='s1_test', type=str) 68 | args = parser.parse_args() 69 | return args 70 | 71 | 72 | def load_grasp_res_file(grasp_res_file): 73 | """Loads a Grasp result file. 74 | 75 | Args: 76 | grasp_res_file: Path to the Grasp result file. 77 | 78 | Returns: 79 | A dictionary holding the loaded Grasp results. 80 | """ 81 | def _convert_keys_to_float(x): 82 | def _try_convert(k): 83 | try: 84 | return float(k) 85 | except ValueError: 86 | return k 87 | return {_try_convert(k): v for k, v in x.items()} 88 | 89 | with open(grasp_res_file, 'r') as f: 90 | results = json.load(f, object_hook=lambda x: _convert_keys_to_float(x)) 91 | 92 | return results 93 | 94 | 95 | def main(): 96 | args = parse_args() 97 | print('Dataset name: {}'.format(args.name)) 98 | 99 | plt.figure() 100 | 101 | for m in methods: 102 | print('Method: {}'.format(m['name']).format(name=args.name)) 103 | 104 | results = load_grasp_res_file(m['path'].format(name=args.name)) 105 | 106 | coverages = [] 107 | precisions = [] 108 | tabular_data = [] 109 | for r in GraspEvaluator.radius: 110 | for a in GraspEvaluator.angles: 111 | for thr in GraspEvaluator.dist_thresholds: 112 | c = np.mean([x['coverage'][r][a][thr] for x in results]) 113 | p = np.mean([x['precision'][r][a][thr] for x in results]) 114 | coverages.append(c) 115 | precisions.append(p) 116 | tabular_data.append([r, a, thr, c, p]) 117 | metrics = [ 118 | 'radius (m)', 'angle (deg)', 'dist th (m)', 'coverage', 'precision' 119 | ] 120 | table = tabulate(tabular_data, 121 | headers=metrics, 122 | tablefmt='pipe', 123 | floatfmt='.4f', 124 | numalign='right') 125 | print('Results: \n' + table) 126 | 127 | plt.plot(coverages, 128 | precisions, 129 | label=m['name'], 130 | linestyle=m['linestyle'], 131 | marker=m['marker']) 132 | 133 | plt.xlabel('Coverage') 134 | plt.ylabel('Precision') 135 | plt.xlim(0, plt.xlim()[1]) 136 | plt.ylim(0, plt.ylim()[1]) 137 | plt.grid() 138 | plt.legend() 139 | plt.tight_layout() 140 | 141 | grasp_curve_file = os.path.join( 142 | res_dir, "grasp_precision_coverage_{name}.pdf".format(name=args.name)) 143 | print('Saving figure to {}'.format(grasp_curve_file)) 144 | plt.savefig(grasp_curve_file) 145 | 146 | plt.show() 147 | 148 | 149 | if __name__ == '__main__': 150 | main() 151 | -------------------------------------------------------------------------------- /dex_ycb_toolkit/layers/ycb_layer.py: -------------------------------------------------------------------------------- 1 | # DexYCB Toolkit 2 | # Copyright (C) 2021 NVIDIA Corporation 3 | # Licensed under the GNU General Public License v3.0 [see LICENSE for details] 4 | 5 | """Layer to transform YCB mesh vertices with SE3 transformation.""" 6 | 7 | import os 8 | import torch 9 | 10 | from torch.nn import Module 11 | 12 | from ..obj import OBJ 13 | 14 | 15 | class YCBLayer(Module): 16 | """Layer to transform YCB mesh vertices with SE3 transformation.""" 17 | 18 | def __init__(self, i): 19 | """Constructor. 20 | 21 | Args: 22 | i: YCB object index. 23 | """ 24 | super(YCBLayer, self).__init__() 25 | 26 | assert 'DEX_YCB_DIR' in os.environ, "environment variable 'DEX_YCB_DIR' is not set" 27 | self._path = os.environ['DEX_YCB_DIR'] + "/models" 28 | self._classes = ('__background__', '002_master_chef_can', '003_cracker_box', 29 | '004_sugar_box', '005_tomato_soup_can', 30 | '006_mustard_bottle', '007_tuna_fish_can', 31 | '008_pudding_box', '009_gelatin_box', 32 | '010_potted_meat_can', '011_banana', '019_pitcher_base', 33 | '021_bleach_cleanser', '024_bowl', '025_mug', 34 | '035_power_drill', '036_wood_block', '037_scissors', 35 | '040_large_marker', '051_large_clamp', 36 | '052_extra_large_clamp', '061_foam_brick') 37 | self._class_name = self._classes[i] 38 | self._obj_file = self._path + '/' + self._class_name + "/textured_simple.obj" 39 | self._obj = OBJ(self._obj_file) 40 | assert len(self._obj.mesh_list) == 1 41 | assert len(self._obj.mesh_list[0].groups) == 1 42 | g = self._obj.mesh_list[0].groups[0] 43 | 44 | self._material = g.material 45 | self._tex_coords = self._obj.t[g.f_t] 46 | 47 | v = torch.from_numpy(self._obj.v).t() 48 | n = torch.from_numpy(self._obj.n).t() 49 | assert (g.f_v == g.f_n).all() 50 | f = torch.from_numpy(g.f_v).view((-1, 3)) 51 | self.register_buffer('v', v) 52 | self.register_buffer('n', n) 53 | self.register_buffer('f', f) 54 | 55 | @property 56 | def obj_file(self): 57 | return self._obj_file 58 | 59 | @property 60 | def material(self): 61 | return self._material 62 | 63 | @property 64 | def tex_coords(self): 65 | return self._tex_coords 66 | 67 | def forward(self, r, t): 68 | """Forward function. 69 | 70 | Args: 71 | r: A tensor of shape [B, 3] containing the rotation in axis-angle. 72 | t: A tensor of shape [B, 3] containing the translation. 73 | 74 | Returns: 75 | v: A tensor of shape [B, N, 3] containing the transformed vertices. 76 | n: A tensor of shape [B, N, 3] containing the transformed normals. 77 | """ 78 | R = rv2dcm(r) 79 | v = torch.matmul(R, self.v).permute(0, 2, 1) + t.unsqueeze(1) 80 | n = torch.matmul(R, self.n).permute(0, 2, 1) 81 | return v, n 82 | 83 | 84 | # https://en.wikipedia.org/wiki/Rodrigues%27_rotation_formula#Matrix_notation 85 | def rv2dcm(rv): 86 | """Converts rotation vectors to direction cosine matrices. 87 | 88 | Args: 89 | rv: A tensor of shape [B, 3] containing the rotation vectors. 90 | 91 | Returns: 92 | A tensor of shape [B, 3, 3] containing the direction cosine matrices. 93 | """ 94 | angle = torch.norm(rv + 1e-8, p=2, dim=1) 95 | axis = rv / angle.unsqueeze(1) 96 | s = torch.sin(angle).unsqueeze(1).unsqueeze(2) 97 | c = torch.cos(angle).unsqueeze(1).unsqueeze(2) 98 | I = torch.eye(3, device=rv.device).expand(rv.size(0), -1, -1) 99 | z = torch.zeros_like(angle) 100 | K = torch.stack( 101 | (torch.stack((z, -axis[:, 2], axis[:, 1]), 102 | dim=1), torch.stack((axis[:, 2], z, -axis[:, 0]), dim=1), 103 | torch.stack((-axis[:, 1], axis[:, 0], z), dim=1)), 104 | dim=1) 105 | dcm = I + s * K + (1 - c) * torch.bmm(K, K) 106 | return dcm 107 | 108 | 109 | # https://en.wikipedia.org/wiki/Rotation_formalisms_in_three_dimensions#Rotation_matrix_%E2%86%94_Euler_axis/angle 110 | # https://github.com/kashif/ceres-solver/blob/087462a90dd1c23ac443501f3314d0fcedaea5f7/include/ceres/rotation.h#L178 111 | # S. Sarabandi and F. Thomas. A Survey on the Computation of Quaternions from Rotation Matrices. J MECH ROBOT, 2019. 112 | # https://en.wikipedia.org/wiki/Rodrigues%27_rotation_formula 113 | def dcm2rv(dcm): 114 | """Converts direction cosine matrices to rotation vectors. 115 | 116 | Args: 117 | dcm: A tensor of shape [B, 3, 3] containing the direction cosine matrices. 118 | 119 | Returns: 120 | A tensor of shape [B, 3] containing the rotation vectors. 121 | """ 122 | X = torch.stack((dcm[:, 2, 1] - dcm[:, 1, 2], dcm[:, 0, 2] - dcm[:, 2, 0], 123 | dcm[:, 1, 0] - dcm[:, 0, 1]), 124 | dim=1) 125 | s = torch.norm(X, p=2, dim=1) / 2 126 | c = (dcm[:, 0, 0] + dcm[:, 1, 1] + dcm[:, 2, 2] - 1) / 2 127 | c = torch.clamp(c, -1, 1) 128 | angle = torch.atan2(s, c) 129 | Y = torch.stack((dcm[:, 0, 0], dcm[:, 1, 1], dcm[:, 2, 2]), dim=1) 130 | Y = torch.sqrt((Y - c.unsqueeze(1)) / (1 - c.unsqueeze(1))) 131 | rv = torch.zeros((dcm.size(0), 3), device=dcm.device) 132 | i1 = s > 1e-3 133 | i2 = (s <= 1e-3) & (c > 0) 134 | i3 = (s <= 1e-3) & (c < 0) 135 | rv[i1] = angle[i1].unsqueeze(1) * X[i1] / (2 * s[i1].unsqueeze(1)) 136 | rv[i2] = X[i2] / 2 137 | rv[i3] = angle[i3].unsqueeze(1) * torch.sign(X[i3]) * Y[i3] 138 | return rv 139 | -------------------------------------------------------------------------------- /dex_ycb_toolkit/obj.py: -------------------------------------------------------------------------------- 1 | # DexYCB Toolkit 2 | # Copyright (C) 2021 NVIDIA Corporation 3 | # Licensed under the GNU General Public License v3.0 [see LICENSE for details] 4 | 5 | """Wavefront OBJ file loader. 6 | 7 | Functions and classes are largely derived from 8 | https://github.com/pyglet/pyglet/blob/f762169c9dd88c22c8d6d2399a129cc23654d99c/contrib/model/model/obj_batch.py 9 | """ 10 | 11 | import os 12 | import logging 13 | import numpy as np 14 | 15 | 16 | class Material: 17 | """Material.""" 18 | diffuse = [.8, .8, .8] 19 | ambient = [.2, .2, .2] 20 | specular = [0., 0., 0.] 21 | emission = [0., 0., 0.] 22 | shininess = 0. 23 | opacity = 1. 24 | texture_path = None 25 | 26 | def __init__(self, name): 27 | """Constructor. 28 | 29 | Args: 30 | name: Material name. 31 | """ 32 | self.name = name 33 | 34 | 35 | class MaterialGroup: 36 | """Material group.""" 37 | 38 | def __init__(self, material): 39 | """Constructor. 40 | 41 | Args: 42 | material: A Material object. 43 | """ 44 | self.material = material 45 | 46 | self.f_v = [] 47 | self.f_n = [] 48 | self.f_t = [] 49 | 50 | 51 | class Mesh: 52 | """Mesh.""" 53 | 54 | def __init__(self, name): 55 | """Constructor. 56 | 57 | Args: 58 | name: Mesh name. 59 | """ 60 | self.name = name 61 | self.groups = [] 62 | 63 | 64 | class OBJ: 65 | """3D data loaded from an OBJ file.""" 66 | 67 | def __init__(self, filename, file=None, path=None): 68 | """Constructor. 69 | 70 | Args: 71 | filename: Path to the OBJ file. 72 | file: An file object. 73 | path: Path to the directory storing the material files. 74 | """ 75 | self.materials = {} 76 | self.meshes = {} 77 | self.mesh_list = [] 78 | 79 | if file is None: 80 | file = open(filename, 'r') 81 | 82 | if path is None: 83 | path = os.path.dirname(filename) 84 | self.path = path 85 | 86 | mesh = None 87 | group = None 88 | material = None 89 | 90 | self.v = [] 91 | self.n = [] 92 | self.t = [] 93 | 94 | for line in file: 95 | if line.startswith('#'): 96 | continue 97 | values = line.split() 98 | if not values: 99 | continue 100 | 101 | if values[0] == 'v': 102 | self.v.append(list(map(float, values[1:4]))) 103 | elif values[0] == 'vn': 104 | self.n.append(list(map(float, values[1:4]))) 105 | elif values[0] == 'vt': 106 | self.t.append(list(map(float, values[1:3]))) 107 | elif values[0] == 'mtllib': 108 | self._load_material_library(values[1]) 109 | elif values[0] in ('usemtl', 'usemat'): 110 | material = self.materials.get(values[1], None) 111 | if material is None: 112 | logging.warn('Unknown material: %s' % values[1]) 113 | if mesh is not None: 114 | group = MaterialGroup(material) 115 | mesh.groups.append(group) 116 | elif values[0] == 'o': 117 | mesh = Mesh(values[1]) 118 | self.meshes[mesh.name] = mesh 119 | self.mesh_list.append(mesh) 120 | group = None 121 | elif values[0] == 'f': 122 | if mesh is None: 123 | mesh = Mesh('') 124 | self.mesh_list.append(mesh) 125 | if material is None: 126 | material = Material("") 127 | if group is None: 128 | group = MaterialGroup(material) 129 | mesh.groups.append(group) 130 | 131 | for i, v in enumerate(values[1:]): 132 | v_index, t_index, n_index = \ 133 | (list(map(int, [j or 0 for j in v.split('/')])) + [0, 0])[:3] 134 | if v_index < 0: 135 | v_index += len(vertices) 136 | if t_index < 0: 137 | t_index += len(tex_coords) 138 | if n_index < 0: 139 | n_index += len(normals) 140 | if i < 3: 141 | group.f_v.append(v_index - 1) 142 | group.f_n.append(n_index - 1) 143 | group.f_t.append(t_index - 1) 144 | else: 145 | # Triangulate. 146 | group.f_v += [group.f_v[-3 * (i - 2)], group.f_v[-1], v_index - 1] 147 | group.f_n += [group.f_n[-3 * (i - 2)], group.f_n[-1], n_index - 1] 148 | group.f_t += [group.f_t[-3 * (i - 2)], group.f_t[-1], t_index - 1] 149 | 150 | self.v = np.array(self.v, dtype=np.float32) 151 | self.n = np.array(self.n, dtype=np.float32) 152 | self.t = np.array(self.t, dtype=np.float32) 153 | 154 | for mesh in self.mesh_list: 155 | for group in mesh.groups: 156 | group.f_v = np.array(group.f_v, dtype=np.int64).reshape(-1, 3) 157 | group.f_n = np.array(group.f_n, dtype=np.int64).reshape(-1, 3) 158 | group.f_t = np.array(group.f_t, dtype=np.int64).reshape(-1, 3) 159 | 160 | def _open_material_file(self, filename): 161 | """Opens a material file. 162 | 163 | Args: 164 | filename: Path to the material file. 165 | 166 | Returns: 167 | A file object. 168 | """ 169 | return open(os.path.join(self.path, filename), 'r') 170 | 171 | def _load_material_library(self, filename): 172 | """Loads the material from a material file. 173 | 174 | Args: 175 | filename: Path to the material file. 176 | """ 177 | material = None 178 | file = self._open_material_file(filename) 179 | 180 | for line in file: 181 | if line.startswith('#'): 182 | continue 183 | values = line.split() 184 | if not values: 185 | continue 186 | 187 | if values[0] == 'newmtl': 188 | material = Material(values[1]) 189 | self.materials[material.name] = material 190 | elif material is None: 191 | logging.warn('Expected "newmtl" in %s' % filename) 192 | continue 193 | 194 | try: 195 | if values[0] == 'Kd': 196 | material.diffuse = list(map(float, values[1:])) 197 | elif values[0] == 'Ka': 198 | material.ambient = list(map(float, values[1:])) 199 | elif values[0] == 'Ks': 200 | material.specular = list(map(float, values[1:])) 201 | elif values[0] == 'Ke': 202 | material.emissive = list(map(float, values[1:])) 203 | elif values[0] == 'Ns': 204 | material.shininess = float(values[1]) 205 | elif values[0] == 'd': 206 | material.opacity = float(values[1]) 207 | elif values[0] == 'map_Kd': 208 | material.texture_path = os.path.abspath(self.path + '/' + values[1]) 209 | except BaseException as ex: 210 | logging.warning('Parse error in %s.' % (filename, ex)) 211 | -------------------------------------------------------------------------------- /dex_ycb_toolkit/hpe_eval.py: -------------------------------------------------------------------------------- 1 | # DexYCB Toolkit 2 | # Copyright (C) 2021 NVIDIA Corporation 3 | # Licensed under the GNU General Public License v3.0 [see LICENSE for details] 4 | 5 | """HPE evaluator.""" 6 | 7 | import os 8 | import sys 9 | import time 10 | import numpy as np 11 | import pickle 12 | 13 | from tabulate import tabulate 14 | 15 | from dex_ycb_toolkit.factory import get_dataset 16 | from dex_ycb_toolkit.logging import get_logger 17 | 18 | freihand_root = os.path.join(os.path.dirname(__file__), "..", "freihand") 19 | sys.path.append(freihand_root) 20 | 21 | from utils.eval_util import EvalUtil 22 | from eval import align_w_scale, curve, createHTML 23 | 24 | _AUC_VAL_MIN = 0.0 25 | _AUC_VAL_MAX = 50.0 26 | _AUC_STEPS = 100 27 | 28 | 29 | class HPEEvaluator(): 30 | """HPE evaluator.""" 31 | 32 | def __init__(self, name): 33 | """Constructor. 34 | 35 | Args: 36 | name: Dataset name. E.g., 's0_test'. 37 | """ 38 | self._name = name 39 | 40 | self._dataset = get_dataset(self._name) 41 | 42 | self._out_dir = os.path.join(os.path.dirname(__file__), "..", "results") 43 | 44 | self._anno_file = os.path.join(self._out_dir, 45 | "anno_hpe_{}.pkl".format(self._name)) 46 | 47 | if os.path.isfile(self._anno_file): 48 | print('Found HPE annotation file.') 49 | else: 50 | print('Cannot find HPE annotation file.') 51 | self._generate_anno_file() 52 | 53 | self._anno = self._load_anno_file() 54 | 55 | def _generate_anno_file(self): 56 | """Generates the annotation file.""" 57 | print('Generating HPE annotation file') 58 | s = time.time() 59 | 60 | joint_3d_gt = {} 61 | 62 | for i in range(len(self._dataset)): 63 | if (i + 1) in np.floor(np.linspace(0, len(self._dataset), 11))[1:]: 64 | print('{:3.0f}% {:6d}/{:6d}'.format(100 * i / len(self._dataset), i, 65 | len(self._dataset))) 66 | 67 | sample = self._dataset[i] 68 | 69 | label = np.load(sample['label_file']) 70 | joint_3d = label['joint_3d'].reshape(21, 3) 71 | 72 | if np.all(joint_3d == -1): 73 | continue 74 | 75 | joint_3d *= 1000 76 | 77 | joint_3d_gt[i] = joint_3d 78 | 79 | print('# total samples: {:6d}'.format(len(self._dataset))) 80 | print('# valid samples: {:6d}'.format(len(joint_3d_gt))) 81 | 82 | anno = { 83 | 'joint_3d': joint_3d_gt, 84 | } 85 | with open(self._anno_file, 'wb') as f: 86 | pickle.dump(anno, f) 87 | 88 | e = time.time() 89 | print('time: {:7.2f}'.format(e - s)) 90 | 91 | def _load_anno_file(self): 92 | """Loads the annotation file. 93 | 94 | Returns: 95 | A dictionary holding the loaded annotation. 96 | """ 97 | with open(self._anno_file, 'rb') as f: 98 | anno = pickle.load(f) 99 | 100 | anno['joint_3d'] = { 101 | k: v.astype(np.float64) for k, v in anno['joint_3d'].items() 102 | } 103 | 104 | return anno 105 | 106 | def _load_results(self, res_file): 107 | """Loads results from a result file. 108 | 109 | Args: 110 | res_file: Path to the result file. 111 | 112 | Returns: 113 | A dictionary holding the loaded results. 114 | 115 | Raises: 116 | ValueError: If a line in the result file does not have 64 comma-seperated 117 | elements. 118 | """ 119 | results = {} 120 | with open(res_file, 'r') as f: 121 | for line in f: 122 | elems = line.split(',') 123 | if len(elems) != 64: 124 | raise ValueError( 125 | 'a line does not have 64 comma-seperated elements: {}'.format( 126 | line)) 127 | image_id = int(elems[0]) 128 | joint_3d = np.array(elems[1:], dtype=np.float64).reshape(21, 3) 129 | results[image_id] = joint_3d 130 | return results 131 | 132 | def evaluate(self, res_file, out_dir=None): 133 | """Evaluates HPE metrics given a result file. 134 | 135 | Args: 136 | res_file: Path to the result file. 137 | out_dir: Path to the output directory. 138 | 139 | Returns: 140 | A dictionary holding the results. 141 | """ 142 | if out_dir is None: 143 | out_dir = self._out_dir 144 | 145 | res_name = os.path.splitext(os.path.basename(res_file))[0] 146 | log_file = os.path.join(out_dir, 147 | "hpe_eval_{}_{}.log".format(self._name, res_name)) 148 | logger = get_logger(log_file) 149 | 150 | res = self._load_results(res_file) 151 | 152 | logger.info('Running evaluation') 153 | 154 | joint_3d_gt = self._anno['joint_3d'] 155 | 156 | eval_util_ab = EvalUtil() 157 | eval_util_rr = EvalUtil() 158 | eval_util_pa = EvalUtil() 159 | 160 | for i, kpt_gt in joint_3d_gt.items(): 161 | assert i in res, "missing image id in result file: {}".format(i) 162 | vis = np.ones_like(kpt_gt[:, 0]) 163 | kpt_pred = res[i] 164 | 165 | eval_util_ab.feed(kpt_gt, vis, kpt_pred) 166 | eval_util_rr.feed(kpt_gt - kpt_gt[0], vis, kpt_pred - kpt_pred[0]) 167 | eval_util_pa.feed(kpt_gt, vis, align_w_scale(kpt_gt, kpt_pred)) 168 | 169 | mean_ab, _, auc_ab, pck_ab, thresh_ab = eval_util_ab.get_measures( 170 | _AUC_VAL_MIN, _AUC_VAL_MAX, _AUC_STEPS) 171 | mean_rr, _, auc_rr, pck_rr, thresh_rr = eval_util_rr.get_measures( 172 | _AUC_VAL_MIN, _AUC_VAL_MAX, _AUC_STEPS) 173 | mean_pa, _, auc_pa, pck_pa, thresh_pa = eval_util_pa.get_measures( 174 | _AUC_VAL_MIN, _AUC_VAL_MAX, _AUC_STEPS) 175 | 176 | tabular_data = [['absolute', mean_ab, auc_ab], 177 | ['root-relative', mean_rr, auc_rr], 178 | ['procrustes', mean_pa, auc_pa]] 179 | metrics = ['alignment', 'MPJPE (mm)', 'AUC'] 180 | table = tabulate(tabular_data, 181 | headers=metrics, 182 | tablefmt='pipe', 183 | floatfmt='.4f', 184 | numalign='right') 185 | logger.info('Results: \n' + table) 186 | 187 | hpe_curve_dir = os.path.join(out_dir, 188 | "hpe_curve_{}_{}".format(self._name, res_name)) 189 | os.makedirs(hpe_curve_dir, exist_ok=True) 190 | 191 | createHTML(hpe_curve_dir, [ 192 | curve(thresh_ab, pck_ab, 'Distance in mm', 193 | 'Percentage of correct keypoints', 194 | 'PCK curve for absolute keypoint error'), 195 | curve(thresh_rr, pck_rr, 'Distance in mm', 196 | 'Percentage of correct keypoints', 197 | 'PCK curve for root-relative keypoint error'), 198 | curve(thresh_pa, pck_pa, 'Distance in mm', 199 | 'Percentage of correct keypoints', 200 | 'PCK curve for Procrustes aligned keypoint error'), 201 | ]) 202 | 203 | results = { 204 | 'absolute': { 205 | 'mpjpe': mean_ab, 206 | 'auc': auc_ab 207 | }, 208 | 'root-relative': { 209 | 'mpjpe': mean_rr, 210 | 'auc': auc_rr 211 | }, 212 | 'procrustes': { 213 | 'mpjpe': mean_pa, 214 | 'auc': auc_pa 215 | }, 216 | } 217 | 218 | logger.info('Evaluation complete.') 219 | 220 | return results 221 | -------------------------------------------------------------------------------- /examples/render_sequence.py: -------------------------------------------------------------------------------- 1 | # DexYCB Toolkit 2 | # Copyright (C) 2021 NVIDIA Corporation 3 | # Licensed under the GNU General Public License v3.0 [see LICENSE for details] 4 | 5 | """Example of rendering a sequence.""" 6 | 7 | import argparse 8 | import torch 9 | import pyrender 10 | import trimesh 11 | import os 12 | import numpy as np 13 | import cv2 14 | 15 | from dex_ycb_toolkit.sequence_loader import SequenceLoader 16 | 17 | _YCB_COLORS = { 18 | 1: (255, 0, 0), # 002_master_chef_can 19 | 2: ( 0, 255, 0), # 003_cracker_box 20 | 3: ( 0, 0, 255), # 004_sugar_box 21 | 4: (255, 255, 0), # 005_tomato_soup_can 22 | 5: (255, 0, 255), # 006_mustard_bottle 23 | 6: ( 0, 255, 255), # 007_tuna_fish_can 24 | 7: (128, 0, 0), # 008_pudding_box 25 | 8: ( 0, 128, 0), # 009_gelatin_box 26 | 9: ( 0, 0, 128), # 010_potted_meat_can 27 | 10: (128, 128, 0), # 011_banana 28 | 11: (128, 0, 128), # 019_pitcher_base 29 | 12: ( 0, 128, 128), # 021_bleach_cleanser 30 | 13: ( 64, 0, 0), # 024_bowl 31 | 14: ( 0, 64, 0), # 025_mug 32 | 15: ( 0, 0, 64), # 035_power_drill 33 | 16: ( 64, 64, 0), # 036_wood_block 34 | 17: ( 64, 0, 64), # 037_scissors 35 | 18: ( 0, 64, 64), # 040_large_marker 36 | 19: (192, 0, 0), # 051_large_clamp 37 | 20: ( 0, 192, 0), # 052_extra_large_clamp 38 | 21: ( 0, 0, 192), # 061_foam_brick 39 | } 40 | _MANO_COLOR = (255, 255, 255) 41 | 42 | 43 | def parse_args(): 44 | parser = argparse.ArgumentParser( 45 | description='Render hand & object poses in camera views.') 46 | parser.add_argument('--name', 47 | help='Name of the sequence', 48 | default=None, 49 | type=str) 50 | parser.add_argument('--device', 51 | help='Device for data loader computation', 52 | default='cuda:0', 53 | type=str) 54 | args = parser.parse_args() 55 | return args 56 | 57 | 58 | class Renderer(): 59 | """Renderer.""" 60 | 61 | def __init__(self, name, device='cuda:0'): 62 | """Constructor. 63 | 64 | Args: 65 | name: Sequence name. 66 | device: A torch.device string argument. The specified device is used only 67 | for certain data loading computations, but not storing the loaded data. 68 | Currently the loaded data is always stored as numpy arrays on cpu. 69 | """ 70 | assert device in ('cuda', 'cpu') or device.split(':')[0] == 'cuda' 71 | self._name = name 72 | self._device = torch.device(device) 73 | 74 | self._loader = SequenceLoader(self._name, 75 | device=device, 76 | preload=False, 77 | app='renderer') 78 | 79 | # Create pyrender cameras. 80 | self._cameras = [] 81 | for c in range(self._loader.num_cameras): 82 | K = self._loader.K[c].cpu().numpy() 83 | fx = K[0][0].item() 84 | fy = K[1][1].item() 85 | cx = K[0][2].item() 86 | cy = K[1][2].item() 87 | cam = pyrender.IntrinsicsCamera(fx, fy, cx, cy) 88 | self._cameras.append(cam) 89 | 90 | # Create meshes for YCB objects. 91 | self._mesh_y = [] 92 | for o in range(self._loader.num_ycb): 93 | obj_file = self._loader.ycb_group_layer.obj_file[o] 94 | mesh = trimesh.load(obj_file) 95 | mesh = pyrender.Mesh.from_trimesh(mesh) 96 | self._mesh_y.append(mesh) 97 | 98 | # Create spheres for MANO joints. 99 | self._mesh_j = [] 100 | for o in range(self._loader.num_mano): 101 | mesh = trimesh.creation.uv_sphere(radius=0.005) 102 | mesh.visual.vertex_colors = [1.0, 0.0, 0.0] 103 | self._mesh_j.append(mesh) 104 | 105 | self._faces = self._loader.mano_group_layer.f.cpu().numpy() 106 | 107 | w = self._loader.dimensions[0] 108 | h = self._loader.dimensions[1] 109 | self._r = pyrender.OffscreenRenderer(viewport_width=w, viewport_height=h) 110 | 111 | self._render_dir = [ 112 | os.path.join(os.path.dirname(__file__), "..", "data", "render", 113 | self._name, self._loader.serials[c]) 114 | for c in range(self._loader.num_cameras) 115 | ] 116 | for d in self._render_dir: 117 | os.makedirs(d, exist_ok=True) 118 | 119 | def _blend(self, im_real, im_render): 120 | """Blends the real and rendered images. 121 | 122 | Args: 123 | im_real: A uint8 numpy array of shape [H, W, 3] containing the real image. 124 | im_render: A uint8 numpy array of shape [H, W, 3] containing the rendered 125 | image. 126 | """ 127 | im = 0.33 * im_real.astype(np.float32) + 0.67 * im_render.astype(np.float32) 128 | im = im.astype(np.uint8) 129 | return im 130 | 131 | def _render_color_seg(self): 132 | """Renders and saves color and segmenetation images.""" 133 | print('Rendering color and segmentation') 134 | for i in range(self._loader.num_frames): 135 | print('{:03d}/{:03d}'.format(i + 1, self._loader.num_frames)) 136 | 137 | self._loader.step() 138 | 139 | for c in range(self._loader.num_cameras): 140 | # Create pyrender scene. 141 | scene = pyrender.Scene(bg_color=np.array([0.0, 0.0, 0.0, 0.0]), 142 | ambient_light=np.array([1.0, 1.0, 1.0])) 143 | 144 | # Add camera. 145 | scene.add(self._cameras[c], pose=np.eye(4)) 146 | 147 | seg_node_map = {} 148 | 149 | pose_y = self._loader.ycb_pose[c] 150 | vert_m = self._loader.mano_vert[c] 151 | 152 | # Add YCB meshes. 153 | for o in range(self._loader.num_ycb): 154 | if np.all(pose_y[o] == 0.0): 155 | continue 156 | pose = pose_y[o].copy() 157 | pose[1] *= -1 158 | pose[2] *= -1 159 | node = scene.add(self._mesh_y[o], pose=pose) 160 | seg_node_map.update({node: _YCB_COLORS[self._loader.ycb_ids[o]]}) 161 | 162 | # Add MANO meshes. 163 | for o in range(self._loader.num_mano): 164 | if np.all(vert_m[o] == 0.0): 165 | continue 166 | vert = vert_m[o].copy() 167 | vert[:, 1] *= -1 168 | vert[:, 2] *= -1 169 | mesh = trimesh.Trimesh(vertices=vert, faces=self._faces) 170 | mesh1 = pyrender.Mesh.from_trimesh(mesh) 171 | mesh1.primitives[0].material.baseColorFactor = [0.7, 0.7, 0.7, 1.0] 172 | mesh2 = pyrender.Mesh.from_trimesh(mesh, wireframe=True) 173 | mesh2.primitives[0].material.baseColorFactor = [0.0, 0.0, 0.0, 1.0] 174 | node1 = scene.add(mesh1) 175 | node2 = scene.add(mesh2) 176 | seg_node_map.update({node1: _MANO_COLOR}) 177 | 178 | color, _ = self._r.render(scene) 179 | color_seg, _ = self._r.render(scene, 180 | pyrender.RenderFlags.SEG, 181 | seg_node_map=seg_node_map) 182 | 183 | im = self._loader.pcd_rgb[c] 184 | b_color = self._blend(im, color) 185 | b_color_seg = self._blend(im, color_seg) 186 | 187 | color_file = self._render_dir[c] + "/color_{:06d}.jpg".format(i) 188 | seg_file = self._render_dir[c] + "/seg_{:06d}.jpg".format(i) 189 | cv2.imwrite(color_file, b_color[:, :, ::-1]) 190 | cv2.imwrite(seg_file, b_color_seg[:, :, ::-1]) 191 | 192 | def _render_joint(self): 193 | """Renders and saves hand joint visualizations.""" 194 | print('Rendering joint') 195 | for i in range(self._loader.num_frames): 196 | print('{:03d}/{:03d}'.format(i + 1, self._loader.num_frames)) 197 | 198 | self._loader.step() 199 | 200 | for c in range(self._loader.num_cameras): 201 | # Create pyrender scene. 202 | scene = pyrender.Scene(bg_color=np.array([0.0, 0.0, 0.0, 0.0]), 203 | ambient_light=np.array([1.0, 1.0, 1.0])) 204 | 205 | # Add camera. 206 | scene.add(self._cameras[c], pose=np.eye(4)) 207 | 208 | joint_3d = self._loader.mano_joint_3d[c] 209 | 210 | # Add MANO joints. 211 | for o in range(self._loader.num_mano): 212 | if np.all(joint_3d[o] == -1): 213 | continue 214 | j = joint_3d[o].copy() 215 | j[:, 1] *= -1 216 | j[:, 2] *= -1 217 | tfs = np.tile(np.eye(4), (21, 1, 1)) 218 | tfs[:, :3, 3] = j 219 | mesh = pyrender.Mesh.from_trimesh(self._mesh_j[o], poses=tfs) 220 | scene.add(mesh) 221 | 222 | color, _ = self._r.render(scene) 223 | 224 | im = self._loader.pcd_rgb[c] 225 | color = self._blend(im, color) 226 | 227 | color_file = self._render_dir[c] + "/joints_{:06d}.jpg".format(i) 228 | cv2.imwrite(color_file, color[:, :, ::-1]) 229 | 230 | def run(self): 231 | """Runs the renderer.""" 232 | self._render_color_seg() 233 | self._render_joint() 234 | 235 | 236 | if __name__ == '__main__': 237 | args = parse_args() 238 | 239 | renderer = Renderer(args.name, args.device) 240 | renderer.run() 241 | -------------------------------------------------------------------------------- /dex_ycb_toolkit/coco_eval.py: -------------------------------------------------------------------------------- 1 | # DexYCB Toolkit 2 | # Copyright (C) 2021 NVIDIA Corporation 3 | # Licensed under the GNU General Public License v3.0 [see LICENSE for details] 4 | 5 | """COCO evaluator.""" 6 | 7 | import os 8 | import time 9 | import numpy as np 10 | import pycocotools.mask 11 | import json 12 | import copy 13 | import itertools 14 | 15 | from pycocotools.coco import COCO 16 | from pycocotools.cocoeval import COCOeval 17 | from tabulate import tabulate 18 | 19 | from dex_ycb_toolkit.factory import get_dataset 20 | from dex_ycb_toolkit.logging import get_logger 21 | 22 | # TODO(ywchao): tune OKS following https://cocodataset.org/#keypoints-eval. 23 | _KPT_OKS_SIGMAS = [0.05] * 21 24 | 25 | 26 | class COCOEvaluator(): 27 | """COCO evaluator.""" 28 | 29 | def __init__(self, name): 30 | """Constructor. 31 | 32 | Args: 33 | name: Dataset name. E.g., 's0_test'. 34 | """ 35 | self._name = name 36 | 37 | self._dataset = get_dataset(self._name) 38 | 39 | self._class_names = {**self._dataset.ycb_classes, 22: 'hand'} 40 | 41 | self._out_dir = os.path.join(os.path.dirname(__file__), "..", "results") 42 | 43 | self._anno_file = os.path.join(self._out_dir, 44 | "anno_coco_{}.json".format(self._name)) 45 | 46 | if os.path.isfile(self._anno_file): 47 | print('Found COCO annnotation file.') 48 | else: 49 | print('Cannot find COCO annnotation file.') 50 | self._generate_anno_file() 51 | 52 | def _generate_anno_file(self): 53 | """Generates the annotation file.""" 54 | print('Generating COCO annotation file') 55 | s = time.time() 56 | 57 | images = [] 58 | annotations = [] 59 | cnt_ann = 0 60 | 61 | for i in range(len(self._dataset)): 62 | if (i + 1) in np.floor(np.linspace(0, len(self._dataset), 11))[1:]: 63 | print('{:3.0f}% {:6d}/{:6d}'.format(100 * i / len(self._dataset), i, 64 | len(self._dataset))) 65 | 66 | sample = self._dataset[i] 67 | 68 | img = { 69 | 'id': i, 70 | 'width': self._dataset.w, 71 | 'height': self._dataset.h, 72 | } 73 | images.append(img) 74 | 75 | label = np.load(sample['label_file']) 76 | 77 | for y in sample['ycb_ids'] + [255]: 78 | mask = label['seg'] == y 79 | if np.count_nonzero(mask) == 0: 80 | continue 81 | mask = np.asfortranarray(mask) 82 | rle = pycocotools.mask.encode(mask) 83 | segmentation = rle 84 | segmentation['counts'] = segmentation['counts'].decode('ascii') 85 | # https://github.com/cocodataset/cocoapi/issues/36 86 | area = pycocotools.mask.area(rle).item() 87 | bbox = pycocotools.mask.toBbox(rle).tolist() 88 | if y == 255: 89 | category_id = 22 90 | keypoints = label['joint_2d'].squeeze(0).tolist() 91 | keypoints = [[0.0, 0.0, 0] if x[0] == -1 and x[1] == -1 else x + [2] 92 | for x in keypoints] 93 | keypoints = [y for x in keypoints for y in x] 94 | num_keypoints = 21 95 | else: 96 | category_id = y 97 | keypoints = [0] * 21 * 3 98 | num_keypoints = 0 99 | ann = { 100 | 'id': cnt_ann + 1, 101 | 'image_id': i, 102 | 'category_id': category_id, 103 | 'segmentation': segmentation, 104 | 'area': area, 105 | 'bbox': bbox, 106 | 'iscrowd': 0, 107 | 'keypoints': keypoints, 108 | 'num_keypoints': num_keypoints, 109 | } 110 | annotations.append(ann) 111 | cnt_ann += 1 112 | 113 | categories = [] 114 | 115 | for i, x in self._class_names.items(): 116 | if x == 'hand': 117 | supercategory = 'mano' 118 | keypoints = self._dataset.mano_joints 119 | skeleton = [[y + 1 for y in x] for x in self._dataset.mano_joint_connect 120 | ] 121 | else: 122 | supercategory = 'ycb' 123 | keypoints = [] 124 | skeleton = [] 125 | cat = { 126 | 'id': i, 127 | 'name': x, 128 | 'supercategory': supercategory, 129 | 'keypoints': keypoints, 130 | 'skeleton': skeleton, 131 | } 132 | categories.append(cat) 133 | 134 | anno = {} 135 | anno['info'] = {} 136 | anno['images'] = images 137 | anno['annotations'] = annotations 138 | anno['categories'] = categories 139 | 140 | print('Saving to {}'.format(self._anno_file)) 141 | 142 | os.makedirs(os.path.dirname(self._anno_file), exist_ok=True) 143 | 144 | with open(self._anno_file, 'w') as f: 145 | json.dump(anno, f) 146 | 147 | e = time.time() 148 | print('time: {:7.2f}'.format(e - s)) 149 | 150 | # https://github.com/facebookresearch/detectron2/blob/492cf9c7bae22d7d528f7f58169fcd52a450a0ca/detectron2/evaluation/coco_evaluation.py#L252 151 | def _derive_coco_results(self, coco_eval, iou_type, logger): 152 | """Derives COCO results. 153 | 154 | Args: 155 | coco_eval: A COCOEval object. 156 | iou_type: 'bbox', 'segm', or 'keypoints'. 157 | logger: Logger. 158 | 159 | Returns: 160 | A dictionary holding the results. 161 | """ 162 | metrics = { 163 | 'bbox': ['AP', 'AP50', 'AP75', 'APs', 'APm', 'APl'], 164 | 'segm': ['AP', 'AP50', 'AP75', 'APs', 'APm', 'APl'], 165 | 'keypoints': ['AP', 'AP50', 'AP75', 'APm', 'APl'], 166 | }[iou_type] 167 | 168 | results = { 169 | metric: float(coco_eval.stats[idx] * 170 | 100 if coco_eval.stats[idx] >= 0 else "nan") 171 | for idx, metric in enumerate(metrics) 172 | } 173 | keys, values = tuple(zip(*results.items())) 174 | table = tabulate( 175 | [values], 176 | headers=keys, 177 | tablefmt='pipe', 178 | floatfmt='.3f', 179 | stralign='center', 180 | numalign='center', 181 | ) 182 | logger.info('Evaluation results for *{}*: \n'.format(iou_type) + table) 183 | if not np.isfinite(sum(results.values())): 184 | logger.info('Some metrics cannot be computed and is shown as NaN.') 185 | 186 | precisions = coco_eval.eval["precision"] 187 | assert len(self._class_names) == precisions.shape[2] 188 | 189 | results_per_category = [] 190 | for idx, (_, name) in enumerate(self._class_names.items()): 191 | precision = precisions[:, :, idx, 0, -1] 192 | precision = precision[precision > -1] 193 | ap = np.mean(precision) if precision.size else float('nan') 194 | results_per_category.append(("{}".format(name), float(ap * 100))) 195 | 196 | n_cols = min(6, len(results_per_category) * 2) 197 | results_flatten = list(itertools.chain(*results_per_category)) 198 | results_2d = itertools.zip_longest( 199 | *[results_flatten[i::n_cols] for i in range(n_cols)]) 200 | table = tabulate( 201 | results_2d, 202 | tablefmt='pipe', 203 | floatfmt='.3f', 204 | headers=['category', 'AP'] * (n_cols // 2), 205 | numalign='left', 206 | ) 207 | logger.info('Per-category *{}* AP: \n'.format(iou_type) + table) 208 | 209 | results.update({'AP-' + name: ap for name, ap in results_per_category}) 210 | return results 211 | 212 | def evaluate(self, 213 | res_file, 214 | out_dir=None, 215 | tasks=('bbox', 'segm', 'keypoints')): 216 | """Evaluates COCO metrics given a result file. 217 | 218 | Args: 219 | res_file: Path to the result file. 220 | out_dir: Path to the output directory. 221 | tasks: A tuple of evaluated tasks. 'bbox', 'segm', and 'keypoints'. 222 | 223 | Returns: 224 | A dictionary holding the results. 225 | """ 226 | if out_dir is None: 227 | out_dir = self._out_dir 228 | 229 | res_name = os.path.splitext(os.path.basename(res_file))[0] 230 | log_file = os.path.join(out_dir, 231 | "coco_eval_{}_{}.log".format(self._name, res_name)) 232 | logger = get_logger(log_file) 233 | 234 | coco_gt = COCO(self._anno_file) 235 | coco_dt = coco_gt.loadRes(res_file) 236 | 237 | results = {} 238 | 239 | for task in tasks: 240 | # https://github.com/facebookresearch/detectron2/blob/492cf9c7bae22d7d528f7f58169fcd52a450a0ca/detectron2/evaluation/coco_evaluation.py#L506 241 | if task == 'segm': 242 | coco_dt_ = copy.deepcopy(coco_dt) 243 | for ann in coco_dt_.loadAnns(coco_dt_.getAnnIds()): 244 | ann.pop('bbox', None) 245 | coco_dt_ = coco_gt.loadRes(coco_dt_.dataset['annotations']) 246 | else: 247 | coco_dt_ = coco_dt 248 | 249 | coco_eval = COCOeval(coco_gt, coco_dt_, task) 250 | 251 | if task == 'keypoints': 252 | coco_eval.params.kpt_oks_sigmas = np.array(_KPT_OKS_SIGMAS) 253 | 254 | coco_eval.evaluate() 255 | coco_eval.accumulate() 256 | coco_eval.summarize() 257 | 258 | results[task] = self._derive_coco_results(coco_eval, task, logger) 259 | 260 | logger.info('Evaluation complete.') 261 | 262 | return results 263 | -------------------------------------------------------------------------------- /dex_ycb_toolkit/dex_ycb.py: -------------------------------------------------------------------------------- 1 | # DexYCB Toolkit 2 | # Copyright (C) 2021 NVIDIA Corporation 3 | # Licensed under the GNU General Public License v3.0 [see LICENSE for details] 4 | 5 | """DexYCB dataset.""" 6 | 7 | import os 8 | import yaml 9 | import numpy as np 10 | 11 | _SUBJECTS = [ 12 | '20200709-subject-01', 13 | '20200813-subject-02', 14 | '20200820-subject-03', 15 | '20200903-subject-04', 16 | '20200908-subject-05', 17 | '20200918-subject-06', 18 | '20200928-subject-07', 19 | '20201002-subject-08', 20 | '20201015-subject-09', 21 | '20201022-subject-10', 22 | ] 23 | 24 | _SERIALS = [ 25 | '836212060125', 26 | '839512060362', 27 | '840412060917', 28 | '841412060263', 29 | '932122060857', 30 | '932122060861', 31 | '932122061900', 32 | '932122062010', 33 | ] 34 | 35 | _YCB_CLASSES = { 36 | 1: '002_master_chef_can', 37 | 2: '003_cracker_box', 38 | 3: '004_sugar_box', 39 | 4: '005_tomato_soup_can', 40 | 5: '006_mustard_bottle', 41 | 6: '007_tuna_fish_can', 42 | 7: '008_pudding_box', 43 | 8: '009_gelatin_box', 44 | 9: '010_potted_meat_can', 45 | 10: '011_banana', 46 | 11: '019_pitcher_base', 47 | 12: '021_bleach_cleanser', 48 | 13: '024_bowl', 49 | 14: '025_mug', 50 | 15: '035_power_drill', 51 | 16: '036_wood_block', 52 | 17: '037_scissors', 53 | 18: '040_large_marker', 54 | 19: '051_large_clamp', 55 | 20: '052_extra_large_clamp', 56 | 21: '061_foam_brick', 57 | } 58 | 59 | _MANO_JOINTS = [ 60 | 'wrist', 61 | 'thumb_mcp', 62 | 'thumb_pip', 63 | 'thumb_dip', 64 | 'thumb_tip', 65 | 'index_mcp', 66 | 'index_pip', 67 | 'index_dip', 68 | 'index_tip', 69 | 'middle_mcp', 70 | 'middle_pip', 71 | 'middle_dip', 72 | 'middle_tip', 73 | 'ring_mcp', 74 | 'ring_pip', 75 | 'ring_dip', 76 | 'ring_tip', 77 | 'little_mcp', 78 | 'little_pip', 79 | 'little_dip', 80 | 'little_tip' 81 | ] 82 | 83 | _MANO_JOINT_CONNECT = [ 84 | [0, 1], [ 1, 2], [ 2, 3], [ 3, 4], 85 | [0, 5], [ 5, 6], [ 6, 7], [ 7, 8], 86 | [0, 9], [ 9, 10], [10, 11], [11, 12], 87 | [0, 13], [13, 14], [14, 15], [15, 16], 88 | [0, 17], [17, 18], [18, 19], [19, 20], 89 | ] 90 | 91 | _BOP_EVAL_SUBSAMPLING_FACTOR = 4 92 | 93 | 94 | class DexYCBDataset(): 95 | """DexYCB dataset.""" 96 | ycb_classes = _YCB_CLASSES 97 | mano_joints = _MANO_JOINTS 98 | mano_joint_connect = _MANO_JOINT_CONNECT 99 | 100 | def __init__(self, setup, split): 101 | """Constructor. 102 | 103 | Args: 104 | setup: Setup name. 's0', 's1', 's2', or 's3'. 105 | split: Split name. 'train', 'val', or 'test'. 106 | """ 107 | self._setup = setup 108 | self._split = split 109 | 110 | assert 'DEX_YCB_DIR' in os.environ, "environment variable 'DEX_YCB_DIR' is not set" 111 | self._data_dir = os.environ['DEX_YCB_DIR'] 112 | self._calib_dir = os.path.join(self._data_dir, "calibration") 113 | self._model_dir = os.path.join(self._data_dir, "models") 114 | 115 | self._color_format = "color_{:06d}.jpg" 116 | self._depth_format = "aligned_depth_to_color_{:06d}.png" 117 | self._label_format = "labels_{:06d}.npz" 118 | self._h = 480 119 | self._w = 640 120 | 121 | self._obj_file = { 122 | k: os.path.join(self._model_dir, v, "textured_simple.obj") 123 | for k, v in _YCB_CLASSES.items() 124 | } 125 | 126 | # Seen subjects, camera views, grasped objects. 127 | if self._setup == 's0': 128 | if self._split == 'train': 129 | subject_ind = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9] 130 | serial_ind = [0, 1, 2, 3, 4, 5, 6, 7] 131 | sequence_ind = [i for i in range(100) if i % 5 != 4] 132 | if self._split == 'val': 133 | subject_ind = [0, 1] 134 | serial_ind = [0, 1, 2, 3, 4, 5, 6, 7] 135 | sequence_ind = [i for i in range(100) if i % 5 == 4] 136 | if self._split == 'test': 137 | subject_ind = [2, 3, 4, 5, 6, 7, 8, 9] 138 | serial_ind = [0, 1, 2, 3, 4, 5, 6, 7] 139 | sequence_ind = [i for i in range(100) if i % 5 == 4] 140 | 141 | # Unseen subjects. 142 | if self._setup == 's1': 143 | if self._split == 'train': 144 | subject_ind = [0, 1, 2, 3, 4, 5, 9] 145 | serial_ind = [0, 1, 2, 3, 4, 5, 6, 7] 146 | sequence_ind = list(range(100)) 147 | if self._split == 'val': 148 | subject_ind = [6] 149 | serial_ind = [0, 1, 2, 3, 4, 5, 6, 7] 150 | sequence_ind = list(range(100)) 151 | if self._split == 'test': 152 | subject_ind = [7, 8] 153 | serial_ind = [0, 1, 2, 3, 4, 5, 6, 7] 154 | sequence_ind = list(range(100)) 155 | 156 | # Unseen camera views. 157 | if self._setup == 's2': 158 | if self._split == 'train': 159 | subject_ind = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9] 160 | serial_ind = [0, 1, 2, 3, 4, 5] 161 | sequence_ind = list(range(100)) 162 | if self._split == 'val': 163 | subject_ind = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9] 164 | serial_ind = [6] 165 | sequence_ind = list(range(100)) 166 | if self._split == 'test': 167 | subject_ind = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9] 168 | serial_ind = [7] 169 | sequence_ind = list(range(100)) 170 | 171 | # Unseen grasped objects. 172 | if self._setup == 's3': 173 | if self._split == 'train': 174 | subject_ind = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9] 175 | serial_ind = [0, 1, 2, 3, 4, 5, 6, 7] 176 | sequence_ind = [ 177 | i for i in range(100) if i // 5 not in (3, 7, 11, 15, 19) 178 | ] 179 | if self._split == 'val': 180 | subject_ind = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9] 181 | serial_ind = [0, 1, 2, 3, 4, 5, 6, 7] 182 | sequence_ind = [i for i in range(100) if i // 5 in (3, 19)] 183 | if self._split == 'test': 184 | subject_ind = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9] 185 | serial_ind = [0, 1, 2, 3, 4, 5, 6, 7] 186 | sequence_ind = [i for i in range(100) if i // 5 in (7, 11, 15)] 187 | 188 | self._subjects = [_SUBJECTS[i] for i in subject_ind] 189 | 190 | self._serials = [_SERIALS[i] for i in serial_ind] 191 | self._intrinsics = [] 192 | for s in self._serials: 193 | intr_file = os.path.join(self._calib_dir, "intrinsics", 194 | "{}_{}x{}.yml".format(s, self._w, self._h)) 195 | with open(intr_file, 'r') as f: 196 | intr = yaml.load(f, Loader=yaml.FullLoader) 197 | intr = intr['color'] 198 | self._intrinsics.append(intr) 199 | 200 | self._sequences = [] 201 | self._mapping = [] 202 | self._ycb_ids = [] 203 | self._ycb_grasp_ind = [] 204 | self._mano_side = [] 205 | self._mano_betas = [] 206 | offset = 0 207 | for n in self._subjects: 208 | seq = sorted(os.listdir(os.path.join(self._data_dir, n))) 209 | seq = [os.path.join(n, s) for s in seq] 210 | assert len(seq) == 100 211 | seq = [seq[i] for i in sequence_ind] 212 | self._sequences += seq 213 | for i, q in enumerate(seq): 214 | meta_file = os.path.join(self._data_dir, q, "meta.yml") 215 | with open(meta_file, 'r') as f: 216 | meta = yaml.load(f, Loader=yaml.FullLoader) 217 | c = np.arange(len(self._serials)) 218 | f = np.arange(meta['num_frames']) 219 | f, c = np.meshgrid(f, c) 220 | c = c.ravel() 221 | f = f.ravel() 222 | s = (offset + i) * np.ones_like(c) 223 | m = np.vstack((s, c, f)).T 224 | self._mapping.append(m) 225 | self._ycb_ids.append(meta['ycb_ids']) 226 | self._ycb_grasp_ind.append(meta['ycb_grasp_ind']) 227 | self._mano_side.append(meta['mano_sides'][0]) 228 | mano_calib_file = os.path.join(self._data_dir, "calibration", 229 | "mano_{}".format(meta['mano_calib'][0]), 230 | "mano.yml") 231 | with open(mano_calib_file, 'r') as f: 232 | mano_calib = yaml.load(f, Loader=yaml.FullLoader) 233 | self._mano_betas.append(mano_calib['betas']) 234 | offset += len(seq) 235 | self._mapping = np.vstack(self._mapping) 236 | 237 | def __len__(self): 238 | return len(self._mapping) 239 | 240 | def __getitem__(self, idx): 241 | s, c, f = self._mapping[idx] 242 | d = os.path.join(self._data_dir, self._sequences[s], self._serials[c]) 243 | sample = { 244 | 'color_file': os.path.join(d, self._color_format.format(f)), 245 | 'depth_file': os.path.join(d, self._depth_format.format(f)), 246 | 'label_file': os.path.join(d, self._label_format.format(f)), 247 | 'intrinsics': self._intrinsics[c], 248 | 'ycb_ids': self._ycb_ids[s], 249 | 'ycb_grasp_ind': self._ycb_grasp_ind[s], 250 | 'mano_side': self._mano_side[s], 251 | 'mano_betas': self._mano_betas[s], 252 | } 253 | if self._split == 'test': 254 | sample['is_bop_target'] = (f % _BOP_EVAL_SUBSAMPLING_FACTOR == 0).item() 255 | id_next = idx + _BOP_EVAL_SUBSAMPLING_FACTOR 256 | is_last = (id_next >= len(self._mapping) or 257 | (np.any(self._mapping[id_next][:2] != [s, c])).item()) 258 | sample['is_grasp_target'] = sample['is_bop_target'] and is_last 259 | return sample 260 | 261 | @property 262 | def data_dir(self): 263 | return self._data_dir 264 | 265 | @property 266 | def h(self): 267 | return self._h 268 | 269 | @property 270 | def w(self): 271 | return self._w 272 | 273 | @property 274 | def obj_file(self): 275 | return self._obj_file 276 | 277 | def get_bop_id_from_idx(self, idx): 278 | """Returns the BOP scene ID and image ID given an index. 279 | 280 | Args: 281 | idx: Index of sample. 282 | 283 | Returns: 284 | scene_id: BOP scene ID. 285 | im_id: BOP image ID. 286 | """ 287 | s, c, f = map(lambda x: x.item(), self._mapping[idx]) 288 | scene_id = s * len(self._serials) + c 289 | im_id = f 290 | return scene_id, im_id 291 | -------------------------------------------------------------------------------- /dex_ycb_toolkit/bop_eval.py: -------------------------------------------------------------------------------- 1 | # DexYCB Toolkit 2 | # Copyright (C) 2021 NVIDIA Corporation 3 | # Licensed under the GNU General Public License v3.0 [see LICENSE for details] 4 | 5 | """BOP evaluator.""" 6 | 7 | import os 8 | import sys 9 | import numpy as np 10 | import subprocess 11 | import itertools 12 | 13 | from collections import defaultdict 14 | from tabulate import tabulate 15 | 16 | from dex_ycb_toolkit.factory import get_dataset 17 | from dex_ycb_toolkit.logging import get_logger 18 | 19 | bop_toolkit_root = os.path.join(os.path.dirname(__file__), "..", "bop_toolkit") 20 | sys.path.append(bop_toolkit_root) 21 | 22 | from bop_toolkit_lib import dataset_params 23 | from bop_toolkit_lib import inout 24 | from bop_toolkit_lib import misc 25 | from bop_toolkit_lib import score 26 | 27 | # See dataset_info.md in http://ptak.felk.cvut.cz/6DB/public/bop_datasets/ycbv_base.zip. 28 | _BOP_TRANSLATIONS = { 29 | '002_master_chef_can': [ 1.3360, -0.5000, 3.5105], 30 | '003_cracker_box': [ 0.5575, 1.7005, 4.8050], 31 | '004_sugar_box': [ -0.9520, 1.4670, 4.3645], 32 | '005_tomato_soup_can': [ -0.0240, -1.5270, 8.4035], 33 | '006_mustard_bottle': [ 1.2995, 2.4870, -11.8290], 34 | '007_tuna_fish_can': [ -0.1565, 0.1150, 4.2625], 35 | '008_pudding_box': [ 1.1645, -4.2015, 3.1190], 36 | '009_gelatin_box': [ 1.4460, -0.5915, 3.6085], 37 | '010_potted_meat_can': [ 2.4195, 0.3075, 8.0715], 38 | '011_banana': [-18.6730, 12.1915, -1.4635], 39 | '019_pitcher_base': [ 5.3370, 5.8855, 25.6115], 40 | '021_bleach_cleanser': [ 4.9290, -2.4800, -13.2920], 41 | '024_bowl': [ -0.2270, 0.7950, -2.9675], 42 | '025_mug': [ -8.4675, -0.6995, -1.6145], 43 | '035_power_drill': [ 9.0710, 20.9360, -2.1190], 44 | '036_wood_block': [ 1.4265, -2.5305, 17.1890], 45 | '037_scissors': [ 7.0535, -28.1320, 0.0420], 46 | '040_large_marker': [ 0.0460, -2.1040, 0.3500], 47 | '051_large_clamp': [ 10.5180, -1.9640, -0.4745], 48 | '052_extra_large_clamp': [ -0.3950, -10.4130, 0.1620], 49 | '061_foam_brick': [ -0.0805, 0.0805, -8.2435], 50 | } 51 | 52 | 53 | class BOPEvaluator(): 54 | """BOP evaluator.""" 55 | 56 | def __init__(self, name): 57 | """Constructor. 58 | 59 | Args: 60 | name: Dataset name. E.g., 's0_test'. 61 | """ 62 | self._name = name 63 | 64 | self._dataset = get_dataset(self._name) 65 | 66 | self._setup = self._name.split('_')[0] 67 | self._split = self._name.split('_')[1] 68 | 69 | self._out_dir = os.path.join(os.path.dirname(__file__), "..", "results") 70 | self._bop_dir = os.path.join(self._dataset.data_dir, "bop") 71 | 72 | self._p = { 73 | 'errors': [ 74 | { 75 | 'n_top': -1, 76 | 'type': 'vsd', 77 | 'vsd_delta': 15, 78 | 'vsd_taus': list(np.arange(0.05, 0.51, 0.05)), 79 | 'correct_th': [[th] for th in np.arange(0.05, 0.51, 0.05)] 80 | }, 81 | { 82 | 'n_top': -1, 83 | 'type': 'mssd', 84 | 'correct_th': [[th] for th in np.arange(0.05, 0.51, 0.05)] 85 | }, 86 | { 87 | 'n_top': -1, 88 | 'type': 'mspd', 89 | 'correct_th': [[th] for th in np.arange(5, 51, 5)] 90 | }, 91 | ], 92 | 'visib_gt_min': -1, 93 | } 94 | 95 | dp_split = dataset_params.get_split_params(self._bop_dir, self._setup, 96 | self._split) 97 | dp_model = dataset_params.get_model_params(self._bop_dir, 98 | self._setup, 99 | model_type='eval') 100 | self._scene_ids = dp_split['scene_ids'] 101 | self._obj_ids = dp_model['obj_ids'] 102 | 103 | self._grasp_id = defaultdict(lambda: {}) 104 | for i in range(len(self._dataset)): 105 | sample = self._dataset[i] 106 | scene_id, im_id = self._dataset.get_bop_id_from_idx(i) 107 | obj_id = sample['ycb_ids'][sample['ycb_grasp_ind']] 108 | self._grasp_id[scene_id][im_id] = obj_id 109 | 110 | def _convert_pose_to_bop(self, est): 111 | """Converts pose from DexYCB models to BOP YCBV models. 112 | 113 | Args: 114 | est: A dictionary holding a single pose estimate. 115 | 116 | Returns: 117 | A dictionary holding the converted pose. 118 | """ 119 | est['t'] -= np.dot( 120 | est['R'], 121 | _BOP_TRANSLATIONS[self._dataset.ycb_classes[est['obj_id']]]).reshape( 122 | 3, 1) 123 | return est 124 | 125 | def _derive_bop_results(self, out_dir, result_name, grasp_only, logger): 126 | """Derives BOP results. 127 | 128 | Args: 129 | out_dir: Path to the output directory. 130 | result_name: BOP result name. Should be the name of a folder under out_dir 131 | that contains output from BOP evaluation. 132 | grasp_only: Whether to derive results on grasped objects only. 133 | logger: Logger. 134 | 135 | Returns: 136 | A dictionary holding the results. 137 | """ 138 | if grasp_only: 139 | set_str = 'grasp only' 140 | else: 141 | set_str = 'all' 142 | 143 | logger.info('Deriving results for *{}*'.format(set_str)) 144 | 145 | average_recalls = {} 146 | average_recalls_obj = defaultdict(lambda: {}) 147 | 148 | for error in self._p['errors']: 149 | 150 | error_dir_paths = {} 151 | if error['type'] == 'vsd': 152 | for vsd_tau in error['vsd_taus']: 153 | error_sign = misc.get_error_signature(error['type'], 154 | error['n_top'], 155 | vsd_delta=error['vsd_delta'], 156 | vsd_tau=vsd_tau) 157 | error_dir_paths[error_sign] = os.path.join(result_name, error_sign) 158 | else: 159 | error_sign = misc.get_error_signature(error['type'], error['n_top']) 160 | error_dir_paths[error_sign] = os.path.join(result_name, error_sign) 161 | 162 | recalls = [] 163 | recalls_obj = defaultdict(lambda: []) 164 | 165 | for error_sign, error_dir_path in error_dir_paths.items(): 166 | for correct_th in error['correct_th']: 167 | 168 | score_sign = misc.get_score_signature(correct_th, 169 | self._p['visib_gt_min']) 170 | matches_filename = "matches_{}.json".format(score_sign) 171 | matches_path = os.path.join(out_dir, error_dir_path, matches_filename) 172 | 173 | matches = inout.load_json(matches_path) 174 | 175 | if grasp_only: 176 | matches = [ 177 | m for m in matches 178 | if m['obj_id'] == self._grasp_id[m['scene_id']][m['im_id']] 179 | ] 180 | 181 | scores = score.calc_localization_scores(self._scene_ids, 182 | self._obj_ids, 183 | matches, 184 | error['n_top'], 185 | do_print=False) 186 | 187 | recalls.append(scores['recall']) 188 | for i, r in scores['obj_recalls'].items(): 189 | recalls_obj[i].append(r) 190 | 191 | average_recalls[error['type']] = np.mean(recalls) 192 | for i, r in recalls_obj.items(): 193 | average_recalls_obj[i][error['type']] = np.mean(r) 194 | 195 | results = {i: r * 100 for i, r in average_recalls.items()} 196 | results['mean'] = np.mean( 197 | [results['vsd'], results['mssd'], results['mspd']]) 198 | 199 | keys, values = tuple(zip(*results.items())) 200 | table = tabulate( 201 | [values], 202 | headers=keys, 203 | tablefmt='pipe', 204 | floatfmt='.3f', 205 | stralign='center', 206 | numalign='center', 207 | ) 208 | logger.info('Evaluation results for *{}*: \n'.format(set_str) + table) 209 | 210 | results_per_object = {} 211 | for i, v in average_recalls_obj.items(): 212 | res = {k: r * 100 for k, r in v.items()} 213 | res['mean'] = np.mean([res['vsd'], res['mssd'], res['mspd']]) 214 | results_per_object[self._dataset.ycb_classes[i]] = res 215 | 216 | n_cols = 5 217 | results_tuple = [(k, v['vsd'], v['mssd'], v['mspd'], v['mean']) 218 | for k, v in results_per_object.items()] 219 | results_flatten = list(itertools.chain(*results_tuple)) 220 | results_2d = itertools.zip_longest( 221 | *[results_flatten[i::n_cols] for i in range(n_cols)]) 222 | table = tabulate( 223 | results_2d, 224 | tablefmt='pipe', 225 | floatfmt='.3f', 226 | headers=['object', 'vsd', 'mssd', 'mspd', 'mean'] * (n_cols // 5), 227 | numalign='right', 228 | ) 229 | logger.info('Per-object scores for *{}*: \n'.format(set_str) + table) 230 | 231 | results['per_obj'] = results_per_object 232 | 233 | return results 234 | 235 | def evaluate(self, res_file, out_dir=None, renderer_type='python'): 236 | """Evaluates BOP metrics given a result file. 237 | 238 | Args: 239 | res_file: Path to the result file. 240 | out_dir: Path to the output directory. 241 | renderer_type: Renderer type. 'python' or 'cpp'. 242 | 243 | Returns: 244 | A dictionary holding the results. 245 | 246 | Raises: 247 | RuntimeError: If BOP evaluation failed. 248 | """ 249 | if out_dir is None: 250 | out_dir = self._out_dir 251 | 252 | ests = inout.load_bop_results(res_file) 253 | ests = [self._convert_pose_to_bop(est) for est in ests] 254 | res_name = os.path.splitext(os.path.basename(res_file))[0] 255 | bop_res_name = 'bop-{}_{}-{}'.format(res_name.replace('_', '-'), 256 | self._setup, self._split) 257 | bop_res_file = os.path.join(out_dir, "{}.csv".format(bop_res_name)) 258 | inout.save_bop_results(bop_res_file, ests) 259 | 260 | eval_cmd = [ 261 | 'python', 262 | os.path.join('scripts', 'eval_bop19.py'), 263 | '--renderer_type={}'.format(renderer_type), 264 | '--result_filenames={}'.format(bop_res_file), 265 | '--results_path={}'.format(out_dir), 266 | '--eval_path={}'.format(out_dir), 267 | ] 268 | cwd = "bop_toolkit" 269 | env = os.environ.copy() 270 | env['PYTHONPATH'] = "." 271 | env['BOP_PATH'] = self._bop_dir 272 | 273 | if subprocess.run(eval_cmd, cwd=cwd, env=env).returncode != 0: 274 | raise RuntimeError('BOP evaluation failed.') 275 | 276 | log_file = os.path.join(out_dir, 277 | "bop_eval_{}_{}.log".format(self._name, res_name)) 278 | logger = get_logger(log_file) 279 | 280 | results = {} 281 | results['all'] = self._derive_bop_results(out_dir, bop_res_name, False, 282 | logger) 283 | results['grasp_only'] = self._derive_bop_results(out_dir, bop_res_name, 284 | True, logger) 285 | 286 | logger.info('Evaluation complete.') 287 | 288 | return results 289 | -------------------------------------------------------------------------------- /examples/all_cvpr2021_results_eval_scripts.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | 3 | # COCO 4 | python examples/evaluate_coco.py \ 5 | --name s0_test \ 6 | --res_file results/cvpr2021_results/coco_maskrcnn_s0_test.json 7 | python examples/evaluate_coco.py \ 8 | --name s1_test \ 9 | --res_file results/cvpr2021_results/coco_maskrcnn_s1_test.json 10 | python examples/evaluate_coco.py \ 11 | --name s2_test \ 12 | --res_file results/cvpr2021_results/coco_maskrcnn_s2_test.json 13 | python examples/evaluate_coco.py \ 14 | --name s3_test \ 15 | --res_file results/cvpr2021_results/coco_maskrcnn_s3_test.json 16 | 17 | # BOP 18 | python examples/evaluate_bop.py \ 19 | --name s0_test \ 20 | --res_file results/cvpr2021_results/bop_posecnn_s0_test.csv 21 | python examples/evaluate_bop.py \ 22 | --name s1_test \ 23 | --res_file results/cvpr2021_results/bop_posecnn_s1_test.csv 24 | python examples/evaluate_bop.py \ 25 | --name s2_test \ 26 | --res_file results/cvpr2021_results/bop_posecnn_s2_test.csv 27 | python examples/evaluate_bop.py \ 28 | --name s3_test \ 29 | --res_file results/cvpr2021_results/bop_posecnn_s3_test.csv 30 | python examples/evaluate_bop.py \ 31 | --name s0_test \ 32 | --res_file results/cvpr2021_results/bop_posecnn_s0_test_refined.csv 33 | python examples/evaluate_bop.py \ 34 | --name s1_test \ 35 | --res_file results/cvpr2021_results/bop_posecnn_s1_test_refined.csv 36 | python examples/evaluate_bop.py \ 37 | --name s2_test \ 38 | --res_file results/cvpr2021_results/bop_posecnn_s2_test_refined.csv 39 | python examples/evaluate_bop.py \ 40 | --name s3_test \ 41 | --res_file results/cvpr2021_results/bop_posecnn_s3_test_refined.csv 42 | python examples/evaluate_bop.py \ 43 | --name s0_test \ 44 | --res_file results/cvpr2021_results/bop_deepim_s0_test_COLOR.csv 45 | python examples/evaluate_bop.py \ 46 | --name s1_test \ 47 | --res_file results/cvpr2021_results/bop_deepim_s1_test_COLOR.csv 48 | python examples/evaluate_bop.py \ 49 | --name s2_test \ 50 | --res_file results/cvpr2021_results/bop_deepim_s2_test_COLOR.csv 51 | python examples/evaluate_bop.py \ 52 | --name s3_test \ 53 | --res_file results/cvpr2021_results/bop_deepim_s3_test_COLOR.csv 54 | python examples/evaluate_bop.py \ 55 | --name s0_test \ 56 | --res_file results/cvpr2021_results/bop_deepim_s0_test_RGBD.csv 57 | python examples/evaluate_bop.py \ 58 | --name s1_test \ 59 | --res_file results/cvpr2021_results/bop_deepim_s1_test_RGBD.csv 60 | python examples/evaluate_bop.py \ 61 | --name s2_test \ 62 | --res_file results/cvpr2021_results/bop_deepim_s2_test_RGBD.csv 63 | python examples/evaluate_bop.py \ 64 | --name s3_test \ 65 | --res_file results/cvpr2021_results/bop_deepim_s3_test_RGBD.csv 66 | python examples/evaluate_bop.py \ 67 | --name s0_test \ 68 | --res_file results/cvpr2021_results/bop_poserbpf_s0_test_rgb.csv 69 | python examples/evaluate_bop.py \ 70 | --name s1_test \ 71 | --res_file results/cvpr2021_results/bop_poserbpf_s1_test_rgb.csv 72 | python examples/evaluate_bop.py \ 73 | --name s2_test \ 74 | --res_file results/cvpr2021_results/bop_poserbpf_s2_test_rgb.csv 75 | python examples/evaluate_bop.py \ 76 | --name s3_test \ 77 | --res_file results/cvpr2021_results/bop_poserbpf_s3_test_rgb.csv 78 | python examples/evaluate_bop.py \ 79 | --name s0_test \ 80 | --res_file results/cvpr2021_results/bop_poserbpf_s0_test_rgbd.csv 81 | python examples/evaluate_bop.py \ 82 | --name s1_test \ 83 | --res_file results/cvpr2021_results/bop_poserbpf_s1_test_rgbd.csv 84 | python examples/evaluate_bop.py \ 85 | --name s2_test \ 86 | --res_file results/cvpr2021_results/bop_poserbpf_s2_test_rgbd.csv 87 | python examples/evaluate_bop.py \ 88 | --name s3_test \ 89 | --res_file results/cvpr2021_results/bop_poserbpf_s3_test_rgbd.csv 90 | python examples/evaluate_bop.py \ 91 | --name s0_test \ 92 | --res_file results/cvpr2021_results/bop_cosypose_s0_test.csv 93 | python examples/evaluate_bop.py \ 94 | --name s1_test \ 95 | --res_file results/cvpr2021_results/bop_cosypose_s1_test.csv 96 | python examples/evaluate_bop.py \ 97 | --name s2_test \ 98 | --res_file results/cvpr2021_results/bop_cosypose_s2_test.csv 99 | python examples/evaluate_bop.py \ 100 | --name s3_test \ 101 | --res_file results/cvpr2021_results/bop_cosypose_s3_test.csv 102 | python examples/evaluate_bop.py \ 103 | --name s1_test \ 104 | --res_file results/cvpr2021_results/bop_dope_s1_test.csv 105 | 106 | # HPE 107 | python examples/evaluate_hpe.py \ 108 | --name s0_test \ 109 | --res_file results/cvpr2021_results/hpe_spurr_hrnet_s0_test.txt 110 | python examples/evaluate_hpe.py \ 111 | --name s1_test \ 112 | --res_file results/cvpr2021_results/hpe_spurr_hrnet_s1_test.txt 113 | python examples/evaluate_hpe.py \ 114 | --name s2_test \ 115 | --res_file results/cvpr2021_results/hpe_spurr_hrnet_s2_test.txt 116 | python examples/evaluate_hpe.py \ 117 | --name s3_test \ 118 | --res_file results/cvpr2021_results/hpe_spurr_hrnet_s3_test.txt 119 | python examples/evaluate_hpe.py \ 120 | --name s0_test \ 121 | --res_file results/cvpr2021_results/hpe_spurr_resnet50_s0_test.txt 122 | python examples/evaluate_hpe.py \ 123 | --name s1_test \ 124 | --res_file results/cvpr2021_results/hpe_spurr_resnet50_s1_test.txt 125 | python examples/evaluate_hpe.py \ 126 | --name s2_test \ 127 | --res_file results/cvpr2021_results/hpe_spurr_resnet50_s2_test.txt 128 | python examples/evaluate_hpe.py \ 129 | --name s3_test \ 130 | --res_file results/cvpr2021_results/hpe_spurr_resnet50_s3_test.txt 131 | python examples/evaluate_hpe.py \ 132 | --name s0_test \ 133 | --res_file results/cvpr2021_results/hpe_a2j_s0_test.txt 134 | python examples/evaluate_hpe.py \ 135 | --name s1_test \ 136 | --res_file results/cvpr2021_results/hpe_a2j_s1_test.txt 137 | python examples/evaluate_hpe.py \ 138 | --name s2_test \ 139 | --res_file results/cvpr2021_results/hpe_a2j_s2_test.txt 140 | python examples/evaluate_hpe.py \ 141 | --name s3_test \ 142 | --res_file results/cvpr2021_results/hpe_a2j_s3_test.txt 143 | 144 | # Grasp 145 | python examples/evaluate_grasp.py \ 146 | --name s0_test \ 147 | --bop_res_file results/cvpr2021_results/bop_posecnn_s0_test.csv \ 148 | --coco_res_file results/cvpr2021_results/coco_maskrcnn_s0_test.json 149 | python examples/evaluate_grasp.py \ 150 | --name s1_test \ 151 | --bop_res_file results/cvpr2021_results/bop_posecnn_s1_test.csv \ 152 | --coco_res_file results/cvpr2021_results/coco_maskrcnn_s1_test.json 153 | python examples/evaluate_grasp.py \ 154 | --name s2_test \ 155 | --bop_res_file results/cvpr2021_results/bop_posecnn_s2_test.csv \ 156 | --coco_res_file results/cvpr2021_results/coco_maskrcnn_s2_test.json 157 | python examples/evaluate_grasp.py \ 158 | --name s3_test \ 159 | --bop_res_file results/cvpr2021_results/bop_posecnn_s3_test.csv \ 160 | --coco_res_file results/cvpr2021_results/coco_maskrcnn_s3_test.json 161 | python examples/evaluate_grasp.py \ 162 | --name s0_test \ 163 | --bop_res_file results/cvpr2021_results/bop_posecnn_s0_test_refined.csv \ 164 | --coco_res_file results/cvpr2021_results/coco_maskrcnn_s0_test.json 165 | python examples/evaluate_grasp.py \ 166 | --name s1_test \ 167 | --bop_res_file results/cvpr2021_results/bop_posecnn_s1_test_refined.csv \ 168 | --coco_res_file results/cvpr2021_results/coco_maskrcnn_s1_test.json 169 | python examples/evaluate_grasp.py \ 170 | --name s2_test \ 171 | --bop_res_file results/cvpr2021_results/bop_posecnn_s2_test_refined.csv \ 172 | --coco_res_file results/cvpr2021_results/coco_maskrcnn_s2_test.json 173 | python examples/evaluate_grasp.py \ 174 | --name s3_test \ 175 | --bop_res_file results/cvpr2021_results/bop_posecnn_s3_test_refined.csv \ 176 | --coco_res_file results/cvpr2021_results/coco_maskrcnn_s3_test.json 177 | python examples/evaluate_grasp.py \ 178 | --name s0_test \ 179 | --bop_res_file results/cvpr2021_results/bop_deepim_s0_test_COLOR.csv \ 180 | --coco_res_file results/cvpr2021_results/coco_maskrcnn_s0_test.json 181 | python examples/evaluate_grasp.py \ 182 | --name s1_test \ 183 | --bop_res_file results/cvpr2021_results/bop_deepim_s1_test_COLOR.csv \ 184 | --coco_res_file results/cvpr2021_results/coco_maskrcnn_s1_test.json 185 | python examples/evaluate_grasp.py \ 186 | --name s2_test \ 187 | --bop_res_file results/cvpr2021_results/bop_deepim_s2_test_COLOR.csv \ 188 | --coco_res_file results/cvpr2021_results/coco_maskrcnn_s2_test.json 189 | python examples/evaluate_grasp.py \ 190 | --name s3_test \ 191 | --bop_res_file results/cvpr2021_results/bop_deepim_s3_test_COLOR.csv \ 192 | --coco_res_file results/cvpr2021_results/coco_maskrcnn_s3_test.json 193 | python examples/evaluate_grasp.py \ 194 | --name s0_test \ 195 | --bop_res_file results/cvpr2021_results/bop_deepim_s0_test_RGBD.csv \ 196 | --coco_res_file results/cvpr2021_results/coco_maskrcnn_s0_test.json 197 | python examples/evaluate_grasp.py \ 198 | --name s1_test \ 199 | --bop_res_file results/cvpr2021_results/bop_deepim_s1_test_RGBD.csv \ 200 | --coco_res_file results/cvpr2021_results/coco_maskrcnn_s1_test.json 201 | python examples/evaluate_grasp.py \ 202 | --name s2_test \ 203 | --bop_res_file results/cvpr2021_results/bop_deepim_s2_test_RGBD.csv \ 204 | --coco_res_file results/cvpr2021_results/coco_maskrcnn_s2_test.json 205 | python examples/evaluate_grasp.py \ 206 | --name s3_test \ 207 | --bop_res_file results/cvpr2021_results/bop_deepim_s3_test_RGBD.csv \ 208 | --coco_res_file results/cvpr2021_results/coco_maskrcnn_s3_test.json 209 | python examples/evaluate_grasp.py \ 210 | --name s0_test \ 211 | --bop_res_file results/cvpr2021_results/bop_poserbpf_s0_test_rgb.csv \ 212 | --coco_res_file results/cvpr2021_results/coco_maskrcnn_s0_test.json 213 | python examples/evaluate_grasp.py \ 214 | --name s1_test \ 215 | --bop_res_file results/cvpr2021_results/bop_poserbpf_s1_test_rgb.csv \ 216 | --coco_res_file results/cvpr2021_results/coco_maskrcnn_s1_test.json 217 | python examples/evaluate_grasp.py \ 218 | --name s2_test \ 219 | --bop_res_file results/cvpr2021_results/bop_poserbpf_s2_test_rgb.csv \ 220 | --coco_res_file results/cvpr2021_results/coco_maskrcnn_s2_test.json 221 | python examples/evaluate_grasp.py \ 222 | --name s3_test \ 223 | --bop_res_file results/cvpr2021_results/bop_poserbpf_s3_test_rgb.csv \ 224 | --coco_res_file results/cvpr2021_results/coco_maskrcnn_s3_test.json 225 | python examples/evaluate_grasp.py \ 226 | --name s0_test \ 227 | --bop_res_file results/cvpr2021_results/bop_poserbpf_s0_test_rgbd.csv \ 228 | --coco_res_file results/cvpr2021_results/coco_maskrcnn_s0_test.json 229 | python examples/evaluate_grasp.py \ 230 | --name s1_test \ 231 | --bop_res_file results/cvpr2021_results/bop_poserbpf_s1_test_rgbd.csv \ 232 | --coco_res_file results/cvpr2021_results/coco_maskrcnn_s1_test.json 233 | python examples/evaluate_grasp.py \ 234 | --name s2_test \ 235 | --bop_res_file results/cvpr2021_results/bop_poserbpf_s2_test_rgbd.csv \ 236 | --coco_res_file results/cvpr2021_results/coco_maskrcnn_s2_test.json 237 | python examples/evaluate_grasp.py \ 238 | --name s3_test \ 239 | --bop_res_file results/cvpr2021_results/bop_poserbpf_s3_test_rgbd.csv \ 240 | --coco_res_file results/cvpr2021_results/coco_maskrcnn_s3_test.json 241 | python examples/evaluate_grasp.py \ 242 | --name s0_test \ 243 | --bop_res_file results/cvpr2021_results/bop_cosypose_s0_test.csv \ 244 | --coco_res_file results/cvpr2021_results/coco_maskrcnn_s0_test.json 245 | python examples/evaluate_grasp.py \ 246 | --name s1_test \ 247 | --bop_res_file results/cvpr2021_results/bop_cosypose_s1_test.csv \ 248 | --coco_res_file results/cvpr2021_results/coco_maskrcnn_s1_test.json 249 | python examples/evaluate_grasp.py \ 250 | --name s2_test \ 251 | --bop_res_file results/cvpr2021_results/bop_cosypose_s2_test.csv \ 252 | --coco_res_file results/cvpr2021_results/coco_maskrcnn_s2_test.json 253 | python examples/evaluate_grasp.py \ 254 | --name s3_test \ 255 | --bop_res_file results/cvpr2021_results/bop_cosypose_s3_test.csv \ 256 | --coco_res_file results/cvpr2021_results/coco_maskrcnn_s3_test.json 257 | -------------------------------------------------------------------------------- /assets/panda_gripper.obj: -------------------------------------------------------------------------------- 1 | v 0.05421821303665638 0.007778378669172524 0.11076592143774033 2 | v 0.039979396890921635 -0.00613801833242178 0.11197762053608895 3 | v 0.042534183906391264 0.010387184098362923 0.0585316962443525 4 | v 0.054274293556809426 -0.00583982700482011 0.11220345135927201 5 | v 0.06640338242053986 -0.01035996433347464 0.05855462865121663 6 | v 0.03997647545023938 0.008730395697057247 0.09456484627127648 7 | v 0.03994319871620974 0.008622935973107815 0.10935282196998597 8 | v 0.05610156148672104 0.010479452088475226 0.07738816600441933 9 | v 0.0652202693372965 0.010388848371803757 0.07758761990964413 10 | v 0.06525340951979161 -0.01040047220885754 0.07743658919036389 11 | v 0.05391758117824793 0.008841173723340033 0.10898935657143594 12 | v 0.05510264165699482 -0.010494819842278959 0.07683558996915818 13 | v 0.03990887965395814 0.005481416825205088 0.11212470989823342 14 | v 0.053988714665174485 -0.008616076782345774 0.1097279139995575 15 | v 0.05429311186075211 0.005277918651700018 0.11224903401136399 16 | v 0.039867356615141035 -0.00869277399033308 0.10856622692346574 17 | v 0.06630941011011601 0.010401263833045956 0.058566509751579725 18 | v 0.04258330418728292 -0.010448622517287731 0.05854680107415188 19 | v -0.05421821303665638 -0.007778378669172525 0.11076592143774033 20 | v -0.039979396890921635 0.00613801833242178 0.11197762053608895 21 | v -0.042534183906391264 -0.010387184098362923 0.0585316962443525 22 | v -0.054274293556809426 0.005839827004820108 0.11220345135927201 23 | v -0.06640338242053986 0.010359964333474636 0.05855462865121663 24 | v -0.03997647545023938 -0.008730395697057247 0.09456484627127648 25 | v -0.03994319871620974 -0.008622935973107815 0.10935282196998597 26 | v -0.05610156148672104 -0.010479452088475227 0.07738816600441933 27 | v -0.0652202693372965 -0.01038884837180376 0.07758761990964413 28 | v -0.06525340951979161 0.010400472208857536 0.07743658919036389 29 | v -0.05391758117824793 -0.008841173723340034 0.10898935657143594 30 | v -0.05510264165699482 0.010494819842278957 0.07683558996915818 31 | v -0.03990887965395814 -0.005481416825205088 0.11212470989823342 32 | v -0.053988714665174485 0.008616076782345772 0.1097279139995575 33 | v -0.05429311186075211 -0.00527791865170002 0.11224903401136399 34 | v -0.039867356615141035 0.00869277399033308 0.10856622692346574 35 | v -0.06630941011011601 -0.01040126383304596 0.058566509751579725 36 | v -0.04258330418728292 0.010448622517287731 0.05854680107415188 37 | v -0.04178430512547493 -0.028010861948132515 -0.0063964021392166615 38 | v -0.09769713878631592 -0.019066786393523216 0.024627480655908585 39 | v -0.09701105952262878 0.01980205439031124 0.020913975313305855 40 | v -0.09099503606557846 -0.00017339896294288337 0.0005767836119048297 41 | v 0.0005787304835394025 -0.031635917723178864 0.005973074585199356 42 | v -0.09544170647859573 -0.021822135895490646 0.021756364032626152 43 | v -0.10028521716594696 0.016218392178416252 0.04586976766586304 44 | v 0.1000911220908165 -0.014017123728990555 0.056026946753263474 45 | v -0.09460194408893585 0.018555866554379463 0.056617286056280136 46 | v -0.09433312714099884 0.021342597901821136 0.008942007087171078 47 | v 0.08797474205493927 -0.0006270006415434182 -0.024961603805422783 48 | v -0.09127645939588547 -0.014421950094401836 0.06568973511457443 49 | v -0.09320925921201706 -0.01604551635682583 0.06361687183380127 50 | v -0.10042587667703629 -0.004494380671530962 0.058119092136621475 51 | v 0.09249575436115265 -0.012987935915589333 0.06571194529533386 52 | v -0.09116631746292114 -0.023555776104331017 0.004607920069247484 53 | v 0.09070632606744766 -0.015471714548766613 0.06519009917974472 54 | v -0.08543447405099869 -0.022732771933078766 -0.004120331723242998 55 | v -0.09289686381816864 -0.023026082664728165 0.009876862168312073 56 | v 0.09790761768817902 -0.01751038245856762 0.05166616290807724 57 | v 0.0005585200269706547 0.03158979117870331 0.006214227061718702 58 | v 0.10119467973709106 0.003602118231356144 -0.012627636082470417 59 | v 0.09665588289499283 0.0004695942625403404 0.06307835876941681 60 | v 0.09384757280349731 0.017607156187295914 0.058448486030101776 61 | v -0.04204234480857849 0.02801508456468582 -0.006349603645503521 62 | v 0.09945143014192581 3.1802206649445e-05 -0.017161205410957336 63 | v 0.04551994055509567 0.0019174328772351146 -0.025659434497356415 64 | v -0.09025220572948456 -0.01676723174750805 0.06433220207691193 65 | v 0.10030148178339005 0.001190581009723246 0.05862313136458397 66 | v -0.07437754422426224 0.024816755205392838 -0.007155700121074915 67 | v 0.10259784758090973 -0.0034295637160539627 -0.00896429643034935 68 | v 0.1027156412601471 0.003494761884212494 -0.008029340766370296 69 | v -0.08589234948158264 -0.02449524775147438 -0.0018327292054891586 70 | v 0.09406288713216782 0.0131387272849679 0.06468318402767181 71 | v -0.08206511288881302 -0.025270354002714157 0.0051970407366752625 72 | v -0.09466791152954102 -0.02065763622522354 0.009537984617054462 73 | v -0.08824997395277023 0.022314293310046196 -0.0019331998191773891 74 | v -0.09747105836868286 -0.0016167220892384648 0.06230733543634415 75 | v 0.09552478045225143 -0.017053674906492233 0.02212120220065117 76 | v -0.08335519582033157 0.022376984357833862 -0.005526112858206034 77 | v -0.09936285763978958 -0.016994841396808624 0.05411478504538536 78 | v 0.0968022570014 0.017033156007528305 0.030322037637233734 79 | v 0.09160291403532028 0.01695552095770836 0.005562833044677973 80 | v -0.09012892097234726 0.016734914854168892 0.06443186104297638 81 | v 0.09177957475185394 0.01571837067604065 0.06491253525018692 82 | v 0.0840023085474968 0.0018427835311740637 -0.02552490122616291 83 | v 0.08797289431095123 0.0030759950168430805 -0.02397872507572174 84 | v 0.09962863475084305 -0.016013137996196747 0.04930143058300018 85 | v -0.09201552718877792 -0.01844867318868637 0.058393169194459915 86 | v -0.08025997132062912 -0.0008337647304870188 -0.007321717217564583 87 | v -0.07576971501111984 -0.025980981066823006 -0.005082232877612114 88 | v -0.10019978880882263 0.015550931915640831 0.05467259883880615 89 | v 0.09941626340150833 0.015804897993803024 0.05497027188539505 90 | v 0.10374269634485245 -1.7281157852266915e-05 -0.00919930450618267 91 | v 0.08254846930503845 -0.0008678357116878033 -0.025870200246572495 92 | v 0.0875329002737999 -0.016812244430184364 -0.0012064524926245213 93 | v 0.08223627507686615 -0.016532698646187782 -0.005118109285831451 94 | v -0.09373555332422256 0.022707808762788773 0.014340022578835487 95 | v -0.09371249377727509 -0.012566703371703625 0.06516847014427185 96 | v -0.07666800171136856 -0.024650242179632187 -0.0069321258924901485 97 | v 0.08927122503519058 0.01713424362242222 0.0636143907904625 98 | v 0.08776598423719406 -0.0032150978222489357 -0.023884393274784088 99 | v -0.09726057201623917 -0.019229214638471603 0.05058842897415161 100 | v -0.09369184076786041 0.020670883357524872 0.04330829158425331 101 | v 0.09740705043077469 0.017585095018148422 0.051984645426273346 102 | v 0.09855398535728455 -0.01663215272128582 0.05473393574357033 103 | v 0.09344169497489929 -0.014617033302783966 0.06450004875659943 104 | v 0.08296618610620499 0.00381033169105649 -0.024449335411190987 105 | v -0.09092690050601959 -0.021324951201677322 0.0009798021055758 106 | v -0.09280849248170853 -0.0001125619382946752 0.06596215069293976 107 | v -0.0917946845293045 0.021482910960912704 0.0026841284707188606 108 | v 0.09998264163732529 -0.009323876351118088 0.058489199727773666 109 | v 0.08358591049909592 -0.0036368216387927532 -0.024606257677078247 110 | v 0.1001875177025795 0.012505676597356796 0.056894149631261826 111 | v -0.09290558844804764 0.015396904200315475 0.06455627083778381 112 | v 0.0851321741938591 0.016558213159441948 -0.0038727361243218184 113 | v 0.09294531494379044 -0.0005056463996879756 0.06595310568809509 114 | v 0.10115781426429749 -0.0036167786456644535 -0.012610324658453465 115 | v -0.07790137827396393 0.02295910380780697 -0.007399038877338171 116 | v -0.0857401043176651 0.024729391559958458 -0.0012316935462877154 117 | v -0.10016821324825287 -0.014623090624809265 0.055734917521476746 118 | v -0.09951794147491455 -0.018192630261182785 0.043814171105623245 119 | v 0.09070031344890594 -0.017254667356610298 0.0630820095539093 120 | v 0.0919061228632927 -0.016804175451397896 0.006295484956353903 121 | v 0.09953752160072327 0.016230100765824318 0.051584091037511826 122 | v -0.08118050545454025 0.025447947904467583 0.0035006047692149878 123 | v -0.09906721860170364 0.017129460349678993 0.05430515855550766 124 | v -0.08656162023544312 -0.00033731618896126747 -0.004163281060755253 125 | v -0.09461534768342972 -0.00031412430689670146 0.007574658375233412 126 | v -0.07529757916927338 0.026034310460090637 -0.005030847620218992 127 | v -0.08017436414957047 -0.02276112325489521 -0.006909539457410574 128 | v 0.0018608596874400973 0.03161578252911568 0.0011797059560194612 129 | v 0.0458698496222496 -0.0015001518186181784 -0.02592480182647705 130 | v -0.0817025899887085 0.024515172466635704 -0.005051423329859972 131 | v -0.10003473609685898 0.009941554628312588 0.05834079533815384 132 | v -0.09267213940620422 0.013539588078856468 0.0656878799200058 133 | v -0.09849356859922409 0.019268833100795746 0.0449417382478714 134 | v -0.09040140360593796 0.023869164288043976 0.004368236754089594 135 | v 0.0019865017384290695 -0.031597502529621124 0.001152931246906519 136 | v -0.09849606454372406 -1.593970591784455e-05 0.027081793174147606 137 | v 0.10398972034454346 -4.109224391868338e-05 0.005690876394510269 138 | v 0.09192700684070587 0.01342480443418026 0.06573130935430527 139 | f 5 18 3 140 | f 3 8 17 141 | f 17 8 9 142 | f 18 6 3 143 | f 10 9 4 144 | f 3 17 5 145 | f 14 12 10 146 | f 12 16 18 147 | f 14 16 12 148 | f 7 3 6 149 | f 18 16 6 150 | f 16 13 6 151 | f 10 12 18 152 | f 10 18 5 153 | f 10 5 9 154 | f 9 15 4 155 | f 1 15 9 156 | f 7 13 1 157 | f 3 7 8 158 | f 10 4 14 159 | f 5 17 9 160 | f 8 11 9 161 | f 11 1 9 162 | f 4 13 2 163 | f 15 13 4 164 | f 16 2 13 165 | f 2 16 14 166 | f 13 7 6 167 | f 1 13 15 168 | f 7 11 8 169 | f 2 14 4 170 | f 1 11 7 171 | f 23 36 21 172 | f 21 26 35 173 | f 35 26 27 174 | f 36 24 21 175 | f 28 27 22 176 | f 21 35 23 177 | f 32 30 28 178 | f 30 34 36 179 | f 32 34 30 180 | f 25 21 24 181 | f 36 34 24 182 | f 34 31 24 183 | f 28 30 36 184 | f 28 36 23 185 | f 28 23 27 186 | f 27 33 22 187 | f 19 33 27 188 | f 25 31 19 189 | f 21 25 26 190 | f 28 22 32 191 | f 23 35 27 192 | f 26 29 27 193 | f 29 19 27 194 | f 22 31 20 195 | f 33 31 22 196 | f 34 20 31 197 | f 20 34 32 198 | f 31 25 24 199 | f 19 31 33 200 | f 25 29 26 201 | f 20 32 22 202 | f 19 29 25 203 | f 80 97 57 204 | f 75 56 135 205 | f 87 135 41 206 | f 78 128 101 207 | f 128 57 101 208 | f 45 80 57 209 | f 120 135 92 210 | f 41 135 56 211 | f 75 135 120 212 | f 121 137 90 213 | f 114 67 120 214 | f 71 87 41 215 | f 37 135 87 216 | f 106 95 48 217 | f 119 53 64 218 | f 79 128 78 219 | f 97 60 57 220 | f 60 101 57 221 | f 110 137 121 222 | f 88 133 43 223 | f 63 61 104 224 | f 100 57 122 225 | f 135 109 93 226 | f 119 41 56 227 | f 44 84 137 228 | f 114 92 98 229 | f 106 48 51 230 | f 51 113 106 231 | f 132 106 113 232 | f 113 138 132 233 | f 71 41 99 234 | f 78 101 121 235 | f 104 128 112 236 | f 68 79 78 237 | f 128 104 61 238 | f 45 100 133 239 | f 100 45 57 240 | f 100 122 134 241 | f 100 94 133 242 | f 94 100 134 243 | f 128 61 126 244 | f 128 126 122 245 | f 128 122 57 246 | f 66 61 63 247 | f 129 86 115 248 | f 86 129 127 249 | f 115 66 63 250 | f 109 135 37 251 | f 62 114 98 252 | f 93 109 98 253 | f 92 93 98 254 | f 92 135 93 255 | f 102 119 56 256 | f 62 58 90 257 | f 84 102 56 258 | f 137 84 90 259 | f 74 106 132 260 | f 106 74 95 261 | f 80 132 81 262 | f 138 81 132 263 | f 119 85 41 264 | f 85 119 64 265 | f 85 99 41 266 | f 47 62 98 267 | f 83 104 112 268 | f 81 97 80 269 | f 43 133 39 270 | f 131 74 132 271 | f 133 123 45 272 | f 125 136 39 273 | f 126 61 66 274 | f 124 73 76 275 | f 105 54 69 276 | f 82 129 63 277 | f 129 96 127 278 | f 75 84 56 279 | f 68 121 90 280 | f 62 90 114 281 | f 87 71 69 282 | f 48 53 51 283 | f 64 53 48 284 | f 113 70 138 285 | f 44 137 108 286 | f 38 136 125 287 | f 74 50 95 288 | f 43 117 50 289 | f 112 79 58 290 | f 58 79 68 291 | f 58 68 90 292 | f 81 89 60 293 | f 89 110 121 294 | f 70 89 81 295 | f 89 70 110 296 | f 78 121 68 297 | f 89 121 101 298 | f 101 60 89 299 | f 50 74 131 300 | f 63 104 82 301 | f 125 39 46 302 | f 136 43 39 303 | f 127 54 124 304 | f 73 40 107 305 | f 124 76 86 306 | f 86 76 115 307 | f 63 129 115 308 | f 37 96 129 309 | f 37 87 96 310 | f 130 73 116 311 | f 109 37 129 312 | f 72 38 125 313 | f 117 95 50 314 | f 79 112 128 315 | f 65 70 59 316 | f 88 43 131 317 | f 131 43 50 318 | f 134 46 94 319 | f 123 133 88 320 | f 132 111 131 321 | f 111 88 131 322 | f 111 45 123 323 | f 54 105 124 324 | f 116 126 130 325 | f 46 134 107 326 | f 105 125 40 327 | f 40 125 107 328 | f 91 129 82 329 | f 109 129 91 330 | f 92 114 120 331 | f 84 44 102 332 | f 67 90 84 333 | f 67 75 120 334 | f 113 59 70 335 | f 65 59 108 336 | f 137 110 65 337 | f 65 108 137 338 | f 52 69 71 339 | f 83 62 47 340 | f 47 82 83 341 | f 62 83 58 342 | f 112 58 83 343 | f 97 81 60 344 | f 70 81 138 345 | f 70 65 110 346 | f 82 104 83 347 | f 111 123 88 348 | f 111 132 80 349 | f 45 111 80 350 | f 125 46 107 351 | f 39 94 46 352 | f 94 39 133 353 | f 115 130 66 354 | f 126 116 122 355 | f 116 134 122 356 | f 105 40 124 357 | f 124 40 73 358 | f 86 127 124 359 | f 84 75 67 360 | f 67 114 90 361 | f 103 108 51 362 | f 108 59 51 363 | f 44 103 102 364 | f 53 103 51 365 | f 55 42 72 366 | f 118 43 136 367 | f 117 43 118 368 | f 52 105 69 369 | f 91 82 47 370 | f 73 130 76 371 | f 126 66 130 372 | f 73 107 134 373 | f 116 73 134 374 | f 44 108 103 375 | f 51 59 113 376 | f 119 103 53 377 | f 49 99 85 378 | f 85 64 49 379 | f 77 99 49 380 | f 52 55 72 381 | f 99 42 71 382 | f 42 55 71 383 | f 55 52 71 384 | f 52 72 105 385 | f 117 49 95 386 | f 64 48 49 387 | f 130 115 76 388 | f 69 54 96 389 | f 96 54 127 390 | f 96 87 69 391 | f 77 117 118 392 | f 72 42 38 393 | f 99 118 42 394 | f 118 99 77 395 | f 105 72 125 396 | f 117 77 49 397 | f 48 95 49 398 | f 98 91 47 399 | f 119 102 103 400 | f 38 118 136 401 | f 98 109 91 402 | f 118 38 42 403 | -------------------------------------------------------------------------------- /dex_ycb_toolkit/sequence_loader.py: -------------------------------------------------------------------------------- 1 | # DexYCB Toolkit 2 | # Copyright (C) 2021 NVIDIA Corporation 3 | # Licensed under the GNU General Public License v3.0 [see LICENSE for details] 4 | 5 | """DexYCB sequence loader.""" 6 | 7 | import torch 8 | import os 9 | import yaml 10 | import numpy as np 11 | import cv2 12 | 13 | from scipy.spatial.transform import Rotation as Rot 14 | 15 | from .layers.ycb_group_layer import YCBGroupLayer 16 | from .layers.mano_group_layer import MANOGroupLayer 17 | from .layers.ycb_layer import dcm2rv, rv2dcm 18 | 19 | 20 | class SequenceLoader(): 21 | """DexYCB sequence loader.""" 22 | 23 | def __init__( 24 | self, 25 | name, 26 | device='cuda:0', 27 | preload=True, 28 | app='viewer', 29 | ): 30 | """Constructor. 31 | 32 | Args: 33 | name: Sequence name. 34 | device: A torch.device string argument. The specified device is used only 35 | for certain data loading computations, but not storing the loaded data. 36 | Currently the loaded data is always stored as numpy arrays on CPU. 37 | preload: Whether to preload the point cloud or load it online. 38 | app: 'viewer' or 'renderer'. 39 | """ 40 | assert device in ('cuda', 'cpu') or device.split(':')[0] == 'cuda' 41 | assert app in ('viewer', 'renderer') 42 | self._name = name 43 | self._device = torch.device(device) 44 | self._preload = preload 45 | self._app = app 46 | 47 | assert 'DEX_YCB_DIR' in os.environ, "environment variable 'DEX_YCB_DIR' is not set" 48 | self._dex_ycb_dir = os.environ['DEX_YCB_DIR'] 49 | 50 | # Load meta. 51 | meta_file = self._dex_ycb_dir + '/' + self._name + "/meta.yml" 52 | with open(meta_file, 'r') as f: 53 | meta = yaml.load(f, Loader=yaml.FullLoader) 54 | 55 | self._serials = meta['serials'] 56 | self._h = 480 57 | self._w = 640 58 | self._num_cameras = len(self._serials) 59 | self._data_dir = [ 60 | self._dex_ycb_dir + '/' + self._name + '/' + s for s in self._serials 61 | ] 62 | self._color_prefix = "color_" 63 | self._depth_prefix = "aligned_depth_to_color_" 64 | self._label_prefix = "labels_" 65 | self._num_frames = meta['num_frames'] 66 | self._ycb_ids = meta['ycb_ids'] 67 | self._mano_sides = meta['mano_sides'] 68 | 69 | # Load intrinsics. 70 | def intr_to_K(x): 71 | return torch.tensor( 72 | [[x['fx'], 0.0, x['ppx']], [0.0, x['fy'], x['ppy']], [0.0, 0.0, 1.0]], 73 | dtype=torch.float32, 74 | device=self._device) 75 | 76 | self._K = [] 77 | for s in self._serials: 78 | intr_file = self._dex_ycb_dir + "/calibration/intrinsics/" + s + '_' + str( 79 | self._w) + 'x' + str(self._h) + ".yml" 80 | with open(intr_file, 'r') as f: 81 | intr = yaml.load(f, Loader=yaml.FullLoader) 82 | K = intr_to_K(intr['color']) 83 | self._K.append(K) 84 | self._K_inv = [torch.inverse(k) for k in self._K] 85 | 86 | # Load extrinsics. 87 | extr_file = self._dex_ycb_dir + "/calibration/extrinsics_" + meta[ 88 | 'extrinsics'] + "/extrinsics.yml" 89 | with open(extr_file, 'r') as f: 90 | extr = yaml.load(f, Loader=yaml.FullLoader) 91 | T = extr['extrinsics'] 92 | T = { 93 | s: torch.tensor(T[s], dtype=torch.float32, 94 | device=self._device).view(3, 4) for s in T 95 | } 96 | self._R = [T[s][:, :3] for s in self._serials] 97 | self._t = [T[s][:, 3] for s in self._serials] 98 | self._R_inv = [torch.inverse(r) for r in self._R] 99 | self._t_inv = [torch.mv(r, -t) for r, t in zip(self._R_inv, self._t)] 100 | self._master_intrinsics = self._K[[ 101 | i for i, s in enumerate(self._serials) if s == extr['master'] 102 | ][0]].cpu().numpy() 103 | self._tag_R = T['apriltag'][:, :3] 104 | self._tag_t = T['apriltag'][:, 3] 105 | self._tag_R_inv = torch.inverse(self._tag_R) 106 | self._tag_t_inv = torch.mv(self._tag_R_inv, -self._tag_t) 107 | self._tag_lim = [-0.00, +1.20, -0.10, +0.70, -0.10, +0.70] 108 | 109 | # Compute texture coordinates. 110 | y, x = torch.meshgrid(torch.arange(self._h), torch.arange(self._w)) 111 | x = x.float() 112 | y = y.float() 113 | s = torch.stack((x / (self._w - 1), y / (self._h - 1)), dim=2) 114 | self._pcd_tex_coord = [s.numpy()] * self._num_cameras 115 | 116 | # Compute rays. 117 | self._p = [] 118 | ones = torch.ones((self._h, self._w), dtype=torch.float32) 119 | xy1s = torch.stack((x, y, ones), dim=2).view(self._w * self._h, 3).t() 120 | xy1s = xy1s.to(self._device) 121 | for c in range(self._num_cameras): 122 | p = torch.mm(self._K_inv[c], xy1s) 123 | self._p.append(p) 124 | 125 | # Load point cloud. 126 | if self._preload: 127 | print('Preloading point cloud') 128 | self._color = [] 129 | self._depth = [] 130 | for c in range(self._num_cameras): 131 | color = [] 132 | depth = [] 133 | for i in range(self._num_frames): 134 | rgb, d = self._load_frame_rgbd(c, i) 135 | color.append(rgb) 136 | depth.append(d) 137 | self._color.append(color) 138 | self._depth.append(depth) 139 | self._color = np.array(self._color, dtype=np.uint8) 140 | self._depth = np.array(self._depth, dtype=np.uint16) 141 | self._pcd_rgb = [x for x in self._color] 142 | self._pcd_vert = [] 143 | self._pcd_mask = [] 144 | for c in range(self._num_cameras): 145 | p, m = self._deproject_depth_and_filter_points(self._depth[c], c) 146 | self._pcd_vert.append(p) 147 | self._pcd_mask.append(m) 148 | else: 149 | print('Loading point cloud online') 150 | self._pcd_rgb = [ 151 | np.zeros((self._h, self._w, 3), dtype=np.uint8) 152 | for _ in range(self._num_cameras) 153 | ] 154 | self._pcd_vert = [ 155 | np.zeros((self._h, self._w, 3), dtype=np.float32) 156 | for _ in range(self._num_cameras) 157 | ] 158 | self._pcd_mask = [ 159 | np.zeros((self._h, self._w), dtype=np.bool) 160 | for _ in range(self._num_cameras) 161 | ] 162 | 163 | # Create YCB group layer. 164 | self._ycb_group_layer = YCBGroupLayer(self._ycb_ids).to(self._device) 165 | 166 | self._ycb_model_dir = self._dex_ycb_dir + "/models" 167 | self._ycb_count = self._ycb_group_layer.count 168 | self._ycb_material = self._ycb_group_layer.material 169 | self._ycb_tex_coords = self._ycb_group_layer.tex_coords 170 | 171 | # Create MANO group layer. 172 | mano_betas = [] 173 | for m in meta['mano_calib']: 174 | mano_calib_file = self._dex_ycb_dir + "/calibration/mano_" + m + "/mano.yml" 175 | with open(mano_calib_file, 'r') as f: 176 | mano_calib = yaml.load(f, Loader=yaml.FullLoader) 177 | betas = np.array(mano_calib['betas'], dtype=np.float32) 178 | mano_betas.append(betas) 179 | 180 | self._mano_group_layer = MANOGroupLayer(self._mano_sides, 181 | mano_betas).to(self._device) 182 | 183 | # Prepare data for viewer. 184 | if app == 'viewer': 185 | s = np.cumsum([0] + self._ycb_group_layer.count[:-1]) 186 | e = np.cumsum(self._ycb_group_layer.count) 187 | self._ycb_seg = list(zip(s, e)) 188 | 189 | ycb_file = self._dex_ycb_dir + '/' + self._name + "/pose.npz" 190 | data = np.load(ycb_file) 191 | ycb_pose = data['pose_y'] 192 | i = np.any(ycb_pose != [0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0], axis=2) 193 | pose = ycb_pose.reshape(-1, 7) 194 | v, n = self.transform_ycb(pose) 195 | self._ycb_vert = [ 196 | np.zeros((self._num_frames, n, 3), dtype=np.float32) 197 | for n in self._ycb_count 198 | ] 199 | self._ycb_norm = [ 200 | np.zeros((self._num_frames, n, 3), dtype=np.float32) 201 | for n in self._ycb_count 202 | ] 203 | for o in range(self._ycb_group_layer.num_obj): 204 | io = i[:, o] 205 | self._ycb_vert[o][io] = v[io, self._ycb_seg[o][0]:self._ycb_seg[o][1]] 206 | self._ycb_norm[o][io] = n[io, self._ycb_seg[o][0]:self._ycb_seg[o][1]] 207 | 208 | mano_file = self._dex_ycb_dir + '/' + self._name + "/pose.npz" 209 | data = np.load(mano_file) 210 | mano_pose = data['pose_m'] 211 | i = np.any(mano_pose != 0.0, axis=2) 212 | pose = torch.from_numpy(mano_pose).to(self._device) 213 | pose = pose.view(-1, self._mano_group_layer.num_obj * 51) 214 | verts, _ = self._mano_group_layer(pose) 215 | # Numpy array is faster than PyTorch Tensor here. 216 | verts = verts.cpu().numpy() 217 | f = self._mano_group_layer.f.cpu().numpy() 218 | v = verts[:, f.ravel()] 219 | n = np.cross(v[:, 1::3, :] - v[:, 0::3, :], v[:, 2::3, :] - v[:, 1::3, :]) 220 | n = np.repeat(n, 3, axis=1) 221 | l = verts[:, f[:, [0, 1, 1, 2, 2, 0]].ravel(), :] 222 | self._mano_vert = [ 223 | np.zeros((self._num_frames, 4614, 3), dtype=np.float32) 224 | for _ in range(self._mano_group_layer.num_obj) 225 | ] 226 | self._mano_norm = [ 227 | np.zeros((self._num_frames, 4614, 3), dtype=np.float32) 228 | for _ in range(self._mano_group_layer.num_obj) 229 | ] 230 | self._mano_line = [ 231 | np.zeros((self._num_frames, 9228, 3), dtype=np.float32) 232 | for _ in range(self._mano_group_layer.num_obj) 233 | ] 234 | for o in range(self._mano_group_layer.num_obj): 235 | io = i[:, o] 236 | self._mano_vert[o][io] = v[io, 4614 * o:4614 * (o + 1), :] 237 | self._mano_norm[o][io] = n[io, 4614 * o:4614 * (o + 1), :] 238 | self._mano_line[o][io] = l[io, 9228 * o:9228 * (o + 1), :] 239 | 240 | # Prepare data for renderer. 241 | if app == 'renderer': 242 | self._ycb_pose = [] 243 | self._mano_vert = [] 244 | self._mano_joint_3d = [] 245 | 246 | for c in range(self._num_cameras): 247 | ycb_pose = [] 248 | mano_pose = [] 249 | mano_joint_3d = [] 250 | for i in range(self._num_frames): 251 | label_file = self._data_dir[ 252 | c] + '/' + self._label_prefix + "{:06d}.npz".format(i) 253 | label = np.load(label_file) 254 | pose_y = np.hstack((label['pose_y'], 255 | np.array([[[0, 0, 0, 1]]] * len(label['pose_y']), 256 | dtype=np.float32))) 257 | pose_m = label['pose_m'] 258 | joint_3d = label['joint_3d'] 259 | ycb_pose.append(pose_y) 260 | mano_pose.append(pose_m) 261 | mano_joint_3d.append(joint_3d) 262 | ycb_pose = np.array(ycb_pose, dtype=np.float32) 263 | mano_pose = np.array(mano_pose, dtype=np.float32) 264 | mano_joint_3d = np.array(mano_joint_3d, dtype=np.float32) 265 | self._ycb_pose.append(ycb_pose) 266 | self._mano_joint_3d.append(mano_joint_3d) 267 | 268 | i = np.any(mano_pose != 0.0, axis=2) 269 | pose = torch.from_numpy(mano_pose).to(self._device) 270 | pose = pose.view(-1, self._mano_group_layer.num_obj * 51) 271 | verts, _ = self._mano_group_layer(pose) 272 | verts = verts.cpu().numpy() 273 | mano_vert = [ 274 | np.zeros((self._num_frames, 778, 3), dtype=np.float32) 275 | for _ in range(self._mano_group_layer.num_obj) 276 | ] 277 | for o in range(self._mano_group_layer.num_obj): 278 | io = i[:, o] 279 | mano_vert[o][io] = verts[io, 778 * o:778 * (o + 1), :] 280 | self._mano_vert.append(mano_vert) 281 | 282 | self._frame = -1 283 | 284 | def _load_frame_rgbd(self, c, i): 285 | """Loads an RGB-D frame. 286 | 287 | Args: 288 | c: Camera index. 289 | i: Frame index. 290 | 291 | Returns: 292 | color: A unit8 numpy array of shape [H, W, 3] containing the color image. 293 | depth: A uint16 numpy array of shape [H, W] containing the depth image. 294 | """ 295 | color_file = self._data_dir[ 296 | c] + '/' + self._color_prefix + "{:06d}.jpg".format(i) 297 | color = cv2.imread(color_file) 298 | color = color[:, :, ::-1] 299 | depth_file = self._data_dir[ 300 | c] + '/' + self._depth_prefix + "{:06d}.png".format(i) 301 | depth = cv2.imread(depth_file, cv2.IMREAD_ANYDEPTH) 302 | return color, depth 303 | 304 | def _deproject_depth_and_filter_points(self, d, c): 305 | """Deprojects a depth image to point cloud and filters points. 306 | 307 | Args: 308 | d: A uint16 numpy array of shape [F, H, W] or [H, W] containing the depth 309 | image in millimeters. 310 | c: Camera index. 311 | 312 | Returns: 313 | p: A float32 numpy array of shape [F, H, W, 3] or [H, W, 3] containing the 314 | point cloud. 315 | m: A bool numpy array of shape [F, H, W] or [H, W] containing the mask for 316 | points within the tag cooridnate limit. 317 | """ 318 | nd = d.ndim 319 | d = d.astype(np.float32) / 1000 320 | d = torch.from_numpy(d).to(self._device) 321 | p = torch.mul( 322 | d.view(1, -1, self._w * self._h).expand(3, -1, -1), 323 | self._p[c].unsqueeze(1)) 324 | p = torch.addmm(self._t[c].unsqueeze(1), self._R[c], p.view(3, -1)) 325 | p_tag = torch.addmm(self._tag_t_inv.unsqueeze(1), self._tag_R_inv, p) 326 | mx1 = p_tag[0, :] > self._tag_lim[0] 327 | mx2 = p_tag[0, :] < self._tag_lim[1] 328 | my1 = p_tag[1, :] > self._tag_lim[2] 329 | my2 = p_tag[1, :] < self._tag_lim[3] 330 | mz1 = p_tag[2, :] > self._tag_lim[4] 331 | mz2 = p_tag[2, :] < self._tag_lim[5] 332 | m = mx1 & mx2 & my1 & my2 & mz1 & mz2 333 | p = p.t().view(-1, self._h, self._w, 3) 334 | m = m.view(-1, self._h, self._w) 335 | if nd == 2: 336 | p = p.squeeze(0) 337 | m = m.squeeze(0) 338 | p = p.cpu().numpy() 339 | m = m.cpu().numpy() 340 | return p, m 341 | 342 | def transform_ycb(self, 343 | pose, 344 | c=None, 345 | camera_to_world=True, 346 | run_ycb_group_layer=True, 347 | return_trans_mat=False): 348 | """Transforms poses in SE3 between world and camera frames. 349 | 350 | Args: 351 | pose: A float32 numpy array of shape [N, 7] or [N, 6] containing the 352 | poses. Each row contains one pose represented by rotation in quaternion 353 | (x, y, z, w) or rotation vector and translation. 354 | c: Camera index. 355 | camera_to_world: Whether from camera to world or from world to camera. 356 | run_ycb_group_layer: Whether to return vertices and normals by running the 357 | YCB group layer or to return poses. 358 | return_trans_mat: Whether to return poses in transformation matrices. 359 | 360 | Returns: 361 | If run_ycb_group_layer is True: 362 | v: A float32 numpy array of shape [F, V, 3] containing the vertices. 363 | n: A float32 numpy array of shape [F, V, 3] containing the normals. 364 | else: 365 | A float32 numpy array of shape [N, 6] containing the transformed poses. 366 | """ 367 | if pose.shape[1] == 7: 368 | q = pose[:, :4] 369 | t = pose[:, 4:] 370 | R = Rot.from_quat(q).as_dcm().astype(np.float32) 371 | R = torch.from_numpy(R).to(self._device) 372 | t = torch.from_numpy(t).to(self._device) 373 | if pose.shape[1] == 6: 374 | r = pose[:, :3] 375 | t = pose[:, 3:] 376 | r = torch.from_numpy(r).to(self._device) 377 | t = torch.from_numpy(t).to(self._device) 378 | R = rv2dcm(r) 379 | if c is not None: 380 | if camera_to_world: 381 | R_c = self._R[c] 382 | t_c = self._t[c] 383 | else: 384 | R_c = self._R_inv[c] 385 | t_c = self._t_inv[c] 386 | R = torch.bmm(R_c.expand(R.size(0), -1, -1), R) 387 | t = torch.addmm(t_c, t, R_c.t()) 388 | if run_ycb_group_layer or not return_trans_mat: 389 | r = dcm2rv(R) 390 | p = torch.cat([r, t], dim=1) 391 | else: 392 | p = torch.cat([R, t.unsqueeze(2)], dim=2) 393 | p = torch.cat([ 394 | p, 395 | torch.tensor([[[0, 0, 0, 1]]] * R.size(0), 396 | dtype=torch.float32, 397 | device=self._device) 398 | ], 399 | dim=1) 400 | if run_ycb_group_layer: 401 | p = p.view(-1, self._ycb_group_layer.num_obj * 6) 402 | v, n = self._ycb_group_layer(p) 403 | v = v[:, self._ycb_group_layer.f.view(-1)] 404 | n = n[:, self._ycb_group_layer.f.view(-1)] 405 | v = v.cpu().numpy() 406 | n = n.cpu().numpy() 407 | return v, n 408 | else: 409 | p = p.cpu().numpy() 410 | return p 411 | 412 | @property 413 | def serials(self): 414 | return self._serials 415 | 416 | @property 417 | def num_cameras(self): 418 | return self._num_cameras 419 | 420 | @property 421 | def num_frames(self): 422 | return self._num_frames 423 | 424 | @property 425 | def dimensions(self): 426 | return self._w, self._h 427 | 428 | @property 429 | def ycb_ids(self): 430 | return self._ycb_ids 431 | 432 | @property 433 | def K(self): 434 | return self._K 435 | 436 | @property 437 | def master_intrinsics(self): 438 | return self._master_intrinsics 439 | 440 | def step(self): 441 | """Steps the frame.""" 442 | self._frame = (self._frame + 1) % self._num_frames 443 | if not self._preload: 444 | self._update_pcd() 445 | 446 | def _update_pcd(self): 447 | """Updates the point cloud.""" 448 | for c in range(self._num_cameras): 449 | rgb, d = self._load_frame_rgbd(c, self._frame) 450 | p, m = self._deproject_depth_and_filter_points(d, c) 451 | self._pcd_rgb[c][:] = rgb 452 | self._pcd_vert[c][:] = p 453 | self._pcd_mask[c][:] = m 454 | 455 | @property 456 | def pcd_rgb(self): 457 | if self._preload: 458 | return [x[self._frame] for x in self._pcd_rgb] 459 | else: 460 | return self._pcd_rgb 461 | 462 | @property 463 | def pcd_vert(self): 464 | if self._preload: 465 | return [x[self._frame] for x in self._pcd_vert] 466 | else: 467 | return self._pcd_vert 468 | 469 | @property 470 | def pcd_tex_coord(self): 471 | return self._pcd_tex_coord 472 | 473 | @property 474 | def pcd_mask(self): 475 | if self._preload: 476 | return [x[self._frame] for x in self._pcd_mask] 477 | else: 478 | return self._pcd_mask 479 | 480 | @property 481 | def ycb_group_layer(self): 482 | return self._ycb_group_layer 483 | 484 | @property 485 | def num_ycb(self): 486 | return self._ycb_group_layer.num_obj 487 | 488 | @property 489 | def ycb_model_dir(self): 490 | return self._ycb_model_dir 491 | 492 | @property 493 | def ycb_count(self): 494 | return self._ycb_count 495 | 496 | @property 497 | def ycb_material(self): 498 | return self._ycb_material 499 | 500 | @property 501 | def ycb_pose(self): 502 | if self._app == 'viewer': 503 | return None 504 | if self._app == 'renderer': 505 | return [x[self._frame] for x in self._ycb_pose] 506 | 507 | @property 508 | def ycb_vert(self): 509 | if self._app == 'viewer': 510 | return [x[self._frame] for x in self._ycb_vert] 511 | if self._app == 'renderer': 512 | return None 513 | 514 | @property 515 | def ycb_norm(self): 516 | if self._app == 'viewer': 517 | return [x[self._frame] for x in self._ycb_norm] 518 | if self._app == 'renderer': 519 | return None 520 | 521 | @property 522 | def ycb_tex_coords(self): 523 | return self._ycb_tex_coords 524 | 525 | @property 526 | def mano_group_layer(self): 527 | return self._mano_group_layer 528 | 529 | @property 530 | def num_mano(self): 531 | return self._mano_group_layer.num_obj 532 | 533 | @property 534 | def mano_vert(self): 535 | if self._app == 'viewer': 536 | return [x[self._frame] for x in self._mano_vert] 537 | if self._app == 'renderer': 538 | return [[y[self._frame] for y in x] for x in self._mano_vert] 539 | 540 | @property 541 | def mano_norm(self): 542 | if self._app == 'viewer': 543 | return [x[self._frame] for x in self._mano_norm] 544 | if self._app == 'renderer': 545 | return None 546 | 547 | @property 548 | def mano_line(self): 549 | if self._app == 'viewer': 550 | return [x[self._frame] for x in self._mano_line] 551 | if self._app == 'renderer': 552 | return None 553 | 554 | @property 555 | def mano_joint_3d(self): 556 | if self._app == 'viewer': 557 | return None 558 | if self._app == 'renderer': 559 | return [x[self._frame] for x in self._mano_joint_3d] 560 | -------------------------------------------------------------------------------- /dex_ycb_toolkit/window.py: -------------------------------------------------------------------------------- 1 | # DexYCB Toolkit 2 | # Copyright (C) 2021 NVIDIA Corporation 3 | # Licensed under the GNU General Public License v3.0 [see LICENSE for details] 4 | 5 | """Interactive 3D scene viewer using pyglet. 6 | 7 | Functions and classes are largely derived from 8 | https://github.com/IntelRealSense/librealsense/blob/81d469db173dd682d3bada9bd7c7570db0f7cf76/wrappers/python/examples/pyglet_pointcloud_viewer.py 9 | 10 | Usage of class Window: 11 | ------ 12 | Mouse: 13 | Drag with left button to rotate around pivot (thick small axes), 14 | with right button to translate and the wheel to zoom. 15 | 16 | Keyboard: 17 | [p] Pause 18 | [r] Reset View 19 | [z] Toggle point scaling 20 | [x] Toggle point distance attenuation 21 | [l] Toggle lighting 22 | [1/2/3/...] Toggle camera switch 23 | [k] Toggle point mask 24 | [m] Toggle YCB/MANO mesh 25 | [SPACE] Step frame during pause 26 | [s] Save PNG (./out.png) 27 | [q/ESC] Quit 28 | """ 29 | 30 | import numpy as np 31 | import math 32 | import pyglet 33 | import os 34 | import logging 35 | 36 | from pyglet.gl import * 37 | 38 | 39 | # https://stackoverflow.com/a/6802723 40 | def rotation_matrix(axis, theta): 41 | """Returns the rotation matrix associated with counterclockwise rotation about 42 | the given axis by theta radians. 43 | 44 | Args: 45 | axis: Axis represented by a tuple (x, y, z). 46 | theta: Theta in radians. 47 | 48 | Returns: 49 | A float64 numpy array of shape [3, 3] containing the rotation matrix. 50 | """ 51 | axis = np.asarray(axis) 52 | axis = axis / math.sqrt(np.dot(axis, axis)) 53 | a = math.cos(theta / 2.0) 54 | b, c, d = -axis * math.sin(theta / 2.0) 55 | aa, bb, cc, dd = a * a, b * b, c * c, d * d 56 | bc, ad, ac, ab, bd, cd = b * c, a * d, a * c, a * b, b * d, c * d 57 | return np.array([[aa + bb - cc - dd, 2 * (bc + ad), 2 * (bd - ac)], 58 | [2 * (bc - ad), aa + cc - bb - dd, 2 * (cd + ab)], 59 | [2 * (bd + ac), 2 * (cd - ab), aa + dd - bb - cc]]) 60 | 61 | 62 | class AppState: 63 | """Viewer window app state.""" 64 | 65 | def __init__(self, *args, **kwargs): 66 | """Constructor. 67 | 68 | Args: 69 | args: Variable length argument list. 70 | kwargs: Arbitrary keyword arguments. 71 | """ 72 | self.pitch, self.yaw = math.radians(-10), math.radians(-15) 73 | self.translation = np.array([0, 0, 1], np.float32) 74 | self.distance = 2 75 | self.mouse_btns = [False, False, False] 76 | self.paused = False 77 | self.scale = True 78 | self.attenuation = False 79 | self.lighting = False 80 | self.camera_off = [False] * kwargs['num_cameras'] 81 | self.mask = 0 82 | self.model_off = False 83 | 84 | def reset(self): 85 | """Resets the app to an initial state.""" 86 | self.pitch, self.yaw, self.distance = 0, 0, 2 87 | self.translation[:] = 0, 0, 1 88 | 89 | @property 90 | def rotation(self): 91 | Rx = rotation_matrix((1, 0, 0), math.radians(-self.pitch)) 92 | Ry = rotation_matrix((0, 1, 0), math.radians(-self.yaw)) 93 | return np.dot(Ry, Rx).astype(np.float32) 94 | 95 | 96 | def reset_pyglet_resource_path(path): 97 | """Resets pyglet's resource path. 98 | 99 | Args: 100 | path: Path to be reset to. 101 | """ 102 | if not os.path.isabs(path): 103 | path = os.path.abspath(path) 104 | pyglet.resource.path = [path] 105 | pyglet.resource.reindex() 106 | 107 | 108 | class Material(pyglet.graphics.Group): 109 | """Material.""" 110 | 111 | def __init__(self, material, **kwargs): 112 | """Constructor. 113 | 114 | Args: 115 | material: A material object loaded from an OBJ file. 116 | kwargs: Arbitrary keyword arguments. 117 | """ 118 | super(Material, self).__init__(**kwargs) 119 | self.material = material 120 | self.texture = None 121 | self.texture_name = None 122 | 123 | if material.texture_path is not None: 124 | texture_name = os.path.relpath(self.material.texture_path, 125 | pyglet.resource.path[0]) 126 | try: 127 | self.texture = pyglet.resource.image(texture_name) 128 | self.texture_name = texture_name 129 | except BaseException as ex: 130 | logging.warn('Could not load texture %s: %s' % (texture_name, ex)) 131 | 132 | def set_state(self, face=GL_FRONT_AND_BACK): 133 | if self.texture: 134 | glEnable(self.texture.target) 135 | glBindTexture(self.texture.target, self.texture.id) 136 | else: 137 | glDisable(GL_TEXTURE_2D) 138 | 139 | glMaterialfv(face, GL_DIFFUSE, (GLfloat * 4)(*(self.material.diffuse + 140 | [self.material.opacity]))) 141 | glMaterialfv(face, GL_AMBIENT, (GLfloat * 4)(*(self.material.ambient + 142 | [self.material.opacity]))) 143 | glMaterialfv(face, GL_SPECULAR, (GLfloat * 4)(*(self.material.specular + 144 | [self.material.opacity]))) 145 | glMaterialfv(face, GL_EMISSION, (GLfloat * 4)(*(self.material.emission + 146 | [self.material.opacity]))) 147 | glMaterialf(face, GL_SHININESS, self.material.shininess) 148 | 149 | def unset_state(self): 150 | if self.texture: 151 | glDisable(self.texture.target) 152 | glDisable(GL_COLOR_MATERIAL) 153 | 154 | def __eq__(self, other): 155 | if self.texture is None: 156 | return super(Material, self).__eq__(other) 157 | return (self.__class__ is other.__class__ and 158 | self.texture.id == other.texture.id and 159 | self.texture.target == other.texture.target and 160 | self.parent == other.parent) 161 | 162 | def __hash__(self): 163 | if self.texture is None: 164 | return super(Material, self).__hash__() 165 | return hash((self.texture.id, self.texture.target)) 166 | 167 | def set_alpha(self, alpha): 168 | """Sets the alpha value. 169 | 170 | Args: 171 | alpha: Alpha value. 172 | """ 173 | if self.texture is None and self.texture_name is None: 174 | logging.warn('Texture was not loaded successfully') 175 | return 176 | assert 0.0 <= alpha <= 1.0 177 | a_val = round(255 * alpha) 178 | 179 | f = pyglet.resource.file(self.texture_name) 180 | image = pyglet.image.load(self.texture_name, file=f) 181 | f.close() 182 | data = image.get_data(image.format, image.pitch) 183 | if image.format == 'RGB': 184 | rgb = np.array(data, dtype=np.uint8).reshape( 185 | (image.height, image.width, 3)) 186 | a = a_val * np.ones((image.height, image.width, 1), dtype=np.uint8) 187 | rgba = np.concatenate((rgb, a), axis=2) 188 | elif image.format == 'RGBA': 189 | rgba = np.array(data, dtype=np.uint8).reshape( 190 | (image.height, image.width, 4)) 191 | rgba[:, :, 3] = a_val 192 | # Somehow need to flip the height dimension. 193 | new_data = rgba[::-1, :, :] 194 | new_data = new_data.ravel().tobytes() 195 | image = pyglet.image.ImageData(image.width, 196 | image.height, 197 | 'RGBA', 198 | new_data, 199 | pitch=image.width * 4) 200 | self.texture = image.get_texture(True) 201 | 202 | 203 | def axes(size=1, width=1): 204 | """Draws 3D axes. 205 | 206 | Args: 207 | size: Axes length. 208 | width: Axes width. 209 | """ 210 | glLineWidth(width) 211 | pyglet.graphics.draw(6, GL_LINES, 212 | ('v3f', (0, 0, 0, size, 0, 0, 213 | 0, 0, 0, 0, size, 0, 214 | 0, 0, 0, 0, 0, size)), 215 | ('c3f', (1, 0, 0, 1, 0, 0, 216 | 0, 1, 0, 0, 1, 0, 217 | 0, 0, 1, 0, 0, 1, 218 | )) 219 | ) 220 | 221 | 222 | def frustum(dimensions, intrinsics): 223 | """Draws the camera's frustum. 224 | 225 | Args: 226 | dimensions: A tuple (w, h) containing the image width and height. 227 | intrinsics: A float32 numpy array of size [3, 3] containing the intrinsic 228 | matrix. 229 | """ 230 | w, h = dimensions[0], dimensions[1] 231 | batch = pyglet.graphics.Batch() 232 | 233 | for d in range(1, 6, 2): 234 | 235 | def get_point(x, y): 236 | p = list(np.linalg.inv(intrinsics).dot([x, y, 1]) * d) 237 | batch.add(2, GL_LINES, None, ('v3f', [0, 0, 0] + p)) 238 | return p 239 | 240 | top_left = get_point(0, 0) 241 | top_right = get_point(w, 0) 242 | bottom_right = get_point(w, h) 243 | bottom_left = get_point(0, h) 244 | 245 | batch.add(2, GL_LINES, None, ('v3f', top_left + top_right)) 246 | batch.add(2, GL_LINES, None, ('v3f', top_right + bottom_right)) 247 | batch.add(2, GL_LINES, None, ('v3f', bottom_right + bottom_left)) 248 | batch.add(2, GL_LINES, None, ('v3f', bottom_left + top_left)) 249 | 250 | batch.draw() 251 | 252 | 253 | def grid(size=1, n=10, width=1): 254 | """Draws a grid on XZ plane. 255 | 256 | Args: 257 | size: Grid line length in X and Z direction. 258 | n: Grid number. 259 | width: Grid line width. 260 | """ 261 | glLineWidth(width) 262 | s = size / float(n) 263 | s2 = 0.5 * size 264 | batch = pyglet.graphics.Batch() 265 | 266 | for i in range(0, n + 1): 267 | x = -s2 + i * s 268 | batch.add(2, GL_LINES, None, ('v3f', (x, 0, -s2, x, 0, s2))) 269 | for i in range(0, n + 1): 270 | z = -s2 + i * s 271 | batch.add(2, GL_LINES, None, ('v3f', (-s2, 0, z, s2, 0, z))) 272 | 273 | batch.draw() 274 | 275 | 276 | class Window(): 277 | """Viewer window.""" 278 | 279 | def __init__(self, dataloader): 280 | """Constructor: 281 | 282 | Args: 283 | dataloader: A SequenceLoader object. 284 | """ 285 | self.dataloader = dataloader 286 | 287 | self.config = Config(double_buffer=True, samples=8) # MSAA 288 | self.window = pyglet.window.Window(config=self.config, resizable=True) 289 | 290 | self.state = AppState(num_cameras=self.dataloader.num_cameras) 291 | 292 | self.pcd_vlist = [] 293 | self.pcd_image = [] 294 | w, h = self.dataloader.dimensions 295 | for _ in range(self.dataloader.num_cameras): 296 | self.pcd_vlist.append( 297 | pyglet.graphics.vertex_list(w * h, 'v3f/stream', 't2f/stream', 298 | 'n3f/stream')) 299 | self.pcd_image.append( 300 | pyglet.image.ImageData(w, h, 'RGB', (GLubyte * (w * h * 3))())) 301 | 302 | reset_pyglet_resource_path(self.dataloader.ycb_model_dir) 303 | self.ycb_batch = pyglet.graphics.Batch() 304 | self.ycb_vlist = [] 305 | for o in range(self.dataloader.num_ycb): 306 | n = self.dataloader.ycb_count[o] 307 | m = self.dataloader.ycb_material[o] 308 | g = Material(m) 309 | g.set_alpha(0.7) 310 | self.ycb_vlist.append( 311 | self.ycb_batch.add(n, GL_TRIANGLES, g, 'v3f/stream', 't2f/stream', 312 | 'n3f/stream')) 313 | 314 | self.ycb_per_view = isinstance(self.dataloader.ycb_vert[0], 315 | list) and isinstance( 316 | self.dataloader.ycb_norm[0], list) 317 | 318 | self.mano_batch = pyglet.graphics.Batch() 319 | self.mano_vlist = [] 320 | self.mano_llist = [] 321 | for _ in range(self.dataloader.num_mano): 322 | self.mano_vlist.append( 323 | self.mano_batch.add(4614, GL_TRIANGLES, None, 'v3f/stream', 324 | 'n3f/stream', 325 | ('c4f/static', [0.7, 0.7, 0.7, 0.7] * 4614))) 326 | self.mano_llist.append( 327 | self.mano_batch.add(9228, GL_LINES, None, 'v3f/stream', 328 | ('c3f/static', [0, 0, 0] * 9228))) 329 | 330 | self.fps_display = pyglet.window.FPSDisplay(self.window) 331 | 332 | @self.window.event 333 | def on_mouse_drag(x, y, dx, dy, buttons, modifiers): 334 | w, h = map(float, self.window.get_size()) 335 | 336 | if buttons & pyglet.window.mouse.LEFT: 337 | self.state.yaw -= dx * 0.5 338 | self.state.pitch -= dy * 0.5 339 | 340 | if buttons & pyglet.window.mouse.RIGHT: 341 | dp = np.array((dx / w, -dy / h, 0), np.float32) 342 | self.state.translation += np.dot(self.state.rotation, dp) 343 | 344 | if buttons & pyglet.window.mouse.MIDDLE: 345 | dz = dy * 0.01 346 | self.state.translation -= (0, 0, dz) 347 | self.state.distance -= dz 348 | 349 | @self.window.event 350 | def handle_mouse_btns(x, y, button, modifiers): 351 | self.state.mouse_btns[0] ^= (button & pyglet.window.mouse.LEFT) 352 | self.state.mouse_btns[1] ^= (button & pyglet.window.mouse.RIGHT) 353 | self.state.mouse_btns[2] ^= (button & pyglet.window.mouse.MIDDLE) 354 | 355 | self.window.on_mouse_press = self.window.on_mouse_release = handle_mouse_btns 356 | 357 | @self.window.event 358 | def on_mouse_scroll(x, y, scroll_x, scroll_y): 359 | dz = scroll_y * 0.1 360 | self.state.translation -= (0, 0, dz) 361 | self.state.distance -= dz 362 | 363 | @self.window.event 364 | def on_key_press(symbol, modifiers): 365 | if symbol == pyglet.window.key.R: 366 | self.state.reset() 367 | 368 | if symbol == pyglet.window.key.P: 369 | self.state.paused ^= True 370 | 371 | if symbol == pyglet.window.key.Z: 372 | self.state.scale ^= True 373 | 374 | if symbol == pyglet.window.key.X: 375 | self.state.attenuation ^= True 376 | 377 | if symbol == pyglet.window.key.L: 378 | self.state.lighting ^= True 379 | self._update_pcd_normals() 380 | 381 | # _1, _2, _3, ... 382 | if 49 <= symbol < 49 + len(self.state.camera_off): 383 | self.state.camera_off[symbol - 49] ^= True 384 | self._update_ycb() 385 | 386 | if symbol == pyglet.window.key.K: 387 | self.state.mask ^= 1 388 | self._update_pcd() 389 | 390 | if symbol == pyglet.window.key.M: 391 | self.state.model_off ^= True 392 | 393 | if symbol == pyglet.window.key.SPACE and self.state.paused: 394 | self.update(ignore_pause=True) 395 | 396 | if symbol == pyglet.window.key.S: 397 | pyglet.image.get_buffer_manager().get_color_buffer().save('out.png') 398 | 399 | if symbol == pyglet.window.key.Q: 400 | self.window.close() 401 | 402 | @self.window.event 403 | def on_draw(): 404 | self.window.clear() 405 | 406 | glEnable(GL_DEPTH_TEST) 407 | glEnable(GL_LINE_SMOOTH) 408 | 409 | width, height = self.window.get_size() 410 | glViewport(0, 0, width, height) 411 | 412 | # Set projection matrix stack. 413 | glMatrixMode(GL_PROJECTION) 414 | glLoadIdentity() 415 | gluPerspective(60, width / float(height), 0.01, 20) 416 | 417 | # Set modelview matrix stack. 418 | glMatrixMode(GL_MODELVIEW) 419 | glLoadIdentity() 420 | 421 | gluLookAt(0, 0, 0, 0, 0, 1, 0, -1, 0) 422 | 423 | glTranslatef(0, 0, self.state.distance) 424 | glRotated(self.state.pitch, 1, 0, 0) 425 | glRotated(self.state.yaw, 0, 1, 0) 426 | 427 | if any(self.state.mouse_btns): 428 | axes(0.1, 4) 429 | 430 | glTranslatef(0, 0, -self.state.distance) 431 | glTranslatef(*self.state.translation) 432 | 433 | # Draw grid. 434 | glColor3f(0.5, 0.5, 0.5) 435 | glPushMatrix() 436 | glTranslatef(0, 0.5, 0.5) 437 | grid() 438 | glPopMatrix() 439 | 440 | # Set point size. 441 | w, h = self.dataloader.dimensions 442 | psz = max(self.window.get_size()) / float(max( 443 | w, h)) if self.state.scale else 1 444 | glPointSize(psz) 445 | distance = (0, 0, 1) if self.state.attenuation else (1, 0, 0) 446 | glPointParameterfv(GL_POINT_DISTANCE_ATTENUATION, 447 | (GLfloat * 3)(*distance)) 448 | 449 | # Set lighting. 450 | if self.state.lighting: 451 | if not self.state.model_off: 452 | ldir = [0.0, 0.0, -1.0] 453 | else: 454 | ldir = np.dot(self.state.rotation, (0, 0, 1)) 455 | ldir = list(ldir) + [0] 456 | glLightfv(GL_LIGHT0, GL_POSITION, (GLfloat * 4)(*ldir)) 457 | glLightfv(GL_LIGHT0, GL_DIFFUSE, (GLfloat * 3)(1.0, 1.0, 1.0)) 458 | glLightfv(GL_LIGHT0, GL_AMBIENT, (GLfloat * 3)(0.75, 0.75, 0.75)) 459 | glLightfv(GL_LIGHT0, GL_SPECULAR, (GLfloat * 3)(0, 0, 0)) 460 | glEnable(GL_LIGHT0) 461 | glEnable(GL_NORMALIZE) 462 | glEnable(GL_LIGHTING) 463 | 464 | glColor3f(1, 1, 1) 465 | 466 | # Draw point cloud for each camera. 467 | for c in range(len(self.pcd_image)): 468 | if self.state.camera_off[c]: 469 | continue 470 | 471 | # Set texture matrix stack. 472 | glMatrixMode(GL_TEXTURE) 473 | glLoadIdentity() 474 | glTranslatef(0.5 / self.pcd_image[c].width, 475 | 0.5 / self.pcd_image[c].height, 0) 476 | image_texture = self.pcd_image[c].get_texture() 477 | tw, th = image_texture.owner.width, image_texture.owner.height 478 | glScalef(self.pcd_image[c].width / float(tw), 479 | self.pcd_image[c].height / float(th), 1) 480 | 481 | # Draw vertices and textures. 482 | texture = self.pcd_image[c].get_texture() 483 | glEnable(texture.target) 484 | glBindTexture(texture.target, texture.id) 485 | glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MAG_FILTER, GL_NEAREST) 486 | 487 | glEnable(GL_POINT_SPRITE) 488 | 489 | if not self.state.scale and not self.state.attenuation: 490 | glDisable(GL_MULTISAMPLE) 491 | self.pcd_vlist[c].draw(GL_POINTS) 492 | glDisable(texture.target) 493 | if not self.state.scale and not self.state.attenuation: 494 | glEnable(GL_MULTISAMPLE) 495 | 496 | # Draw YCB mesh. 497 | if not self.state.model_off: 498 | glEnable(GL_BLEND) 499 | glBlendFunc(GL_SRC_ALPHA, GL_ONE_MINUS_SRC_ALPHA) 500 | 501 | glMatrixMode(GL_TEXTURE) 502 | glLoadIdentity() 503 | 504 | self.ycb_batch.draw() 505 | 506 | # Draw MANO mesh. 507 | if not self.state.model_off: 508 | self.mano_batch.draw() 509 | 510 | glDisable(GL_LIGHTING) 511 | 512 | # Draw frustum and axes. 513 | glColor3f(0.25, 0.25, 0.25) 514 | frustum(self.dataloader.dimensions, self.dataloader.master_intrinsics) 515 | axes() 516 | 517 | # Reset matrix stacks. 518 | glMatrixMode(GL_PROJECTION) 519 | glLoadIdentity() 520 | glOrtho(0, width, 0, height, -1, 1) 521 | glMatrixMode(GL_MODELVIEW) 522 | glLoadIdentity() 523 | glMatrixMode(GL_TEXTURE) 524 | glLoadIdentity() 525 | 526 | glDisable(GL_DEPTH_TEST) 527 | 528 | self.fps_display.draw() 529 | 530 | def update(self, ignore_pause=False): 531 | """Updates the viewer window. 532 | 533 | Args: 534 | ignore_pause: Whether to update under a pause. 535 | """ 536 | if not ignore_pause and self.state.paused: 537 | return 538 | 539 | # Need to call `step()` at the start of `update()`. If called at the end 540 | # `self.pcd_image` will change as `pcd_rgb` changes since they share data. 541 | # This will cause incorrect texture in drawing. 542 | self.dataloader.step() 543 | 544 | self._update_pcd() 545 | self._update_pcd_normals() 546 | self._update_ycb() 547 | self._update_mano() 548 | 549 | def _copy(self, dst, src): 550 | """Copies a numpy array to a pyglet array. 551 | 552 | Args: 553 | dst: The pyglet array to copy to. 554 | src: The numpy array to copy from. 555 | """ 556 | np.array(dst, copy=False)[:] = src.ravel() 557 | 558 | def _update_pcd(self): 559 | """Updates point clouds.""" 560 | pcd_rgb = self.dataloader.pcd_rgb 561 | pcd_vert = self.dataloader.pcd_vert 562 | pcd_tex_coord = self.dataloader.pcd_tex_coord 563 | pcd_mask = self.dataloader.pcd_mask 564 | for c in range(len(self.pcd_image)): 565 | self.pcd_image[c].set_data('RGB', pcd_rgb[c].strides[0], 566 | pcd_rgb[c].ctypes.data) 567 | self._copy(self.pcd_vlist[c].vertices, pcd_vert[c]) 568 | self._copy(self.pcd_vlist[c].tex_coords, pcd_tex_coord[c]) 569 | if self.state.mask == 1: 570 | vertices = np.array(self.pcd_vlist[c].vertices, copy=False) 571 | for i in range(3): 572 | vertices[i::3][np.logical_not(pcd_mask[c]).ravel()] = 0 573 | 574 | def _update_pcd_normals(self): 575 | """Updates point cloud normals.""" 576 | if self.state.lighting: 577 | pcd_vert = self.dataloader.pcd_vert 578 | for c in range(len(self.pcd_image)): 579 | dy, dx = np.gradient(pcd_vert[c], axis=(0, 1)) 580 | n = np.cross(dx, dy) 581 | self._copy(self.pcd_vlist[c].normals, n) 582 | 583 | def _update_ycb(self): 584 | """Updates YCB objects.""" 585 | ycb_vert = self.dataloader.ycb_vert 586 | ycb_norm = self.dataloader.ycb_norm 587 | ycb_tex_coords = self.dataloader.ycb_tex_coords 588 | if self.ycb_per_view: 589 | for c, v in enumerate(self.state.camera_off): 590 | if not v: 591 | ycb_vert = ycb_vert[c] 592 | ycb_norm = ycb_norm[c] 593 | break 594 | for o in range(self.dataloader.num_ycb): 595 | self._copy(self.ycb_vlist[o].vertices, ycb_vert[o]) 596 | self._copy(self.ycb_vlist[o].normals, ycb_norm[o]) 597 | self._copy(self.ycb_vlist[o].tex_coords, ycb_tex_coords[o]) 598 | 599 | def _update_mano(self): 600 | """Updates MANO hands.""" 601 | mano_vert = self.dataloader.mano_vert 602 | mano_norm = self.dataloader.mano_norm 603 | mano_line = self.dataloader.mano_line 604 | for o in range(self.dataloader.num_mano): 605 | self._copy(self.mano_vlist[o].vertices, mano_vert[o]) 606 | self._copy(self.mano_vlist[o].normals, mano_norm[o]) 607 | self._copy(self.mano_llist[o].vertices, mano_line[o]) 608 | -------------------------------------------------------------------------------- /dex_ycb_toolkit/grasp_eval.py: -------------------------------------------------------------------------------- 1 | # DexYCB Toolkit 2 | # Copyright (C) 2021 NVIDIA Corporation 3 | # Licensed under the GNU General Public License v3.0 [see LICENSE for details] 4 | 5 | """Grasp evaluator.""" 6 | 7 | import os 8 | import sys 9 | import copy 10 | import trimesh 11 | import trimesh.transformations as tra 12 | import json 13 | import numpy as np 14 | import pyrender 15 | import time 16 | import torch 17 | import pickle 18 | import pycocotools.mask 19 | import cv2 20 | 21 | from manopth.manolayer import ManoLayer 22 | from scipy.spatial.distance import cdist 23 | from scipy.spatial import cKDTree 24 | from tabulate import tabulate 25 | 26 | from dex_ycb_toolkit.factory import get_dataset 27 | from dex_ycb_toolkit.logging import get_logger 28 | 29 | bop_toolkit_root = os.path.join(os.path.dirname(__file__), "..", "bop_toolkit") 30 | sys.path.append(bop_toolkit_root) 31 | 32 | from bop_toolkit_lib import inout 33 | 34 | _RADIUS = [0.05] 35 | _ANGLES = [15] 36 | _DIST_THRESHOLDS = [0.00, 0.01, 0.02, 0.03, 0.04, 0.05, 0.06, 0.07] 37 | 38 | _OBJ_MESH_GT_COLOR = (1.0, 1.0, 1.0, 0.6) 39 | _HAND_MESH_GT_COLOR = (0.47, 0.29, 0.21, 0.60) 40 | _HAND_MESH_PC_COLOR = (0.01, 0.02, 0.37, 1.00) 41 | 42 | _COVERED_GRASP_COLOR = (0.00, 1.00, 0.04, 1.00) 43 | _COLLIDE_GRASP_COLOR = (0.62, 0.02, 0.01, 1.00) 44 | _FAILURE_GRASP_COLOR = (0.7, 0.7, 0.7, 0.8) 45 | 46 | 47 | class GraspEvaluator(): 48 | """Grasp evaluator.""" 49 | radius = _RADIUS 50 | angles = _ANGLES 51 | dist_thresholds = _DIST_THRESHOLDS 52 | 53 | def __init__(self, name): 54 | """Constructor. 55 | 56 | Args: 57 | name: Dataset name. E.g., 's0_test'. 58 | """ 59 | self._name = name 60 | 61 | self._dataset = get_dataset(self._name) 62 | 63 | self._ycb_meshes = {} 64 | 65 | self._ycb_grasp_file = os.path.join("assets", 66 | "ycb_farthest_100_grasps.json") 67 | with open(self._ycb_grasp_file, 'r') as f: 68 | self._ycb_grasps = json.load(f) 69 | 70 | self._mano_layer_r = ManoLayer(flat_hand_mean=False, 71 | ncomps=45, 72 | side='right', 73 | mano_root='manopth/mano/models', 74 | use_pca=True) 75 | self._mano_layer_l = ManoLayer(flat_hand_mean=False, 76 | ncomps=45, 77 | side='left', 78 | mano_root='manopth/mano/models', 79 | use_pca=True) 80 | 81 | self._gripper_mesh_file = os.path.join("assets", "panda_gripper.obj") 82 | self._gripper_mesh = trimesh.load(self._gripper_mesh_file) 83 | 84 | self._gripper_mesh_vis_file = os.path.join("assets", "panda_tubes.obj") 85 | self._gripper_mesh_vis = trimesh.load(self._gripper_mesh_vis_file) 86 | 87 | self._gripper_pc_file = os.path.join("assets", "panda_pc.npy") 88 | gripper_pc = np.load(self._gripper_pc_file, allow_pickle=True) 89 | gripper_pc = gripper_pc.item()['points'][:100, :3] 90 | self._gripper_pc = gripper_pc 91 | 92 | self._h = 480 93 | self._w = 640 94 | x = np.linspace(0, self._w - 1, self._w) 95 | y = np.linspace(0, self._h - 1, self._h) 96 | self._xmap, self._ymap = np.meshgrid(x, y) 97 | 98 | self._default_coverage = {} 99 | self._default_precision = {} 100 | for r in _RADIUS: 101 | for a in _ANGLES: 102 | for thr in _DIST_THRESHOLDS: 103 | self._default_coverage.setdefault(r, {}).setdefault(a, {}).setdefault( 104 | thr, 0.0) 105 | self._default_precision.setdefault(r, {}).setdefault(a, 106 | {}).setdefault( 107 | thr, 0.0) 108 | 109 | self._tf_to_opengl = np.eye(4) 110 | self._tf_to_opengl[1, 1] = -1 111 | self._tf_to_opengl[2, 2] = -1 112 | 113 | self._covered_grasp_material = pyrender.MetallicRoughnessMaterial( 114 | alphaMode="BLEND", 115 | doubleSided=True, 116 | baseColorFactor=_COVERED_GRASP_COLOR, 117 | metallicFactor=0.0) 118 | self._collide_grasp_material = pyrender.MetallicRoughnessMaterial( 119 | alphaMode="BLEND", 120 | doubleSided=True, 121 | baseColorFactor=_COLLIDE_GRASP_COLOR, 122 | metallicFactor=0.0) 123 | self._failure_grasp_material = pyrender.MetallicRoughnessMaterial( 124 | alphaMode="BLEND", 125 | doubleSided=True, 126 | baseColorFactor=_FAILURE_GRASP_COLOR, 127 | metallicFactor=0.0) 128 | 129 | self._r = pyrender.OffscreenRenderer(viewport_width=self._w, 130 | viewport_height=self._h) 131 | 132 | self._out_dir = os.path.join(os.path.dirname(__file__), "..", "results") 133 | 134 | self._anno_file = os.path.join(self._out_dir, 135 | "anno_grasp_{}.pkl".format(self._name)) 136 | if os.path.isfile(self._anno_file): 137 | print('Found Grasp annotation file.') 138 | else: 139 | print('Cannot find Grasp annotation file.') 140 | self._generate_anno_file() 141 | 142 | self._anno = self._load_anno_file() 143 | 144 | def _generate_anno_file(self): 145 | """Generates the annotation file.""" 146 | print('Generating Grasp annotation file') 147 | s = time.time() 148 | 149 | anno = {} 150 | 151 | for i in range(len(self._dataset)): 152 | if (i + 1) in np.floor(np.linspace(0, len(self._dataset), 11))[1:]: 153 | print('{:3.0f}% {:6d}/{:6d}'.format(100 * i / len(self._dataset), i, 154 | len(self._dataset))) 155 | 156 | sample = self._dataset[i] 157 | 158 | # Skip samples not in the eval set. 159 | if not sample['is_grasp_target']: 160 | continue 161 | 162 | # Skip objects without pre-generated grasps. 163 | ycb_class = self._dataset.ycb_classes[sample['ycb_ids'][ 164 | sample['ycb_grasp_ind']]] 165 | if ycb_class not in self._ycb_grasps: 166 | continue 167 | 168 | label = np.load(sample['label_file']) 169 | 170 | pose_y = label['pose_y'][sample['ycb_grasp_ind']] 171 | pose_y = np.vstack((pose_y, np.array([[0, 0, 0, 1]]))) 172 | pose_m = label['pose_m'] 173 | 174 | if np.all(pose_m == 0.0): 175 | verts_m = None 176 | faces_m = None 177 | else: 178 | mano_side = sample['mano_side'] 179 | if mano_side == 'right': 180 | mano_layer = self._mano_layer_r 181 | if mano_side == 'left': 182 | mano_layer = self._mano_layer_l 183 | 184 | pose_m = torch.from_numpy(pose_m) 185 | mano_pose = pose_m[:, 0:48] 186 | mano_trans = pose_m[:, 48:51] 187 | 188 | mano_betas = np.array(sample['mano_betas'], dtype=np.float32) 189 | mano_betas = torch.from_numpy(mano_betas).unsqueeze(0) 190 | 191 | verts_m, _ = mano_layer(mano_pose, mano_betas, mano_trans) 192 | verts_m = verts_m[0] / 1000 193 | verts_m = verts_m.numpy() 194 | faces_m = mano_layer.th_faces 195 | 196 | anno[i] = { 197 | 'pose_y': pose_y, 198 | 'verts_m': verts_m, 199 | 'faces_m': faces_m, 200 | } 201 | 202 | print('# total samples: {:6d}'.format(len(self._dataset))) 203 | print('# valid samples: {:6d}'.format(len(anno))) 204 | print('# valid samples w hand: {:6d}'.format( 205 | sum([ 206 | x['verts_m'] is not None and x['faces_m'] is not None 207 | for x in anno.values() 208 | ]))) 209 | 210 | with open(self._anno_file, 'wb') as f: 211 | pickle.dump(anno, f) 212 | 213 | e = time.time() 214 | print('time: {:7.2f}'.format(e - s)) 215 | 216 | def _load_anno_file(self): 217 | """Loads the annotation file. 218 | 219 | Returns: 220 | A dictionary holding the loaded annotation. 221 | """ 222 | with open(self._anno_file, 'rb') as f: 223 | anno = pickle.load(f) 224 | 225 | return anno 226 | 227 | def _load_ycb_mesh(self, ycb_id): 228 | """Loads the mesh of a YCB object given a YCB object ID. 229 | 230 | Args: 231 | ycb_id: A YCB object ID. 232 | 233 | Returns: 234 | A dictionary holding the meshes for visualizing ground-truth pose and 235 | predicted pose. 236 | """ 237 | obj_file = self._dataset.obj_file[ycb_id] 238 | ycb_mesh_pred = trimesh.load(obj_file) 239 | ycb_mesh_gt = copy.deepcopy(ycb_mesh_pred) 240 | ycb_mesh_gt.visual = ycb_mesh_gt.visual.to_color() 241 | ycb_mesh_gt.visual.face_colors = np.tile(_OBJ_MESH_GT_COLOR, 242 | (len(ycb_mesh_gt.faces), 1)) 243 | ycb_mesh = { 244 | 'gt': ycb_mesh_gt, 245 | 'pred': ycb_mesh_pred, 246 | } 247 | return ycb_mesh 248 | 249 | def _get_hand_pc_from_det(self, dets, sample, hand_cat_id=22, radius=0.2): 250 | """Gets hand point cloud from hand detection and depth image. 251 | 252 | Args: 253 | dets: A dictionary holding the object and hand detections of an image. 254 | sample: A dictionary holding an image sample. 255 | hand_cat_id: Category ID for the hand class. 256 | radius: Radius threshold for filtering points. 257 | 258 | Returns: 259 | A float64 numpy array of shape [N, 3] containing the hand point cloud. 260 | """ 261 | if hand_cat_id not in dets: 262 | # Use empty point cloud if hand is not detected. 263 | hand_pc = np.zeros((0, 3)) 264 | else: 265 | # Use hand detection with highest score. 266 | max_score = 0 267 | max_score_ind = 0 268 | for dt_ind, dt in enumerate(dets[hand_cat_id]): 269 | if dt['score'] > max_score: 270 | max_score = dt['score'] 271 | max_score_ind = dt_ind 272 | 273 | segmentation = dets[hand_cat_id][max_score_ind]['segmentation'] 274 | mask = pycocotools.mask.decode(segmentation) 275 | 276 | depth_file = sample['depth_file'] 277 | depth = cv2.imread(depth_file, cv2.IMREAD_ANYDEPTH) 278 | depth = depth.astype(np.float32) / 1000 279 | 280 | fx = sample['intrinsics']['fx'] 281 | fy = sample['intrinsics']['fy'] 282 | ppx = sample['intrinsics']['ppx'] 283 | ppy = sample['intrinsics']['ppy'] 284 | 285 | pt0 = (self._xmap - ppx) * depth / fx 286 | pt1 = (self._ymap - ppy) * depth / fy 287 | pt2 = depth 288 | 289 | mask &= depth > 0 290 | 291 | choose = mask.ravel().nonzero()[0] 292 | pt0 = pt0.ravel()[choose][:, None] 293 | pt1 = pt1.ravel()[choose][:, None] 294 | pt2 = pt2.ravel()[choose][:, None] 295 | hand_pc = np.hstack((pt0, pt1, pt2)) 296 | 297 | if len(hand_pc) > 0: 298 | hand_center = np.median(hand_pc, axis=0, keepdims=True) 299 | dist = cdist(hand_pc, hand_center).squeeze(axis=1) 300 | choose = dist < radius 301 | hand_pc = hand_pc[choose] 302 | 303 | return hand_pc 304 | 305 | def _compute_grasp_coverage(self, samples, gt_poses, neighborhood_radius, 306 | neighborhood_angle): 307 | """Computes coverage rate of two sets of grasps. 308 | 309 | Args: 310 | samples: A float64 numpy array of shape [S, 7] containing the grasps to 311 | cover the other set. Each row contains one grasp represented by 312 | translation and rotation in quaternion (w, x, y, z). 313 | gt_poses: A float64 numpy array of shape [G, 7] containing the grasps to 314 | be covered. Each row contains one grasp represented by translation and 315 | rotation in quaternion (w, x, y, z). 316 | neighborhood_radius: Radius Threshold. 317 | neighborhood_angle: A float64 numpy array of shape [] containing the angle 318 | threshold. 319 | 320 | Returns: 321 | num_covered_poses: A int64 numpy array of shape [] containing the number 322 | of covered grasps. 323 | covered_sample_id: A int32 numpy array of shape [C] containing the indices 324 | of covered grasps. 325 | """ 326 | if len(samples) == 0: 327 | return 0.0, np.array([], dtype=np.int32) 328 | 329 | # Build kdtree. 330 | samples_tree = cKDTree(samples[:, :3]) 331 | 332 | # Find nearest neighbor. 333 | gt_neighbors_within_radius = samples_tree.query_ball_point( 334 | gt_poses[:, :3], r=neighborhood_radius) 335 | 336 | # Check orientation distance for these pairs. 337 | gt_num_neighbors = np.zeros(len(gt_poses)) 338 | gt_ind_neighbors = [] 339 | for i, (p, n) in enumerate(zip(gt_poses, gt_neighbors_within_radius)): 340 | conj_p = tra.quaternion_conjugate(p[3:]) 341 | angles = [ 342 | tra.rotation_from_matrix( 343 | tra.quaternion_matrix( 344 | tra.quaternion_multiply(conj_p, samples[b, 3:])))[0] 345 | for b in n 346 | ] 347 | within_neighborhood = np.abs(angles) < neighborhood_angle 348 | num_neighbors = np.sum(within_neighborhood) 349 | ind_neighbors = np.array(n)[within_neighborhood] 350 | gt_num_neighbors[i] = num_neighbors 351 | gt_ind_neighbors.append(ind_neighbors) 352 | 353 | num_covered_poses = np.sum(gt_num_neighbors > 0) 354 | covered_sample_id = np.unique(np.concatenate(gt_ind_neighbors)).astype( 355 | np.int32) 356 | 357 | return num_covered_poses, covered_sample_id 358 | 359 | def _visualize(self, fx, fy, cx, cy, obj_mesh_gt, hand_mesh_gt, obj_mesh_pred, 360 | hand_pc, pred_grasps, collision_free, covered_grasp_id, 361 | im_real, vis_file): 362 | """Visualizes predicted grasps and saves to a image file. 363 | 364 | Args: 365 | fx: Focal length in X direction. 366 | fy: Focal length in Y direction. 367 | cx: Principal point offset in X direction. 368 | cy: Principal point offset in Y direction. 369 | obj_mesh_gt: Ground-truth object mesh. 370 | hand_mesh_gt: Ground-truth hand mesh. 371 | obj_mesh_pred: Predicted object mesh. 372 | hand_pc: A float64 numpy array of shape [N, 3] containing the hand point 373 | cloud. 374 | pred_grasps: A list of float64 numpy arrays of shape [4, 4] containing the 375 | predicted grasps. 376 | collision_free: A bool numpy array of shape [G] indicating whether each 377 | predicted grasp is collision free. 378 | covered_grasp_id: An int32 numpy array of shape [C] containing the indices 379 | of collision free grasps that are covered. 380 | im_real: A uint8 numpy array of shape [H, W, 3] containing the color 381 | image. 382 | vis_file: Path to the visualization file. 383 | """ 384 | scene = pyrender.Scene(bg_color=np.array([0.0, 0.0, 0.0, 0.0]), 385 | ambient_light=np.array([1.0, 1.0, 1.0])) 386 | 387 | camera = pyrender.camera.IntrinsicsCamera(fx, fy, cx, cy) 388 | scene.add(camera) 389 | 390 | scene.add(pyrender.Mesh.from_trimesh(obj_mesh_gt, smooth=False)) 391 | if hand_mesh_gt is not None: 392 | scene.add(pyrender.Mesh.from_trimesh(hand_mesh_gt, smooth=False)) 393 | scene.add(pyrender.Mesh.from_trimesh(obj_mesh_pred)) 394 | 395 | hand_mesh_pc = trimesh.creation.uv_sphere(radius=0.003) 396 | hand_mesh_pc.visual.vertex_colors = _HAND_MESH_PC_COLOR 397 | tfs = np.tile(np.eye(4), (len(hand_pc), 1, 1)) 398 | tfs[:, :3, 3] = hand_pc 399 | scene.add(pyrender.Mesh.from_trimesh(hand_mesh_pc, poses=tfs)) 400 | 401 | for i, grasp in enumerate(pred_grasps): 402 | if not collision_free[i]: 403 | material = self._collide_grasp_material 404 | elif i in np.where(collision_free)[0][covered_grasp_id]: 405 | material = self._covered_grasp_material 406 | else: 407 | material = self._failure_grasp_material 408 | 409 | scene.add( 410 | pyrender.Mesh.from_trimesh( 411 | copy.deepcopy(self._gripper_mesh_vis).apply_transform( 412 | grasp).apply_transform(self._tf_to_opengl), material)) 413 | 414 | im_render, _ = self._r.render(scene) 415 | im = 0.2 * im_real.astype(np.float32) + 0.8 * im_render[:, :, ::-1].astype( 416 | np.float32) 417 | im = im.astype(np.uint8) 418 | 419 | cv2.imwrite(vis_file, im) 420 | 421 | def evaluate(self, bop_res_file, coco_res_file, out_dir=None, 422 | visualize=False): 423 | """Evaluates Grasp metrics given a BOP result file and a COCO result file. 424 | 425 | Args: 426 | bop_res_file: Path to the BOP result file. 427 | coco_res_file: Path to the COCO result file. 428 | out_dir: Path to the output directory. 429 | visualize: Whether to run visualization. 430 | 431 | Returns: 432 | A dictionary of results. 433 | """ 434 | if out_dir is None: 435 | out_dir = self._out_dir 436 | 437 | bop_res_name = os.path.splitext(os.path.basename(bop_res_file))[0] 438 | coco_res_name = os.path.splitext(os.path.basename(coco_res_file))[0] 439 | log_file = os.path.join( 440 | out_dir, "grasp_eval_{}_{}_{}.log".format(self._name, bop_res_name, 441 | coco_res_name)) 442 | logger = get_logger(log_file) 443 | 444 | grasp_res_file = os.path.join( 445 | out_dir, "grasp_res_{}_{}_{}.json".format(self._name, bop_res_name, 446 | coco_res_name)) 447 | if visualize: 448 | grasp_vis_dir = os.path.join( 449 | out_dir, "grasp_vis_{}_{}_{}".format(self._name, bop_res_name, 450 | coco_res_name)) 451 | os.makedirs(grasp_vis_dir, exist_ok=True) 452 | 453 | ests = inout.load_bop_results(bop_res_file) 454 | ests_org = {} 455 | for est in ests: 456 | ests_org.setdefault(est['scene_id'], 457 | {}).setdefault(est['im_id'], 458 | {}).setdefault(est['obj_id'], 459 | []).append(est) 460 | 461 | with open(coco_res_file, 'r') as f: 462 | dets = json.load(f) 463 | dets_org = {} 464 | for det in dets: 465 | dets_org.setdefault(det['image_id'], 466 | {}).setdefault(det['category_id'], []).append(det) 467 | 468 | m = trimesh.collision.CollisionManager() 469 | 470 | logger.info('Running evaluation') 471 | 472 | results = [] 473 | 474 | for ind, i in enumerate(self._anno): 475 | sample = self._dataset[i] 476 | 477 | ycb_id = sample['ycb_ids'][sample['ycb_grasp_ind']] 478 | ycb_class = self._dataset.ycb_classes[ycb_id] 479 | 480 | # Handle grasped object not being detected. 481 | scene_id, im_id = self._dataset.get_bop_id_from_idx(i) 482 | if scene_id not in ests_org or im_id not in ests_org[ 483 | scene_id] or ycb_id not in ests_org[scene_id][im_id]: 484 | results.append({ 485 | 'coverage': self._default_coverage, 486 | 'precision': self._default_precision, 487 | }) 488 | continue 489 | 490 | # Load YCB mesh. 491 | if ycb_id not in self._ycb_meshes: 492 | self._ycb_meshes[ycb_id] = self._load_ycb_mesh(ycb_id) 493 | 494 | # Get ground-truth object mesh. 495 | obj_pose_gt = self._anno[i]['pose_y'] 496 | obj_mesh_gt = copy.deepcopy(self._ycb_meshes[ycb_id]['gt']) 497 | obj_mesh_gt.apply_transform(obj_pose_gt) 498 | m.add_object('gt_obj', obj_mesh_gt) 499 | 500 | # Get ground-truth hand mesh. 501 | hand_verts = self._anno[i]['verts_m'] 502 | hand_faces = self._anno[i]['faces_m'] 503 | if hand_verts is not None and hand_faces is not None: 504 | hand_mesh_gt = trimesh.Trimesh(hand_verts, hand_faces) 505 | if visualize: 506 | hand_mesh_gt.visual.face_colors = np.tile( 507 | _HAND_MESH_GT_COLOR, (len(hand_mesh_gt.faces), 1)) 508 | else: 509 | # Leave hand out of evaluation if ground-truth is missing. 510 | hand_mesh_gt = None 511 | if hand_mesh_gt is not None: 512 | m.add_object('gt_hand', hand_mesh_gt) 513 | 514 | # Calculate ground-truth grasps based on ground-truth object and hand mesh. 515 | gt_grasps_q = [] 516 | for grasp_o in self._ycb_grasps[ycb_class]: 517 | grasp_w = np.matmul(obj_pose_gt, grasp_o) 518 | hit = m.in_collision_single(self._gripper_mesh, transform=grasp_w) 519 | if not hit: 520 | q = tra.quaternion_from_matrix(grasp_w, isprecise=True) 521 | t = grasp_w[:3, 3] 522 | g = np.hstack((t, q)) 523 | gt_grasps_q.append(g) 524 | 525 | m.remove_object('gt_obj') 526 | if hand_mesh_gt is not None: 527 | m.remove_object('gt_hand') 528 | 529 | # Get predicted object mesh. 530 | ests_sorted = sorted(ests_org[scene_id][im_id][ycb_id], 531 | key=lambda e: e['score'], 532 | reverse=True) 533 | est = ests_sorted[0] 534 | obj_pose_pred = np.eye(4) 535 | obj_pose_pred[:3, :3] = est['R'] 536 | obj_pose_pred[:3, 3] = est['t'][:, 0] / 1000 537 | obj_mesh_pred = copy.deepcopy(self._ycb_meshes[ycb_id]['pred']) 538 | obj_mesh_pred.apply_transform(obj_pose_pred) 539 | m.add_object('pred_obj', obj_mesh_pred) 540 | 541 | # Get predicted hand point cloud. 542 | if hand_mesh_gt is not None and i in dets_org: 543 | hand_pc = self._get_hand_pc_from_det(dets_org[i], sample) 544 | else: 545 | hand_pc = np.zeros((0, 3)) 546 | 547 | # Calculate predicted grasps based on predicted object mesh and hand point cloud. 548 | pred_grasps_m = {thr: [] for thr in _DIST_THRESHOLDS} 549 | pred_grasps_q = {thr: [] for thr in _DIST_THRESHOLDS} 550 | for grasp_o in self._ycb_grasps[ycb_class]: 551 | grasp_w = np.matmul(obj_pose_pred, grasp_o) 552 | hit = m.in_collision_single(self._gripper_mesh, transform=grasp_w) 553 | if not hit: 554 | r = grasp_w[:3, :3] 555 | t = grasp_w[:3, 3] 556 | gripper_pc = np.matmul(self._gripper_pc, r.T) + t 557 | if len(hand_pc) == 0: 558 | min_dist = max(_DIST_THRESHOLDS) + 1 559 | else: 560 | min_dist = cdist(gripper_pc, hand_pc).min() 561 | 562 | for thr in _DIST_THRESHOLDS: 563 | if min_dist > thr: 564 | pred_grasps_m[thr].append(grasp_w) 565 | q = tra.quaternion_from_matrix(grasp_w, isprecise=True) 566 | t = grasp_w[:3, 3] 567 | g = np.hstack((t, q)) 568 | pred_grasps_q[thr].append(g) 569 | 570 | m.remove_object('pred_obj') 571 | 572 | m.add_object('gt_obj', obj_mesh_gt) 573 | if hand_mesh_gt is not None: 574 | m.add_object('gt_hand', hand_mesh_gt) 575 | 576 | if visualize: 577 | color_file = sample['color_file'] 578 | color = cv2.imread(color_file) 579 | 580 | obj_mesh_gt.apply_transform(self._tf_to_opengl) 581 | if hand_mesh_gt is not None: 582 | hand_mesh_gt.apply_transform(self._tf_to_opengl) 583 | obj_mesh_pred.apply_transform(self._tf_to_opengl) 584 | hand_pc[:, 1:] *= -1 585 | 586 | # Compute coverage and precision. 587 | coverage = copy.deepcopy(self._default_coverage) 588 | precision = copy.deepcopy(self._default_precision) 589 | 590 | if len(gt_grasps_q) > 0: 591 | for r in _RADIUS: 592 | for a in _ANGLES: 593 | for thr in _DIST_THRESHOLDS: 594 | # Check collision with ground-trtuh object and hand mesh. 595 | collision_free = np.ones(len(pred_grasps_m[thr]), dtype=bool) 596 | for g_ind, g in enumerate(pred_grasps_m[thr]): 597 | hit = m.in_collision_single(self._gripper_mesh, transform=g) 598 | if hit: 599 | collision_free[g_ind] = 0 600 | 601 | if collision_free.sum() > 0: 602 | num_covered_gt_grasps, covered_pred_grasp_id = self._compute_grasp_coverage( 603 | np.array(pred_grasps_q[thr])[collision_free, :], 604 | np.array(gt_grasps_q), r, np.deg2rad(a)) 605 | 606 | num_covered_pred_grasps, _ = self._compute_grasp_coverage( 607 | np.array(gt_grasps_q), 608 | np.array(pred_grasps_q[thr])[collision_free, :], r, 609 | np.deg2rad(a)) 610 | 611 | coverage[r][a][thr] = num_covered_gt_grasps / len(gt_grasps_q) 612 | precision[r][a][thr] = num_covered_pred_grasps / len( 613 | pred_grasps_q[thr]) 614 | 615 | if visualize: 616 | vis_dir = os.path.join( 617 | grasp_vis_dir, 618 | "radius_{:4.2f}.angle_{:2d}.min-dist-threshold_{:4.2f}". 619 | format(r, a, thr)) 620 | os.makedirs(vis_dir, exist_ok=True) 621 | vis_file = os.path.join(vis_dir, "{:06d}.jpg".format(i)) 622 | self._visualize(sample['intrinsics']['fx'], 623 | sample['intrinsics']['fy'], 624 | sample['intrinsics']['ppx'], 625 | sample['intrinsics']['ppy'], obj_mesh_gt, 626 | hand_mesh_gt, obj_mesh_pred, hand_pc, 627 | pred_grasps_m[thr], collision_free, 628 | covered_pred_grasp_id, color, vis_file) 629 | 630 | results.append({ 631 | 'coverage': coverage, 632 | 'precision': precision, 633 | }) 634 | 635 | logger.info('{:04d}/{:04d} {:6d} {:21s} # gt grasps: {:3d}'.format( 636 | ind + 1, len(self._anno), i, ycb_class, len(gt_grasps_q))) 637 | 638 | m.remove_object('gt_obj') 639 | if hand_mesh_gt is not None: 640 | m.remove_object('gt_hand') 641 | 642 | tabular_data = [] 643 | for r in _RADIUS: 644 | for a in _ANGLES: 645 | for thr in _DIST_THRESHOLDS: 646 | coverage = np.mean([x['coverage'][r][a][thr] for x in results]) 647 | precision = np.mean([x['precision'][r][a][thr] for x in results]) 648 | tabular_data.append([r, a, thr, coverage, precision]) 649 | metrics = [ 650 | 'radius (m)', 'angle (deg)', 'dist th (m)', 'coverage', 'precision' 651 | ] 652 | table = tabulate(tabular_data, 653 | headers=metrics, 654 | tablefmt='pipe', 655 | floatfmt='.4f', 656 | numalign='right') 657 | logger.info('Results: \n' + table) 658 | 659 | with open(grasp_res_file, 'w') as f: 660 | json.dump(results, f) 661 | 662 | logger.info('Evaluation complete.') 663 | 664 | return results 665 | --------------------------------------------------------------------------------