lordliang / ga-sketching

(PG2023/CGF) This is the official PyTorch implementation of PG2023/CGF paper: GA-Sketching: Shape Modeling from Multi-View Sketching with Geometry-Aligned Deep Implicit Functions

GA-Sketching: Shape Modeling from Multi-View Sketching with Geometry-Aligned Deep Implicit Functions

Introduction

Sketch-based shape modeling aims to bridge the gap between 2D drawing and 3D modeling by providing an intuitive and accessible approach to create 3D shapes from 2D sketches. However, existing methods still suffer from limitations in reconstruction quality and multi-view interaction friendliness, hindering their practical application. This paper proposes a faithful and user-friendly iterative solution to tackle these limitations by learning geometry-aligned deep implicit functions from one or multiple sketches. Our method lifts 2D sketches to volume-based feature tensors, which align strongly with the output 3D shape, enabling accurate reconstruction and faithful editing. Such a geometry-aligned feature encoding technique is well-suited to iterative modeling since features from different viewpoints can be easily memorized or aggregated. Based on these advantages, we design a unified interactive system for sketch-based shape modeling. It enables users to generate the desired geometry iteratively by drawing sketches from any number of viewpoints. In addition, it allows users to edit the generated surface by making few local modifications.

Setup

The code is tested on Ubuntu 18.04 with PyTorch 1.12.1 CUDA 11.6 installed. Please follow the following steps to install PyTorch and PyTorch3D first. All experiments are conducted on a single NVIDIA GeForce RTX 2080 Ti gpu.

# create and activate the conda environment
conda create -n gas python=3.9.15
conda activate gas

# install necessary packages
conda install pytorch==1.12.1 torchvision==0.13.1 torchaudio==0.12.1 cudatoolkit=11.6 -c pytorch
conda install -c fvcore -c iopath -c conda-forge fvcore iopath
conda install -c bottler nvidiacub
conda install pytorch3d -c pytorch3d

pip install -r requirements.txt

cd tools/libmesh/
python setup.py build_ext --inplace

Data generation

First you can obtain watertight and simplified meshes via https://github.com/davidstutz/mesh-fusion and put them into 'datasets/shapenet/watertight_simplified_off/'. Then run script:

# render depth map and normal map
python tools/sketch_render/gen_dn_from_mesh.py
# render sketch
python tools/sketch_render/gen_sketch_from_dn.py
# generate boundary sampling
python tools/boundary_sampling.py

Training and Evaluation

Please download our predtrained models from OneDrive.

# Training single-view network
python train_single.py -cat airplane

# Evaluating single-view network
python generate_single.py  -cat airplane -checkpoint 200
python tools/evaluate.py -cat airplane -generation_path experiments/GASv64_airplane/evaluation_200@128_1v/generation/

# Training multi-view network
python train_multi.py -cat airplane

# Evaluating multi-view network
python generate_multi.py  -cat airplane -checkpoint 200 -n 2
python tools/evaluate.py -cat airplane -generation_path experiments/GASv64_airplane/evaluation_200@128_2v/generation/
python generate_multi.py  -cat airplane -checkpoint 200 -n 3
python tools/evaluate.py -cat airplane -generation_path experiments/GASv64_airplane/evaluation_200@128_3v/generation/

User Interface

python GA-Sketching-UI/main.py

The generated sketches and meshes will be saved in 'cache/'.

Demo Video

demo.mp4

ga-sketching's Issues

About the rendered results in the paper.

Thanks for your excellent work!!! Could you please tell me the method for rendering the mesh results in the paper?

render 2d for comparison to generated sketch.

Dear author：
After I generated the sketches by your script, I want to compare to 2d color-img rendered by 3D mesh. Do you have the script to get the rendered 2d img? I would greatly appreciate it.

issue when evaluate using pretrained model

Sry to trouble you again. When I try to do evaluations based on the pretrained model on OneDrive, KeyError occurs:

(gas) vic@PC:~/GAS/GA-Sketching$ python generate_single.py  -cat airplane -checkpoint 200
initialize network with normal
Loaded checkpoint from: /home/vic/GAS/GA-Sketching/core/../experiments/GASv64_airplane/checkpoints/checkpoint_epoch_200.tar
**checkpoint keys**: odict_keys(['volume_encoder.conv_in_0.weight', 'volume_encoder.conv_in_0.bias', 'volume_encoder.conv_in_1.weight', 'volume_encoder.conv_in_1.bias', 'volume_encoder.conv_in_bn.weight', 'volume_encoder.conv_in_bn.bias', 'volume_encoder.conv_in_bn.running_mean', 'volume_encoder.conv_in_bn.running_var', 'volume_encoder.conv_in_bn.num_batches_tracked', 'ifnet_encoder.conv_0.weight', 'ifnet_encoder.conv_0.bias', 'ifnet_encoder.conv_0_1.weight', 'ifnet_encoder.conv_0_1.bias', 'ifnet_encoder.conv_1.weight', 'ifnet_encoder.conv_1.bias', 'ifnet_encoder.conv_1_1.weight', 'ifnet_encoder.conv_1_1.bias', 'ifnet_encoder.conv_2.weight', 'ifnet_encoder.conv_2.bias', 'ifnet_encoder.conv_2_1.weight', 'ifnet_encoder.conv_2_1.bias', 'ifnet_encoder.conv_3.weight', 'ifnet_encoder.conv_3.bias', 'ifnet_encoder.conv_3_1.weight', 'ifnet_encoder.conv_3_1.bias', 'ifnet_encoder.conv0_1_bn.weight', 'ifnet_encoder.conv0_1_bn.bias', 'ifnet_encoder.conv0_1_bn.running_mean', 'ifnet_encoder.conv0_1_bn.running_var', 'ifnet_encoder.conv0_1_bn.num_batches_tracked', 'ifnet_encoder.conv1_1_bn.weight', 'ifnet_encoder.conv1_1_bn.bias', 'ifnet_encoder.conv1_1_bn.running_mean', 'ifnet_encoder.conv1_1_bn.running_var', 'ifnet_encoder.conv1_1_bn.num_batches_tracked', 'ifnet_encoder.conv2_1_bn.weight', 'ifnet_encoder.conv2_1_bn.bias', 'ifnet_encoder.conv2_1_bn.running_mean', 'ifnet_encoder.conv2_1_bn.running_var', 'ifnet_encoder.conv2_1_bn.num_batches_tracked', 'ifnet_encoder.conv3_1_bn.weight', 'ifnet_encoder.conv3_1_bn.bias', 'ifnet_encoder.conv3_1_bn.running_mean', 'ifnet_encoder.conv3_1_bn.running_var', 'ifnet_encoder.conv3_1_bn.num_batches_tracked', 'fc_0.weight', 'fc_0.bias', 'fc_1.weight', 'fc_1.bias', 'fc_2.weight', 'fc_2.bias', 'fc_out.weight', 'fc_out.bias'])
Traceback (most recent call last):
  File "/home/vic/GAS/GA-Sketching/generate_single.py", line 23, in <module>
    gen = Generator(net,
  File "/home/vic/GAS/GA-Sketching/core/generator.py", line 17, in __init__
    self.load_checkpoint(checkpoint)
  File "/home/vic/GAS/GA-Sketching/core/generator.py", line 91, in load_checkpoint
    self.model.load_state_dict(checkpoint['model_state_dict'])
KeyError: 'model_state_dict'

I print the checkpoint keys, is there something wrong with the way I load checkpoint?
I then change the code for loading the checkpoint to:

        #self.model.load_state_dict(checkpoint['model_state_dict'])
        state_dict = self.model.state_dict()
        for k1, k2 in zip(state_dict.keys(), checkpoint.keys()):
            state_dict[k1] = checkpoint[k2]
        self.model.load_state_dict(state_dict)

Then seems everything works fine, but then another issue occurs:

(gas) vic@PC:~/GAS/GA-Sketching$ python generate_single.py  -cat airplane -checkpoint 200
initialize network with normal
Loaded checkpoint from: /home/vic/GAS/GA-Sketching/core/../experiments/GASv64_airplane/checkpoints/checkpoint_epoch_200.tar
experiments/GASv64_airplane/evaluation_200@128_1v/
0it [00:00, ?it/s]
Traceback (most recent call last):
  File "/home/vic/GAS/GA-Sketching/generate_single.py", line 31, in <module>
    gen_iterator(out_path, dataset, gen)    
  File "/home/vic/GAS/GA-Sketching/core/generation_iterator.py", line 19, in gen_iterator
    for i, data in tqdm(enumerate(loader)):
  File "/home/vic/anaconda3/envs/gas/lib/python3.9/site-packages/tqdm/std.py", line 1182, in __iter__
    for obj in iterable:
  File "/home/vic/anaconda3/envs/gas/lib/python3.9/site-packages/torch/utils/data/dataloader.py", line 681, in __next__
    data = self._next_data()
  File "/home/vic/anaconda3/envs/gas/lib/python3.9/site-packages/torch/utils/data/dataloader.py", line 721, in _next_data
    data = self._dataset_fetcher.fetch(index)  # may raise StopIteration
  File "/home/vic/anaconda3/envs/gas/lib/python3.9/site-packages/torch/utils/data/_utils/fetch.py", line 49, in fetch
    data = [self.dataset[idx] for idx in possibly_batched_index]
  File "/home/vic/anaconda3/envs/gas/lib/python3.9/site-packages/torch/utils/data/_utils/fetch.py", line 49, in <listcomp>
    data = [self.dataset[idx] for idx in possibly_batched_index]
  File "/home/vic/GAS/GA-Sketching/core/data/voxelized_data_shapenet.py", line 225, in __getitem__
    a0, e0, a1, e1, a2, e2 = self.dict[uid]
KeyError: '02691156/1b8e84935fdc3ec82be289de70e8db31'

Here I use the sketch you provided yesterday, I copy their uid to create a .lst file, I thought it might point to the sketch data you provided. My understanding about this program: the input is image and output is a corresponding mesh, so actually what I am confused here is where should I put my reference image for the model? What is the format of the input data?
Thanks!

mesh_type and pytorch3d version

Dear author:
I want to ask the pytorch3d version if I use pip(ERROR: Could not find a version that satisfies the requirement pytorch3d==2.31 (from versions: 0.1.1, 0.2.0, 0.2.5, 0.3.0));
And I want to use my own mesh data, should I use the the 'mesh-fusion' to change my mesh to 'xxx.off' ? The normal '.obj' mesh can not work?
I am looking forward to your reply!

problem about running the UI interface

I find that in edit_widget.py, there are some functions about shadow, and I also check the paper ShadowDraw. My problem is, is the Data Generation step necessary? Is it for building a shadow dataset to guide the users? Can you provide thoes images or I have to build it myself?
Many thanks!

reproduction error

self.model.merger.load_state_dict(checkpoint['model_state_dict'])
KeyError: 'model_state_dict'