Update README.md

README.md

@@ -1,55 +1,57 @@
 # [Pre-train, Self-train, Distill: A simple recipe for Supersizing 3D Reconstruction](https://shubhtuls.github.io/ss3d/)
 Kalyan Vasudev Alwala, Abhinav Gupta, Shubham Tulsiani
 
-* [Paper](todo)
-* [Project Page](https://shubhtuls.github.io/ss3d/)
+[[Paper](todo)] [[Project Page](https://shubhtuls.github.io/ss3d/)]
 
-<img src="https://shubhtuls.github.io/ss3d/resources/overview.jpg" width="100%">
+<img src="https://shubhtuls.github.io/ss3d/resources/teaser.gif" width="50%">
 
-## Requirements
+## Setup
+Download the final distilled model from [here](https://dl.fbaipublicfiles.com/ss3d/distilled_model.torch).
+
+Install the following pre-requisites:
 * Python >=3.6
 * PyTorch tested with `1.10.0` 
 * TorchVision tested with `0.11.1`
 * Trimesh
 * pymcubes
 
-## 3D reconstruction example, 
+## 3D Reconstruction Interface
 
-Download the final distilled model from [here](https://dl.fbaipublicfiles.com/ss3d/distilled_model.torch).
+Reconstruct 3D in 3 simple simple steps! Please see the [demo notebook](demo.ipynb) for a working example.
 
 ```python
-import torch
-from model import VNet
-from data_utils import generate_input_img, extract_trimesh
 
-# Create Model and Load the pre-trained chekcpoint
+# 1. Load the pre-trained checkpoint
 model_3d = VNet()
 model_3d.load_state_dict(torch.load("<Path to the Model>"))
 model_3d.eval()
 
 
-# Load and process image for inference
-# NOTE: Expects mask and rgb images of same height and width and 
-# belonging only to a single object.
+# 2. Preprocess an RGB image with associated object mask according to our model's input interface
 inp_img = generate_input_img(
-    "<Path to your RGB image>",
-    "<Path to your corresponding Mask image>",
+    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()
 ```
 
-## Coming soon
-* Full training pipeline.
+## Training and Evaluation
+* coming soon
 
-<!-- # Citation
+
+## Citation
+If you find the project useful for your research, please consider citing:-
 ```
-TODO
+@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)}
+}
 ```
-
-# License
-TODO -->