Feature/single cat no sym

.gitignore

@@ -1,5 +1,6 @@
 __pycache__
-.ipynb_checkpoints
+wandb/
+ipynb_checkpoints
 .vscode
 logs
 *.pyc
@@ -19,4 +20,4 @@ local_data
 dist/
 *.so
 build/
-*.temp
+*.temp

cosypose/bop_config.py

@@ -38,6 +38,38 @@
     test_ds_name=['lmo.bop19'],
 )
 
+BOP_CONFIG['bracket_assembly'] = dict(
+    input_resize=(640, 480),
+    urdf_ds_name='bracket_assembly',
+    obj_ds_name='bracket_assembly',
+    train_pbr_ds_name=['bracket_assembly'],
+    train_pbr_real_ds_names=[('bracket_assembly', 1), ()],
+    inference_ds_name=['bracket_assembly'],
+    test_ds_name=[],
+)
+BOP_CONFIG['bracket_assembly_noaug'] = dict(
+    input_resize=(640, 480),
+    urdf_ds_name='bracket_assembly',
+    obj_ds_name='bracket_assembly',
+    train_pbr_ds_name=['bracket_assembly'],
+    train_pbr_real_ds_names=[('bracket_assembly', 1), ()],
+    inference_ds_name=['bracket_assembly'],
+    test_ds_name=[],
+    rgb_augmentation=False,
+    background_augmentation=False,
+    gray_augmentation=False
+)
+
+BOP_CONFIG['bracket_assembly_debug'] = dict(
+    input_resize=(640, 480),
+    urdf_ds_name='bracket_assembly_debug',
+    obj_ds_name='bracket_assembly_debug',
+    train_pbr_ds_name=['bracket_assembly_debug'],
+    train_pbr_real_ds_names=[('bracket_assembly_debug', 1), ()],
+    inference_ds_name=['bracket_assembly_debug'],
+    test_ds_name=[],
+)
+
 
 BOP_CONFIG['tless'] = dict(
     input_resize=(720, 540),

cosypose/config.py

@@ -29,7 +29,6 @@
 
 DEPS_DIR = PROJECT_DIR / 'deps'
 CACHE_DIR = LOCAL_DATA_DIR / 'joblib_cache'
-
 assert LOCAL_DATA_DIR.exists()
 CACHE_DIR.mkdir(exist_ok=True)
 TEST_DATA_DIR.mkdir(exist_ok=True)

cosypose/datasets/bop.py

@@ -39,6 +39,7 @@ def build_index(ds_dir, save_file, split, save_file_annotations):
         for f in ('scene_camera.json', 'scene_gt_info.json', 'scene_gt.json'):
             path = (scene_dir / f)
             if path.exists():
+                print(path, f)
                 annotations_scene[f.split('.')[0]] = json.loads(path.read_text())
         annotations[scene_id] = annotations_scene
         # for view_id in annotations_scene['scene_gt_info'].keys():
@@ -56,9 +57,10 @@ def build_index(ds_dir, save_file, split, save_file_annotations):
 
 
 class BOPDataset:
-    def __init__(self, ds_dir, split='train', load_depth=False):
+    def __init__(self, ds_dir, split='train', load_depth=False, train_classes=None, visib_fract_thres=0.5):
         ds_dir = Path(ds_dir)
         self.ds_dir = ds_dir
+        self.train_classes = train_classes
         assert ds_dir.exists(), 'Dataset does not exists.'
 
         self.split = split
@@ -73,9 +75,31 @@ def __init__(self, ds_dir, split='train', load_depth=False):
             split=split)
         self.frame_index = pd.read_feather(save_file_index).reset_index(drop=True)
         self.annotations = pickle.loads(save_file_annotations.read_bytes())
-
+        # TODO
         models_infos = json.loads((ds_dir / 'models' / 'models_info.json').read_text())
-        self.all_labels = [f'obj_{int(obj_id):06d}' for obj_id in models_infos.keys()]
+        # filter detection model, only load frames_debug
+        if 'debug' in str(ds_dir):
+            frames_debug = [('000000', '62')]
+            annotations_debug = {}
+            frame_df_debug = pd.DataFrame()
+            for scene_index, t_index in frames_debug:
+                frame_df_debug = frame_df_debug.append(
+                    self.frame_index.loc[
+                        (self.frame_index['scene_id']==int(scene_index))
+                            & (self.frame_index['view_id']==int(t_index))
+                        ]
+                        )
+                annotations_debug[scene_index] = annotations_debug.get(scene_index, {})
+                for key, value_dict in self.annotations[scene_index].items():                
+                    annotations_debug[scene_index][key] = annotations_debug[scene_index].get(key, {})
+                    annotations_debug[scene_index][key][t_index] = value_dict[t_index]
+            self.annotations = annotations_debug
+            self.frame_index = frame_df_debug
+        self.visib_fract_thres = visib_fract_thres
+        if train_classes is not None:
+            self.all_labels = [f'obj_{int(obj_id):06d}' for obj_id in models_infos.keys() if str(obj_id) in train_classes]
+        else:
+            self.all_labels = [f'obj_{int(obj_id):06d}' for obj_id in models_infos.keys()]
         self.load_depth = load_depth
 
     def __len__(self):
@@ -108,7 +132,7 @@ def __getitem__(self, frame_id):
         cam_annotation = self.annotations[scene_id_str]['scene_camera'][str(view_id)]
         if 'cam_R_w2c' in cam_annotation:
             RC0 = np.array(cam_annotation['cam_R_w2c']).reshape(3, 3)
-            tC0 = np.array(cam_annotation['cam_t_w2c']) * 0.001
+            tC0 = np.array(cam_annotation['cam_t_w2c'])# * 0.001
             TC0 = Transform(RC0, tC0)
         else:
             TC0 = Transform(np.eye(3), np.zeros(3))
@@ -126,12 +150,16 @@ def __getitem__(self, frame_id):
             n_objects = len(annotation)
             visib = self.annotations[scene_id_str]['scene_gt_info'][str(view_id)]
             for n in range(n_objects):
+                if visib[n]['visib_fract'] < self.visib_fract_thres:
+                    continue
                 RCO = np.array(annotation[n]['cam_R_m2c']).reshape(3, 3)
-                tCO = np.array(annotation[n]['cam_t_m2c']) * 0.001
+                tCO = np.array(annotation[n]['cam_t_m2c']) #* 0.001
                 TCO = Transform(RCO, tCO)
                 T0O = T0C * TCO
                 T0O = T0O.toHomogeneousMatrix()
                 obj_id = annotation[n]['obj_id']
+                if self.train_classes and str(obj_id) not in self.train_classes:
+                    continue
                 name = f'obj_{int(obj_id):06d}'
                 bbox_visib = np.array(visib[n]['bbox_visib'])
                 x, y, w, h = bbox_visib
@@ -159,7 +187,7 @@ def __getitem__(self, frame_id):
             if not depth_path.exists():
                 depth_path = depth_path.with_suffix('.tif')
             depth = np.array(inout.load_depth(depth_path))
-            camera['depth'] = depth * cam_annotation['depth_scale'] / 1000
+            camera['depth'] = depth * cam_annotation['depth_scale'] #/ 1000
 
         obs = dict(
             objects=objects,

cosypose/datasets/bop_object_datasets.py

@@ -3,26 +3,34 @@
 
 
 class BOPObjectDataset:
-    def __init__(self, ds_dir):
+    def __init__(self, ds_dir, mesh_units='mm', ignore_symmetric=False, train_classes=None):
+        assert mesh_units in ("mm", "m") # 'mm' for tless, 'm' for bracket assembly
         ds_dir = Path(ds_dir)
         infos_file = ds_dir / 'models_info.json'
         infos = json.loads(infos_file.read_text())
         objects = []
         for obj_id, bop_info in infos.items():
+            # just load one model (nut: obj_id = obj_000005)
+            # print("obj_id", train_classes, obj_id)
+            if train_classes and str(obj_id) not in train_classes:
+                continue
             obj_id = int(obj_id)
             obj_label = f'obj_{obj_id:06d}'
             mesh_path = (ds_dir / obj_label).with_suffix('.ply').as_posix()
             obj = dict(
                 label=obj_label,
                 category=None,
                 mesh_path=mesh_path,
-                mesh_units='mm',
+                mesh_units=mesh_units, # mm for tless
             )
+            # TODO: for bracket assembly
+            # set everything to false
             is_symmetric = False
-            for k in ('symmetries_discrete', 'symmetries_continuous'):
-                obj[k] = bop_info.get(k, [])
-                if len(obj[k]) > 0:
-                    is_symmetric = True
+            if not ignore_symmetric:
+                for k in ('symmetries_discrete', 'symmetries_continuous'):
+                    obj[k] = bop_info.get(k, [])
+                    if len(obj[k]) > 0:
+                        is_symmetric = True
             obj['is_symmetric'] = is_symmetric
             obj['diameter'] = bop_info['diameter']
             scale = 0.001 if obj['mesh_units'] == 'mm' else 1.0

cosypose/datasets/datasets_cfg.py

@@ -30,6 +30,7 @@ def keep_bop19(ds):
 
 
 def make_scene_dataset(ds_name, n_frames=None):
+    '''return a dataframe? with index for each frame'''
     # TLESS
     if ds_name == 'tless.primesense.train':
         ds = _make_tless_dataset('train_primesense')
@@ -129,15 +130,34 @@ def make_scene_dataset(ds_name, n_frames=None):
     elif ds_name == 'tudl.train.real':
         ds_dir = BOP_DS_DIR / 'tudl'
         ds = BOPDataset(ds_dir, split='train_real')
-
+    elif ds_name == 'bracket_assembly_debug':
+        ds_dir = BOP_DS_DIR / 'bracket_assembly_debug'
+        ds = BOPDataset(ds_dir, split='train_pbr')
+    elif ds_name == 'bracket_assembly_debug_nut' or ds_name == 'bracket_assembly_debug_nut_nosym':
+        ds_dir = BOP_DS_DIR / 'bracket_assembly_debug'
+        ds = BOPDataset(ds_dir, split='train_pbr', train_classes=['5'])
+    elif ds_name == 'bracket_assembly' or ds_name == 'bracket_assembly_nosym':
+        ds_dir = BOP_DS_DIR / 'bracket_assembly'
+        ds = BOPDataset(ds_dir, split='train_pbr')
+    elif ds_name == 'bracket_assembly_debug_nut' or ds_name == 'bracket_assembly_debug_nut_nosym':
+        ds_dir = BOP_DS_DIR / 'bracket_assembly_debug_nut'
+        ds = BOPDataset(ds_dir, split='train_pbr')
+    elif ds_name == 'bracket_assembly_04_22' or ds_name == 'bracket_assembly_04_22_nosym':
+        ds_dir = BOP_DS_DIR / 'bracket_assembly'
+        ds = BOPDataset(ds_dir, split='train_pbr')
+    elif ds_name == 'bracket_assembly_nut' or ds_name == 'bracket_assembly_nut_nosym':
+        ds_dir = BOP_DS_DIR / 'bracket_assembly'
+        ds = BOPDataset(ds_dir, split='train_pbr', train_classes=['5'])
+    elif ds_name == 'bracket_assembly_04_22_nut' or ds_name == 'bracket_assembly_04_22_nut_nosym':
+        ds_dir = BOP_DS_DIR / 'bracket_assembly'
+        ds = BOPDataset(ds_dir, split='train_pbr', train_classes=['5'])
     # Synthetic datasets
     elif 'synthetic.' in ds_name:
         from .synthetic_dataset import SyntheticSceneDataset
         assert '.train' in ds_name or '.val' in ds_name
         is_train = 'train' in ds_name.split('.')[-1]
         ds_name = ds_name.split('.')[1]
         ds = SyntheticSceneDataset(ds_dir=LOCAL_DATA_DIR / 'synt_datasets' / ds_name, train=is_train)
-
     else:
         raise ValueError(ds_name)
 
@@ -176,7 +196,28 @@ def make_object_dataset(ds_name):
         ds = BOPObjectDataset(BOP_DS_DIR / 'lm/models')
     elif ds_name == 'tudl':
         ds = BOPObjectDataset(BOP_DS_DIR / 'tudl/models')
-
+    elif ds_name == 'bracket_assembly_debug':
+        ds = BOPObjectDataset(BOP_DS_DIR / 'bracket_assembly_debug/models', mesh_units='m', ignore_symmetric=False)
+    elif ds_name == 'bracket_assembly_debug_nut':
+        ds = BOPObjectDataset(BOP_DS_DIR / 'bracket_assembly_debug/models', mesh_units='m', ignore_symmetric=False, train_classes=['5'])
+    elif ds_name == 'bracket_assembly_debug_nut_nosym':
+        ds = BOPObjectDataset(BOP_DS_DIR / 'bracket_assembly_debug/models', mesh_units='m', ignore_symmetric=True, train_classes=['5'])
+    elif ds_name == 'bracket_assembly':
+        ds = BOPObjectDataset(BOP_DS_DIR / 'bracket_assembly/models', mesh_units='m', ignore_symmetric=False)
+    elif ds_name == 'bracket_assembly_04_22':
+        ds = BOPObjectDataset(BOP_DS_DIR / 'bracket_assembly_04_22/models', mesh_units='m', ignore_symmetric=False)
+    elif ds_name == 'bracket_assembly_nosym':
+        ds = BOPObjectDataset(BOP_DS_DIR / 'bracket_assembly/models', mesh_units='m', ignore_symmetric=True)
+    elif ds_name == 'bracket_assembly_04_22_nosym':
+        ds = BOPObjectDataset(BOP_DS_DIR / 'bracket_assembly_04_22/models', mesh_units='m', ignore_symmetric=True)
+    elif ds_name == 'bracket_assembly_nut':
+        ds = BOPObjectDataset(BOP_DS_DIR / 'bracket_assembly/models', mesh_units='m', ignore_symmetric=False, train_classes=['5'])
+    elif ds_name == 'bracket_assembly_04_22_nut':
+        ds = BOPObjectDataset(BOP_DS_DIR / 'bracket_assembly_04_22/models', mesh_units='m', ignore_symmetric=False, train_classes=['5'])
+    elif ds_name == 'bracket_assembly_nut_nosym':
+        ds = BOPObjectDataset(BOP_DS_DIR / 'bracket_assembly/models', mesh_units='m', ignore_symmetric=True, train_classes=['5'])
+    elif ds_name == 'bracket_assembly_04_22_nut_nosym':
+        ds = BOPObjectDataset(BOP_DS_DIR / 'bracket_assembly_04_22/models', mesh_units='m', ignore_symmetric=True, train_classes=['5'])
     else:
         raise ValueError(ds_name)
     return ds
@@ -216,6 +257,8 @@ def make_urdf_dataset(ds_name):
 
     elif ds_name == 'camera':
         ds = OneUrdfDataset(ASSET_DIR / 'camera/model.urdf', 'camera')
+    elif ds_name == 'bracket_assembly':
+        ds = BOPUrdfDataset(LOCAL_DATA_DIR / 'urdfs' / 'bracket_assembly')
     else:
         raise ValueError(ds_name)
     return ds
@@ -226,4 +269,4 @@ def make_texture_dataset(ds_name):
         ds = TextureDataset(LOCAL_DATA_DIR / 'texture_datasets' / 'shapenet')
     else:
         raise ValueError(ds_name)
-    return ds
+    return ds

cosypose/datasets/pose_dataset.py

@@ -90,6 +90,7 @@ def get_data(self, idx):
         mask_uniqs = set(np.unique(mask))
         objects_visible = []
         for obj in state['objects']:
+            # TODO: only load nut, filter by obj['id_in_segm']
             add = False
             if obj['id_in_segm'] in mask_uniqs and np.all(np.array(obj['bbox']) >= 0):
                 add = True
@@ -105,7 +106,6 @@ def get_data(self, idx):
                 objects_visible.append(obj)
         if len(objects_visible) == 0:
             raise NoObjectError
-        # assert len(objects_visible) > 0, idx
 
         rgb = torch.as_tensor(rgb).permute(2, 0, 1).to(torch.uint8)
         assert rgb.shape[0] == 3

cosypose/datasets/urdf_dataset.py

@@ -28,7 +28,7 @@ def __len__(self):
 class BOPUrdfDataset(UrdfDataset):
     def __init__(self, ds_dir):
         super().__init__(ds_dir)
-        self.index['scale'] = 0.001
+        self.index['scale'] = 1.0
 
 
 class OneUrdfDataset:

cosypose/datasets/utils.py

@@ -32,6 +32,7 @@ def make_detections_from_segmentation(masks):
 
     for mask_n in masks:
         dets_n = dict()
+        # print(torch.unique(mask_n, sorted=True))
         for uniq in torch.unique(mask_n, sorted=True):
             ids = np.where((mask_n == uniq).cpu().numpy())
             x1, y1, x2, y2 = np.min(ids[1]), np.min(ids[0]), np.max(ids[1]), np.max(ids[0])

cosypose/datasets/wrappers/multiview_wrapper.py

@@ -36,6 +36,7 @@ def __init__(self, scene_ds, n_views=4):
     def __getitem__(self, idx):
         row = self.frame_index.iloc[idx]
         ds_ids = row['scene_ds_ids']
+        # print('multiview wrapper get item', idx, ds_ids, row)
         rgbs, masks, obss = [], [], []
         for ds_id in ds_ids:
             rgb, mask, obs = self.scene_ds[ds_id]

cosypose/evaluation/eval_runner/pose_eval.py

@@ -61,6 +61,13 @@ def collate_fn(self, batch):
     def evaluate(self, obj_predictions, device='cuda'):
         for meter in self.meters.values():
             meter.reset()
+        # TODO: ADD PER CATEGORY INFERENCE
+        # all_categories = []
+        # obj_predictions_copy = obj_predictions.copy()
+        # print(obj_data_gt.infos['frame_obj_id'], obj_data_gt.infos['label'])
+        # for obj_category in all_categories:
+        #     obj_predictions_copy.frame_index = obj_predictions.frame_index.loc[obj_predictions['']==obj_category]
+
         obj_predictions = obj_predictions.to(device)
         for obj_data_gt in tqdm(self.dataloader):
             for k, meter in self.meters.items():
@@ -69,11 +76,15 @@ def evaluate(self, obj_predictions, device='cuda'):
 
     def summary(self):
         summary, dfs = dict(), dict()
+        # all_categories = []
         for meter_k, meter in sorted(self.meters.items()):
+
+            # print(obj_data_gt.infos['frame_obj_id'], obj_data_gt.infos['label'])
+            # for obj_category in all_categories:
             meter.gather_distributed(tmp_dir=self.tmp_dir)
             if get_rank() == 0 and len(meter.datas) > 0:
                 summary_, df_ = meter.summary()
                 dfs[meter_k] = df_
                 for k, v in summary_.items():
                     summary[meter_k + '/' + k] = v
-        return summary, dfs
+        return summary, dfs

cosypose/evaluation/meters/detection_meters.py

@@ -177,7 +177,7 @@ def compute_ap(label_df, label_n_gt):
 
         df = pred_df[['label', valid_k, 'score']].to_dataframe().set_index(['label'])
         for label, label_n_gt in n_gts.items():
-            if df.index.contains(label):
+            if label in df.index:
                 label_df = df.loc[[label]]
                 if label_df[valid_k].sum() > 0:
                     ap, label_df = compute_ap(label_df, label_n_gt)

cosypose/evaluation/meters/pose_meters.py

@@ -275,7 +275,7 @@ def compute_ap(label_df, label_n_gt):
 
         df = pred_df[['label', valid_k, 'score']].to_dataframe().set_index(['label'])
         for label, label_n_gt in n_gts.items():
-            if df.index.contains(label):
+            if label in df.index:
                 label_df = df.loc[[label]]
                 if label_df[valid_k].sum() > 0:
                     ap, label_df = compute_ap(label_df, label_n_gt)