search

RuyiLian / CheckerPose

The implementation of the paper 'CheckerPose: Progressive Dense Keypoint Localization for Object Pose Estimation with Graph Neural Network' (ICCV2023).
9 stars 3 forks source link

Custom Datasets #4

Open WHHHHY opened 6 days ago

WHHHHY commented 6 days ago

Thanks for your excellent work! And if I want train my customized datasets made in linemod format，how should I modify the paths and codes?

RuyiLian commented 4 days ago

Thanks for your interest of our work! I assume that you would like to use a single network to predict multiple objects simultaneously, like our Linemod experiments reported in the paper. For this dataset we have separate files for dataset loading , training and testing .

Dataset codes are mainly in checkerpose/tools_for_LM/get_lm_datasets.py and checkerpose/lm_dataset_pytorch.py. The dataset paths can be modified in checkerpose/tools_for_LM/get_lm_datasets.py. Training codes can be found in checkerpose/pretrain_lm.py and checkerpose/train_lm.py. (The name of pretrain is kind of misleading, it actually means training only the low level layers of the whole network.) And test file is checkerpose/test_lm.py. For these files, reset the dataset related codes. Also, write new configs similar to checkerpose/config/lm/init_gnn2_hrnetw18_npt512_lm.txt (for pretrain) and checkerpose/config/lm/hr18GNN2_res6_gnn3Skip_mlpQuery_lm.txt (for train).

For Linemod experiments, we use two training subsets. One is the images from BOP website, and the other one is a synthetic dataset downloaded from GDR-Net GitHub repo. You may need to write

training_data_folder_2 = none

in your config files if you only have one single training subset.

Also, in our default setting, we randomly reset the background images https://github.com/RuyiLian/CheckerPose/blob/fb725bed2a9eb6c3646c476c017d9ad00aa63a28/checkerpose/config/lm/hr18GNN2_res6_gnn3Skip_mlpQuery_lm.txt#L11 We use VOC 2012 as background images. The details can be found in our README (data preparation, 2 and 4) You can just discard codes of changing background if it is unnecessary.