Merge pull request #14 from swtyree/pr_bop_inference_single_dataset

BOP inference on individual datasets

Author: ylabbe

cosypose/scripts/run_bop_inference.py

@@ -193,6 +193,7 @@ def main():
     parser.add_argument('--id', default=-1, type=int)
     parser.add_argument('--config', default='bop-pbr', type=str)
     parser.add_argument('--nviews', dest='n_views', default=1, type=int)
+    parser.add_argument('--dataset', default=None, type=str)
     parser.add_argument('--icp', action='store_true')
     args = parser.parse_args()
 
@@ -246,7 +247,11 @@ def main():
         ds_names = ['hb', 'tless', 'ycbv']
     else:
         ds_names = ['hb', 'icbin', 'itodd', 'lmo', 'tless', 'tudl', 'ycbv']
-
+    
+    if args.dataset is not None:
+        assert args.dataset in ds_names, f'{args.dataset} not in {ds_names}'
+        ds_names = [args.dataset]
+    
     for ds_name in ds_names:
         this_cfg = deepcopy(cfg)
         this_cfg.ds_name = BOP_CONFIG[ds_name]['inference_ds_name'][0]