├── .gitignore ├── README.md ├── apps ├── eval_3d.py ├── train_dmc.py └── train_normal_net.py ├── assets ├── obj │ └── skirt │ │ └── skirt.obj ├── smplx │ └── skirt │ │ └── smplx.obj ├── teaser.jpg └── tex │ └── skirt.jpg ├── configs ├── dmc_demo.yaml ├── dmc_train.yaml └── normal_train.yaml ├── dataset └── multihuman │ ├── img │ ├── 20_0 │ │ ├── 0.jpg │ │ ├── 180.jpg │ │ ├── 270.jpg │ │ └── 90.jpg │ └── 20_1 │ │ ├── 0.jpg │ │ ├── 180.jpg │ │ ├── 270.jpg │ │ └── 90.jpg │ ├── mask │ ├── 20_0 │ │ ├── 0.jpg │ │ ├── 180.jpg │ │ ├── 270.jpg │ │ └── 90.jpg │ └── 20_1 │ │ ├── 0.jpg │ │ ├── 180.jpg │ │ ├── 270.jpg │ │ └── 90.jpg │ ├── parameter │ ├── 20_0 │ │ ├── 0_extrinsic.npy │ │ ├── 0_intrinsic.npy │ │ ├── 180_extrinsic.npy │ │ ├── 180_intrinsic.npy │ │ ├── 270_extrinsic.npy │ │ ├── 270_intrinsic.npy │ │ ├── 90_extrinsic.npy │ │ └── 90_intrinsic.npy │ └── 20_1 │ │ ├── 0_extrinsic.npy │ │ ├── 0_intrinsic.npy │ │ ├── 180_extrinsic.npy │ │ ├── 180_intrinsic.npy │ │ ├── 270_extrinsic.npy │ │ ├── 270_intrinsic.npy │ │ ├── 90_extrinsic.npy │ │ └── 90_intrinsic.npy │ └── smplx │ ├── 20_0 │ └── smplx.obj │ └── 20_1 │ └── smplx.obj ├── demo.sh ├── lib ├── __init__.py ├── data │ ├── DMCDataset.py │ ├── NormalDataset.py │ ├── __init__.py │ ├── smplx_correction.obj │ ├── smplx_fine.obj │ ├── smplx_lz_faces.obj │ └── smplx_multi.obj ├── ext_transform.py ├── geometry.py ├── loss_util.py ├── mesh_util.py ├── model │ ├── BaseNet.py │ ├── DMCNet.py │ ├── HGFilters.py │ ├── NormalNet.py │ ├── PositionEncoder.py │ ├── ResNet3d.py │ ├── SurfaceClassifier.py │ ├── __init__.py │ └── attention.py ├── net_util.py ├── options.py ├── sample_util.py ├── sdf.py ├── smpl_util.py └── train_util.py ├── render_dataset.sh ├── render_two.sh ├── taichi_render_gpu ├── README.md ├── render_multi.py ├── render_smpl.py ├── setup.py └── taichi_three │ ├── __init__.py │ ├── common.py │ ├── geometry.py │ ├── light.py │ ├── loader.py │ ├── meshgen.py │ ├── model.py │ ├── raycast.py │ ├── scatter.py │ ├── scene.py │ ├── shading.py │ ├── transform.py │ └── version.py └── train.sh /.gitignore: -------------------------------------------------------------------------------- 1 | # Byte-compiled / optimized / DLL files 2 | __pycache__/ 3 | *.py[cod] 4 | *$py.class 5 | 6 | # C extensions 7 | *.so 8 | 9 | # Distribution / packaging 10 | .Python 11 | build/ 12 | develop-eggs/ 13 | dist/ 14 | downloads/ 15 | eggs/ 16 | .eggs/ 17 | #lib/ 18 | lib64/ 19 | parts/ 20 | sdist/ 21 | var/ 22 | wheels/ 23 | pip-wheel-metadata/ 24 | share/python-wheels/ 25 | *.egg-info/ 26 | .installed.cfg 27 | *.egg 28 | MANIFEST 29 | 30 | # PyInstaller 31 | # Usually these files are written by a python script from a template 32 | # before PyInstaller builds the exe, so as to inject date/other infos into it. 33 | *.manifest 34 | *.spec 35 | 36 | # Installer logs 37 | pip-log.txt 38 | pip-delete-this-directory.txt 39 | 40 | # Unit test / coverage reports 41 | htmlcov/ 42 | .tox/ 43 | .nox/ 44 | .coverage 45 | .coverage.* 46 | .cache 47 | nosetests.xml 48 | coverage.xml 49 | *.cover 50 | *.py,cover 51 | .hypothesis/ 52 | .pytest_cache/ 53 | 54 | # Translations 55 | *.mo 56 | *.pot 57 | 58 | # Django stuff: 59 | *.log 60 | local_settings.py 61 | db.sqlite3 62 | db.sqlite3-journal 63 | 64 | # Flask stuff: 65 | instance/ 66 | .webassets-cache 67 | 68 | # Scrapy stuff: 69 | .scrapy 70 | 71 | # Sphinx documentation 72 | docs/_build/ 73 | 74 | # PyBuilder 75 | target/ 76 | 77 | # Jupyter Notebook 78 | .ipynb_checkpoints 79 | 80 | # IPython 81 | profile_default/ 82 | ipython_config.py 83 | 84 | # pyenv 85 | .python-version 86 | 87 | # pipenv 88 | # According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control. 89 | # However, in case of collaboration, if having platform-specific dependencies or dependencies 90 | # having no cross-platform support, pipenv may install dependencies that don't work, or not 91 | # install all needed dependencies. 92 | #Pipfile.lock 93 | 94 | # PEP 582; used by e.g. github.com/David-OConnor/pyflow 95 | __pypackages__/ 96 | 97 | # Celery stuff 98 | celerybeat-schedule 99 | celerybeat.pid 100 | 101 | # SageMath parsed files 102 | *.sage.py 103 | 104 | # Environments 105 | .env 106 | .venv 107 | env/ 108 | venv/ 109 | ENV/ 110 | env.bak/ 111 | venv.bak/ 112 | 113 | # Spyder project settings 114 | .spyderproject 115 | .spyproject 116 | 117 | # Rope project settings 118 | .ropeproject 119 | 120 | # mkdocs documentation 121 | /site 122 | 123 | # mypy 124 | .mypy_cache/ 125 | .dmypy.json 126 | dmypy.json 127 | 128 | # Pyre type checker 129 | .pyre/ 130 | 131 | # vscode 132 | .vscode 133 | 134 | # python wheel 135 | .whl 136 | 137 | # dataset 138 | dataset 139 | checkpoints 140 | log 141 | results 142 | output 143 | view* 144 | show 145 | metric 146 | 3rd 147 | freeview 148 | .zip 149 | .pylintrc 150 | .idea 151 | dmc_demo 152 | logs -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # [DeepMultiCap: Performance Capture of Multiple Characters Using Sparse Multiview Cameras (ICCV 2021)](http://www.liuyebin.com/dmc/dmc.html) 2 | [Yang Zheng*](https://y-zheng18.github.io/zy.github.io/), Ruizhi Shao*, [Yuxiang Zhang](https://zhangyux15.github.io/), [Tao Yu](http://ytrock.com/), [Zerong Zheng](http://zhengzerong.github.io/), [Qionghai Dai](http://media.au.tsinghua.edu.cn/english/team/qhdai.html), [Yebin Liu](http://www.liuyebin.com/). 3 | 4 | [![report](https://img.shields.io/badge/arxiv-report-red)](https://arxiv.org/abs/2105.00261) 5 | 6 | This repository contains the official pytorch implementation of ”*DeepMultiCap: Performance Capture of Multiple Characters Using Sparse Multiview Cameras*“. 7 | 8 | ![Teaser Image](assets/teaser.jpg) 9 | 10 | ### News 11 | * **[2021/9/26]** We added more scans to [MultiHuman dataset](https://github.com/y-zheng18/MultiHuman-Dataset). You can use MultiHuman to train/fine-tune our model or your own models! 12 | * **[2021/9/18]** [MultiHuman dataset](https://github.com/y-zheng18/MultiHuman-Dataset) for evaluation purpose is available! 13 | 14 | ## Requirements 15 | - [PyTorch](https://pytorch.org/) 16 | - torchvision 17 | - trimesh 18 | - numpy 19 | - matplotlib 20 | - PIL 21 | - skimage 22 | - tqdm 23 | - cv2 24 | - json 25 | - taichi==0.6.39 or 0.7.15 26 | - taichi_three 27 | - taichi_glsl==0.0.10 28 | 29 | ## Pretrained Model 30 | We have provided pretrained model in [One Drive](https://mailstsinghuaeducn-my.sharepoint.com/:u:/g/personal/shaorz20_mails_tsinghua_edu_cn/EdVJtlpRplRHvGzQENV8ESQB4E_0ZY3B9l76XHuEowj1YA?e=MZqUxM) 31 | 32 | 33 | ## Training 34 | To train DeepMultiCap model, please run the following code. 35 | ``` 36 | sh render_dataset.sh 37 | sh train.sh 38 | ``` 39 | 40 | ## Demo 41 | Run the following code to inference on the MultiHuman dataset. 42 | ``` 43 | sh render_two.sh 44 | sh demo.sh 45 | ``` 46 | 47 | ## Citation 48 | ``` 49 | @inproceedings{zheng2021deepmulticap, 50 | title={DeepMultiCap: Performance Capture of Multiple Characters Using Sparse Multiview Cameras}, 51 | author={Zheng, Yang and Shao, Ruizhi and Zhang, Yuxiang and Yu, Tao and Zheng, Zerong and Dai, Qionghai and Liu, Yebin}, 52 | booktitle={IEEE Conference on Computer Vision (ICCV 2021)}, 53 | year={2021}, 54 | } 55 | ``` 56 | -------------------------------------------------------------------------------- /apps/eval_3d.py: -------------------------------------------------------------------------------- 1 | import sys 2 | import os 3 | 4 | sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..'))) 5 | ROOT_PATH = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) 6 | 7 | import time 8 | import json 9 | import numpy as np 10 | import cv2 11 | import random 12 | import torch 13 | from torch.utils.data import DataLoader 14 | from torch.nn import DataParallel 15 | from tqdm import tqdm 16 | import matplotlib.pyplot as plt 17 | from torch.utils.tensorboard import SummaryWriter 18 | import torch.nn.functional as F 19 | from multiprocessing import Process, Manager, Lock 20 | 21 | from lib.options import parse_config 22 | from lib.mesh_util import * 23 | from lib.sample_util import * 24 | from lib.train_util import * 25 | from lib.data import * 26 | from lib.model import * 27 | from lib.geometry import index 28 | 29 | import matplotlib 30 | matplotlib.use('AGG') 31 | 32 | # get options 33 | opt = parse_config() 34 | 35 | def inference_nm(opt): 36 | # set cuda 37 | cuda = torch.device('cuda:%s' % opt.gpu_ids[0]) 38 | netG = DMCNet(opt, projection_mode='perspective').to(device=cuda) 39 | netN = NormalNet().to(cuda) 40 | print('Using Network: ', netG.name) 41 | gpu_ids = [int(i) for i in opt.gpu_ids.split(',')] 42 | netG = DataParallel(netG, device_ids=gpu_ids) 43 | netN = DataParallel(netN, device_ids=gpu_ids) 44 | # load checkpoints 45 | if opt.load_netG_checkpoint_path is not None: 46 | print('loading for net G ...', opt.load_netG_checkpoint_path) 47 | netG.load_state_dict(torch.load(opt.load_netG_checkpoint_path, map_location=cuda), strict=False) 48 | 49 | if opt.load_netN_checkpoint_path is not None: 50 | print('loading for net N ...', opt.load_netN_checkpoint_path) 51 | netN.load_state_dict(torch.load(opt.load_netN_checkpoint_path, map_location=cuda), strict=False) 52 | 53 | print("loaded finished!") 54 | 55 | test_netG = netG.module 56 | netN = netN.module 57 | dataset = DMCDataset(opt, phase='inference', yaw_list=opt.yaw_list, cache_data = Manager().dict(), cache_data_lock=Lock()) 58 | print(dataset.__len__()) 59 | 60 | os.makedirs(opt.results_path, exist_ok=True) 61 | os.makedirs('%s/%s' % (opt.results_path, opt.name), exist_ok=True) 62 | 63 | opt_log = os.path.join(opt.results_path, opt.name, 'opt.txt') 64 | with open(opt_log, 'w') as outfile: 65 | outfile.write(json.dumps(vars(opt), indent=2)) 66 | 67 | with torch.no_grad(): 68 | test_netG.eval() 69 | dic = {} 70 | print('generate mesh (inference) ...') 71 | for i in tqdm(range(0, len(dataset) // opt.num_views)): 72 | test_data = dataset[i] 73 | save_path = '%s/%s/inference_eval_%s_%d.obj' % ( 74 | opt.results_path, opt.name, test_data['name'], test_data['yid']) 75 | if dic.__contains__(test_data['name']): 76 | break 77 | dic[test_data['name']] = True 78 | image_tensor = test_data['image'].to(device=cuda).unsqueeze(0) 79 | mask_tensor = test_data['mask'].to(device=cuda).unsqueeze(0) 80 | res = netN.forward(image_tensor) 81 | res = res * mask_tensor 82 | # test_data['normal'] = torch.nn.Upsample(size=[1024, 1024], mode='bilinear')(res[0]) 83 | test_data['normal'] = res[0] 84 | print('saving to ' + save_path) 85 | gen_mesh_dmc(opt, test_netG, cuda, test_data, save_path, threshold=opt.mc_threshold, use_octree=True) 86 | 87 | 88 | if __name__ == '__main__': 89 | np.random.seed(int(time.time())) 90 | random.seed(int(time.time())) 91 | torch.manual_seed(int(time.time())) 92 | 93 | inference_nm(opt) 94 | -------------------------------------------------------------------------------- /apps/train_dmc.py: -------------------------------------------------------------------------------- 1 | import sys 2 | import os 3 | 4 | sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..'))) 5 | ROOT_PATH = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) 6 | 7 | import threading 8 | import time 9 | import json 10 | import numpy as np 11 | import cv2 12 | import random 13 | import torch 14 | from torch.utils.data import DataLoader 15 | from torch.nn import DataParallel 16 | from tqdm import tqdm 17 | import torch.nn.functional as F 18 | from torch.utils.tensorboard import SummaryWriter 19 | import matplotlib.pyplot as plt 20 | from multiprocessing import Process, Manager, Lock 21 | 22 | from lib.options import parse_config 23 | from lib.mesh_util import * 24 | from lib.sample_util import * 25 | from lib.train_util import * 26 | from lib.data import * 27 | from lib.model import * 28 | 29 | # get options 30 | opt = parse_config() 31 | 32 | def train(opt): 33 | np.random.seed(int(time.time())) 34 | random.seed(int(time.time())) 35 | torch.manual_seed(int(time.time())) 36 | # set cuda 37 | log = SummaryWriter(opt.log_path) 38 | total_iteration = 0 39 | cuda = torch.device('cuda:%s' % opt.gpu_ids[0]) 40 | netG = DMCNet(opt, projection_mode='perspective').to(cuda) 41 | netN = NormalNet().to(cuda) 42 | print('Using Network: ', netG.name, netN.name) 43 | gpu_ids = [int(i) for i in opt.gpu_ids.split(',')] 44 | netG = DataParallel(netG, device_ids=gpu_ids) 45 | netN = DataParallel(netN, device_ids=gpu_ids) 46 | 47 | optimizerG = torch.optim.Adam(netG.parameters(), lr=opt.learning_rate) 48 | lr = opt.learning_rate 49 | 50 | def set_train(): 51 | netG.train() 52 | 53 | if opt.load_netG_checkpoint_path is not None: 54 | print('loading for net G ...', opt.load_netG_checkpoint_path) 55 | netG.load_state_dict(torch.load(opt.load_netG_checkpoint_path, map_location=cuda), strict=False) 56 | 57 | if opt.load_netN_checkpoint_path is not None: 58 | print('loading for net N ...', opt.load_netN_checkpoint_path) 59 | netN.load_state_dict(torch.load(opt.load_netN_checkpoint_path, map_location=cuda), strict=False) 60 | 61 | print("loaded finished!") 62 | 63 | yaw_list = sorted(np.random.choice(range(360), 30)) 64 | print(yaw_list) 65 | 66 | train_dataset = DMCDataset(opt, cache_data=Manager().dict(), cache_data_lock=Lock(), phase='train', yaw_list=yaw_list) 67 | test_dataset = DMCDataset(opt, cache_data=Manager().dict(), cache_data_lock=Lock(), phase='test', yaw_list=yaw_list) 68 | 69 | projection_mode = train_dataset.projection_mode 70 | print('projection_mode:', projection_mode) 71 | # create data loader 72 | train_data_loader = DataLoader(train_dataset, 73 | batch_size=opt.batch_size, shuffle=not opt.serial_batches, 74 | num_workers=opt.num_threads, pin_memory=opt.pin_memory) 75 | print('train data size: ', len(train_data_loader)) 76 | 77 | # NOTE: batch size should be 1 and use all the points for evaluation 78 | test_data_loader = DataLoader(test_dataset, 79 | batch_size=1, shuffle=False, 80 | num_workers=opt.num_threads, pin_memory=opt.pin_memory) 81 | print('test data size: ', len(test_data_loader)) 82 | 83 | os.makedirs(opt.checkpoints_path, exist_ok=True) 84 | os.makedirs(opt.results_path, exist_ok=True) 85 | os.makedirs('%s/%s' % (opt.checkpoints_path, opt.name), exist_ok=True) 86 | os.makedirs('%s/%s' % (opt.results_path, opt.name), exist_ok=True) 87 | 88 | opt_log = os.path.join(opt.results_path, opt.name, 'opt.txt') 89 | with open(opt_log, 'w') as outfile: 90 | outfile.write(json.dumps(vars(opt), indent=2)) 91 | 92 | # training 93 | start_epoch = 0 94 | print("start training......") 95 | 96 | for epoch in range(start_epoch, opt.num_epoch): 97 | epoch_start_time = time.time() 98 | set_train() 99 | iter_data_time = time.time() 100 | np.random.seed(int(time.time())) 101 | random.seed(int(time.time())) 102 | torch.manual_seed(int(time.time())) 103 | train_bar = tqdm(enumerate(train_data_loader)) 104 | for train_idx, train_data in train_bar: 105 | total_iteration += 1 106 | iter_start_time = time.time() 107 | # retrieve the data 108 | for key in train_data: 109 | if torch.is_tensor(train_data[key]): 110 | train_data[key] = train_data[key].to(device=cuda) 111 | 112 | # predict normal 113 | with torch.no_grad(): 114 | net_normal = netN.forward(train_data['image']) 115 | net_normal = net_normal * train_data['mask'] 116 | 117 | train_data['normal'] = net_normal.detach() 118 | res, error = netG.forward(train_data) 119 | optimizerG.zero_grad() 120 | if len(gpu_ids) > 1: 121 | error = error.sum() 122 | error.backward() 123 | optimizerG.step() 124 | 125 | iter_net_time = time.time() 126 | eta = ((iter_net_time - epoch_start_time) / (train_idx + 1)) * len(train_data_loader) - ( 127 | iter_net_time - epoch_start_time) 128 | 129 | log.add_scalar('loss', error.item() / len(gpu_ids), total_iteration) 130 | if train_idx % opt.freq_plot == 0: 131 | descrip = 'Name: {0} | Epoch: {1} | {2}/{3} | Err: {4:.06f} | LR: {5:.06f} | Sigma: {6:.02f} | dataT: {7:.05f} | netT: {8:.05f} | ETA: {9:02d}:{10:02d}'.format( 132 | opt.name, epoch, train_idx, len(train_data_loader), error.item() / len(gpu_ids), lr, opt.sigma, 133 | iter_start_time - iter_data_time, 134 | iter_net_time - iter_start_time, int(eta // 60), 135 | int(eta - 60 * (eta // 60))) 136 | train_bar.set_description(descrip) 137 | 138 | if train_idx % opt.freq_save == 0: 139 | torch.save(netG.state_dict(), '%s/%s/netG_latest' % (opt.checkpoints_path, opt.name)) 140 | torch.save(netG.state_dict(), '%s/%s/netG_epoch_%d' % (opt.checkpoints_path, opt.name, epoch)) 141 | torch.save(optimizerG.state_dict(), '%s/%s/optim_latest' % (opt.checkpoints_path, opt.name)) 142 | torch.save(optimizerG.state_dict(), '%s/%s/optim_epoch_%d' % (opt.checkpoints_path, opt.name, epoch)) 143 | 144 | if train_idx % opt.freq_save_ply == 0: 145 | save_path = '%s/%s/pred.ply' % (opt.results_path, opt.name) 146 | r = res[0].cpu() 147 | points = train_data['samples'][0].transpose(0, 1).cpu() 148 | save_samples_truncted_prob(save_path, points.detach().numpy(), r.detach().numpy()) 149 | 150 | iter_data_time = time.time() 151 | 152 | # update learning rate 153 | lr = adjust_learning_rate(optimizerG, epoch, lr, opt.schedule, opt.gamma) 154 | train_dataset.clear_cache() 155 | 156 | yaw_list = sorted(np.random.choice(range(360), 30)) 157 | train_dataset.yaw_list = yaw_list 158 | log.close() 159 | 160 | 161 | if __name__ == '__main__': 162 | train(opt) -------------------------------------------------------------------------------- /apps/train_normal_net.py: -------------------------------------------------------------------------------- 1 | import sys 2 | import os 3 | 4 | sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..'))) 5 | ROOT_PATH = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) 6 | 7 | import threading 8 | import time 9 | import json 10 | import numpy as np 11 | import cv2 12 | import random 13 | import torch 14 | from torch.utils.data import DataLoader 15 | from torch.nn import DataParallel 16 | from tqdm import tqdm 17 | from torch.utils.tensorboard import SummaryWriter 18 | import torch.nn.functional as F 19 | import matplotlib.pyplot as plt 20 | 21 | from lib.mesh_util import * 22 | from lib.sample_util import * 23 | from lib.train_util import * 24 | from lib.data import * 25 | from lib.model import * 26 | from lib.geometry import index 27 | from lib.loss_util import VGGPerceptualLoss 28 | from lib.options import parse_config 29 | 30 | # get options 31 | opt = parse_config() 32 | log = SummaryWriter(opt.log_path) 33 | 34 | def train(opt): 35 | 36 | gpu_ids = [int(i) for i in opt.gpu_ids.split(',')] 37 | cuda = torch.device("cuda:%d" % gpu_ids[0]) 38 | netN = NormalNet() 39 | dataset = NormalDataset(opt) 40 | train_data_loader = DataLoader(dataset, batch_size=opt.batch_size, shuffle=not opt.serial_batches, 41 | num_workers=opt.num_threads, pin_memory=opt.pin_memory) 42 | 43 | netN.to(cuda) 44 | netN = DataParallel(netN, device_ids=gpu_ids) 45 | 46 | if opt.load_netN_checkpoint_path is not None: 47 | netN.load_state_dict(torch.load(opt.load_netN_checkpoint_path), strict=False) 48 | 49 | lr = opt.learning_rate 50 | optimizerN = torch.optim.Adam(netN.parameters(), lr=lr) 51 | 52 | os.makedirs(opt.checkpoints_path, exist_ok=True) 53 | os.makedirs(opt.results_path, exist_ok=True) 54 | os.makedirs('%s/%s' % (opt.checkpoints_path, opt.name), exist_ok=True) 55 | os.makedirs('%s/%s' % (opt.results_path, opt.name), exist_ok=True) 56 | 57 | EPOCH = 100 58 | total_iteration = 0 59 | 60 | perceptual_loss = VGGPerceptualLoss().to(cuda) 61 | for epoch in range(EPOCH): 62 | epoch_start_time = time.time() 63 | iter_data_time = time.time() 64 | train_bar = tqdm(enumerate(train_data_loader)) 65 | for train_idx, train_data in train_bar: 66 | 67 | total_iteration += 1 68 | iter_start_time = time.time() 69 | 70 | # retrieve the data 71 | image_tensor = train_data['img'].to(device=cuda) 72 | normal_tensor = train_data['normal'].to(device=cuda) 73 | mask_tensor = train_data['mask'].to(device=cuda) 74 | 75 | res = netN.forward(image_tensor) 76 | res = res * mask_tensor 77 | error = F.l1_loss(normal_tensor, res) 78 | perceptual_error = perceptual_loss(normal_tensor.squeeze(1), res.squeeze(1)) 79 | 80 | error = 5*error + perceptual_error 81 | 82 | optimizerN.zero_grad() 83 | error.backward() 84 | optimizerN.step() 85 | 86 | iter_net_time = time.time() 87 | eta = ((iter_net_time - epoch_start_time) / (train_idx + 1)) * len(train_data_loader) - ( 88 | iter_net_time - epoch_start_time) 89 | 90 | log.add_scalar('loss', error.item(), total_iteration) 91 | if train_idx % opt.freq_plot == 0: 92 | descrip = 'Name: {0} | Epoch: {1} | {2}/{3} | Err: {4:.06f} | LR: {5:.06f} | Sigma: {6:.02f} | dataT: {7:.05f} | netT: {8:.05f} | ETA: {9:02d}:{10:02d}'.format( 93 | opt.name, epoch, train_idx, len(train_data_loader), error.item(), lr, opt.sigma, 94 | iter_start_time - iter_data_time, 95 | iter_net_time - iter_start_time, int(eta // 60), 96 | int(eta - 60 * (eta // 60))) 97 | train_bar.set_description(descrip) 98 | 99 | if train_idx % opt.freq_save == 0 and train_idx != 0: 100 | torch.save(netN.state_dict(), '%s/%s/netN_latest' % (opt.checkpoints_path, opt.name)) 101 | torch.save(netN.state_dict(), '%s/%s/netN_epoch_%d' % (opt.checkpoints_path, opt.name, epoch)) 102 | 103 | if train_idx % opt.freq_normal_show == 0: 104 | show_img = (image_tensor[0][0].cpu().detach().permute(1, 2, 0) + 0.5).numpy() 105 | net_normal = (res[0][0].cpu().detach().permute(1, 2, 0) + 0.5).numpy() 106 | gt_normal = (normal_tensor[0][0].cpu().detach().permute(1, 2, 0) + 0.5).numpy() 107 | plt.subplot(131) 108 | plt.imshow(show_img) 109 | plt.subplot(132) 110 | plt.imshow(net_normal) 111 | plt.subplot(133) 112 | plt.imshow(gt_normal) 113 | plt.savefig('%s/%s/epoch%03d_%05d.jpg' % (opt.results_path, opt.name, epoch, train_idx)) 114 | plt.close('all') 115 | 116 | iter_data_time = time.time() 117 | 118 | if __name__ == "__main__": 119 | train(opt) -------------------------------------------------------------------------------- /assets/teaser.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/DSaurus/DeepMultiCap/4e93adfbdd60520d4edaa8bd8e461fe41d01c489/assets/teaser.jpg -------------------------------------------------------------------------------- /assets/tex/skirt.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/DSaurus/DeepMultiCap/4e93adfbdd60520d4edaa8bd8e461fe41d01c489/assets/tex/skirt.jpg -------------------------------------------------------------------------------- /configs/dmc_demo.yaml: -------------------------------------------------------------------------------- 1 | name: dmc_demo 2 | log_path: logs/dmc_demo 3 | gpu_ids: 0 4 | num_threads: 1 5 | batch_size: 1 6 | infer: True 7 | infer_reverse: True 8 | 9 | # data 10 | num_views: 4 11 | yaw_list: [0, 90, 180, 270] 12 | flip_x: False 13 | 14 | # network 15 | fine_part: True 16 | mlp_dim: [323, 1024, 512, 256, 128, 1] 17 | fine_mlp_dim: [355, 512, 256, 128, 1] 18 | 19 | # path 20 | load_netG_checkpoint_path: checkpoints/demo/netG_latest 21 | load_netN_checkpoint_path: checkpoints/demo/netN_latest 22 | 23 | # smpl 24 | smpl_faces: lib/data/smplx_multi.obj 25 | 26 | # marching cube 27 | resolution: 512 28 | mc_threshold: 0.5 -------------------------------------------------------------------------------- /configs/dmc_train.yaml: -------------------------------------------------------------------------------- 1 | name: dmc_demo 2 | log_path: logs/dmc_demo 3 | gpu_ids: 0 4 | num_threads: 1 5 | batch_size: 1 6 | infer: False 7 | 8 | # data 9 | num_views: 4 10 | flip_x: False 11 | 12 | # network 13 | mlp_dim: [323, 1024, 512, 256, 128, 1] 14 | fine_mlp_dim: [355, 512, 256, 128, 1] 15 | 16 | # path 17 | load_netG_checkpoint_path: checkpoints/demo/netG_latest 18 | load_netN_checkpoint_path: checkpoints/demo/netN_latest 19 | 20 | # smpl 21 | smpl_faces: lib/data/smplx_fine.obj 22 | 23 | # marching cube 24 | resolution: 256 25 | mc_threshold: 0.6 -------------------------------------------------------------------------------- /configs/normal_train.yaml: -------------------------------------------------------------------------------- 1 | name: dmc_normal 2 | log_path: logs/dmc_normal 3 | gpu_ids: 0 4 | num_threads: 1 5 | batch_size: 4 6 | 7 | # path 8 | load_netN_checkpoint_path: checkpoints/demo/netN_latest -------------------------------------------------------------------------------- /dataset/multihuman/img/20_0/0.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/DSaurus/DeepMultiCap/4e93adfbdd60520d4edaa8bd8e461fe41d01c489/dataset/multihuman/img/20_0/0.jpg -------------------------------------------------------------------------------- /dataset/multihuman/img/20_0/180.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/DSaurus/DeepMultiCap/4e93adfbdd60520d4edaa8bd8e461fe41d01c489/dataset/multihuman/img/20_0/180.jpg -------------------------------------------------------------------------------- /dataset/multihuman/img/20_0/270.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/DSaurus/DeepMultiCap/4e93adfbdd60520d4edaa8bd8e461fe41d01c489/dataset/multihuman/img/20_0/270.jpg -------------------------------------------------------------------------------- /dataset/multihuman/img/20_0/90.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/DSaurus/DeepMultiCap/4e93adfbdd60520d4edaa8bd8e461fe41d01c489/dataset/multihuman/img/20_0/90.jpg -------------------------------------------------------------------------------- /dataset/multihuman/img/20_1/0.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/DSaurus/DeepMultiCap/4e93adfbdd60520d4edaa8bd8e461fe41d01c489/dataset/multihuman/img/20_1/0.jpg -------------------------------------------------------------------------------- /dataset/multihuman/img/20_1/180.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/DSaurus/DeepMultiCap/4e93adfbdd60520d4edaa8bd8e461fe41d01c489/dataset/multihuman/img/20_1/180.jpg -------------------------------------------------------------------------------- /dataset/multihuman/img/20_1/270.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/DSaurus/DeepMultiCap/4e93adfbdd60520d4edaa8bd8e461fe41d01c489/dataset/multihuman/img/20_1/270.jpg -------------------------------------------------------------------------------- /dataset/multihuman/img/20_1/90.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/DSaurus/DeepMultiCap/4e93adfbdd60520d4edaa8bd8e461fe41d01c489/dataset/multihuman/img/20_1/90.jpg -------------------------------------------------------------------------------- /dataset/multihuman/mask/20_0/0.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/DSaurus/DeepMultiCap/4e93adfbdd60520d4edaa8bd8e461fe41d01c489/dataset/multihuman/mask/20_0/0.jpg -------------------------------------------------------------------------------- /dataset/multihuman/mask/20_0/180.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/DSaurus/DeepMultiCap/4e93adfbdd60520d4edaa8bd8e461fe41d01c489/dataset/multihuman/mask/20_0/180.jpg -------------------------------------------------------------------------------- /dataset/multihuman/mask/20_0/270.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/DSaurus/DeepMultiCap/4e93adfbdd60520d4edaa8bd8e461fe41d01c489/dataset/multihuman/mask/20_0/270.jpg -------------------------------------------------------------------------------- /dataset/multihuman/mask/20_0/90.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/DSaurus/DeepMultiCap/4e93adfbdd60520d4edaa8bd8e461fe41d01c489/dataset/multihuman/mask/20_0/90.jpg -------------------------------------------------------------------------------- /dataset/multihuman/mask/20_1/0.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/DSaurus/DeepMultiCap/4e93adfbdd60520d4edaa8bd8e461fe41d01c489/dataset/multihuman/mask/20_1/0.jpg -------------------------------------------------------------------------------- /dataset/multihuman/mask/20_1/180.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/DSaurus/DeepMultiCap/4e93adfbdd60520d4edaa8bd8e461fe41d01c489/dataset/multihuman/mask/20_1/180.jpg -------------------------------------------------------------------------------- /dataset/multihuman/mask/20_1/270.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/DSaurus/DeepMultiCap/4e93adfbdd60520d4edaa8bd8e461fe41d01c489/dataset/multihuman/mask/20_1/270.jpg -------------------------------------------------------------------------------- /dataset/multihuman/mask/20_1/90.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/DSaurus/DeepMultiCap/4e93adfbdd60520d4edaa8bd8e461fe41d01c489/dataset/multihuman/mask/20_1/90.jpg -------------------------------------------------------------------------------- /dataset/multihuman/parameter/20_0/0_extrinsic.npy: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/DSaurus/DeepMultiCap/4e93adfbdd60520d4edaa8bd8e461fe41d01c489/dataset/multihuman/parameter/20_0/0_extrinsic.npy -------------------------------------------------------------------------------- /dataset/multihuman/parameter/20_0/0_intrinsic.npy: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/DSaurus/DeepMultiCap/4e93adfbdd60520d4edaa8bd8e461fe41d01c489/dataset/multihuman/parameter/20_0/0_intrinsic.npy -------------------------------------------------------------------------------- /dataset/multihuman/parameter/20_0/180_extrinsic.npy: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/DSaurus/DeepMultiCap/4e93adfbdd60520d4edaa8bd8e461fe41d01c489/dataset/multihuman/parameter/20_0/180_extrinsic.npy -------------------------------------------------------------------------------- /dataset/multihuman/parameter/20_0/180_intrinsic.npy: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/DSaurus/DeepMultiCap/4e93adfbdd60520d4edaa8bd8e461fe41d01c489/dataset/multihuman/parameter/20_0/180_intrinsic.npy -------------------------------------------------------------------------------- /dataset/multihuman/parameter/20_0/270_extrinsic.npy: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/DSaurus/DeepMultiCap/4e93adfbdd60520d4edaa8bd8e461fe41d01c489/dataset/multihuman/parameter/20_0/270_extrinsic.npy -------------------------------------------------------------------------------- /dataset/multihuman/parameter/20_0/270_intrinsic.npy: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/DSaurus/DeepMultiCap/4e93adfbdd60520d4edaa8bd8e461fe41d01c489/dataset/multihuman/parameter/20_0/270_intrinsic.npy -------------------------------------------------------------------------------- /dataset/multihuman/parameter/20_0/90_extrinsic.npy: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/DSaurus/DeepMultiCap/4e93adfbdd60520d4edaa8bd8e461fe41d01c489/dataset/multihuman/parameter/20_0/90_extrinsic.npy -------------------------------------------------------------------------------- /dataset/multihuman/parameter/20_0/90_intrinsic.npy: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/DSaurus/DeepMultiCap/4e93adfbdd60520d4edaa8bd8e461fe41d01c489/dataset/multihuman/parameter/20_0/90_intrinsic.npy -------------------------------------------------------------------------------- /dataset/multihuman/parameter/20_1/0_extrinsic.npy: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/DSaurus/DeepMultiCap/4e93adfbdd60520d4edaa8bd8e461fe41d01c489/dataset/multihuman/parameter/20_1/0_extrinsic.npy -------------------------------------------------------------------------------- /dataset/multihuman/parameter/20_1/0_intrinsic.npy: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/DSaurus/DeepMultiCap/4e93adfbdd60520d4edaa8bd8e461fe41d01c489/dataset/multihuman/parameter/20_1/0_intrinsic.npy -------------------------------------------------------------------------------- /dataset/multihuman/parameter/20_1/180_extrinsic.npy: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/DSaurus/DeepMultiCap/4e93adfbdd60520d4edaa8bd8e461fe41d01c489/dataset/multihuman/parameter/20_1/180_extrinsic.npy -------------------------------------------------------------------------------- /dataset/multihuman/parameter/20_1/180_intrinsic.npy: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/DSaurus/DeepMultiCap/4e93adfbdd60520d4edaa8bd8e461fe41d01c489/dataset/multihuman/parameter/20_1/180_intrinsic.npy -------------------------------------------------------------------------------- /dataset/multihuman/parameter/20_1/270_extrinsic.npy: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/DSaurus/DeepMultiCap/4e93adfbdd60520d4edaa8bd8e461fe41d01c489/dataset/multihuman/parameter/20_1/270_extrinsic.npy -------------------------------------------------------------------------------- /dataset/multihuman/parameter/20_1/270_intrinsic.npy: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/DSaurus/DeepMultiCap/4e93adfbdd60520d4edaa8bd8e461fe41d01c489/dataset/multihuman/parameter/20_1/270_intrinsic.npy -------------------------------------------------------------------------------- /dataset/multihuman/parameter/20_1/90_extrinsic.npy: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/DSaurus/DeepMultiCap/4e93adfbdd60520d4edaa8bd8e461fe41d01c489/dataset/multihuman/parameter/20_1/90_extrinsic.npy -------------------------------------------------------------------------------- /dataset/multihuman/parameter/20_1/90_intrinsic.npy: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/DSaurus/DeepMultiCap/4e93adfbdd60520d4edaa8bd8e461fe41d01c489/dataset/multihuman/parameter/20_1/90_intrinsic.npy -------------------------------------------------------------------------------- /demo.sh: -------------------------------------------------------------------------------- 1 | python apps/eval_3d.py --config configs/dmc_demo.yaml --dataroot dataset/multihuman -------------------------------------------------------------------------------- /lib/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/DSaurus/DeepMultiCap/4e93adfbdd60520d4edaa8bd8e461fe41d01c489/lib/__init__.py -------------------------------------------------------------------------------- /lib/data/NormalDataset.py: -------------------------------------------------------------------------------- 1 | from torch.utils.data import Dataset 2 | from multiprocessing import Process, Manager, Lock 3 | import numpy as np 4 | import os 5 | import random 6 | import torchvision.transforms as transforms 7 | from PIL import Image, ImageOps 8 | import cv2 9 | import torch 10 | from PIL.ImageFilter import GaussianBlur 11 | import trimesh 12 | import logging 13 | from tqdm import tqdm 14 | import matplotlib.pyplot as plt 15 | 16 | class NormalDataset(Dataset): 17 | def __init__(self, opt, phase = 'train'): 18 | self.opt = opt 19 | 20 | self.root = opt.dataroot 21 | self.NORMAL = os.path.join(self.root, 'normal') 22 | self.RENDER = os.path.join(self.root, 'img') 23 | self.MASK = os.path.join(self.root, 'mask') 24 | self.phase = phase 25 | self.is_train = self.phase == 'train' 26 | self.load_size = 512 27 | 28 | self.subjects = self.get_subjects() 29 | self.yaw_list = range(0, 360, 6) 30 | 31 | # PIL to tensor 32 | self.to_tensor = transforms.Compose([ 33 | transforms.Resize(512), 34 | transforms.ToTensor(), 35 | transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)) 36 | ]) 37 | 38 | # augmentation 39 | self.aug_trans = transforms.Compose([ 40 | transforms.ColorJitter(brightness=opt.aug_bri, contrast=opt.aug_con, saturation=opt.aug_sat, 41 | hue=opt.aug_hue) 42 | ]) 43 | 44 | 45 | def __len__(self): 46 | return len(self.subjects) * len(self.yaw_list) 47 | 48 | def get_subjects(self): 49 | all_subjects = os.listdir(self.RENDER) 50 | if os.path.isfile(os.path.join(self.root, 'val.txt')): 51 | var_subjects = np.loadtxt(os.path.join(self.root, 'val.txt'), dtype=str) 52 | else: 53 | var_subjects = [] 54 | if len(var_subjects) == 0: 55 | return all_subjects 56 | 57 | if self.is_train: 58 | return sorted(list(set(all_subjects) - set(var_subjects))) 59 | else: 60 | return sorted(list(var_subjects)) 61 | 62 | def get_render(self, subject, num_views, yid=0, random_sample=False): 63 | ''' 64 | Return the render data 65 | :param subject: subject name 66 | :param num_views: how many views to return 67 | :param view_id: the first view_id. If None, select a random one. 68 | :return: 69 | 'img': [num_views, C, W, H] images 70 | 'calib': [num_views, 4, 4] calibration matrix 71 | 'extrinsic': [num_views, 4, 4] extrinsic matrix 72 | 'mask': [num_views, 1, W, H] masks 73 | ''' 74 | 75 | # The ids are an even distribution of num_views around view_id 76 | view_ids = [self.yaw_list[(yid + len(self.yaw_list) // num_views * offset) % len(self.yaw_list)] 77 | for offset in range(1)] 78 | if random_sample and self.is_train: 79 | view_ids = np.random.choice(self.yaw_list, num_views, replace=False) 80 | 81 | render_list = [] 82 | normal_list = [] 83 | mask_list = [] 84 | 85 | for vid in view_ids: 86 | render_path = os.path.join(self.RENDER, subject, '{}.jpg'.format(vid)) 87 | if not os.path.exists(render_path): 88 | render_path = os.path.join(self.RENDER, subject, '{}.png'.format(vid)) 89 | normal_path = os.path.join(self.NORMAL, subject, '{}.png'.format(vid)) 90 | if not os.path.exists(normal_path): 91 | normal_path = os.path.join(self.NORMAL, subject, '{}.jpg'.format(vid)) 92 | mask_path = os.path.join(self.MASK, subject, '{}.jpg'.format(vid)) 93 | if not os.path.exists(mask_path): 94 | mask_path = os.path.join(self.MASK, subject, '{}.png'.format(vid)) 95 | 96 | mask = Image.open(mask_path) 97 | render = Image.open(render_path).convert('RGB') 98 | normal = Image.open(normal_path).convert('RGB') 99 | 100 | if self.is_train: 101 | # Pad images 102 | pad_size = int(0.1 * self.load_size) 103 | render = ImageOps.expand(render, pad_size, fill=0) 104 | normal = ImageOps.expand(normal, pad_size, fill=0) 105 | mask = ImageOps.expand(mask, pad_size, fill=0) 106 | 107 | w, h = render.size 108 | th, tw = self.load_size, self.load_size 109 | 110 | # random scale 111 | if self.opt.random_scale: 112 | rand_scale = random.uniform(0.9, 1.1) 113 | w = int(rand_scale * w) 114 | h = int(rand_scale * h) 115 | render = render.resize((w, h), Image.BILINEAR) 116 | normal = normal.resize((w, h), Image.BILINEAR) 117 | mask = mask.resize((w, h), Image.BILINEAR) 118 | 119 | # random translate in the pixel space 120 | if self.opt.random_trans: 121 | dx = random.randint(-int(round((w - tw) / 10)), 122 | int(round((w - tw) / 10))) 123 | dy = random.randint(-int(round((h - th) / 10)), 124 | int(round((h - th) / 10))) 125 | else: 126 | dx = 0 127 | dy = 0 128 | 129 | x1 = int(round((w - tw) / 2.)) + dx 130 | y1 = int(round((h - th) / 2.)) + dy 131 | 132 | render = render.crop((x1, y1, x1 + tw, y1 + th)) 133 | normal = normal.crop((x1, y1, x1 + tw, y1 + th)) 134 | mask = mask.crop((x1, y1, x1 + tw, y1 + th)) 135 | 136 | render = self.aug_trans(render) 137 | 138 | # random blur 139 | if self.opt.aug_blur > 0.00001: 140 | blur = GaussianBlur(np.random.uniform(0, self.opt.aug_blur)) 141 | render = render.filter(blur) 142 | 143 | mask = torch.FloatTensor(np.array(mask))[:, :, 0] / 255 144 | render = self.to_tensor(render) * mask.reshape(1, 512, 512) 145 | normal = self.to_tensor(normal) * mask.reshape(1, 512, 512) 146 | 147 | render_list.append(render) 148 | normal_list.append(normal) 149 | mask_list.append(mask.reshape(1, 512, 512)) 150 | 151 | return { 152 | 'img': torch.stack(render_list, dim=0), 153 | 'normal': torch.stack(normal_list, dim=0), 154 | 'mask': torch.stack(mask_list, dim=0) 155 | } 156 | 157 | def get_item(self, index): 158 | sid = index % len(self.subjects) 159 | tmp = index // len(self.subjects) 160 | yid = tmp % len(self.yaw_list) 161 | 162 | subject = self.subjects[sid] 163 | res = { 164 | 'name': subject, 165 | 'sid': sid, 166 | 'yid': yid, 167 | } 168 | render_data = self.get_render(subject, num_views=1, yid=yid, 169 | random_sample=self.opt.random_multiview) 170 | res.update(render_data) 171 | return res 172 | 173 | def __getitem__(self, index): 174 | return self.get_item(index) 175 | -------------------------------------------------------------------------------- /lib/data/__init__.py: -------------------------------------------------------------------------------- 1 | from .DMCDataset import DMCDataset 2 | from .NormalDataset import NormalDataset -------------------------------------------------------------------------------- /lib/data/smplx_correction.obj: -------------------------------------------------------------------------------- 1 | f 3490 3549 3840 2 | f 3852 3963 3964 3 | f 3970 3971 3547 4 | f 3967 3966 3548 5 | f 3555 3554 5402 6 | f 3492 3491 5511 7 | f 3969 5402 3978 8 | f 3969 3968 3964 9 | f 3552 3553 3554 10 | f 3489 3490 3491 11 | f 4401 3967 3547 12 | f 4292 4293 3970 13 | f 3491 3840 5514 14 | f 3548 3550 3549 15 | f 3490 3489 3547 16 | f 5511 5514 5726 17 | f 5512 5511 5723 18 | f 5940 5941 4401 19 | f 3555 3964 3963 20 | f 4401 5941 5944 21 | f 3970 4293 5940 22 | f 5944 5949 3968 23 | f 5512 5947 5943 24 | f 5724 5616 5947 25 | f 5949 3853 3964 26 | f 3492 5943 4292 27 | f 3966 3978 3550 28 | f 3966 3967 3968 29 | f 6250 6253 6252 30 | f 6251 6309 6308 31 | f 6313 6316 6315 32 | f 6251 6252 6595 33 | f 3852 3853 6712 34 | f 6714 6717 6716 35 | f 6718 6250 6308 36 | f 6715 6308 6309 37 | f 7137 6719 6308 38 | f 4292 6250 6718 39 | f 6316 6717 8136 40 | f 6316 6313 6711 41 | f 6253 8234 8233 42 | f 6252 8233 8236 43 | f 8234 8418 8417 44 | f 8233 8417 8420 45 | f 6717 6316 6712 46 | f 6717 6714 6726 47 | f 6309 6251 6310 48 | f 6714 6309 6311 49 | f 5940 6719 7137 50 | f 6253 6250 4292 51 | f 7137 6715 5944 52 | f 6718 6719 5940 53 | f 8234 6253 5943 54 | f 8418 8234 5947 55 | f 5949 6716 6712 56 | f 5944 6715 6716 57 | f 3490 3840 3491 58 | f 3852 3964 3853 59 | f 3970 3547 3489 60 | f 3967 3548 3547 61 | f 3555 5402 3969 62 | f 3492 5511 5512 63 | f 3969 3978 3966 64 | f 3969 3964 3555 65 | f 3552 3554 3555 66 | f 3489 3491 3492 67 | f 4401 3547 3971 68 | f 4292 3970 3489 69 | f 3491 5514 5511 70 | f 3548 3549 3490 71 | f 3490 3547 3548 72 | f 5511 5726 5723 73 | f 5512 5723 5724 74 | f 5940 4401 3971 75 | f 3555 3963 3552 76 | f 4401 5944 3967 77 | f 3970 5940 3971 78 | f 5944 3968 3967 79 | f 5512 5943 3492 80 | f 5724 5947 5512 81 | f 5949 3964 3968 82 | f 3492 4292 3489 83 | f 3966 3550 3548 84 | f 3966 3968 3969 85 | f 6250 6252 6251 86 | f 6251 6308 6250 87 | f 6313 6315 6314 88 | f 6251 6595 6310 89 | f 3852 6712 6711 90 | f 6714 6716 6715 91 | f 6718 6308 6719 92 | f 6715 6309 6714 93 | f 7137 6308 6715 94 | f 4292 6718 4293 95 | f 6316 8136 6315 96 | f 6316 6711 6712 97 | f 6253 8233 6252 98 | f 6252 8236 6595 99 | f 8234 8417 8233 100 | f 8233 8420 8236 101 | f 6717 6712 6716 102 | f 6717 6726 8136 103 | f 6309 6310 6311 104 | f 6714 6311 6726 105 | f 5940 7137 5941 106 | f 6253 4292 5943 107 | f 7137 5944 5941 108 | f 6718 5940 4293 109 | f 8234 5943 5947 110 | f 8418 5947 5616 111 | f 5949 6712 3853 112 | f 5944 6716 5949 -------------------------------------------------------------------------------- /lib/data/smplx_fine.obj: -------------------------------------------------------------------------------- 1 | 2 | f 3490 10863 10862 3 | f 3840 10861 10863 4 | f 3549 10862 10861 5 | f 10861 10862 10863 6 | f 3852 10910 10909 7 | f 3964 10908 10910 8 | f 3963 10909 10908 9 | f 10908 10909 10910 10 | f 3970 10952 10951 11 | f 3547 10950 10952 12 | f 3971 10951 10950 13 | f 10950 10951 10952 14 | f 3967 11054 11053 15 | f 3548 11052 11054 16 | f 3966 11053 11052 17 | f 11052 11053 11054 18 | f 3555 11604 11603 19 | f 5402 11602 11604 20 | f 3554 11603 11602 21 | f 11602 11603 11604 22 | f 5427 11777 11776 23 | f 3552 11775 11777 24 | f 3556 11776 11775 25 | f 11775 11776 11777 26 | f 3492 11899 11898 27 | f 5511 11897 11899 28 | f 3491 11898 11897 29 | f 11897 11898 11899 30 | f 3969 12576 12575 31 | f 3978 12574 12576 32 | f 5402 12575 12574 33 | f 12574 12575 12576 34 | f 3969 12950 12949 35 | f 3964 12948 12950 36 | f 3968 12949 12948 37 | f 12948 12949 12950 38 | f 3552 13050 13049 39 | f 3554 13048 13050 40 | f 3553 13049 13048 41 | f 13048 13049 13050 42 | f 3553 13057 13049 43 | f 3556 11775 13057 44 | f 3552 13049 11775 45 | f 11775 13049 13057 46 | f 3489 13097 13096 47 | f 3491 13095 13097 48 | f 3490 13096 13095 49 | f 13095 13096 13097 50 | f 4401 13180 13179 51 | f 3547 13178 13180 52 | f 3967 13179 13178 53 | f 13178 13179 13180 54 | f 4292 13194 13193 55 | f 3970 13192 13194 56 | f 4293 13193 13192 57 | f 13192 13193 13194 58 | f 3491 13825 13824 59 | f 5514 13823 13825 60 | f 3840 13824 13823 61 | f 13823 13824 13825 62 | f 3548 14016 14015 63 | f 3549 10615 14016 64 | f 3550 14015 10615 65 | f 10615 14015 14016 66 | f 3490 14210 13096 67 | f 3547 14209 14210 68 | f 3489 13096 14209 69 | f 14209 13096 14210 70 | f 5511 14355 14354 71 | f 5726 13970 14355 72 | f 5514 14354 13970 73 | f 13970 14354 14355 74 | f 5512 14372 14371 75 | f 5723 14370 14372 76 | f 5511 14371 14370 77 | f 14370 14371 14372 78 | f 5940 14557 14556 79 | f 4401 14555 14557 80 | f 5941 14556 14555 81 | f 14555 14556 14557 82 | f 3555 14589 14588 83 | f 3963 10908 14589 84 | f 3964 14588 10908 85 | f 10908 14588 14589 86 | f 4401 14893 14555 87 | f 5944 14892 14893 88 | f 5941 14555 14892 89 | f 14892 14555 14893 90 | f 3970 14898 13192 91 | f 5940 14897 14898 92 | f 4293 13192 14897 93 | f 14897 13192 14898 94 | f 5944 14911 14910 95 | f 3968 14909 14911 96 | f 5949 14910 14909 97 | f 14909 14910 14911 98 | f 3963 14913 10909 99 | f 5951 14912 14913 100 | f 3852 10909 14912 101 | f 14912 10909 14913 102 | f 5512 14922 14921 103 | f 5943 14920 14922 104 | f 5947 14921 14920 105 | f 14920 14921 14922 106 | f 5724 14928 14927 107 | f 5947 14926 14928 108 | f 5616 14927 14926 109 | f 14926 14927 14928 110 | f 5949 14931 14930 111 | f 3964 14929 14931 112 | f 3853 14930 14929 113 | f 14929 14930 14931 114 | f 3492 14953 14952 115 | f 4292 14951 14953 116 | f 5943 14952 14951 117 | f 14951 14952 14953 118 | f 3966 15095 15094 119 | f 3550 14196 15095 120 | f 3978 15094 14196 121 | f 14196 15094 15095 122 | f 3966 15412 11053 123 | f 3968 15411 15412 124 | f 3967 11053 15411 125 | f 15411 11053 15412 126 | f 6250 15708 15707 127 | f 6252 15706 15708 128 | f 6253 15707 15706 129 | f 15706 15707 15708 130 | f 6251 15761 15760 131 | f 6308 15759 15761 132 | f 6309 15760 15759 133 | f 15759 15760 15761 134 | f 6313 15767 15766 135 | f 6316 15766 15765 136 | f 15765 15766 15767 137 | f 6317 15768 15770 138 | f 6251 16104 16103 139 | f 6252 16103 16102 140 | f 16102 16103 16104 141 | f 3852 16248 16247 142 | f 6712 16246 16248 143 | f 3853 16247 16246 144 | f 16246 16247 16248 145 | f 6714 16254 16253 146 | f 6716 16252 16254 147 | f 6717 16253 16252 148 | f 16252 16253 16254 149 | f 6718 16257 16256 150 | f 6308 16255 16257 151 | f 6250 16256 16255 152 | f 16255 16256 16257 153 | f 6715 16614 16613 154 | f 6309 15759 16614 155 | f 6308 16613 15759 156 | f 15759 16613 16614 157 | f 7137 17040 17039 158 | f 6308 17038 17040 159 | f 6719 17039 17038 160 | f 17038 17039 17040 161 | f 4292 17042 17041 162 | f 6718 16256 17042 163 | f 6250 17041 16256 164 | f 16256 17041 17042 165 | f 6316 17612 17611 166 | f 6717 17611 17610 167 | f 17610 17611 17612 168 | f 6316 17624 15766 169 | f 6711 17623 17624 170 | f 6313 15766 17623 171 | f 17623 15766 17624 172 | f 8161 17655 17654 173 | f 6313 17623 17655 174 | f 6711 17654 17623 175 | f 17623 17654 17655 176 | f 6253 18035 18034 177 | f 8233 18033 18035 178 | f 8234 18034 18033 179 | f 18033 18034 18035 180 | f 6252 18044 18043 181 | f 8233 18043 18042 182 | f 18042 18043 18044 183 | f 8234 18842 18841 184 | f 8417 18840 18842 185 | f 8418 18841 18840 186 | f 18840 18841 18842 187 | f 8233 18850 18849 188 | f 8417 18849 18848 189 | f 18848 18849 18850 190 | f 6717 18866 17611 191 | f 6712 18865 18866 192 | f 6316 17611 18865 193 | f 18865 17611 18866 194 | f 6717 18872 16253 195 | f 6714 16253 18871 196 | f 18871 16253 18872 197 | f 6309 20153 15760 198 | f 6251 15760 20152 199 | f 20152 15760 20153 200 | f 6714 20172 20171 201 | f 6309 20171 20170 202 | f 20170 20171 20172 203 | f 5940 20449 20448 204 | f 7137 17039 20449 205 | f 6719 20448 17039 206 | f 17039 20448 20449 207 | f 6253 20455 15707 208 | f 4292 17041 20455 209 | f 6250 15707 17041 210 | f 17041 15707 20455 211 | f 7137 20458 20457 212 | f 5944 20456 20458 213 | f 6715 20457 20456 214 | f 20456 20457 20458 215 | f 6718 20468 20467 216 | f 5940 20448 20468 217 | f 6719 20467 20448 218 | f 20448 20467 20468 219 | f 8234 20485 18034 220 | f 5943 20484 20485 221 | f 6253 18034 20484 222 | f 20484 18034 20485 223 | f 8418 20493 18841 224 | f 5947 20492 20493 225 | f 8234 18841 20492 226 | f 20492 18841 20493 227 | f 5949 20498 20497 228 | f 6712 20496 20498 229 | f 6716 20497 20496 230 | f 20496 20497 20498 231 | f 5944 20500 20456 232 | f 6716 20499 20500 233 | f 6715 20456 20499 234 | f 20499 20456 20500 235 | f 6711 20510 17654 236 | f 5951 20509 20510 237 | f 8161 17654 20509 238 | f 20509 17654 20510 239 | f 3490 13095 10863 240 | f 3491 13824 13095 241 | f 3840 10863 13824 242 | f 13824 10863 13095 243 | f 3852 16247 10910 244 | f 3853 14929 16247 245 | f 3964 10910 14929 246 | f 14929 10910 16247 247 | f 3970 20673 10952 248 | f 3489 14209 20673 249 | f 3547 10952 14209 250 | f 14209 10952 20673 251 | f 3967 13178 11054 252 | f 3547 20693 13178 253 | f 3548 11054 20693 254 | f 20693 11054 13178 255 | f 3555 20827 11604 256 | f 3969 12575 20827 257 | f 5402 11604 12575 258 | f 12575 11604 20827 259 | f 5427 20873 11777 260 | f 3963 20872 20873 261 | f 3552 11777 20872 262 | f 20872 11777 20873 263 | f 3492 20898 11899 264 | f 5512 14371 20898 265 | f 5511 11899 14371 266 | f 14371 11899 20898 267 | f 3969 21032 12576 268 | f 3966 15094 21032 269 | f 3978 12576 15094 270 | f 15094 12576 21032 271 | f 3969 20827 12950 272 | f 3555 14588 20827 273 | f 3964 12950 14588 274 | f 14588 12950 20827 275 | f 3552 21136 13050 276 | f 3555 11603 21136 277 | f 3554 13050 11603 278 | f 11603 13050 21136 279 | f 3553 14067 13057 280 | f 3557 21138 14067 281 | f 3556 13057 21138 282 | f 21138 13057 14067 283 | f 3489 21149 13097 284 | f 3492 11898 21149 285 | f 3491 13097 11898 286 | f 11898 13097 21149 287 | f 4401 21158 13180 288 | f 3971 10950 21158 289 | f 3547 13180 10950 290 | f 10950 13180 21158 291 | f 4292 21161 13194 292 | f 3489 20673 21161 293 | f 3970 13194 20673 294 | f 20673 13194 21161 295 | f 3491 11897 13825 296 | f 5511 14354 11897 297 | f 5514 13825 14354 298 | f 14354 13825 11897 299 | f 3548 21289 14016 300 | f 3490 10862 21289 301 | f 3549 14016 10862 302 | f 10862 14016 21289 303 | f 3490 21289 14210 304 | f 3548 20693 21289 305 | f 3547 14210 20693 306 | f 20693 14210 21289 307 | f 5511 14370 14355 308 | f 5723 15017 14370 309 | f 5726 14355 15017 310 | f 15017 14355 14370 311 | f 5512 21334 14372 312 | f 5724 21333 21334 313 | f 5723 14372 21333 314 | f 21333 14372 21334 315 | f 5940 21348 14557 316 | f 3971 21158 21348 317 | f 4401 14557 21158 318 | f 21158 14557 21348 319 | f 3555 21136 14589 320 | f 3552 20872 21136 321 | f 3963 14589 20872 322 | f 20872 14589 21136 323 | f 4401 13179 14893 324 | f 3967 21363 13179 325 | f 5944 14893 21363 326 | f 21363 14893 13179 327 | f 3970 10951 14898 328 | f 3971 21348 10951 329 | f 5940 14898 21348 330 | f 21348 14898 10951 331 | f 5944 21363 14911 332 | f 3967 15411 21363 333 | f 3968 14911 15411 334 | f 15411 14911 21363 335 | f 3963 20873 14913 336 | f 5427 14918 20873 337 | f 5951 14913 14918 338 | f 14918 14913 20873 339 | f 5512 20898 14922 340 | f 3492 14952 20898 341 | f 5943 14922 14952 342 | f 14952 14922 20898 343 | f 5724 21334 14928 344 | f 5512 14921 21334 345 | f 5947 14928 14921 346 | f 14921 14928 21334 347 | f 5949 14909 14931 348 | f 3968 12948 14909 349 | f 3964 14931 12948 350 | f 12948 14931 14909 351 | f 3492 21149 14953 352 | f 3489 21161 21149 353 | f 4292 14953 21161 354 | f 21161 14953 21149 355 | f 3966 11052 15095 356 | f 3548 14015 11052 357 | f 3550 15095 14015 358 | f 14015 15095 11052 359 | f 3966 21032 15412 360 | f 3969 12949 21032 361 | f 3968 15412 12949 362 | f 12949 15412 21032 363 | f 6250 21461 15708 364 | f 6251 16103 21461 365 | f 6252 15708 16103 366 | f 16103 15708 21461 367 | f 6251 21461 15761 368 | f 6250 16255 21461 369 | f 6308 15761 16255 370 | f 16255 15761 21461 371 | f 6313 21480 15767 372 | f 6313 21481 21480 373 | f 6317 15770 21481 374 | f 21481 15770 21480 375 | f 6251 20152 16104 376 | f 3852 21608 16248 377 | f 6711 21607 21608 378 | f 6712 16248 21607 379 | f 21607 16248 21608 380 | f 6714 21609 16254 381 | f 6715 20499 21609 382 | f 6716 16254 20499 383 | f 20499 16254 21609 384 | f 6718 20467 16257 385 | f 6719 17038 20467 386 | f 6308 16257 17038 387 | f 17038 16257 20467 388 | f 6715 21609 16614 389 | f 6714 20171 21609 390 | f 6309 16614 20171 391 | f 20171 16614 21609 392 | f 7137 20457 17040 393 | f 6715 16613 20457 394 | f 6308 17040 16613 395 | f 16613 17040 20457 396 | f 4292 13193 17042 397 | f 4293 21804 13193 398 | f 6718 17042 21804 399 | f 21804 17042 13193 400 | f 6316 15765 17612 401 | f 6316 18865 17624 402 | f 6712 21607 18865 403 | f 6711 17624 21607 404 | f 21607 17624 18865 405 | f 8161 17919 17655 406 | f 6317 21481 17919 407 | f 6313 17655 21481 408 | f 21481 17655 17919 409 | f 6253 15706 18035 410 | f 6252 18043 15706 411 | f 8233 18035 18043 412 | f 18043 18035 15706 413 | f 6252 16102 18044 414 | f 8234 18033 18842 415 | f 8233 18849 18033 416 | f 8417 18842 18849 417 | f 18849 18842 18033 418 | f 8233 18042 18850 419 | f 6717 16252 18866 420 | f 6716 20496 16252 421 | f 6712 18866 20496 422 | f 20496 18866 16252 423 | f 6717 17610 18872 424 | f 6309 20170 20153 425 | f 6714 18871 20172 426 | f 5940 14556 20449 427 | f 5941 22228 14556 428 | f 7137 20449 22228 429 | f 22228 20449 14556 430 | f 6253 20484 20455 431 | f 5943 14951 20484 432 | f 4292 20455 14951 433 | f 14951 20455 20484 434 | f 7137 22228 20458 435 | f 5941 14892 22228 436 | f 5944 20458 14892 437 | f 14892 20458 22228 438 | f 6718 21804 20468 439 | f 4293 14897 21804 440 | f 5940 20468 14897 441 | f 14897 20468 21804 442 | f 8234 20492 20485 443 | f 5947 14920 20492 444 | f 5943 20485 14920 445 | f 14920 20485 20492 446 | f 8418 18911 20493 447 | f 5616 14926 18911 448 | f 5947 20493 14926 449 | f 14926 20493 18911 450 | f 5949 14930 20498 451 | f 3853 16246 14930 452 | f 6712 20498 16246 453 | f 16246 20498 14930 454 | f 5944 14910 20500 455 | f 5949 20497 14910 456 | f 6716 20500 20497 457 | f 20497 20500 14910 458 | f 6711 21608 20510 459 | f 3852 14912 21608 460 | f 5951 20510 14912 461 | f 14912 20510 21608 -------------------------------------------------------------------------------- /lib/data/smplx_lz_faces.obj: -------------------------------------------------------------------------------- 1 | 2 | f 857 858 856 3 | f 856 858 680 4 | f 860 861 859 5 | f 859 861 862 6 | f 917 918 862 7 | f 862 918 859 8 | f 919 920 856 9 | f 856 920 921 10 | f 923 924 922 11 | f 922 924 925 12 | f 926 927 925 13 | f 925 927 922 14 | f 861 1204 862 15 | f 862 1204 921 16 | f 1327 1328 1326 17 | f 1326 1328 1329 18 | f 1334 1335 1333 19 | f 1333 1335 1336 20 | f 1337 1338 918 21 | f 918 1338 859 22 | f 1345 1346 919 23 | f 919 1346 920 24 | f 918 917 1336 25 | f 1336 917 1333 26 | f 1337 918 1768 27 | f 1768 918 1336 28 | f 1769 1770 1338 29 | f 1338 1770 859 30 | f 1334 2824 923 31 | f 923 2824 924 32 | f 922 1329 923 33 | f 923 1329 1328 34 | f 1329 922 2849 35 | f 2849 922 927 36 | f 923 1328 1334 37 | f 1334 1328 1335 38 | f 1333 1346 1334 39 | f 1334 1346 2824 40 | f 919 666 1345 41 | f 1345 666 665 42 | f 856 680 919 43 | f 919 680 666 44 | f 1345 665 3477 45 | f 3477 665 1483 46 | f 3479 2832 3478 47 | f 3478 2832 2831 48 | f 3477 1483 3479 49 | f 3479 1483 2832 50 | f 3478 2831 2846 51 | f 2846 2831 617 52 | f 862 921 917 53 | f 917 921 920 54 | f 3478 925 3479 55 | f 3479 925 924 56 | f 2846 926 3478 57 | f 3478 926 925 58 | f 856 921 857 59 | f 857 921 1204 60 | f 917 920 1333 61 | f 1333 920 1346 62 | f 3479 924 3477 63 | f 3477 924 2824 64 | f 3477 2824 1345 65 | f 1345 2824 1346 66 | f 1337 1768 3501 67 | f 3501 1768 3502 68 | f 859 1770 860 69 | f 860 1770 3504 70 | f 1336 3505 1768 71 | f 1768 3505 3502 72 | f 1337 3501 1338 73 | f 1338 3501 1769 74 | f 1335 1328 3510 75 | f 3510 1328 1327 76 | f 1336 1335 3505 77 | f 3505 1335 3510 78 | f 2849 3512 1329 79 | f 1329 3512 1326 80 | f 4168 4343 4342 81 | f 4342 4343 4344 82 | f 4346 4347 4345 83 | f 4345 4347 4348 84 | f 4345 4403 4346 85 | f 4346 4403 4404 86 | f 4405 4406 4342 87 | f 4342 4406 4407 88 | f 4409 4410 4408 89 | f 4408 4410 4411 90 | f 4408 4412 4409 91 | f 4409 4412 4413 92 | f 4405 4690 4346 93 | f 4346 4690 4347 94 | f 4806 4807 1326 95 | f 1326 4807 1327 96 | f 4810 4811 4809 97 | f 4809 4811 4812 98 | f 4813 4403 4814 99 | f 4814 4403 4345 100 | f 4406 4821 4407 101 | f 4407 4821 4822 102 | f 4809 4404 4810 103 | f 4810 4404 4403 104 | f 4810 4403 5235 105 | f 5235 4403 4813 106 | f 4345 1770 4814 107 | f 4814 1770 1769 108 | f 4410 6285 4411 109 | f 4411 6285 4812 110 | f 4807 4806 4411 111 | f 4411 4806 4408 112 | f 4412 4408 6310 113 | f 6310 4408 4806 114 | f 4811 4807 4812 115 | f 4812 4807 4411 116 | f 6285 4821 4812 117 | f 4812 4821 4809 118 | f 4153 4152 4822 119 | f 4822 4152 4407 120 | f 4152 4168 4407 121 | f 4407 4168 4342 122 | f 4954 4153 6875 123 | f 6875 4153 4822 124 | f 6293 6292 6876 125 | f 6876 6292 6877 126 | f 6292 4954 6877 127 | f 6877 4954 6875 128 | f 4105 6293 6307 129 | f 6307 6293 6876 130 | f 4406 4405 4404 131 | f 4404 4405 4346 132 | f 4410 4409 6877 133 | f 6877 4409 6876 134 | f 4409 4413 6876 135 | f 6876 4413 6307 136 | f 4690 4405 4344 137 | f 4344 4405 4342 138 | f 4821 4406 4809 139 | f 4809 4406 4404 140 | f 6285 4410 6875 141 | f 6875 4410 6877 142 | f 4821 6285 4822 143 | f 4822 6285 6875 144 | f 3502 5235 3501 145 | f 3501 5235 4813 146 | f 3504 1770 4348 147 | f 4348 1770 4345 148 | f 3502 3505 5235 149 | f 5235 3505 4810 150 | f 1769 3501 4814 151 | f 4814 3501 4813 152 | f 1327 4807 3510 153 | f 3510 4807 4811 154 | f 3510 4811 3505 155 | f 3505 4811 4810 156 | f 1326 3512 4806 157 | f 4806 3512 6310 -------------------------------------------------------------------------------- /lib/data/smplx_multi.obj: -------------------------------------------------------------------------------- 1 | 2 | f 3486 3550 3549 3 | f 3486 3487 3310 4 | f 3490 3840 3549 5 | f 3852 3964 3963 6 | f 3970 3547 3971 7 | f 3967 3548 3966 8 | f 3555 5402 3554 9 | f 5427 3552 3556 10 | f 3969 3978 5402 11 | f 3969 3964 3968 12 | f 3552 3554 3553 13 | f 3553 3556 3552 14 | f 3489 3491 3490 15 | f 4401 3547 3967 16 | f 4292 3970 4293 17 | f 3979 3294 3295 18 | f 3551 3310 3294 19 | f 9024 5409 5410 20 | f 9025 4115 5409 21 | f 5424 5410 3246 22 | f 3488 3549 3840 23 | f 9023 3554 5402 24 | f 3548 3549 3550 25 | f 9024 3557 3553 26 | f 3551 3978 3550 27 | f 3490 3547 3489 28 | f 9025 3553 3554 29 | f 5940 4401 5941 30 | f 3555 3963 3964 31 | f 4401 5944 5941 32 | f 3970 5940 4293 33 | f 5944 3968 5949 34 | f 3963 5951 3852 35 | f 5949 3964 3853 36 | f 3492 4292 5943 37 | f 3979 5402 3978 38 | f 3966 3550 3978 39 | f 9023 3295 4115 40 | f 3966 3968 3967 41 | f 6247 6248 6249 42 | f 6250 6252 6253 43 | f 6251 6308 6309 44 | f 6247 6311 6312 45 | f 6313 6315 6316 46 | f 6314 6317 6318 47 | f 6251 6595 6252 48 | f 3852 6712 3853 49 | f 6714 6716 6717 50 | f 6718 6308 6250 51 | f 6312 6726 6727 52 | f 6715 6309 6308 53 | f 7137 6308 6719 54 | f 4292 6718 6250 55 | f 6316 8136 6717 56 | f 6316 6711 6313 57 | f 8161 6313 6711 58 | f 6717 6712 6316 59 | f 6717 6726 6714 60 | f 6727 6057 6312 61 | f 6312 6073 6247 62 | f 8723 6058 6727 63 | f 8724 8143 8725 64 | f 8725 6859 8723 65 | f 8158 8144 8724 66 | f 6309 6310 6251 67 | f 8725 6314 8724 68 | f 8724 6318 8158 69 | f 6249 6310 6247 70 | f 6714 6311 6309 71 | f 8723 6315 8725 72 | f 6727 8136 8723 73 | f 5940 7137 6719 74 | f 6253 4292 6250 75 | f 7137 5944 6715 76 | f 6718 5940 6719 77 | f 5949 6712 6716 78 | f 5944 6716 6715 79 | f 6711 5951 8161 80 | f 3486 3551 3550 81 | f 3486 3488 3487 82 | f 3490 3491 3840 83 | f 3852 3853 3964 84 | f 3970 3489 3547 85 | f 3967 3547 3548 86 | f 3555 3969 5402 87 | f 5427 3963 3552 88 | f 3969 3966 3978 89 | f 3969 3555 3964 90 | f 3552 3555 3554 91 | f 3553 3557 3556 92 | f 3489 3492 3491 93 | f 4401 3971 3547 94 | f 4292 3489 3970 95 | f 3979 3551 3294 96 | f 3551 3486 3310 97 | f 9024 9025 5409 98 | f 9025 9023 4115 99 | f 5424 9024 5410 100 | f 3488 3486 3549 101 | f 9023 9025 3554 102 | f 3548 3490 3549 103 | f 9024 5424 3557 104 | f 3551 3979 3978 105 | f 3490 3548 3547 106 | f 9025 9024 3553 107 | f 5940 3971 4401 108 | f 3555 3552 3963 109 | f 4401 3967 5944 110 | f 3970 3971 5940 111 | f 5944 3967 3968 112 | f 3963 5427 5951 113 | f 5949 3968 3964 114 | f 3492 3489 4292 115 | f 3979 9023 5402 116 | f 3966 3548 3550 117 | f 9023 3979 3295 118 | f 3966 3969 3968 119 | f 6247 6073 6248 120 | f 6250 6251 6252 121 | f 6251 6250 6308 122 | f 6247 6310 6311 123 | f 6313 6314 6315 124 | f 6314 6313 6317 125 | f 6251 6310 6595 126 | f 3852 6711 6712 127 | f 6714 6715 6716 128 | f 6718 6719 6308 129 | f 6312 6311 6726 130 | f 6715 6714 6309 131 | f 7137 6715 6308 132 | f 4292 4293 6718 133 | f 6316 6315 8136 134 | f 6316 6712 6711 135 | f 8161 6317 6313 136 | f 6717 6716 6712 137 | f 6717 8136 6726 138 | f 6727 6058 6057 139 | f 6312 6057 6073 140 | f 8723 6859 6058 141 | f 8724 8144 8143 142 | f 8725 8143 6859 143 | f 8158 6009 8144 144 | f 6309 6311 6310 145 | f 8725 6315 6314 146 | f 8724 6314 6318 147 | f 6249 6595 6310 148 | f 6714 6726 6311 149 | f 8723 8136 6315 150 | f 6727 6726 8136 151 | f 5940 5941 7137 152 | f 6253 5943 4292 153 | f 7137 5941 5944 154 | f 6718 4293 5940 155 | f 5949 3853 6712 156 | f 5944 5949 6716 157 | f 6711 3852 5951 -------------------------------------------------------------------------------- /lib/ext_transform.py: -------------------------------------------------------------------------------- 1 | import random 2 | 3 | import numpy as np 4 | from skimage.filters import gaussian 5 | import torch 6 | from PIL import Image, ImageFilter 7 | 8 | 9 | class RandomVerticalFlip(object): 10 | def __call__(self, img): 11 | if random.random() < 0.5: 12 | return img.transpose(Image.FLIP_TOP_BOTTOM) 13 | return img 14 | 15 | 16 | class DeNormalize(object): 17 | def __init__(self, mean, std): 18 | self.mean = mean 19 | self.std = std 20 | 21 | def __call__(self, tensor): 22 | for t, m, s in zip(tensor, self.mean, self.std): 23 | t.mul_(s).add_(m) 24 | return tensor 25 | 26 | 27 | class MaskToTensor(object): 28 | def __call__(self, img): 29 | return torch.from_numpy(np.array(img, dtype=np.int32)).long() 30 | 31 | 32 | class FreeScale(object): 33 | def __init__(self, size, interpolation=Image.BILINEAR): 34 | self.size = tuple(reversed(size)) # size: (h, w) 35 | self.interpolation = interpolation 36 | 37 | def __call__(self, img): 38 | return img.resize(self.size, self.interpolation) 39 | 40 | 41 | class FlipChannels(object): 42 | def __call__(self, img): 43 | img = np.array(img)[:, :, ::-1] 44 | return Image.fromarray(img.astype(np.uint8)) 45 | 46 | 47 | class RandomGaussianBlur(object): 48 | def __call__(self, img): 49 | sigma = 0.15 + random.random() * 1.15 50 | blurred_img = gaussian(np.array(img), sigma=sigma, multichannel=True) 51 | blurred_img *= 255 52 | return Image.fromarray(blurred_img.astype(np.uint8)) 53 | 54 | # Lighting data augmentation take from here - https://github.com/eladhoffer/convNet.pytorch/blob/master/preprocess.py 55 | 56 | 57 | class Lighting(object): 58 | """Lighting noise(AlexNet - style PCA - based noise)""" 59 | 60 | def __init__(self, alphastd, 61 | eigval=(0.2175, 0.0188, 0.0045), 62 | eigvec=((-0.5675, 0.7192, 0.4009), 63 | (-0.5808, -0.0045, -0.8140), 64 | (-0.5836, -0.6948, 0.4203))): 65 | self.alphastd = alphastd 66 | self.eigval = torch.Tensor(eigval) 67 | self.eigvec = torch.Tensor(eigvec) 68 | 69 | def __call__(self, img): 70 | if self.alphastd == 0: 71 | return img 72 | 73 | alpha = img.new().resize_(3).normal_(0, self.alphastd) 74 | rgb = self.eigvec.type_as(img).clone()\ 75 | .mul(alpha.view(1, 3).expand(3, 3))\ 76 | .mul(self.eigval.view(1, 3).expand(3, 3))\ 77 | .sum(1).squeeze() 78 | return img.add(rgb.view(3, 1, 1).expand_as(img)) 79 | -------------------------------------------------------------------------------- /lib/geometry.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import numpy as np 3 | import math 4 | 5 | def rotationX(angle): 6 | return [ 7 | [1, 0, 0], 8 | [0, math.cos(angle), -math.sin(angle)], 9 | [0, math.sin(angle), math.cos(angle)], 10 | ] 11 | 12 | 13 | def rotationY(angle): 14 | return [ 15 | [math.cos(angle), 0, math.sin(angle)], 16 | [0, 1, 0], 17 | [-math.sin(angle), 0, math.cos(angle)], 18 | ] 19 | 20 | 21 | def rotationZ(angle): 22 | return [ 23 | [math.cos(angle), -math.sin(angle), 0], 24 | [math.sin(angle), math.cos(angle), 0], 25 | [0, 0, 1], 26 | ] 27 | 28 | def batch_cross_3d(a, b): 29 | c = torch.zeros(a.shape[0], 3) 30 | c[:, 0], c[:, 1], c[:, 2] = a[:, 1]*b[:, 2]-a[:, 2]*b[:, 1], b[:, 0]*a[:, 2]-a[:, 0]*b[:, 2], a[:, 0]*b[:, 1]-b[:, 0]*a[:, 1] 31 | return c 32 | 33 | def cross_3d(a, b): 34 | return np.array([a[1]*b[2]-a[2]*b[1], b[0]*a[2]-a[0]*b[2], a[0]*b[1]-b[0]*a[1]]) 35 | 36 | def index(feat, uv, size=None): 37 | ''' 38 | :param feat: [B, C, H, W] image features 39 | :param uv: [B, 2, N] uv coordinates in the image plane, range [-1, 1] 40 | :return: [B, C, N] image features at the uv coordinates 41 | ''' 42 | uv = uv.transpose(1, 2) # [B, N, 2] 43 | uv = uv.unsqueeze(2) # [B, N, 1, 2] 44 | if size != None: 45 | uv = (uv - size / 2) / (size / 2) 46 | # NOTE: for newer PyTorch, it seems that training results are degraded due to implementation diff in F.grid_sample 47 | # for old versions, simply remove the aligned_corners argument. 48 | samples = torch.nn.functional.grid_sample(feat, uv, align_corners=True) # [B, C, N, 1] 49 | return samples[:, :, :, 0] # [B, C, N] 50 | 51 | 52 | def orthogonal(points, calibrations, transforms=None): 53 | ''' 54 | Compute the orthogonal projections of 3D points into the image plane by given projection matrix 55 | :param points: [B, 3, N] Tensor of 3D points 56 | :param calibrations: [B, 4, 4] Tensor of projection matrix 57 | :param transforms: [B, 2, 3] Tensor of image transform matrix 58 | :return: xyz: [B, 3, N] Tensor of xyz coordinates in the image plane 59 | ''' 60 | rot = calibrations[:, :3, :3] 61 | trans = calibrations[:, :3, 3:4] 62 | pts = torch.baddbmm(trans, rot, points) # [B, 3, N] 63 | if transforms is not None: 64 | scale = transforms[:2, :2] 65 | shift = transforms[:2, 2:3] 66 | pts[:, :2, :] = torch.baddbmm(shift, scale, pts[:, :2, :]) 67 | return pts 68 | 69 | 70 | def perspective(points, calibrations, transforms=None): 71 | ''' 72 | Compute the perspective projections of 3D points into the image plane by given projection matrix 73 | :param points: [Bx3xN] Tensor of 3D points 74 | :param calibrations: [Bx4x4] Tensor of projection matrix 75 | :param transforms: [Bx2x3] Tensor of image transform matrix 76 | :return: xy: [Bx2xN] Tensor of xy coordinates in the image plane 77 | ''' 78 | B, _, N = points.shape 79 | device = points.device 80 | points = torch.cat([points, torch.ones((B, 1, N), device=device)], dim=1) 81 | points = calibrations @ points 82 | points[:, :2, :] /= points[:, 2:, :] 83 | return points 84 | 85 | def rotationMatrixToAngles(R): 86 | """ 87 | R : (bs, 3, 3) 88 | """ 89 | # print(R.shape) 90 | sy = torch.sqrt(R[:, 0, 0] * R[:, 0, 0] + R[:, 1, 0] * R[:, 1, 0]) 91 | singular = sy < 1e-6 92 | mask = ~singular 93 | x = torch.zeros(R.shape[0]) 94 | y = torch.zeros(R.shape[0]) 95 | z = torch.zeros(R.shape[0]) 96 | if torch.sum(mask): 97 | x[mask] = torch.atan2(R[mask, 2, 1], R[mask, 2, 2]) 98 | y[mask] = torch.atan2(-R[mask, 2, 0], sy[mask]) 99 | z[mask] = torch.atan2(R[mask, 1, 0], R[mask, 0, 0]) 100 | if torch.sum(singular): 101 | x[singular] = math.atan2(-R[singular, 1, 2], R[singular, 1, 1]) 102 | y[singular] = torch.atan2(-R[singular, 2, 0], sy[singular]) 103 | z[singular] = 0 104 | return torch.cat([x.unsqueeze(1), y.unsqueeze(1), z.unsqueeze(1)], dim=1)# np.array([x, y, z]) -------------------------------------------------------------------------------- /lib/loss_util.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torchvision 3 | 4 | class VGGPerceptualLoss(torch.nn.Module): 5 | def __init__(self, resize=True): 6 | super(VGGPerceptualLoss, self).__init__() 7 | blocks = [] 8 | blocks.append(torchvision.models.vgg16(pretrained=True).features[:4].eval()) 9 | blocks.append(torchvision.models.vgg16(pretrained=True).features[4:9].eval()) 10 | blocks.append(torchvision.models.vgg16(pretrained=True).features[9:16].eval()) 11 | blocks.append(torchvision.models.vgg16(pretrained=True).features[16:23].eval()) 12 | for bl in blocks: 13 | for p in bl: 14 | p.requires_grad = False 15 | self.blocks = torch.nn.ModuleList(blocks) 16 | self.transform = torch.nn.functional.interpolate 17 | self.mean = torch.nn.Parameter(torch.tensor([0.485, 0.456, 0.406]).view(1,3,1,1)) 18 | self.std = torch.nn.Parameter(torch.tensor([0.229, 0.224, 0.225]).view(1,3,1,1)) 19 | self.resize = resize 20 | 21 | def forward(self, input, target): 22 | if input.shape[1] != 3: 23 | input = input.repeat(1, 3, 1, 1) 24 | target = target.repeat(1, 3, 1, 1) 25 | if self.resize: 26 | input = self.transform(input, mode='bilinear', size=(224, 224), align_corners=False) 27 | target = self.transform(target, mode='bilinear', size=(224, 224), align_corners=False) 28 | loss = 0.0 29 | x = input 30 | y = target 31 | for block in self.blocks: 32 | x = block(x) 33 | y = block(y) 34 | loss += torch.nn.functional.l1_loss(x, y) 35 | return loss -------------------------------------------------------------------------------- /lib/mesh_util.py: -------------------------------------------------------------------------------- 1 | from skimage import measure 2 | import numpy as np 3 | import torch 4 | from PIL import Image 5 | import torch.nn.functional as F 6 | import trimesh 7 | 8 | from .sdf import create_grid, eval_grid_octree, eval_grid 9 | from .net_util import reshape_sample_tensor 10 | from .geometry import index 11 | 12 | 13 | def gen_mesh_dmc(opt, net, cuda, data, save_path, use_octree=True, threshold=0.5): 14 | image_tensor = data['image'].to(device=cuda) 15 | calib_tensor = data['calib'].to(device=cuda) 16 | extrinsic = data['extrinsic'].to(device=cuda).unsqueeze(0) 17 | vox_tensor = data['vox'].to(device=cuda).unsqueeze(0) 18 | smpl_normal = data['smpl_normal'].to(device=cuda).unsqueeze(0) 19 | save_smpl_normal = smpl_normal.clone() 20 | normal_tensor = data['normal'].to(device=cuda) 21 | scale, center = data['scale'].to(device=cuda).unsqueeze(0), data['center'].to(device=cuda).unsqueeze(0) 22 | mask, ero_mask = data['mask'].to(device=cuda).unsqueeze(0), data['ero_mask'].to(device=cuda).unsqueeze(0) 23 | 24 | net.mask_init(mask, ero_mask) 25 | net.norm_init(scale, center) 26 | net.smpl_init(smpl_normal) 27 | 28 | net.filter2d(torch.cat([image_tensor.unsqueeze(0), smpl_normal], dim=2)) 29 | if opt.fine_part: 30 | if normal_tensor.shape[2] == 1024: 31 | print('1024') 32 | smpl_normal = torch.nn.Upsample(size=[1024, 1024], mode='bilinear')(smpl_normal.squeeze(0)).unsqueeze(0) 33 | net.filter_normal(torch.cat([normal_tensor.unsqueeze(0), smpl_normal], dim=2)) 34 | 35 | net.filter3d(vox_tensor) 36 | 37 | b_min = data['b_min'] 38 | b_max = data['b_max'] 39 | 40 | save_img_path = save_path[:-4] + '.png' 41 | save_img_list = [] 42 | for v in range(image_tensor.shape[0]): 43 | save_img = (np.transpose(image_tensor[v].detach().cpu().numpy(), (1, 2, 0)) * 0.5 + 0.5)[:, :, ::-1] * 255.0 44 | save_smpl = (np.transpose(save_smpl_normal[0][v].detach().cpu().numpy(), (1, 2, 0)) * 0.5 + 0.5)[:, :, ::-1] * 255.0 45 | save_img_list.append(save_img / 2 + save_smpl / 2) 46 | for v in range(normal_tensor.shape[0]): 47 | save_nm = normal_tensor[v] 48 | save_nm = F.interpolate(save_nm.unsqueeze(0), size=[512, 512], mode='bilinear')[0] 49 | save_nm = (np.transpose(save_nm.detach().cpu().numpy(), (1, 2, 0)) * 0.5 + 0.5)[:, :, ::-1] * 255.0 50 | save_img_list.append(save_nm) 51 | # save_nm = smpl_normal[0:, v] 52 | # save_nm = F.interpolate(save_nm, size=[512, 512], mode='bilinear')[0] 53 | # save_nm = (np.transpose(save_nm.detach().cpu().numpy(), (1, 2, 0)) * 0.5 + 0.5)[:, :, ::-1] * 255.0 54 | # save_img_list.append(save_nm) 55 | save_img = np.concatenate(save_img_list, axis=1) 56 | Image.fromarray(np.uint8(save_img[:,:,::-1])).save(save_img_path) 57 | 58 | try: 59 | verts, faces, _, _ = reconstruction_3d( 60 | net, cuda, calib_tensor.unsqueeze(0), extrinsic, opt.resolution, b_min, b_max, use_octree=use_octree, threshold=threshold) 61 | verts_tensor = torch.from_numpy(verts.T).unsqueeze(0).to(device=cuda).float() 62 | xyz_tensor = net.projection(verts_tensor, calib_tensor[:1]) 63 | uv = xyz_tensor[:, :2, :] 64 | color = index(image_tensor[:1], uv).detach().cpu().numpy()[0].T 65 | color = color * 0.5 + 0.5 66 | save_obj_mesh_with_color(save_path, verts, faces, color) 67 | except Exception as e: 68 | print(e) 69 | 70 | 71 | def reconstruction_3d(net, cuda, calib_tensor, extrinsic, 72 | resolution, b_min, b_max, 73 | net_3d=False, use_octree=False, num_samples=30000, threshold=0.5, transform=None): 74 | ''' 75 | Reconstruct meshes from sdf predicted by the network. 76 | :param net: a BasePixImpNet object. call image filter beforehead. 77 | :param cuda: cuda device 78 | :param calib_tensor: calibration tensor 79 | :param resolution: resolution of the grid cell 80 | :param b_min: bounding box corner [x_min, y_min, z_min] 81 | :param b_max: bounding box corner [x_max, y_max, z_max] 82 | :param use_octree: whether to use octree acceleration 83 | :param num_samples: how many points to query each gpu iteration 84 | :return: marching cubes results. 85 | ''' 86 | # First we create a grid by resolution 87 | # and transforming matrix for grid coordinates to real world xyz 88 | #print(b_min, b_max) 89 | coords, mat = create_grid(resolution, resolution, resolution, 90 | b_min, b_max, transform=transform) 91 | # print(coords.shape, mat.shape) 92 | # Then we define the lambda function for cell evaluation 93 | def eval_func(points): 94 | # print(points.shape) 95 | points = np.expand_dims(points, axis=0) 96 | samples = torch.from_numpy(points).to(device=cuda).float() 97 | net.query(samples, calib_tensor, extrinsic) 98 | pred = net.get_preds()[0][0] 99 | return pred.detach().cpu().numpy() 100 | 101 | # Then we evaluate the grid 102 | if use_octree: 103 | sdf = eval_grid_octree(coords, eval_func, num_samples=num_samples) 104 | else: 105 | sdf = eval_grid(coords, eval_func, num_samples=num_samples) 106 | 107 | # Finally we do marching cubes 108 | #try: 109 | verts, faces, normals, values = measure.marching_cubes_lewiner(sdf, threshold) 110 | # transform verts into world coordinate system 111 | verts = np.matmul(mat[:3, :3], verts.T) + mat[:3, 3:4] 112 | verts = verts.T 113 | return verts, faces, normals, values 114 | #except: 115 | # print('error cannot marching cubes') 116 | # return -1 117 | 118 | def save_obj_mesh(mesh_path, verts, faces): 119 | file = open(mesh_path, 'w') 120 | 121 | for v in verts: 122 | file.write('v %.4f %.4f %.4f\n' % (v[0], v[1], v[2])) 123 | for f in faces: 124 | f_plus = f + 1 125 | file.write('f %d %d %d\n' % (f_plus[0], f_plus[2], f_plus[1])) 126 | file.close() 127 | 128 | 129 | def save_obj_mesh_with_color(mesh_path, verts, faces, colors, reverse=False): 130 | file = open(mesh_path, 'w') 131 | 132 | for idx, v in enumerate(verts): 133 | c = colors[idx] 134 | file.write('v %.4f %.4f %.4f %.4f %.4f %.4f\n' % (v[0], v[1], v[2], c[0], c[1], c[2])) 135 | for f in faces: 136 | f_plus = f + 1 137 | if reverse: 138 | file.write('f %d %d %d\n' % (f_plus[0], f_plus[1], f_plus[2])) 139 | else: 140 | file.write('f %d %d %d\n' % (f_plus[0], f_plus[2], f_plus[1])) 141 | file.close() 142 | 143 | 144 | def save_obj_mesh_with_uv(mesh_path, verts, faces, uvs): 145 | file = open(mesh_path, 'w') 146 | 147 | for idx, v in enumerate(verts): 148 | vt = uvs[idx] 149 | file.write('v %.4f %.4f %.4f\n' % (v[0], v[1], v[2])) 150 | file.write('vt %.4f %.4f\n' % (vt[0], vt[1])) 151 | 152 | for f in faces: 153 | f_plus = f + 1 154 | file.write('f %d/%d %d/%d %d/%d\n' % (f_plus[0], f_plus[0], 155 | f_plus[2], f_plus[2], 156 | f_plus[1], f_plus[1])) 157 | file.close() 158 | 159 | 160 | 161 | def _append(faces, indices): 162 | if len(indices) == 4: 163 | faces.append([indices[0], indices[1], indices[2]]) 164 | faces.append([indices[2], indices[3], indices[0]]) 165 | elif len(indices) == 3: 166 | faces.append(indices) 167 | else: 168 | assert False, len(indices) 169 | 170 | 171 | def readobj(path, scale=1): 172 | vi = [] 173 | vt = [] 174 | vn = [] 175 | faces = [] 176 | 177 | with open(path, 'r') as myfile: 178 | lines = myfile.readlines() 179 | 180 | # cache vertices 181 | for line in lines: 182 | try: 183 | type, fields = line.split(maxsplit=1) 184 | fields = [float(_) for _ in fields.split()] 185 | except ValueError: 186 | continue 187 | 188 | if type == 'v': 189 | vi.append(fields) 190 | elif type == 'vt': 191 | vt.append(fields) 192 | elif type == 'vn': 193 | vn.append(fields) 194 | 195 | # cache faces 196 | for line in lines: 197 | try: 198 | type, fields = line.split(maxsplit=1) 199 | fields = fields.split() 200 | except ValueError: 201 | continue 202 | 203 | # line looks like 'f 5/1/1 1/2/1 4/3/1' 204 | # or 'f 314/380/494 382/400/494 388/550/494 506/551/494' for quads 205 | if type != 'f': 206 | continue 207 | 208 | # a field should look like '5/1/1' 209 | # for vertex/vertex UV coords/vertex Normal (indexes number in the list) 210 | # the index in 'f 5/1/1 1/2/1 4/3/1' STARTS AT 1 !!! 211 | indices = [[int(_) - 1 if _ != '' else 0 for _ in field.split('/')] for field in fields] 212 | 213 | if len(indices) == 4: 214 | faces.append([indices[0], indices[1], indices[2]]) 215 | faces.append([indices[2], indices[3], indices[0]]) 216 | elif len(indices) == 3: 217 | faces.append(indices) 218 | else: 219 | assert False, len(indices) 220 | 221 | ret = {} 222 | ret['vi'] = None if len(vi) == 0 else np.array(vi).astype(np.float32) * scale 223 | ret['vt'] = None if len(vt) == 0 else np.array(vt).astype(np.float32) 224 | ret['vn'] = None if len(vn) == 0 else np.array(vn).astype(np.float32) 225 | ret['f'] = None if len(faces) == 0 else np.array(faces).astype(np.int32) 226 | return ret 227 | 228 | -------------------------------------------------------------------------------- /lib/model/BaseNet.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | import torch.nn.functional as F 4 | 5 | from ..geometry import index, orthogonal, perspective 6 | 7 | class BaseNet(nn.Module): 8 | def __init__(self, 9 | projection_mode='orthogonal', 10 | error_term=nn.MSELoss(), 11 | ): 12 | """ 13 | :param projection_mode: 14 | Either orthogonal or perspective. 15 | It will call the corresponding function for projection. 16 | :param error_term: 17 | nn Loss between the predicted [B, Res, N] and the label [B, Res, N] 18 | """ 19 | super(BaseNet, self).__init__() 20 | self.name = 'base' 21 | 22 | self.error_term = error_term 23 | 24 | self.index = index 25 | self.projection = orthogonal if projection_mode == 'orthogonal' else perspective 26 | 27 | self.preds = None 28 | self.labels = None 29 | 30 | def forward(self, points, images, calibs, transforms=None): 31 | ''' 32 | :param points: [B, 3, N] world space coordinates of points 33 | :param images: [B, C, H, W] input images 34 | :param calibs: [B, 3, 4] calibration matrices for each image 35 | :param transforms: Optional [B, 2, 3] image space coordinate transforms 36 | :return: [B, Res, N] predictions for each point 37 | ''' 38 | self.filter(images) 39 | self.query(points, calibs, transforms) 40 | return self.get_preds() 41 | 42 | def filter3d(self, points): 43 | ''' 44 | Filter the input point cloud 45 | store all intermediate 3d features. 46 | :param points: [B, 3, N] input point cloud 47 | ''' 48 | None 49 | 50 | def filter2d(self, images): 51 | ''' 52 | Filter the input images 53 | store all intermediate 2d features. 54 | :param images: [B, C, H, W] input images 55 | ''' 56 | None 57 | 58 | def query(self, points, calibs, transforms=None, labels=None): 59 | ''' 60 | Given 3D points, query the network predictions for each point. 61 | Image features should be pre-computed before this call. 62 | store all intermediate features. 63 | query() function may behave differently during training/testing. 64 | :param points: [B, 3, N] world space coordinates of points 65 | :param calibs: [B, 3, 4] calibration matrices for each image 66 | :param transforms: Optional [B, 2, 3] image space coordinate transforms 67 | :param labels: Optional [B, Res, N] gt labeling 68 | :return: [B, Res, N] predictions for each point 69 | ''' 70 | None 71 | 72 | def get_preds(self): 73 | ''' 74 | Get the predictions from the last query 75 | :return: [B, Res, N] network prediction for the last query 76 | ''' 77 | return self.preds 78 | 79 | def get_error(self): 80 | ''' 81 | Get the network loss from the last query 82 | :return: loss term 83 | ''' 84 | return self.error_term(self.preds, self.labels) 85 | -------------------------------------------------------------------------------- /lib/model/HGFilters.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | import torch.nn.functional as F 4 | from ..net_util import * 5 | 6 | 7 | class HourGlass(nn.Module): 8 | def __init__(self, num_modules, depth, num_features, norm='batch', base=64): 9 | super(HourGlass, self).__init__() 10 | self.num_modules = num_modules 11 | self.depth = depth 12 | self.features = num_features 13 | self.norm = norm 14 | self.base = base 15 | 16 | self._generate_network(self.depth) 17 | 18 | def _generate_network(self, level): 19 | base = self.base 20 | self.add_module('b1_' + str(level), ConvBlock(self.features, self.features, norm=self.norm, base=base)) 21 | 22 | self.add_module('b2_' + str(level), ConvBlock(self.features, self.features, norm=self.norm, base=base)) 23 | 24 | if level > 1: 25 | self._generate_network(level - 1) 26 | else: 27 | self.add_module('b2_plus_' + str(level), ConvBlock(self.features, self.features, norm=self.norm, base=base)) 28 | 29 | self.add_module('b3_' + str(level), ConvBlock(self.features, self.features, norm=self.norm, base=base)) 30 | 31 | def _forward(self, level, inp): 32 | # Upper branch 33 | up1 = inp 34 | up1 = self._modules['b1_' + str(level)](up1) 35 | 36 | # Lower branch 37 | low1 = F.avg_pool2d(inp, 2, stride=2) 38 | low1 = self._modules['b2_' + str(level)](low1) 39 | 40 | if level > 1: 41 | low2 = self._forward(level - 1, low1) 42 | else: 43 | low2 = low1 44 | low2 = self._modules['b2_plus_' + str(level)](low2) 45 | 46 | low3 = low2 47 | low3 = self._modules['b3_' + str(level)](low3) 48 | 49 | # NOTE: for newer PyTorch (1.3~), it seems that training results are degraded due to implementation diff in F.grid_sample 50 | # if the pretrained model behaves weirdly, switch with the commented line. 51 | # NOTE: I also found that "bicubic" works better. 52 | up2 = F.interpolate(low3, scale_factor=2, mode='bicubic', align_corners=True) 53 | # up2 = F.interpolate(low3, scale_factor=2, mode='nearest) 54 | 55 | return up1 + up2 56 | 57 | def forward(self, x): 58 | return self._forward(self.depth, x) 59 | 60 | 61 | class HGFilter(nn.Module): 62 | def __init__(self, opt, base=64, downsample=1): 63 | super(HGFilter, self).__init__() 64 | self.num_modules = opt.num_stack 65 | 66 | self.opt = opt 67 | self.downsample = downsample 68 | # Base part 69 | if downsample == 0: 70 | self.conv1_h = nn.Conv2d(3, base, kernel_size=3, stride=1, padding=1) 71 | self.conv1_nh = nn.Conv2d(6, base, kernel_size=3, stride=1, padding=1) 72 | else: 73 | self.conv1 = nn.Conv2d(3, base, kernel_size=7, stride=2, padding=3) 74 | self.conv1_n = nn.Conv2d(6, base, kernel_size=7, stride=2, padding=3) 75 | 76 | if self.opt.norm == 'batch': 77 | self.bn1 = nn.BatchNorm2d(base) 78 | elif self.opt.norm == 'group': 79 | self.bn1 = nn.GroupNorm(base // 2, base) 80 | 81 | if self.opt.hg_down == 'conv64': 82 | self.conv2 = ConvBlock(base, base, self.opt.norm, base) 83 | self.down_conv2 = nn.Conv2d(base, base * 2, kernel_size=3, stride=2, padding=1) 84 | elif self.opt.hg_down == 'conv128': 85 | self.conv2 = ConvBlock(base, base * 2, self.opt.norm, base) 86 | self.down_conv2 = nn.Conv2d(base * 2, base * 2, kernel_size=3, stride=2, padding=1) 87 | elif self.opt.hg_down == 'ave_pool': 88 | self.conv2 = ConvBlock(base, base * 2, self.opt.norm, base) 89 | else: 90 | raise NameError('Unknown Fan Filter setting!') 91 | 92 | self.conv3 = ConvBlock(base * 2, base * 2, self.opt.norm, base) 93 | self.conv4 = ConvBlock(base * 2, base * 4, self.opt.norm, base) 94 | 95 | # Stacking part 96 | for hg_module in range(self.num_modules): 97 | self.add_module('m' + str(hg_module), HourGlass(1, opt.num_hourglass, base * 4, self.opt.norm, base=base)) 98 | 99 | self.add_module('top_m_' + str(hg_module), ConvBlock(base * 4, base * 4, self.opt.norm, base)) 100 | self.add_module('conv_last' + str(hg_module), 101 | nn.Conv2d(base * 4, base * 4, kernel_size=1, stride=1, padding=0)) 102 | if self.opt.norm == 'batch': 103 | self.add_module('bn_end' + str(hg_module), nn.BatchNorm2d(base * 4)) 104 | elif self.opt.norm == 'group': 105 | self.add_module('bn_end' + str(hg_module), nn.GroupNorm(base // 2, base * 4)) 106 | 107 | self.add_module('l' + str(hg_module), nn.Conv2d(base * 4, 108 | opt.hourglass_dim, kernel_size=1, stride=1, padding=0)) 109 | 110 | if hg_module < self.num_modules - 1: 111 | self.add_module( 112 | 'bl' + str(hg_module), nn.Conv2d(base * 4, base * 4, kernel_size=1, stride=1, padding=0)) 113 | self.add_module('al' + str(hg_module), nn.Conv2d(opt.hourglass_dim, 114 | base * 4, kernel_size=1, stride=1, padding=0)) 115 | 116 | def forward(self, x): 117 | if self.downsample == 0: 118 | x = F.relu(self.bn1(self.conv1_nh(x)), True) 119 | else: 120 | x = F.relu(self.bn1(self.conv1_n(x)), True) 121 | 122 | tmpx = x 123 | if self.opt.hg_down == 'ave_pool': 124 | if self.downsample == 0: 125 | x = self.conv2(x) 126 | else: 127 | x = F.avg_pool2d(self.conv2(x), 2, stride=2) 128 | elif self.opt.hg_down in ['conv64', 'conv128']: 129 | x = self.conv2(x) 130 | x = self.down_conv2(x) 131 | else: 132 | raise NameError('Unknown Fan Filter setting!') 133 | 134 | normx = x 135 | 136 | x = self.conv3(x) 137 | x = self.conv4(x) 138 | 139 | previous = x 140 | 141 | outputs = [] 142 | for i in range(self.num_modules): 143 | hg = self._modules['m' + str(i)](previous) 144 | 145 | ll = hg 146 | ll = self._modules['top_m_' + str(i)](ll) 147 | 148 | ll = F.relu(self._modules['bn_end' + str(i)] 149 | (self._modules['conv_last' + str(i)](ll)), True) 150 | 151 | # Predict heatmaps 152 | tmp_out = self._modules['l' + str(i)](ll) 153 | outputs.append(tmp_out) 154 | 155 | if i < self.num_modules - 1: 156 | ll = self._modules['bl' + str(i)](ll) 157 | tmp_out_ = self._modules['al' + str(i)](tmp_out) 158 | previous = previous + ll + tmp_out_ 159 | 160 | return outputs, tmpx.detach(), normx 161 | -------------------------------------------------------------------------------- /lib/model/NormalNet.py: -------------------------------------------------------------------------------- 1 | ''' 2 | Copyright (C) 2019 NVIDIA Corporation. Ting-Chun Wang, Ming-Yu Liu, Jun-Yan Zhu. 3 | BSD License. All rights reserved. 4 | Redistribution and use in source and binary forms, with or without 5 | modification, are permitted provided that the following conditions are met: 6 | * Redistributions of source code must retain the above copyright notice, this 7 | list of conditions and the following disclaimer. 8 | * Redistributions in binary form must reproduce the above copyright notice, 9 | this list of conditions and the following disclaimer in the documentation 10 | and/or other materials provided with the distribution. 11 | THE AUTHOR DISCLAIMS ALL WARRANTIES WITH REGARD TO THIS SOFTWARE, INCLUDING ALL 12 | IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR ANY PARTICULAR PURPOSE. 13 | IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR ANY SPECIAL, INDIRECT OR CONSEQUENTIAL 14 | DAMAGES OR ANY DAMAGES WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, 15 | WHETHER IN AN ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING 16 | OUT OF OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE. 17 | ''' 18 | import torch 19 | import torch.nn as nn 20 | import functools 21 | from torch.autograd import Variable 22 | import numpy as np 23 | 24 | ############################################################################### 25 | # Functions 26 | ############################################################################### 27 | def weights_init(m): 28 | classname = m.__class__.__name__ 29 | if classname.find('Conv') != -1: 30 | m.weight.data.normal_(0.0, 0.02) 31 | elif classname.find('BatchNorm2d') != -1: 32 | m.weight.data.normal_(1.0, 0.02) 33 | m.bias.data.fill_(0) 34 | 35 | def get_norm_layer(norm_type='instance'): 36 | if norm_type == 'batch': 37 | norm_layer = functools.partial(nn.BatchNorm2d, affine=True) 38 | elif norm_type == 'instance': 39 | norm_layer = functools.partial(nn.InstanceNorm2d, affine=False) 40 | else: 41 | raise NotImplementedError('normalization layer [%s] is not found' % norm_type) 42 | return norm_layer 43 | 44 | def define_G(input_nc, output_nc, ngf, netG, n_downsample_global=3, n_blocks_global=9, n_local_enhancers=1, 45 | n_blocks_local=3, norm='instance', gpu_ids=[], last_op=nn.Tanh()): 46 | norm_layer = get_norm_layer(norm_type=norm) 47 | if netG == 'global': 48 | netG = GlobalGenerator(input_nc, output_nc, ngf, n_downsample_global, n_blocks_global, norm_layer, last_op=last_op) 49 | elif netG == 'local': 50 | netG = LocalEnhancer(input_nc, output_nc, ngf, n_downsample_global, n_blocks_global, 51 | n_local_enhancers, n_blocks_local, norm_layer) 52 | elif netG == 'encoder': 53 | netG = Encoder(input_nc, output_nc, ngf, n_downsample_global, norm_layer) 54 | else: 55 | raise('generator not implemented!') 56 | # print(netG) 57 | if len(gpu_ids) > 0: 58 | assert(torch.cuda.is_available()) 59 | netG.cuda(gpu_ids[0]) 60 | netG.apply(weights_init) 61 | return netG 62 | 63 | def print_network(net): 64 | if isinstance(net, list): 65 | net = net[0] 66 | num_params = 0 67 | for param in net.parameters(): 68 | num_params += param.numel() 69 | print(net) 70 | print('Total number of parameters: %d' % num_params) 71 | 72 | ############################################################################## 73 | # Generator 74 | ############################################################################## 75 | class LocalEnhancer(nn.Module): 76 | def __init__(self, input_nc, output_nc, ngf=32, n_downsample_global=3, n_blocks_global=9, 77 | n_local_enhancers=1, n_blocks_local=3, norm_layer=nn.BatchNorm2d, padding_type='reflect'): 78 | super(LocalEnhancer, self).__init__() 79 | self.n_local_enhancers = n_local_enhancers 80 | 81 | ###### global generator model ##### 82 | ngf_global = ngf * (2**n_local_enhancers) 83 | model_global = GlobalGenerator(input_nc, output_nc, ngf_global, n_downsample_global, n_blocks_global, norm_layer).model 84 | model_global = [model_global[i] for i in range(len(model_global)-3)] # get rid of final convolution layers 85 | self.model = nn.Sequential(*model_global) 86 | 87 | ###### local enhancer layers ##### 88 | for n in range(1, n_local_enhancers+1): 89 | ### downsample 90 | ngf_global = ngf * (2**(n_local_enhancers-n)) 91 | model_downsample = [nn.ReflectionPad2d(3), nn.Conv2d(input_nc, ngf_global, kernel_size=7, padding=0), 92 | norm_layer(ngf_global), nn.ReLU(True), 93 | nn.Conv2d(ngf_global, ngf_global * 2, kernel_size=3, stride=2, padding=1), 94 | norm_layer(ngf_global * 2), nn.ReLU(True)] 95 | ### residual blocks 96 | model_upsample = [] 97 | for i in range(n_blocks_local): 98 | model_upsample += [ResnetBlock(ngf_global * 2, padding_type=padding_type, norm_layer=norm_layer)] 99 | 100 | ### upsample 101 | model_upsample += [nn.ConvTranspose2d(ngf_global * 2, ngf_global, kernel_size=3, stride=2, padding=1, output_padding=1), 102 | norm_layer(ngf_global), nn.ReLU(True)] 103 | 104 | ### final convolution 105 | if n == n_local_enhancers: 106 | model_upsample += [nn.ReflectionPad2d(3), nn.Conv2d(ngf, output_nc, kernel_size=7, padding=0), nn.Tanh()] 107 | 108 | setattr(self, 'model'+str(n)+'_1', nn.Sequential(*model_downsample)) 109 | setattr(self, 'model'+str(n)+'_2', nn.Sequential(*model_upsample)) 110 | 111 | self.downsample = nn.AvgPool2d(3, stride=2, padding=[1, 1], count_include_pad=False) 112 | 113 | def forward(self, input): 114 | ### create input pyramid 115 | input_downsampled = [input] 116 | for i in range(self.n_local_enhancers): 117 | input_downsampled.append(self.downsample(input_downsampled[-1])) 118 | 119 | ### output at coarest level 120 | output_prev = self.model(input_downsampled[-1]) 121 | ### build up one layer at a time 122 | for n_local_enhancers in range(1, self.n_local_enhancers+1): 123 | model_downsample = getattr(self, 'model'+str(n_local_enhancers)+'_1') 124 | model_upsample = getattr(self, 'model'+str(n_local_enhancers)+'_2') 125 | input_i = input_downsampled[self.n_local_enhancers-n_local_enhancers] 126 | output_prev = model_upsample(model_downsample(input_i) + output_prev) 127 | return output_prev 128 | 129 | class NormalNet(nn.Module): 130 | def __init__(self, name='normalnet', input_nc=3, output_nc=3, ngf=64, n_downsampling=4, n_blocks=9, norm_layer=nn.BatchNorm2d, 131 | padding_type='reflect', last_op=nn.Tanh()): 132 | assert(n_blocks >= 0) 133 | super(NormalNet, self).__init__() 134 | self.name = name 135 | activation = nn.ReLU(True) 136 | 137 | model = [nn.ReflectionPad2d(3), nn.Conv2d(input_nc, ngf, kernel_size=7, padding=0), nn.GroupNorm(ngf//4, ngf), activation] 138 | ### downsample 139 | for i in range(n_downsampling): 140 | mult = 2**i 141 | model += [nn.Conv2d(ngf * mult, ngf * mult * 2, kernel_size=3, stride=2, padding=1), 142 | nn.GroupNorm(ngf * mult * 2 // 4, ngf * mult * 2), activation] 143 | 144 | ### resnet blocks 145 | mult = 2**n_downsampling 146 | for i in range(n_blocks): 147 | model += [ResnetBlock(ngf * mult, padding_type=padding_type, activation=activation, norm_layer=norm_layer)] 148 | 149 | ### upsample 150 | for i in range(n_downsampling): 151 | mult = 2**(n_downsampling - i) 152 | # model += [nn.Upsample(scale_factor=2), nn.Conv2d(ngf * mult, int(ngf * mult / 2), kernel_size=3, stride=1, padding=1), 153 | # nn.GroupNorm(int(ngf * mult / 2) // 4, int(ngf * mult / 2)), activation] 154 | model += [nn.ConvTranspose2d(ngf * mult, int(ngf * mult / 2), kernel_size=3, stride=2, padding=1, output_padding=1), 155 | nn.GroupNorm(int(ngf * mult / 2) // 4, int(ngf * mult / 2)), activation] 156 | model += [nn.ReflectionPad2d(3), nn.Conv2d(ngf, output_nc, kernel_size=7, padding=0)] 157 | if last_op is not None: 158 | model += [last_op] 159 | self.model = nn.Sequential(*model) 160 | 161 | def forward(self, in_x, label=None): 162 | res_list = [] 163 | if len(in_x.shape) == 5: 164 | for view in range(in_x.shape[1]): 165 | x = in_x[:, view, :, :, :] 166 | x = self.model(x) 167 | res_list.append(x.unsqueeze(1)) 168 | return torch.cat(res_list, dim=1) 169 | else: 170 | return [self.model(in_x)], 0, 0 171 | 172 | # Define a resnet block 173 | class ResnetBlock(nn.Module): 174 | def __init__(self, dim, padding_type, norm_layer, activation=nn.ReLU(True), use_dropout=False): 175 | super(ResnetBlock, self).__init__() 176 | self.conv_block = self.build_conv_block(dim, padding_type, norm_layer, activation, use_dropout) 177 | 178 | def build_conv_block(self, dim, padding_type, norm_layer, activation, use_dropout): 179 | conv_block = [] 180 | p = 0 181 | if padding_type == 'reflect': 182 | conv_block += [nn.ReflectionPad2d(1)] 183 | elif padding_type == 'replicate': 184 | conv_block += [nn.ReplicationPad2d(1)] 185 | elif padding_type == 'zero': 186 | p = 1 187 | else: 188 | raise NotImplementedError('padding [%s] is not implemented' % padding_type) 189 | 190 | conv_block += [nn.Conv2d(dim, dim, kernel_size=3, padding=p), 191 | nn.GroupNorm(dim // 4, dim), 192 | activation] 193 | if use_dropout: 194 | conv_block += [nn.Dropout(0.5)] 195 | 196 | p = 0 197 | if padding_type == 'reflect': 198 | conv_block += [nn.ReflectionPad2d(1)] 199 | elif padding_type == 'replicate': 200 | conv_block += [nn.ReplicationPad2d(1)] 201 | elif padding_type == 'zero': 202 | p = 1 203 | else: 204 | raise NotImplementedError('padding [%s] is not implemented' % padding_type) 205 | conv_block += [nn.Conv2d(dim, dim, kernel_size=3, padding=p), 206 | nn.GroupNorm(dim // 4, dim)] 207 | 208 | return nn.Sequential(*conv_block) 209 | 210 | def forward(self, x): 211 | out = x + self.conv_block(x) 212 | return out 213 | 214 | class Encoder(nn.Module): 215 | def __init__(self, input_nc, output_nc, ngf=32, n_downsampling=4, norm_layer=nn.BatchNorm2d): 216 | super(Encoder, self).__init__() 217 | self.output_nc = output_nc 218 | 219 | model = [nn.ReflectionPad2d(3), nn.Conv2d(input_nc, ngf, kernel_size=7, padding=0), 220 | norm_layer(ngf), nn.ReLU(True)] 221 | ### downsample 222 | for i in range(n_downsampling): 223 | mult = 2**i 224 | model += [nn.Conv2d(ngf * mult, ngf * mult * 2, kernel_size=3, stride=2, padding=1), 225 | norm_layer(ngf * mult * 2), nn.ReLU(True)] 226 | 227 | ### upsample 228 | for i in range(n_downsampling): 229 | mult = 2**(n_downsampling - i) 230 | model += [nn.ConvTranspose2d(ngf * mult, int(ngf * mult / 2), kernel_size=3, stride=2, padding=1, output_padding=1), 231 | norm_layer(int(ngf * mult / 2)), nn.ReLU(True)] 232 | 233 | model += [nn.ReflectionPad2d(3), nn.Conv2d(ngf, output_nc, kernel_size=7, padding=0), nn.Tanh()] 234 | self.model = nn.Sequential(*model) 235 | 236 | def forward(self, input, inst): 237 | outputs = self.model(input) 238 | 239 | # instance-wise average pooling 240 | outputs_mean = outputs.clone() 241 | inst_list = np.unique(inst.cpu().numpy().astype(int)) 242 | for i in inst_list: 243 | for b in range(input.size()[0]): 244 | indices = (inst[b:b+1] == int(i)).nonzero() # n x 4 245 | for j in range(self.output_nc): 246 | output_ins = outputs[indices[:,0] + b, indices[:,1] + j, indices[:,2], indices[:,3]] 247 | mean_feat = torch.mean(output_ins).expand_as(output_ins) 248 | outputs_mean[indices[:,0] + b, indices[:,1] + j, indices[:,2], indices[:,3]] = mean_feat 249 | return outputs_mean -------------------------------------------------------------------------------- /lib/model/PositionEncoder.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | 4 | # Positional encoding 5 | class Embedder: 6 | def __init__(self, **kwargs): 7 | self.kwargs = kwargs 8 | self.create_embedding_fn() 9 | 10 | def create_embedding_fn(self): 11 | embed_fns = [] 12 | d = self.kwargs['input_dims'] 13 | out_dim = 0 14 | if self.kwargs['include_input']: 15 | embed_fns.append(lambda x : x) 16 | out_dim += d 17 | 18 | max_freq = self.kwargs['max_freq_log2'] 19 | N_freqs = self.kwargs['num_freqs'] 20 | 21 | if self.kwargs['log_sampling']: 22 | freq_bands = 2.**torch.linspace(0., max_freq, steps=N_freqs) 23 | else: 24 | freq_bands = torch.linspace(2.**0., 2.**max_freq, steps=N_freqs) 25 | 26 | for freq in freq_bands: 27 | for p_fn in self.kwargs['periodic_fns']: 28 | embed_fns.append(lambda x, p_fn=p_fn, freq=freq : p_fn(x * freq)) 29 | out_dim += d 30 | 31 | self.embed_fns = embed_fns 32 | self.out_dim = out_dim 33 | 34 | def embed(self, inputs): 35 | return torch.cat([fn(inputs) for fn in self.embed_fns], -1) 36 | 37 | def get_embedder(multires, i=0): 38 | if i == -1: 39 | return nn.Identity(), 3 40 | 41 | embed_kwargs = { 42 | 'include_input' : True, 43 | 'input_dims' : 1, 44 | 'max_freq_log2' : multires-1, 45 | 'num_freqs' : multires, 46 | 'log_sampling' : True, 47 | 'periodic_fns' : [torch.sin, torch.cos], 48 | } 49 | 50 | embedder_obj = Embedder(**embed_kwargs) 51 | embed = lambda x, eo=embedder_obj : eo.embed(x) 52 | return embed, embedder_obj.out_dim -------------------------------------------------------------------------------- /lib/model/ResNet3d.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | 4 | class IdentityBlock(nn.Module): 5 | def __init__(self, channels, filters): 6 | super(IdentityBlock, self).__init__() 7 | self.channels = channels 8 | self.filters = filters 9 | self.net = nn.Sequential( 10 | nn.Conv3d(channels, filters[0], 1), 11 | nn.BatchNorm3d(filters[0]), 12 | nn.ReLU(True), 13 | nn.Conv3d(filters[0], filters[1], 3, padding=1), 14 | nn.BatchNorm3d(filters[1]), 15 | nn.ReLU(True), 16 | nn.Conv3d(filters[1], filters[2], 1), 17 | nn.BatchNorm3d(filters[2]) 18 | ) 19 | self.channel_net = nn.Conv3d(channels, filters[2], 1, 1, 0) 20 | 21 | self.relu = nn.ReLU(True) 22 | 23 | def forward(self, x): 24 | y = self.net(x) 25 | if self.channels != self.filters[2]: 26 | y = self.channel_net(x) + y 27 | else: 28 | y = x + y 29 | y = self.relu(y) 30 | 31 | return y 32 | 33 | 34 | class ConvBlock(nn.Module): 35 | def __init__(self, channels, filters, stride=2): 36 | super(ConvBlock, self).__init__() 37 | self.channels = channels 38 | self.filters = filters 39 | self.net = nn.Sequential( 40 | nn.Conv3d(channels, filters[0], 1, stride=stride), 41 | nn.BatchNorm3d(filters[0]), 42 | nn.ReLU(True), 43 | nn.Conv3d(filters[0], filters[1], 3, padding=1), 44 | nn.BatchNorm3d(filters[1]), 45 | nn.ReLU(True), 46 | nn.Conv3d(filters[1], filters[2], 1), 47 | nn.BatchNorm3d(filters[2]) 48 | ) 49 | self.downsample = nn.Sequential( 50 | nn.Conv3d(channels, filters[2], 1, stride=stride), 51 | nn.BatchNorm3d(filters[2]) 52 | ) 53 | self.relu = nn.ReLU(True) 54 | 55 | def forward(self, x): 56 | y = self.net(x) 57 | x = self.downsample(x) 58 | y = x + y 59 | y = self.relu(y) 60 | 61 | return y 62 | 63 | class ResNet3d(nn.Module): 64 | def __init__(self): 65 | super(ResNet3d, self).__init__() 66 | self.stage1 = nn.Sequential( 67 | ConvBlock(1, [16, 16, 32]), 68 | IdentityBlock(32, [16, 16, 32]) 69 | ) 70 | self.stage_d = nn.Sequential( 71 | ConvBlock(32, [16, 16, 32]), 72 | nn.BatchNorm3d(32), 73 | nn.Conv3d(32, 64, 1) 74 | ) 75 | self.stage2 = nn.Sequential( 76 | ConvBlock(32, [16, 16, 64]), 77 | IdentityBlock(64, [32, 32, 64]), 78 | IdentityBlock(64, [32, 32, 64]) 79 | ) 80 | # self.stage3 = nn.Sequential( 81 | # ConvBlock(32, [16, 16, 64]), 82 | # IdentityBlock(64, [32, 32, 64]), 83 | # IdentityBlock(64, [32, 32, 64]), 84 | # IdentityBlock(64, [32, 32, 64]) 85 | # ) 86 | def forward(self, x): 87 | x1 = self.stage1(x) 88 | x2 = self.stage2(x1) 89 | # x3 = self.stage3(x2) 90 | return x2 -------------------------------------------------------------------------------- /lib/model/SurfaceClassifier.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | import torch.nn.functional as F 4 | 5 | 6 | class SurfaceClassifier(nn.Module): 7 | def __init__(self, filter_channels, num_views=1, no_residual=True, last_op=None): 8 | super(SurfaceClassifier, self).__init__() 9 | 10 | self.filters = [] 11 | self.num_views = num_views 12 | self.no_residual = no_residual 13 | filter_channels = filter_channels 14 | self.last_op = last_op 15 | 16 | if self.no_residual: 17 | for l in range(0, len(filter_channels) - 1): 18 | self.filters.append(nn.Conv1d( 19 | filter_channels[l], 20 | filter_channels[l + 1], 21 | 1)) 22 | self.add_module("conv%d" % l, self.filters[l]) 23 | else: 24 | for l in range(0, len(filter_channels) - 1): 25 | if 0 != l: 26 | self.filters.append( 27 | nn.Conv1d( 28 | filter_channels[l] + filter_channels[0], 29 | filter_channels[l + 1], 30 | 1)) 31 | else: 32 | self.filters.append(nn.Conv1d( 33 | filter_channels[l], 34 | filter_channels[l + 1], 35 | 1)) 36 | 37 | self.add_module("conv%d" % l, self.filters[l]) 38 | 39 | def forward(self, feature, xy=None, ret_list=False): 40 | ''' 41 | 42 | :param feature: list of [BxC_inxN] tensors of image features 43 | :param xy: [Bx3xN] tensor of (x,y) coodinates in the image plane 44 | :return: [BxC_outxN] tensor of features extracted at the coordinates 45 | ''' 46 | y = feature 47 | tmpy = feature 48 | y_list = [] 49 | for i, f in enumerate(self.filters): 50 | if self.no_residual: 51 | y = self._modules['conv' + str(i)](y) 52 | else: 53 | y = self._modules['conv' + str(i)]( 54 | y if i == 0 55 | else torch.cat([y, tmpy], 1) 56 | ) 57 | if i != len(self.filters) - 1: 58 | y = F.leaky_relu(y) 59 | y_list.append(y) 60 | if self.last_op: 61 | y = self.last_op(y) 62 | if ret_list: 63 | return y_list 64 | else: 65 | return y 66 | -------------------------------------------------------------------------------- /lib/model/__init__.py: -------------------------------------------------------------------------------- 1 | import cv2 2 | from .DMCNet import DMCNet 3 | from .NormalNet import NormalNet 4 | -------------------------------------------------------------------------------- /lib/model/attention.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | import torch.nn.functional as F 4 | import numpy as np 5 | 6 | 7 | class ScaledDotProductAttention(nn.Module): 8 | def __init__(self, temperature, attn_dropout=0.1): 9 | super().__init__() 10 | self.temperature = temperature 11 | self.dropout = nn.Dropout(attn_dropout) 12 | 13 | def forward(self, q, k, v): 14 | """ 15 | Scaled Dot-Product Attention 16 | out = softmax(QK^T/temperature)V 17 | :param q: (bs, lenq, d_k) 18 | :param k: (bs, lenv, d_k) 19 | :param v: (bs, lenv, d_v) 20 | :param mask: None 21 | :return: 22 | """ 23 | attn = torch.matmul(q / self.temperature, k.transpose(2, 3)) 24 | attn = self.dropout(F.softmax(attn, dim=-1)) 25 | output = torch.matmul(attn, v) 26 | return output, attn 27 | 28 | 29 | class MultiHeadAttention(nn.Module): 30 | """ 31 | Multi-Head Attention module 32 | """ 33 | 34 | def __init__(self, n_head, d_model, d_k, d_v, dropout=0.1, norm=False): 35 | super().__init__() 36 | 37 | self.n_head = n_head 38 | self.d_k = d_k 39 | self.d_v = d_v 40 | 41 | self.w_qs = nn.Linear(d_model, n_head * d_k, bias=False) 42 | self.w_ks = nn.Linear(d_model, n_head * d_k, bias=False) 43 | self.w_vs = nn.Linear(d_model, n_head * d_v, bias=False) 44 | self.fc = nn.Linear(n_head * d_v, d_model, bias=False) 45 | 46 | self.attention = ScaledDotProductAttention(temperature=d_k ** 0.5) 47 | 48 | self.dropout = nn.Dropout(dropout) 49 | # self.bn = nn.BatchNorm1d(d_model) 50 | self.norm = norm 51 | if self.norm: 52 | self.layer_norm = nn.LayerNorm(d_model, eps=1e-6) 53 | 54 | def forward(self, q, k, v): 55 | """ 56 | :param q: (bs, lenq, d_model) 57 | :param k: (bs, lenv, d_model) 58 | :param v: (bs, lenv, d_model) 59 | :param mask: 60 | :return: (bs, lenq, d_model) 61 | """ 62 | d_k, d_v, n_head = self.d_k, self.d_v, self.n_head 63 | bs, len_q, len_k, len_v = q.size(0), q.size(1), k.size(1), v.size(1) 64 | 65 | residual = q 66 | 67 | # Pass through the pre-attention projection: bs x lq x (n*dv) 68 | # Separate different heads: bs x lq x n x dv 69 | q = self.w_qs(q).view(bs, len_q, n_head, d_k) 70 | k = self.w_ks(k).view(bs, len_k, n_head, d_k) 71 | v = self.w_vs(v).view(bs, len_v, n_head, d_v) 72 | 73 | # Transpose for attention dot product: bs x n x lq x dv 74 | q, k, v = q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2) 75 | q, attn = self.attention(q, k, v) 76 | 77 | # Transpose to move the head dimension back: bs x lq x n x dv 78 | # Combine the last two dimensions to concatenate all the heads together: b sx lq x (n*dv) 79 | q = q.transpose(1, 2).contiguous().view(bs, len_q, -1) 80 | q = self.dropout(self.fc(q)) 81 | q += residual 82 | # q = self.bn(q.permute((0, 2, 1))).permute((0, 2, 1)) 83 | # q = self.layer_norm(q) 84 | if self.norm: 85 | q = self.layer_norm(q) 86 | return q, attn 87 | 88 | 89 | class PositionwiseFeedForward(nn.Module): 90 | """ A two-feed-forward-layer module """ 91 | 92 | def __init__(self, d_in, d_hid, dropout=0.1, norm=False): 93 | super().__init__() 94 | self.w_1 = nn.Linear(d_in, d_hid) # position-wise 95 | self.w_2 = nn.Linear(d_hid, d_in) # position-wise 96 | self.norm = norm 97 | if self.norm: 98 | self.layer_norm = nn.LayerNorm(d_in, eps=1e-6) 99 | # self.bn = nn.BatchNorm1d(d_in) 100 | self.dropout = nn.Dropout(dropout) 101 | 102 | def forward(self, x): 103 | residual = x 104 | x = self.w_2(F.relu(self.w_1(x))) 105 | x = self.dropout(x) 106 | x += residual 107 | if self.norm: 108 | # x = self.bn(x.permute((0, 2, 1))).permute((0, 2, 1)) 109 | x = self.layer_norm(x) 110 | return x 111 | 112 | class PositionalEncoding(nn.Module): 113 | 114 | def __init__(self, d_hid, n_position=200): 115 | super(PositionalEncoding, self).__init__() 116 | 117 | # Not a parameter 118 | self.register_buffer('pos_table', self._get_sinusoid_encoding_table(n_position, d_hid)) 119 | 120 | def _get_sinusoid_encoding_table(self, n_position, d_hid): 121 | ''' Sinusoid position encoding table ''' 122 | # TODO: make it with torch instead of numpy 123 | 124 | def get_position_angle_vec(position): 125 | return [position / np.power(10000, 2 * (hid_j // 2) / d_hid) for hid_j in range(d_hid)] 126 | 127 | sinusoid_table = np.array([get_position_angle_vec(pos_i) for pos_i in range(n_position)]) 128 | sinusoid_table[:, 0::2] = np.sin(sinusoid_table[:, 0::2]) # dim 2i 129 | sinusoid_table[:, 1::2] = np.cos(sinusoid_table[:, 1::2]) # dim 2i+1 130 | 131 | return torch.FloatTensor(sinusoid_table).unsqueeze(0) 132 | 133 | def forward(self, x): 134 | return x + self.pos_table[:, :x.size(1)].clone().detach() 135 | 136 | class EncoderLayer(nn.Module): 137 | """ 138 | multi-attention + position feed forward 139 | """ 140 | 141 | def __init__(self, d_model, d_inner, n_head, d_k, d_v, dropout=0.1, norm=False): 142 | super(EncoderLayer, self).__init__() 143 | self.slf_attn = MultiHeadAttention(n_head, d_model, d_k, d_v, dropout=dropout, norm=norm) 144 | self.pos_ffn = PositionwiseFeedForward(d_model, d_inner, dropout=dropout, norm=norm) 145 | 146 | def forward(self, enc_input): 147 | enc_output, enc_slf_attn = self.slf_attn( 148 | enc_input, enc_input, enc_input) 149 | enc_output = self.pos_ffn(enc_output) 150 | return enc_output, enc_slf_attn 151 | 152 | 153 | class DecoderLayer(nn.Module): 154 | """ Compose with three layers """ 155 | 156 | def __init__(self, d_model, d_inner, n_head, d_k, d_v, dropout=0.1, norm=False): 157 | super(DecoderLayer, self).__init__() 158 | self.slf_attn = MultiHeadAttention(n_head, d_model, d_k, d_v, dropout=dropout, norm=norm) 159 | self.enc_attn = MultiHeadAttention(n_head, d_model, d_k, d_v, dropout=dropout, norm=norm) 160 | self.pos_ffn = PositionwiseFeedForward(d_model, d_inner, dropout=dropout, norm=norm) 161 | 162 | def forward(self, dec_input, enc_output): 163 | dec_output, dec_slf_attn = self.slf_attn(dec_input, dec_input, dec_input) 164 | dec_output, dec_enc_attn = self.enc_attn(dec_output, enc_output, enc_output) 165 | dec_output = self.pos_ffn(dec_output) 166 | return dec_output, dec_slf_attn, dec_enc_attn 167 | 168 | 169 | class Encoder(nn.Module): 170 | """ A encoder model with self attention mechanism. """ 171 | 172 | def __init__(self, n_layers, n_head, d_k, d_v, d_model, d_inner, dropout=0.1, norm=False, n_position=200): 173 | 174 | super().__init__() 175 | 176 | # self.src_word_emb = nn.Embedding(n_src_vocab, d_word_vec, padding_idx=pad_idx) 177 | self.position_enc = PositionalEncoding(d_model, n_position=n_position) 178 | self.dropout = nn.Dropout(p=dropout) 179 | self.layer_stack = nn.ModuleList([ 180 | EncoderLayer(d_model, d_inner, n_head, d_k, d_v, dropout=dropout, norm=norm) 181 | for _ in range(n_layers)]) 182 | self.norm = norm 183 | if self.norm: 184 | self.layer_norm = nn.LayerNorm(d_model, eps=1e-6) 185 | # self.bn = nn.BatchNorm1d(d_model) 186 | 187 | def forward(self, feats_embedding, return_attns=False): 188 | """ 189 | :param feats_embedding: (bs, num_views, dim) 190 | :param return_attns: 191 | :return: 192 | """ 193 | enc_slf_attn_list = [] 194 | # enc_output = self.dropout(self.position_enc(feats_embedding)) 195 | enc_output = feats_embedding 196 | if self.norm: 197 | enc_output = self.layer_norm(enc_output) 198 | for enc_layer in self.layer_stack: 199 | enc_output, enc_slf_attn = enc_layer(enc_output) 200 | enc_slf_attn_list += [enc_slf_attn] if return_attns else [] 201 | 202 | if return_attns: 203 | return enc_output, enc_slf_attn_list 204 | return enc_output, 205 | 206 | 207 | class Decoder(nn.Module): 208 | """ A decoder model with self attention mechanism. """ 209 | 210 | def __init__(self, n_layers, n_head, d_k, d_v, d_model, d_inner, dropout=0.1, norm=False, n_position=200): 211 | 212 | super().__init__() 213 | # self.trg_word_emb = nn.Embedding(n_trg_vocab, d_word_vec, padding_idx=pad_idx) 214 | self.position_enc = PositionalEncoding(d_model, n_position=n_position) 215 | self.dropout = nn.Dropout(p=dropout) 216 | self.layer_stack = nn.ModuleList([ 217 | DecoderLayer(d_model, d_inner, n_head, d_k, d_v, dropout=dropout, norm=norm) 218 | for _ in range(n_layers)]) 219 | self.norm = norm 220 | if self.norm: 221 | self.layer_norm = nn.LayerNorm(d_model, eps=1e-6) 222 | # self.layer_norm = nn.LayerNorm(d_model, eps=1e-6) 223 | # self.bn = nn.BatchNorm1d(d_model) 224 | 225 | def forward(self, dec_input, enc_output, return_attns=False): 226 | """ 227 | :param dec_input: (bs, 1, dim) 228 | :param enc_output: (bs, num_views, dim) 229 | :param return_attns: 230 | :return: 231 | """ 232 | dec_output = dec_input 233 | # dec_output = self.dropout(self.position_enc(dec_output)) 234 | if self.norm: 235 | dec_output = self.layer_norm(dec_output) 236 | dec_slf_attn_list, dec_enc_attn_list = [], [] 237 | for dec_layer in self.layer_stack: 238 | dec_output, dec_slf_attn, dec_enc_attn = dec_layer( 239 | dec_output, enc_output) 240 | dec_slf_attn_list += [dec_slf_attn] if return_attns else [] 241 | dec_enc_attn_list += [dec_enc_attn] if return_attns else [] 242 | 243 | if return_attns: 244 | return dec_output, dec_slf_attn_list, dec_enc_attn_list 245 | return dec_output, 246 | 247 | 248 | class Transformer(nn.Module): 249 | ''' A sequence to sequence model with attention mechanism. ''' 250 | 251 | def __init__(self, d_model=512, d_inner=2048, 252 | n_layers=6, n_head=8, d_k=64, d_v=64, dropout=0.1): 253 | 254 | super().__init__() 255 | self.encoder = Encoder(d_model=d_model, d_inner=d_inner, n_layers=n_layers, 256 | n_head=n_head, d_k=d_k, d_v=d_v, dropout=dropout) 257 | 258 | self.decoder = Decoder(d_model=d_model, d_inner=d_inner, n_layers=n_layers, 259 | n_head=n_head, d_k=d_k, d_v=d_v, dropout=dropout) 260 | for p in self.parameters(): 261 | if p.dim() > 1: 262 | nn.init.xavier_uniform_(p) 263 | 264 | def forward(self, src_feats, trg_feats): 265 | """ 266 | :param src_feats: (bs, num_views, dim) 267 | :param trg_feats: (bs, 1, dim) 268 | :return: fused feats: (bs, 1, dim) 269 | """ 270 | enc_output, *_ = self.encoder(src_feats) 271 | dec_output, *_ = self.decoder(trg_feats, enc_output) 272 | return dec_output -------------------------------------------------------------------------------- /lib/net_util.py: -------------------------------------------------------------------------------- 1 | import torch 2 | from torch.nn import init 3 | import torch.nn as nn 4 | import torch.nn.functional as F 5 | import functools 6 | import cv2 7 | from PIL import Image 8 | from tqdm import tqdm 9 | import numpy as np 10 | 11 | from .mesh_util import * 12 | from .sample_util import * 13 | from .geometry import index 14 | 15 | 16 | def check_tensor(t, name='test'): 17 | if torch.sum(torch.isnan(t) + torch.isinf(t)) > 0: 18 | print(name, torch.sum(torch.isnan(t) + torch.isinf(t))) 19 | exit(0) 20 | 21 | def reshape_multiview_tensors(image_tensor, calib_tensor): 22 | # Careful here! Because we put single view and multiview together, 23 | # the returned tensor.shape is 5-dim: [B, num_views, C, W, H] 24 | # So we need to convert it back to 4-dim [B*num_views, C, W, H] 25 | # Don't worry classifier will handle multi-view cases 26 | image_tensor = image_tensor.view( 27 | image_tensor.shape[0] * image_tensor.shape[1], 28 | image_tensor.shape[2], 29 | image_tensor.shape[3], 30 | image_tensor.shape[4] 31 | ) 32 | calib_tensor = calib_tensor.view( 33 | calib_tensor.shape[0] * calib_tensor.shape[1], 34 | calib_tensor.shape[2], 35 | calib_tensor.shape[3] 36 | ) 37 | 38 | return image_tensor, calib_tensor 39 | 40 | 41 | def reshape_sample_tensor(sample_tensor, num_views): 42 | if num_views == 1: 43 | return sample_tensor 44 | # Need to repeat sample_tensor along the batch dim num_views times 45 | sample_tensor = sample_tensor.unsqueeze(dim=1) 46 | sample_tensor = sample_tensor.repeat(1, num_views, 1, 1) 47 | sample_tensor = sample_tensor.view( 48 | sample_tensor.shape[0] * sample_tensor.shape[1], 49 | sample_tensor.shape[2], 50 | sample_tensor.shape[3] 51 | ) 52 | return sample_tensor 53 | 54 | 55 | def adjust_learning_rate(optimizer, epoch, lr, schedule, gamma): 56 | """Sets the learning rate to the initial LR decayed by schedule""" 57 | if epoch in schedule: 58 | lr *= gamma 59 | for param_group in optimizer.param_groups: 60 | param_group['lr'] = lr 61 | return lr 62 | 63 | 64 | def compute_acc(pred, gt, thresh=0.5): 65 | ''' 66 | return: 67 | IOU, precision, and recall 68 | ''' 69 | with torch.no_grad(): 70 | vol_pred = pred > thresh 71 | vol_gt = gt > thresh 72 | 73 | union = vol_pred | vol_gt 74 | inter = vol_pred & vol_gt 75 | 76 | true_pos = inter.sum().float() 77 | 78 | union = union.sum().float() 79 | if union == 0: 80 | union = 1 81 | vol_pred = vol_pred.sum().float() 82 | if vol_pred == 0: 83 | vol_pred = 1 84 | vol_gt = vol_gt.sum().float() 85 | if vol_gt == 0: 86 | vol_gt = 1 87 | return true_pos / union, true_pos / vol_pred, true_pos / vol_gt 88 | 89 | def conv3x3(in_planes, out_planes, strd=1, padding=1, bias=False): 90 | "3x3 convolution with padding" 91 | return nn.Conv2d(in_planes, out_planes, kernel_size=3, 92 | stride=strd, padding=padding, bias=bias) 93 | 94 | def init_weights(net, init_type='normal', init_gain=0.02): 95 | """Initialize network weights. 96 | 97 | Parameters: 98 | net (network) -- network to be initialized 99 | init_type (str) -- the name of an initialization method: normal | xavier | kaiming | orthogonal 100 | init_gain (float) -- scaling factor for normal, xavier and orthogonal. 101 | 102 | We use 'normal' in the original pix2pix and CycleGAN paper. But xavier and kaiming might 103 | work better for some applications. Feel free to try yourself. 104 | """ 105 | 106 | def init_func(m): # define the initialization function 107 | classname = m.__class__.__name__ 108 | if hasattr(m, 'weight') and (classname.find('Conv') != -1 or classname.find('Linear') != -1): 109 | if init_type == 'normal': 110 | init.normal_(m.weight.data, 0.0, init_gain) 111 | elif init_type == 'xavier': 112 | init.xavier_normal_(m.weight.data, gain=init_gain) 113 | elif init_type == 'kaiming': 114 | init.kaiming_normal_(m.weight.data, a=0, mode='fan_in') 115 | elif init_type == 'orthogonal': 116 | init.orthogonal_(m.weight.data, gain=init_gain) 117 | else: 118 | raise NotImplementedError('initialization method [%s] is not implemented' % init_type) 119 | if hasattr(m, 'bias') and m.bias is not None: 120 | init.constant_(m.bias.data, 0.0) 121 | elif classname.find( 122 | 'BatchNorm2d') != -1: # BatchNorm Layer's weight is not a matrix; only normal distribution applies. 123 | init.normal_(m.weight.data, 1.0, init_gain) 124 | init.constant_(m.bias.data, 0.0) 125 | 126 | print('initialize network with %s' % init_type) 127 | net.apply(init_func) # apply the initialization function 128 | 129 | 130 | def init_net(net, init_type='normal', init_gain=0.02, gpu_ids=[]): 131 | """Initialize a network: 1. register CPU/GPU device (with multi-GPU support); 2. initialize the network weights 132 | Parameters: 133 | net (network) -- the network to be initialized 134 | init_type (str) -- the name of an initialization method: normal | xavier | kaiming | orthogonal 135 | gain (float) -- scaling factor for normal, xavier and orthogonal. 136 | gpu_ids (int list) -- which GPUs the network runs on: e.g., 0,1,2 137 | 138 | Return an initialized network. 139 | """ 140 | if len(gpu_ids) > 0: 141 | assert (torch.cuda.is_available()) 142 | net.to(gpu_ids[0]) 143 | net = torch.nn.DataParallel(net, gpu_ids) # multi-GPUs 144 | init_weights(net, init_type, init_gain=init_gain) 145 | return net 146 | 147 | 148 | def imageSpaceRotation(xy, rot): 149 | ''' 150 | args: 151 | xy: (B, 2, N) input 152 | rot: (B, 2) x,y axis rotation angles 153 | 154 | rotation center will be always image center (other rotation center can be represented by additional z translation) 155 | ''' 156 | disp = rot.unsqueeze(2).sin().expand_as(xy) 157 | return (disp * xy).sum(dim=1) 158 | 159 | 160 | def cal_gradient_penalty(netD, real_data, fake_data, device, type='mixed', constant=1.0, lambda_gp=10.0): 161 | """Calculate the gradient penalty loss, used in WGAN-GP paper https://arxiv.org/abs/1704.00028 162 | 163 | Arguments: 164 | netD (network) -- discriminator network 165 | real_data (tensor array) -- real images 166 | fake_data (tensor array) -- generated images from the generator 167 | device (str) -- GPU / CPU: from torch.device('cuda:{}'.format(self.gpu_ids[0])) if self.gpu_ids else torch.device('cpu') 168 | type (str) -- if we mix real and fake data or not [real | fake | mixed]. 169 | constant (float) -- the constant used in formula ( | |gradient||_2 - constant)^2 170 | lambda_gp (float) -- weight for this loss 171 | 172 | Returns the gradient penalty loss 173 | """ 174 | if lambda_gp > 0.0: 175 | if type == 'real': # either use real images, fake images, or a linear interpolation of two. 176 | interpolatesv = real_data 177 | elif type == 'fake': 178 | interpolatesv = fake_data 179 | elif type == 'mixed': 180 | alpha = torch.rand(real_data.shape[0], 1) 181 | alpha = alpha.expand(real_data.shape[0], real_data.nelement() // real_data.shape[0]).contiguous().view( 182 | *real_data.shape) 183 | alpha = alpha.to(device) 184 | interpolatesv = alpha * real_data + ((1 - alpha) * fake_data) 185 | else: 186 | raise NotImplementedError('{} not implemented'.format(type)) 187 | interpolatesv.requires_grad_(True) 188 | disc_interpolates = netD(interpolatesv) 189 | gradients = torch.autograd.grad(outputs=disc_interpolates, inputs=interpolatesv, 190 | grad_outputs=torch.ones(disc_interpolates.size()).to(device), 191 | create_graph=True, retain_graph=True, only_inputs=True) 192 | gradients = gradients[0].view(real_data.size(0), -1) # flat the data 193 | gradient_penalty = (((gradients + 1e-16).norm(2, dim=1) - constant) ** 2).mean() * lambda_gp # added eps 194 | return gradient_penalty, gradients 195 | else: 196 | return 0.0, None 197 | 198 | def get_norm_layer(norm_type='instance'): 199 | """Return a normalization layer 200 | Parameters: 201 | norm_type (str) -- the name of the normalization layer: batch | instance | none 202 | For BatchNorm, we use learnable affine parameters and track running statistics (mean/stddev). 203 | For InstanceNorm, we do not use learnable affine parameters. We do not track running statistics. 204 | """ 205 | if norm_type == 'batch': 206 | norm_layer = functools.partial(nn.BatchNorm2d, affine=True, track_running_stats=True) 207 | elif norm_type == 'instance': 208 | norm_layer = functools.partial(nn.InstanceNorm2d, affine=False, track_running_stats=False) 209 | elif norm_type == 'group': 210 | norm_layer = functools.partial(nn.GroupNorm, 32) 211 | elif norm_type == 'none': 212 | norm_layer = None 213 | else: 214 | raise NotImplementedError('normalization layer [%s] is not found' % norm_type) 215 | return norm_layer 216 | 217 | class Flatten(nn.Module): 218 | def forward(self, input): 219 | return input.view(input.size(0), -1) 220 | 221 | class ConvBlock(nn.Module): 222 | def __init__(self, in_planes, out_planes, norm='batch', base=64): 223 | super(ConvBlock, self).__init__() 224 | self.conv1 = conv3x3(in_planes, int(out_planes / 2)) 225 | self.conv2 = conv3x3(int(out_planes / 2), int(out_planes / 4)) 226 | self.conv3 = conv3x3(int(out_planes / 4), int(out_planes / 4)) 227 | 228 | if norm == 'batch': 229 | self.bn1 = nn.BatchNorm2d(in_planes) 230 | self.bn2 = nn.BatchNorm2d(int(out_planes / 2)) 231 | self.bn3 = nn.BatchNorm2d(int(out_planes / 4)) 232 | self.bn4 = nn.BatchNorm2d(in_planes) 233 | elif norm == 'group': 234 | self.bn1 = nn.GroupNorm(base // 2, in_planes) 235 | self.bn2 = nn.GroupNorm(base // 2, int(out_planes / 2)) 236 | self.bn3 = nn.GroupNorm(base // 2, int(out_planes / 4)) 237 | self.bn4 = nn.GroupNorm(base // 2, in_planes) 238 | 239 | if in_planes != out_planes: 240 | self.downsample = nn.Sequential( 241 | self.bn4, 242 | nn.ReLU(True), 243 | nn.Conv2d(in_planes, out_planes, 244 | kernel_size=1, stride=1, bias=False), 245 | ) 246 | else: 247 | self.downsample = None 248 | 249 | def forward(self, x): 250 | residual = x 251 | 252 | out1 = self.bn1(x) 253 | out1 = F.relu(out1, True) 254 | out1 = self.conv1(out1) 255 | 256 | out2 = self.bn2(out1) 257 | out2 = F.relu(out2, True) 258 | out2 = self.conv2(out2) 259 | 260 | out3 = self.bn3(out2) 261 | out3 = F.relu(out3, True) 262 | out3 = self.conv3(out3) 263 | 264 | out3 = torch.cat((out1, out2, out3), 1) 265 | 266 | if self.downsample is not None: 267 | residual = self.downsample(residual) 268 | 269 | out3 += residual 270 | 271 | return out3 272 | -------------------------------------------------------------------------------- /lib/options.py: -------------------------------------------------------------------------------- 1 | import argparse 2 | import configargparse 3 | import os 4 | 5 | def print_options(self, opt): 6 | message = '' 7 | message += '----------------- Options ---------------\n' 8 | for k, v in sorted(vars(opt).items()): 9 | comment = '' 10 | default = self.parser.get_default(k) 11 | if v != default: 12 | comment = '\t[default: %s]' % str(default) 13 | message += '{:>25}: {:<30}{}\n'.format(str(k), str(v), comment) 14 | message += '----------------- End -------------------' 15 | print(message) 16 | 17 | def parse_config(argv=None): 18 | arg_formatter = configargparse.ArgumentDefaultsHelpFormatter 19 | cfg_parser = configargparse.DefaultConfigFileParser 20 | description = 'project' 21 | parser = configargparse.ArgParser(formatter_class=arg_formatter, 22 | config_file_parser_class=cfg_parser, 23 | description=description, 24 | prog='deepmulticap') 25 | 26 | # general settings 27 | parser.add_argument('--config', is_config_file=True, help='config file path') 28 | parser.add_argument('--name', type=str, default='deepmulticap', help='name of a model/experiment.') 29 | 30 | # Datasets related 31 | parser.add_argument('--log_path', type=str, default='./train_log') 32 | parser.add_argument('--dataroot', type=str, default='./data', 33 | help='path to images (data folder)') 34 | parser.add_argument('--obj_path', type=str) 35 | parser.add_argument('--smpl_path', type=str) 36 | parser.add_argument('--tex_path', type=str) 37 | parser.add_argument('--loadSize', type=int, default=512, help='load size of input image') 38 | parser.add_argument('--b_min', nargs='+', default=[-3, -2, -3], type=float) 39 | parser.add_argument('--b_max', nargs='+', default=[3, 14, 3], type=float) 40 | parser.add_argument('--smpl_faces', type=str, default='/media/data1/shaoruizhi/Multiview_Pair/lib/data/smplx_fine.obj') 41 | 42 | # Experiment related 43 | parser.add_argument('--num_views', type=int, default=1, help='How many views to use for multiview network.') 44 | parser.add_argument('--random_multiview', action='store_true', help='Select random multiview combination.') 45 | parser.add_argument('--random_fix_view', action='store_true', help='Select random multiview combination.') 46 | 47 | # Training related 48 | parser.add_argument('--gpu_ids', type=str, default='0', help='gpu ids: e.g. 0 0,1,2, 0,2, -1 for CPU mode') 49 | 50 | parser.add_argument('--num_threads', default=1, type=int, help='# sthreads for loading data') 51 | parser.add_argument('--serial_batches', action='store_true', 52 | help='if true, takes images in order to make batches, otherwise takes them randomly') 53 | parser.add_argument('--pin_memory', action='store_true', help='pin_memory') 54 | 55 | parser.add_argument('--batch_size', type=int, default=2, help='input batch size') 56 | parser.add_argument('--learning_rate', type=float, default=1e-3, help='adam learning rate') 57 | parser.add_argument('--num_epoch', type=int, default=1000, help='num epoch to train') 58 | 59 | 60 | parser.add_argument('--freq_plot', type=int, default=10, help='freqency of the save_checkpoints') 61 | parser.add_argument('--freq_save', type=int, default=50, help='freqency of the save_checkpoints') 62 | parser.add_argument('--freq_save_ply', type=int, default=100, help='freqency of the save ply') 63 | parser.add_argument('--freq_normal_show', type=int, default=1000, help='freqency of the save ply') 64 | 65 | # Testing related 66 | parser.add_argument('--resolution', type=int, default=256, help='# of grid in mesh reconstruction') 67 | parser.add_argument('--mc_threshold', type=float, default=0.5, help='marching cube threshold') 68 | 69 | # Sampling related 70 | parser.add_argument('--sigma', type=float, default=0.25, help='perturbation standard deviation for positions') 71 | parser.add_argument('--num_sample_inout', type=int, default=5000, help='# of sampling points') 72 | parser.add_argument('--visibility_sample', action='store_true') 73 | 74 | # Model related 75 | # General 76 | parser.add_argument('--norm', type=str, default='group', 77 | help='instance normalization or batch normalization or group normalization') 78 | 79 | # hg filter specify 80 | parser.add_argument('--num_stack', type=int, default=4, help='# of hourglass') 81 | parser.add_argument('--num_hourglass', type=int, default=2, help='# of stacked layer of hourglass') 82 | parser.add_argument('--fine_num_stack', type=int, default=1, help='# of hourglass') 83 | parser.add_argument('--fine_num_hourglass', type=int, default=2, help='# of stacked layer of hourglass') 84 | parser.add_argument('--fine_hourglass_dim', type=int, default=32, help='feature dimension') 85 | parser.add_argument('--skip_hourglass', action='store_true', help='skip connection in hourglass') 86 | parser.add_argument('--hg_down', type=str, default='ave_pool', help='ave pool || conv64 || conv128') 87 | parser.add_argument('--hourglass_dim', type=int, default='256', help='256 | 512') 88 | 89 | # Classification General 90 | parser.add_argument('--position_encoding', action='store_true', help='using position encoding') 91 | parser.add_argument('--mlp_dim', nargs='+', default=[323, 1024, 512, 256, 128, 1], type=int, 92 | help='# of dimensions of mlp') 93 | parser.add_argument('--fine_mlp_dim', nargs='+', default=[355, 512, 256, 128, 1], type=int, help='# of dimensions of fine mlp') 94 | parser.add_argument('--use_tanh', action='store_true', 95 | help='using tanh after last conv of image_filter network') 96 | 97 | # for train 98 | parser.add_argument('--random_flip', action='store_true', help='if random flip') 99 | parser.add_argument('--random_trans', action='store_true', help='if random flip') 100 | parser.add_argument('--random_scale', action='store_true', help='if random flip') 101 | parser.add_argument('--random_rotation', action='store_true', help='if random flip') 102 | parser.add_argument('--flip_normal', action='store_true', help='if smpl normal flip') 103 | parser.add_argument('--flip_smpl', action='store_true', help='if smpl 3d flip') 104 | parser.add_argument('--no_residual', action='store_true', help='no skip connection in mlp') 105 | 106 | parser.add_argument('--fine_part', action='store_true') 107 | parser.add_argument('--coarse_part', action='store_true') 108 | parser.add_argument('--mask_part', action='store_true') 109 | parser.add_argument('--preserve_single', action='store_true') 110 | 111 | # for infer 112 | parser.add_argument('--infer', action='store_true') 113 | parser.add_argument('--no_correct', action='store_true') 114 | parser.add_argument('--yaw_list', type=int, nargs='+', default=[0, 1, 2, 3]) 115 | parser.add_argument('--infer_reverse', action='store_true', help='infer rotation') 116 | parser.add_argument('--flip_x', action='store_true') 117 | 118 | # path 119 | parser.add_argument('--checkpoints_path', type=str, default='./checkpoints', help='path to save checkpoints') 120 | parser.add_argument('--load_netG_checkpoint_path', type=str, default=None, help='path to save checkpoints') 121 | parser.add_argument('--load_netN_checkpoint_path', type=str, default=None, help='path to save checkpoints') 122 | parser.add_argument('--load_optim_checkpoint_path', type=str, default=None, help='path to save checkpoints') 123 | parser.add_argument('--results_path', type=str, default='./results', help='path to save results ply') 124 | # for single image reconstruction 125 | parser.add_argument('--mask_path', type=str, help='path for input mask') 126 | parser.add_argument('--img_path', type=str, help='path for input image') 127 | 128 | # aug 129 | parser.add_argument('--aug_alstd', type=float, default=0.0, help='augmentation pca lighting alpha std') 130 | parser.add_argument('--aug_bri', type=float, default=0.0, help='augmentation brightness') 131 | parser.add_argument('--aug_con', type=float, default=0.0, help='augmentation contrast') 132 | parser.add_argument('--aug_sat', type=float, default=0.0, help='augmentation saturation') 133 | parser.add_argument('--aug_hue', type=float, default=0.0, help='augmentation hue') 134 | parser.add_argument('--aug_blur', type=float, default=0.0, help='augmentation blur') 135 | 136 | # debug 137 | parser.add_argument('--debug_3d', action='store_true') 138 | parser.add_argument('--debug_data', action='store_true') 139 | 140 | args, _ = parser.parse_known_args() 141 | 142 | return args -------------------------------------------------------------------------------- /lib/sample_util.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import torch 3 | from .geometry import batch_cross_3d 4 | 5 | def max_min_norm(x, b_min, b_max): 6 | ''' 7 | :param x [B, 3, N] 8 | :param b_min [3] list 9 | :param b_max [3] list 10 | ''' 11 | b_min = torch.FloatTensor(b_min).to(x.device).reshape(1, 3, 1) 12 | b_max = torch.FloatTensor(b_max).to(x.device).reshape(1, 3, 1) 13 | return (x - b_min) / (b_max - b_min) 14 | 15 | def obj_sample_surface_wo_tex(obj, count): 16 | with torch.no_grad(): 17 | vert = torch.FloatTensor(obj['vi'][:, :3]) 18 | faces = torch.LongTensor(obj['f']).permute(0, 2, 1) 19 | v_tri = torch.zeros((faces.shape[0], 3, 3)) 20 | for i in range(3): 21 | v_tri[:, i, :] = vert[faces[:, 0, i], :] 22 | v_ori = v_tri[:, 0, :].clone() 23 | v_vec = v_tri[:, 1:, :].clone() - v_ori.unsqueeze(1) 24 | 25 | area = torch.zeros((faces.shape[0], 1)) 26 | cross = batch_cross_3d(v_vec[:, 0, :], v_vec[:, 1, :]) 27 | area = torch.abs(torch.sum(cross, dim=1).unsqueeze(1)).numpy() 28 | 29 | area_sum = np.sum(area) 30 | area_cum = np.cumsum(area) 31 | face_pick = np.random.random(count) * area_sum 32 | face_index = np.searchsorted(area_cum, face_pick)-1 33 | 34 | v_ori = v_ori[face_index, :] 35 | v_vec = v_vec[face_index, :] 36 | 37 | random_lengths = np.random.random((count, 2, 1)) 38 | 39 | # points will be distributed on a quadrilateral if we use 2 0-1 samples 40 | # if the two scalar components sum less than 1.0 the point will be 41 | # inside the triangle, so we find vectors longer than 1.0 and 42 | # transform them to be inside the triangle 43 | random_test = random_lengths.sum(axis=1).reshape(-1) > 1.0 44 | random_lengths[random_test] -= 1.0 45 | random_lengths = torch.FloatTensor(np.abs(random_lengths)) 46 | 47 | # multiply triangle edge vectors by the random lengths and sum 48 | sample_v_vec = (v_vec * random_lengths).sum(dim=1) 49 | 50 | # finally, offset by the origin to generate 51 | v_samples = sample_v_vec + v_ori 52 | 53 | return v_samples 54 | 55 | def obj_sample_surface(obj, count): 56 | with torch.no_grad(): 57 | vert = torch.FloatTensor(obj['vi']) 58 | tex = torch.FloatTensor(obj['vt']) 59 | faces = torch.LongTensor(obj['f']).permute(0, 2, 1) 60 | v_tri = torch.zeros((faces.shape[0], 3, 3)) 61 | t_tri = torch.zeros((faces.shape[0], 3, 2)) 62 | for i in range(3): 63 | v_tri[:, i, :] = vert[faces[:, 0, i], :] 64 | t_tri[:, i, :] = tex[faces[:, 1, i], :] 65 | v_ori = v_tri[:, 0, :].clone() 66 | t_ori = t_tri[:, 0, :].clone() 67 | v_vec = v_tri[:, 1:, :].clone() - v_ori.unsqueeze(1) 68 | t_vec = t_tri[:, 1:, :].clone() - t_ori.unsqueeze(1) 69 | 70 | area = torch.zeros((faces.shape[0], 1)) 71 | cross = batch_cross_3d(v_vec[:, 0, :], v_vec[:, 1, :]) 72 | area = torch.abs(torch.sum(cross, dim=1).unsqueeze(1)).numpy() 73 | normal = cross / torch.sum(torch.sqrt(cross**2), dim=1).unsqueeze(1) 74 | 75 | area_sum = np.sum(area) 76 | # cumulative area (len(mesh.faces)) 77 | area_cum = np.cumsum(area) 78 | face_pick = np.random.random(count) * area_sum 79 | face_index = np.searchsorted(area_cum, face_pick)-1 80 | 81 | v_ori, t_ori = v_ori[face_index, :], t_ori[face_index, :] 82 | v_vec, t_vec = v_vec[face_index, :], t_vec[face_index, :] 83 | 84 | random_lengths = np.random.random((count, 2, 1)) 85 | 86 | # points will be distributed on a quadrilateral if we use 2 0-1 samples 87 | # if the two scalar components sum less than 1.0 the point will be 88 | # inside the triangle, so we find vectors longer than 1.0 and 89 | # transform them to be inside the triangle 90 | random_test = random_lengths.sum(axis=1).reshape(-1) > 1.0 91 | random_lengths[random_test] -= 1.0 92 | random_lengths = torch.FloatTensor(np.abs(random_lengths)) 93 | 94 | # multiply triangle edge vectors by the random lengths and sum 95 | sample_v_vec = (v_vec * random_lengths).sum(dim=1) 96 | sample_t_vec = (t_vec * random_lengths).sum(dim=1) 97 | 98 | # finally, offset by the origin to generate 99 | v_samples = sample_v_vec + v_ori 100 | t_samples = sample_t_vec + t_ori 101 | n_samples = normal[face_index, :] 102 | 103 | return v_samples, t_samples, n_samples 104 | 105 | def save_samples_truncted_prob_obj(fname, points, prob, continue_write=False, pts_num=0): 106 | ''' 107 | Save the visualization of sampling to a obj file. 108 | Red points represent positive predictions. 109 | Green points represent negative predictions. 110 | :param fname: File name to save 111 | :param points: [N, 3] array of points 112 | :param prob: [N, 1] array of predictions in the range [0~1] 113 | :return: 114 | ''' 115 | r = (prob).reshape([-1, 1]) * 255 116 | g = (1-prob).reshape([-1, 1]) * 255 117 | b = np.zeros(r.shape) 118 | if pts_num == 0: 119 | pts_num = points.shape[0] 120 | 121 | to_save = np.concatenate([points, r, g, b], axis=-1) 122 | if continue_write: 123 | f = open(fname, mode='a') 124 | else: 125 | f = open(fname, mode='w') 126 | # f.write('ply\nformat ascii 1.0\nelement vertex %d\nproperty float x\nproperty float y\nproperty float z\nproperty uchar red\nproperty uchar green\nproperty uchar blue\nend_header\n' % pts_num) 127 | for i in range(points.shape[0]): 128 | f.write("v %f %f %f %d %d %d\n" % (points[i, 0], points[i, 1], points[i, 2], r[i], g[i], b[i])) 129 | 130 | def save_samples_truncted_prob(fname, points, prob): 131 | ''' 132 | Save the visualization of sampling to a ply file. 133 | Red points represent positive predictions. 134 | Green points represent negative predictions. 135 | :param fname: File name to save 136 | :param points: [N, 3] array of points 137 | :param prob: [N, 1] array of predictions in the range [0~1] 138 | :return: 139 | ''' 140 | r = (prob).reshape([-1, 1]) * 255 141 | g = (1-prob).reshape([-1, 1]) * 255 142 | b = np.zeros(r.shape) 143 | 144 | to_save = np.concatenate([points, r, g, b], axis=-1) 145 | return np.savetxt(fname, 146 | to_save, 147 | fmt='%.6f %.6f %.6f %d %d %d', 148 | comments='', 149 | header=( 150 | 'ply\nformat ascii 1.0\nelement vertex {:d}\nproperty float x\nproperty float y\nproperty float z\nproperty uchar red\nproperty uchar green\nproperty uchar blue\nend_header').format( 151 | points.shape[0]) 152 | ) 153 | 154 | 155 | def save_samples_rgb(fname, points, rgb): 156 | ''' 157 | Save the visualization of sampling to a ply file. 158 | Red points represent positive predictions. 159 | Green points represent negative predictions. 160 | :param fname: File name to save 161 | :param points: [N, 3] array of points 162 | :param rgb: [N, 3] array of rgb values in the range [0~1] 163 | :return: 164 | ''' 165 | to_save = np.concatenate([points, rgb * 255], axis=-1) 166 | return np.savetxt(fname, 167 | to_save, 168 | fmt='%.6f %.6f %.6f %d %d %d', 169 | comments='', 170 | header=( 171 | 'ply\nformat ascii 1.0\nelement vertex {:d}\nproperty float x\nproperty float y\nproperty float z\nproperty uchar red\nproperty uchar green\nproperty uchar blue\nend_header').format( 172 | points.shape[0]) 173 | ) 174 | -------------------------------------------------------------------------------- /lib/sdf.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | 3 | 4 | def create_grid(resX, resY, resZ, b_min=np.array([0, 0, 0]), b_max=np.array([1, 1, 1]), transform=None): 5 | ''' 6 | Create a dense grid of given resolution and bounding box 7 | :param resX: resolution along X axis 8 | :param resY: resolution along Y axis 9 | :param resZ: resolution along Z axis 10 | :param b_min: vec3 (x_min, y_min, z_min) bounding box corner 11 | :param b_max: vec3 (x_max, y_max, z_max) bounding box corner 12 | :return: [3, resX, resY, resZ] coordinates of the grid, and transform matrix from mesh index 13 | ''' 14 | # print('start creating grid') 15 | coords = np.mgrid[:resX, :resY, :resZ] 16 | coords = coords.reshape(3, -1) 17 | coords_matrix = np.eye(4) 18 | length = b_max - b_min 19 | coords_matrix[0, 0] = length[0] / resX 20 | coords_matrix[1, 1] = length[1] / resY 21 | coords_matrix[2, 2] = length[2] / resZ 22 | coords_matrix[0:3, 3] = b_min 23 | coords = np.matmul(coords_matrix[:3, :3], coords) + coords_matrix[:3, 3:4] 24 | if transform is not None: 25 | coords = np.matmul(transform[:3, :3], coords) + transform[:3, 3:4] 26 | coords_matrix = np.matmul(transform, coords_matrix) 27 | coords = coords.reshape(3, resX, resY, resZ) 28 | # print('creating_grid_done') 29 | return coords, coords_matrix 30 | 31 | 32 | def batch_eval(points, eval_func, num_samples=512 * 512 * 512): 33 | num_pts = points.shape[1] 34 | sdf = np.zeros(num_pts) 35 | 36 | num_batches = num_pts // num_samples 37 | for i in range(num_batches): 38 | sdf[i * num_samples:i * num_samples + num_samples] = eval_func( 39 | points[:, i * num_samples:i * num_samples + num_samples]) 40 | if num_pts % num_samples: 41 | sdf[num_batches * num_samples:] = eval_func(points[:, num_batches * num_samples:]) 42 | 43 | return sdf 44 | 45 | 46 | def eval_grid(coords, eval_func, num_samples=512 * 512 * 512): 47 | resolution = coords.shape[1:4] 48 | coords = coords.reshape([3, -1]) 49 | sdf = batch_eval(coords, eval_func, num_samples=num_samples) 50 | return sdf.reshape(resolution) 51 | 52 | 53 | def eval_grid_octree(coords, eval_func, 54 | init_resolution=64, threshold=0.01, 55 | num_samples=512 * 512 * 512): 56 | resolution = coords.shape[1:4] 57 | 58 | sdf = np.zeros(resolution) 59 | 60 | dirty = np.ones(resolution, dtype=np.bool) 61 | grid_mask = np.zeros(resolution, dtype=np.bool) 62 | 63 | reso = resolution[0] // init_resolution 64 | 65 | while reso > 0: 66 | # subdivide the grid 67 | grid_mask[0:resolution[0]:reso, 0:resolution[1]:reso, 0:resolution[2]:reso] = True 68 | # test samples in this iteration 69 | test_mask = np.logical_and(grid_mask, dirty) 70 | #print('step size:', reso, 'test sample size:', test_mask.sum()) 71 | points = coords[:, test_mask] 72 | 73 | sdf[test_mask] = batch_eval(points, eval_func, num_samples=num_samples) 74 | dirty[test_mask] = False 75 | 76 | # do interpolation 77 | if reso <= 1: 78 | break 79 | for x in range(0, resolution[0] - reso, reso): 80 | for y in range(0, resolution[1] - reso, reso): 81 | for z in range(0, resolution[2] - reso, reso): 82 | # if center marked, return 83 | if not dirty[x + reso // 2, y + reso // 2, z + reso // 2]: 84 | continue 85 | v0 = sdf[x, y, z] 86 | v1 = sdf[x, y, z + reso] 87 | v2 = sdf[x, y + reso, z] 88 | v3 = sdf[x, y + reso, z + reso] 89 | v4 = sdf[x + reso, y, z] 90 | v5 = sdf[x + reso, y, z + reso] 91 | v6 = sdf[x + reso, y + reso, z] 92 | v7 = sdf[x + reso, y + reso, z + reso] 93 | v = np.array([v0, v1, v2, v3, v4, v5, v6, v7]) 94 | v_min = v.min() 95 | v_max = v.max() 96 | # this cell is all the same 97 | if (v_max - v_min) < threshold: 98 | sdf[x:x + reso, y:y + reso, z:z + reso] = (v_max + v_min) / 2 99 | dirty[x:x + reso, y:y + reso, z:z + reso] = False 100 | reso //= 2 101 | 102 | return sdf.reshape(resolution) 103 | -------------------------------------------------------------------------------- /lib/smpl_util.py: -------------------------------------------------------------------------------- 1 | from .geometry import * 2 | 3 | def calc_smpl_normal(smpl): 4 | vi = smpl['vi'] 5 | vn = np.zeros((vi.shape[0], 3)) 6 | for f in smpl['f']: 7 | a, b, c = vi[f[0, 0]], vi[f[1, 0]], vi[f[2, 0]] 8 | n = cross_3d(c-a, b-a) 9 | vn[f[0, 0]] += n 10 | vn[f[1, 0]] += n 11 | vn[f[2, 0]] += n 12 | vn = vn / np.sqrt(np.sum(vn**2, axis=1)).reshape((-1, 1)) 13 | return vn 14 | 15 | def smpl_normalize(smpl, norm_faces, flip_normal=False, init_rot=None): 16 | b_min = np.zeros(3) 17 | b_max = np.zeros(3) 18 | scale = 1.0 19 | center = np.zeros(3) 20 | 21 | vi = smpl['vi'].copy()[:, :3] 22 | b0 = np.min(vi, axis=0) 23 | b1 = np.max(vi, axis=0) 24 | center = (b0 + b1) / 2 25 | scale = np.min(1.0 / (b1 - b0)) * 0.9 26 | b_min = center - 0.5 / scale 27 | b_max = center + 0.5 / scale 28 | 29 | normal = np.zeros((3)) 30 | for f in norm_faces: 31 | a, b, c = vi[f[0]][0], vi[f[1]][0], vi[f[2]][0] 32 | normal += cross_3d(c - a, b - a) 33 | if flip_normal: 34 | normal = -normal 35 | 36 | x, z = normal[0], normal[2] 37 | theta = math.acos(z / math.sqrt(z*z +x*x)) 38 | if x < 0: 39 | theta = 2*math.acos(-1) - theta 40 | rot = rotationY(-theta) 41 | 42 | vi -= center 43 | vi *= scale 44 | if init_rot is not None: 45 | x, y, z = init_rot 46 | rot = np.array(rotationX(x / 180 * math.acos(-1))) @ rot 47 | rot = np.array(rotationY(y / 180 * math.acos(-1))) @ rot 48 | rot = np.array(rotationZ(z / 180 * math.acos(-1))) @ rot 49 | 50 | vi = (rot @ vi.T).T 51 | 52 | return { 53 | 'b_min' : b_min, 54 | 'b_max' : b_max, 55 | 'scale' : scale, 56 | 'center' : center, 57 | 'direction' : normal, 58 | 'smpl' : vi 59 | } -------------------------------------------------------------------------------- /lib/train_util.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import numpy as np 3 | from .mesh_util import * 4 | from .sample_util import * 5 | from .geometry import * 6 | import cv2 7 | from PIL import Image 8 | from tqdm import tqdm 9 | from PIL.ImageFilter import MinFilter 10 | 11 | def find_border(img): 12 | img = img.filter(MinFilter(11)) 13 | img = np.array(img) 14 | img_1 = np.sum(img, axis=2) 15 | img_x = np.sum(img_1, axis=0) 16 | img_y = np.sum(img_1, axis=1) 17 | x_min = img_x.shape[0] 18 | x_max = 0 19 | y_min = img_y.shape[0] 20 | y_max = 0 21 | for x in range(img_x.shape[0]): 22 | if img_x[x] > 0: 23 | x_min = x 24 | break 25 | for x in range(img_x.shape[0] - 1, 0, -1): 26 | if img_x[x] > 0: 27 | x_max = x 28 | break 29 | for y in range(img_y.shape[0]): 30 | if img_y[y] > 0: 31 | y_min = y 32 | break 33 | for y in range(img_y.shape[0] - 1, 0, -1): 34 | if img_y[y] > 0: 35 | y_max = y 36 | break 37 | return x_min, x_max, y_min, y_max 38 | 39 | def reshape_multiview_tensors(image_tensor, calib_tensor): 40 | # Careful here! Because we put single view and multiview together, 41 | # the returned tensor.shape is 5-dim: [B, num_views, C, W, H] 42 | # So we need to convert it back to 4-dim [B*num_views, C, W, H] 43 | # Don't worry classifier will handle multi-view cases 44 | image_tensor = image_tensor.view( 45 | image_tensor.shape[0] * image_tensor.shape[1], 46 | image_tensor.shape[2], 47 | image_tensor.shape[3], 48 | image_tensor.shape[4] 49 | ) 50 | calib_tensor = calib_tensor.view( 51 | calib_tensor.shape[0] * calib_tensor.shape[1], 52 | calib_tensor.shape[2], 53 | calib_tensor.shape[3] 54 | ) 55 | 56 | return image_tensor, calib_tensor 57 | 58 | def reshape_multiview_tensors_3d(image_tensor, calib_tensor, extrinsic): 59 | # Careful here! Because we put single view and multiview together, 60 | # the returned tensor.shape is 5-dim: [B, num_views, C, W, H] 61 | # So we need to convert it back to 4-dim [B*num_views, C, W, H] 62 | # Don't worry classifier will handle multi-view cases 63 | image_tensor = image_tensor.view( 64 | image_tensor.shape[0] * image_tensor.shape[1], 65 | image_tensor.shape[2], 66 | image_tensor.shape[3], 67 | image_tensor.shape[4] 68 | ) 69 | calib_tensor = calib_tensor.view( 70 | calib_tensor.shape[0] * calib_tensor.shape[1], 71 | calib_tensor.shape[2], 72 | calib_tensor.shape[3] 73 | ) 74 | extrinsic = extrinsic.view( 75 | extrinsic.shape[0] * extrinsic.shape[1], 76 | extrinsic.shape[2], 77 | extrinsic.shape[3] 78 | ) 79 | 80 | return image_tensor, calib_tensor, extrinsic 81 | 82 | 83 | def reshape_sample_tensor(sample_tensor, num_views): 84 | if num_views == 1: 85 | return sample_tensor 86 | # Need to repeat sample_tensor along the batch dim num_views times 87 | sample_tensor = sample_tensor.unsqueeze(dim=1) 88 | sample_tensor = sample_tensor.repeat(1, num_views, 1, 1) 89 | sample_tensor = sample_tensor.view( 90 | sample_tensor.shape[0] * sample_tensor.shape[1], 91 | sample_tensor.shape[2], 92 | sample_tensor.shape[3] 93 | ) 94 | return sample_tensor 95 | 96 | 97 | def adjust_learning_rate(optimizer, epoch, lr, schedule, gamma): 98 | """Sets the learning rate to the initial LR decayed by schedule""" 99 | if epoch in schedule: 100 | lr *= gamma 101 | for param_group in optimizer.param_groups: 102 | param_group['lr'] = lr 103 | return lr -------------------------------------------------------------------------------- /render_dataset.sh: -------------------------------------------------------------------------------- 1 | cd taichi_render_gpu 2 | python render_multi.py --data_root ../assets/obj --texture_root ../assets/tex --save_path ../dataset/example 3 | python render_smpl.py --dataroot ../dataset/example --obj_path ../assets/smplx --faces_path ../lib/data/smplx_fine.obj 4 | -------------------------------------------------------------------------------- /render_two.sh: -------------------------------------------------------------------------------- 1 | cd taichi_render_gpu 2 | 3 | python render_smpl.py --dataroot ../dataset/multihuman \ 4 | --obj_path ../dataset/multihuman/smplx \ 5 | --faces_path ../lib/data/smplx_multi.obj --yaw_list 0 90 180 270 6 | 7 | -------------------------------------------------------------------------------- /taichi_render_gpu/README.md: -------------------------------------------------------------------------------- 1 | Taichi THREE 2 | ============ 3 | 4 | [![Downloads](https://pepy.tech/badge/taichi-three)](https://pepy.tech/project/taichi-three) 5 | [![Latest Release](https://img.shields.io/github/v/release/taichi-dev/taichi_three)](https://github.com/taichi-dev/taichi_three/releases) 6 | 7 | Taichi THREE is an extension library of the [Taichi Programming Language](https://github.com/taichi-dev/taichi) to render 3D scenes into nice-looking 2D images in real-time (work in progress). 8 | 9 | 10 | ![Example](https://github.com/taichi-dev/taichi_three/raw/master/assets/monkey.png) 11 | 12 | 13 | Installation 14 | ------------ 15 | 16 | Install Taichi THREE Dev: 17 | 18 | ```bash 19 | # Python 3.6/3.7/3.8 (64-bit) 20 | pip install taichi taichi_glsl 21 | python setup.py build install 22 | ``` 23 | 24 | This should also install its dependencies `taichi` and `taichi_glsl` as well. 25 | 26 | 27 | How to play 28 | ----------- 29 | 30 | First, import Taichi and Taichi THREE: 31 | ```py 32 | import taichi as ti 33 | import taichi_three as t3 34 | 35 | ti.init(ti.gpu) 36 | ``` 37 | 38 | --- 39 | 40 | Then, create a scene using: 41 | ```py 42 | scene = t3.Scene() 43 | ``` 44 | 45 | NOTE: scene creation should be put before any kernel invocation or host access, 46 | i.e. before materialization, so that `Scene.__init__` could define its internal 47 | tensors without an error. 48 | 49 | TL;DR: Put this line as forward as possible! Ideally right below `ti.init()`. 50 | 51 | --- 52 | 53 | After that, load the model(s), and feed them into `scene`: 54 | 55 | ```py 56 | model = t3.Model(t3.readobj('assets/monkey.obj', scale=0.6)) 57 | scene.add_model(model) 58 | ``` 59 | 60 | If you want to add texture, read the texture image and feed it into `model`: 61 | 62 | ```py 63 | texture = ti.imread('assets/cloth.jpg') 64 | model = t3.Model(t3.readobj('assets/monkey.obj', scale=0.6), tex=texture) 65 | ``` 66 | 67 | NOTE: model creations should also be put as forward as possible too. 68 | 69 | --- 70 | 71 | Then, create the camera(s), and put it into `scene`: 72 | 73 | ```py 74 | camera = t3.Camera() 75 | scene.add_camera(camera) 76 | ``` 77 | 78 | NOTE: camera creations should also be put as forward as possible. 79 | 80 | --- 81 | 82 | Also don't forget to set the light: 83 | ```py 84 | light = t3.Light() 85 | scene.add_light(light) 86 | ``` 87 | 88 | --- 89 | 90 | Finally, create a GUI. And here goes the main loop: 91 | 92 | ```py 93 | gui = ti.GUI('Monkey') 94 | while gui.running: 95 | scene.render() # render the model(s) into image 96 | gui.set_image(camera.img) # display the result image 97 | gui.show() 98 | ``` 99 | 100 | --- 101 | 102 | Checkout the [`examples/`](https://github.com/taichi-dev/taichi_three/tree/master/examples) directory for more runnable examples. 103 | -------------------------------------------------------------------------------- /taichi_render_gpu/render_multi.py: -------------------------------------------------------------------------------- 1 | import taichi as ti 2 | import taichi_three as t3 3 | import numpy as np 4 | from taichi_three.transform import * 5 | from tqdm import tqdm 6 | import os 7 | import time 8 | import cv2 9 | import json 10 | os.environ["KMP_DUPLICATE_LIB_OK"]="True" 11 | 12 | def find_border(img): 13 | img_1 = np.sum(img, axis=2) 14 | img_x = np.sum(img_1, axis=0) 15 | img_y = np.sum(img_1, axis=1) 16 | x_min = img_x.shape[0] 17 | x_max = 0 18 | y_min = img_y.shape[0] 19 | y_max = 0 20 | for x in range(img_x.shape[0]): 21 | if img_x[x] > 0: 22 | x_min = x 23 | break 24 | for x in range(img_x.shape[0]-1, 0, -1): 25 | if img_x[x] > 0: 26 | x_max = x 27 | break 28 | for y in range(img_y.shape[0]): 29 | if img_y[y] > 0: 30 | y_min = y 31 | break 32 | for y in range(img_y.shape[0]-1, 0, -1): 33 | if img_y[y] > 0: 34 | y_max = y 35 | break 36 | return x_min, x_max, y_min, y_max 37 | 38 | class StaticRenderer: 39 | def __init__(self): 40 | ti.init(ti.cpu) 41 | self.scene = t3.Scene() 42 | self.N = 10 43 | 44 | def change_all(self): 45 | save_obj = [] 46 | save_tex = [] 47 | for model in self.scene.models: 48 | save_obj.append(model.init_obj) 49 | save_tex.append(model.init_tex) 50 | ti.init(ti.cpu) 51 | print('init') 52 | self.scene = t3.Scene() 53 | for i in range(len(save_obj)): 54 | model = t3.StaticModel(self.N, obj=save_obj[i], tex=save_tex[i]) 55 | self.scene.add_model(model) 56 | 57 | def check_update(self, obj): 58 | temp_n = self.N 59 | self.N = max(obj['vi'].shape[0], self.N) 60 | self.N = max(obj['f'].shape[0], self.N) 61 | if not (obj['vt'] is None): 62 | self.N = max(obj['vt'].shape[0], self.N) 63 | 64 | if self.N > temp_n: 65 | self.N *= 2 66 | self.change_all() 67 | self.camera_light() 68 | 69 | def add_model(self, obj, tex=None): 70 | self.check_update(obj) 71 | model = t3.StaticModel(self.N, obj=obj, tex=tex) 72 | self.scene.add_model(model) 73 | 74 | def modify_model(self, index, obj, tex=None): 75 | self.check_update(obj) 76 | self.scene.models[index].init_obj = obj 77 | self.scene.models[index].init_tex = tex 78 | self.scene.models[index]._init() 79 | 80 | def camera_light(self): 81 | camera = t3.Camera(res=res) 82 | camera1 = t3.Camera(res=(512, 512)) 83 | self.scene.add_camera(camera) 84 | self.scene.add_camera(camera1) 85 | light_dir = np.array([0, 0, 1]) 86 | for l in range(6): 87 | rotate = np.matmul(rotationX(math.radians(np.random.uniform(-30, 30))), 88 | rotationY(math.radians(360 // 6 * l))) 89 | dir = [*np.matmul(rotate, light_dir)] 90 | light = t3.Light(dir, color=[1.0, 1.0, 1.0]) 91 | self.scene.add_light(light) 92 | 93 | def render_mv_random_mask(renderer, data_path, texture_path, data_id, save_path, res=(1024, 1024), enable_gpu=False, dis_scale=1, 94 | ran_mask_num = 0): 95 | img_path = os.path.join(texture_path, data_id + '.jpg') 96 | obj_path = os.path.join(data_path, data_id, data_id + '.obj') 97 | obj_names = os.listdir(data_path) 98 | 99 | img_save_path = os.path.join(save_path, 'img', data_id) 100 | depth_save_path = os.path.join(save_path, 'depth', data_id) 101 | normal_save_path = os.path.join(save_path, 'normal', data_id) 102 | mask_save_path = os.path.join(save_path, 'mask', data_id) 103 | parameter_save_path = os.path.join(save_path, 'parameter', data_id) 104 | if not os.path.exists(img_save_path): 105 | os.makedirs(img_save_path) 106 | if not os.path.exists(parameter_save_path): 107 | os.makedirs(parameter_save_path) 108 | if not os.path.exists(mask_save_path): 109 | os.makedirs(mask_save_path) 110 | if not os.path.exists(normal_save_path): 111 | os.makedirs(normal_save_path) 112 | if not os.path.exists(depth_save_path): 113 | os.makedirs(depth_save_path) 114 | 115 | texture = ti.imread(img_path) 116 | obj = t3.readobj(obj_path, scale=1) 117 | 118 | if len(renderer.scene.models) >= 1: 119 | renderer.modify_model(0, obj, texture) 120 | else: 121 | renderer.add_model(obj, texture) 122 | 123 | mask_shapes = [] 124 | mask_objs = [] 125 | for i in range(ran_mask_num): 126 | mask_name = np.random.choice(obj_names, 1)[0] 127 | mask_obj_path = os.path.join(data_path, mask_name, mask_name + '.obj') 128 | mask_tex_path = os.path.join(texture_path, mask_name + '.jpg') 129 | 130 | msk_tex = ti.imread(mask_tex_path) 131 | msk_obj = t3.readobj(mask_obj_path, scale=1) 132 | rotation = np.array(rotationY(math.radians(np.random.uniform(0, 360)))) 133 | msk_obj['vi'][:, :3] = (rotation @ msk_obj['vi'][:, :3].T).T 134 | obj_bbx = np.max(obj['vi'], axis=0) - np.min(obj['vi'], axis=0) 135 | mask_obj_bbx = np.max(msk_obj['vi'], axis=0) - np.min(msk_obj['vi'], axis=0) 136 | r = np.sqrt( (obj_bbx[0] + mask_obj_bbx[0])**2 + (obj_bbx[2] + mask_obj_bbx[2])**2) / 2 137 | ran = np.random.rand() * math.acos(-1) * 2 138 | shift = np.array([r*math.sin(ran), 0, r*math.cos(ran)]) 139 | msk_obj['vi'][:, :3] += shift.reshape((1, 3)) 140 | if len(renderer.scene.models) >= 2+i: 141 | renderer.modify_model(i+1, msk_obj, msk_tex) 142 | else: 143 | renderer.add_model(msk_obj, msk_tex) 144 | mask_shapes.append(msk_obj['vi'].shape[0]) 145 | mask_objs.append(msk_obj) 146 | 147 | vi = obj['vi'] 148 | median = np.median(vi, axis=0) # + (np.random.randn(3) - 0.5) * 0.2 149 | vmin = vi.min(0) 150 | vmax = vi.max(0) 151 | median[1] = (vmax[1] * 4 + vmin[1] * 3) / 7 152 | 153 | r_color = np.zeros((obj['vi'].shape[0], 3)) 154 | r_color[:, 0] = 1 155 | renderer.scene.models[0].modify_color(r_color) 156 | for i in range(1, ran_mask_num+1): 157 | b_color = np.zeros((mask_shapes[i-1], 3)) 158 | b_color[:, 2] = 1 159 | renderer.scene.models[i].modify_color(b_color) 160 | 161 | angle_mul = 1 162 | for angle in tqdm(range(360), desc='angle'): 163 | for i in range(ran_mask_num + 1): 164 | renderer.scene.models[i].type[None] = 0 165 | # if (os.path.exists(os.path.join(img_save_path, '{}.jpg'.format(angle)))): 166 | # continue 167 | dis = vmax[1] - vmin[1] 168 | dis *= dis_scale 169 | ori_vec = np.array([0, 0, dis]) 170 | p = np.random.uniform(-30, 10) 171 | rotate = np.matmul(rotationY(math.radians(angle*angle_mul)), rotationX(math.radians(p))) 172 | fwd = np.matmul(rotate, ori_vec) 173 | fx = res[0] * 0.5 174 | fy = res[1] * 0.5 175 | cx = fx 176 | cy = fy 177 | target = median 178 | pos = target + fwd 179 | renderer.scene.cameras[0].set(pos=pos, target=target) 180 | renderer.scene.cameras[0].set_intrinsic(fx, fy, cx, cy) 181 | renderer.scene.cameras[0]._init() 182 | renderer.scene.single_render(0) 183 | 184 | img = renderer.scene.cameras[0].img.to_numpy() 185 | 186 | x_min, x_max, y_min, y_max = find_border(img) 187 | 188 | x_min -= 20 189 | x_max += 20 190 | x_len = x_max - x_min 191 | y_min = (y_max + y_min - x_len) // 2 192 | scale = 512.0 / x_len 193 | fx = 512 * scale 194 | fy = 512 * scale 195 | cx = scale * (cx - y_min) 196 | cy = scale * (cy - x_min) 197 | renderer.scene.cameras[1].set_intrinsic(fx, fy, cx, cy) 198 | renderer.scene.cameras[1].set(pos=pos,target=target) 199 | renderer.scene.cameras[1]._init() 200 | renderer.scene.single_render(0) 201 | camera1 = renderer.scene.cameras[1] 202 | depth_map = camera1.zbuf.to_numpy().swapaxes(0, 1)[::-1, :] 203 | np.savez(os.path.join(depth_save_path, '{}.npz'.format(angle)), depth_map) 204 | 205 | renderer.scene.render() 206 | camera1 = renderer.scene.cameras[1] 207 | np.save(os.path.join(parameter_save_path, '{}_extrinsic.npy'.format(angle)), 208 | camera1.export_extrinsic()) 209 | np.save(os.path.join(parameter_save_path, '{}_intrinsic.npy'.format(angle)), 210 | camera1.export_intrinsic()) 211 | ti.imwrite(camera1.img, os.path.join(img_save_path, '{}.png'.format(angle))) 212 | ti.imwrite(camera1.normal_map, os.path.join(normal_save_path, '{}.png'.format(angle))) 213 | 214 | for i in range(ran_mask_num + 1): 215 | renderer.scene.models[i].type[None] = 1 216 | renderer.scene.render() 217 | camera1 = renderer.scene.cameras[1] 218 | mask = camera1.img.to_numpy() 219 | # mask[:, :, 2] = mask[:, :, 0] 220 | # mask[:, :, 1] = mask[:, :, 0] 221 | ti.imwrite(mask, os.path.join(mask_save_path, '{}.png'.format(angle))) 222 | 223 | # for i in range(ran_mask_num+1): 224 | # del scene.models[i] 225 | # for i in range(ran_mask_num): 226 | # del mask_objs[i] 227 | 228 | if __name__ == '__main__': 229 | res = (1024, 1024) 230 | 231 | import argparse 232 | parser = argparse.ArgumentParser() 233 | parser.add_argument("--data_root", type=str) 234 | parser.add_argument("--texture_root", type=str) 235 | parser.add_argument("--save_path", type=str) 236 | parser.add_argument("--ran_mask_num", type=int, default=0) 237 | args = parser.parse_args() 238 | data_root = args.data_root 239 | texture_root = args.texture_root 240 | save_path = args.save_path 241 | ran_mask_num = args.ran_mask_num 242 | it = 0 243 | render_num = 0 244 | renderer = StaticRenderer() 245 | for data_id in tqdm(os.listdir(data_root), desc='data_id'): 246 | render_mv_random_mask(renderer, data_root, texture_root, data_id, save_path, res, False, dis_scale=2, ran_mask_num=ran_mask_num) 247 | render_num += 1 248 | if render_num > 50: 249 | break 250 | -------------------------------------------------------------------------------- /taichi_render_gpu/render_smpl.py: -------------------------------------------------------------------------------- 1 | import taichi as ti 2 | import taichi_three as t3 3 | import numpy as np 4 | from taichi_three.transform import * 5 | from tqdm import tqdm 6 | import os 7 | import sys 8 | import cv2 9 | import trimesh 10 | import matplotlib.pyplot as plt 11 | 12 | sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..'))) 13 | ROOT_PATH = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) 14 | 15 | from lib.smpl_util import * 16 | from lib.mesh_util import * 17 | 18 | def read_norm_smpl(path, smpl_faces, flip_normal=False, init_rot=None): 19 | obj = t3.readobj(path, scale=1) 20 | faces = t3.readobj(smpl_faces)['f'] 21 | o_vi = obj['vi'].copy() 22 | norm_vi = smpl_normalize(obj, faces, flip_normal, init_rot)['smpl'] 23 | obj['vi'] = norm_vi 24 | vn = calc_smpl_normal(obj) 25 | obj['vi'] = o_vi 26 | obj['vp'] = norm_vi 27 | obj['vn'] = vn 28 | return obj 29 | 30 | 31 | def render_smpl_global_normal(dataroot, obj_path, faces_path, res=(1024, 1024), angles=range(360), flip_y=False, flip_normal=False, init_rot=None): 32 | ti.init(ti.cpu) 33 | pos_save_root = os.path.join(dataroot, 'smpl_pos') 34 | os.makedirs(pos_save_root, exist_ok=True) 35 | parameter_path = os.path.join(dataroot, 'parameter') 36 | obj_list = os.listdir(obj_path) 37 | obj = read_norm_smpl(os.path.join(obj_path, obj_list[0], 'smplx.obj'), faces_path, flip_normal, init_rot) 38 | model = t3.Model(obj=obj, col_n=obj['vi'].shape[0]) 39 | 40 | scene = t3.Scene() 41 | scene.add_model(model) 42 | light_dir = np.array([0, 0, 1]) 43 | for l in range(4): 44 | rotate = np.matmul(rotationX(math.radians(np.random.uniform(-30, 30))), 45 | rotationY(math.radians(360 // 4 * l))) 46 | dir = [*np.matmul(rotate, light_dir)] 47 | light = t3.Light(dir, color=[1.0, 1.0, 1.0]) 48 | scene.add_light(light) 49 | 50 | camera = t3.Camera(res=res) 51 | scene.add_camera(camera) 52 | scene.init() 53 | for obj_name in tqdm(os.listdir(parameter_path)): 54 | pos_save_path = os.path.join(pos_save_root, obj_name) 55 | if os.path.exists(pos_save_path) and len(os.listdir(os.path.join(pos_save_path))) == len(angles): 56 | continue 57 | if not os.path.exists(os.path.join(obj_path, obj_name, 'smplx.obj')): 58 | continue 59 | obj = read_norm_smpl(os.path.join(obj_path, obj_name, 'smplx.obj'), faces_path, flip_normal, init_rot) 60 | os.makedirs(pos_save_path, exist_ok=True) 61 | for angle in angles: 62 | intrinsic = np.load(os.path.join(parameter_path, obj_name, '{}_intrinsic.npy'.format(angle))) 63 | extrinsic = np.load(os.path.join(parameter_path, obj_name, '{}_extrinsic.npy'.format(angle))) 64 | 65 | if flip_y: 66 | camera.set_intrinsic(fx=intrinsic[0, 0], fy=-intrinsic[1, 1], cx=intrinsic[0, 2], cy=res[0]-intrinsic[1, 2]) 67 | else: 68 | camera.set_intrinsic(fx=intrinsic[0, 0], fy=intrinsic[1, 1], cx=intrinsic[0, 2], cy=intrinsic[1, 2]) 69 | 70 | trans = extrinsic[:, :3] 71 | T = extrinsic[:, 3] 72 | p = -trans.T @ T 73 | camera.set_extrinsic(trans.T, p) 74 | color = obj['vn'] 75 | model.from_obj(obj) 76 | model.vc.from_numpy(color) 77 | model.type[None] = 1 78 | camera._init() 79 | scene.render() 80 | 81 | ti.imwrite( (camera.img.to_numpy() + 1)/2, os.path.join(pos_save_path, '{}.jpg'.format(angle))) 82 | 83 | 84 | if __name__ == '__main__': 85 | res = (512, 512) 86 | 87 | import argparse 88 | parser = argparse.ArgumentParser() 89 | parser.add_argument("--dataroot", type=str) 90 | parser.add_argument("--obj_path", type=str) 91 | parser.add_argument("--faces_path", type=str) 92 | parser.add_argument("--yaw_list", type=int, nargs='+', default=[i for i in range(360)]) 93 | parser.add_argument("--flip_y", action="store_true") 94 | parser.add_argument("--flip_normal", action="store_true") 95 | args = parser.parse_args() 96 | 97 | render_smpl_global_normal(args.dataroot, args.obj_path, args.faces_path, res, args.yaw_list, args.flip_y, args.flip_normal) 98 | 99 | -------------------------------------------------------------------------------- /taichi_render_gpu/setup.py: -------------------------------------------------------------------------------- 1 | project_name = 'taichi_three' 2 | from taichi_three.version import version, taichi_version, taiglsl_version 3 | version = '.'.join(map(str, version)) 4 | taichi_version = '.'.join(map(str, taichi_version)) 5 | taiglsl_version = '.'.join(map(str, taiglsl_version)) 6 | description = 'A Taichi extension library helps you rendering 3D scenes' 7 | long_description = ''' 8 | Taichi THREE 9 | ============ 10 | 11 | Taichi THREE is an extension library of the `Taichi Programming Language `_, that helps rendering your 3D scenes into nice-looking 2D images to display in GUI. 12 | ''' 13 | classifiers = [ 14 | 'Intended Audience :: Developers', 15 | 'Intended Audience :: Science/Research', 16 | 'Topic :: Multimedia :: Graphics', 17 | 'Topic :: Games/Entertainment :: Simulation', 18 | 'Operating System :: OS Independent', 19 | ] 20 | python_requires = '>=3.6' 21 | install_requires = [ 22 | 'taichi==' + taichi_version, 23 | 'taichi-glsl==' + taiglsl_version, 24 | ] 25 | 26 | import setuptools 27 | 28 | setuptools.setup( 29 | name=project_name, 30 | version=version, 31 | author='邵睿智', 32 | author_email='jia1saurus@gmail.com', 33 | description=description, 34 | long_description=long_description, 35 | classifiers=classifiers, 36 | python_requires=python_requires, 37 | install_requies=install_requires, 38 | packages=setuptools.find_packages(), 39 | ) 40 | -------------------------------------------------------------------------------- /taichi_render_gpu/taichi_three/__init__.py: -------------------------------------------------------------------------------- 1 | import taichi_glsl as ts 2 | 3 | from .version import version as __version__ 4 | from .scene import * 5 | from .model import * 6 | from .scatter import * 7 | from .geometry import * 8 | from .loader import * 9 | from .raycast import * 10 | from .meshgen import * 11 | from .light import * 12 | -------------------------------------------------------------------------------- /taichi_render_gpu/taichi_three/common.py: -------------------------------------------------------------------------------- 1 | class AutoInit: 2 | def init(self): 3 | if not hasattr(self, '_AutoInit_had_init'): 4 | self._init() 5 | self._AutoInit_had_init = True 6 | 7 | def _init(self): 8 | raise NotImplementedError 9 | -------------------------------------------------------------------------------- /taichi_render_gpu/taichi_three/geometry.py: -------------------------------------------------------------------------------- 1 | import math 2 | import taichi as ti 3 | import taichi_glsl as ts 4 | 5 | 6 | @ti.func 7 | def render_triangle(model, camera, face): 8 | scene = model.scene 9 | L2W = model.L2W 10 | _1 = ti.static(min(1, model.faces.m - 1)) 11 | _2 = ti.static(min(2, model.faces.m - 1)) 12 | ia, ib, ic = model.vi[face[0, 0]], model.vi[face[1, 0]], model.vi[face[2, 0]] 13 | ca, cb, cc = ia, ib, ic 14 | if model.type[None] >= 1: 15 | ca, cb, cc = model.vc[face[0, 0]], model.vc[face[1, 0]], model.vc[face[2, 0]] 16 | ta, tb, tc = model.vt[face[0, _1]], model.vt[face[1, _1]], model.vt[face[2, _1]] 17 | na, nb, nc = model.vn[face[0, _2]], model.vn[face[1, _2]], model.vn[face[2, _2]] 18 | a = camera.untrans_pos(L2W @ ia) 19 | b = camera.untrans_pos(L2W @ ib) 20 | c = camera.untrans_pos(L2W @ ic) 21 | 22 | # NOTE: the normal computation indicates that # a front-facing face should 23 | # be COUNTER-CLOCKWISE, i.e., glFrontFace(GL_CCW); 24 | # this is to be compatible with obj model loading. 25 | normal = ts.normalize(ts.cross(a - b, a - c)) 26 | pos = (a + b + c) / 3 27 | view_pos = (a + b + c) / 3 28 | shading = (-ts.dot(ts.normalize(view_pos), normal) - 0.5)*2.0 29 | if ti.static(camera.type == camera.ORTHO): 30 | view_pos = ts.vec3(0.0, 0.0, 1.0) 31 | reverse_v = 1 32 | if ti.static(model.reverse == True): 33 | reverse_v = -1 34 | if ts.dot(view_pos, normal)*reverse_v <= 0: 35 | # shading 36 | color = ts.vec3(0.0) 37 | for light in ti.static(scene.lights): 38 | color += scene.opt.render_func(pos, normal, ts.vec3(0.0), light) 39 | color = scene.opt.pre_process(color) 40 | A = camera.uncook(a) 41 | B = camera.uncook(b) 42 | C = camera.uncook(c) 43 | scr_norm = 1 / ts.cross(A - C, B - A) 44 | B_A = (B - A) * scr_norm 45 | C_B = (C - B) * scr_norm 46 | A_C = (A - C) * scr_norm 47 | 48 | W = 1 49 | # screen space bounding box 50 | M, N = int(ti.floor(min(A, B, C) - W)), int(ti.ceil(max(A, B, C) + W)) 51 | M.x, N.x = min(max(M.x, 0), camera.img.shape[0]), min(max(N.x, 0), camera.img.shape[1]) 52 | M.y, N.y = min(max(M.y, 0), camera.img.shape[0]), min(max(N.y, 0), camera.img.shape[1]) 53 | for X in ti.grouped(ti.ndrange((M.x, N.x), (M.y, N.y))): 54 | # barycentric coordinates using the area method 55 | X_A = X - A 56 | w_C = ts.cross(B_A, X_A) 57 | w_B = ts.cross(A_C, X_A) 58 | w_A = 1 - w_C - w_B 59 | # draw 60 | in_screen = w_A >= 0 and w_B >= 0 and w_C >= 0 and 0 < X[0] < camera.img.shape[0] and 0 < X[1] < camera.img.shape[1] 61 | if not in_screen: 62 | continue 63 | zindex = 0.0 64 | if ti.static(model.reverse == True): 65 | zindex = (a.z * w_A + b.z * w_B + c.z * w_C) 66 | else: 67 | zindex = 1 / (a.z * w_A + b.z * w_B + c.z * w_C) 68 | if zindex < ti.atomic_max(camera.zbuf[X], zindex): 69 | continue 70 | 71 | coor = (ta * w_A + tb * w_B + tc * w_C) 72 | if model.type[None] == 1: 73 | camera.img[X] = (ca * w_A + cb * w_B + cc * w_C) 74 | elif model.type[None] == 2: 75 | camera.img[X] = color * (ca * w_A + cb * w_B + cc * w_C) 76 | else: 77 | camera.img[X] = color * model.texSample(coor) 78 | camera.mask[X] = ts.vec3(1) 79 | camera.normal_map[X] = -normal / 2.0 + 0.5 80 | # camera.normal_map[X] = ts.vec3(shading / 2.0 + 0.5) 81 | 82 | @ti.func 83 | def render_particle(model, camera, vertex, radius): 84 | scene = model.scene 85 | L2W = model.L2W 86 | a = camera.untrans_pos(L2W @ vertex) 87 | A = camera.uncook(a) 88 | 89 | M = int(ti.floor(A - radius)) 90 | N = int(ti.ceil(A + radius)) 91 | 92 | for X in ti.grouped(ti.ndrange((M.x, N.x), (M.y, N.y))): 93 | if X.x < 0 or X.x >= camera.res[0] or X.y < 0 or X.y >= camera.res[1]: 94 | continue 95 | if (X - A).norm_sqr() > radius**2: 96 | continue 97 | 98 | camera.img[X] = ts.vec3(1) 99 | -------------------------------------------------------------------------------- /taichi_render_gpu/taichi_three/light.py: -------------------------------------------------------------------------------- 1 | import taichi as ti 2 | import taichi_glsl as ts 3 | from .common import * 4 | import math 5 | 6 | ''' 7 | The base light class represents a directional light. 8 | ''' 9 | @ti.data_oriented 10 | class Light(AutoInit): 11 | 12 | def __init__(self, dir=None, color=None): 13 | dir = dir or [0, 0, 1] 14 | norm = math.sqrt(sum(x ** 2 for x in dir)) 15 | dir = [x / norm for x in dir] 16 | 17 | self.dir_py = [-x for x in dir] 18 | self.color_py = color or [1, 1, 1] 19 | 20 | self.dir = ti.Vector.field(3, ti.float32, ()) 21 | self.color = ti.Vector.field(3, ti.float32, ()) 22 | # store the current light direction in the view space 23 | # so that we don't have to compute it for each vertex 24 | self.viewdir = ti.Vector.field(3, ti.float32, ()) 25 | 26 | def set(self, dir=[0, 0, 1], color=[1, 1, 1]): 27 | norm = math.sqrt(sum(x**2 for x in dir)) 28 | dir = [x / norm for x in dir] 29 | self.dir_py = dir 30 | self.color = color 31 | 32 | def _init(self): 33 | self.dir[None] = self.dir_py 34 | self.color[None] = self.color_py 35 | 36 | @ti.func 37 | def intensity(self, pos): 38 | return 1 39 | 40 | @ti.func 41 | def get_color(self, pos): 42 | return self.color[None] * self.intensity(pos) 43 | 44 | @ti.func 45 | def get_dir(self, pos): 46 | return self.viewdir 47 | 48 | @ti.func 49 | def set_view(self, camera): 50 | self.viewdir[None] = camera.untrans_dir(self.dir[None]) 51 | 52 | 53 | 54 | class PointLight(Light): 55 | pass 56 | 57 | -------------------------------------------------------------------------------- /taichi_render_gpu/taichi_three/loader.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | 3 | 4 | 5 | def _append(faces, indices): 6 | if len(indices) == 4: 7 | faces.append([indices[0], indices[1], indices[2]]) 8 | faces.append([indices[2], indices[3], indices[0]]) 9 | elif len(indices) == 3: 10 | faces.append(indices) 11 | else: 12 | assert False, len(indices) 13 | 14 | 15 | def readobj(path, scale=1): 16 | vi = [] 17 | vt = [] 18 | vn = [] 19 | faces = [] 20 | 21 | with open(path, 'r') as myfile: 22 | lines = myfile.readlines() 23 | 24 | # cache vertices 25 | for line in lines: 26 | try: 27 | type, fields = line.split(maxsplit=1) 28 | fields = [float(_) for _ in fields.split()] 29 | except ValueError: 30 | continue 31 | 32 | if type == 'v': 33 | vi.append(fields) 34 | elif type == 'vt': 35 | vt.append(fields) 36 | elif type == 'vn': 37 | vn.append(fields) 38 | 39 | # cache faces 40 | for line in lines: 41 | try: 42 | type, fields = line.split(maxsplit=1) 43 | fields = fields.split() 44 | except ValueError: 45 | continue 46 | 47 | # line looks like 'f 5/1/1 1/2/1 4/3/1' 48 | # or 'f 314/380/494 382/400/494 388/550/494 506/551/494' for quads 49 | if type != 'f': 50 | continue 51 | 52 | # a field should look like '5/1/1' 53 | # for vertex/vertex UV coords/vertex Normal (indexes number in the list) 54 | # the index in 'f 5/1/1 1/2/1 4/3/1' STARTS AT 1 !!! 55 | 56 | indices = [[int(_) - 1 if _ != '' else 0 for _ in field.split('/')] for field in fields] 57 | 58 | if len(indices) == 4: 59 | faces.append([indices[0], indices[1], indices[2]]) 60 | faces.append([indices[2], indices[3], indices[0]]) 61 | elif len(indices) == 3: 62 | faces.append(indices) 63 | else: 64 | assert False, len(indices) 65 | 66 | ret = {} 67 | ret['vi'] = None if len(vi) == 0 else np.array(vi).astype(np.float32) * scale 68 | ret['vt'] = None if len(vt) == 0 else np.array(vt).astype(np.float32) 69 | ret['vn'] = None if len(vn) == 0 else np.array(vn).astype(np.float32) 70 | ret['f'] = None if len(faces) == 0 else np.array(faces).astype(np.int32) 71 | return ret 72 | -------------------------------------------------------------------------------- /taichi_render_gpu/taichi_three/meshgen.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import taichi as ti 3 | import taichi_glsl as tl 4 | import math 5 | 6 | 7 | def _pre(x): 8 | if not isinstance(x, ti.Matrix): 9 | x = ti.Vector(x) 10 | return x 11 | 12 | 13 | def _ser(foo): 14 | def wrapped(self, *args, **kwargs): 15 | foo(self, *args, **kwargs) 16 | return self 17 | 18 | return wrapped 19 | 20 | 21 | def _mparg(foo): 22 | def wrapped(self, *args): 23 | if len(args) > 1: 24 | return [foo(self, x) for x in args] 25 | else: 26 | return foo(self, args[0]) 27 | 28 | return wrapped 29 | 30 | 31 | class MeshGen: 32 | def __init__(self): 33 | self.v = [] 34 | self.f = [] 35 | 36 | @_ser 37 | def quad(self, a, b, c, d): 38 | a, b, c, d = self.add_v(a, b, c, d) 39 | self.add_f([a, b, c], [c, d, a]) 40 | 41 | @_ser 42 | def cube(self, a, b): 43 | aaa = self.add_v(tl.mix(a, b, tl.D.yyy)) 44 | baa = self.add_v(tl.mix(a, b, tl.D.xyy)) 45 | aba = self.add_v(tl.mix(a, b, tl.D.yxy)) 46 | aab = self.add_v(tl.mix(a, b, tl.D.yyx)) 47 | bba = self.add_v(tl.mix(a, b, tl.D.xxy)) 48 | abb = self.add_v(tl.mix(a, b, tl.D.yxx)) 49 | bab = self.add_v(tl.mix(a, b, tl.D.xyx)) 50 | bbb = self.add_v(tl.mix(a, b, tl.D.xxx)) 51 | 52 | self.add_f4([aaa, aba, bba, baa]) # back 53 | self.add_f4([aab, bab, bbb, abb]) # front 54 | self.add_f4([aaa, aab, abb, aba]) # left 55 | self.add_f4([baa, bba, bbb, bab]) # right 56 | self.add_f4([aaa, baa, bab, aab]) # bottom 57 | self.add_f4([aba, abb, bbb, bba]) # top 58 | 59 | @_ser 60 | def cylinder(self, bottom, top, dir1, dir2, N): 61 | bottom = _pre(bottom) 62 | top = _pre(top) 63 | dir1 = _pre(dir1) 64 | dir2 = _pre(dir2) 65 | 66 | B, T = [], [] 67 | for i in range(N): 68 | disp = tl.mat(dir1.entries, dir2.entries).T() @ tl.vecAngle(tl.math.tau * i / N) 69 | B.append(self.add_v(bottom + disp)) 70 | T.append(self.add_v(top + disp)) 71 | 72 | BC = self.add_v(bottom) 73 | TC = self.add_v(top) 74 | 75 | for i in range(N): 76 | j = (i + 1) % N 77 | self.add_f4([B[i], B[j], T[j], T[i]]) 78 | 79 | for i in range(N): 80 | j = (i + 1) % N 81 | self.add_f([B[j], B[i], BC]) 82 | self.add_f([T[i], T[j], TC]) 83 | 84 | @_ser 85 | def tri(self, a, b, c): 86 | a, b, c = self.add_v(a, b, c) 87 | self.add_f([a, b, c]) 88 | 89 | @_mparg 90 | def add_v(self, v): 91 | if isinstance(v, ti.Matrix): 92 | v = v.entries 93 | ret = len(self.v) 94 | self.v.append(v) 95 | return ret 96 | 97 | @_mparg 98 | def add_f(self, f): 99 | ret = len(self.f) 100 | self.f.append(f) 101 | return ret 102 | 103 | @_mparg 104 | def add_f4(self, f): 105 | a, b, c, d = f 106 | return self.add_f([a, b, c], [c, d, a]) 107 | 108 | 109 | def __getitem__(self, key): 110 | if key == 'v': 111 | return np.array(self.v) 112 | if key == 'f': 113 | return np.array(self.f) 114 | -------------------------------------------------------------------------------- /taichi_render_gpu/taichi_three/model.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import taichi as ti 3 | import taichi_glsl as ts 4 | from .geometry import * 5 | from .transform import * 6 | from .common import * 7 | import math 8 | 9 | 10 | @ti.data_oriented 11 | class Model(AutoInit): 12 | TEX = 0 13 | COLOR = 1 14 | 15 | def __init__(self, f_n=None, f_m=None, 16 | vi_n=None, vt_n=None, vn_n=None, tex_n=None, col_n=None, 17 | obj=None, tex=None): 18 | self.L2W = Affine.field(()) 19 | 20 | self.faces = None 21 | self.vi = None 22 | self.vt = None 23 | self.vn = None 24 | self.tex = None 25 | self.type = ti.field(dtype=ti.int32, shape=()) 26 | self.reverse = False 27 | 28 | if obj is not None: 29 | f_n = None if obj['f'] is None else obj['f'].shape[0] 30 | vi_n = None if obj['vi'] is None else obj['vi'].shape[0] 31 | vt_n = None if obj['vt'] is None else obj['vt'].shape[0] 32 | vn_n = None if obj['vn'] is None else obj['vn'].shape[0] 33 | 34 | if tex is not None: 35 | tex_n = tex.shape[:2] 36 | 37 | if f_m is None: 38 | f_m = 1 39 | if vt_n is not None: 40 | f_m = 2 41 | if vn_n is not None: 42 | f_m = 3 43 | 44 | if vi_n is None: 45 | vi_n = 1 46 | if vt_n is None: 47 | vt_n = 1 48 | if vn_n is None: 49 | vn_n = 1 50 | if col_n is None: 51 | col_n = 1 52 | 53 | if f_n is not None: 54 | self.faces = ti.Matrix.field(3, f_m, ti.i32, f_n) 55 | if vi_n is not None: 56 | self.vi = ti.Vector.field(3, ti.f32, vi_n) 57 | if vt_n is not None: 58 | self.vt = ti.Vector.field(2, ti.f32, vt_n) 59 | if vn_n is not None: 60 | self.vn = ti.Vector.field(3, ti.f32, vn_n) 61 | if tex_n is not None: 62 | self.tex = ti.Vector.field(3, ti.f32, tex_n) 63 | if col_n is not None: 64 | self.vc = ti.Vector.field(3, ti.f32, col_n) 65 | 66 | if obj is not None: 67 | self.init_obj = obj 68 | if tex is not None: 69 | self.init_tex = tex 70 | 71 | def from_obj(self, obj): 72 | if obj['f'] is not None: 73 | self.faces.from_numpy(obj['f']) 74 | if obj['vi'] is not None: 75 | self.vi.from_numpy(obj['vi']) 76 | if obj['vt'] is not None: 77 | self.vt.from_numpy(obj['vt']) 78 | if obj['vn'] is not None: 79 | self.vn.from_numpy(obj['vn']) 80 | 81 | def _init(self): 82 | self.L2W.init() 83 | self.type[None] = 0 84 | if hasattr(self, 'init_obj'): 85 | self.from_obj(self.init_obj) 86 | if hasattr(self, 'init_tex'): 87 | self.tex.from_numpy(self.init_tex.astype(np.float32) / 255) 88 | 89 | @ti.func 90 | def render(self, camera): 91 | for i in ti.grouped(self.faces): 92 | render_triangle(self, camera, self.faces[i]) 93 | 94 | @ti.func 95 | def texSample(self, coor): 96 | if ti.static(self.tex is not None): 97 | return ts.bilerp(self.tex, coor * ts.vec(*self.tex.shape)) 98 | else: 99 | return 1 100 | 101 | 102 | @ti.data_oriented 103 | class StaticModel(AutoInit): 104 | TEX = 0 105 | COLOR = 1 106 | 107 | def __init__(self, N, f_m=None, col_n=None, 108 | obj=None, tex=None): 109 | self.L2W = Affine.field(()) 110 | 111 | self.faces = None 112 | self.vi = None 113 | self.vt = None 114 | self.vn = None 115 | self.tex = None 116 | # 0 origin 1 pure color 2 shader color 117 | self.type = ti.field(dtype=ti.int32, shape=()) 118 | self.f_n = ti.field(dtype=ti.int32, shape=()) 119 | self.reverse = False 120 | self.N = N 121 | 122 | if obj is not None: 123 | f_n = None if obj['f'] is None else obj['f'].shape[0] 124 | vi_n = None if obj['vi'] is None else obj['vi'].shape[0] 125 | vt_n = None if obj['vt'] is None else obj['vt'].shape[0] 126 | vn_n = None if obj['vn'] is None else obj['vn'].shape[0] 127 | 128 | if not (tex is None): 129 | tex_n = tex.shape[:2] 130 | else: 131 | tex_n = None 132 | 133 | if f_m is None: 134 | f_m = 1 135 | if vt_n is not None: 136 | f_m = 2 137 | if vn_n is not None: 138 | f_m = 3 139 | 140 | if vi_n is None: 141 | vi_n = 1 142 | if vt_n is None: 143 | vt_n = 1 144 | if vn_n is None: 145 | vn_n = 1 146 | if col_n is None: 147 | col_n = 1 148 | 149 | if f_n is not None: 150 | self.faces = ti.Matrix.field(3, f_m, ti.i32, N) 151 | if vi_n is not None: 152 | self.vi = ti.Vector.field(3, ti.f32, N) 153 | if vt_n is not None: 154 | self.vt = ti.Vector.field(2, ti.f32, N) 155 | if vn_n is not None: 156 | self.vn = ti.Vector.field(3, ti.f32, N) 157 | if not (tex_n is None): 158 | self.tex = ti.Vector.field(3, ti.f32, tex_n) 159 | if col_n is not None: 160 | self.vc = ti.Vector.field(3, ti.f32, N) 161 | 162 | if obj is not None: 163 | self.init_obj = obj 164 | if tex is not None: 165 | self.init_tex = tex 166 | 167 | def modify_color(self, color): 168 | s_color = np.zeros((self.N, 3)).astype(np.float32) 169 | s_color[:color.shape[0]] = color 170 | self.vc.from_numpy(s_color) 171 | 172 | def from_obj(self, obj): 173 | N = self.N 174 | if obj['f'] is not None: 175 | s_faces = np.zeros((N, obj['f'].shape[1], obj['f'].shape[2])).astype(int) 176 | s_faces[:obj['f'].shape[0]] = obj['f'] 177 | self.f_n[None] = obj['f'].shape[0] 178 | self.faces.from_numpy(s_faces) 179 | if obj['vi'] is not None: 180 | s_vi = np.zeros((N, 3)).astype(np.float32) 181 | s_vi[:obj['vi'].shape[0]] = obj['vi'][:, :3] 182 | self.vi.from_numpy(s_vi) 183 | if obj['vt'] is not None: 184 | s_vt = np.zeros((N, 2)).astype(np.float32) 185 | s_vt[:obj['vt'].shape[0]] = obj['vt'] 186 | self.vt.from_numpy(s_vt) 187 | if obj['vn'] is not None: 188 | s_vn = np.zeros((N, 3)).astype(np.float32) 189 | s_vn[:obj['vn'].shape[0]] = obj['vn'] 190 | self.vn.from_numpy(s_vn) 191 | 192 | def _init(self): 193 | self.L2W.init() 194 | self.type[None] = 0 195 | if hasattr(self, 'init_obj'): 196 | self.from_obj(self.init_obj) 197 | if hasattr(self, 'init_tex') and (self.init_tex is not None): 198 | self.tex.from_numpy(self.init_tex.astype(np.float32) / 255) 199 | 200 | @ti.func 201 | def render(self, camera): 202 | for i in ti.grouped(self.faces): 203 | render_triangle(self, camera, self.faces[i]) 204 | 205 | @ti.func 206 | def texSample(self, coor): 207 | if ti.static(self.tex is not None): 208 | return ts.bilerp(self.tex, coor * ts.vec(*self.tex.shape)) 209 | else: 210 | return 1 211 | -------------------------------------------------------------------------------- /taichi_render_gpu/taichi_three/raycast.py: -------------------------------------------------------------------------------- 1 | import taichi as ti 2 | import taichi_glsl as ts 3 | from .scene import * 4 | import math 5 | 6 | 7 | EPS = 1e-3 8 | INF = 1e3 9 | 10 | 11 | @ti.data_oriented 12 | class ObjectRT(ts.TaichiClass): 13 | @ti.func 14 | def calc_sdf(self, p): 15 | ret = INF 16 | for I in ti.grouped(ti.ndrange(*self.pos.shape())): 17 | ret = min(ret, self.make_one(I).do_calc_sdf(p)) 18 | return ret 19 | 20 | @ti.func 21 | def intersect(self, orig, dir): 22 | ret, normal = INF, ts.vec3(0.0) 23 | for I in ti.grouped(ti.ndrange(*self.pos.shape())): 24 | t, n = self.make_one(I).do_intersect(orig, dir) 25 | if t < ret: 26 | ret, normal = t, n 27 | return ret, normal 28 | 29 | def do_calc_sdf(self, p): 30 | raise NotImplementedError 31 | 32 | def do_intersect(self, orig, dir): 33 | raise NotImplementedError 34 | 35 | 36 | @ti.data_oriented 37 | class Ball(ObjectRT): 38 | def __init__(self, pos, radius): 39 | self.pos = pos 40 | self.radius = radius 41 | 42 | @ti.func 43 | def make_one(self, I): 44 | return Ball(self.pos[I], self.radius[I]) 45 | 46 | @ti.func 47 | def do_calc_sdf(self, p): 48 | return ts.distance(self.pos, p) - self.radius 49 | 50 | @ti.func 51 | def do_intersect(self, orig, dir): 52 | op = self.pos - orig 53 | b = op.dot(dir) 54 | det = b ** 2 - op.norm_sqr() + self.radius ** 2 55 | ret = INF 56 | if det > 0.0: 57 | det = ti.sqrt(det) 58 | t = b - det 59 | if t > EPS: 60 | ret = t 61 | else: 62 | t = b + det 63 | if t > EPS: 64 | ret = t 65 | return ret, ts.normalize(dir * ret - op) 66 | 67 | 68 | @ti.data_oriented 69 | class SceneRTBase(Scene): 70 | def __init__(self): 71 | super(SceneRTBase, self).__init__() 72 | self.balls = [] 73 | 74 | def trace(self, pos, dir): 75 | raise NotImplementedError 76 | 77 | @ti.func 78 | def color_at(self, coor, camera): 79 | orig, dir = camera.generate(coor) 80 | 81 | pos, normal = self.trace(orig, dir) 82 | light_dir = self.light_dir[None] 83 | 84 | color = self.opt.render_func(pos, normal, dir, light_dir) 85 | color = self.opt.pre_process(color) 86 | return color 87 | 88 | @ti.kernel 89 | def _render(self): 90 | if ti.static(len(self.cameras)): 91 | for camera in ti.static(self.cameras): 92 | for I in ti.grouped(camera.img): 93 | coor = self.cook_coor(I, camera) 94 | color = self.color_at(coor, camera) 95 | camera.img[I] = color 96 | 97 | def add_ball(self, pos, radius): 98 | b = Ball(pos, radius) 99 | self.balls.append(b) 100 | 101 | 102 | @ti.data_oriented 103 | class SceneRT(SceneRTBase): 104 | @ti.func 105 | def intersect(self, orig, dir): 106 | ret, normal = INF, ts.vec3(0.0) 107 | for b in ti.static(self.balls): 108 | t, n = b.intersect(orig, dir) 109 | if t < ret: 110 | ret, normal = t, n 111 | return ret, normal 112 | 113 | @ti.func 114 | def trace(self, orig, dir): 115 | depth, normal = self.intersect(orig, dir) 116 | pos = orig + dir * depth 117 | return pos, normal 118 | 119 | 120 | @ti.data_oriented 121 | class SceneSDF(SceneRTBase): 122 | @ti.func 123 | def calc_sdf(self, p): 124 | ret = INF 125 | for b in ti.static(self.balls): 126 | ret = min(ret, b.calc_sdf(p)) 127 | return ret 128 | 129 | @ti.func 130 | def calc_grad(self, p): 131 | return ts.vec( 132 | self.calc_sdf(p + ts.vec(EPS, 0, 0)), 133 | self.calc_sdf(p + ts.vec(0, EPS, 0)), 134 | self.calc_sdf(p + ts.vec(0, 0, EPS))) 135 | 136 | @ti.func 137 | def trace(self, orig, dir): 138 | pos = orig 139 | color = ts.vec3(0.0) 140 | normal = ts.vec3(0.0) 141 | for s in range(100): 142 | t = self.calc_sdf(pos) 143 | if t <= 0: 144 | normal = ts.normalize(self.calc_grad(pos) - t) 145 | break 146 | pos += dir * t 147 | return pos, normal 148 | -------------------------------------------------------------------------------- /taichi_render_gpu/taichi_three/scatter.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import taichi as ti 3 | import taichi_glsl as ts 4 | from .geometry import * 5 | from .transform import * 6 | from .common import * 7 | import math 8 | 9 | 10 | @ti.data_oriented 11 | class ScatterModel(AutoInit): 12 | def __init__(self, num=None, radius=2): 13 | self.L2W = Affine.field(()) 14 | 15 | self.num = num 16 | self.radius = radius 17 | 18 | if num is not None: 19 | self.particles = ti.Vector.field(3, ti.i32, num) 20 | 21 | def _init(self): 22 | self.L2W.init() 23 | 24 | @ti.func 25 | def render(self, camera): 26 | for i in ti.grouped(self.particles): 27 | render_particle(self, camera, self.particles[i], self.radius) 28 | -------------------------------------------------------------------------------- /taichi_render_gpu/taichi_three/scene.py: -------------------------------------------------------------------------------- 1 | import taichi as ti 2 | import taichi_glsl as ts 3 | from .transform import * 4 | from .shading import * 5 | from .light import * 6 | 7 | 8 | @ti.data_oriented 9 | class Scene(AutoInit): 10 | def __init__(self): 11 | self.lights = [] 12 | self.cameras = [] 13 | self.opt = Shading() 14 | self.models = [] 15 | 16 | def set_light_dir(self, ldir): 17 | # changes light direction input to the direction 18 | # from the light towards the object 19 | # to be consistent with future light types 20 | if not self.lights: 21 | light = Light(ldir) 22 | self.add_light(light) 23 | else: 24 | self.light[0].set(ldir) 25 | 26 | @ti.func 27 | def cook_coor(self, I, camera): 28 | scale = ti.static(2 / min(*camera.img.shape())) 29 | coor = (I - ts.vec2(*camera.img.shape()) / 2) * scale 30 | return coor 31 | 32 | @ti.func 33 | def uncook_coor(self, coor, camera): 34 | scale = ti.static(min(*camera.img.shape()) / 2) 35 | I = coor.xy * scale + ts.vec2(*camera.img.shape()) / 2 36 | return I 37 | 38 | def add_model(self, model): 39 | model.scene = self 40 | self.models.append(model) 41 | 42 | def add_camera(self, camera): 43 | camera.scene = self 44 | self.cameras.append(camera) 45 | 46 | def add_light(self, light): 47 | light.scene = self 48 | self.lights.append(light) 49 | 50 | def _init(self): 51 | for light in self.lights: 52 | light.init() 53 | for camera in self.cameras: 54 | camera.init() 55 | for model in self.models: 56 | model.init() 57 | 58 | 59 | @ti.kernel 60 | def _single_render(self, num : ti.template()): 61 | if ti.static(len(self.cameras)): 62 | for camera in ti.static(self.cameras): 63 | camera.clear_buffer() 64 | # sets up light directions 65 | for light in ti.static(self.lights): 66 | light.set_view(camera) 67 | self.models[num].render(camera) 68 | 69 | def single_render(self, num): 70 | for light in self.lights: 71 | light.init() 72 | for camera in self.cameras: 73 | camera.init() 74 | self.models[num].init() 75 | self._single_render(num) 76 | 77 | def render(self): 78 | self.init() 79 | self._render() 80 | 81 | @ti.kernel 82 | def _render(self): 83 | if ti.static(len(self.cameras)): 84 | for camera in ti.static(self.cameras): 85 | camera.clear_buffer() 86 | # sets up light directions 87 | for light in ti.static(self.lights): 88 | light.set_view(camera) 89 | if ti.static(len(self.models)): 90 | for model in ti.static(self.models): 91 | model.render(camera) 92 | -------------------------------------------------------------------------------- /taichi_render_gpu/taichi_three/shading.py: -------------------------------------------------------------------------------- 1 | import taichi as ti 2 | import taichi_glsl as ts 3 | from .transform import * 4 | import math 5 | 6 | class Shading: 7 | def __init__(self, **kwargs): 8 | self.is_normal_map = False 9 | self.lambert = 0.58 10 | self.half_lambert = 0.04 11 | self.blinn_phong = 0.3 12 | self.phong = 0.0 13 | self.shineness = 10 14 | self.__dict__.update(kwargs) 15 | 16 | @ti.func 17 | def render_func(self, pos, normal, dir, light): 18 | color = ts.vec3(0.0) 19 | light_dir = light.get_dir(pos) 20 | shineness = self.shineness 21 | half_lambert = ts.dot(normal, light_dir) * 0.5 + 0.5 22 | lambert = max(0, ts.dot(normal, light_dir)) 23 | blinn_phong = ts.dot(normal, ts.mix(light_dir, -dir, 0.5)) 24 | blinn_phong = pow(max(blinn_phong, 0), shineness) 25 | refl_dir = ts.reflect(light_dir, normal) 26 | phong = -ts.dot(normal, refl_dir) 27 | phong = pow(max(phong, 0), shineness) 28 | 29 | strength = 0.0 30 | if ti.static(self.lambert != 0.0): 31 | strength += lambert * self.lambert 32 | if ti.static(self.half_lambert != 0.0): 33 | strength += half_lambert * self.half_lambert 34 | if ti.static(self.blinn_phong != 0.0): 35 | strength += blinn_phong * self.blinn_phong 36 | if ti.static(self.phong != 0.0): 37 | strength += phong * self.phong 38 | color = ts.vec3(strength) 39 | 40 | if ti.static(self.is_normal_map): 41 | color = normal * 0.5 + 0.5 42 | return color * light.get_color(pos) 43 | 44 | @ti.func 45 | def pre_process(self, color): 46 | blue = ts.vec3(0.00, 0.01, 0.05) 47 | orange = ts.vec3(1.19, 1.04, 0.98) 48 | return ti.sqrt(ts.mix(blue, orange, color)) 49 | -------------------------------------------------------------------------------- /taichi_render_gpu/taichi_three/transform.py: -------------------------------------------------------------------------------- 1 | import taichi as ti 2 | import taichi_glsl as ts 3 | from .common import * 4 | import math 5 | 6 | 7 | def rotationX(angle): 8 | return [ 9 | [1, 0, 0], 10 | [0, math.cos(angle), -math.sin(angle)], 11 | [0, math.sin(angle), math.cos(angle)], 12 | ] 13 | 14 | def rotationY(angle): 15 | return [ 16 | [ math.cos(angle), 0, math.sin(angle)], 17 | [ 0, 1, 0], 18 | [-math.sin(angle), 0, math.cos(angle)], 19 | ] 20 | 21 | def rotationZ(angle): 22 | return [ 23 | [math.cos(angle), -math.sin(angle), 0], 24 | [math.sin(angle), math.cos(angle), 0], 25 | [ 0, 0, 1], 26 | ] 27 | 28 | 29 | @ti.data_oriented 30 | class Affine(ts.TaichiClass, AutoInit): 31 | @property 32 | def matrix(self): 33 | return self.entries[0] 34 | 35 | @property 36 | def offset(self): 37 | return self.entries[1] 38 | 39 | @classmethod 40 | def _field(cls, shape=None): 41 | return ti.Matrix.field(3, 3, ti.f32, shape), ti.Vector.field(3, ti.f32, shape) 42 | 43 | @ti.func 44 | def loadIdentity(self): 45 | self.matrix = [[1, 0, 0], [0, 1, 0], [0, 0, 1]] 46 | self.offset = [0, 0, 0] 47 | 48 | @ti.kernel 49 | def _init(self): 50 | self.loadIdentity() 51 | 52 | @ti.func 53 | def __matmul__(self, other): 54 | return self.matrix @ other + self.offset 55 | 56 | @ti.func 57 | def inverse(self): 58 | # TODO: incorrect: 59 | return Affine(self.matrix.inverse(), -self.offset) 60 | 61 | def loadOrtho(self, fwd=[0, 0, 1], up=[0, 1, 0]): 62 | # fwd = target - pos 63 | # fwd = fwd.normalized() 64 | fwd_len = math.sqrt(sum(x**2 for x in fwd)) 65 | fwd = [x / fwd_len for x in fwd] 66 | # right = fwd.cross(up) 67 | right = [ 68 | fwd[2] * up[1] - fwd[1] * up[2], 69 | fwd[0] * up[2] - fwd[2] * up[0], 70 | fwd[1] * up[0] - fwd[0] * up[1], 71 | ] 72 | # right = right.normalized() 73 | right_len = math.sqrt(sum(x**2 for x in right)) 74 | right = [x / right_len for x in right] 75 | # up = right.cross(fwd) 76 | up = [ 77 | right[2] * fwd[1] - right[1] * fwd[2], 78 | right[0] * fwd[2] - right[2] * fwd[0], 79 | right[1] * fwd[0] - right[0] * fwd[1], 80 | ] 81 | 82 | # trans = ti.Matrix.cols([right, up, fwd]) 83 | trans = [right, up, fwd] 84 | trans = [[trans[i][j] for i in range(3)] for j in range(3)] 85 | self.matrix[None] = trans 86 | 87 | def from_mouse(self, mpos): 88 | if isinstance(mpos, ti.GUI): 89 | if mpos.is_pressed(ti.GUI.LMB): 90 | mpos = mpos.get_cursor_pos() 91 | else: 92 | mpos = (0, 0) 93 | a, t = mpos 94 | if a != 0 or t != 0: 95 | a, t = a * math.tau - math.pi, t * math.pi - math.pi / 2 96 | c = math.cos(t) 97 | self.loadOrtho(fwd=[c * math.sin(a), math.sin(t), c * math.cos(a)]) 98 | 99 | 100 | @ti.data_oriented 101 | class Camera(AutoInit): 102 | ORTHO = 'Orthogonal' 103 | TAN_FOV = 'Tangent Perspective' # rectilinear perspective 104 | COS_FOV = 'Cosine Perspective' # curvilinear perspective, see en.wikipedia.org/wiki/Curvilinear_perspective 105 | 106 | def __init__(self, res=None, fx=None, fy=None, cx=None, cy=None, 107 | pos=[0, 0, -2], target=[0, 0, 0], up=[0, 1, 0], fov=30): 108 | self.res = res or (512, 512) 109 | self.img = ti.Vector.field(3, ti.f32, self.res) 110 | self.zbuf = ti.field(ti.f32, self.res) 111 | self.mask = ti.Vector.var(3, ti.f32, self.res) 112 | self.normal_map = ti.Vector.var(3, ti.f32, self.res) 113 | self.trans = ti.Matrix.field(3, 3, ti.f32, ()) 114 | self.pos = ti.Vector.field(3, ti.f32, ()) 115 | self.target = ti.Vector.field(3, ti.f32, ()) 116 | self.intrinsic = ti.Matrix.field(3, 3, ti.f32, ()) 117 | self.type = self.TAN_FOV 118 | self.fov = math.radians(fov) 119 | 120 | self.cx = cx or self.res[0] // 2 121 | self.cy = cy or self.res[1] // 2 122 | self.fx = fx or self.cx / math.tan(self.fov) 123 | self.fy = fy or self.cy / math.tan(self.fov) 124 | # python scope camera transformations 125 | self.pos_py = pos 126 | self.target_py = target 127 | self.trans_py = None 128 | self.up_py = up 129 | self.set(init=True) 130 | # mouse position for camera control 131 | self.mpos = (0, 0) 132 | 133 | def set_intrinsic(self, fx=None, fy=None, cx=None, cy=None): 134 | # see http://ais.informatik.uni-freiburg.de/teaching/ws09/robotics2/pdfs/rob2-08-camera-calibration.pdf 135 | self.fx = fx or self.fx 136 | self.fy = fy or self.fy 137 | self.cx = self.cx if cx is None else cx 138 | self.cy = self.cy if cy is None else cy 139 | 140 | ''' 141 | NOTE: taichi_three uses a LEFT HANDED coordinate system. 142 | that is, the +Z axis points FROM the camera TOWARDS the scene, 143 | with X, Y being device coordinates 144 | ''' 145 | def set(self, pos=None, target=None, up=None, init=False): 146 | pos = self.pos_py if pos is None else pos 147 | target = self.target_py if target is None else target 148 | up = self.up_py if up is None else up 149 | # fwd = target - pos 150 | fwd = [target[i] - pos[i] for i in range(3)] 151 | # fwd = fwd.normalized() 152 | fwd_len = math.sqrt(sum(x**2 for x in fwd)) 153 | fwd = [x / fwd_len for x in fwd] 154 | # right = fwd.cross(up) 155 | right = [ 156 | fwd[2] * up[1] - fwd[1] * up[2], 157 | fwd[0] * up[2] - fwd[2] * up[0], 158 | fwd[1] * up[0] - fwd[0] * up[1], 159 | ] 160 | # right = right.normalized() 161 | right_len = math.sqrt(sum(x**2 for x in right)) 162 | right = [x / right_len for x in right] 163 | # up = right.cross(fwd) 164 | up = [ 165 | right[2] * fwd[1] - right[1] * fwd[2], 166 | right[0] * fwd[2] - right[2] * fwd[0], 167 | right[1] * fwd[0] - right[0] * fwd[1], 168 | ] 169 | 170 | # trans = ti.Matrix.cols([right, up, fwd]) 171 | trans = [right, up, fwd] 172 | self.trans_py = [[trans[i][j] for i in range(3)] for j in range(3)] 173 | self.pos_py = pos 174 | self.target_py = target 175 | if not init: 176 | self.pos[None] = self.pos_py 177 | self.trans[None] = self.trans_py 178 | self.target[None] = self.target_py 179 | 180 | def set_extrinsic(self, trans, pose): 181 | trans = [[trans[i][j] for j in range(3)] for i in range(3)] 182 | self.trans_py = trans 183 | self.pos_py = pose 184 | 185 | def _init(self): 186 | self.pos[None] = self.pos_py 187 | self.trans[None] = self.trans_py 188 | self.target[None] = self.target_py 189 | self.intrinsic[None][0, 0] = self.fx 190 | self.intrinsic[None][0, 2] = self.cx 191 | self.intrinsic[None][1, 1] = self.fy 192 | self.intrinsic[None][1, 2] = self.cy 193 | self.intrinsic[None][2, 2] = 1.0 194 | 195 | @ti.func 196 | def clear_buffer(self): 197 | for I in ti.grouped(self.img): 198 | self.img[I] = ts.vec3(0.0) 199 | self.zbuf[I] = 0.0 200 | self.mask[I] = ts.vec3(0.0) 201 | self.normal_map[I] = ts.vec3(0.5) 202 | 203 | def from_mouse(self, gui): 204 | is_alter_move = gui.is_pressed(ti.GUI.CTRL) 205 | if gui.is_pressed(ti.GUI.LMB): 206 | mpos = gui.get_cursor_pos() 207 | if self.mpos != (0, 0): 208 | self.orbit((mpos[0] - self.mpos[0], mpos[1] - self.mpos[1]), 209 | pov=is_alter_move) 210 | self.mpos = mpos 211 | elif gui.is_pressed(ti.GUI.RMB): 212 | mpos = gui.get_cursor_pos() 213 | if self.mpos != (0, 0): 214 | self.zoom_by_mouse(mpos, (mpos[0] - self.mpos[0], mpos[1] - self.mpos[1]), 215 | dolly=is_alter_move) 216 | self.mpos = mpos 217 | elif gui.is_pressed(ti.GUI.MMB): 218 | mpos = gui.get_cursor_pos() 219 | if self.mpos != (0, 0): 220 | self.pan((mpos[0] - self.mpos[0], mpos[1] - self.mpos[1])) 221 | self.mpos = mpos 222 | else: 223 | if gui.event and gui.event.key == ti.GUI.WHEEL: 224 | # one mouse wheel unit is (0, 120) 225 | self.zoom(-gui.event.delta[1] / 1200, 226 | dolly=is_alter_move) 227 | gui.event = None 228 | mpos = (0, 0) 229 | self.mpos = mpos 230 | 231 | 232 | def orbit(self, delta, sensitivity=5, pov=False): 233 | ds, dt = delta 234 | if ds != 0 or dt != 0: 235 | dis = math.sqrt(sum((self.target_py[i] - self.pos_py[i]) ** 2 for i in range(3))) 236 | fov = self.fov 237 | ds, dt = ds * fov * sensitivity, dt * fov * sensitivity 238 | newdir = ts.vec3(ds, dt, 1).normalized() 239 | newdir = [sum(self.trans[None][i, j] * newdir[j] for j in range(3))\ 240 | for i in range(3)] 241 | if pov: 242 | newtarget = [self.pos_py[i] + dis * newdir[i] for i in range(3)] 243 | self.set(target=newtarget) 244 | else: 245 | newpos = [self.target_py[i] - dis * newdir[i] for i in range(3)] 246 | self.set(pos=newpos) 247 | 248 | def zoom_by_mouse(self, pos, delta, sensitivity=3, dolly=False): 249 | ds, dt = delta 250 | if ds != 0 or dt != 0: 251 | z = math.sqrt(ds ** 2 + dt ** 2) * sensitivity 252 | if (pos[0] - 0.5) * ds + (pos[1] - 0.5) * dt > 0: 253 | z *= -1 254 | self.zoom(z, dolly) 255 | 256 | def zoom(self, z, dolly=False): 257 | newpos = [(1 + z) * self.pos_py[i] - z * self.target_py[i] for i in range(3)] 258 | if dolly: 259 | newtarget = [z * self.pos_py[i] + (1 - z) * self.target_py[i] for i in range(3)] 260 | self.set(pos=newpos, target=newtarget) 261 | else: 262 | self.set(pos=newpos) 263 | 264 | def pan(self, delta, sensitivity=3): 265 | ds, dt = delta 266 | if ds != 0 or dt != 0: 267 | dis = math.sqrt(sum((self.target_py[i] - self.pos_py[i]) ** 2 for i in range(3))) 268 | fov = self.fov 269 | ds, dt = ds * fov * sensitivity, dt * fov * sensitivity 270 | newdir = ts.vec3(-ds, -dt, 1).normalized() 271 | newdir = [sum(self.trans[None][i, j] * newdir[j] for j in range(3))\ 272 | for i in range(3)] 273 | newtarget = [self.pos_py[i] + dis * newdir[i] for i in range(3)] 274 | newpos = [self.pos_py[i] + newtarget[i] - self.target_py[i] for i in range(3)] 275 | self.set(pos=newpos, target=newtarget) 276 | 277 | @ti.func 278 | def trans_pos(self, pos): 279 | return self.trans[None] @ pos + self.pos[None] 280 | 281 | @ti.func 282 | def trans_dir(self, pos): 283 | return self.trans[None] @ pos 284 | 285 | @ti.func 286 | def untrans_pos(self, pos): 287 | return self.trans[None].inverse() @ (pos - self.pos[None]) 288 | 289 | @ti.func 290 | def untrans_dir(self, pos): 291 | return self.trans[None].inverse() @ pos 292 | 293 | @ti.func 294 | def uncook(self, pos): 295 | if ti.static(self.type == self.ORTHO): 296 | pos[0] *= self.intrinsic[None][0, 0] 297 | pos[1] *= self.intrinsic[None][1, 1] 298 | pos[0] += self.intrinsic[None][0, 2] 299 | pos[1] += self.intrinsic[None][1, 2] 300 | elif ti.static(self.type == self.TAN_FOV): 301 | pos = self.intrinsic[None] @ pos 302 | pos[0] /= abs(pos[2]) 303 | pos[1] /= abs(pos[2]) 304 | else: 305 | raise NotImplementedError("Curvilinear projection matrix not implemented!") 306 | return ts.vec2(pos[0], pos[1]) 307 | 308 | def export_intrinsic(self): 309 | import numpy as np 310 | intrinsic = np.zeros((3, 3)) 311 | intrinsic[0, 0] = self.fx 312 | intrinsic[1, 1] = self.fy 313 | intrinsic[0, 2] = self.cx 314 | intrinsic[1, 2] = self.cy 315 | intrinsic[2, 2] = 1 316 | return intrinsic 317 | 318 | def export_extrinsic(self): 319 | import numpy as np 320 | trans = np.array(self.trans_py) 321 | pos = np.array(self.pos_py) 322 | extrinsic = np.zeros((3, 4)) 323 | 324 | trans = np.transpose(trans) 325 | for i in range(3): 326 | for j in range(3): 327 | extrinsic[i][j] = trans[i, j] 328 | pos = -trans @ pos 329 | for i in range(3): 330 | extrinsic[i][3] = pos[i] 331 | return extrinsic 332 | 333 | @ti.func 334 | def generate(self, coor): 335 | fov = ti.static(self.fov) 336 | tan_fov = ti.static(math.tan(fov)) 337 | 338 | orig = ts.vec3(0.0) 339 | dir = ts.vec3(0.0, 0.0, 1.0) 340 | 341 | if ti.static(self.type == self.ORTHO): 342 | orig = ts.vec3(coor, 0.0) 343 | elif ti.static(self.type == self.TAN_FOV): 344 | uv = coor * fov 345 | dir = ts.normalize(ts.vec3(uv, 1)) 346 | elif ti.static(self.type == self.COS_FOV): 347 | uv = coor * fov 348 | dir = ts.vec3(ti.sin(uv), ti.cos(uv.norm())) 349 | 350 | orig = self.trans_pos(orig) 351 | dir = self.trans_dir(dir) 352 | 353 | return orig, dir 354 | -------------------------------------------------------------------------------- /taichi_render_gpu/taichi_three/version.py: -------------------------------------------------------------------------------- 1 | version = (0, 0, 4) 2 | taiglsl_version = (0, 0, 9) 3 | taichi_version = (0, 6, 28) 4 | 5 | print(f'[Tai3D] version {".".join(map(str, version))}') 6 | print(f'[Tai3D] Inputs are welcomed at https://github.com/taichi-dev/taichi_three') 7 | print(f'[Tai3D] Camera control hints: LMB to orbit, MMB to move, RMB to scale') 8 | -------------------------------------------------------------------------------- /train.sh: -------------------------------------------------------------------------------- 1 | python apps/train_dmc.py --config configs/dmc_train.yaml --dataroot dataset/example --name=train \ 2 | --smpl_path=assets/smplx --obj_path=assets/obj --coarse_part 3 | 4 | # python apps/train_dmc.py --config configs/dmc_train.yaml --dataroot dataset/example --name=train \ 5 | # --smpl_path=assets/smplx --obj_path=assets/obj --fine_part 6 | 7 | # python apps/train_normal_net.py --config configs/normal_train.yaml --dataroot dataset/example --------------------------------------------------------------------------------