Add baseDataloader

.gitignore

@@ -2,6 +2,7 @@
 configs/user/nguyen.yaml
 .hydra/
 outputs/
+None/
 bop_viz_kit/
 huggingface/
 slurm_logs/

README.md

@@ -30,15 +30,10 @@ 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)
+- 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/datasets/)
 - 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)
-
-##  TODO
-- Tutorial of training/testing on custom datasets
-- Gradio demo with similarity visualization
-- Release universal pretrained models
 <details><summary>Click to expand</summary>
 
 ```bash
@@ -103,4 +98,9 @@ The code is adapted from [Nope](https://github.com/nv-nguyen/nope), [Temos](http
 The authors thank Martin Sundermeyer, Paul Wohlhart and Shreyas Hampali for their fast reply, feedback!
 
 ## Contact
-If you have any question, feel free to create an issue or contact the first author at van-nguyen.nguyen@enpc.fr
+If you have any question, feel free to create an issue or contact the first author at van-nguyen.nguyen@enpc.fr
+
+##  TODO
+- Tutorial of training/testing on custom datasets
+- Gradio demo with similarity visualization
+- Release universal pretrained models

src/dataloader/base.py

@@ -24,6 +24,7 @@
 import trimesh
 from src.utils.visualization_utils import render_pts_to_image
 from src.poses.utils import crop_frame
+
 # set level logging
 logging.basicConfig(level=logging.INFO)
 
@@ -64,8 +65,7 @@ def load_list_scene(self, split=None):
                     if os.path.isdir(osp.join(self.root_dir, scene))
                 ]
             )
-            print(osp.join(self.root_dir, split[0]))
-        logging.info(f"Found {self.list_scenes} scenes")
+        logging.info(f"Found {len(self.list_scenes)} scenes")
 
     def load_scene(self, path, use_visible_mask=True):
         # Load rgb and mask images
@@ -223,7 +223,13 @@ def load_metaData(self, reset_metaData, mode="query", split="test", level=2):
 
     def load_cad(self, cad_name="models"):
         cad_dir = f"{self.root_dir}/models/{cad_name}"
-        cad_names = sorted([x for x in os.listdir(cad_dir) if x.endswith(".ply")])
+        cad_names = sorted(
+            [
+                x
+                for x in os.listdir(cad_dir)
+                if x.endswith(".ply") and not x.endswith("_old.ply")
+            ]
+        )
         models_info = load_json(osp.join(cad_dir, "models_info.json"))
         self.cads = {}
         for cad_name in cad_names:
@@ -299,12 +305,38 @@ def crop_with_gt_pose(self, img, mask, pose, K, virtual_bbox_size):
 
 if __name__ == "__main__":
     logging.basicConfig(level=logging.INFO)
-    root_dir = "/home/nguyen/Documents/datasets/templateV2/o-lm/"
-    reset_metaData = True
-    dataset = BaseBOP(root_dir)
-    dataset.load_list_scene("test")
-    dataset.load_metaData(reset_metaData=reset_metaData)
-    dataset.load_cad()
-    for scene_path in dataset.list_scenes:
-        scene_id = scene_path.split("/")[-1]
-        dataset.check_scene(scene_id, "./tmp/")
+    root_dir = "/gpfsscratch/rech/xjd/uyb58rn/datasets/template-pose"
+    dataset_names = ["hb", "hope", "icbin", "lm", "o-lm", "ruapc", "tudl"]
+    # tless is special
+    for dataset_name, split in zip(
+        ["tless/test", "tless/train"], ["test_primesense", "train_primesense"]
+    ):
+        dataset = BaseBOP(os.path.join(root_dir, dataset_name), split)
+        dataset.load_list_scene(split=split)
+        dataset.load_metaData(reset_metaData=True)
+        dataset.load_cad(cad_name="models_cad")
+        for scene_path in dataset.list_scenes:
+            scene_id = scene_path.split("/")[-1]
+            dataset.check_scene(scene_id, f"./tmp/{dataset_name}")
+
+    for dataset_name in tqdm(dataset_names):
+        splits = [
+            split
+            for split in os.listdir(os.path.join(root_dir, dataset_name))
+            if os.path.isdir(os.path.join(root_dir, dataset_name, split))
+        ]
+        splits = [
+            split
+            for split in splits
+            if split.startswith("train")
+            or split.startswith("val")
+            or split.startswith("test")
+        ]
+        for split in splits:
+            dataset = BaseBOP(os.path.join(root_dir, dataset_name), split)
+            dataset.load_list_scene(split=split)
+            dataset.load_metaData(reset_metaData=True)
+            dataset.load_cad()
+            for scene_path in dataset.list_scenes:
+                scene_id = scene_path.split("/")[-1]
+                dataset.check_scene(scene_id, f"./tmp/{dataset_name}")

src/dataloader/bop.py

@@ -0,0 +1,53 @@
+import os, random
+import numpy as np
+from PIL import Image, ImageFilter
+import pandas as pd
+import torch
+import torchvision.transforms as transforms
+import torch.utils.data as data
+import json
+from src.dataloader.base import BaseBOP
+import logging
+import cv2
+import os.path as osp
+
+# set level logging
+logging.basicConfig(level=logging.INFO)
+
+
+class BOPDataset(BaseBOP):
+    def __init__(
+        self,
+        root_dir,
+        template_dir,
+        obj_ids,
+        img_size,
+        cropping_with_bbox=None,
+        reset_metaData=False,
+        **kwargs,
+    ):
+        self.root_dir = root_dir
+        self.template_dir = template_dir
+        self.img_size = img_size
+        self.mask_size = 25 if img_size == 64 else int(img_size // 8)
+
+        self.cropping_with_bbox = cropping_with_bbox
+        self.obj_ids = obj_ids
+        self.load_template_poses(template_dir=template_dir)
+        self.load_list_scene(obj_ids)
+        self.load_metaData(
+            reset_metaData=reset_metaData,
+            mode="query",
+        )
+        self.im_transform = transforms.Compose(
+            [
+                transforms.ToTensor(),
+                transforms.Normalize(
+                    mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]
+                ),
+            ]
+        )
+        logging.info(f"Length of dataloader: {self.__len__()}")
+
+    def load_template_poses(self, template_dir):
+        self.templates_poses = np.load(osp.join(template_dir, "obj_poses.npy"))

src/scripts/compute_neighbors.py

@@ -34,14 +34,19 @@ def compute_neighbors(cfg: DictConfig) -> None:
     OmegaConf.set_struct(cfg, False)
     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 os.listdir(data_cfg.root_dir)
+            if os.path.isdir(os.path.join(data_cfg.root_dir, split))
+        ]
         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(

src/scripts/process_mesh.py

@@ -70,7 +70,6 @@ def render(cfg: DictConfig) -> None:
     # convert mesh format of ycbv or hope from textured CAD to vertex color CAD
     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") and not name.endswith("_old.ply")]
         )

src/scripts/sanity_check_datasets.py

@@ -0,0 +1,86 @@
+import os, sys
+import numpy as np
+import shutil
+from tqdm import tqdm
+import time
+from functools import partial
+import multiprocessing
+import logging
+import os, sys
+import os.path as osp
+
+# set level logging
+logging.basicConfig(level=logging.INFO)
+import logging
+import hydra
+from omegaconf import DictConfig, OmegaConf
+import numpy as np
+from src.scripts.render_template import call_blender_proc
+
+
+@hydra.main(
+    version_base=None,
+    config_path="../../configs",
+    config_name="render",
+)
+def render(cfg: DictConfig) -> None:
+    OmegaConf.set_struct(cfg, False)
+    save_dir = osp.join(osp.dirname(cfg.data.lm.root_dir), "templates")
+    dataset_name = cfg.dataset_to_render
+    if dataset_name == "tless":
+        data_cfg = cfg.data["tless_test"]
+    else:
+        data_cfg = cfg.data[dataset_name]
+    obj_pose_path = f"{save_dir}/{dataset_name}/obj_poses.npy"
+    logging.info(f"Checking {data_cfg.dataset_name} ...")
+    if dataset_name in ["tless"]:
+        cad_dir = os.path.join(data_cfg.root_dir, "models/models_cad")
+    else:
+        cad_dir = os.path.join(data_cfg.root_dir, "models/models")
+    object_ids = sorted(
+        [
+            int(name[4:][:-4])
+            for name in os.listdir(cad_dir)
+            if name.endswith(".ply") and not name.endswith("old.ply")
+        ]
+    )
+    cad_paths, template_dirs = [], []
+    for object_id in object_ids:
+        cad_path = os.path.join(
+            cad_dir,
+            "obj_{:06d}.ply".format(object_id),
+        )
+        template_dir = os.path.join(save_dir, f"{dataset_name}/obj_{object_id:06d}")
+        num_templates = len(os.listdir(template_dir))
+        if num_templates != 642:
+            logging.info(
+                f"Dataset {data_cfg.dataset_name}, obj {object_id} failed, found only {num_templates}"
+            )
+            cad_paths.append(cad_path)
+            template_dirs.append(template_dir)
+
+    os.environ["CUDA_VISIBLE_DEVICES"] = cfg.gpus
+    start_time = time.time()
+    pool = multiprocessing.Pool(processes=int(cfg.num_workers))
+    call_blender_proc_with_index = partial(
+        call_blender_proc,
+        list_cad_path=cad_paths,
+        list_output_dir=template_dirs,
+        obj_pose_path=obj_pose_path,
+        disable_output=cfg.disable_output,
+        gpus_devices=cfg.gpus,
+    )
+    mapped_values = list(
+        tqdm(
+            pool.imap_unordered(
+                call_blender_proc_with_index, range(len(template_dirs))
+            ),
+            total=len(template_dirs),
+        )
+    )
+    finish_time = time.time()
+    print("Total time to render templates for query:", finish_time - start_time)
+
+
+if __name__ == "__main__":
+    render()