Add train.py run succesfully on all BOP datasets

README.md

@@ -70,7 +70,7 @@ conda env create -f environment.yml
 conda activate template
 ```
 
-### 2. Datasets (from scratch) :smiley_cat: :electric_plug:
+### 2. Datasets
 First, create template poses from icosahedron:
 ```
 blenderproc run src/poses/create_poses.py
@@ -90,6 +90,24 @@ TODO
 ```
 </details>
 
+ ##  Launch a training 
+
+<details><summary>Click to expand</summary>
+
+### 1. Mixing all BOP datasets except LINEMOD and T-LESS (only objects 19-30)
+```
+python train.py name_exp=train_all
+```
+
+The parsing is done with Hydra library. You can override anything in the configuration by passing arguments. For example:
+```
+# experiment 1: change batch_size, using data augmentation, update name_exp
+python train.py machine.batch_size=2 use_augmentation=True name_exp=train_augmentation
+
+# experiment 2: change batch_size, using data augmentation, update name_exp, update_lr
+python train.py machine.batch_size=2 use_augmentation=True model.lr=0.001 name_exp=train_augmentation_lr0.001
+```
+</details>
 
 ## Acknowledgement
 

configs/data/hb.yaml

@@ -7,3 +7,11 @@ source:
   unzip_mode: unzip
   processing:
 
+dataloader:
+  _target_: src.dataloader.bop.BOPDataset
+  root_dir: ${machine.root_dir}/hb
+  template_dir: ${machine.root_dir}/templates/hb
+  obj_ids: 
+  reset_metaData: True
+  img_size: 256
+  cropping_with_bbox: True

configs/data/hope.yaml

@@ -7,3 +7,11 @@ source:
   unzip_mode: unzip
   processing:
 
+dataloader:
+  _target_: src.dataloader.bop.BOPDataset
+  root_dir: ${machine.root_dir}/hope
+  template_dir: ${machine.root_dir}/templates/hope
+  obj_ids:
+  reset_metaData: True
+  img_size: 256
+  cropping_with_bbox: True

configs/data/icbin.yaml

@@ -7,3 +7,11 @@ source:
   unzip_mode: unzip
   processing:
 
+dataloader:
+  _target_: src.dataloader.bop.BOPDataset
+  root_dir: ${machine.root_dir}/icbin
+  template_dir: ${machine.root_dir}/templates/icbin
+  obj_ids:
+  reset_metaData: True
+  img_size: 256
+  cropping_with_bbox: True

configs/data/icmi.yaml

@@ -7,3 +7,11 @@ source:
   unzip_mode: unzip
   processing:
 
+dataloader:
+  _target_: src.dataloader.bop.BOPDataset
+  root_dir: ${machine.root_dir}/icmi
+  template_dir: ${machine.root_dir}/templates/icmi
+  obj_ids:
+  reset_metaData: True
+  img_size: 256
+  cropping_with_bbox: True

configs/data/lm.yaml

@@ -9,10 +9,11 @@ source:
   processing: 
 
 dataloader:
-  _target_: src.dataloader.lm.LINEMOD
+  _target_: src.dataloader.bop.LINEMOD
   root_dir: ${machine.root_dir}/lm
   template_dir: ${machine.root_dir}/templates/linemod
   split_name: ${split_name}
   reset_metaData: True
   mode:
-  img_size: 256
+  img_size: 256
+  cropping_with_bbox: True

configs/data/ruapc.yaml

@@ -7,3 +7,11 @@ source:
   unzip_mode: unzip
   processing:
 
+dataloader:
+  _target_: src.dataloader.bop.BOPDataset
+  root_dir: ${machine.root_dir}/ruapc
+  template_dir: ${machine.root_dir}/templates/ruapc
+  obj_ids: 
+  reset_metaData: True
+  img_size: 256
+  cropping_with_bbox: True

configs/data/tless_train.yaml

@@ -7,3 +7,11 @@ source:
   unzip_mode: unzip
   processing:
 
+dataloader:
+  _target_: src.dataloader.bop.BOPDataset
+  root_dir: ${machine.root_dir}/tless/train
+  template_dir: ${machine.root_dir}/templates/tless
+  obj_ids:
+  reset_metaData: True
+  img_size: 256
+  cropping_with_bbox: True

configs/data/tudl.yaml

@@ -5,4 +5,13 @@ source:
   cad_url: https://bop.felk.cvut.cz/media/data/bop_datasets/tudl_models.zip
   http: True
   unzip_mode: unzip
-  processing:
+  processing:
+
+dataloader:
+  _target_: src.dataloader.bop.BOPDataset
+  root_dir: ${machine.root_dir}/tudl
+  template_dir: ${machine.root_dir}/templates/tudl
+  obj_ids:
+  reset_metaData: True
+  img_size: 256
+  cropping_with_bbox: True

configs/model/resnet50.yaml

@@ -1,7 +1,7 @@
 modelname: FeatureExtractor
 _target_: src.model.network.FeatureExtractor
 
-descriptor_size: 16
+descriptor_size: 8
 threshold: 0.2
 pretrained_weight: ${machine.root_dir}/pretrained/moco_v2_800ep_pretrain.pth
 

configs/train.yaml

@@ -10,6 +10,7 @@ defaults:
 save_dir: ${machine.root_dir}/results/${name_exp}
 name_exp: train
 use_pretrained: True
+use_augmentation: False
 train_datasets:
   - tless_train
   - hb 

src/dataloader/base.py

@@ -47,8 +47,6 @@ def load_list_scene(self, split=None):
         if isinstance(split, str):
             if split is not None:
                 split_folder = osp.join(self.root_dir, split)
-            else:
-                split_folder = self.root_dir
             self.list_scenes = sorted(
                 [
                     osp.join(split_folder, scene)
@@ -65,6 +63,8 @@ def load_list_scene(self, split=None):
                     if os.path.isdir(osp.join(self.root_dir, scene))
                 ]
             )
+        else:
+            raise NotImplementedError
         logging.info(f"Found {len(self.list_scenes)} scenes")
 
     def load_scene(self, path, use_visible_mask=True):
@@ -100,6 +100,8 @@ def load_metaData(self, reset_metaData, mode="query", split="test", level=2):
                 "mask_visib_path": [],
                 "depth_path": [],
                 "visib_fract": [],
+                "bbox_obj": [],
+                "bbox_visib": [],
                 "intrinsic": [],
                 "idx_template": [],
                 "inplane": [],
@@ -182,7 +184,19 @@ def load_metaData(self, reset_metaData, mode="query", split="test", level=2):
                                 video_metaData["scene_camera"][f"{id_frame}"]["cam_K"]
                             )
                             metaData["visib_fract"].append(
-                                video_metaData["scene_gt_info"][f"{id_frame}"][idx_obj]
+                                video_metaData["scene_gt_info"][f"{id_frame}"][idx_obj][
+                                    "visib_fract"
+                                ]
+                            )
+                            metaData["bbox_obj"].append(
+                                video_metaData["scene_gt_info"][f"{id_frame}"][idx_obj][
+                                    "bbox_obj"
+                                ]
+                            )
+                            metaData["bbox_visib"].append(
+                                video_metaData["scene_gt_info"][f"{id_frame}"][idx_obj][
+                                    "bbox_visib"
+                                ]
                             )
                             metaData["idx_template"].append(
                                 templates_frame[idx_obj]["idx_template"]