smks17 / ss3d

Code release for "Pre-train, Self-train, Distill A simple recipe for Supersizing 3D Reconstruction" but with VAE

NOTE (For the forked)

This fork changes the model to Variational Autoencoder to test improvement.

Also uses ModelNet datasets because of not available to access ShapeNet.

Pre-train, Self-train, Distill: A simple recipe for Supersizing 3D Reconstruction

Kalyan Vasudev Alwala, Abhinav Gupta, Shubham Tulsiani

[Paper] [Project Page]

Setup

Download the final distilled model from here.

Install the following pre-requisites:

• Python >=3.6
• PyTorch tested with 1.10.0
• TorchVision tested with 0.11.1
• Trimesh
• pymcubes

3D Reconstruction Interface

Reconstruct 3D in 3 simple simple steps! Please see the demo notebook for a working example.

# 1. Load the pre-trained checkpoint
model_3d = VNet()
model_3d.load_state_dict(torch.load("<Path to the Model>"))
model_3d.eval()


# 2. Preprocess an RGB image with associated object mask according to our model's input interface
inp_img = generate_input_img(
    img_rgb,
    img_mask,
)

# 3. Obtain 3D prediction!
out_mesh = extract_trimesh(model_3d, inp_img, "cuda")
# To save the mesh
out_mesh.export("out_mesh_pymcubes.obj")
# To visualize the mesh
out_mesh.show()

Training and Evaluation

Please to README_TRAINING.md for more details.

Citation

If you find the project useful for your research, please consider citing:-

@inproceedings{vasudev2022ss3d,
  title={Pre-train, Self-train, Distill: A simple recipe for Supersizing 3D Reconstruction},
  author={Vasudev, Kalyan Alwala and  Gupta, Abhinav and Tulsiani, Shubham},
  year={2022},
  booktitle={Computer Vision and Pattern Recognition (CVPR)}
}

Contributing

We welcome your pull requests! Please see CONTRIBUTING and CODE_OF_CONDUCT for more information.

License

ss3d is released under the CC-BY-NC 4.0 license. See LICENSE for additional details. However the Sire implementation is additionally licensed under the MIT license (see NOTICE for additional details).