Add all posible datasets of BOP and train script

README.md

@@ -30,6 +30,7 @@ If you like this project, check out related works from our group:
 
 ##  Updates (WIP)
 We have introduced additional features and updates to the codebase:
+- Releasing ready-to-use universal model pretrained on different datasets of BOP challenge [Linemod, HomebrewedDB, HOPE, RU-APC, IC-BIN, IC-MI, TUD-L, T-LESS](https://bop.felk.cvut.cz/media/data/bop_datasets/icmi_train.zip)
 - Adding code to generate poses (OpenCV coordinate) from icosahedron with Blender
 - Parsing with [hydra](https://github.com/facebookresearch/hydra) library, simplifying training_step, testing_step with [pytorch lightning](https://lightning.ai/)
 - Path structure (of pretrained models, dataset) is defined as in our recent project [NOPE](https://github.com/nv-nguyen/nope)
@@ -79,14 +80,18 @@ First, create template poses from icosahedron:
 ```
 blenderproc run src/poses/create_poses.py
 ```
-There are two options for next steps:
+Next, download and process BOP datasets
+```
+./src/scripts/download_and_process_datasets.sh
+```
+There are two options for final step (rendering synthetic templates from CAD models):
 #### Option 1: Download preprocessed dataset:
 ```
-./src/scripts/download_preprocessed_data.sh
+TODO
 ```
 #### Option 2: Create/download dataset from scratch
 ```
-./src/scripts/download_and_process_from_scratch.sh
+./src/scripts/render_all.sh
 ```
 </details>
 

configs/data/all.yaml

@@ -3,5 +3,9 @@ defaults:
   - /data/olm@olm
   - /data/tless_train@tless_train
   - /data/tless_test@tless_test
-  - /data/hb@hb # additional training sets to have universial model
-  - /data/hope@hope # additional training sets to have universial model
+  - /data/hb@hb # additional training sets wrt paper to have universial model
+  - /data/hope@hope # additional training sets wrt paper to have universial model
+  - /data/icmi@icmi # additional training sets wrt paper to have universial model
+  - /data/icbin@icbin # additional training sets wrt paper to have universial model
+  - /data/ruapc@ruapc # additional training sets wrt paper to have universial model
+  - /data/tudl@tudl # additional training sets wrt paper to have universial model

configs/data/icbin.yaml

@@ -0,0 +1,9 @@
+dataset_name: IC-BIN
+root_dir: ${machine.root_dir}/icbin
+source: 
+  url: https://bop.felk.cvut.cz/media/data/bop_datasets/icbin_train.zip
+  cad_url: https://bop.felk.cvut.cz/media/data/bop_datasets/icbin_models.zip
+  http: True
+  unzip_mode: unzip
+  processing:
+

configs/data/icmi.yaml

@@ -0,0 +1,9 @@
+dataset_name: IC-MI
+root_dir: ${machine.root_dir}/icmi
+source: 
+  url: https://bop.felk.cvut.cz/media/data/bop_datasets/icmi_train.zip
+  cad_url: https://bop.felk.cvut.cz/media/data/bop_datasets/icmi_models.zip
+  http: True
+  unzip_mode: unzip
+  processing:
+

configs/data/lm.yaml

@@ -1,4 +1,4 @@
-dataset_name: query
+dataset_name: LINEMOD
 root_dir: ${machine.root_dir}/lm
 obj_names: "ape, benchvise, cam, can, cat, driller, duck, eggbox, glue, holepuncher, iron, lamp, phone"
 source:

configs/data/olm.yaml

@@ -1,4 +1,4 @@
-dataset_name: occlusionquery
+dataset_name: occlusionLINEMOD
 root_dir: ${machine.root_dir}/o-lm
 obj_names: "ape, can, cat, driller, duck, eggbox, glue, holepuncher"
 source:

configs/data/ruapc.yaml

@@ -0,0 +1,9 @@
+dataset_name: RU-APC
+root_dir: ${machine.root_dir}/ruapc
+source: 
+  url: https://bop.felk.cvut.cz/media/data/bop_datasets/ruapc_train.zip
+  cad_url: https://bop.felk.cvut.cz/media/data/bop_datasets/ruapc_models.zip
+  http: True
+  unzip_mode: unzip
+  processing:
+

configs/data/tudl.yaml

@@ -0,0 +1,8 @@
+dataset_name: TUD-L
+root_dir: ${machine.root_dir}/tudl
+source: 
+  url: https://bop.felk.cvut.cz/media/data/bop_datasets/tudl_train_real.zip
+  cad_url: https://bop.felk.cvut.cz/media/data/bop_datasets/tudl_models.zip
+  http: True
+  unzip_mode: unzip
+  processing:

configs/model/base.yaml

@@ -0,0 +1,14 @@
+modelname: BaseFeatureExtractor
+_target_: src.model.base_network.BaseFeatureExtractor
+
+descriptor_size: 16
+threshold: 0.2
+pretrained_weight: ${machine.root_dir}/pretrained/moco_v2_800ep_pretrain.pth
+
+# config_optim
+lr: 0.01
+weight_decay: 0.0005
+warm_up_steps: 1000
+log_interval: 1000
+log_dir: ${save_dir}/media
+use_all_gather: false

configs/model/resnet50.yaml

@@ -0,0 +1,14 @@
+modelname: FeatureExtractor
+_target_: src.model.network.FeatureExtractor
+
+descriptor_size: 16
+threshold: 0.2
+pretrained_weight: ${machine.root_dir}/pretrained/moco_v2_800ep_pretrain.pth
+
+# config_optim
+lr: 0.0001
+weight_decay: 0.0005
+warm_up_steps: 1000
+log_interval: 1000
+log_dir: ${save_dir}/media
+use_all_gather: false

configs/train.yaml

@@ -0,0 +1,24 @@
+# Composing nested config with default
+defaults:
+  - user: default
+  - machine: local
+  - callback: base 
+  - model: resnet50
+  - data: all
+  - _self_
+
+save_dir: ${machine.root_dir}/results/${name_exp}
+name_exp: train
+use_pretrained: True
+train_datasets:
+  - tless_train
+  - hb 
+  - hope 
+  - icmi 
+  - icbin
+  - ruapc
+  - tudl
+test_datasets:
+  - lm 
+  - olm
+  - tless_test

src/scripts/compute_neighbors.py

@@ -28,66 +28,63 @@
 @hydra.main(
     version_base=None,
     config_path="../../configs",
-    config_name="process_data",
+    config_name="render",
 )
 def compute_neighbors(cfg: DictConfig) -> None:
     OmegaConf.set_struct(cfg, False)
-
-    # query
-    level = 2
-    list_root_dir = [
-        cfg.data.lm.root_dir,
-        cfg.data.olm.root_dir,
-        cfg.data.tless_train.root_dir,
-        cfg.data.tless_test.root_dir,
-    ]
-    list_split = ["test", "test", "train_primesense", "test_primesense"]
-    list_pose_distribution = ["upper", "upper", "all", "all"]
-    for root_dir, split, pose_distribution in zip(
-        list_root_dir, list_split, list_pose_distribution
-    ):
-        start_time = time.time()
-        finder = NearestTemplateFinder(
-            level_templates=level,
-            pose_distribution=pose_distribution,
-            return_inplane=True,
-        )
-        bop_dataset = BaseBOP(root_dir=root_dir)
-        bop_dataset.load_list_scene(split=split)
-        for scene_path in tqdm(bop_dataset.list_scenes):
-            templates_infos = {}
-            scene_data = bop_dataset.load_scene(scene_path)
-            save_template_path = osp.join(scene_path, f"template_level{level}.json")
-            rgbs_path = scene_data["rgb_paths"]
-            for idx_frame in range(len(rgbs_path)):
-                rgb_path = scene_data["rgb_paths"][idx_frame]
-                id_frame = int(str(rgb_path).split("/")[-1].split(".")[0])
-                frame_poses = (
-                    scene_data["scene_gt"][f"{id_frame}"]
-                    if f"{id_frame}" in scene_data["scene_gt"]
-                    else scene_data["scene_gt"][id_frame]
-                )
-                frame_poses = (
-                    frame_poses if isinstance(frame_poses, list) else [frame_poses]
-                )
-                cad_ids = [x["obj_id"] for x in frame_poses]
-                cad_poses = np.array(
-                    [
-                        combine_R_and_T(x["cam_R_m2c"], x["cam_t_m2c"])
-                        for x in frame_poses
+    for data_cfg in cfg.data.values():
+        logging.info(f"Compute neighbors template for {data_cfg.dataset_name}")
+        splits = os.listdir(data_cfg.root_dir)
+        splits = [
+            split
+            for split in splits
+            if split.startswith("train")
+            or split.startswith("val")
+            or split.startswith("test")
+        ]
+        for split in splits:
+            start_time = time.time()
+            finder = NearestTemplateFinder(
+                level_templates=2,
+                pose_distribution="all",
+                return_inplane=True,
+            )
+            bop_dataset = BaseBOP(root_dir=data_cfg.root_dir, split=split)
+            bop_dataset.load_list_scene(split=split)
+            for scene_path in tqdm(bop_dataset.list_scenes):
+                templates_infos = {}
+                scene_data = bop_dataset.load_scene(scene_path)
+                save_template_path = osp.join(scene_path, f"template_level2.json")
+                rgbs_path = scene_data["rgb_paths"]
+                for idx_frame in range(len(rgbs_path)):
+                    rgb_path = scene_data["rgb_paths"][idx_frame]
+                    id_frame = int(str(rgb_path).split("/")[-1].split(".")[0])
+                    frame_poses = (
+                        scene_data["scene_gt"][f"{id_frame}"]
+                        if f"{id_frame}" in scene_data["scene_gt"]
+                        else scene_data["scene_gt"][id_frame]
+                    )
+                    frame_poses = (
+                        frame_poses if isinstance(frame_poses, list) else [frame_poses]
+                    )
+                    cad_ids = [x["obj_id"] for x in frame_poses]
+                    cad_poses = np.array(
+                        [
+                            combine_R_and_T(x["cam_R_m2c"], x["cam_t_m2c"])
+                            for x in frame_poses
+                        ]
+                    )
+                    idx_templates, inplanes = finder.search_nearest_template(cad_poses)
+                    templates_infos[f"{id_frame}"] = [
+                        {
+                            "obj_id": cad_ids[idx_obj],
+                            "idx_template": int(idx_templates[idx_obj]),
+                            "inplane": float(inplanes[idx_obj]),
+                        }
+                        for idx_obj in range(len(cad_ids))
                     ]
-                )
-                idx_templates, inplanes = finder.search_nearest_template(cad_poses)
-                templates_infos[f"{id_frame}"] = [
-                    {
-                        "obj_id": cad_ids[idx_obj],
-                        "idx_template": int(idx_templates[idx_obj]),
-                        "inplane": float(inplanes[idx_obj]),
-                    }
-                    for idx_obj in range(len(cad_ids))
-                ]
-            save_json(save_template_path, templates_infos)
-        logging.info(f"Time to compute neighbors: {time.time() - start_time}")
+                save_json(save_template_path, templates_infos)
+            logging.info(f"Time to compute neighbors: {time.time() - start_time}")
 
 
 if __name__ == "__main__":

src/scripts/download.py

@@ -57,13 +57,12 @@ def run_download(config: DictConfig) -> None:
 @hydra.main(
     version_base=None,
     config_path="../../configs",
-    config_name="download",
+    config_name="render",
 )
 def download(cfg: DictConfig) -> None:
     OmegaConf.set_struct(cfg, False)
     for data_cfg in cfg.data.values():
         logging.info(f"Downloading {data_cfg.dataset_name}")
-        print(data_cfg)
         run_download(data_cfg)
         logging.info(f"---" * 100)
 

src/scripts/download_and_process_datasets.sh

@@ -0,0 +1,4 @@
+python -m src.scripts.download
+python -m src.scripts.process_lm_gt
+python -m src.scripts.process_mesh 
+python -m src.scripts.compute_neighbors

src/scripts/process_mesh.py

@@ -13,20 +13,49 @@
 import numpy as np
 from src.utils.trimesh_utils import load_mesh
 import trimesh
+from src.utils.inout import write_txt
 
 
-def convert_ply_to_obj(idx, ply_paths, obj_paths):
+def manual_formatting(save_path, vertices, faces, vertex_colors):
+    new_texts = [
+        "ply",
+        "format ascii 1.0",
+        "comment Created by Blender 2.77 (sub 0) - www.blender.org, source file: ''",
+        f"element vertex {len(vertices)}",
+        "property float x",
+        "property float y",
+        "property float z",
+        # "property float nx",
+        # "property float ny",
+        # "property float nz",
+        "property uchar red",
+        "property uchar green",
+        "property uchar blue",
+        f"element face {len(faces)}",
+        "property list uchar uint vertex_indices",
+        "end_header",
+    ]
+    assert len(vertices) == len(vertex_colors)
+    for i in range(len(vertices)):
+        new_texts.append(
+            f"{vertices[i][0]} {vertices[i][1]} {vertices[i][2]} {vertex_colors[i][0]} {vertex_colors[i][1]} {vertex_colors[i][2]}"
+        )
+    for i in range(len(faces)):
+        new_texts.append(f"3 {faces[i][0]} {faces[i][1]} {faces[i][2]}")
+    write_txt(save_path, new_texts)
+    print(f"Finish formatting {save_path}")
+
+
+def convert_ply_to_obj(idx, ply_paths):
     ply_path = ply_paths[idx]
-    obj_path = obj_paths[idx]
 
     # open mesh
     mesh = load_mesh(ply_path)
     texture = mesh.visual.material.image
     vertex_colors = trimesh.visual.uv_to_color(mesh.visual.uv, texture)
-    new_mesh = trimesh.Trimesh(
-        vertices=mesh.vertices, faces=mesh.faces, vertex_colors=vertex_colors
-    )
-    new_mesh.export(obj_path)
+    # rename the older one and the new_one
+    os.rename(ply_path, ply_path.replace(".ply", "_old.ply"))
+    manual_formatting(ply_path, mesh.vertices, mesh.faces, vertex_colors=vertex_colors)
 
 
 @hydra.main(
@@ -39,19 +68,15 @@ def render(cfg: DictConfig) -> None:
 
     start_time = time.time()
     # convert mesh format of ycbv or hope from textured CAD to vertex color CAD
-    for dataset_name in ["hope"]:
+    for dataset_name in ["hope", "ruapc"]:
         cad_dir = os.path.join(cfg.data[dataset_name].root_dir, "models/models")
+        print(cad_dir)
         cad_names = sorted(
-            [name for name in os.listdir(cad_dir) if name.endswith(".ply")]
+            [name for name in os.listdir(cad_dir) if name.endswith(".ply") and not name.endswith("_old.ply")]
         )
-        save_paths = [
-            osp.join(cad_dir, f"{cad_path[:-4]}.obj") for cad_path in cad_names
-        ]
         cad_paths = [osp.join(cad_dir, name) for name in cad_names]
 
-        convert_ply_to_obj_with_index = partial(
-            convert_ply_to_obj, ply_paths=cad_paths, obj_paths=save_paths
-        )
+        convert_ply_to_obj_with_index = partial(convert_ply_to_obj, ply_paths=cad_paths)
         pool = multiprocessing.Pool(processes=1)
         mapped_values = list(
             tqdm(

src/scripts/render_all.sh

@@ -0,0 +1,8 @@
+python -m src.scripts.render_template dataset_to_render=lm
+python -m src.scripts.render_template dataset_to_render=tless
+python -m src.scripts.render_template dataset_to_render=hb
+python -m src.scripts.render_template dataset_to_render=hope
+python -m src.scripts.render_template dataset_to_render=ruapc
+python -m src.scripts.render_template dataset_to_render=tudl
+python -m src.scripts.render_template dataset_to_render=icmi
+python -m src.scripts.render_template dataset_to_render=icbin