From a17defcee105917e1b176298710e0e85056e96f3 Mon Sep 17 00:00:00 2001 From: Van Nguyen NGUYEN Date: Mon, 8 May 2023 17:33:42 +0200 Subject: [PATCH] Add all posible datasets of BOP and train script --- README.md | 11 +- configs/data/all.yaml | 8 +- configs/data/icbin.yaml | 9 ++ configs/data/icmi.yaml | 9 ++ configs/data/lm.yaml | 2 +- configs/data/olm.yaml | 2 +- configs/data/ruapc.yaml | 9 ++ configs/data/tudl.yaml | 8 ++ configs/download.yaml | 5 - configs/model/base.yaml | 14 +++ configs/model/resnet50.yaml | 14 +++ configs/train.yaml | 24 ++++ src/scripts/compute_neighbors.py | 109 +++++++++--------- src/scripts/download.py | 3 +- src/scripts/download_and_process_datasets.sh | 4 + .../download_and_process_from_scratch.sh | 7 -- src/scripts/process_mesh.py | 53 ++++++--- src/scripts/render_all.sh | 8 ++ src/scripts/render_template.py | 8 +- train.py | 96 +++++++++++++++ 20 files changed, 310 insertions(+), 93 deletions(-) create mode 100644 configs/data/icbin.yaml create mode 100644 configs/data/icmi.yaml create mode 100644 configs/data/ruapc.yaml create mode 100644 configs/data/tudl.yaml delete mode 100644 configs/download.yaml create mode 100644 configs/model/base.yaml create mode 100644 configs/model/resnet50.yaml create mode 100644 configs/train.yaml create mode 100644 src/scripts/download_and_process_datasets.sh delete mode 100644 src/scripts/download_and_process_from_scratch.sh create mode 100644 src/scripts/render_all.sh create mode 100644 train.py diff --git a/README.md b/README.md index ae296bd..1e8141a 100644 --- a/README.md +++ b/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 ``` diff --git a/configs/data/all.yaml b/configs/data/all.yaml index 60880c1..790e8fa 100644 --- a/configs/data/all.yaml +++ b/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 \ No newline at end of file + - /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 \ No newline at end of file diff --git a/configs/data/icbin.yaml b/configs/data/icbin.yaml new file mode 100644 index 0000000..5b38cd2 --- /dev/null +++ b/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: + diff --git a/configs/data/icmi.yaml b/configs/data/icmi.yaml new file mode 100644 index 0000000..e8d78dd --- /dev/null +++ b/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: + diff --git a/configs/data/lm.yaml b/configs/data/lm.yaml index 77fff12..98f39d2 100644 --- a/configs/data/lm.yaml +++ b/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: diff --git a/configs/data/olm.yaml b/configs/data/olm.yaml index b5c81a3..8be46d1 100644 --- a/configs/data/olm.yaml +++ b/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: diff --git a/configs/data/ruapc.yaml b/configs/data/ruapc.yaml new file mode 100644 index 0000000..6d93e0d --- /dev/null +++ b/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: + diff --git a/configs/data/tudl.yaml b/configs/data/tudl.yaml new file mode 100644 index 0000000..cf52f37 --- /dev/null +++ b/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: \ No newline at end of file diff --git a/configs/download.yaml b/configs/download.yaml deleted file mode 100644 index 0a97555..0000000 --- a/configs/download.yaml +++ /dev/null @@ -1,5 +0,0 @@ -defaults: - - user: default - - machine: local - - data: all - - _self_ \ No newline at end of file diff --git a/configs/model/base.yaml b/configs/model/base.yaml new file mode 100644 index 0000000..73b6bbf --- /dev/null +++ b/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 \ No newline at end of file diff --git a/configs/model/resnet50.yaml b/configs/model/resnet50.yaml new file mode 100644 index 0000000..9dd9099 --- /dev/null +++ b/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 \ No newline at end of file diff --git a/configs/train.yaml b/configs/train.yaml new file mode 100644 index 0000000..1025a39 --- /dev/null +++ b/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 diff --git a/src/scripts/compute_neighbors.py b/src/scripts/compute_neighbors.py index 8b6ba7b..2a5346e 100644 --- a/src/scripts/compute_neighbors.py +++ b/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__": diff --git a/src/scripts/download.py b/src/scripts/download.py index 1819e4c..8918d9f 100644 --- a/src/scripts/download.py +++ b/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) diff --git a/src/scripts/download_and_process_datasets.sh b/src/scripts/download_and_process_datasets.sh new file mode 100644 index 0000000..7dc1936 --- /dev/null +++ b/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 diff --git a/src/scripts/download_and_process_from_scratch.sh b/src/scripts/download_and_process_from_scratch.sh deleted file mode 100644 index d23fef5..0000000 --- a/src/scripts/download_and_process_from_scratch.sh +++ /dev/null @@ -1,7 +0,0 @@ -python -m src.scripts.download -python -m src.scripts.process_lm_gt -python -m src.scripts.render_template dataset_to_render=lm -python -m src.scripts.render_template dataset_to_render=tless -python -m src.scripts.process_mesh -python -m src.scripts.render_template dataset_to_render=hb -python -m src.scripts.compute_neighbors \ No newline at end of file diff --git a/src/scripts/process_mesh.py b/src/scripts/process_mesh.py index fe419c8..1c43b66 100644 --- a/src/scripts/process_mesh.py +++ b/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( diff --git a/src/scripts/render_all.sh b/src/scripts/render_all.sh new file mode 100644 index 0000000..ecdfe02 --- /dev/null +++ b/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 \ No newline at end of file diff --git a/src/scripts/render_template.py b/src/scripts/render_template.py index 4cc9193..109e517 100644 --- a/src/scripts/render_template.py +++ b/src/scripts/render_template.py @@ -43,7 +43,6 @@ def call_blender_proc( command += " false" os.system(command) - @hydra.main( version_base=None, config_path="../../configs", @@ -68,7 +67,7 @@ def render(cfg: DictConfig) -> None: cad_paths = [] output_dirs = [] object_ids = sorted( - [int(name[4:][:-4]) for name in os.listdir(cad_dir) if name.endswith(".ply")] + [int(name[4:][:-4]) for name in os.listdir(cad_dir) if name.endswith(".ply") and not name.endswith("old.ply")] ) for object_id in object_ids: cad_paths.append( @@ -77,6 +76,11 @@ def render(cfg: DictConfig) -> None: "obj_{:06d}.ply".format(object_id), ) ) + # hope and ycbv cad format is different which make the render is always black + # use obj format instead (make sure you have used python -m src.scripts.process_mesh to convert ply to obj) + # if cfg.dataset_to_render in ["hope", "ycbv"]: + # cad_paths[-1] = cad_paths[-1].replace(".ply", ".obj") + output_dirs.append( os.path.join( save_dir, diff --git a/train.py b/train.py new file mode 100644 index 0000000..ecf099b --- /dev/null +++ b/train.py @@ -0,0 +1,96 @@ +import logging +import os +import logging +import hydra +from omegaconf import DictConfig, OmegaConf +from torch.utils.data import DataLoader +import torch.nn as nn +from src.utils.weight import load_checkpoint +from src.dataloader.lm_utils import get_list_id_obj_from_split_name +import pytorch_lightning as pl + +pl.seed_everything(2022) +# set level logging +logging.basicConfig(level=logging.INFO) + + +@hydra.main(version_base=None, config_path="configs", config_name="train") +def train(cfg: DictConfig): + OmegaConf.set_struct(cfg, False) + hydra_cfg = hydra.core.hydra_config.HydraConfig.get() + output_path = hydra_cfg["runtime"]["output_dir"] + os.makedirs(cfg.callback.checkpoint.dirpath, exist_ok=True) + logging.info( + f"Training script. The outputs of hydra will be stored in: {output_path}" + ) + logging.info(f"Checkpoints will be stored in: {cfg.callback.checkpoint.dirpath}") + + # Delayed imports to get faster parsing + from hydra.utils import instantiate + + logging.info("Initializing logger, callbacks and trainer") + os.environ["WANDB_API_KEY"] = cfg.user.wandb_api_key + if cfg.machine.dryrun: + os.environ["WANDB_MODE"] = "offline" + logging.info(f"Wandb logger initialized at {cfg.save_dir}") + + if cfg.machine.name == "slurm": + cfg.machine.trainer.devices = int(os.environ["SLURM_GPUS_ON_NODE"]) + cfg.machine.trainer.num_nodes = int(os.environ["SLURM_NNODES"]) + trainer = instantiate(cfg.machine.trainer) + logging.info(f"Trainer initialized") + + model = instantiate(cfg.model) + logging.info(f"Model '{cfg.model.modelname}' loaded") + if cfg.model.pretrained_weight is not None: + load_checkpoint( + model, + cfg.model.pretrained_weight, + prefix="backbone.", + state_dict_key="model", + ) + + train_dataloaders = {} + for data_name in cfg.train_datasets: + config_dataloader = cfg.data[data_name] + train_dataloader = DataLoader( + instantiate(config_dataloader), + batch_size=cfg.machine.batch_size, + num_workers=cfg.machine.num_workers, + shuffle=True, + ) + logging.info( + f"Loading train dataloader with {data_name}, size {len(train_dataloader)} done!" + ) + train_dataloaders[data_name] = train_dataloader + from src.utils.dataloader import concat_dataloader + train_dataloaders = concat_dataloader(train_dataloaders) + + val_dataloaders = {} + for data_name in cfg.test_datasets: + config_dataloader = cfg.data[data_name] + val_dataloader = DataLoader( + instantiate(config_dataloader), + batch_size=cfg.machine.batch_size, + num_workers=cfg.machine.num_workers, + shuffle=False, + ) + val_dataloaders[data_name] = val_dataloader + logging.info( + f"Loading validation dataloader with {data_name}, size {len(val_dataloader)} done!" + ) + val_dataloaders = concat_dataloader(val_dataloaders) + logging.info("Fitting the model..") + trainer.fit( + model, + train_dataloaders=train_dataloaders, + val_dataloaders=val_dataloaders, + ckpt_path=cfg.model.checkpoint_path + if cfg.model.checkpoint_path is not None and cfg.use_pretrained + else None, + ) + logging.info(f"Fitting done") + + +if __name__ == "__main__": + train()