This is a Python3 / Pytorch implementation of TM-NET.
To run this code you need the following:
-
A machine with multiple GPUs(memory >= 12GB)
-
Python packages in the requirements.txt
pip install -r requirements.txt
- Run
GetTransformedCube.m
to get transformed a mini bounding box which will be used as source shape in non-rigid registration for each partial 3D model. - Run
SupportAnalysis.m
to extract structure information from the partialobj
files producing a correspondingcode.mat
for each 3D model. - Run
register.m
to perform non-rigid registrations from transformed mini bounding boxes to original partial 3D models. - Run
GenerateData.m
to extract deformation information between source shapes and registered shapes which will be used as the input ofTM-NET
. - Run
TransferColorPerPixelScript.m
to generate texture images for registered parts. - Run
PrepareForTraining.m
to split geometry, structure, image data to training or test dir. It will also divide texture image to six patches.
An example is shown in Pipeline.m
. After you successfully run the code, the dir structure will be like follows:
├─box50
│ ├─37b6df64a97a5c29369151623ac3890b
│ └─d374912c3a9fca96c141a04b2a487fd9
├─Chair
│ ├─37b6df64a97a5c29369151623ac3890b
│ | └─models
│ └─d374912c3a9fca96c141a04b2a487fd9
│ └─models
├─final50
│ ├─test
│ │ └─37b6df64a97a5c29369151623ac3890b
│ └─train
│ └─d374912c3a9fca96c141a04b2a487fd9
└─vaenew50
├─37b6df64a97a5c29369151623ac3890b
│ ├─back
│ ├─leg_ver_1
│ ├─leg_ver_2
│ ├─leg_ver_3
│ ├─leg_ver_4
│ └─seat
└─d374912c3a9fca96c141a04b2a487fd9
├─back
├─leg_ver_1
├─leg_ver_2
├─leg_ver_3
├─leg_ver_4
└─seat
Folder final50
is all we need for training and test.
- Train PartVAE for each part
python ./python/train.py --yaml ./python/yaml/table/surface/geovae.yml
python ./python/train.py --yaml ./python/yaml/table/leg/geovae.yml
- Train VQVAE
python ./python/train.py --yaml ./python/yaml/table/vqvae.yml
- Extract discrete code for the seed part
python ./python/extract_latents_central_part.py \
--image_dir ../data/table/ \
--mat_dir ../data/table \
--vqvae_ckpt ./table_vqvae/latest.pth \
--vqvae_yaml ./python/yaml/table/vqvae.yml \
--geovae_ckpt ./table_geovae/surface/latest.pth \
--geovae_yaml ./python/yaml/table/surface/geovae.yml \
--category table \
--save_path ./table_latents \
--device 0 \
--mode 'train' or 'test' or 'val'
- Train conditional PixelSNAIL for the seed part
python ./python/train.py --yaml ./python/yaml/table/surface/pixelsnail_top.yml
python ./python/train.py --yaml ./python/yaml/table/surface/pixelsnail_bottom.yml
- Extract discrete code for other parts
python ./python/extract_latents_other_parts.py \
--image_dir ../data/table/ \
--mat_dir ../data/table \
--vqvae_ckpt ./table_vqvae/latest.pth \
--vqvae_yaml ./python/yaml/table/vqvae.yml \
--geovae_ckpt_dir ./table_geovae \
--geovae_yaml ./python/yaml/table/geovae.yml \
--category table \
--save_path ./table_latents \
--device 0 \
--mode 'train' or 'test' or 'val'
- Train conditional PixelSNAIL for other parts
python ./python/train.py --yaml ./python/yaml/table/leg/pixelsnail_top.yml
python ./python/train.py --yaml ./python/yaml/table/leg/pixelsnail_bottom.yml
- Sample texture for the seed part
python ./python/conditional_sample_2levels_central_part.py \
--path ./table_latents \
--part_name surface \
--vqvae ./table_vqvae/latest.pth \
--vqvae_yaml ./python/yaml/table/vqvae.yml \
--top ./table_pixelsnail/top_16/latest.pth \
--top_yaml ./python/yaml/table/pixelsnail_top_center_16.yml \
--bottom ./table_pixelsnail/bottom/latest.pth \
--bottom_yaml ./python/yaml/table/pixelsnail_bottom_center.yml \
--device 0 \
--batch 1
- Sample texture for other parts
python ./python/conditional_sample_2levels_other_parts.py \
--path ./table_latents \
--central_part_name surface \
--part_name leg \
--vqvae ./table_vqvae/latest.pth \
--vqvae_yaml ./python/yaml/table/vqvae.yml \
--top ./table_pixelsnail/leg/top_16/latest.pth \
--top_yaml ./python/yaml/table/leg/pixelsnail_top_center_16.yml \
--bottom ./table_pixelsnail/leg/bottom/latest.pth \
--bottom_yaml ./python/yaml/table/leg/pixelsnail_bottom_center.yml \
--central_part_sample_dir ./table_pixelsnail/top_16/auto_texture \
--device 0 \
--batch 1