A-template is all you need: 2D to 3D reconstruction with templates learned by contrastive learning
- Presentation link (mandarin)
- Report link
- Slide link
You must have a GUI for 2D image and SDF rendering!!
Environment of Ubuntu with GUI is recommended.
If you are unable to generate required 2D images and SDF ground truth, please jump to next section. We do provide some examples for you to download.
- Please prepare
secret.json
, which should include:
{"pascal_root":"path-to-save-PASCAL3D+_release1.1.zip",
"pascal_folder":"path-of-PASCAL3D+_release1.1-folder",
"shapenet_url":"download-link-of-ShapeNetCore.v2.zip",
"shapenet_root":"path-to-save-ShapeNetCore.v2.zip",
"shapenet_folder":"path-of-ShapeNetCore.v2-folder"}
- Utilize
initialize.py
to prepare everyting required.
python3 initialize.py
The code is heavily based on link. However, only images
and albedo
defined there will be generated. You should build the environment accurately first to avoid any erroneous 2D image rendering!
- Download blender v2.79. You must use v2.79 because some functions are deprecated in the latest version.
- Test if the command
blender
works or not.
blender
- Install OpenCV in the blender python env. Please (cd) to [blender]/2.79/python/bin
./python3.5m -m ensurepip
./python3.5m -m pip install --upgrade pip
./python3.5m -m pip install opencv-python
./python3.5m -m pip install opencv-contrib-python
- Try to
import numpy
in this python env. If there is any error, re-install numpy by deleting the numpy folder in the following path[blender]/2.79/python/lib/site-packages/
. Then,
./python3.5m -m pip install numpy
- Start generate the 2D images by the following command. Please notice that the
jsons_root
denotes./examples/splits/
,data_root
denotesShapeNetCore.v2/
, andoutput_root
is where you want to save the 2D images.
python batch.py --jsons_root $jsons_root --data_root $data_root --output_root $output_root
Please refer to link. In short, dependencies of CLI11
, Pangolin
, nanoflann
, and Eigen3
should be built first. However, the whole process is arduous and various bugs do exist!
python rgba2rgb.py -d [data_source] -c 'chairs', --level 'easy'
gdown is recommended for downloading these pre-generated 2D images and SDF ground truth.
pip install gdown
gdown --id 1BORVFh5ODnyhQqGTvCFtSUFrD7XCcSOr #02691156 planes
gdown --id 1WkyrOmvw2G3JgRutWvl_WucYn9RmFA8O #02958343 cars
gdown --id 1M53cyqhkyPo8IZgQaDjo4qNZc9hhkSsi #03001627 chairs
gdown --id 1NBwU309G4FT2SP_gxZzne2uADqbD4Rce #03636649 lamps
gdown --id 1DplfrOr98VEW5xAV8jEy2yNrlOetjNGF #04256520 sofas
gdown --id 1GetjpIDWa7NOIN06tuQQJmR0MuVKwoal #04379243 tables
gdown --id 11y3A8RI4n4vlUVTDu56FO2mc7rhASjr6 #02691156 planes
gdown --id 1_ogLvFOZTg6YFak-_WIvLneMoYYcLRvs #02958343 cars
gdown --id 1lLm0yyrGfEjT7ye3xjbrgW7p2neYHgZ2 #03001627 chairs
gdown --id 1ObsXX4TJjADnlyuFPVJxz00w-fZFK2ak #03636649 lamps
gdown --id 18k8LqWlc2AYr5-gn2u_KQsodGg9macqx #04256520 sofas
gdown --id 1TzUSNe4kAB3plnUf_ZUsxTnVgNq0gqzJ #04379243 tables
gdown --id 1W1oP4zHWP0gPXFAGWGm9rIhCrUlXfhtA #02691156 planes
gdown --id 17UsiUaeSF2O9BbOoCONg-BEfW6_3GTZm #03001627 chairs
gdown --id 1FKSDTitAmIQNODEj5mSQUMFyoPJOpnY0 #04256520 sofas
gdown --id 156WckbItDBjmDL89J7yrNxuL2MOQSmam #02691156 planes
gdown --id 1F7xs-3v0yA0IcIhEn5AaeKdaAEaZYlOA #03001627 chairs
gdown --id 1vI0qsAtJ2kwSx8KpyH9sJu_2fCGh1DOu #04256520 sofas
gdown --id 1cLW_B_EmhQNeM-pJ2iJIKGFdf7Cyzk6E #planes #end2end with contrastive
gdown --id 11kGrEJqINyfqwcT0aIEoHRC2uwgcz_GF #planes #two phases
gdown --id 1GXpNWaz5njg8x1zC9Ntf83CDfsRtSGdS #chairs #end2end
gdown --id 1bN3B4ws-PXd1kmCrkGSl_4VAcgRgfK7I #chairs #two phases
gdown --id 1lq1-cK86g10EK53E4gskUwnqc-zd9n7p #sofas #end2end
gdown --id 10I4sWk6XjfV6koe00jUkUmhi792YQ-Ot #sofas #two phases
The training data of chairs 03001627_train.tar.gz
can be downloaded from
this link.
There will be a folder named 03001627
when the file is extracted.
To download the file and merge all the extracted data to your dataset folder, run
bash download_chairs_train.sh [data_source_name]
where [data_source_name]
should be the path to yor dataset folder.
To be more explicit, it should be the parent folder of SdfSamples
.
Therefore, before running this script, please make sure that all the other data has been prepared
and the folder structure is the same as the one specified in Data Layouts.
For example, run
bash download_chairs_train.sh data/
The data must be organized as the following hierarchical structure.
<data_source_name>/
.datasources.json
SdfSamples/
<dataset_name>/
<class_name>/
<instance_name>/
<instance_name>.npz
SurfaceSamples/
<dataset_name>/
<class_name>/
<instance_name>/
<instance_name>.ply
NormalizationParameters/
<dataset_name>/
<class_name>/
<instance_name>/
<instance_name>.npz
<class_name>/
<instance_name>/
<easy/hard>
<image>/
XXXXXXXX.png
XXXXXXXX.png
.
<image_rgb>/
XXXXXXXX.png
XXXXXXXX.png
.
Please use extract_embedding.py
in the environment of Deep Implicit Template. This code will automatically extract the embedding from the pretrained Embedding
weights and store in a dictionary, which can be accessed by the instance_name
.
python generate_training_meshes.py -e ./pretrained/${obj}_dit
We also provided the pre-extracted one in ./pretrained_embedding/
folder.
Phase-1 training can also be conducted by the original code in deep implicit template. However, the required embedding obtained in this step has been provided in ./pretrained_embedding/
and our program will directly load the embedding there. Thus, you don't need to do the phase-1 training but need to copy the pretrained models of wrapping and SDF network provided here. Please prepare /ModelParameters/latest.pth
for every object category.
python train_solution1.py -e pretrained/${obj}_dit/ -d ./data/
python train_deep_implicit_templates.py -e examples/${obj}_dit --debug --batch_split 2 -d ./data
To expedite training, the mixed precision mode is provided:
Package apex is required! Please make sure that you have installed it first!
python train_deep_implicit_templates.py -e examples/${obj}_dit --debug --batch_split 2 -d ./data --mixed_precision --mixed_precision_level O1
Plese note that mixed_precision_level
has four options, O0
, O1
, O2
, O3
, corresponding to different settings.
O0
: FP32 trainingO1
: Mixed Precision (recommended for typical use)O2
: “Almost FP16” Mixed PrecisionO3
: FP16 training
For more information, please refer to the official document.
If you cannot successfully install apex
from pypl. Please refer to link to build from source, and add argument to indicate where your apex
is.
--apex_path path_to_apex
python train_deep_implicit_templates_v2.py \
-e examples/${obj}_dit \
--debug \
--batch_split 2 \
-d ./data \
[--mixed_precision] \
[--pretrained_weights exps/checkpoints/checkpont_latest.pt]
train_deep_implicit_templates_v2.py
contains a few updates compared to the original one, including:
- start training with weights obtained from contrastive learning.
- please use
--pretrained_weights
to specify the path to the checkpoint obtained fromtrain_contrastive.py
. - here a slightly different encoder
_Encoder
(check it out innetworks/encoder.py
) is used for easy weight transfer.
- please use
- use torch.cuda.amp for mixed precision training.
- please add
--mixed_precision
to enable such feature.
- please add
- train with a dataset that includes both hard and easy 2d images.
- if you want to train with certain level of difficulty instead, please modify
*
toeasy/hard
in this line.
- if you want to train with certain level of difficulty instead, please modify
We can eventually replace train_deep_implicit_templates.py
with train_deep_implicit_templates_v2.py
if everyone has
no problem running this version of training code.
Encoders that apply different contrastive learning frameworks can be found in
the folder contrastive/encoders
.
For more details please refer to ENCODER.md.
Run the following to perform contrastive learning:
python train_contrastive.py -e experiments/planes
Note that the experiment directory should include a config.yaml
to specify the configurations.
An example of config.yaml
can be found here.
To use the pretrained weights for the later deep implicit templates training, please refer to ENCODER.md.
python generate_template_mesh.py -e examples/${obj}_dit --debug
view_id
starts from 0
to 50
. If it is set to -1
, random view_id
will be selected for each instance.
python generate_training_meshes.py -e examples/${obj}_dit --debug --start_id 0 --end_id 20 --octree --keep_normalization --view_id $id
view_id
starts from 0
to 50
. If it is set to -1
, random view_id
will be selected for each instance.
python generate_training_meshes.py -e examples/${obj}_dit --debug --start_id 0 --end_id 20 --octree --keep_normalization --mode test --view_id $id
python generate_general_meshes.py -e examples/${obj}_dit --octree --image_path $where_you_put_the_real_world_images
python evaluate.py -c $class_name -d $data_src -i $mesh_output_folder
e.g.
python evaluate.py -c sofas -d ./data/ -i examples/sofas_dit/TrainingMeshes/2000/ShapeNetV2/04256520/
- To analyze the pretrained embedding
python analyze.py -e examples/${obj}_dit -p --thread 16
- To analyze the embedding yield from own encoder
early_stop
is to prevent the time-consuming inference process and only extract some data for analysis
python analyze.py -e examples/${obj}_dit -d $data_path -c latest --early_stop 20 --thread 16
This code repo is heavily based on Deep Implicit Template. We thank the authors for their great job!