├── README.md ├── assets ├── perspnet.pdf ├── result.jpg └── test.jpg ├── data ├── __init__.py ├── arkit_dataset.py ├── augmentation.py ├── base_dataset.py ├── biwi_dataset.py ├── biwi_preprocess.py └── triangles.npy ├── demo.py ├── lib ├── Sim3DR │ ├── Sim3DR.py │ ├── Sim3DR_Cython.cpython-36m-x86_64-linux-gnu.so │ ├── __init__.py │ ├── _init_paths.py │ ├── build_sim3dr.sh │ ├── lib │ │ ├── rasterize.cpp │ │ ├── rasterize.h │ │ ├── rasterize.pyx │ │ └── rasterize_kernel.cpp │ ├── lighting.py │ ├── readme.md │ ├── setup.py │ └── tests │ │ ├── io.cpp │ │ ├── io.h │ │ └── test.cpp ├── eyemouth_index.py ├── mesh │ ├── __init__.py │ ├── cython │ │ ├── mesh_core.cpp │ │ ├── mesh_core.h │ │ ├── mesh_core_cython.cpp │ │ ├── mesh_core_cython.cpython-36m-x86_64-linux-gnu.so │ │ ├── mesh_core_cython.pyx │ │ └── setup.py │ ├── io.py │ ├── light.py │ ├── readme.md │ ├── render.py │ ├── transform.py │ └── vis.py └── mesh_p │ ├── __init__.py │ ├── cython │ ├── mesh_core.cpp │ ├── mesh_core.h │ ├── mesh_core_cython.cpp │ ├── mesh_core_cython.pyx │ └── setup.py │ ├── io.py │ ├── light.py │ ├── readme.md │ ├── render.py │ ├── transform.py │ └── vis.py ├── models ├── __init__.py ├── base_model.py ├── loss.py ├── model_repository.py ├── network.py ├── networks.py ├── perspnet_model.py └── resnet.py ├── npy ├── kpt_ind.npy ├── projection_matrix.txt ├── tris_2388x3.npy ├── tris_2500x3_202110.npy └── uv_coords_std_202109.npy ├── options ├── __init__.py ├── base_options.py ├── test_options.py └── train_options.py ├── requirement.txt ├── test.py ├── train.py └── util ├── __init__.py ├── box_utils.py ├── face_deformnet_utils.py ├── get_data.py ├── html.py ├── image_pool.py ├── prior_box.py ├── py_cpu_nms.py ├── renderer.py ├── util.py └── visualizer.py /README.md: -------------------------------------------------------------------------------- 1 | # Toward 3D Face Reconstruction in Perspective Projection: Estimating 6DoF Face Pose From Monocular Image 2 | 3 | [![Paper](https://img.shields.io/badge/arXiv-Paper-b31b1b.svg)](https://arxiv.org/pdf/2205.04126v2.pdf) 4 | 5 | This repository contains the official PyTorch implementation of: 6 | 7 | **Toward 3D Face Reconstruction in Perspective Projection: Estimating 6DoF Face Pose From Monocular Image** 8 | 9 | ![](assets/perspnet.pdf) 10 | 11 | ## Installation 12 | 13 | **Setup python environment** 14 | 15 | * **CUDA>=10.1** 16 | * Python = 3.7 17 | * PyTorch >= 1.7.1 18 | 19 | **Install other dependencies** 20 | 21 | pip install -r requirement.txt 22 | 23 | **Install libs** 24 | 25 | 1. Please follow `./lib/mesh/README.md`. 26 | 27 | 2. Please follow `./lib/mesh_p/README.md`. 28 | 29 | 3. Please follow `./lib/Sim3DR/README.md`. 30 | 31 | ## Demo 32 | 33 | Download our pretraind model from [here](https://drive.google.com/file/d/1K9rAmQ7Hduz1on9SWnyLHf_ZxfLwlTOy/view?usp=drive_link) and put it in the folder `./checkpoint/run1/latest_net_R.pth`. 34 | 35 | ```sh 36 | $ python demo.py 37 | ``` 38 | 39 | ## Download Dataset 40 | 41 | ** Download ARKitFace dataset** 42 | 43 | Please contact email cbsropenproject@outlook.com. We will send you an agreement for signing, and after our verification, the download link for ARKitFace dataset will be sent to you. Please ensure compliance with agreement and do not use this dataset for any non research purposes. For more information about the dataset, please refer to this [here](https://github.com/o0Helloworld0o-ustc/ECCV2022_WCPA_track2). 44 | 45 | Then put the download dataset in the folder `./dataset/ARKitFace`. 46 | 47 | ** Download BIWI dataset** 48 | 49 | Please download the BIWI dataset from https://www.kaggle.com/datasets/kmader/biwi-kinect-head-pose-database and transform GT pose to RGB image, and make a biwi_test.csv file based on that of ARKitFace dataset and the code of./data/biwi_preprocess.py. 50 | 51 | 52 | ## Testing and training 53 | 54 | Download our pretraind model from [here](https://drive.google.com/file/d/1K9rAmQ7Hduz1on9SWnyLHf_ZxfLwlTOy/view?usp=drive_link) and put it in the folder `./checkpoint/run1/latest_net_R.pth`. 55 | 56 | **Testing for ARKitFace dataset** 57 | 58 | ```sh 59 | $ python -u test.py --csv_path_test test.csv 60 | ``` 61 | 62 | 63 | **Testing for BIWI dataset** 64 | 65 | ```sh 66 | $ python -u test.py --csv_path_test biwi_test.csv --dataset_mode biwi 67 | ``` 68 | 69 | 70 | **Training** 71 | 72 | ```sh 73 | $ python -u train.py --csv_path_train train.csv --csv_path_test test.csv 74 | ``` 75 | 76 | ## Citation 77 | 78 | If you find our work usefull in your research, please use the following BibTex entry. 79 | 80 | ```latex 81 | @ARTICLE{10127617, 82 | author={Kao, Yueying and Pan, Bowen and Xu, Miao and Lyu, Jiangjing and Zhu, Xiangyu and Chang, Yuanzhang and Li, Xiaobo and Lei, Zhen}, 83 | journal={IEEE Transactions on Image Processing}, 84 | title={Toward 3D Face Reconstruction in Perspective Projection: Estimating 6DoF Face Pose From Monocular Image}, 85 | year={2023}, 86 | volume={32}, 87 | number={}, 88 | pages={3080-3091}, 89 | doi={10.1109/TIP.2023.3275535}} 90 | ``` 91 | 92 | 93 | 94 | 95 | 96 | 97 | -------------------------------------------------------------------------------- /assets/perspnet.pdf: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/cbsropenproject/6dof_face/c80ca80125b2059e8f2984688814d59d0ac8e10c/assets/perspnet.pdf -------------------------------------------------------------------------------- /assets/result.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/cbsropenproject/6dof_face/c80ca80125b2059e8f2984688814d59d0ac8e10c/assets/result.jpg -------------------------------------------------------------------------------- /assets/test.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/cbsropenproject/6dof_face/c80ca80125b2059e8f2984688814d59d0ac8e10c/assets/test.jpg -------------------------------------------------------------------------------- /data/__init__.py: -------------------------------------------------------------------------------- 1 | """This package includes all the modules related to data loading and preprocessing 2 | 3 | To add a custom dataset class called 'dummy', you need to add a file called 'dummy_dataset.py' and define a subclass 'DummyDataset' inherited from BaseDataset. 4 | You need to implement four functions: 5 | -- <__init__>: initialize the class, first call BaseDataset.__init__(self, opt). 6 | -- <__len__>: return the size of dataset. 7 | -- <__getitem__>: get a data point from data loader. 8 | -- : (optionally) add dataset-specific options and set default options. 9 | 10 | Now you can use the dataset class by specifying flag '--dataset_mode dummy'. 11 | See our template dataset class 'template_dataset.py' for more details. 12 | """ 13 | import importlib 14 | import torch.utils.data 15 | from data.base_dataset import BaseDataset 16 | 17 | 18 | def find_dataset_using_name(dataset_name): 19 | """Import the module "data/[dataset_name]_dataset.py". 20 | 21 | In the file, the class called DatasetNameDataset() will 22 | be instantiated. It has to be a subclass of BaseDataset, 23 | and it is case-insensitive. 24 | """ 25 | dataset_filename = "data." + dataset_name + "_dataset" 26 | datasetlib = importlib.import_module(dataset_filename) 27 | 28 | dataset = None 29 | target_dataset_name = dataset_name.replace('_', '') + 'dataset' 30 | for name, cls in datasetlib.__dict__.items(): 31 | if name.lower() == target_dataset_name.lower() \ 32 | and issubclass(cls, BaseDataset): 33 | dataset = cls 34 | 35 | if dataset is None: 36 | raise NotImplementedError("In %s.py, there should be a subclass of BaseDataset with class name that matches %s in lowercase." % (dataset_filename, target_dataset_name)) 37 | 38 | return dataset 39 | 40 | 41 | def get_option_setter(dataset_name): 42 | """Return the static method of the dataset class.""" 43 | dataset_class = find_dataset_using_name(dataset_name) 44 | return dataset_class.modify_commandline_options 45 | 46 | 47 | def create_dataset(opt): 48 | """Create a dataset given the option. 49 | 50 | This function wraps the class CustomDatasetDataLoader. 51 | This is the main interface between this package and 'train.py'/'test.py' 52 | 53 | Example: 54 | >>> from data import create_dataset 55 | >>> dataset = create_dataset(opt) 56 | """ 57 | data_loader = CustomDatasetDataLoader(opt) 58 | dataset = data_loader.load_data() 59 | return dataset 60 | 61 | 62 | class CustomDatasetDataLoader(): 63 | """Wrapper class of Dataset class that performs multi-threaded data loading""" 64 | 65 | def __init__(self, opt): 66 | """Initialize this class 67 | 68 | Step 1: create a dataset instance given the name [dataset_mode] 69 | Step 2: create a multi-threaded data loader. 70 | """ 71 | self.opt = opt 72 | dataset_class = find_dataset_using_name(opt.dataset_mode) 73 | self.dataset = dataset_class(opt) 74 | print("dataset [%s] was created" % type(self.dataset).__name__) 75 | self.dataloader = torch.utils.data.DataLoader( 76 | self.dataset, 77 | batch_size=opt.batch_size, 78 | shuffle=not opt.serial_batches, 79 | num_workers=int(opt.num_threads), 80 | drop_last=True, 81 | ) 82 | #drop_last=True if opt.isTrain else False, 83 | 84 | def set_epoch(self, epoch): 85 | self.dataset.current_epoch = epoch 86 | 87 | def load_data(self): 88 | return self 89 | 90 | def __len__(self): 91 | """Return the number of data in the dataset""" 92 | return min(len(self.dataset), self.opt.max_dataset_size) 93 | 94 | def __iter__(self): 95 | """Return a batch of data""" 96 | for i, data in enumerate(self.dataloader): 97 | if i * self.opt.batch_size >= self.opt.max_dataset_size: 98 | break 99 | yield data 100 | -------------------------------------------------------------------------------- /data/base_dataset.py: -------------------------------------------------------------------------------- 1 | """This module implements an abstract base class (ABC) 'BaseDataset' for datasets. 2 | 3 | It also includes common transformation functions (e.g., get_transform, __scale_width), which can be later used in subclasses. 4 | """ 5 | import random 6 | import numpy as np 7 | import torch.utils.data as data 8 | from PIL import Image 9 | import torchvision.transforms as transforms 10 | from abc import ABC, abstractmethod 11 | 12 | 13 | class BaseDataset(data.Dataset, ABC): 14 | """This class is an abstract base class (ABC) for datasets. 15 | 16 | To create a subclass, you need to implement the following four functions: 17 | -- <__init__>: initialize the class, first call BaseDataset.__init__(self, opt). 18 | -- <__len__>: return the size of dataset. 19 | -- <__getitem__>: get a data point. 20 | -- : (optionally) add dataset-specific options and set default options. 21 | """ 22 | 23 | def __init__(self, opt): 24 | """Initialize the class; save the options in the class 25 | 26 | Parameters: 27 | opt (Option class)-- stores all the experiment flags; needs to be a subclass of BaseOptions 28 | """ 29 | self.opt = opt 30 | self.root = opt.dataroot 31 | self.current_epoch = 0 32 | 33 | @staticmethod 34 | def modify_commandline_options(parser, is_train): 35 | """Add new dataset-specific options, and rewrite default values for existing options. 36 | 37 | Parameters: 38 | parser -- original option parser 39 | is_train (bool) -- whether training phase or test phase. You can use this flag to add training-specific or test-specific options. 40 | 41 | Returns: 42 | the modified parser. 43 | """ 44 | return parser 45 | 46 | @abstractmethod 47 | def __len__(self): 48 | """Return the total number of images in the dataset.""" 49 | return 0 50 | 51 | @abstractmethod 52 | def __getitem__(self, index): 53 | """Return a data point and its metadata information. 54 | 55 | Parameters: 56 | index - - a random integer for data indexing 57 | 58 | Returns: 59 | a dictionary of data with their names. It ususally contains the data itself and its metadata information. 60 | """ 61 | pass 62 | 63 | 64 | def get_params(opt, size): 65 | w, h = size 66 | new_h = h 67 | new_w = w 68 | if opt.preprocess == 'resize_and_crop': 69 | new_h = new_w = opt.load_size 70 | elif opt.preprocess == 'scale_width_and_crop': 71 | new_w = opt.load_size 72 | new_h = opt.load_size * h // w 73 | 74 | x = random.randint(0, np.maximum(0, new_w - opt.crop_size)) 75 | y = random.randint(0, np.maximum(0, new_h - opt.crop_size)) 76 | 77 | flip = random.random() > 0.5 78 | 79 | return {'crop_pos': (x, y), 'flip': flip} 80 | 81 | 82 | def get_transform(opt, params=None, grayscale=False, method=Image.BICUBIC, convert=True): 83 | transform_list = [] 84 | if grayscale: 85 | transform_list.append(transforms.Grayscale(1)) 86 | if 'fixsize' in opt.preprocess: 87 | transform_list.append(transforms.Resize(params["size"], method)) 88 | if 'resize' in opt.preprocess: 89 | osize = [opt.load_size, opt.load_size] 90 | if "gta2cityscapes" in opt.dataroot: 91 | osize[0] = opt.load_size // 2 92 | transform_list.append(transforms.Resize(osize, method)) 93 | elif 'scale_width' in opt.preprocess: 94 | transform_list.append(transforms.Lambda(lambda img: __scale_width(img, opt.load_size, opt.crop_size, method))) 95 | elif 'scale_shortside' in opt.preprocess: 96 | transform_list.append(transforms.Lambda(lambda img: __scale_shortside(img, opt.load_size, opt.crop_size, method))) 97 | 98 | if 'zoom' in opt.preprocess: 99 | if params is None: 100 | transform_list.append(transforms.Lambda(lambda img: __random_zoom(img, opt.load_size, opt.crop_size, method))) 101 | else: 102 | transform_list.append(transforms.Lambda(lambda img: __random_zoom(img, opt.load_size, opt.crop_size, method, factor=params["scale_factor"]))) 103 | 104 | if 'crop' in opt.preprocess: 105 | if params is None or 'crop_pos' not in params: 106 | transform_list.append(transforms.RandomCrop(opt.crop_size)) 107 | else: 108 | transform_list.append(transforms.Lambda(lambda img: __crop(img, params['crop_pos'], opt.crop_size))) 109 | 110 | if 'patch' in opt.preprocess: 111 | transform_list.append(transforms.Lambda(lambda img: __patch(img, params['patch_index'], opt.crop_size))) 112 | 113 | if 'trim' in opt.preprocess: 114 | transform_list.append(transforms.Lambda(lambda img: __trim(img, opt.crop_size))) 115 | 116 | # if opt.preprocess == 'none': 117 | transform_list.append(transforms.Lambda(lambda img: __make_power_2(img, base=4, method=method))) 118 | 119 | if not opt.no_flip: 120 | if params is None or 'flip' not in params: 121 | transform_list.append(transforms.RandomHorizontalFlip()) 122 | elif 'flip' in params: 123 | transform_list.append(transforms.Lambda(lambda img: __flip(img, params['flip']))) 124 | 125 | if convert: 126 | transform_list += [transforms.ToTensor()] 127 | if grayscale: 128 | transform_list += [transforms.Normalize((0.5,), (0.5,))] 129 | else: 130 | transform_list += [transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))] 131 | return transforms.Compose(transform_list) 132 | 133 | 134 | def __make_power_2(img, base, method=Image.BICUBIC): 135 | ow, oh = img.size 136 | h = int(round(oh / base) * base) 137 | w = int(round(ow / base) * base) 138 | if h == oh and w == ow: 139 | return img 140 | 141 | return img.resize((w, h), method) 142 | 143 | 144 | def __random_zoom(img, target_width, crop_width, method=Image.BICUBIC, factor=None): 145 | if factor is None: 146 | zoom_level = np.random.uniform(0.8, 1.0, size=[2]) 147 | else: 148 | zoom_level = (factor[0], factor[1]) 149 | iw, ih = img.size 150 | zoomw = max(crop_width, iw * zoom_level[0]) 151 | zoomh = max(crop_width, ih * zoom_level[1]) 152 | img = img.resize((int(round(zoomw)), int(round(zoomh))), method) 153 | return img 154 | 155 | 156 | def __scale_shortside(img, target_width, crop_width, method=Image.BICUBIC): 157 | ow, oh = img.size 158 | shortside = min(ow, oh) 159 | if shortside >= target_width: 160 | return img 161 | else: 162 | scale = target_width / shortside 163 | return img.resize((round(ow * scale), round(oh * scale)), method) 164 | 165 | 166 | def __trim(img, trim_width): 167 | ow, oh = img.size 168 | if ow > trim_width: 169 | xstart = np.random.randint(ow - trim_width) 170 | xend = xstart + trim_width 171 | else: 172 | xstart = 0 173 | xend = ow 174 | if oh > trim_width: 175 | ystart = np.random.randint(oh - trim_width) 176 | yend = ystart + trim_width 177 | else: 178 | ystart = 0 179 | yend = oh 180 | return img.crop((xstart, ystart, xend, yend)) 181 | 182 | 183 | def __scale_width(img, target_width, crop_width, method=Image.BICUBIC): 184 | ow, oh = img.size 185 | if ow == target_width and oh >= crop_width: 186 | return img 187 | w = target_width 188 | h = int(max(target_width * oh / ow, crop_width)) 189 | return img.resize((w, h), method) 190 | 191 | 192 | def __crop(img, pos, size): 193 | ow, oh = img.size 194 | x1, y1 = pos 195 | tw = th = size 196 | if (ow > tw or oh > th): 197 | return img.crop((x1, y1, x1 + tw, y1 + th)) 198 | return img 199 | 200 | 201 | def __patch(img, index, size): 202 | ow, oh = img.size 203 | nw, nh = ow // size, oh // size 204 | roomx = ow - nw * size 205 | roomy = oh - nh * size 206 | startx = np.random.randint(int(roomx) + 1) 207 | starty = np.random.randint(int(roomy) + 1) 208 | 209 | index = index % (nw * nh) 210 | ix = index // nh 211 | iy = index % nh 212 | gridx = startx + ix * size 213 | gridy = starty + iy * size 214 | return img.crop((gridx, gridy, gridx + size, gridy + size)) 215 | 216 | 217 | def __flip(img, flip): 218 | if flip: 219 | return img.transpose(Image.FLIP_LEFT_RIGHT) 220 | return img 221 | 222 | 223 | def __print_size_warning(ow, oh, w, h): 224 | """Print warning information about image size(only print once)""" 225 | if not hasattr(__print_size_warning, 'has_printed'): 226 | print("The image size needs to be a multiple of 4. " 227 | "The loaded image size was (%d, %d), so it was adjusted to " 228 | "(%d, %d). This adjustment will be done to all images " 229 | "whose sizes are not multiples of 4" % (ow, oh, w, h)) 230 | __print_size_warning.has_printed = True 231 | -------------------------------------------------------------------------------- /data/biwi_dataset.py: -------------------------------------------------------------------------------- 1 | import pdb 2 | import cv2 3 | import numpy as np 4 | import pandas as pd 5 | from PIL import Image 6 | from pathlib import Path 7 | 8 | import torch 9 | import torchvision.transforms as transforms 10 | 11 | from data.base_dataset import BaseDataset 12 | from data.augmentation import EulerAugmentor, HorizontalFlipAugmentor 13 | 14 | use_jpeg4py = False 15 | data_root = Path('./dataset/ARKitFace') 16 | 17 | 18 | npz_path = 'npy/kpt_ind.npy' 19 | kpt_ind = np.load(npz_path) 20 | 21 | 22 | 23 | class BIWIDataset(BaseDataset): 24 | 25 | def __init__(self, opt): 26 | 27 | self.data_root = data_root 28 | self.opt = opt 29 | self.is_train = False 30 | self.img_size = opt.img_size 31 | 32 | csv_path = data_root / 'csv/metadata_biwi.csv' 33 | 34 | self.df = pd.read_csv(csv_path, dtype={'subject_id': str, 'facial_action': str, 'img_id': str}, 35 | nrows=2721 if opt.debug else None) 36 | self.df['data_batch'] += '_v1' 37 | 38 | 39 | img_mean = np.array([0.485, 0.456, 0.406]) 40 | img_std = np.array([0.229, 0.224, 0.225]) 41 | self.tfm_train = transforms.Compose([ 42 | transforms.ToTensor(), 43 | transforms.Normalize(mean=img_mean, std=img_std) 44 | ]) 45 | self.tfm_test = transforms.Compose([ 46 | transforms.ToTensor(), 47 | transforms.Normalize(mean=img_mean, std=img_std) 48 | ]) 49 | 50 | self.kpt_ind = kpt_ind 51 | self.dst_pts = np.float32([ 52 | [0, 0], 53 | [0, opt.img_size - 1], 54 | [opt.img_size - 1, 0] 55 | ]) 56 | 57 | 58 | 59 | 60 | 61 | def __getitem__(self, index): 62 | data_batch = str(self.df['data_batch'][index]) 63 | subject_id = str(self.df['subject_id'][index]) 64 | facial_action = str(self.df['facial_action'][index]) 65 | img_id = str(self.df['img_id'][index]) 66 | 67 | img_path = self.data_root / 'image' / data_batch / subject_id / facial_action / f'{img_id}_ar.jpg' 68 | npz_path = self.data_root / 'info' / data_batch / subject_id / facial_action / f'{img_id}_info.npz' 69 | M = np.load(npz_path) 70 | 71 | if use_jpeg4py: 72 | img_raw = jpeg.JPEG(img_path).decode() 73 | else: 74 | img_raw = cv2.imread(str(img_path)) 75 | img_raw = cv2.cvtColor(img_raw, cv2.COLOR_BGR2RGB) 76 | 77 | img_h, img_w, _ = img_raw.shape 78 | 79 | 80 | 81 | 82 | if self.opt.use_gt_bbox: 83 | points2d_68 = M['points2d'][self.kpt_ind] # (68, 2) 84 | else: 85 | raise ValueError('please add --use_gt_bbox') 86 | 87 | 88 | 89 | x_min, y_min = points2d_68.min(axis=0) 90 | x_max, y_max = points2d_68.max(axis=0) 91 | x_center = (x_min + x_max) / 2 92 | y_center = (y_min + y_max) / 2 93 | w, h = x_max - x_min, y_max - y_min 94 | size = max(w, h) 95 | ss = np.array([0.75, 0.75, 0.75, 0.75]) 96 | 97 | if self.is_train: 98 | rnd_size = 0.15 * size 99 | dx = np.random.uniform(-rnd_size, rnd_size) 100 | dy = np.random.uniform(-rnd_size, rnd_size) 101 | x_center, y_center = x_center + dx, y_center + dy 102 | ss *= np.random.uniform(0.95, 1.05) 103 | 104 | left = x_center - ss[0] * size 105 | right = x_center + ss[1] * size 106 | top = y_center - ss[2] * size 107 | bottom = y_center + ss[3] * size 108 | 109 | 110 | src_pts = np.float32([ 111 | [left, top], 112 | [left, bottom], 113 | [right, top] 114 | ]) 115 | 116 | tform = cv2.getAffineTransform(src_pts, self.dst_pts) 117 | img = cv2.warpAffine(img_raw, tform, (self.img_size,)*2) 118 | 119 | if self.is_train: 120 | img = self.tfm_train(Image.fromarray(img)) 121 | else: 122 | img = self.tfm_test(Image.fromarray(img)) 123 | 124 | 125 | 126 | points2d_68 = M['points2d'][self.kpt_ind] # 127 | W, b = tform.T[:2], tform.T[2] 128 | local_points2d_68 = points2d_68 @ W + b 129 | pts68 = torch.tensor(local_points2d_68, dtype=torch.float32) / self.img_size * 2 - 1 # 范围[-1, 1] 130 | 131 | 132 | verts3d = torch.tensor(M['verts_gt'], dtype=torch.float32) * 4.5 * 2 # [-1, 1] 133 | choose = torch.tensor(np.random.rand(1024).astype(np.int64)) 134 | mask = torch.tensor(np.zeros((192,192)).astype(np.int64)) 135 | 136 | nocs = torch.tensor(np.random.rand(1024,3).astype(np.float32)) 137 | model = torch.tensor(np.random.rand(1220,3).astype(np.float32)) ### for biwi, please change it to GT model! 138 | corr_mat = torch.tensor(np.random.rand(1024,1220).astype(np.float32)) 139 | uvmap = torch.tensor(np.random.rand(192,192,3), dtype=torch.float32).permute(2, 0, 1) 140 | d = { 141 | 'img': img, 142 | 'choose': choose, 143 | 'model': model, 144 | 'uvmap': uvmap, 145 | 'nocs': nocs, 146 | 'mask': mask, 147 | 'corr_mat': corr_mat 148 | } 149 | 150 | if hasattr(self.opt, 'eval'): 151 | tform_inv = cv2.getAffineTransform(self.dst_pts, src_pts) 152 | d['tform_inv'] = tform_inv 153 | d['R_t'] = M['R_t'] 154 | d['img_path'] = str(img_path) 155 | 156 | return d 157 | 158 | 159 | 160 | def __len__(self): 161 | return len(self.df) 162 | 163 | 164 | 165 | 166 | 167 | def compute_metrics(self, inference_data): 168 | bs_list = np.array(inference_data['batch_size']) 169 | loss1 = np.array(inference_data['loss_total']) 170 | loss2 = np.array(inference_data['loss_corr']) 171 | loss3 = np.array(inference_data['loss_recon3d']) 172 | loss4 = np.array(inference_data['loss_uv']) 173 | loss5 = np.array(inference_data['loss_mat']) 174 | loss6 = np.array(inference_data['loss_seg']) 175 | loss_total = (loss1 * bs_list).sum() / bs_list.sum() 176 | loss_corr = (loss2 * bs_list).sum() / bs_list.sum() 177 | loss_recon3d = (loss3 * bs_list).sum() / bs_list.sum() 178 | loss_entropy = (loss4 * bs_list).sum() / bs_list.sum() 179 | loss_mat = (loss5 * bs_list).sum() / bs_list.sum() 180 | loss_seg = (loss6 * bs_list).sum() / bs_list.sum() 181 | d = { 182 | 'loss_total': loss_total, 183 | 'loss_corr': loss_corr, 184 | 'loss_recon3d': loss_recon3d, 185 | 'loss_entropy': loss_entropy, 186 | 'loss_mat': loss_mat, 187 | 'loss_seg': loss_seg 188 | } 189 | if hasattr(self.opt, 'eval'): 190 | d['pnp_fail'] = inference_data['pnp_fail'] 191 | d['3DRecon'] = np.mean(inference_data['3DRecon']) 192 | d['ADD'] = np.mean(inference_data['ADD']) 193 | d['pitch_mae'] = np.mean(inference_data['pitch_mae']) 194 | d['yaw_mae'] = np.mean(inference_data['yaw_mae']) 195 | d['roll_mae'] = np.mean(inference_data['roll_mae']) 196 | d['tx_mae'] = np.mean(inference_data['tx_mae']) 197 | d['ty_mae'] = np.mean(inference_data['ty_mae']) 198 | d['tz_mae'] = np.mean(inference_data['tz_mae']) 199 | d['5°5cm'] = inference_data['strict_success'] / inference_data['total_count'] 200 | d['5°10cm'] = inference_data['easy_success'] / inference_data['total_count'] 201 | d['mean_IoU'] = np.mean(inference_data['IoU']) 202 | 203 | # 单位转换 204 | d['3DRecon'] = '%.2f mm' % (d['3DRecon'] * 1000) 205 | d['ADD'] = '%.2f mm' % (d['ADD'] * 1000) 206 | d['pitch_mae'] = '%.2f °' % d['pitch_mae'] 207 | d['yaw_mae'] = '%.2f °' % d['yaw_mae'] 208 | d['roll_mae'] = '%.2f °' % d['roll_mae'] 209 | d['tx_mae'] = '%.2f mm' % (d['tx_mae'] * 1000) 210 | d['ty_mae'] = '%.2f mm' % (d['ty_mae'] * 1000) 211 | d['tz_mae'] = '%.2f mm' % (d['tz_mae'] * 1000) 212 | d['5°5cm'] = '%.2f' % (d['5°5cm'] * 100) 213 | d['5°10cm'] = '%.2f' % (d['5°10cm'] * 100) 214 | d['mean_IoU'] = '%.4f' % d['mean_IoU'] 215 | 216 | return d 217 | -------------------------------------------------------------------------------- /data/triangles.npy: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/cbsropenproject/6dof_face/c80ca80125b2059e8f2984688814d59d0ac8e10c/data/triangles.npy -------------------------------------------------------------------------------- /demo.py: -------------------------------------------------------------------------------- 1 | import matplotlib 2 | matplotlib.use('Agg') 3 | import os 4 | import sys 5 | import cv2 6 | import numpy as np 7 | from PIL import Image 8 | import torch 9 | import torchvision 10 | import torchvision.transforms as transforms 11 | from options.test_options import TestOptions 12 | from models import create_model 13 | from util.renderer import Renderer 14 | import torch.nn.functional as F 15 | 16 | renderer = Renderer(img_size=800) 17 | 18 | def resize_img(img, bbox, target_size=None): 19 | img_h, img_w, _ = img.shape 20 | if img_w > img_h: 21 | max_size = img_w 22 | diff = img_w - img_h 23 | if diff % 2 == 1: 24 | pad1 = (diff + 1) // 2 25 | pad2 = (diff - 1) // 2 26 | else: 27 | pad1 = diff // 2 28 | pad2 = pad1 29 | img_pad = np.pad(img, ((pad1, pad2), (0, 0), (0, 0)),mode='constant') 30 | bbox[:,1] +=pad1 31 | bbox[:,3] +=pad1 32 | elif img_w == img_h: 33 | img_pad = img 34 | elif img_h > img_w: 35 | max_size = img_h 36 | diff = img_h - img_w 37 | if diff % 2 == 1: 38 | pad1 = (diff + 1) // 2 39 | pad2 = (diff - 1) // 2 40 | else: 41 | pad1 = diff // 2 42 | pad2 = pad1 43 | img_pad = np.pad(img, ((0, 0), (pad1, pad2), (0, 0)),mode='constant') 44 | bbox[:,0] +=pad1 45 | bbox[:,2] +=pad1 46 | 47 | img_resize = cv2.resize(img_pad, (target_size,)*2) 48 | scale = target_size / img_pad.shape[0] 49 | # lms *= scale 50 | bbox*=scale 51 | 52 | return img_resize, bbox 53 | 54 | 55 | npz_path = './npy/kpt_ind.npy' 56 | kpt_ind = np.load(npz_path) 57 | 58 | npy_path = './npy/uv_coords_std_202109.npy' 59 | uv_coords_std = np.load(npy_path)[:1220] 60 | uv_coords = uv_coords_std * 2 - 1 61 | grid = torch.tensor(uv_coords).unsqueeze(0).unsqueeze(1).to('cuda') 62 | 63 | 64 | if __name__ == '__main__': 65 | print('【process_id】', os.getpid()) 66 | print('【command】python -u ' + ' '.join(sys.argv) + '\n') 67 | 68 | opt = TestOptions().parse() # get test options 69 | 70 | model = create_model(opt) 71 | model.setup(opt) 72 | model.netR.eval() 73 | 74 | img_path = 'assets/test.jpg' 75 | save_img_r = 'assets/result.jpg' 76 | #save_img_d = 'assets/test_d.jpg' 77 | 78 | img_raw = cv2.imread(str(img_path)) 79 | img_raw = cv2.cvtColor(img_raw, cv2.COLOR_BGR2RGB) 80 | 81 | bboxes = [[74.90,250.16,104.35,286.82], 82 | [871.42,259.26,900.92,297.38], 83 | [451.48,264.31,482.92,305.67], 84 | [618.63,282.38,648.07,316.87], 85 | [348.44,264.21,375.59,299.17], 86 | [734.70,265.40,763.80,301.03], 87 | [495.54,267.46,523.81,300.98] 88 | ] 89 | bboxes = np.array(bboxes) 90 | num_faces = len(bboxes) 91 | print('\nnum_faces =', num_faces) 92 | 93 | img, bboxes = resize_img(img_raw, bboxes, target_size=800) 94 | 95 | img_mean = np.array([0.485, 0.456, 0.406]) 96 | img_std = np.array([0.229, 0.224, 0.225]) 97 | tfm_test = transforms.Compose([ 98 | transforms.ToTensor(), 99 | transforms.Normalize(mean=img_mean, std=img_std) 100 | ]) 101 | 102 | dst_pts = np.float32([ 103 | [0, 0], 104 | [0, opt.img_size - 1], 105 | [opt.img_size - 1, 0] 106 | ]) 107 | 108 | img_crop_list = [] 109 | img_raw_list = [] 110 | input_tensor_list = [] 111 | input_tensor_raw_list = [] 112 | tform_inv_list = [] 113 | left_list = [] 114 | right_list = [] 115 | top_list = [] 116 | bottom_list = [] 117 | 118 | for index in range(num_faces): 119 | 120 | x_min = bboxes[index][0] 121 | y_min = bboxes[index][1] 122 | x_max = bboxes[index][2] 123 | y_max = bboxes[index][3] 124 | x_center = (x_min + x_max) / 2 125 | y_center = (y_min + y_max) / 2 126 | w, h = x_max - x_min, y_max - y_min 127 | size = max(w, h) 128 | ss = np.array([0.75, 0.75, 0.75, 0.75]) 129 | 130 | left = x_center - ss[0] * size 131 | right = x_center + ss[1] * size 132 | top = y_center - ss[2] * size 133 | bottom = y_center + ss[3] * size 134 | 135 | left_list.append(left) 136 | right_list.append(right) 137 | top_list.append(top) 138 | bottom_list.append(bottom) 139 | 140 | src_pts = np.float32([ 141 | [left, top], 142 | [left, bottom], 143 | [right, top] 144 | ]) 145 | 146 | tform = cv2.getAffineTransform(src_pts, dst_pts) 147 | tform_inv = cv2.getAffineTransform(dst_pts, src_pts) 148 | img_crop = cv2.warpAffine(img, tform, (opt.img_size,) * 2) 149 | img_raw = cv2.resize(img, (opt.img_size * 2, opt.img_size * 2)) 150 | 151 | tform_inv_list.append(tform_inv) 152 | img_crop_list.append(img_crop) 153 | img_raw_list.append(img_raw) 154 | 155 | input_tensor = tfm_test(Image.fromarray(img_crop)).unsqueeze(0) 156 | input_raw_tensor = tfm_test(Image.fromarray(img_raw)).unsqueeze(0) 157 | input_tensor_list.append(input_tensor) 158 | input_tensor_raw_list.append(input_raw_tensor) 159 | 160 | left = left_list 161 | right = right_list 162 | top = top_list 163 | bottom = bottom_list 164 | 165 | img_crop_display = torchvision.utils.make_grid( 166 | torch.tensor(img_crop_list).permute(0, 3, 1, 2), 167 | nrow=4 168 | ).permute(1, 2, 0).numpy() 169 | 170 | if num_faces > 0: 171 | model.img = torch.cat(input_tensor_list, dim=0).to(model.device) 172 | with torch.no_grad(): 173 | _, seg_pred, _, _ = model.netR.cnn(model.img) 174 | bs = seg_pred.shape[0] 175 | mask_pr = torch.argmax(seg_pred, 1).cpu().detach().numpy() 176 | chooses = np.zeros((bs, opt.n_pts)) 177 | 178 | for i in range(bs): 179 | choose = mask_pr[i].flatten().nonzero()[0] 180 | if len(choose) > opt.n_pts: 181 | c_mask = np.zeros(len(choose), dtype=int) 182 | c_mask[:opt.n_pts] = 1 183 | np.random.shuffle(c_mask) 184 | choose = choose[c_mask.nonzero()] 185 | else: 186 | choose = np.pad(choose, (0, opt.n_pts - len(choose)), 'wrap') 187 | chooses[i, :] = choose 188 | 189 | chooses = chooses.astype(np.int64) 190 | chooses = torch.LongTensor(chooses).to(device='cuda') 191 | 192 | assign_mat, seg_pred, uv_pred = model.netR(model.img, chooses) 193 | verts3d_pred = F.grid_sample(uv_pred, grid.expand(bs, -1, -1, -1), align_corners=False) 194 | verts3d_pred = verts3d_pred.squeeze(2).permute(0, 2, 1).contiguous() 195 | 196 | temp2 = img.copy() 197 | temp3 = img.copy() 198 | 199 | img_size = 800 200 | f = 1.574437 * img_size / 2 201 | K_img = np.array([ 202 | [f, 0, img_size / 2.0], 203 | [0, f, img_size / 2.0], 204 | [0, 0, 1] 205 | ]) 206 | T = np.array([ 207 | [1.0, 0, 0, 0], 208 | [0, -1, 0, 0], 209 | [0, 0, -1, 0], 210 | [0, 0, 0, 1] 211 | ]) 212 | 213 | inst_shape = verts3d_pred / 9.0 214 | verts3d_pred = inst_shape.cpu().numpy().reshape(-1, 1220, 3) 215 | assign_mat = F.softmax(assign_mat, dim=2) 216 | nocs_coords = torch.bmm(assign_mat, inst_shape) 217 | nocs_coords = nocs_coords.detach().cpu().numpy().reshape(-1, opt.n_pts, 3) 218 | 219 | for i in range(num_faces): 220 | choose = chooses.cpu().numpy()[i] 221 | choose, choose_idx = np.unique(choose, return_index=True) 222 | nocs_coord = nocs_coords[i, choose_idx, :] 223 | col_idx = choose % opt.img_size 224 | row_idx = choose // opt.img_size 225 | local_pts2d = np.concatenate((col_idx.reshape(-1, 1), row_idx.reshape(-1, 1)), axis=1) 226 | 227 | tform_inv = tform_inv_list[i] 228 | 229 | W, b = tform_inv.T[:2], tform_inv.T[2] 230 | global_pts68_pred = local_pts2d @ W + b 231 | 232 | # for p in global_pts68_pred: 233 | # cv2.circle(temp2, (int(p[0]), int(p[1])), radius=2, color=(0, 255, 0), thickness=-1) 234 | 235 | _, rvecs, tvecs, _ = cv2.solvePnPRansac( 236 | nocs_coord, 237 | global_pts68_pred, 238 | K_img, 239 | None 240 | ) 241 | 242 | rotM = cv2.Rodrigues(rvecs)[0].T 243 | tvecs = tvecs.squeeze(axis=1) 244 | 245 | # to GL style 246 | R_temp = np.identity(4) 247 | R_temp[:3, :3] = rotM 248 | R_temp[3, :3] = tvecs 249 | R_t_pred = R_temp @ T 250 | 251 | temp3 = renderer(verts3d_pred[i], R_t_pred, temp3) 252 | 253 | #temp2 = temp2[:, :, :: -1] 254 | temp3 = temp3[:, :, :: -1] 255 | 256 | cv2.imwrite(save_img_r, temp3) 257 | #cv2.imwrite(save_img_d, temp2) 258 | 259 | 260 | 261 | -------------------------------------------------------------------------------- /lib/Sim3DR/Sim3DR.py: -------------------------------------------------------------------------------- 1 | # coding: utf-8 2 | 3 | from . import _init_paths 4 | import numpy as np 5 | import Sim3DR_Cython 6 | 7 | 8 | def get_normal(vertices, triangles): 9 | normal = np.zeros_like(vertices, dtype=np.float32) 10 | Sim3DR_Cython.get_normal(normal, vertices, triangles, vertices.shape[0], triangles.shape[0]) 11 | return normal 12 | 13 | 14 | def rasterize(vertices, triangles, colors, bg=None, 15 | height=None, width=None, channel=None, 16 | reverse=False): 17 | if bg is not None: 18 | height, width, channel = bg.shape 19 | else: 20 | assert height is not None and width is not None and channel is not None 21 | bg = np.zeros((height, width, channel), dtype=np.uint8) 22 | 23 | buffer = np.zeros((height, width), dtype=np.float32) - 1e8 24 | 25 | if colors.dtype != np.float32: 26 | colors = colors.astype(np.float32) 27 | Sim3DR_Cython.rasterize(bg, vertices, triangles, colors, buffer, triangles.shape[0], height, width, channel, 28 | reverse=reverse) 29 | return bg 30 | -------------------------------------------------------------------------------- /lib/Sim3DR/Sim3DR_Cython.cpython-36m-x86_64-linux-gnu.so: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/cbsropenproject/6dof_face/c80ca80125b2059e8f2984688814d59d0ac8e10c/lib/Sim3DR/Sim3DR_Cython.cpython-36m-x86_64-linux-gnu.so -------------------------------------------------------------------------------- /lib/Sim3DR/__init__.py: -------------------------------------------------------------------------------- 1 | # coding: utf-8 2 | 3 | from .Sim3DR import get_normal, rasterize 4 | from .lighting import RenderPipeline 5 | -------------------------------------------------------------------------------- /lib/Sim3DR/_init_paths.py: -------------------------------------------------------------------------------- 1 | # coding: utf-8 2 | 3 | import os.path as osp 4 | import sys 5 | 6 | 7 | def add_path(path): 8 | if path not in sys.path: 9 | sys.path.insert(0, path) 10 | 11 | 12 | this_dir = osp.dirname(__file__) 13 | lib_path = osp.join(this_dir, '.') 14 | add_path(lib_path) 15 | -------------------------------------------------------------------------------- /lib/Sim3DR/build_sim3dr.sh: -------------------------------------------------------------------------------- 1 | python3 setup.py build_ext --inplace -------------------------------------------------------------------------------- /lib/Sim3DR/lib/rasterize.h: -------------------------------------------------------------------------------- 1 | #ifndef MESH_CORE_HPP_ 2 | #define MESH_CORE_HPP_ 3 | 4 | #include 5 | #include 6 | #include 7 | #include 8 | #include 9 | #include 10 | 11 | using namespace std; 12 | 13 | class Point3D { 14 | public: 15 | float x; 16 | float y; 17 | float z; 18 | 19 | public: 20 | Point3D() : x(0.f), y(0.f), z(0.f) {} 21 | Point3D(float x_, float y_, float z_) : x(x_), y(y_), z(z_) {} 22 | 23 | void initialize(float x_, float y_, float z_){ 24 | this->x = x_; this->y = y_; this->z = z_; 25 | } 26 | 27 | Point3D cross(Point3D &p){ 28 | Point3D c; 29 | c.x = this->y * p.z - this->z * p.y; 30 | c.y = this->z * p.x - this->x * p.z; 31 | c.z = this->x * p.y - this->y * p.x; 32 | return c; 33 | } 34 | 35 | float dot(Point3D &p) { 36 | return this->x * p.x + this->y * p.y + this->z * p.z; 37 | } 38 | 39 | Point3D operator-(const Point3D &p) { 40 | Point3D np; 41 | np.x = this->x - p.x; 42 | np.y = this->y - p.y; 43 | np.z = this->z - p.z; 44 | return np; 45 | } 46 | 47 | }; 48 | 49 | class Point { 50 | public: 51 | float x; 52 | float y; 53 | 54 | public: 55 | Point() : x(0.f), y(0.f) {} 56 | Point(float x_, float y_) : x(x_), y(y_) {} 57 | float dot(Point p) { 58 | return this->x * p.x + this->y * p.y; 59 | } 60 | 61 | Point operator-(const Point &p) { 62 | Point np; 63 | np.x = this->x - p.x; 64 | np.y = this->y - p.y; 65 | return np; 66 | } 67 | 68 | Point operator+(const Point &p) { 69 | Point np; 70 | np.x = this->x + p.x; 71 | np.y = this->y + p.y; 72 | return np; 73 | } 74 | 75 | Point operator*(float s) { 76 | Point np; 77 | np.x = s * this->x; 78 | np.y = s * this->y; 79 | return np; 80 | } 81 | }; 82 | 83 | 84 | bool is_point_in_tri(Point p, Point p0, Point p1, Point p2); 85 | 86 | void get_point_weight(float *weight, Point p, Point p0, Point p1, Point p2); 87 | 88 | void _get_tri_normal(float *tri_normal, float *vertices, int *triangles, int ntri, bool norm_flg); 89 | 90 | void _get_ver_normal(float *ver_normal, float *tri_normal, int *triangles, int nver, int ntri); 91 | 92 | void _get_normal(float *ver_normal, float *vertices, int *triangles, int nver, int ntri); 93 | 94 | void _rasterize_triangles( 95 | float *vertices, int *triangles, float *depth_buffer, int *triangle_buffer, float *barycentric_weight, 96 | int ntri, int h, int w); 97 | 98 | void _rasterize( 99 | unsigned char *image, float *vertices, int *triangles, float *colors, 100 | float *depth_buffer, int ntri, int h, int w, int c, float alpha, bool reverse); 101 | 102 | void _render_texture_core( 103 | float *image, float *vertices, int *triangles, 104 | float *texture, float *tex_coords, int *tex_triangles, 105 | float *depth_buffer, 106 | int nver, int tex_nver, int ntri, 107 | int h, int w, int c, 108 | int tex_h, int tex_w, int tex_c, 109 | int mapping_type); 110 | 111 | void _write_obj_with_colors_texture(string filename, string mtl_name, 112 | float *vertices, int *triangles, float *colors, float *uv_coords, 113 | int nver, int ntri, int ntexver); 114 | 115 | #endif -------------------------------------------------------------------------------- /lib/Sim3DR/lib/rasterize.pyx: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | cimport numpy as np 3 | # from libcpp.string cimport string 4 | cimport cython 5 | from libcpp cimport bool 6 | 7 | # from cpython import bool 8 | 9 | # use the Numpy-C-API from Cython 10 | np.import_array() 11 | 12 | # cdefine the signature of our c function 13 | cdef extern from "rasterize.h": 14 | void _rasterize_triangles( 15 | float*vertices, int*triangles, float*depth_buffer, int*triangle_buffer, float*barycentric_weight, 16 | int ntri, int h, int w 17 | ) 18 | 19 | void _rasterize( 20 | unsigned char*image, float*vertices, int*triangles, float*colors, float*depth_buffer, 21 | int ntri, int h, int w, int c, float alpha, bool reverse 22 | ) 23 | 24 | # void _render_texture_core( 25 | # float* image, float* vertices, int* triangles, 26 | # float* texture, float* tex_coords, int* tex_triangles, 27 | # float* depth_buffer, 28 | # int nver, int tex_nver, int ntri, 29 | # int h, int w, int c, 30 | # int tex_h, int tex_w, int tex_c, 31 | # int mapping_type) 32 | 33 | void _get_tri_normal(float *tri_normal, float *vertices, int *triangles, int nver, bool norm_flg) 34 | void _get_ver_normal(float *ver_normal, float*tri_normal, int*triangles, int nver, int ntri) 35 | void _get_normal(float *ver_normal, float *vertices, int *triangles, int nver, int ntri) 36 | 37 | 38 | # void _write_obj_with_colors_texture(string filename, string mtl_name, 39 | # float* vertices, int* triangles, float* colors, float* uv_coords, 40 | # int nver, int ntri, int ntexver) 41 | 42 | @cython.boundscheck(False) 43 | @cython.wraparound(False) 44 | def get_tri_normal(np.ndarray[float, ndim=2, mode="c"] tri_normal not None, 45 | np.ndarray[float, ndim=2, mode = "c"] vertices not None, 46 | np.ndarray[int, ndim=2, mode="c"] triangles not None, 47 | int ntri, bool norm_flg = False): 48 | _get_tri_normal( np.PyArray_DATA(tri_normal), np.PyArray_DATA(vertices), 49 | np.PyArray_DATA(triangles), ntri, norm_flg) 50 | 51 | @cython.boundscheck(False) # turn off bounds-checking for entire function 52 | @cython.wraparound(False) # turn off negative index wrapping for entire function 53 | def get_ver_normal(np.ndarray[float, ndim=2, mode = "c"] ver_normal not None, 54 | np.ndarray[float, ndim=2, mode = "c"] tri_normal not None, 55 | np.ndarray[int, ndim=2, mode="c"] triangles not None, 56 | int nver, int ntri): 57 | _get_ver_normal( 58 | np.PyArray_DATA(ver_normal), np.PyArray_DATA(tri_normal), np.PyArray_DATA(triangles), 59 | nver, ntri) 60 | 61 | @cython.boundscheck(False) # turn off bounds-checking for entire function 62 | @cython.wraparound(False) # turn off negative index wrapping for entire function 63 | def get_normal(np.ndarray[float, ndim=2, mode = "c"] ver_normal not None, 64 | np.ndarray[float, ndim=2, mode = "c"] vertices not None, 65 | np.ndarray[int, ndim=2, mode="c"] triangles not None, 66 | int nver, int ntri): 67 | _get_normal( 68 | np.PyArray_DATA(ver_normal), np.PyArray_DATA(vertices), np.PyArray_DATA(triangles), 69 | nver, ntri) 70 | 71 | 72 | @cython.boundscheck(False) # turn off bounds-checking for entire function 73 | @cython.wraparound(False) # turn off negative index wrapping for entire function 74 | def rasterize_triangles( 75 | np.ndarray[float, ndim=2, mode = "c"] vertices not None, 76 | np.ndarray[int, ndim=2, mode="c"] triangles not None, 77 | np.ndarray[float, ndim=2, mode = "c"] depth_buffer not None, 78 | np.ndarray[int, ndim=2, mode = "c"] triangle_buffer not None, 79 | np.ndarray[float, ndim=2, mode = "c"] barycentric_weight not None, 80 | int ntri, int h, int w 81 | ): 82 | _rasterize_triangles( 83 | np.PyArray_DATA(vertices), np.PyArray_DATA(triangles), 84 | np.PyArray_DATA(depth_buffer), np.PyArray_DATA(triangle_buffer), 85 | np.PyArray_DATA(barycentric_weight), 86 | ntri, h, w) 87 | 88 | @cython.boundscheck(False) # turn off bounds-checking for entire function 89 | @cython.wraparound(False) # turn off negative index wrapping for entire function 90 | def rasterize(np.ndarray[unsigned char, ndim=3, mode = "c"] image not None, 91 | np.ndarray[float, ndim=2, mode = "c"] vertices not None, 92 | np.ndarray[int, ndim=2, mode="c"] triangles not None, 93 | np.ndarray[float, ndim=2, mode = "c"] colors not None, 94 | np.ndarray[float, ndim=2, mode = "c"] depth_buffer not None, 95 | int ntri, int h, int w, int c, float alpha = 1, bool reverse = False 96 | ): 97 | _rasterize( 98 | np.PyArray_DATA(image), np.PyArray_DATA(vertices), 99 | np.PyArray_DATA(triangles), 100 | np.PyArray_DATA(colors), 101 | np.PyArray_DATA(depth_buffer), 102 | ntri, h, w, c, alpha, reverse) 103 | 104 | # def render_texture_core(np.ndarray[float, ndim=3, mode = "c"] image not None, 105 | # np.ndarray[float, ndim=2, mode = "c"] vertices not None, 106 | # np.ndarray[int, ndim=2, mode="c"] triangles not None, 107 | # np.ndarray[float, ndim=3, mode = "c"] texture not None, 108 | # np.ndarray[float, ndim=2, mode = "c"] tex_coords not None, 109 | # np.ndarray[int, ndim=2, mode="c"] tex_triangles not None, 110 | # np.ndarray[float, ndim=2, mode = "c"] depth_buffer not None, 111 | # int nver, int tex_nver, int ntri, 112 | # int h, int w, int c, 113 | # int tex_h, int tex_w, int tex_c, 114 | # int mapping_type 115 | # ): 116 | # _render_texture_core( 117 | # np.PyArray_DATA(image), np.PyArray_DATA(vertices), np.PyArray_DATA(triangles), 118 | # np.PyArray_DATA(texture), np.PyArray_DATA(tex_coords), np.PyArray_DATA(tex_triangles), 119 | # np.PyArray_DATA(depth_buffer), 120 | # nver, tex_nver, ntri, 121 | # h, w, c, 122 | # tex_h, tex_w, tex_c, 123 | # mapping_type) 124 | # 125 | # def write_obj_with_colors_texture_core(string filename, string mtl_name, 126 | # np.ndarray[float, ndim=2, mode = "c"] vertices not None, 127 | # np.ndarray[int, ndim=2, mode="c"] triangles not None, 128 | # np.ndarray[float, ndim=2, mode = "c"] colors not None, 129 | # np.ndarray[float, ndim=2, mode = "c"] uv_coords not None, 130 | # int nver, int ntri, int ntexver 131 | # ): 132 | # _write_obj_with_colors_texture(filename, mtl_name, 133 | # np.PyArray_DATA(vertices), np.PyArray_DATA(triangles), np.PyArray_DATA(colors), np.PyArray_DATA(uv_coords), 134 | # nver, ntri, ntexver) 135 | -------------------------------------------------------------------------------- /lib/Sim3DR/lighting.py: -------------------------------------------------------------------------------- 1 | # coding: utf-8 2 | 3 | import numpy as np 4 | from .Sim3DR import get_normal, rasterize 5 | 6 | _norm = lambda arr: arr / np.sqrt(np.sum(arr ** 2, axis=1))[:, None] 7 | 8 | 9 | def norm_vertices(vertices): 10 | vertices -= vertices.min(0)[None, :] 11 | vertices /= vertices.max() 12 | vertices *= 2 13 | vertices -= vertices.max(0)[None, :] / 2 14 | return vertices 15 | 16 | 17 | def convert_type(obj): 18 | if isinstance(obj, tuple) or isinstance(obj, list): 19 | return np.array(obj, dtype=np.float32)[None, :] 20 | return obj 21 | 22 | 23 | class RenderPipeline(object): 24 | def __init__(self, **kwargs): 25 | self.intensity_ambient = convert_type(kwargs.get('intensity_ambient', 0.3)) 26 | self.intensity_directional = convert_type(kwargs.get('intensity_directional', 0.6)) 27 | self.intensity_specular = convert_type(kwargs.get('intensity_specular', 0.1)) 28 | self.specular_exp = kwargs.get('specular_exp', 5) 29 | self.color_ambient = convert_type(kwargs.get('color_ambient', (1, 1, 1))) 30 | self.color_directional = convert_type(kwargs.get('color_directional', (1, 1, 1))) 31 | self.light_pos = convert_type(kwargs.get('light_pos', (0, 0, 5))) 32 | self.view_pos = convert_type(kwargs.get('view_pos', (0, 0, 5))) 33 | 34 | def update_light_pos(self, light_pos): 35 | self.light_pos = convert_type(light_pos) 36 | 37 | def __call__(self, vertices, triangles, bg, texture=None): 38 | normal = get_normal(vertices, triangles) 39 | 40 | # 2. lighting 41 | light = np.zeros_like(vertices, dtype=np.float32) 42 | # ambient component 43 | if self.intensity_ambient > 0: 44 | light += self.intensity_ambient * self.color_ambient 45 | 46 | vertices_n = norm_vertices(vertices.copy()) 47 | if self.intensity_directional > 0: 48 | # diffuse component 49 | direction = _norm(self.light_pos - vertices_n) 50 | cos = np.sum(normal * direction, axis=1)[:, None] 51 | # cos = np.clip(cos, 0, 1) 52 | # todo: check below 53 | light += self.intensity_directional * (self.color_directional * np.clip(cos, 0, 1)) 54 | 55 | # specular component 56 | if self.intensity_specular > 0: 57 | v2v = _norm(self.view_pos - vertices_n) 58 | reflection = 2 * cos * normal - direction 59 | spe = np.sum((v2v * reflection) ** self.specular_exp, axis=1)[:, None] 60 | spe = np.where(cos != 0, np.clip(spe, 0, 1), np.zeros_like(spe)) 61 | light += self.intensity_specular * self.color_directional * np.clip(spe, 0, 1) 62 | light = np.clip(light, 0, 1) 63 | 64 | # 2. rasterization, [0, 1] 65 | if texture is None: 66 | render_img = rasterize(vertices, triangles, light, bg=bg) 67 | return render_img 68 | else: 69 | texture *= light 70 | render_img = rasterize(vertices, triangles, texture, bg=bg) 71 | return render_img 72 | 73 | 74 | def main(): 75 | pass 76 | 77 | 78 | if __name__ == '__main__': 79 | main() 80 | -------------------------------------------------------------------------------- /lib/Sim3DR/readme.md: -------------------------------------------------------------------------------- 1 | ## Forked from https://github.com/cleardusk/3DDFA_V2/tree/master/Sim3DR 2 | 3 | ## Sim3DR 4 | This is a simple 3D render, written by c++ and cython. 5 | 6 | ### Build Sim3DR 7 | 8 | ```shell script 9 | python3 setup.py build_ext --inplace 10 | ``` 11 | 12 | -------------------------------------------------------------------------------- /lib/Sim3DR/setup.py: -------------------------------------------------------------------------------- 1 | ''' 2 | python setup.py build_ext -i 3 | to compile 4 | ''' 5 | 6 | from distutils.core import setup, Extension 7 | from Cython.Build import cythonize 8 | from Cython.Distutils import build_ext 9 | import numpy 10 | 11 | setup( 12 | name='Sim3DR_Cython', # not the package name 13 | cmdclass={'build_ext': build_ext}, 14 | ext_modules=[Extension("Sim3DR_Cython", 15 | sources=["lib/rasterize.pyx", "lib/rasterize_kernel.cpp"], 16 | language='c++', 17 | include_dirs=[numpy.get_include()], 18 | extra_compile_args=["-std=c++11"])], 19 | ) 20 | -------------------------------------------------------------------------------- /lib/Sim3DR/tests/io.cpp: -------------------------------------------------------------------------------- 1 | #include "io.h" 2 | 3 | //void load_obj(const string obj_fp, float* vertices, float* colors, float* triangles){ 4 | // string line; 5 | // ifstream in(obj_fp); 6 | // 7 | // if(in.is_open()){ 8 | // while (getline(in, line)){ 9 | // stringstream ss(line); 10 | // 11 | // char t; // type: v, f 12 | // ss >> t; 13 | // if (t == 'v'){ 14 | // 15 | // } 16 | // } 17 | // } 18 | //} 19 | 20 | void load_obj(const char *obj_fp, float *vertices, float *colors, int *triangles, int nver, int ntri) { 21 | FILE *fp; 22 | fp = fopen(obj_fp, "r"); 23 | 24 | char t; // type: v or f 25 | if (fp != nullptr) { 26 | for (int i = 0; i < nver; ++i) { 27 | fscanf(fp, "%c", &t); 28 | for (int j = 0; j < 3; ++j) 29 | fscanf(fp, " %f", &vertices[3 * i + j]); 30 | for (int j = 0; j < 3; ++j) 31 | fscanf(fp, " %f", &colors[3 * i + j]); 32 | fscanf(fp, "\n"); 33 | } 34 | // fscanf(fp, "%c", &t); 35 | for (int i = 0; i < ntri; ++i) { 36 | fscanf(fp, "%c", &t); 37 | for (int j = 0; j < 3; ++j) { 38 | fscanf(fp, " %d", &triangles[3 * i + j]); 39 | triangles[3 * i + j] -= 1; 40 | } 41 | fscanf(fp, "\n"); 42 | } 43 | 44 | fclose(fp); 45 | } 46 | } 47 | 48 | void load_ply(const char *ply_fp, float *vertices, int *triangles, int nver, int ntri) { 49 | FILE *fp; 50 | fp = fopen(ply_fp, "r"); 51 | 52 | // char s[256]; 53 | char t; 54 | if (fp != nullptr) { 55 | // for (int i = 0; i < 9; ++i) 56 | // fscanf(fp, "%s", s); 57 | for (int i = 0; i < nver; ++i) 58 | fscanf(fp, "%f %f %f\n", &vertices[3 * i], &vertices[3 * i + 1], &vertices[3 * i + 2]); 59 | 60 | for (int i = 0; i < ntri; ++i) 61 | fscanf(fp, "%c %d %d %d\n", &t, &triangles[3 * i], &triangles[3 * i + 1], &triangles[3 * i + 2]); 62 | 63 | fclose(fp); 64 | } 65 | } 66 | 67 | void write_ppm(const char *filename, unsigned char *img, int h, int w, int c) { 68 | FILE *fp; 69 | //open file for output 70 | fp = fopen(filename, "wb"); 71 | if (!fp) { 72 | fprintf(stderr, "Unable to open file '%s'\n", filename); 73 | exit(1); 74 | } 75 | 76 | //write the header file 77 | //image format 78 | fprintf(fp, "P6\n"); 79 | 80 | //image size 81 | fprintf(fp, "%d %d\n", w, h); 82 | 83 | // rgb component depth 84 | fprintf(fp, "%d\n", MAX_PXL_VALUE); 85 | 86 | // pixel data 87 | fwrite(img, sizeof(unsigned char), size_t(h * w * c), fp); 88 | fclose(fp); 89 | } -------------------------------------------------------------------------------- /lib/Sim3DR/tests/io.h: -------------------------------------------------------------------------------- 1 | #ifndef IO_H_ 2 | #define IO_H_ 3 | 4 | #include 5 | #include 6 | #include 7 | #include 8 | #include 9 | 10 | using namespace std; 11 | 12 | #define MAX_PXL_VALUE 255 13 | 14 | void load_obj(const char* obj_fp, float* vertices, float* colors, int* triangles, int nver, int ntri); 15 | void load_ply(const char* ply_fp, float* vertices, int* triangles, int nver, int ntri); 16 | 17 | 18 | void write_ppm(const char *filename, unsigned char *img, int h, int w, int c); 19 | 20 | #endif -------------------------------------------------------------------------------- /lib/Sim3DR/tests/test.cpp: -------------------------------------------------------------------------------- 1 | /* 2 | * Tesing cases 3 | */ 4 | 5 | #include 6 | #include 7 | #include "rasterize.h" 8 | #include "io.h" 9 | 10 | void test_isPointInTri() { 11 | Point p0(0, 0); 12 | Point p1(1, 0); 13 | Point p2(1, 1); 14 | 15 | Point p(0.2, 0.2); 16 | 17 | if (is_point_in_tri(p, p0, p1, p2)) 18 | std::cout << "In"; 19 | else 20 | std::cout << "Out"; 21 | std::cout << std::endl; 22 | } 23 | 24 | void test_getPointWeight() { 25 | Point p0(0, 0); 26 | Point p1(1, 0); 27 | Point p2(1, 1); 28 | 29 | Point p(0.2, 0.2); 30 | 31 | float weight[3]; 32 | get_point_weight(weight, p, p0, p1, p2); 33 | std::cout << weight[0] << " " << weight[1] << " " << weight[2] << std::endl; 34 | } 35 | 36 | void test_get_tri_normal() { 37 | float tri_normal[3]; 38 | // float vertices[9] = {1, 0, 0, 0, 0, 0, 0, 1, 0}; 39 | float vertices[9] = {1, 1.1, 0, 0, 0, 0, 0, 0.6, 0.7}; 40 | int triangles[3] = {0, 1, 2}; 41 | int ntri = 1; 42 | 43 | _get_tri_normal(tri_normal, vertices, triangles, ntri); 44 | 45 | for (int i = 0; i < 3; ++i) 46 | std::cout << tri_normal[i] << ", "; 47 | std::cout << std::endl; 48 | } 49 | 50 | void test_load_obj() { 51 | const char *fp = "../data/vd005_mesh.obj"; 52 | int nver = 35709; 53 | int ntri = 70789; 54 | 55 | auto *vertices = new float[nver]; 56 | auto *colors = new float[nver]; 57 | auto *triangles = new int[ntri]; 58 | load_obj(fp, vertices, colors, triangles, nver, ntri); 59 | 60 | delete[] vertices; 61 | delete[] colors; 62 | delete[] triangles; 63 | } 64 | 65 | void test_render() { 66 | // 1. loading obj 67 | // const char *fp = "/Users/gjz/gjzprojects/Sim3DR/data/vd005_mesh.obj"; 68 | const char *fp = "/Users/gjz/gjzprojects/Sim3DR/data/face1.obj"; 69 | int nver = 35709; //53215; //35709; 70 | int ntri = 70789; //105840;//70789; 71 | 72 | auto *vertices = new float[3 * nver]; 73 | auto *colors = new float[3 * nver]; 74 | auto *triangles = new int[3 * ntri]; 75 | load_obj(fp, vertices, colors, triangles, nver, ntri); 76 | 77 | // 2. rendering 78 | int h = 224, w = 224, c = 3; 79 | 80 | // enlarging 81 | int scale = 4; 82 | h *= scale; 83 | w *= scale; 84 | for (int i = 0; i < nver * 3; ++i) vertices[i] *= scale; 85 | 86 | auto *image = new unsigned char[h * w * c](); 87 | auto *depth_buffer = new float[h * w](); 88 | 89 | for (int i = 0; i < h * w; ++i) depth_buffer[i] = -999999; 90 | 91 | clock_t t; 92 | t = clock(); 93 | 94 | _rasterize(image, vertices, triangles, colors, depth_buffer, ntri, h, w, c, true); 95 | t = clock() - t; 96 | double time_taken = ((double) t) / CLOCKS_PER_SEC; // in seconds 97 | printf("Render took %f seconds to execute \n", time_taken); 98 | 99 | 100 | // auto *image_char = new u_char[h * w * c](); 101 | // for (int i = 0; i < h * w * c; ++i) 102 | // image_char[i] = u_char(255 * image[i]); 103 | write_ppm("res.ppm", image, h, w, c); 104 | 105 | // delete[] image_char; 106 | delete[] vertices; 107 | delete[] colors; 108 | delete[] triangles; 109 | delete[] image; 110 | delete[] depth_buffer; 111 | } 112 | 113 | void test_light() { 114 | // 1. loading obj 115 | const char *fp = "/Users/gjz/gjzprojects/Sim3DR/data/emma_input_0_noheader.ply"; 116 | int nver = 53215; //35709; 117 | int ntri = 105840; //70789; 118 | 119 | auto *vertices = new float[3 * nver]; 120 | auto *colors = new float[3 * nver]; 121 | auto *triangles = new int[3 * ntri]; 122 | load_ply(fp, vertices, triangles, nver, ntri); 123 | 124 | // 2. rendering 125 | // int h = 1901, w = 3913, c = 3; 126 | int h = 2000, w = 4000, c = 3; 127 | 128 | // enlarging 129 | // int scale = 1; 130 | // h *= scale; 131 | // w *= scale; 132 | // for (int i = 0; i < nver * 3; ++i) vertices[i] *= scale; 133 | 134 | auto *image = new unsigned char[h * w * c](); 135 | auto *depth_buffer = new float[h * w](); 136 | 137 | for (int i = 0; i < h * w; ++i) depth_buffer[i] = -999999; 138 | for (int i = 0; i < 3 * nver; ++i) colors[i] = 0.8; 139 | 140 | clock_t t; 141 | t = clock(); 142 | 143 | _rasterize(image, vertices, triangles, colors, depth_buffer, ntri, h, w, c, true); 144 | t = clock() - t; 145 | double time_taken = ((double) t) / CLOCKS_PER_SEC; // in seconds 146 | printf("Render took %f seconds to execute \n", time_taken); 147 | 148 | 149 | // auto *image_char = new u_char[h * w * c](); 150 | // for (int i = 0; i < h * w * c; ++i) 151 | // image_char[i] = u_char(255 * image[i]); 152 | write_ppm("emma.ppm", image, h, w, c); 153 | 154 | // delete[] image_char; 155 | delete[] vertices; 156 | delete[] colors; 157 | delete[] triangles; 158 | delete[] image; 159 | delete[] depth_buffer; 160 | } 161 | 162 | int main(int argc, char *argv[]) { 163 | // std::cout << "Hello CMake!" << std::endl; 164 | 165 | // test_isPointInTri(); 166 | // test_getPointWeight(); 167 | // test_get_tri_normal(); 168 | // test_load_obj(); 169 | // test_render(); 170 | test_light(); 171 | return 0; 172 | } -------------------------------------------------------------------------------- /lib/eyemouth_index.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | # eye and mouth index and faces (made by the) in face obj 1220 3 | vert_index=np.array([1101, 1100, 1099, 1098, 1097, 1096, 1095, 1094, 1093, 1092, 1091, 1090, 1089, 1088, 1087, 1086, 1085, 1108, 1107, 1106, 1105, 1104, 1103, 1102, 1081, 1080, 1079, 1078, 1077, 1076, 1075, 1074, 1073, 1072, 1071, 1070, 1069, 1068, 1067, 1066, 1065, 1064, 1063, 1062, 1061, 1084, 1083, 1082, 249, 393, 250, 251, 252, 253, 254, 255, 256, 24, 691, 690, 689, 688, 687, 686, 685, 823, 684, 834, 740, 683, 682, 710, 725, 709, 700, 25, 265, 274, 290, 275, 247, 248, 305, 404]) 4 | 5 | face_em = np.array([ 6 | # left eye 7 | [ 0, 1, 23 ], 8 | [ 1, 2, 23 ], 9 | [ 2, 22,23 ], 10 | [ 2, 3, 22 ], 11 | [ 3, 22,21 ], 12 | [ 3, 4, 21 ], 13 | [ 4, 21,20 ], 14 | [ 4, 5, 20 ], 15 | [ 5, 20,19 ], 16 | [ 5, 6, 19 ], 17 | [ 6, 19,18 ], 18 | [ 6, 7, 18 ], 19 | [ 7, 18,17 ], 20 | [ 7, 8, 17 ], 21 | [ 8, 17,16 ], 22 | [ 8, 9, 16 ], 23 | [ 9, 16,15 ], 24 | [ 9, 10,15 ], 25 | [ 10,15,14 ], 26 | [10, 11, 14], 27 | [11, 14, 13], 28 | [11, 12, 13], 29 | # right eye 30 | [24, 25, 47], 31 | [25, 26, 47], 32 | [26, 46, 47], 33 | [26, 27, 46], 34 | [27, 46, 45], 35 | [27, 28, 45], 36 | [28, 45, 44], 37 | [28, 29, 44], 38 | [29, 44, 43], 39 | [29, 30, 43], 40 | [30, 43, 42], 41 | [30, 31, 42], 42 | [31, 42, 41], 43 | [31, 32, 41], 44 | [32, 41, 40], 45 | [32, 33, 40], 46 | [33, 40, 39], 47 | [33, 34, 39], 48 | [34, 39, 38], 49 | [34, 35, 38], 50 | [35, 38, 37], 51 | [35, 36, 37], 52 | 53 | # mouth 54 | [48, 49, 83], 55 | [49, 50, 83], 56 | [50, 83, 82], 57 | [50, 51, 82], 58 | [51, 82, 81], 59 | [51, 52, 81], 60 | [52, 81, 80], 61 | [52, 53, 80], 62 | [53, 80, 79], 63 | [53, 54, 79], 64 | [54, 79, 78], 65 | [54, 55, 78], 66 | [55, 78, 77], 67 | [55, 56, 77], 68 | [56, 77, 76], 69 | [56, 57, 76], 70 | [57, 76, 75], 71 | [57, 58, 75], 72 | [58, 75, 74], 73 | [58, 59, 74], 74 | [59, 74, 73], 75 | [59, 60, 73], 76 | [60, 73, 72], 77 | [60, 61, 72], 78 | [61, 72, 71], 79 | [61, 62, 71], 80 | [62, 71, 70], 81 | [62, 63, 70], 82 | [63, 70, 69], 83 | [63, 64, 69], 84 | [64, 69, 68], 85 | [64, 65, 68], 86 | [65, 68, 67], 87 | [65, 67, 66] 88 | ], dtype=np.int32) 89 | -------------------------------------------------------------------------------- /lib/mesh/__init__.py: -------------------------------------------------------------------------------- 1 | from .cython import mesh_core_cython 2 | from . import io 3 | from . import vis 4 | from . import transform 5 | from . import light 6 | from . import render 7 | 8 | -------------------------------------------------------------------------------- /lib/mesh/cython/mesh_core.h: -------------------------------------------------------------------------------- 1 | #ifndef MESH_CORE_HPP_ 2 | #define MESH_CORE_HPP_ 3 | 4 | #include 5 | #include 6 | #include 7 | #include 8 | #include 9 | #include 10 | 11 | using namespace std; 12 | 13 | class point 14 | { 15 | public: 16 | float x; 17 | float y; 18 | 19 | float dot(point p) 20 | { 21 | return this->x * p.x + this->y * p.y; 22 | } 23 | 24 | point operator-(const point& p) 25 | { 26 | point np; 27 | np.x = this->x - p.x; 28 | np.y = this->y - p.y; 29 | return np; 30 | } 31 | 32 | point operator+(const point& p) 33 | { 34 | point np; 35 | np.x = this->x + p.x; 36 | np.y = this->y + p.y; 37 | return np; 38 | } 39 | 40 | point operator*(float s) 41 | { 42 | point np; 43 | np.x = s * this->x; 44 | np.y = s * this->y; 45 | return np; 46 | } 47 | }; 48 | 49 | 50 | bool isPointInTri(point p, point p0, point p1, point p2, int h, int w); 51 | void get_point_weight(float* weight, point p, point p0, point p1, point p2); 52 | 53 | void _get_normal_core( 54 | float* normal, float* tri_normal, int* triangles, 55 | int ntri); 56 | 57 | void _rasterize_triangles_core( 58 | float* vertices, int* triangles, 59 | float* depth_buffer, int* triangle_buffer, float* barycentric_weight, 60 | int nver, int ntri, 61 | int h, int w); 62 | 63 | void _render_colors_core( 64 | float* image, float* vertices, int* triangles, 65 | float* colors, 66 | float* depth_buffer, 67 | int nver, int ntri, 68 | int h, int w, int c); 69 | 70 | void _render_texture_core( 71 | float* image, float* vertices, int* triangles, 72 | float* texture, float* tex_coords, int* tex_triangles, 73 | float* depth_buffer, 74 | int nver, int tex_nver, int ntri, 75 | int h, int w, int c, 76 | int tex_h, int tex_w, int tex_c, 77 | int mapping_type); 78 | 79 | void _write_obj_with_colors_texture(string filename, string mtl_name, 80 | float* vertices, int* triangles, float* colors, float* uv_coords, 81 | int nver, int ntri, int ntexver); 82 | 83 | #endif -------------------------------------------------------------------------------- /lib/mesh/cython/mesh_core_cython.cpython-36m-x86_64-linux-gnu.so: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/cbsropenproject/6dof_face/c80ca80125b2059e8f2984688814d59d0ac8e10c/lib/mesh/cython/mesh_core_cython.cpython-36m-x86_64-linux-gnu.so -------------------------------------------------------------------------------- /lib/mesh/cython/mesh_core_cython.pyx: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | cimport numpy as np 3 | from libcpp.string cimport string 4 | 5 | # use the Numpy-C-API from Cython 6 | np.import_array() 7 | 8 | # cdefine the signature of our c function 9 | cdef extern from "mesh_core.h": 10 | void _rasterize_triangles_core( 11 | float* vertices, int* triangles, 12 | float* depth_buffer, int* triangle_buffer, float* barycentric_weight, 13 | int nver, int ntri, 14 | int h, int w) 15 | 16 | void _render_colors_core( 17 | float* image, float* vertices, int* triangles, 18 | float* colors, 19 | float* depth_buffer, 20 | int nver, int ntri, 21 | int h, int w, int c) 22 | 23 | void _render_texture_core( 24 | float* image, float* vertices, int* triangles, 25 | float* texture, float* tex_coords, int* tex_triangles, 26 | float* depth_buffer, 27 | int nver, int tex_nver, int ntri, 28 | int h, int w, int c, 29 | int tex_h, int tex_w, int tex_c, 30 | int mapping_type) 31 | 32 | void _get_normal_core( 33 | float* normal, float* tri_normal, int* triangles, 34 | int ntri) 35 | 36 | void _write_obj_with_colors_texture(string filename, string mtl_name, 37 | float* vertices, int* triangles, float* colors, float* uv_coords, 38 | int nver, int ntri, int ntexver) 39 | 40 | def get_normal_core(np.ndarray[float, ndim=2, mode = "c"] normal not None, 41 | np.ndarray[float, ndim=2, mode = "c"] tri_normal not None, 42 | np.ndarray[int, ndim=2, mode="c"] triangles not None, 43 | int ntri 44 | ): 45 | _get_normal_core( 46 | np.PyArray_DATA(normal), np.PyArray_DATA(tri_normal), np.PyArray_DATA(triangles), 47 | ntri) 48 | 49 | def rasterize_triangles_core( 50 | np.ndarray[float, ndim=2, mode = "c"] vertices not None, 51 | np.ndarray[int, ndim=2, mode="c"] triangles not None, 52 | np.ndarray[float, ndim=2, mode = "c"] depth_buffer not None, 53 | np.ndarray[int, ndim=2, mode = "c"] triangle_buffer not None, 54 | np.ndarray[float, ndim=2, mode = "c"] barycentric_weight not None, 55 | int nver, int ntri, 56 | int h, int w 57 | ): 58 | _rasterize_triangles_core( 59 | np.PyArray_DATA(vertices), np.PyArray_DATA(triangles), 60 | np.PyArray_DATA(depth_buffer), np.PyArray_DATA(triangle_buffer), np.PyArray_DATA(barycentric_weight), 61 | nver, ntri, 62 | h, w) 63 | 64 | def render_colors_core(np.ndarray[float, ndim=3, mode = "c"] image not None, 65 | np.ndarray[float, ndim=2, mode = "c"] vertices not None, 66 | np.ndarray[int, ndim=2, mode="c"] triangles not None, 67 | np.ndarray[float, ndim=2, mode = "c"] colors not None, 68 | np.ndarray[float, ndim=2, mode = "c"] depth_buffer not None, 69 | int nver, int ntri, 70 | int h, int w, int c 71 | ): 72 | _render_colors_core( 73 | np.PyArray_DATA(image), np.PyArray_DATA(vertices), np.PyArray_DATA(triangles), 74 | np.PyArray_DATA(colors), 75 | np.PyArray_DATA(depth_buffer), 76 | nver, ntri, 77 | h, w, c) 78 | 79 | def render_texture_core(np.ndarray[float, ndim=3, mode = "c"] image not None, 80 | np.ndarray[float, ndim=2, mode = "c"] vertices not None, 81 | np.ndarray[int, ndim=2, mode="c"] triangles not None, 82 | np.ndarray[float, ndim=3, mode = "c"] texture not None, 83 | np.ndarray[float, ndim=2, mode = "c"] tex_coords not None, 84 | np.ndarray[int, ndim=2, mode="c"] tex_triangles not None, 85 | np.ndarray[float, ndim=2, mode = "c"] depth_buffer not None, 86 | int nver, int tex_nver, int ntri, 87 | int h, int w, int c, 88 | int tex_h, int tex_w, int tex_c, 89 | int mapping_type 90 | ): 91 | _render_texture_core( 92 | np.PyArray_DATA(image), np.PyArray_DATA(vertices), np.PyArray_DATA(triangles), 93 | np.PyArray_DATA(texture), np.PyArray_DATA(tex_coords), np.PyArray_DATA(tex_triangles), 94 | np.PyArray_DATA(depth_buffer), 95 | nver, tex_nver, ntri, 96 | h, w, c, 97 | tex_h, tex_w, tex_c, 98 | mapping_type) 99 | 100 | def write_obj_with_colors_texture_core(string filename, string mtl_name, 101 | np.ndarray[float, ndim=2, mode = "c"] vertices not None, 102 | np.ndarray[int, ndim=2, mode="c"] triangles not None, 103 | np.ndarray[float, ndim=2, mode = "c"] colors not None, 104 | np.ndarray[float, ndim=2, mode = "c"] uv_coords not None, 105 | int nver, int ntri, int ntexver 106 | ): 107 | _write_obj_with_colors_texture(filename, mtl_name, 108 | np.PyArray_DATA(vertices), np.PyArray_DATA(triangles), np.PyArray_DATA(colors), np.PyArray_DATA(uv_coords), 109 | nver, ntri, ntexver) 110 | -------------------------------------------------------------------------------- /lib/mesh/cython/setup.py: -------------------------------------------------------------------------------- 1 | ''' 2 | python setup.py build_ext -i 3 | to compile 4 | ''' 5 | 6 | # setup.py 7 | from distutils.core import setup, Extension 8 | from Cython.Build import cythonize 9 | from Cython.Distutils import build_ext 10 | import numpy 11 | 12 | setup( 13 | name = 'mesh_core_cython', 14 | cmdclass={'build_ext': build_ext}, 15 | ext_modules=[Extension("mesh_core_cython", 16 | sources=["mesh_core_cython.pyx", "mesh_core.cpp"], 17 | language='c++', 18 | include_dirs=[numpy.get_include()])], 19 | ) 20 | 21 | -------------------------------------------------------------------------------- /lib/mesh/io.py: -------------------------------------------------------------------------------- 1 | from __future__ import absolute_import 2 | from __future__ import division 3 | from __future__ import print_function 4 | 5 | import numpy as np 6 | import os 7 | from skimage import io 8 | from time import time 9 | 10 | from .cython import mesh_core_cython 11 | 12 | ## TODO 13 | ## TODO: c++ version 14 | def read_obj(obj_name): 15 | ''' read mesh 16 | ''' 17 | return 0 18 | 19 | # ------------------------- write 20 | def write_asc(path, vertices): 21 | ''' 22 | Args: 23 | vertices: shape = (nver, 3) 24 | ''' 25 | if path.split('.')[-1] == 'asc': 26 | np.savetxt(path, vertices) 27 | else: 28 | np.savetxt(path + '.asc', vertices) 29 | 30 | def write_obj_with_colors(obj_name, vertices, triangles, colors): 31 | ''' Save 3D face model with texture represented by colors. 32 | Args: 33 | obj_name: str 34 | vertices: shape = (nver, 3) 35 | triangles: shape = (ntri, 3) 36 | colors: shape = (nver, 3) 37 | ''' 38 | triangles = triangles.copy() 39 | triangles += 1 # meshlab start with 1 40 | 41 | if obj_name.split('.')[-1] != 'obj': 42 | obj_name = obj_name + '.obj' 43 | 44 | # write obj 45 | with open(obj_name, 'w') as f: 46 | 47 | # write vertices & colors 48 | for i in range(vertices.shape[0]): 49 | # s = 'v {} {} {} \n'.format(vertices[0,i], vertices[1,i], vertices[2,i]) 50 | s = 'v {} {} {} {} {} {}\n'.format(vertices[i, 0], vertices[i, 1], vertices[i, 2], colors[i, 0], colors[i, 1], colors[i, 2]) 51 | f.write(s) 52 | 53 | # write f: ver ind/ uv ind 54 | [k, ntri] = triangles.shape 55 | for i in range(triangles.shape[0]): 56 | # s = 'f {} {} {}\n'.format(triangles[i, 0], triangles[i, 1], triangles[i, 2]) 57 | s = 'f {} {} {}\n'.format(triangles[i, 2], triangles[i, 1], triangles[i, 0]) 58 | f.write(s) 59 | 60 | ## TODO: c++ version 61 | def write_obj_with_texture(obj_name, vertices, triangles, texture, uv_coords): 62 | ''' Save 3D face model with texture represented by texture map. 63 | Ref: https://github.com/patrikhuber/eos/blob/bd00155ebae4b1a13b08bf5a991694d682abbada/include/eos/core/Mesh.hpp 64 | Args: 65 | obj_name: str 66 | vertices: shape = (nver, 3) 67 | triangles: shape = (ntri, 3) 68 | texture: shape = (256,256,3) 69 | uv_coords: shape = (nver, 3) max value<=1 70 | ''' 71 | if obj_name.split('.')[-1] != 'obj': 72 | obj_name = obj_name + '.obj' 73 | mtl_name = obj_name.replace('.obj', '.mtl') 74 | texture_name = obj_name.replace('.obj', '_texture.png') 75 | 76 | triangles = triangles.copy() 77 | triangles += 1 # mesh lab start with 1 78 | 79 | # write obj 80 | with open(obj_name, 'w') as f: 81 | # first line: write mtlib(material library) 82 | s = "mtllib {}\n".format(os.path.abspath(mtl_name)) 83 | f.write(s) 84 | 85 | # write vertices 86 | for i in range(vertices.shape[0]): 87 | s = 'v {} {} {}\n'.format(vertices[i, 0], vertices[i, 1], vertices[i, 2]) 88 | f.write(s) 89 | 90 | # write uv coords 91 | for i in range(uv_coords.shape[0]): 92 | s = 'vt {} {}\n'.format(uv_coords[i,0], 1 - uv_coords[i,1]) 93 | f.write(s) 94 | 95 | f.write("usemtl FaceTexture\n") 96 | 97 | # write f: ver ind/ uv ind 98 | for i in range(triangles.shape[0]): 99 | s = 'f {}/{} {}/{} {}/{}\n'.format(triangles[i,2], triangles[i,2], triangles[i,1], triangles[i,1], triangles[i,0], triangles[i,0]) 100 | f.write(s) 101 | 102 | # write mtl 103 | with open(mtl_name, 'w') as f: 104 | f.write("newmtl FaceTexture\n") 105 | s = 'map_Kd {}\n'.format(os.path.abspath(texture_name)) # map to image 106 | f.write(s) 107 | 108 | # write texture as png 109 | imsave(texture_name, texture) 110 | 111 | # c++ version 112 | def write_obj_with_colors_texture(obj_name, vertices, triangles, colors, texture, uv_coords): 113 | ''' Save 3D face model with texture. 114 | Ref: https://github.com/patrikhuber/eos/blob/bd00155ebae4b1a13b08bf5a991694d682abbada/include/eos/core/Mesh.hpp 115 | Args: 116 | obj_name: str 117 | vertices: shape = (nver, 3) 118 | triangles: shape = (ntri, 3) 119 | colors: shape = (nver, 3) 120 | texture: shape = (256,256,3) 121 | uv_coords: shape = (nver, 3) max value<=1 122 | ''' 123 | if obj_name.split('.')[-1] != 'obj': 124 | obj_name = obj_name + '.obj' 125 | mtl_name = obj_name.replace('.obj', '.mtl') 126 | texture_name = obj_name.replace('.obj', '_texture.png') 127 | 128 | triangles = triangles.copy() 129 | triangles += 1 # mesh lab start with 1 130 | 131 | # write obj 132 | vertices, colors, uv_coords = vertices.astype(np.float32).copy(), colors.astype(np.float32).copy(), uv_coords.astype(np.float32).copy() 133 | mesh_core_cython.write_obj_with_colors_texture_core(str.encode(obj_name), str.encode(os.path.abspath(mtl_name)), vertices, triangles, colors, uv_coords, vertices.shape[0], triangles.shape[0], uv_coords.shape[0]) 134 | 135 | # write mtl 136 | with open(mtl_name, 'w') as f: 137 | f.write("newmtl FaceTexture\n") 138 | s = 'map_Kd {}\n'.format(os.path.abspath(texture_name)) # map to image 139 | f.write(s) 140 | 141 | # write texture as png 142 | io.imsave(texture_name, texture) -------------------------------------------------------------------------------- /lib/mesh/light.py: -------------------------------------------------------------------------------- 1 | ''' 2 | Functions about lighting mesh(changing colors/texture of mesh). 3 | 1. add light to colors/texture (shade each vertex) 4 | 2. fit light according to colors/texture & image. 5 | ''' 6 | 7 | from __future__ import absolute_import 8 | from __future__ import division 9 | from __future__ import print_function 10 | 11 | import numpy as np 12 | from .cython import mesh_core_cython 13 | 14 | def get_normal(vertices, triangles): 15 | ''' calculate normal direction in each vertex 16 | Args: 17 | vertices: [nver, 3] 18 | triangles: [ntri, 3] 19 | Returns: 20 | normal: [nver, 3] 21 | ''' 22 | pt0 = vertices[triangles[:, 0], :] # [ntri, 3] 23 | pt1 = vertices[triangles[:, 1], :] # [ntri, 3] 24 | pt2 = vertices[triangles[:, 2], :] # [ntri, 3] 25 | tri_normal = np.cross(pt0 - pt1, pt0 - pt2) # [ntri, 3]. normal of each triangle 26 | 27 | normal = np.zeros_like(vertices, dtype = np.float32).copy() # [nver, 3] 28 | # for i in range(triangles.shape[0]): 29 | # normal[triangles[i, 0], :] = normal[triangles[i, 0], :] + tri_normal[i, :] 30 | # normal[triangles[i, 1], :] = normal[triangles[i, 1], :] + tri_normal[i, :] 31 | # normal[triangles[i, 2], :] = normal[triangles[i, 2], :] + tri_normal[i, :] 32 | mesh_core_cython.get_normal_core(normal, tri_normal.astype(np.float32).copy(), triangles.copy(), triangles.shape[0]) 33 | 34 | # normalize to unit length 35 | mag = np.sum(normal**2, 1) # [nver] 36 | zero_ind = (mag == 0) 37 | mag[zero_ind] = 1; 38 | normal[zero_ind, 0] = np.ones((np.sum(zero_ind))) 39 | 40 | normal = normal/np.sqrt(mag[:,np.newaxis]) 41 | 42 | return normal 43 | 44 | # TODO: test 45 | def add_light_sh(vertices, triangles, colors, sh_coeff): 46 | ''' 47 | In 3d face, usually assume: 48 | 1. The surface of face is Lambertian(reflect only the low frequencies of lighting) 49 | 2. Lighting can be an arbitrary combination of point sources 50 | --> can be expressed in terms of spherical harmonics(omit the lighting coefficients) 51 | I = albedo * (sh(n) x sh_coeff) 52 | 53 | albedo: n x 1 54 | sh_coeff: 9 x 1 55 | Y(n) = (1, n_x, n_y, n_z, n_xn_y, n_xn_z, n_yn_z, n_x^2 - n_y^2, 3n_z^2 - 1)': n x 9 56 | # Y(n) = (1, n_x, n_y, n_z)': n x 4 57 | 58 | Args: 59 | vertices: [nver, 3] 60 | triangles: [ntri, 3] 61 | colors: [nver, 3] albedo 62 | sh_coeff: [9, 1] spherical harmonics coefficients 63 | 64 | Returns: 65 | lit_colors: [nver, 3] 66 | ''' 67 | assert vertices.shape[0] == colors.shape[0] 68 | nver = vertices.shape[0] 69 | normal = get_normal(vertices, triangles) # [nver, 3] 70 | sh = np.array((np.ones(nver), n[:,0], n[:,1], n[:,2], n[:,0]*n[:,1], n[:,0]*n[:,2], n[:,1]*n[:,2], n[:,0]**2 - n[:,1]**2, 3*(n[:,2]**2) - 1)) # [nver, 9] 71 | ref = sh.dot(sh_coeff) #[nver, 1] 72 | lit_colors = colors*ref 73 | return lit_colors 74 | 75 | 76 | def add_light(vertices, triangles, colors, light_positions = 0, light_intensities = 0): 77 | ''' Gouraud shading. add point lights. 78 | In 3d face, usually assume: 79 | 1. The surface of face is Lambertian(reflect only the low frequencies of lighting) 80 | 2. Lighting can be an arbitrary combination of point sources 81 | 3. No specular (unless skin is oil, 23333) 82 | 83 | Ref: https://cs184.eecs.berkeley.edu/lecture/pipeline 84 | Args: 85 | vertices: [nver, 3] 86 | triangles: [ntri, 3] 87 | light_positions: [nlight, 3] 88 | light_intensities: [nlight, 3] 89 | Returns: 90 | lit_colors: [nver, 3] 91 | ''' 92 | nver = vertices.shape[0] 93 | normals = get_normal(vertices, triangles) # [nver, 3] 94 | 95 | # ambient 96 | # La = ka*Ia 97 | 98 | # diffuse 99 | # Ld = kd*(I/r^2)max(0, nxl) 100 | direction_to_lights = vertices[np.newaxis, :, :] - light_positions[:, np.newaxis, :] # [nlight, nver, 3] 101 | direction_to_lights_n = np.sqrt(np.sum(direction_to_lights**2, axis = 2)) # [nlight, nver] 102 | direction_to_lights = direction_to_lights/direction_to_lights_n[:, :, np.newaxis] 103 | normals_dot_lights = normals[np.newaxis, :, :]*direction_to_lights # [nlight, nver, 3] 104 | normals_dot_lights = np.sum(normals_dot_lights, axis = 2) # [nlight, nver] 105 | diffuse_output = colors[np.newaxis, :, :]*normals_dot_lights[:, :, np.newaxis]*light_intensities[:, np.newaxis, :] 106 | diffuse_output = np.sum(diffuse_output, axis = 0) # [nver, 3] 107 | 108 | # specular 109 | # h = (v + l)/(|v + l|) bisector 110 | # Ls = ks*(I/r^2)max(0, nxh)^p 111 | # increasing p narrows the reflectionlob 112 | 113 | lit_colors = diffuse_output # only diffuse part here. 114 | lit_colors = np.minimum(np.maximum(lit_colors, 0), 1) 115 | return lit_colors 116 | 117 | 118 | 119 | ## TODO. estimate light(sh coeff) 120 | ## -------------------------------- estimate. can not use now. 121 | def fit_light(image, vertices, colors, triangles, vis_ind, lamb = 10, max_iter = 3): 122 | [h, w, c] = image.shape 123 | 124 | # surface normal 125 | norm = get_normal(vertices, triangles) 126 | 127 | nver = vertices.shape[1] 128 | 129 | # vertices --> corresponding image pixel 130 | pt2d = vertices[:2, :] 131 | 132 | pt2d[0,:] = np.minimum(np.maximum(pt2d[0,:], 0), w - 1) 133 | pt2d[1,:] = np.minimum(np.maximum(pt2d[1,:], 0), h - 1) 134 | pt2d = np.round(pt2d).astype(np.int32) # 2 x nver 135 | 136 | image_pixel = image[pt2d[1,:], pt2d[0,:], :] # nver x 3 137 | image_pixel = image_pixel.T # 3 x nver 138 | 139 | # vertices --> corresponding mean texture pixel with illumination 140 | # Spherical Harmonic Basis 141 | harmonic_dim = 9 142 | nx = norm[0,:]; 143 | ny = norm[1,:]; 144 | nz = norm[2,:]; 145 | harmonic = np.zeros((nver, harmonic_dim)) 146 | 147 | pi = np.pi 148 | harmonic[:,0] = np.sqrt(1/(4*pi)) * np.ones((nver,)); 149 | harmonic[:,1] = np.sqrt(3/(4*pi)) * nx; 150 | harmonic[:,2] = np.sqrt(3/(4*pi)) * ny; 151 | harmonic[:,3] = np.sqrt(3/(4*pi)) * nz; 152 | harmonic[:,4] = 1/2. * np.sqrt(3/(4*pi)) * (2*nz**2 - nx**2 - ny**2); 153 | harmonic[:,5] = 3 * np.sqrt(5/(12*pi)) * (ny*nz); 154 | harmonic[:,6] = 3 * np.sqrt(5/(12*pi)) * (nx*nz); 155 | harmonic[:,7] = 3 * np.sqrt(5/(12*pi)) * (nx*ny); 156 | harmonic[:,8] = 3/2. * np.sqrt(5/(12*pi)) * (nx*nx - ny*ny); 157 | 158 | ''' 159 | I' = sum(albedo * lj * hj) j = 0:9 (albedo = tex) 160 | set A = albedo*h (n x 9) 161 | alpha = lj (9 x 1) 162 | Y = I (n x 1) 163 | Y' = A.dot(alpha) 164 | 165 | opt function: 166 | ||Y - A*alpha|| + lambda*(alpha'*alpha) 167 | result: 168 | A'*(Y - A*alpha) + lambda*alpha = 0 169 | ==> 170 | (A'*A*alpha - lambda)*alpha = A'*Y 171 | left: 9 x 9 172 | right: 9 x 1 173 | ''' 174 | n_vis_ind = len(vis_ind) 175 | n = n_vis_ind*c 176 | 177 | Y = np.zeros((n, 1)) 178 | A = np.zeros((n, 9)) 179 | light = np.zeros((3, 1)) 180 | 181 | for k in range(c): 182 | Y[k*n_vis_ind:(k+1)*n_vis_ind, :] = image_pixel[k, vis_ind][:, np.newaxis] 183 | A[k*n_vis_ind:(k+1)*n_vis_ind, :] = texture[k, vis_ind][:, np.newaxis] * harmonic[vis_ind, :] 184 | Ac = texture[k, vis_ind][:, np.newaxis] 185 | Yc = image_pixel[k, vis_ind][:, np.newaxis] 186 | light[k] = (Ac.T.dot(Yc))/(Ac.T.dot(Ac)) 187 | 188 | for i in range(max_iter): 189 | 190 | Yc = Y.copy() 191 | for k in range(c): 192 | Yc[k*n_vis_ind:(k+1)*n_vis_ind, :] /= light[k] 193 | 194 | # update alpha 195 | equation_left = np.dot(A.T, A) + lamb*np.eye(harmonic_dim); # why + ? 196 | equation_right = np.dot(A.T, Yc) 197 | alpha = np.dot(np.linalg.inv(equation_left), equation_right) 198 | 199 | # update light 200 | for k in range(c): 201 | Ac = A[k*n_vis_ind:(k+1)*n_vis_ind, :].dot(alpha) 202 | Yc = Y[k*n_vis_ind:(k+1)*n_vis_ind, :] 203 | light[k] = (Ac.T.dot(Yc))/(Ac.T.dot(Ac)) 204 | 205 | appearance = np.zeros_like(texture) 206 | for k in range(c): 207 | tmp = np.dot(harmonic*texture[k, :][:, np.newaxis], alpha*light[k]) 208 | appearance[k,:] = tmp.T 209 | 210 | appearance = np.minimum(np.maximum(appearance, 0), 1) 211 | 212 | return appearance 213 | 214 | -------------------------------------------------------------------------------- /lib/mesh/readme.md: -------------------------------------------------------------------------------- 1 | ## Forked from https://github.com/yfeng95/face3d 2 | 3 | 4 | ### Compile c++ files to .so for python use 5 | 6 | ```bash 7 | cd mesh/cython 8 | python setup.py build_ext -i 9 | ``` 10 | 11 | -------------------------------------------------------------------------------- /lib/mesh/render.py: -------------------------------------------------------------------------------- 1 | ''' 2 | functions about rendering mesh(from 3d obj to 2d image). 3 | only use rasterization render here. 4 | Note that: 5 | 1. Generally, render func includes camera, light, raterize. Here no camera and light(I write these in other files) 6 | 2. Generally, the input vertices are normalized to [-1,1] and cetered on [0, 0]. (in world space) 7 | Here, the vertices are using image coords, which centers on [w/2, h/2] with the y-axis pointing to oppisite direction. 8 | Means: render here only conducts interpolation.(I just want to make the input flexible) 9 | 10 | Author: Yao Feng 11 | Mail: yaofeng1995@gmail.com 12 | ''' 13 | from __future__ import absolute_import 14 | from __future__ import division 15 | from __future__ import print_function 16 | 17 | import numpy as np 18 | from time import time 19 | 20 | from .cython import mesh_core_cython 21 | 22 | def rasterize_triangles(vertices, triangles, h, w): 23 | ''' 24 | Args: 25 | vertices: [nver, 3] 26 | triangles: [ntri, 3] 27 | h: height 28 | w: width 29 | Returns: 30 | depth_buffer: [h, w] saves the depth, here, the bigger the z, the fronter the point. 31 | triangle_buffer: [h, w] saves the tri id(-1 for no triangle). 32 | barycentric_weight: [h, w, 3] saves corresponding barycentric weight. 33 | 34 | # Each triangle has 3 vertices & Each vertex has 3 coordinates x, y, z. 35 | # h, w is the size of rendering 36 | ''' 37 | 38 | # initial 39 | depth_buffer = np.zeros([h, w]) - 999999. #set the initial z to the farest position 40 | triangle_buffer = np.zeros([h, w], dtype = np.int32) - 1 # if tri id = -1, the pixel has no triangle correspondance 41 | barycentric_weight = np.zeros([h, w, 3], dtype = np.float32) # 42 | 43 | vertices = vertices.astype(np.float32).copy() 44 | triangles = triangles.astype(np.int32).copy() 45 | 46 | mesh_core_cython.rasterize_triangles_core( 47 | vertices, triangles, 48 | depth_buffer, triangle_buffer, barycentric_weight, 49 | vertices.shape[0], triangles.shape[0], 50 | h, w) 51 | 52 | def render_colors(vertices, triangles, colors, h, w, c = 3, BG = None): 53 | ''' render mesh with colors 54 | Args: 55 | vertices: [nver, 3] 56 | triangles: [ntri, 3] 57 | colors: [nver, 3] 58 | h: height 59 | w: width 60 | c: channel 61 | BG: background image 62 | Returns: 63 | image: [h, w, c]. rendered image./rendering. 64 | ''' 65 | 66 | # initial 67 | if BG is None: 68 | image = np.zeros((h, w, c), dtype = np.float32) 69 | else: 70 | assert BG.shape[0] == h and BG.shape[1] == w and BG.shape[2] == c 71 | image = BG 72 | depth_buffer = np.zeros([h, w], dtype = np.float32, order = 'C') - 999999. 73 | 74 | # change orders. --> C-contiguous order(column major) 75 | vertices = vertices.astype(np.float32).copy() 76 | triangles = triangles.astype(np.int32).copy() 77 | colors = colors.astype(np.float32).copy() 78 | ### 79 | st = time() 80 | mesh_core_cython.render_colors_core( 81 | image, vertices, triangles, 82 | colors, 83 | depth_buffer, 84 | vertices.shape[0], triangles.shape[0], 85 | h, w, c) 86 | return image 87 | 88 | 89 | def render_texture(vertices, triangles, texture, tex_coords, tex_triangles, h, w, c = 3, mapping_type = 'nearest', BG = None): 90 | ''' render mesh with texture map 91 | Args: 92 | vertices: [3, nver] 93 | triangles: [3, ntri] 94 | texture: [tex_h, tex_w, 3] 95 | tex_coords: [ntexcoords, 3] 96 | tex_triangles: [ntri, 3] 97 | h: height of rendering 98 | w: width of rendering 99 | c: channel 100 | mapping_type: 'bilinear' or 'nearest' 101 | ''' 102 | # initial 103 | if BG is None: 104 | image = np.zeros((h, w, c), dtype = np.float32) 105 | else: 106 | assert BG.shape[0] == h and BG.shape[1] == w and BG.shape[2] == c 107 | image = BG 108 | 109 | depth_buffer = np.zeros([h, w], dtype = np.float32, order = 'C') - 999999. 110 | 111 | tex_h, tex_w, tex_c = texture.shape 112 | if mapping_type == 'nearest': 113 | mt = int(0) 114 | elif mapping_type == 'bilinear': 115 | mt = int(1) 116 | else: 117 | mt = int(0) 118 | 119 | # -> C order 120 | vertices = vertices.astype(np.float32).copy() 121 | triangles = triangles.astype(np.int32).copy() 122 | texture = texture.astype(np.float32).copy() 123 | tex_coords = tex_coords.astype(np.float32).copy() 124 | tex_triangles = tex_triangles.astype(np.int32).copy() 125 | 126 | mesh_core_cython.render_texture_core( 127 | image, vertices, triangles, 128 | texture, tex_coords, tex_triangles, 129 | depth_buffer, 130 | vertices.shape[0], tex_coords.shape[0], triangles.shape[0], 131 | h, w, c, 132 | tex_h, tex_w, tex_c, 133 | mt) 134 | return image 135 | 136 | -------------------------------------------------------------------------------- /lib/mesh/vis.py: -------------------------------------------------------------------------------- 1 | from __future__ import absolute_import 2 | from __future__ import division 3 | from __future__ import print_function 4 | 5 | import numpy as np 6 | import matplotlib.pyplot as plt 7 | from skimage import measure 8 | from mpl_toolkits.mplot3d import Axes3D 9 | 10 | def plot_mesh(vertices, triangles, subplot = [1,1,1], title = 'mesh', el = 90, az = -90, lwdt=.1, dist = 6, color = "grey"): 11 | ''' 12 | plot the mesh 13 | Args: 14 | vertices: [nver, 3] 15 | triangles: [ntri, 3] 16 | ''' 17 | ax = plt.subplot(subplot[0], subplot[1], subplot[2], projection = '3d') 18 | ax.plot_trisurf(vertices[:, 0], vertices[:, 1], vertices[:, 2], triangles = triangles, lw = lwdt, color = color, alpha = 1) 19 | ax.axis("off") 20 | ax.view_init(elev = el, azim = az) 21 | ax.dist = dist 22 | plt.title(title) 23 | 24 | ### -------------- Todo: use vtk to visualize mesh? or visvis? or VisPy? 25 | -------------------------------------------------------------------------------- /lib/mesh_p/__init__.py: -------------------------------------------------------------------------------- 1 | from .cython import mesh_core_cython 2 | from . import io 3 | from . import vis 4 | from . import transform 5 | from . import light 6 | from . import render 7 | 8 | -------------------------------------------------------------------------------- /lib/mesh_p/cython/mesh_core.h: -------------------------------------------------------------------------------- 1 | #ifndef MESH_CORE_HPP_ 2 | #define MESH_CORE_HPP_ 3 | 4 | #include 5 | #include 6 | #include 7 | #include 8 | #include 9 | #include 10 | 11 | using namespace std; 12 | 13 | class point 14 | { 15 | public: 16 | float x; 17 | float y; 18 | 19 | float dot(point p) 20 | { 21 | return this->x * p.x + this->y * p.y; 22 | } 23 | 24 | point operator-(const point& p) 25 | { 26 | point np; 27 | np.x = this->x - p.x; 28 | np.y = this->y - p.y; 29 | return np; 30 | } 31 | 32 | point operator+(const point& p) 33 | { 34 | point np; 35 | np.x = this->x + p.x; 36 | np.y = this->y + p.y; 37 | return np; 38 | } 39 | 40 | point operator*(float s) 41 | { 42 | point np; 43 | np.x = s * this->x; 44 | np.y = s * this->y; 45 | return np; 46 | } 47 | }; 48 | 49 | 50 | bool isPointInTri(point p, point p0, point p1, point p2, int h, int w); 51 | void get_point_weight(float* weight, point p, point p0, point p1, point p2); 52 | 53 | void _get_normal_core( 54 | float* normal, float* tri_normal, int* triangles, 55 | int ntri); 56 | 57 | void _rasterize_triangles_core( 58 | float* vertices, int* triangles, 59 | float* depth_buffer, int* triangle_buffer, float* barycentric_weight, 60 | int nver, int ntri, 61 | int h, int w); 62 | 63 | void _render_colors_core( 64 | float* image, float* corrs, 65 | float* vertices, int* triangles, 66 | float* colors, 67 | float* depth_buffer, 68 | int nver, int ntri, 69 | int h, int w, int c); 70 | 71 | void _render_texture_core( 72 | float* image, float* vertices, int* triangles, 73 | float* texture, float* tex_coords, int* tex_triangles, 74 | float* depth_buffer, 75 | int nver, int tex_nver, int ntri, 76 | int h, int w, int c, 77 | int tex_h, int tex_w, int tex_c, 78 | int mapping_type); 79 | 80 | void _write_obj_with_colors_texture(string filename, string mtl_name, 81 | float* vertices, int* triangles, float* colors, float* uv_coords, 82 | int nver, int ntri, int ntexver); 83 | 84 | #endif 85 | -------------------------------------------------------------------------------- /lib/mesh_p/cython/mesh_core_cython.pyx: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | cimport numpy as np 3 | from libcpp.string cimport string 4 | 5 | # use the Numpy-C-API from Cython 6 | np.import_array() 7 | 8 | # cdefine the signature of our c function 9 | cdef extern from "mesh_core.h": 10 | void _rasterize_triangles_core( 11 | float* vertices, int* triangles, 12 | float* depth_buffer, int* triangle_buffer, float* barycentric_weight, 13 | int nver, int ntri, 14 | int h, int w) 15 | 16 | void _render_colors_core( 17 | float* image, float* corrs, float* vertices, int* triangles, 18 | float* colors, 19 | float* depth_buffer, 20 | int nver, int ntri, 21 | int h, int w, int c) 22 | 23 | void _render_texture_core( 24 | float* image, float* vertices, int* triangles, 25 | float* texture, float* tex_coords, int* tex_triangles, 26 | float* depth_buffer, 27 | int nver, int tex_nver, int ntri, 28 | int h, int w, int c, 29 | int tex_h, int tex_w, int tex_c, 30 | int mapping_type) 31 | 32 | void _get_normal_core( 33 | float* normal, float* tri_normal, int* triangles, 34 | int ntri) 35 | 36 | void _write_obj_with_colors_texture(string filename, string mtl_name, 37 | float* vertices, int* triangles, float* colors, float* uv_coords, 38 | int nver, int ntri, int ntexver) 39 | 40 | def get_normal_core(np.ndarray[float, ndim=2, mode = "c"] normal not None, 41 | np.ndarray[float, ndim=2, mode = "c"] tri_normal not None, 42 | np.ndarray[int, ndim=2, mode="c"] triangles not None, 43 | int ntri 44 | ): 45 | _get_normal_core( 46 | np.PyArray_DATA(normal), np.PyArray_DATA(tri_normal), np.PyArray_DATA(triangles), 47 | ntri) 48 | 49 | def rasterize_triangles_core( 50 | np.ndarray[float, ndim=2, mode = "c"] vertices not None, 51 | np.ndarray[int, ndim=2, mode="c"] triangles not None, 52 | np.ndarray[float, ndim=2, mode = "c"] depth_buffer not None, 53 | np.ndarray[int, ndim=2, mode = "c"] triangle_buffer not None, 54 | np.ndarray[float, ndim=2, mode = "c"] barycentric_weight not None, 55 | int nver, int ntri, 56 | int h, int w 57 | ): 58 | _rasterize_triangles_core( 59 | np.PyArray_DATA(vertices), np.PyArray_DATA(triangles), 60 | np.PyArray_DATA(depth_buffer), np.PyArray_DATA(triangle_buffer), np.PyArray_DATA(barycentric_weight), 61 | nver, ntri, 62 | h, w) 63 | 64 | def render_colors_core(np.ndarray[float, ndim=3, mode = "c"] image not None, 65 | np.ndarray[float, ndim=3, mode = "c"] corrs not None, 66 | np.ndarray[float, ndim=2, mode = "c"] vertices not None, 67 | np.ndarray[int, ndim=2, mode="c"] triangles not None, 68 | np.ndarray[float, ndim=2, mode = "c"] colors not None, 69 | np.ndarray[float, ndim=2, mode = "c"] depth_buffer not None, 70 | int nver, int ntri, 71 | int h, int w, int c 72 | ): 73 | _render_colors_core( 74 | np.PyArray_DATA(image), np.PyArray_DATA(corrs), np.PyArray_DATA(vertices), np.PyArray_DATA(triangles), 75 | np.PyArray_DATA(colors), 76 | np.PyArray_DATA(depth_buffer), 77 | nver, ntri, 78 | h, w, c) 79 | 80 | def render_texture_core(np.ndarray[float, ndim=3, mode = "c"] image not None, 81 | np.ndarray[float, ndim=2, mode = "c"] vertices not None, 82 | np.ndarray[int, ndim=2, mode="c"] triangles not None, 83 | np.ndarray[float, ndim=3, mode = "c"] texture not None, 84 | np.ndarray[float, ndim=2, mode = "c"] tex_coords not None, 85 | np.ndarray[int, ndim=2, mode="c"] tex_triangles not None, 86 | np.ndarray[float, ndim=2, mode = "c"] depth_buffer not None, 87 | int nver, int tex_nver, int ntri, 88 | int h, int w, int c, 89 | int tex_h, int tex_w, int tex_c, 90 | int mapping_type 91 | ): 92 | _render_texture_core( 93 | np.PyArray_DATA(image), np.PyArray_DATA(vertices), np.PyArray_DATA(triangles), 94 | np.PyArray_DATA(texture), np.PyArray_DATA(tex_coords), np.PyArray_DATA(tex_triangles), 95 | np.PyArray_DATA(depth_buffer), 96 | nver, tex_nver, ntri, 97 | h, w, c, 98 | tex_h, tex_w, tex_c, 99 | mapping_type) 100 | 101 | def write_obj_with_colors_texture_core(string filename, string mtl_name, 102 | np.ndarray[float, ndim=2, mode = "c"] vertices not None, 103 | np.ndarray[int, ndim=2, mode="c"] triangles not None, 104 | np.ndarray[float, ndim=2, mode = "c"] colors not None, 105 | np.ndarray[float, ndim=2, mode = "c"] uv_coords not None, 106 | int nver, int ntri, int ntexver 107 | ): 108 | _write_obj_with_colors_texture(filename, mtl_name, 109 | np.PyArray_DATA(vertices), np.PyArray_DATA(triangles), np.PyArray_DATA(colors), np.PyArray_DATA(uv_coords), 110 | nver, ntri, ntexver) 111 | -------------------------------------------------------------------------------- /lib/mesh_p/cython/setup.py: -------------------------------------------------------------------------------- 1 | ''' 2 | python setup.py build_ext -i 3 | to compile 4 | ''' 5 | 6 | # setup.py 7 | from distutils.core import setup, Extension 8 | from Cython.Build import cythonize 9 | from Cython.Distutils import build_ext 10 | import numpy 11 | 12 | setup( 13 | name = 'mesh_core_cython', 14 | cmdclass={'build_ext': build_ext}, 15 | ext_modules=[Extension("mesh_core_cython", 16 | sources=["mesh_core_cython.pyx", "mesh_core.cpp"], 17 | language='c++', 18 | include_dirs=[numpy.get_include()])], 19 | ) 20 | 21 | -------------------------------------------------------------------------------- /lib/mesh_p/io.py: -------------------------------------------------------------------------------- 1 | from __future__ import absolute_import 2 | from __future__ import division 3 | from __future__ import print_function 4 | 5 | import numpy as np 6 | import os 7 | from skimage import io 8 | from time import time 9 | 10 | from .cython import mesh_core_cython 11 | 12 | ## TODO 13 | ## TODO: c++ version 14 | def read_obj(obj_name): 15 | ''' read mesh 16 | ''' 17 | return 0 18 | 19 | # ------------------------- write 20 | def write_asc(path, vertices): 21 | ''' 22 | Args: 23 | vertices: shape = (nver, 3) 24 | ''' 25 | if path.split('.')[-1] == 'asc': 26 | np.savetxt(path, vertices) 27 | else: 28 | np.savetxt(path + '.asc', vertices) 29 | 30 | def write_obj_with_colors(obj_name, vertices, triangles, colors): 31 | ''' Save 3D face model with texture represented by colors. 32 | Args: 33 | obj_name: str 34 | vertices: shape = (nver, 3) 35 | triangles: shape = (ntri, 3) 36 | colors: shape = (nver, 3) 37 | ''' 38 | triangles = triangles.copy() 39 | triangles += 1 # meshlab start with 1 40 | 41 | if obj_name.split('.')[-1] != 'obj': 42 | obj_name = obj_name + '.obj' 43 | 44 | # write obj 45 | with open(obj_name, 'w') as f: 46 | 47 | # write vertices & colors 48 | for i in range(vertices.shape[0]): 49 | # s = 'v {} {} {} \n'.format(vertices[0,i], vertices[1,i], vertices[2,i]) 50 | s = 'v {} {} {} {} {} {}\n'.format(vertices[i, 0], vertices[i, 1], vertices[i, 2], colors[i, 0], colors[i, 1], colors[i, 2]) 51 | f.write(s) 52 | 53 | # write f: ver ind/ uv ind 54 | [k, ntri] = triangles.shape 55 | for i in range(triangles.shape[0]): 56 | # s = 'f {} {} {}\n'.format(triangles[i, 0], triangles[i, 1], triangles[i, 2]) 57 | s = 'f {} {} {}\n'.format(triangles[i, 2], triangles[i, 1], triangles[i, 0]) 58 | f.write(s) 59 | 60 | ## TODO: c++ version 61 | def write_obj_with_texture(obj_name, vertices, triangles, texture, uv_coords): 62 | ''' Save 3D face model with texture represented by texture map. 63 | Ref: https://github.com/patrikhuber/eos/blob/bd00155ebae4b1a13b08bf5a991694d682abbada/include/eos/core/Mesh.hpp 64 | Args: 65 | obj_name: str 66 | vertices: shape = (nver, 3) 67 | triangles: shape = (ntri, 3) 68 | texture: shape = (256,256,3) 69 | uv_coords: shape = (nver, 3) max value<=1 70 | ''' 71 | if obj_name.split('.')[-1] != 'obj': 72 | obj_name = obj_name + '.obj' 73 | mtl_name = obj_name.replace('.obj', '.mtl') 74 | texture_name = obj_name.replace('.obj', '_texture.png') 75 | 76 | triangles = triangles.copy() 77 | triangles += 1 # mesh lab start with 1 78 | 79 | # write obj 80 | with open(obj_name, 'w') as f: 81 | # first line: write mtlib(material library) 82 | s = "mtllib {}\n".format(os.path.abspath(mtl_name)) 83 | f.write(s) 84 | 85 | # write vertices 86 | for i in range(vertices.shape[0]): 87 | s = 'v {} {} {}\n'.format(vertices[i, 0], vertices[i, 1], vertices[i, 2]) 88 | f.write(s) 89 | 90 | # write uv coords 91 | for i in range(uv_coords.shape[0]): 92 | s = 'vt {} {}\n'.format(uv_coords[i,0], 1 - uv_coords[i,1]) 93 | f.write(s) 94 | 95 | f.write("usemtl FaceTexture\n") 96 | 97 | # write f: ver ind/ uv ind 98 | for i in range(triangles.shape[0]): 99 | s = 'f {}/{} {}/{} {}/{}\n'.format(triangles[i,2], triangles[i,2], triangles[i,1], triangles[i,1], triangles[i,0], triangles[i,0]) 100 | f.write(s) 101 | 102 | # write mtl 103 | with open(mtl_name, 'w') as f: 104 | f.write("newmtl FaceTexture\n") 105 | s = 'map_Kd {}\n'.format(os.path.abspath(texture_name)) # map to image 106 | f.write(s) 107 | 108 | # write texture as png 109 | imsave(texture_name, texture) 110 | 111 | # c++ version 112 | def write_obj_with_colors_texture(obj_name, vertices, triangles, colors, texture, uv_coords): 113 | ''' Save 3D face model with texture. 114 | Ref: https://github.com/patrikhuber/eos/blob/bd00155ebae4b1a13b08bf5a991694d682abbada/include/eos/core/Mesh.hpp 115 | Args: 116 | obj_name: str 117 | vertices: shape = (nver, 3) 118 | triangles: shape = (ntri, 3) 119 | colors: shape = (nver, 3) 120 | texture: shape = (256,256,3) 121 | uv_coords: shape = (nver, 3) max value<=1 122 | ''' 123 | if obj_name.split('.')[-1] != 'obj': 124 | obj_name = obj_name + '.obj' 125 | mtl_name = obj_name.replace('.obj', '.mtl') 126 | texture_name = obj_name.replace('.obj', '_texture.png') 127 | 128 | triangles = triangles.copy() 129 | triangles += 1 # mesh lab start with 1 130 | 131 | # write obj 132 | vertices, colors, uv_coords = vertices.astype(np.float32).copy(), colors.astype(np.float32).copy(), uv_coords.astype(np.float32).copy() 133 | mesh_core_cython.write_obj_with_colors_texture_core(str.encode(obj_name), str.encode(os.path.abspath(mtl_name)), vertices, triangles, colors, uv_coords, vertices.shape[0], triangles.shape[0], uv_coords.shape[0]) 134 | 135 | # write mtl 136 | with open(mtl_name, 'w') as f: 137 | f.write("newmtl FaceTexture\n") 138 | s = 'map_Kd {}\n'.format(os.path.abspath(texture_name)) # map to image 139 | f.write(s) 140 | 141 | # write texture as png 142 | io.imsave(texture_name, texture) -------------------------------------------------------------------------------- /lib/mesh_p/light.py: -------------------------------------------------------------------------------- 1 | ''' 2 | Functions about lighting mesh(changing colors/texture of mesh). 3 | 1. add light to colors/texture (shade each vertex) 4 | 2. fit light according to colors/texture & image. 5 | ''' 6 | 7 | from __future__ import absolute_import 8 | from __future__ import division 9 | from __future__ import print_function 10 | 11 | import numpy as np 12 | from .cython import mesh_core_cython 13 | 14 | def get_normal(vertices, triangles): 15 | ''' calculate normal direction in each vertex 16 | Args: 17 | vertices: [nver, 3] 18 | triangles: [ntri, 3] 19 | Returns: 20 | normal: [nver, 3] 21 | ''' 22 | pt0 = vertices[triangles[:, 0], :] # [ntri, 3] 23 | pt1 = vertices[triangles[:, 1], :] # [ntri, 3] 24 | pt2 = vertices[triangles[:, 2], :] # [ntri, 3] 25 | tri_normal = np.cross(pt0 - pt1, pt0 - pt2) # [ntri, 3]. normal of each triangle 26 | 27 | normal = np.zeros_like(vertices, dtype = np.float32).copy() # [nver, 3] 28 | # for i in range(triangles.shape[0]): 29 | # normal[triangles[i, 0], :] = normal[triangles[i, 0], :] + tri_normal[i, :] 30 | # normal[triangles[i, 1], :] = normal[triangles[i, 1], :] + tri_normal[i, :] 31 | # normal[triangles[i, 2], :] = normal[triangles[i, 2], :] + tri_normal[i, :] 32 | mesh_core_cython.get_normal_core(normal, tri_normal.astype(np.float32).copy(), triangles.copy(), triangles.shape[0]) 33 | 34 | # normalize to unit length 35 | mag = np.sum(normal**2, 1) # [nver] 36 | zero_ind = (mag == 0) 37 | mag[zero_ind] = 1; 38 | normal[zero_ind, 0] = np.ones((np.sum(zero_ind))) 39 | 40 | normal = normal/np.sqrt(mag[:,np.newaxis]) 41 | 42 | return normal 43 | 44 | # TODO: test 45 | def add_light_sh(vertices, triangles, colors, sh_coeff): 46 | ''' 47 | In 3d face, usually assume: 48 | 1. The surface of face is Lambertian(reflect only the low frequencies of lighting) 49 | 2. Lighting can be an arbitrary combination of point sources 50 | --> can be expressed in terms of spherical harmonics(omit the lighting coefficients) 51 | I = albedo * (sh(n) x sh_coeff) 52 | 53 | albedo: n x 1 54 | sh_coeff: 9 x 1 55 | Y(n) = (1, n_x, n_y, n_z, n_xn_y, n_xn_z, n_yn_z, n_x^2 - n_y^2, 3n_z^2 - 1)': n x 9 56 | # Y(n) = (1, n_x, n_y, n_z)': n x 4 57 | 58 | Args: 59 | vertices: [nver, 3] 60 | triangles: [ntri, 3] 61 | colors: [nver, 3] albedo 62 | sh_coeff: [9, 1] spherical harmonics coefficients 63 | 64 | Returns: 65 | lit_colors: [nver, 3] 66 | ''' 67 | assert vertices.shape[0] == colors.shape[0] 68 | nver = vertices.shape[0] 69 | normal = get_normal(vertices, triangles) # [nver, 3] 70 | sh = np.array((np.ones(nver), n[:,0], n[:,1], n[:,2], n[:,0]*n[:,1], n[:,0]*n[:,2], n[:,1]*n[:,2], n[:,0]**2 - n[:,1]**2, 3*(n[:,2]**2) - 1)) # [nver, 9] 71 | ref = sh.dot(sh_coeff) #[nver, 1] 72 | lit_colors = colors*ref 73 | return lit_colors 74 | 75 | 76 | def add_light(vertices, triangles, colors, light_positions = 0, light_intensities = 0): 77 | ''' Gouraud shading. add point lights. 78 | In 3d face, usually assume: 79 | 1. The surface of face is Lambertian(reflect only the low frequencies of lighting) 80 | 2. Lighting can be an arbitrary combination of point sources 81 | 3. No specular (unless skin is oil, 23333) 82 | 83 | Ref: https://cs184.eecs.berkeley.edu/lecture/pipeline 84 | Args: 85 | vertices: [nver, 3] 86 | triangles: [ntri, 3] 87 | light_positions: [nlight, 3] 88 | light_intensities: [nlight, 3] 89 | Returns: 90 | lit_colors: [nver, 3] 91 | ''' 92 | nver = vertices.shape[0] 93 | normals = get_normal(vertices, triangles) # [nver, 3] 94 | 95 | # ambient 96 | # La = ka*Ia 97 | 98 | # diffuse 99 | # Ld = kd*(I/r^2)max(0, nxl) 100 | direction_to_lights = vertices[np.newaxis, :, :] - light_positions[:, np.newaxis, :] # [nlight, nver, 3] 101 | direction_to_lights_n = np.sqrt(np.sum(direction_to_lights**2, axis = 2)) # [nlight, nver] 102 | direction_to_lights = direction_to_lights/direction_to_lights_n[:, :, np.newaxis] 103 | normals_dot_lights = normals[np.newaxis, :, :]*direction_to_lights # [nlight, nver, 3] 104 | normals_dot_lights = np.sum(normals_dot_lights, axis = 2) # [nlight, nver] 105 | diffuse_output = colors[np.newaxis, :, :]*normals_dot_lights[:, :, np.newaxis]*light_intensities[:, np.newaxis, :] 106 | diffuse_output = np.sum(diffuse_output, axis = 0) # [nver, 3] 107 | 108 | # specular 109 | # h = (v + l)/(|v + l|) bisector 110 | # Ls = ks*(I/r^2)max(0, nxh)^p 111 | # increasing p narrows the reflectionlob 112 | 113 | lit_colors = diffuse_output # only diffuse part here. 114 | lit_colors = np.minimum(np.maximum(lit_colors, 0), 1) 115 | return lit_colors 116 | 117 | 118 | 119 | ## TODO. estimate light(sh coeff) 120 | ## -------------------------------- estimate. can not use now. 121 | def fit_light(image, vertices, colors, triangles, vis_ind, lamb = 10, max_iter = 3): 122 | [h, w, c] = image.shape 123 | 124 | # surface normal 125 | norm = get_normal(vertices, triangles) 126 | 127 | nver = vertices.shape[1] 128 | 129 | # vertices --> corresponding image pixel 130 | pt2d = vertices[:2, :] 131 | 132 | pt2d[0,:] = np.minimum(np.maximum(pt2d[0,:], 0), w - 1) 133 | pt2d[1,:] = np.minimum(np.maximum(pt2d[1,:], 0), h - 1) 134 | pt2d = np.round(pt2d).astype(np.int32) # 2 x nver 135 | 136 | image_pixel = image[pt2d[1,:], pt2d[0,:], :] # nver x 3 137 | image_pixel = image_pixel.T # 3 x nver 138 | 139 | # vertices --> corresponding mean texture pixel with illumination 140 | # Spherical Harmonic Basis 141 | harmonic_dim = 9 142 | nx = norm[0,:]; 143 | ny = norm[1,:]; 144 | nz = norm[2,:]; 145 | harmonic = np.zeros((nver, harmonic_dim)) 146 | 147 | pi = np.pi 148 | harmonic[:,0] = np.sqrt(1/(4*pi)) * np.ones((nver,)); 149 | harmonic[:,1] = np.sqrt(3/(4*pi)) * nx; 150 | harmonic[:,2] = np.sqrt(3/(4*pi)) * ny; 151 | harmonic[:,3] = np.sqrt(3/(4*pi)) * nz; 152 | harmonic[:,4] = 1/2. * np.sqrt(3/(4*pi)) * (2*nz**2 - nx**2 - ny**2); 153 | harmonic[:,5] = 3 * np.sqrt(5/(12*pi)) * (ny*nz); 154 | harmonic[:,6] = 3 * np.sqrt(5/(12*pi)) * (nx*nz); 155 | harmonic[:,7] = 3 * np.sqrt(5/(12*pi)) * (nx*ny); 156 | harmonic[:,8] = 3/2. * np.sqrt(5/(12*pi)) * (nx*nx - ny*ny); 157 | 158 | ''' 159 | I' = sum(albedo * lj * hj) j = 0:9 (albedo = tex) 160 | set A = albedo*h (n x 9) 161 | alpha = lj (9 x 1) 162 | Y = I (n x 1) 163 | Y' = A.dot(alpha) 164 | 165 | opt function: 166 | ||Y - A*alpha|| + lambda*(alpha'*alpha) 167 | result: 168 | A'*(Y - A*alpha) + lambda*alpha = 0 169 | ==> 170 | (A'*A*alpha - lambda)*alpha = A'*Y 171 | left: 9 x 9 172 | right: 9 x 1 173 | ''' 174 | n_vis_ind = len(vis_ind) 175 | n = n_vis_ind*c 176 | 177 | Y = np.zeros((n, 1)) 178 | A = np.zeros((n, 9)) 179 | light = np.zeros((3, 1)) 180 | 181 | for k in range(c): 182 | Y[k*n_vis_ind:(k+1)*n_vis_ind, :] = image_pixel[k, vis_ind][:, np.newaxis] 183 | A[k*n_vis_ind:(k+1)*n_vis_ind, :] = texture[k, vis_ind][:, np.newaxis] * harmonic[vis_ind, :] 184 | Ac = texture[k, vis_ind][:, np.newaxis] 185 | Yc = image_pixel[k, vis_ind][:, np.newaxis] 186 | light[k] = (Ac.T.dot(Yc))/(Ac.T.dot(Ac)) 187 | 188 | for i in range(max_iter): 189 | 190 | Yc = Y.copy() 191 | for k in range(c): 192 | Yc[k*n_vis_ind:(k+1)*n_vis_ind, :] /= light[k] 193 | 194 | # update alpha 195 | equation_left = np.dot(A.T, A) + lamb*np.eye(harmonic_dim); # why + ? 196 | equation_right = np.dot(A.T, Yc) 197 | alpha = np.dot(np.linalg.inv(equation_left), equation_right) 198 | 199 | # update light 200 | for k in range(c): 201 | Ac = A[k*n_vis_ind:(k+1)*n_vis_ind, :].dot(alpha) 202 | Yc = Y[k*n_vis_ind:(k+1)*n_vis_ind, :] 203 | light[k] = (Ac.T.dot(Yc))/(Ac.T.dot(Ac)) 204 | 205 | appearance = np.zeros_like(texture) 206 | for k in range(c): 207 | tmp = np.dot(harmonic*texture[k, :][:, np.newaxis], alpha*light[k]) 208 | appearance[k,:] = tmp.T 209 | 210 | appearance = np.minimum(np.maximum(appearance, 0), 1) 211 | 212 | return appearance 213 | 214 | -------------------------------------------------------------------------------- /lib/mesh_p/readme.md: -------------------------------------------------------------------------------- 1 | ## Extension from https://github.com/yfeng95/face3d 2 | 3 | 4 | ### Compile c++ files to .so for python use 5 | 6 | ```bash 7 | cd mesh_p/cython 8 | python setup.py build_ext -i 9 | ``` 10 | 11 | -------------------------------------------------------------------------------- /lib/mesh_p/render.py: -------------------------------------------------------------------------------- 1 | ''' 2 | functions about rendering mesh(from 3d obj to 2d image). 3 | only use rasterization render here. 4 | Note that: 5 | 1. Generally, render func includes camera, light, raterize. Here no camera and light(I write these in other files) 6 | 2. Generally, the input vertices are normalized to [-1,1] and cetered on [0, 0]. (in world space) 7 | Here, the vertices are using image coords, which centers on [w/2, h/2] with the y-axis pointing to oppisite direction. 8 | Means: render here only conducts interpolation.(I just want to make the input flexible) 9 | 10 | Author: Yao Feng 11 | Mail: yaofeng1995@gmail.com 12 | ''' 13 | from __future__ import absolute_import 14 | from __future__ import division 15 | from __future__ import print_function 16 | 17 | import numpy as np 18 | from time import time 19 | 20 | from .cython import mesh_core_cython 21 | 22 | def rasterize_triangles(vertices, triangles, h, w): 23 | ''' 24 | Args: 25 | vertices: [nver, 3] 26 | triangles: [ntri, 3] 27 | h: height 28 | w: width 29 | Returns: 30 | depth_buffer: [h, w] saves the depth, here, the bigger the z, the fronter the point. 31 | triangle_buffer: [h, w] saves the tri id(-1 for no triangle). 32 | barycentric_weight: [h, w, 3] saves corresponding barycentric weight. 33 | 34 | # Each triangle has 3 vertices & Each vertex has 3 coordinates x, y, z. 35 | # h, w is the size of rendering 36 | ''' 37 | 38 | # initial 39 | depth_buffer = np.zeros([h, w]) - 999999. #set the initial z to the farest position 40 | triangle_buffer = np.zeros([h, w], dtype = np.int32) - 1 # if tri id = -1, the pixel has no triangle correspondance 41 | barycentric_weight = np.zeros([h, w, 3], dtype = np.float32) # 42 | 43 | vertices = vertices.astype(np.float32).copy() 44 | triangles = triangles.astype(np.int32).copy() 45 | 46 | mesh_core_cython.rasterize_triangles_core( 47 | vertices, triangles, 48 | depth_buffer, triangle_buffer, barycentric_weight, 49 | vertices.shape[0], triangles.shape[0], 50 | h, w) 51 | 52 | def render_colors(vertices, triangles, colors, h, w, c = 3, BG = None): 53 | ''' render mesh with colors 54 | Args: 55 | vertices: [nver, 3] 56 | triangles: [ntri, 3] 57 | colors: [nver, 3] 58 | h: height 59 | w: width 60 | c: channel 61 | BG: background image 62 | Returns: 63 | image: [h, w, c]. rendered image./rendering. 64 | corr_mat: [h,w,nver] 65 | ''' 66 | 67 | # initial 68 | if BG is None: 69 | image = np.zeros((h, w, c), dtype = np.float32) 70 | else: 71 | assert BG.shape[0] == h and BG.shape[1] == w and BG.shape[2] == c 72 | image = BG 73 | depth_buffer = np.zeros([h, w], dtype = np.float32, order = 'C') - 999999. 74 | corr_mat = np.zeros((h, w, 6), dtype = np.float32) 75 | # change orders. --> C-contiguous order(column major) 76 | vertices = vertices.astype(np.float32).copy() 77 | triangles = triangles.astype(np.int32).copy() 78 | colors = colors.astype(np.float32).copy() 79 | ### 80 | st = time() 81 | mesh_core_cython.render_colors_core( 82 | image, corr_mat, vertices, triangles, 83 | colors, 84 | depth_buffer, 85 | vertices.shape[0], triangles.shape[0], 86 | h, w, c) 87 | return image, corr_mat 88 | 89 | 90 | def render_texture(vertices, triangles, texture, tex_coords, tex_triangles, h, w, c = 3, mapping_type = 'nearest', BG = None): 91 | ''' render mesh with texture map 92 | Args: 93 | vertices: [3, nver] 94 | triangles: [3, ntri] 95 | texture: [tex_h, tex_w, 3] 96 | tex_coords: [ntexcoords, 3] 97 | tex_triangles: [ntri, 3] 98 | h: height of rendering 99 | w: width of rendering 100 | c: channel 101 | mapping_type: 'bilinear' or 'nearest' 102 | ''' 103 | # initial 104 | if BG is None: 105 | image = np.zeros((h, w, c), dtype = np.float32) 106 | else: 107 | assert BG.shape[0] == h and BG.shape[1] == w and BG.shape[2] == c 108 | image = BG 109 | 110 | depth_buffer = np.zeros([h, w], dtype = np.float32, order = 'C') - 999999. 111 | 112 | tex_h, tex_w, tex_c = texture.shape 113 | if mapping_type == 'nearest': 114 | mt = int(0) 115 | elif mapping_type == 'bilinear': 116 | mt = int(1) 117 | else: 118 | mt = int(0) 119 | 120 | # -> C order 121 | vertices = vertices.astype(np.float32).copy() 122 | triangles = triangles.astype(np.int32).copy() 123 | texture = texture.astype(np.float32).copy() 124 | tex_coords = tex_coords.astype(np.float32).copy() 125 | tex_triangles = tex_triangles.astype(np.int32).copy() 126 | 127 | mesh_core_cython.render_texture_core( 128 | image, vertices, triangles, 129 | texture, tex_coords, tex_triangles, 130 | depth_buffer, 131 | vertices.shape[0], tex_coords.shape[0], triangles.shape[0], 132 | h, w, c, 133 | tex_h, tex_w, tex_c, 134 | mt) 135 | return image 136 | 137 | -------------------------------------------------------------------------------- /lib/mesh_p/vis.py: -------------------------------------------------------------------------------- 1 | from __future__ import absolute_import 2 | from __future__ import division 3 | from __future__ import print_function 4 | 5 | import numpy as np 6 | import matplotlib.pyplot as plt 7 | from skimage import measure 8 | from mpl_toolkits.mplot3d import Axes3D 9 | 10 | def plot_mesh(vertices, triangles, subplot = [1,1,1], title = 'mesh', el = 90, az = -90, lwdt=.1, dist = 6, color = "grey"): 11 | ''' 12 | plot the mesh 13 | Args: 14 | vertices: [nver, 3] 15 | triangles: [ntri, 3] 16 | ''' 17 | ax = plt.subplot(subplot[0], subplot[1], subplot[2], projection = '3d') 18 | ax.plot_trisurf(vertices[:, 0], vertices[:, 1], vertices[:, 2], triangles = triangles, lw = lwdt, color = color, alpha = 1) 19 | ax.axis("off") 20 | ax.view_init(elev = el, azim = az) 21 | ax.dist = dist 22 | plt.title(title) 23 | 24 | ### -------------- Todo: use vtk to visualize mesh? or visvis? or VisPy? 25 | -------------------------------------------------------------------------------- /models/__init__.py: -------------------------------------------------------------------------------- 1 | """This package contains modules related to objective functions, optimizations, and network architectures. 2 | 3 | To add a custom model class called 'dummy', you need to add a file called 'dummy_model.py' and define a subclass DummyModel inherited from BaseModel. 4 | You need to implement the following five functions: 5 | -- <__init__>: initialize the class; first call BaseModel.__init__(self, opt). 6 | -- : unpack data from dataset and apply preprocessing. 7 | -- : produce intermediate results. 8 | -- : calculate loss, gradients, and update network weights. 9 | -- : (optionally) add model-specific options and set default options. 10 | 11 | In the function <__init__>, you need to define four lists: 12 | -- self.loss_names (str list): specify the training losses that you want to plot and save. 13 | -- self.model_names (str list): define networks used in our training. 14 | -- self.visual_names (str list): specify the images that you want to display and save. 15 | -- self.optimizers (optimizer list): define and initialize optimizers. You can define one optimizer for each network. If two networks are updated at the same time, you can use itertools.chain to group them. See cycle_gan_model.py for an usage. 16 | 17 | Now you can use the model class by specifying flag '--model dummy'. 18 | See our template model class 'template_model.py' for more details. 19 | """ 20 | 21 | import importlib 22 | from models.base_model import BaseModel 23 | 24 | 25 | def find_model_using_name(model_name): 26 | """Import the module "models/[model_name]_model.py". 27 | 28 | In the file, the class called DatasetNameModel() will 29 | be instantiated. It has to be a subclass of BaseModel, 30 | and it is case-insensitive. 31 | """ 32 | model_filename = "models." + model_name + "_model" 33 | modellib = importlib.import_module(model_filename) 34 | model = None 35 | target_model_name = model_name.replace('_', '') + 'model' 36 | for name, cls in modellib.__dict__.items(): 37 | if name.lower() == target_model_name.lower() \ 38 | and issubclass(cls, BaseModel): 39 | model = cls 40 | 41 | if model is None: 42 | print("In %s.py, there should be a subclass of BaseModel with class name that matches %s in lowercase." % (model_filename, target_model_name)) 43 | exit(0) 44 | 45 | return model 46 | 47 | 48 | def get_option_setter(model_name): 49 | """Return the static method of the model class.""" 50 | model_class = find_model_using_name(model_name) 51 | return model_class.modify_commandline_options 52 | 53 | 54 | def create_model(opt): 55 | """Create a model given the option. 56 | 57 | This function warps the class CustomDatasetDataLoader. 58 | This is the main interface between this package and 'train.py'/'test.py' 59 | 60 | Example: 61 | >>> from models import create_model 62 | >>> model = create_model(opt) 63 | """ 64 | model = find_model_using_name(opt.model) 65 | instance = model(opt) 66 | print("model [%s] was created" % type(instance).__name__) 67 | return instance 68 | -------------------------------------------------------------------------------- /models/loss.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | import torch.nn.functional as F 4 | 5 | 6 | class Loss(nn.Module): 7 | def __init__(self, corr_wt=1.0, l1_wt=1.0, uv_wt=1.0, mat_wt=0.1, seg_wt=0.1): 8 | super(Loss, self).__init__() 9 | self.threshold = 0.01 10 | self.SegLoss=nn.CrossEntropyLoss(reduce=False) 11 | self.UVloss = torch.nn.L1Loss() 12 | self.corr_wt = corr_wt 13 | self.l1_wt = l1_wt 14 | self.uv_wt = uv_wt 15 | self.seg_wt = seg_wt 16 | self.mat_wt = mat_wt 17 | mask = torch.ones(3660) 18 | idx = [600,601,602,591,592,593,981,982,983,987,988,989,1044,1045,1046,2343,2344,2345,2289,2290,2291,2286,2287,2288,1938,1939,1940,1947,1948,1949,45,46,47,39,40,41,30,31,32,24,25,26,924,925,926,228,229,230,12,13,14,1575,1576,1577,2229,2230,2231,3303,3304,3305,3288,3289,3290,3276,3277,3278,3579,3580,3581,3588,3589,3590,3318,3319,3320,3504,3505,3506,3234,3235,3236,3222,3223,3224,3207,3208,3209,3483,3484,3485,3495,3496,3497,570,571,572,324,325,326,282,283,284,63,64,65,1629,1630,1631,1671,1672,1673,1917,1918,1919,2139,2140,2141,2118,2119,2120,84,85,86,813,814,815,834,835,836,1218,1219,1220,285,286,287,69,70,71,1632,1633,1634,2508,2509,2510,2124,2125,2126,78,79,80,819,820,821] 19 | mask[idx] = 10 20 | self.register_buffer('mask', torch.FloatTensor(mask).to(device='cuda')) 21 | 22 | def forward(self, assign_mat, seg_pred, ver3d_pred, uv_pred, nocs, model, seg_gt, mat_gt, uv_gt): 23 | """ 24 | Args: 25 | assign_mat: bs x n_pts x nv 26 | deltas: bs x nv x 3 27 | prior: bs x nv x 3 28 | """ 29 | inst_shape = ver3d_pred/9.0 30 | bs,n_pts,_=inst_shape.shape 31 | # smooth L1 loss for correspondences 32 | soft_assign = F.softmax(assign_mat, dim=2) 33 | #mat_loss=mat_wt* torch.mean(torch.sum(torch.abs(soft_assign-mat_gt),dim=2)) 34 | coords = torch.bmm(soft_assign, inst_shape*4.5) # bs x n_pts x 3 35 | diff = torch.abs(coords - nocs) 36 | less = torch.pow(diff, 2) / (2.0 * self.threshold) 37 | higher = diff - self.threshold / 2.0 38 | corr_loss = torch.where(diff > self.threshold, higher, less) 39 | corr_loss = torch.mean(torch.sum(corr_loss, dim=2)) 40 | corr_loss = self.corr_wt * corr_loss 41 | # kl loss to for distribution 42 | log_assign = F.log_softmax(assign_mat, dim=2) 43 | #kl_loss =self.klLoss(log_assign.view(-1, n_pts),F.softmax(mat_gt, dim=2).view(-1,n_pts)) 44 | #kl_loss =self.klLoss(log_assign,F.softmax(mat_gt, dim=2)) 45 | #print(corr_loss.item()) 46 | if corr_loss.item()<0.01: 47 | self.mat_wt=0.001 48 | self.uv_wt=0.5 49 | self.l1_wt=0.5 50 | elif corr_loss.item()<0.1: 51 | self.mat_wt =0.01 52 | self.seg_wt=0.01 53 | self.l1_wt =1.0 54 | kl_loss =F.kl_div(log_assign, mat_gt, size_average = False) 55 | kl_loss =self.mat_wt * kl_loss/(bs*n_pts) 56 | # entropy loss to encourage peaked distribution 57 | entropy_loss = torch.mean(-torch.sum(soft_assign * log_assign, 2)) 58 | entropy_loss = 0.001 * entropy_loss 59 | 60 | uv_loss= self.UVloss(uv_pred, uv_gt)*self.uv_wt 61 | l1_loss= self.UVloss(ver3d_pred, model*9.0) 62 | 63 | l1_loss = self.l1_wt * l1_loss 64 | seg_loss = self.SegLoss(seg_pred, seg_gt) 65 | 66 | seg_loss =self.seg_wt * seg_loss.mean() 67 | 68 | # total loss 69 | total_loss = corr_loss + l1_loss + kl_loss + seg_loss+ entropy_loss 70 | return total_loss, corr_loss, l1_loss, uv_loss, kl_loss, seg_loss 71 | 72 | 73 | 74 | -------------------------------------------------------------------------------- /models/model_repository.py: -------------------------------------------------------------------------------- 1 | from torch import nn 2 | import torch 3 | from torch.nn import functional as F 4 | from .resnet import resnet18, resnet50, resnet34 5 | 6 | 7 | 8 | class Resnet18_8s(nn.Module): 9 | def __init__(self, fcdim=256, s8dim=128, s4dim=64, s2dim=32, raw_dim=32): 10 | super(Resnet18_8s, self).__init__() 11 | 12 | # Load the pretrained weights, remove avg pool 13 | # layer and get the output stride of 8 14 | resnet18_8s = resnet18(fully_conv=True, 15 | pretrained=True, 16 | output_stride=32, 17 | remove_avg_pool_layer=True) 18 | 19 | 20 | # Randomly initialize the 1x1 Conv scoring layer 21 | resnet18_8s.fc = nn.Sequential( 22 | nn.Conv2d(resnet18_8s.inplanes, fcdim, 3, 1, 1, bias=False), 23 | nn.BatchNorm2d(fcdim), 24 | nn.ReLU(True) 25 | ) 26 | self.resnet18_8s = resnet18_8s 27 | 28 | self.up32sto16s=nn.UpsamplingBilinear2d(scale_factor=2) 29 | # x16s->256 30 | self.conv16s=nn.Sequential( 31 | nn.Conv2d(fcdim+256, fcdim, 3, 1, 1, bias=False), 32 | nn.BatchNorm2d(fcdim), 33 | nn.LeakyReLU(0.1,True) 34 | ) 35 | self.up16sto8s=nn.UpsamplingBilinear2d(scale_factor=2) 36 | #128 37 | self.conv16s_uv=nn.Sequential( 38 | nn.Conv2d(fcdim, fcdim, 3, 1, 1, bias=False), 39 | nn.BatchNorm2d(fcdim), 40 | nn.LeakyReLU(0.1,True) 41 | ) 42 | self.up16sto8s_uv=nn.UpsamplingBilinear2d(scale_factor=2) 43 | 44 | # x8s->128 45 | self.conv8s=nn.Sequential( 46 | nn.Conv2d(128+fcdim, s8dim, 3, 1, 1, bias=False), 47 | nn.BatchNorm2d(s8dim), 48 | nn.LeakyReLU(0.1,True) 49 | ) 50 | self.up8sto4s=nn.UpsamplingBilinear2d(scale_factor=2) 51 | self.conv8s_uv=nn.Sequential( 52 | nn.Conv2d(fcdim, s8dim, 3, 1, 1, bias=False), 53 | nn.BatchNorm2d(s8dim), 54 | nn.LeakyReLU(0.1,True) 55 | ) 56 | self.up8sto4s_uv=nn.UpsamplingBilinear2d(scale_factor=2) 57 | 58 | # x4s->64 59 | self.conv4s=nn.Sequential( 60 | nn.Conv2d(64+s8dim, s4dim, 3, 1, 1, bias=False), 61 | nn.BatchNorm2d(s4dim), 62 | nn.LeakyReLU(0.1,True) 63 | ) 64 | self.up4sto2s=nn.UpsamplingBilinear2d(scale_factor=2) 65 | self.conv4s_uv=nn.Sequential( 66 | nn.Conv2d(s8dim, s4dim, 3, 1, 1, bias=False), 67 | nn.BatchNorm2d(s4dim), 68 | nn.LeakyReLU(0.1,True) 69 | ) 70 | self.up4sto2s_uv=nn.UpsamplingBilinear2d(scale_factor=2) 71 | 72 | # x2s->64 73 | self.conv2s=nn.Sequential( 74 | nn.Conv2d(64+s4dim, s2dim, 3, 1, 1, bias=False), 75 | nn.BatchNorm2d(s2dim), 76 | nn.LeakyReLU(0.1,True) 77 | ) 78 | self.up2storaw = nn.UpsamplingBilinear2d(scale_factor=2) 79 | self.conv2s_uv=nn.Sequential( 80 | nn.Conv2d(s4dim, s2dim, 3, 1, 1, bias=False), 81 | nn.BatchNorm2d(s2dim), 82 | nn.LeakyReLU(0.1,True) 83 | ) 84 | self.up2storaw_uv = nn.UpsamplingBilinear2d(scale_factor=2) 85 | 86 | self.final = nn.Sequential( 87 | nn.Conv2d(3+s2dim, raw_dim, 3, 1, 1, bias=False), 88 | nn.BatchNorm2d(raw_dim), 89 | nn.LeakyReLU(0.1,True) 90 | ) 91 | self.final_uv = nn.Sequential( 92 | nn.Conv2d(s2dim, raw_dim, 3, 1, 1, bias=False), 93 | nn.BatchNorm2d(raw_dim), 94 | nn.LeakyReLU(0.1,True) 95 | ) 96 | 97 | self.uptorawsize=nn.UpsamplingBilinear2d(scale_factor=2) 98 | self.seg=nn.Conv2d(raw_dim, 2, 1, 1) 99 | self.uv=nn.Conv2d(raw_dim, 3, 1, 1) 100 | self.drop = nn.Dropout2d(p=0.15) 101 | 102 | def _normal_initialization(self, layer): 103 | layer.weight.data.normal_(0, 0.01) 104 | layer.bias.data.zero_() 105 | 106 | def forward(self, x, feature_alignment=False): 107 | x2s, x4s, x8s, x16s, x32s, xfc = self.resnet18_8s(x) 108 | 109 | 110 | x16=self.up32sto16s(xfc) #256, 16s 111 | 112 | 113 | fm=self.conv16s(torch.cat([x16,x16s],1)) 114 | fm=self.up16sto8s(fm) 115 | #pdb.set_trace() 116 | fm=self.conv8s(torch.cat([fm,x8s],1)) 117 | fm=self.up8sto4s(fm) 118 | fm=self.drop(fm) 119 | 120 | 121 | fm=self.conv4s(torch.cat([fm,x4s],1)) 122 | fm=self.up4sto2s(fm) 123 | fm=self.drop(fm) 124 | 125 | fm=self.conv2s(torch.cat([fm,x2s],1)) 126 | fm=self.up2storaw(fm) 127 | 128 | 129 | x=self.final(torch.cat([fm,x],1)) 130 | # fm=self.uptorawsize(fm) 131 | # x=self.drop(fm) 132 | 133 | #uv 134 | fm=self.conv16s_uv(x16) # 128 135 | fm=self.up16sto8s_uv(fm) 136 | 137 | fm=self.conv8s_uv(fm) 138 | fm=self.up8sto4s_uv(fm) 139 | fm=self.drop(fm) 140 | 141 | fm=self.conv4s_uv(fm) 142 | fm=self.up4sto2s_uv(fm) 143 | fm=self.drop(fm) 144 | 145 | fm=self.conv2s_uv(fm) 146 | fm=self.up2storaw_uv(fm) 147 | 148 | f_uv=self.final_uv(fm) 149 | uv=self.uv(f_uv) 150 | 151 | return x, self.seg(x), f_uv, torch.tanh(uv) 152 | 153 | 154 | -------------------------------------------------------------------------------- /models/network.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | from .model_repository import Resnet18_8s 4 | import math 5 | import torch.nn.functional as F 6 | 7 | class DeformNetUV2(nn.Module): 8 | def __init__(self, grid, nv_prior=1220, img_size=192): 9 | super(DeformNetUV2, self).__init__() 10 | self.cnn = Resnet18_8s() 11 | self.pos_encod=self.positionalencoding2d(64, 192, 192) 12 | self.grid = grid 13 | self.instance_color = nn.Sequential( 14 | nn.Conv1d(32, 64, 1), 15 | nn.ReLU(), 16 | ) 17 | self.instance_global = nn.Sequential( 18 | nn.Conv1d(128, 128, 1), 19 | nn.ReLU(), 20 | nn.Conv1d(128, 1024, 1), 21 | nn.ReLU(), 22 | nn.AdaptiveAvgPool1d(1), 23 | ) 24 | 25 | self.category_global = nn.Sequential( 26 | nn.Conv1d(32, 128, 1), 27 | nn.ReLU(), 28 | nn.Conv1d(128, 1024, 1), 29 | nn.ReLU(), 30 | nn.AdaptiveAvgPool1d(1), 31 | ) 32 | self.assignment = nn.Sequential( 33 | nn.Conv1d(2176, 512, 1), 34 | nn.ReLU(), 35 | nn.Conv1d(512, 256, 1), 36 | nn.ReLU(), 37 | nn.Conv1d(256, nv_prior, 1), 38 | ) 39 | 40 | def forward(self, img, choose): 41 | """ 42 | Args: 43 | img: bs x 3 x H x W 44 | choose: bs x n_pts 45 | cat_id: bs 46 | prior: bs x nv x 3 47 | 48 | Returns: 49 | assign_mat: bs x n_pts x nv 50 | inst_shape: bs x nv x 3 51 | deltas: bs x nv x 3 52 | log_assign: bs x n_pts x nv, for numerical stability 53 | 54 | """ 55 | #bs, n_pts = points.size()[:2] 56 | bs,n_pts =choose.size() 57 | 58 | nv = self.grid.shape[2] 59 | 60 | out_img, seg_pred, uv_feat, uv_pred = self.cnn(img) 61 | di = out_img.size()[1] 62 | emb = out_img.view(bs, di, -1) 63 | 64 | # print(seg_pred.size(),out_img.size()) 65 | choose = choose.unsqueeze(1).repeat(1, di*2, 1) 66 | # print(choose.size()) 67 | emb = torch.gather(emb, 2, choose[:,0:di,:]).contiguous() 68 | # print(emb.size()) 69 | emb = self.instance_color(emb) #bs*64*n_pts 70 | pos_encod=self.pos_encod.unsqueeze(0).repeat(bs,1,1,1) 71 | #dj = pos_encod.size()[1] 72 | pos_encod=pos_encod.view(bs,di*2,-1) 73 | pos_encod = torch.gather(pos_encod, 2, choose).contiguous() #bs*64*n_pts 74 | inst_local = torch.cat((emb, pos_encod), dim=1) # bs x 128 x n_pts 75 | inst_global = self.instance_global(inst_local) # bs x 1024 x 1 76 | uv_local = F.grid_sample(uv_feat, self.grid.expand(bs, -1, -1, -1), align_corners=False) 77 | uv_local = uv_local.squeeze(2) 78 | 79 | 80 | cat_global = self.category_global(uv_local) # bs x 1024 x 1 81 | # print('cat_global',cat_global.size()) 82 | 83 | # assignemnt matrix 84 | assign_feat = torch.cat((inst_local, inst_global.repeat(1, 1, n_pts), cat_global.repeat(1, 1, n_pts)), dim=1) # bs x 2176 x n_pts 85 | assign_mat = self.assignment(assign_feat) 86 | assign_mat = assign_mat.view(-1, nv, n_pts).contiguous() # bs, nc*nv, n_pts -> bs*nc, nv, n_pts 87 | #index = cat_id + torch.arange(bs, dtype=torch.long).cuda() * self.n_cat 88 | #assign_mat = torch.index_select(assign_mat, 0, index) # bs x nv x n_pts 89 | assign_mat = assign_mat.permute(0, 2, 1).contiguous() # bs x n_pts x nv 90 | 91 | return assign_mat, seg_pred, uv_pred 92 | 93 | def positionalencoding2d(self, d_model, height, width): 94 | """ 95 | :param d_model: dimension of the model 96 | :param height: height of the positions 97 | :param width: width of the positions 98 | :return: d_model*height*width position matrix 99 | """ 100 | if d_model % 4 != 0: 101 | raise ValueError("Cannot use sin/cos positional encoding with " 102 | "odd dimension (got dim={:d})".format(d_model)) 103 | pe = torch.zeros(d_model, height, width).to(device='cuda') 104 | # Each dimension use half of d_model 105 | d_model = int(d_model / 2) 106 | div_term = torch.exp(torch.arange(0., d_model, 2) * 107 | -(math.log(10000.0) / d_model)) 108 | pos_w = torch.arange(0., width).unsqueeze(1) 109 | pos_h = torch.arange(0., height).unsqueeze(1) 110 | pe[0:d_model:2, :, :] = torch.sin(pos_w * div_term).transpose(0, 1).unsqueeze(1).repeat(1, height, 1) 111 | pe[1:d_model:2, :, :] = torch.cos(pos_w * div_term).transpose(0, 1).unsqueeze(1).repeat(1, height, 1) 112 | pe[d_model::2, :, :] = torch.sin(pos_h * div_term).transpose(0, 1).unsqueeze(2).repeat(1, 1, width) 113 | pe[d_model + 1::2, :, :] = torch.cos(pos_h * div_term).transpose(0, 1).unsqueeze(2).repeat(1, 1, width) 114 | return pe 115 | 116 | 117 | -------------------------------------------------------------------------------- /models/networks.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | import torch.nn.functional as F 4 | from torch.nn import init 5 | import functools 6 | from torch.optim import lr_scheduler 7 | import numpy as np 8 | 9 | ############################################################################### 10 | # Helper Functions 11 | ############################################################################### 12 | 13 | 14 | def get_filter(filt_size=3): 15 | if(filt_size == 1): 16 | a = np.array([1., ]) 17 | elif(filt_size == 2): 18 | a = np.array([1., 1.]) 19 | elif(filt_size == 3): 20 | a = np.array([1., 2., 1.]) 21 | elif(filt_size == 4): 22 | a = np.array([1., 3., 3., 1.]) 23 | elif(filt_size == 5): 24 | a = np.array([1., 4., 6., 4., 1.]) 25 | elif(filt_size == 6): 26 | a = np.array([1., 5., 10., 10., 5., 1.]) 27 | elif(filt_size == 7): 28 | a = np.array([1., 6., 15., 20., 15., 6., 1.]) 29 | 30 | filt = torch.Tensor(a[:, None] * a[None, :]) 31 | filt = filt / torch.sum(filt) 32 | 33 | return filt 34 | 35 | 36 | class Downsample(nn.Module): 37 | def __init__(self, channels, pad_type='reflect', filt_size=3, stride=2, pad_off=0): 38 | super(Downsample, self).__init__() 39 | self.filt_size = filt_size 40 | self.pad_off = pad_off 41 | self.pad_sizes = [int(1. * (filt_size - 1) / 2), int(np.ceil(1. * (filt_size - 1) / 2)), int(1. * (filt_size - 1) / 2), int(np.ceil(1. * (filt_size - 1) / 2))] 42 | self.pad_sizes = [pad_size + pad_off for pad_size in self.pad_sizes] 43 | self.stride = stride 44 | self.off = int((self.stride - 1) / 2.) 45 | self.channels = channels 46 | 47 | filt = get_filter(filt_size=self.filt_size) 48 | self.register_buffer('filt', filt[None, None, :, :].repeat((self.channels, 1, 1, 1))) 49 | 50 | self.pad = get_pad_layer(pad_type)(self.pad_sizes) 51 | 52 | def forward(self, inp): 53 | if(self.filt_size == 1): 54 | if(self.pad_off == 0): 55 | return inp[:, :, ::self.stride, ::self.stride] 56 | else: 57 | return self.pad(inp)[:, :, ::self.stride, ::self.stride] 58 | else: 59 | return F.conv2d(self.pad(inp), self.filt, stride=self.stride, groups=inp.shape[1]) 60 | 61 | 62 | class Upsample2(nn.Module): 63 | def __init__(self, scale_factor, mode='nearest'): 64 | super().__init__() 65 | self.factor = scale_factor 66 | self.mode = mode 67 | 68 | def forward(self, x): 69 | return torch.nn.functional.interpolate(x, scale_factor=self.factor, mode=self.mode) 70 | 71 | 72 | class Upsample(nn.Module): 73 | def __init__(self, channels, pad_type='repl', filt_size=4, stride=2): 74 | super(Upsample, self).__init__() 75 | self.filt_size = filt_size 76 | self.filt_odd = np.mod(filt_size, 2) == 1 77 | self.pad_size = int((filt_size - 1) / 2) 78 | self.stride = stride 79 | self.off = int((self.stride - 1) / 2.) 80 | self.channels = channels 81 | 82 | filt = get_filter(filt_size=self.filt_size) * (stride**2) 83 | self.register_buffer('filt', filt[None, None, :, :].repeat((self.channels, 1, 1, 1))) 84 | 85 | self.pad = get_pad_layer(pad_type)([1, 1, 1, 1]) 86 | 87 | def forward(self, inp): 88 | ret_val = F.conv_transpose2d(self.pad(inp), self.filt, stride=self.stride, padding=1 + self.pad_size, groups=inp.shape[1])[:, :, 1:, 1:] 89 | if(self.filt_odd): 90 | return ret_val 91 | else: 92 | return ret_val[:, :, :-1, :-1] 93 | 94 | 95 | def get_pad_layer(pad_type): 96 | if(pad_type in ['refl', 'reflect']): 97 | PadLayer = nn.ReflectionPad2d 98 | elif(pad_type in ['repl', 'replicate']): 99 | PadLayer = nn.ReplicationPad2d 100 | elif(pad_type == 'zero'): 101 | PadLayer = nn.ZeroPad2d 102 | else: 103 | print('Pad type [%s] not recognized' % pad_type) 104 | return PadLayer 105 | 106 | 107 | class Identity(nn.Module): 108 | def forward(self, x): 109 | return x 110 | 111 | 112 | def get_norm_layer(norm_type='instance'): 113 | """Return a normalization layer 114 | 115 | Parameters: 116 | norm_type (str) -- the name of the normalization layer: batch | instance | none 117 | 118 | For BatchNorm, we use learnable affine parameters and track running statistics (mean/stddev). 119 | For InstanceNorm, we do not use learnable affine parameters. We do not track running statistics. 120 | """ 121 | if norm_type == 'batch': 122 | norm_layer = functools.partial(nn.BatchNorm2d, affine=True, track_running_stats=True) 123 | elif norm_type == 'instance': 124 | norm_layer = functools.partial(nn.InstanceNorm2d, affine=False, track_running_stats=False) 125 | elif norm_type == 'none': 126 | def norm_layer(x): 127 | return Identity() 128 | else: 129 | raise NotImplementedError('normalization layer [%s] is not found' % norm_type) 130 | return norm_layer 131 | 132 | 133 | def get_scheduler(optimizer, opt): 134 | """Return a learning rate scheduler 135 | 136 | Parameters: 137 | optimizer -- the optimizer of the network 138 | opt (option class) -- stores all the experiment flags; needs to be a subclass of BaseOptions．　 139 | opt.lr_policy is the name of learning rate policy: linear | step | plateau | cosine 140 | 141 | For 'linear', we keep the same learning rate for the first epochs 142 | and linearly decay the rate to zero over the next epochs. 143 | For other schedulers (step, plateau, and cosine), we use the default PyTorch schedulers. 144 | See https://pytorch.org/docs/stable/optim.html for more details. 145 | """ 146 | if opt.lr_policy == 'linear': 147 | def lambda_rule(epoch): 148 | lr_l = 1.0 - max(0, epoch + opt.epoch_count - opt.n_epochs) / float(opt.n_epochs_decay + 1) 149 | return lr_l 150 | scheduler = lr_scheduler.LambdaLR(optimizer, lr_lambda=lambda_rule) 151 | elif opt.lr_policy == 'step': 152 | scheduler = lr_scheduler.StepLR(optimizer, step_size=opt.lr_decay_iters, gamma=0.1) 153 | elif opt.lr_policy == 'plateau': 154 | scheduler = lr_scheduler.ReduceLROnPlateau(optimizer, mode='min', factor=0.2, threshold=0.01, patience=5) 155 | elif opt.lr_policy == 'cosine': 156 | scheduler = lr_scheduler.CosineAnnealingLR(optimizer, T_max=opt.n_epochs, eta_min=0) 157 | else: 158 | return NotImplementedError('learning rate policy [%s] is not implemented', opt.lr_policy) 159 | return scheduler 160 | 161 | 162 | def init_weights(net, init_type='normal', init_gain=0.02, debug=False): 163 | """Initialize network weights. 164 | 165 | Parameters: 166 | net (network) -- network to be initialized 167 | init_type (str) -- the name of an initialization method: normal | xavier | kaiming | orthogonal 168 | init_gain (float) -- scaling factor for normal, xavier and orthogonal. 169 | 170 | We use 'normal' in the original pix2pix and CycleGAN paper. But xavier and kaiming might 171 | work better for some applications. Feel free to try yourself. 172 | """ 173 | def init_func(m): # define the initialization function 174 | classname = m.__class__.__name__ 175 | if hasattr(m, 'weight') and (classname.find('Conv') != -1 or classname.find('Linear') != -1): 176 | if debug: 177 | print(classname) 178 | if init_type == 'normal': 179 | init.normal_(m.weight.data, 0.0, init_gain) 180 | elif init_type == 'xavier': 181 | init.xavier_normal_(m.weight.data, gain=init_gain) 182 | elif init_type == 'kaiming': 183 | init.kaiming_normal_(m.weight.data, a=0, mode='fan_in') 184 | elif init_type == 'orthogonal': 185 | init.orthogonal_(m.weight.data, gain=init_gain) 186 | else: 187 | raise NotImplementedError('initialization method [%s] is not implemented' % init_type) 188 | if hasattr(m, 'bias') and m.bias is not None: 189 | init.constant_(m.bias.data, 0.0) 190 | elif classname.find('BatchNorm2d') != -1: # BatchNorm Layer's weight is not a matrix; only normal distribution applies. 191 | init.normal_(m.weight.data, 1.0, init_gain) 192 | init.constant_(m.bias.data, 0.0) 193 | 194 | net.apply(init_func) # apply the initialization function 195 | 196 | 197 | def init_net(net, init_type='normal', init_gain=0.02, gpu_ids=[], debug=False, initialize_weights=True): 198 | """Initialize a network: 1. register CPU/GPU device (with multi-GPU support); 2. initialize the network weights 199 | Parameters: 200 | net (network) -- the network to be initialized 201 | init_type (str) -- the name of an initialization method: normal | xavier | kaiming | orthogonal 202 | gain (float) -- scaling factor for normal, xavier and orthogonal. 203 | gpu_ids (int list) -- which GPUs the network runs on: e.g., 0,1,2 204 | 205 | Return an initialized network. 206 | """ 207 | if len(gpu_ids) > 0: 208 | assert(torch.cuda.is_available()) 209 | net.to(gpu_ids[0]) 210 | # if not amp: 211 | # net = torch.nn.DataParallel(net, gpu_ids) # multi-GPUs for non-AMP training 212 | if initialize_weights: 213 | init_weights(net, init_type, init_gain=init_gain, debug=debug) 214 | return net 215 | 216 | 217 | -------------------------------------------------------------------------------- /models/perspnet_model.py: -------------------------------------------------------------------------------- 1 | import matplotlib 2 | matplotlib.use('Agg') 3 | import matplotlib.pyplot as plt 4 | 5 | import pdb 6 | import cv2 7 | import socket 8 | import numpy as np 9 | from pathlib import Path 10 | from scipy.spatial.transform import Rotation 11 | 12 | import torch 13 | import torch.nn as nn 14 | import torch.nn.functional as F 15 | from torchvision.models.resnet import BasicBlock 16 | 17 | from . import networks 18 | from .base_model import BaseModel 19 | from .network import DeformNetUV2 20 | from .loss import Loss 21 | from util.face_deformnet_utils import compute_sRT_errors 22 | from shutil import copyfile 23 | 24 | np.set_printoptions(suppress=True) 25 | class PerspnetModel(BaseModel): 26 | 27 | ''' 28 | @staticmethod 29 | def modify_commandline_options(parser, is_train=True): 30 | return parser 31 | ''' 32 | 33 | 34 | def __init__(self, opt): 35 | BaseModel.__init__(self, opt) 36 | 37 | self.loss_names = ['total', 'corr', 'recon3d', 'uv', 'mat', 'seg'] 38 | self.model_names = ['R'] 39 | 40 | npy_path = 'npy/uv_coords_std_202109.npy' 41 | uv_coords_std = np.load(npy_path)[:1220] 42 | uv_coords = uv_coords_std * 2 - 1 43 | self.grid = torch.tensor(uv_coords).unsqueeze(0).unsqueeze(1).to(self.device) 44 | 45 | 46 | # networks.init_net 47 | self.netR = networks.init_net( 48 | DeformNetUV2(self.grid), 49 | gpu_ids=self.gpu_ids 50 | ) 51 | #self.netR.cuda() 52 | self.criterion = Loss() 53 | if self.isTrain: 54 | self.optimizer_R = torch.optim.Adam(self.netR.parameters(), lr=opt.lr, betas=(opt.beta1, 0.999)) 55 | self.optimizers.append(self.optimizer_R) 56 | 57 | 58 | 59 | def set_input(self, input): 60 | self.img = input['img'].to(self.device) 61 | self.choose = input['choose'].to(self.device) 62 | self.model= input['model'].to(self.device) 63 | self.uvmap = input['uvmap'].to(self.device) 64 | self.nocs = input['nocs'].to(self.device) 65 | self.mask = input['mask'].to(self.device) 66 | self.corr_mat = input['corr_mat'].to(self.device) 67 | 68 | if hasattr(self.opt, 'eval'): 69 | self.tform_inv_list = input['tform_inv'].numpy() 70 | self.R_t_list = input['R_t'].numpy() 71 | self.img_path_list = input['img_path'] 72 | 73 | 74 | def forward(self): 75 | #print(self.choose.max(),self.choose.min()) 76 | self.assign_mat, self.seg_pred, self.uv_pred = self.netR(self.img, self.choose) ### assign_mat, deltas, seg_pred 77 | bs = self.uv_pred.size(0) 78 | verts3d_pred = F.grid_sample(self.uv_pred, self.grid.expand(bs, -1, -1, -1), align_corners=False) # 低版本pytorch没有align_corners=False这个选项 79 | self.verts3d_pred = verts3d_pred.squeeze(2).permute(0, 2, 1).contiguous() 80 | 81 | def forward_seg(self): 82 | _, self.seg_pred, _, _=self.netR.module.cnn(self.img) 83 | 84 | 85 | def backward_R(self): 86 | self.loss_total, self.loss_corr, self.loss_recon3d, self.loss_uv, self.loss_mat, self.loss_seg = self.criterion(self.assign_mat, self.seg_pred, self.verts3d_pred, self.uv_pred, self.nocs, self.model, self.mask, self.corr_mat,self.uvmap) 87 | self.loss_total.backward() 88 | 89 | 90 | 91 | 92 | def optimize_parameters(self): 93 | self.netR.train() 94 | self.forward() 95 | self.optimizer_R.zero_grad() 96 | self.backward_R() 97 | self.optimizer_R.step() 98 | 99 | 100 | def init_evaluation(self): 101 | self.inference_data = { 102 | 'loss_total': [], 'loss_corr': [], 'loss_recon3d': [], 'loss_uv':[], 'loss_mat':[], 'loss_seg':[], 103 | 'batch_size': [], 104 | 'pitch_mae': [], 'yaw_mae': [], 'roll_mae': [], 105 | 'tx_mae': [], 'ty_mae': [], 'tz_mae': [], 106 | '3DRecon': [], 'ADD': [], 107 | 'strict_success': 0, 'easy_success': 0, 'total_count': 0, 108 | 'IoU': [], 109 | 'pnp_fail': 0 110 | } 111 | self.netR.eval() 112 | 113 | npz_path = 'npy/kpt_ind.npy' 114 | self.kpt_ind = np.load(npz_path) 115 | 116 | 117 | def inference_curr_batch(self): 118 | bs = self.img.size(0) 119 | with torch.no_grad(): 120 | if self.opt.seg: 121 | self.forward_seg() 122 | mask_pr = torch.argmax(self.seg_pred, 1).cpu().detach().numpy() 123 | chooses=np.zeros((bs,self.opt.n_pts)) 124 | for i in range(bs): 125 | choose = mask_pr[i].flatten().nonzero()[0] 126 | if len(choose) > self.opt.n_pts: 127 | c_mask = np.zeros(len(choose), dtype=int) 128 | c_mask[:self.opt.n_pts] = 1 129 | np.random.shuffle(c_mask) 130 | choose = choose[c_mask.nonzero()] 131 | else: 132 | choose = np.pad(choose, (0, self.opt.n_pts-len(choose)), 'wrap') 133 | chooses[i,:]=choose 134 | self.choose =chooses.astype(np.int64) 135 | self.choose=torch.LongTensor(self.choose).to(device='cuda') 136 | self.forward() 137 | 138 | cur_loss_total, cur_loss_corr, cur_loss_recon3d, cur_loss_uv, cur_loss_mat, cur_loss_seg = self.criterion(self.assign_mat, self.seg_pred, self.verts3d_pred, self.uv_pred, self.nocs, self.model, self.mask, self.corr_mat, self.uvmap) 139 | self.inference_data['loss_total'].append(cur_loss_total.item()) 140 | self.inference_data['loss_corr'].append(cur_loss_corr.item()) 141 | self.inference_data['loss_recon3d'].append(cur_loss_recon3d.item()) 142 | self.inference_data['loss_uv'].append(cur_loss_uv.item()) 143 | self.inference_data['loss_mat'].append(cur_loss_mat.item()) 144 | self.inference_data['loss_seg'].append(cur_loss_seg.item()) 145 | self.inference_data['batch_size'].append(bs) 146 | 147 | 148 | 149 | if hasattr(self.opt, 'eval'): 150 | # 3D vertices 151 | inst_shape=self.verts3d_pred/9.0 152 | verts3d_pred = inst_shape.cpu().numpy().reshape(-1, 1220, 3) 153 | verts3d_gt = (self.model).cpu().numpy().reshape(-1, 1220, 3) 154 | 155 | #pts68_pred = ((self.out1 * 0.5 + 0.5) * self.opt.img_size).cpu().numpy().reshape(-1, 68, 2) 156 | assign_mat = F.softmax(self.assign_mat, dim=2) 157 | nocs_coords = torch.bmm(assign_mat, inst_shape) 158 | nocs_coords = nocs_coords.detach().cpu().numpy().reshape(-1, self.opt.n_pts, 3) 159 | 160 | img_size = 800 161 | f = 1.574437 * img_size / 2 162 | K_img = np.array([ 163 | [f, 0, img_size / 2.0], 164 | [0, f, img_size / 2.0], 165 | [0, 0, 1] 166 | ]) 167 | T = np.array([ 168 | [1.0, 0, 0, 0], 169 | [ 0, -1, 0, 0], 170 | [ 0, 0, -1, 0], 171 | [ 0, 0, 0, 1] 172 | ]) 173 | 174 | for i in range(bs): 175 | choose = self.choose.cpu().numpy()[i] 176 | choose, choose_idx = np.unique(choose, return_index=True) 177 | nocs_coord = nocs_coords[i, choose_idx, :] 178 | col_idx = choose % self.opt.img_size 179 | row_idx = choose // self.opt.img_size 180 | local_pts2d = np.concatenate((col_idx.reshape(-1,1), row_idx.reshape(-1,1)), axis=1) 181 | 182 | tfm = self.tform_inv_list[i] 183 | W, b = tfm.T[:2], tfm.T[2] 184 | global_pts_pred = local_pts2d @ W + b 185 | 186 | if 1: 187 | _, rvecs, tvecs, _ = cv2.solvePnPRansac( 188 | nocs_coord, 189 | global_pts_pred, 190 | K_img, 191 | None 192 | ) 193 | else: 194 | retval, rvecs, tvecs = cv2.solvePnP( 195 | nocs_coord, 196 | global_pts_pred, 197 | K_img, 198 | None, 199 | flags=cv2.SOLVEPNP_ITERATIVE 200 | ) 201 | 202 | rotM = cv2.Rodrigues(rvecs)[0].T 203 | tvecs = tvecs.squeeze(axis=1) 204 | 205 | # GL 206 | R_temp = np.identity(4) 207 | R_temp[:3, :3] = rotM 208 | R_temp[3, :3] = tvecs 209 | R_t_pred = R_temp @ T 210 | 211 | 212 | # pnp 213 | if R_t_pred[3, 2] > 0: 214 | print('【pnp_fail tz_pred > 0】', self.img_path_list[i]) 215 | self.inference_data['pnp_fail'] += 1 216 | R_t_pred = np.identity(4) 217 | 218 | elif R_t_pred[3, 2] < -100: 219 | print('【pnp_fail tz_pred < -100】', self.img_path_list[i]) 220 | self.inference_data['pnp_fail'] += 1 221 | R_t_pred = np.identity(4) 222 | 223 | 224 | R_t_gt = self.R_t_list[i] 225 | img_path = self.img_path_list[i] 226 | 227 | yaw_gt, pitch_gt, roll_gt = Rotation.from_matrix(R_t_gt[:3, :3].T).as_euler('yxz', degrees=True) 228 | yaw_pred, pitch_pred, roll_pred = Rotation.from_matrix(R_t_pred[:3, :3].T).as_euler('yxz', degrees=True) 229 | 230 | tx_gt, ty_gt, tz_gt = R_t_gt[3, :3] 231 | tx_pred, ty_pred, tz_pred = R_t_pred[3, :3] 232 | 233 | 234 | cur_pitch_mae = np.abs(pitch_gt - pitch_pred) 235 | cur_yaw_mae = np.abs(yaw_gt - yaw_pred) 236 | cur_roll_mae = np.abs(roll_gt - roll_pred) 237 | 238 | cur_tx_mae = np.abs(tx_gt - tx_pred) 239 | cur_ty_mae = np.abs(ty_gt - ty_pred) 240 | cur_tz_mae = np.abs(tz_gt - tz_pred) 241 | 242 | self.inference_data['pitch_mae'].append(cur_pitch_mae) 243 | self.inference_data['yaw_mae'].append(cur_yaw_mae) 244 | self.inference_data['roll_mae'].append(cur_roll_mae) 245 | self.inference_data['tx_mae'].append(cur_tx_mae) 246 | self.inference_data['ty_mae'].append(cur_ty_mae) 247 | self.inference_data['tz_mae'].append(cur_tz_mae) 248 | 249 | # 3DRecon 250 | cur_3drecon = np.sqrt(((verts3d_pred[i] - verts3d_gt[i]) ** 2).sum(axis=1)).mean() 251 | self.inference_data['3DRecon'].append(cur_3drecon) 252 | 253 | # ADD 254 | verts_origin = verts3d_gt[i] 255 | ones = np.ones([verts_origin.shape[0], 1]) 256 | verts_homo = np.concatenate([verts_origin, ones], axis=1) 257 | 258 | verts_cam_pred = verts_homo @ R_t_pred 259 | verts_cam_gt = verts_homo @ R_t_gt 260 | 261 | cur_ADD = np.sqrt(np.sum((verts_cam_pred[:, :3] - verts_cam_gt[:, :3]) ** 2, 1)).mean() 262 | self.inference_data['ADD'].append(cur_ADD) 263 | 264 | # IOU, 5°5cm, 10°5cm 265 | R_error, T_error, IoU = compute_sRT_errors(R_t_pred.T, R_t_gt.T) 266 | if R_error < 5 and T_error < 0.05: 267 | self.inference_data['strict_success'] += 1 268 | if R_error < 10 and T_error < 0.05: 269 | self.inference_data['easy_success'] += 1 270 | 271 | self.inference_data['IoU'].append(IoU) 272 | self.inference_data['total_count'] += 1 273 | 274 | 275 | 276 | 277 | -------------------------------------------------------------------------------- /models/resnet.py: -------------------------------------------------------------------------------- 1 | #from lib.utils.config import cfg as proj_cfg 2 | import torch.nn as nn 3 | import math 4 | import torch.utils.model_zoo as model_zoo 5 | import numpy as np 6 | import torch 7 | 8 | 9 | __all__ = ['ResNet', 'resnet18', 'resnet34', 'resnet50', 'resnet101', 10 | 'resnet152'] 11 | 12 | 13 | model_urls = { 14 | 'resnet18': 'https://download.pytorch.org/models/resnet18-5c106cde.pth', 15 | 'resnet34': 'https://download.pytorch.org/models/resnet34-333f7ec4.pth', 16 | 'resnet50': 'https://download.pytorch.org/models/resnet50-19c8e357.pth', 17 | 'resnet101': 'https://download.pytorch.org/models/resnet101-5d3b4d8f.pth', 18 | 'resnet152': 'https://download.pytorch.org/models/resnet152-b121ed2d.pth', 19 | } 20 | 21 | 22 | def conv3x3(in_planes, out_planes, stride=1, dilation=1): 23 | "3x3 convolution with padding" 24 | 25 | kernel_size = np.asarray((3, 3)) 26 | 27 | # Compute the size of the upsampled filter with 28 | # a specified dilation rate. 29 | upsampled_kernel_size = (kernel_size - 1) * (dilation - 1) + kernel_size 30 | 31 | # Determine the padding that is necessary for full padding, 32 | # meaning the output spatial size is equal to input spatial size 33 | full_padding = (upsampled_kernel_size - 1) // 2 34 | 35 | # Conv2d doesn't accept numpy arrays as arguments 36 | full_padding, kernel_size = tuple(full_padding), tuple(kernel_size) 37 | 38 | return nn.Conv2d(in_planes, out_planes, kernel_size=kernel_size, stride=stride, 39 | padding=full_padding, dilation=dilation, bias=False) 40 | 41 | 42 | class BasicBlock(nn.Module): 43 | expansion = 1 44 | 45 | def __init__(self, inplanes, planes, stride=1, downsample=None, dilation=1): 46 | super(BasicBlock, self).__init__() 47 | self.conv1 = conv3x3(inplanes, planes, stride, dilation=dilation) 48 | self.bn1 = nn.BatchNorm2d(planes) 49 | self.relu = nn.ReLU(inplace=True) 50 | self.conv2 = conv3x3(planes, planes, dilation=dilation) 51 | self.bn2 = nn.BatchNorm2d(planes) 52 | self.downsample = downsample 53 | self.stride = stride 54 | 55 | def forward(self, x): 56 | residual = x 57 | 58 | out = self.conv1(x) 59 | out = self.bn1(out) 60 | out = self.relu(out) 61 | 62 | out = self.conv2(out) 63 | out = self.bn2(out) 64 | 65 | if self.downsample is not None: 66 | residual = self.downsample(x) 67 | 68 | out += residual 69 | out = self.relu(out) 70 | 71 | return out 72 | 73 | 74 | class Bottleneck(nn.Module): 75 | expansion = 4 76 | 77 | def __init__(self, inplanes, planes, stride=1, downsample=None, dilation=1): 78 | super(Bottleneck, self).__init__() 79 | self.conv1 = nn.Conv2d(inplanes, planes, kernel_size=1, bias=False) 80 | self.bn1 = nn.BatchNorm2d(planes) 81 | 82 | self.conv2 = conv3x3(planes, planes, stride=stride, dilation=dilation) 83 | 84 | #self.conv2 = nn.Conv2d(planes, planes, kernel_size=3, stride=stride, 85 | # padding=1, bias=False) 86 | 87 | self.bn2 = nn.BatchNorm2d(planes) 88 | self.conv3 = nn.Conv2d(planes, planes * 4, kernel_size=1, bias=False) 89 | self.bn3 = nn.BatchNorm2d(planes * 4) 90 | self.relu = nn.ReLU(inplace=True) 91 | self.downsample = downsample 92 | self.stride = stride 93 | 94 | def forward(self, x): 95 | residual = x 96 | 97 | out = self.conv1(x) 98 | out = self.bn1(out) 99 | out = self.relu(out) 100 | 101 | out = self.conv2(out) 102 | out = self.bn2(out) 103 | out = self.relu(out) 104 | 105 | out = self.conv3(out) 106 | out = self.bn3(out) 107 | 108 | if self.downsample is not None: 109 | residual = self.downsample(x) 110 | 111 | out += residual 112 | out = self.relu(out) 113 | 114 | return out 115 | 116 | 117 | class ResNet(nn.Module): 118 | 119 | def __init__(self, 120 | block, 121 | layers, 122 | num_classes=1000, 123 | fully_conv=False, 124 | remove_avg_pool_layer=False, 125 | output_stride=32): 126 | 127 | # Add additional variables to track 128 | # output stride. Necessary to achieve 129 | # specified output stride. 130 | self.output_stride = output_stride 131 | self.current_stride = 4 132 | self.current_dilation = 1 133 | 134 | self.remove_avg_pool_layer = remove_avg_pool_layer 135 | 136 | self.inplanes = 64 137 | self.fully_conv = fully_conv 138 | super(ResNet, self).__init__() 139 | self.conv1 = nn.Conv2d(3, 64, kernel_size=7, stride=2, padding=3, 140 | bias=False) 141 | self.bn1 = nn.BatchNorm2d(64) 142 | self.relu = nn.ReLU(inplace=True) 143 | self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1) 144 | 145 | self.layer1 = self._make_layer(block, 64, layers[0]) 146 | self.layer2 = self._make_layer(block, 128, layers[1], stride=2) 147 | self.layer3 = self._make_layer(block, 256, layers[2], stride=2) 148 | self.layer4 = self._make_layer(block, 512, layers[3], stride=2) 149 | 150 | self.avgpool = nn.AvgPool2d(7) 151 | 152 | self.fc = nn.Linear(512 * block.expansion, num_classes) 153 | 154 | if self.fully_conv: 155 | self.avgpool = nn.AvgPool2d(7, padding=3, stride=1) 156 | # In the latest unstable torch 4.0 the tensor.copy_ 157 | # method was changed and doesn't work as it used to be 158 | #self.fc = nn.Conv2d(512 * block.expansion, num_classes, 1) 159 | 160 | for m in self.modules(): 161 | if isinstance(m, nn.Conv2d): 162 | n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels 163 | m.weight.data.normal_(0, math.sqrt(2. / n)) 164 | elif isinstance(m, nn.BatchNorm2d): 165 | m.weight.data.fill_(1) 166 | m.bias.data.zero_() 167 | 168 | def _make_layer(self, block, planes, blocks, stride=1, dilation=1): 169 | downsample = None 170 | 171 | if stride != 1 or self.inplanes != planes * block.expansion: 172 | 173 | # Check if we already achieved desired output stride. 174 | if self.current_stride == self.output_stride: 175 | 176 | # If so, replace subsampling with a dilation to preserve 177 | # current spatial resolution. 178 | self.current_dilation = self.current_dilation * stride 179 | stride = 1 180 | else: 181 | 182 | # If not, perform subsampling and update current 183 | # new output stride. 184 | self.current_stride = self.current_stride * stride 185 | 186 | # We don't dilate 1x1 convolution. 187 | downsample = nn.Sequential( 188 | nn.Conv2d(self.inplanes, planes * block.expansion, 189 | kernel_size=1, stride=stride, bias=False), 190 | nn.BatchNorm2d(planes * block.expansion), 191 | ) 192 | 193 | layers = [] 194 | layers.append(block(self.inplanes, planes, stride, downsample, dilation=self.current_dilation)) 195 | self.inplanes = planes * block.expansion 196 | for i in range(1, blocks): 197 | layers.append(block(self.inplanes, planes, dilation=self.current_dilation)) 198 | 199 | return nn.Sequential(*layers) 200 | 201 | def forward(self, x): 202 | x = self.conv1(x) 203 | x = self.bn1(x) 204 | x2s = self.relu(x) 205 | x = self.maxpool(x2s) 206 | 207 | x4s = self.layer1(x) 208 | x8s = self.layer2(x4s) 209 | x16s = self.layer3(x8s) 210 | x32s = self.layer4(x16s) 211 | x = x32s 212 | 213 | 214 | if not self.remove_avg_pool_layer: 215 | x = self.avgpool(x) 216 | 217 | if not self.fully_conv: 218 | x = x.view(x.size(0), -1) 219 | 220 | xfc = self.fc(x) 221 | 222 | return x2s,x4s,x8s,x16s,x32s,xfc 223 | 224 | 225 | def resnet18(pretrained=False, **kwargs): 226 | """Constructs a ResNet-18 model. 227 | 228 | Args: 229 | pretrained (bool): If True, returns a model pre-trained on ImageNet 230 | """ 231 | model = ResNet(BasicBlock, [2, 2, 2, 2], **kwargs) 232 | if pretrained: 233 | model.load_state_dict(model_zoo.load_url(model_urls['resnet18'], model_dir='./models')) 234 | #model.load_state_dict(torch.load('./models/resnet18-5c106cde.pth')['model_state_dict']) 235 | return model 236 | 237 | 238 | def resnet34(pretrained=False, **kwargs): 239 | """Constructs a ResNet-34 model. 240 | 241 | Args: 242 | pretrained (bool): If True, returns a model pre-trained on ImageNet 243 | """ 244 | model = ResNet(BasicBlock, [3, 4, 6, 3], **kwargs) 245 | if pretrained: 246 | model.load_state_dict(model_zoo.load_url(model_urls['resnet34'],model_dir=proj_cfg.MODEL_DIR)) 247 | return model 248 | 249 | 250 | def resnet50(pretrained=False, **kwargs): 251 | """Constructs a ResNet-50 model. 252 | 253 | Args: 254 | pretrained (bool): If True, returns a model pre-trained on ImageNet 255 | """ 256 | model = ResNet(Bottleneck, [3, 4, 6, 3], **kwargs) 257 | if pretrained: 258 | model.load_state_dict(model_zoo.load_url(model_urls['resnet50'],model_dir=proj_cfg.MODEL_DIR)) 259 | return model 260 | 261 | 262 | def resnet101(pretrained=False, **kwargs): 263 | """Constructs a ResNet-101 model. 264 | 265 | Args: 266 | pretrained (bool): If True, returns a model pre-trained on ImageNet 267 | """ 268 | model = ResNet(Bottleneck, [3, 4, 23, 3], **kwargs) 269 | if pretrained: 270 | model.load_state_dict(model_zoo.load_url(model_urls['resnet101'],model_dir=proj_cfg.MODEL_DIR)) 271 | return model 272 | 273 | 274 | def resnet152(pretrained=False, **kwargs): 275 | """Constructs a ResNet-152 model. 276 | 277 | Args: 278 | pretrained (bool): If True, returns a model pre-trained on ImageNet 279 | """ 280 | model = ResNet(Bottleneck, [3, 8, 36, 3], **kwargs) 281 | if pretrained: 282 | model.load_state_dict(model_zoo.load_url(model_urls['resnet152'],model_dir=proj_cfg.MODEL_DIR)) 283 | return model 284 | -------------------------------------------------------------------------------- /npy/kpt_ind.npy: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/cbsropenproject/6dof_face/c80ca80125b2059e8f2984688814d59d0ac8e10c/npy/kpt_ind.npy -------------------------------------------------------------------------------- /npy/projection_matrix.txt: -------------------------------------------------------------------------------- 1 | 1.57443701 0.0 0.0 0.0 2 | 0.0 1.57443701 0.0 0.0 3 | 0.0 0.0 -0.99999976 -1.0 4 | 0.0 0.0 -0.001 0.0 5 | -------------------------------------------------------------------------------- /npy/tris_2388x3.npy: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/cbsropenproject/6dof_face/c80ca80125b2059e8f2984688814d59d0ac8e10c/npy/tris_2388x3.npy -------------------------------------------------------------------------------- /npy/tris_2500x3_202110.npy: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/cbsropenproject/6dof_face/c80ca80125b2059e8f2984688814d59d0ac8e10c/npy/tris_2500x3_202110.npy -------------------------------------------------------------------------------- /npy/uv_coords_std_202109.npy: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/cbsropenproject/6dof_face/c80ca80125b2059e8f2984688814d59d0ac8e10c/npy/uv_coords_std_202109.npy -------------------------------------------------------------------------------- /options/__init__.py: -------------------------------------------------------------------------------- 1 | """This package options includes option modules: training options, test options, and basic options (used in both training and test).""" 2 | -------------------------------------------------------------------------------- /options/base_options.py: -------------------------------------------------------------------------------- 1 | import argparse 2 | import os 3 | from util import util 4 | import torch 5 | import models 6 | import data 7 | 8 | 9 | class BaseOptions(): 10 | """This class defines options used during both training and test time. 11 | 12 | It also implements several helper functions such as parsing, printing, and saving the options. 13 | It also gathers additional options defined in functions in both dataset class and model class. 14 | """ 15 | 16 | def __init__(self, cmd_line=None): 17 | """Reset the class; indicates the class hasn't been initailized""" 18 | self.initialized = False 19 | self.cmd_line = None 20 | if cmd_line is not None: 21 | self.cmd_line = cmd_line.split() 22 | 23 | def initialize(self, parser): 24 | """Define the common options that are used in both training and test.""" 25 | 26 | # 新增opt参数 27 | parser.add_argument('--csv_path_train', type=str, default='cache/train_list.csv') 28 | parser.add_argument('--csv_path_test', type=str, default='cache/test_list.csv') 29 | parser.add_argument('--img_size', type=int, default=192) 30 | parser.add_argument('--debug', action='store_true', default=False) 31 | parser.add_argument('--test_visualization', action='store_true', default=False) 32 | parser.add_argument('--memo', type=str, default='') 33 | parser.add_argument('--use_gt_bbox', action='store_true', default=False) 34 | 35 | 36 | 37 | # basic parameters 38 | #parser.add_argument('--dataroot', default='placeholder', help='path to images (should have subfolders trainA, trainB, valA, valB, etc)') 39 | parser.add_argument('--name', type=str, default='run1', help='name of the experiment. It decides where to store samples and models') 40 | parser.add_argument('--easy_label', type=str, default='experiment_name', help='Interpretable name') 41 | parser.add_argument('--gpu_ids', type=str, default='0', help='gpu ids: e.g. 0 0,1,2, 0,2. use -1 for CPU') 42 | parser.add_argument('--checkpoints_dir', type=str, default='./checkpoint', help='models are saved here') 43 | # model parameters 44 | parser.add_argument('--model', type=str, default='perspnet', help='chooses which model to use.') 45 | parser.add_argument('--init_type', type=str, default='xavier', choices=['normal', 'xavier', 'kaiming', 'orthogonal'], help='network initialization') 46 | parser.add_argument('--init_gain', type=float, default=0.02, help='scaling factor for normal, xavier and orthogonal.') 47 | # dataset parameters 48 | parser.add_argument('--dataset_mode', type=str, default='arkit', help='chooses how datasets are loaded. [unaligned | aligned | single | colorization]') 49 | #parser.add_argument('--direction', type=str, default='AtoB', help='AtoB or BtoA') 50 | parser.add_argument('--serial_batches', action='store_true', help='if true, takes images in order to make batches, otherwise takes them randomly') 51 | parser.add_argument('--num_threads', default=6, type=int, help='# threads for loading data') 52 | parser.add_argument('--batch_size', type=int, default=64, help='input batch size') 53 | parser.add_argument('--n_pts', type=int, default=1024, help='number of foreground points') 54 | #parser.add_argument('--crop_size', type=int, default=256, help='then crop to this size') 55 | parser.add_argument('--max_dataset_size', type=int, default=float("inf"), help='Maximum number of samples allowed per dataset. If the dataset directory contains more than max_dataset_size, only a subset is loaded.') 56 | #parser.add_argument('--preprocess', type=str, default='resize_and_crop', help='scaling and cropping of images at load time [resize_and_crop | crop | scale_width | scale_width_and_crop | none]') 57 | #parser.add_argument('--no_flip', action='store_true', help='if specified, do not flip the images for data augmentation') 58 | parser.add_argument('--display_winsize', type=int, default=256, help='display window size for both visdom and HTML') 59 | #parser.add_argument('--random_scale_max', type=float, default=3.0, 60 | # help='(used for single image translation) Randomly scale the image by the specified factor as data augmentation.') 61 | # additional parameters 62 | parser.add_argument('--epoch', type=str, default='latest', help='which epoch to load? set to latest to use latest cached model') 63 | parser.add_argument('--verbose', action='store_true', help='if specified, print more debugging information') 64 | parser.add_argument('--suffix', default='', type=str, help='customized suffix: opt.name = opt.name + suffix: e.g., {model}_{netG}_size{load_size}') 65 | 66 | 67 | self.initialized = True 68 | return parser 69 | 70 | def gather_options(self): 71 | """Initialize our parser with basic options(only once). 72 | Add additional model-specific and dataset-specific options. 73 | These options are defined in the function 74 | in model and dataset classes. 75 | """ 76 | if not self.initialized: # check if it has been initialized 77 | parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter) 78 | parser = self.initialize(parser) 79 | 80 | # get the basic options 81 | if self.cmd_line is None: 82 | opt, _ = parser.parse_known_args() 83 | else: 84 | opt, _ = parser.parse_known_args(self.cmd_line) 85 | 86 | # modify model-related parser options 87 | model_name = opt.model 88 | model_option_setter = models.get_option_setter(model_name) 89 | parser = model_option_setter(parser, self.isTrain) 90 | if self.cmd_line is None: 91 | opt, _ = parser.parse_known_args() # parse again with new defaults 92 | else: 93 | opt, _ = parser.parse_known_args(self.cmd_line) # parse again with new defaults 94 | 95 | # modify dataset-related parser options 96 | dataset_name = opt.dataset_mode 97 | dataset_option_setter = data.get_option_setter(dataset_name) 98 | parser = dataset_option_setter(parser, self.isTrain) 99 | 100 | # save and return the parser 101 | self.parser = parser 102 | if self.cmd_line is None: 103 | return parser.parse_args() 104 | else: 105 | return parser.parse_args(self.cmd_line) 106 | 107 | def print_options(self, opt): 108 | """Print and save options 109 | 110 | It will print both current options and default values(if different). 111 | It will save options into a text file / [checkpoints_dir] / opt.txt 112 | """ 113 | message = '' 114 | message += '----------------- Options ---------------\n' 115 | for k, v in sorted(vars(opt).items()): 116 | comment = '' 117 | default = self.parser.get_default(k) 118 | if v != default: 119 | comment = '\t[default: %s]' % str(default) 120 | message += '{:>25}: {:<30}{}\n'.format(str(k), str(v), comment) 121 | message += '----------------- End -------------------' 122 | print(message) 123 | 124 | # save to the disk 125 | expr_dir = os.path.join(opt.checkpoints_dir, opt.name) 126 | util.mkdirs(expr_dir) 127 | file_name = os.path.join(expr_dir, '{}_opt.txt'.format(opt.phase)) 128 | try: 129 | with open(file_name, 'wt') as opt_file: 130 | opt_file.write(message) 131 | opt_file.write('\n') 132 | except PermissionError as error: 133 | print("permission error {}".format(error)) 134 | pass 135 | 136 | def parse(self): 137 | """Parse our options, create checkpoints directory suffix, and set up gpu device.""" 138 | opt = self.gather_options() 139 | opt.isTrain = self.isTrain # train or test 140 | 141 | # process opt.suffix 142 | if opt.suffix: 143 | suffix = ('_' + opt.suffix.format(**vars(opt))) if opt.suffix != '' else '' 144 | opt.name = opt.name + suffix 145 | 146 | self.print_options(opt) 147 | 148 | # set gpu ids 149 | str_ids = opt.gpu_ids.split(',') 150 | opt.gpu_ids = [] 151 | for str_id in str_ids: 152 | id = int(str_id) 153 | if id >= 0: 154 | opt.gpu_ids.append(id) 155 | if len(opt.gpu_ids) > 0: 156 | torch.cuda.set_device(opt.gpu_ids[0]) 157 | 158 | 159 | #【debug】 160 | if opt.debug: 161 | opt.print_freq = opt.batch_size 162 | opt.display_freq = opt.batch_size 163 | opt.update_html_freq = opt.batch_size 164 | 165 | self.opt = opt 166 | return self.opt 167 | -------------------------------------------------------------------------------- /options/test_options.py: -------------------------------------------------------------------------------- 1 | from .base_options import BaseOptions 2 | 3 | 4 | class TestOptions(BaseOptions): 5 | """This class includes test options. 6 | 7 | It also includes shared options defined in BaseOptions. 8 | """ 9 | 10 | def initialize(self, parser): 11 | parser = BaseOptions.initialize(self, parser) # define shared options 12 | parser.add_argument('--results_dir', type=str, default='./results/', help='saves results here.') 13 | parser.add_argument('--phase', type=str, default='test', help='train, val, test, etc') 14 | # Dropout and Batchnorm has different behavioir during training and test. 15 | parser.add_argument('--eval', action='store_true', default=True, help='use eval mode during test time.') 16 | parser.add_argument('--seg', action='store_true', help='use seg_pred during test time.') 17 | #parser.add_argument('--num_test', type=int, default=50, help='how many test images to run') 18 | 19 | # To avoid cropping, the load_size should be the same as crop_size 20 | parser.set_defaults(load_size=parser.get_default('crop_size')) 21 | self.isTrain = False 22 | return parser 23 | -------------------------------------------------------------------------------- /options/train_options.py: -------------------------------------------------------------------------------- 1 | from .base_options import BaseOptions 2 | 3 | 4 | class TrainOptions(BaseOptions): 5 | """This class includes training options. 6 | 7 | It also includes shared options defined in BaseOptions. 8 | """ 9 | 10 | def initialize(self, parser): 11 | parser = BaseOptions.initialize(self, parser) 12 | # visdom and HTML visualization parameters 13 | parser.add_argument('--display_freq', type=int, default=8192, help='frequency of showing training results on screen') 14 | parser.add_argument('--display_ncols', type=int, default=4, help='if positive, display all images in a single visdom web panel with certain number of images per row.') 15 | parser.add_argument('--display_id', type=int, default=-1, help='window id of the web display. Default is random window id') 16 | parser.add_argument('--display_server', type=str, default="http://localhost", help='visdom server of the web display') 17 | parser.add_argument('--display_env', type=str, default='main', help='visdom display environment name (default is "main")') 18 | parser.add_argument('--display_port', type=int, default=8097, help='visdom port of the web display') 19 | parser.add_argument('--update_html_freq', type=int, default=8192, help='frequency of saving training results to html') 20 | parser.add_argument('--print_freq', type=int, default=8192, help='frequency of showing training results on console') 21 | parser.add_argument('--no_html', action='store_true', default=False, help='do not save intermediate training results to [opt.checkpoints_dir]/[opt.name]/web/') 22 | # network saving and loading parameters 23 | parser.add_argument('--save_latest_freq', type=int, default=5000, help='frequency of saving the latest results') 24 | parser.add_argument('--save_epoch_freq', type=int, default=1, help='frequency of saving checkpoints at the end of epochs') 25 | #parser.add_argument('--evaluation_freq', type=int, default=5000, help='evaluation freq') 26 | parser.add_argument('--save_by_iter', action='store_true', help='whether saves model by iteration') 27 | parser.add_argument('--continue_train', action='store_true', help='continue training: load the latest model') 28 | parser.add_argument('--seg', action='store_true', help='continue training: load the latest model') 29 | parser.add_argument('--epoch_count', type=int, default=1, help='the starting epoch count, we save the model by , +, ...') 30 | parser.add_argument('--phase', type=str, default='train', help='train, val, test, etc') 31 | parser.add_argument('--pretrained_name', type=str, default=None, help='resume training from another checkpoint') 32 | 33 | # training parameters 34 | parser.add_argument('--n_epochs', type=int, default=10, help='number of epochs with the initial learning rate') 35 | parser.add_argument('--n_epochs_decay', type=int, default=10, help='number of epochs to linearly decay learning rate to zero') 36 | parser.add_argument('--beta1', type=float, default=0.9, help='momentum term of adam') 37 | parser.add_argument('--beta2', type=float, default=0.999, help='momentum term of adam') 38 | parser.add_argument('--lr', type=float, default=0.0001, help='initial learning rate for adam') 39 | #parser.add_argument('--gan_mode', type=str, default='lsgan', help='the type of GAN objective. [vanilla| lsgan | wgangp]. vanilla GAN loss is the cross-entropy objective used in the original GAN paper.') 40 | #parser.add_argument('--pool_size', type=int, default=50, help='the size of image buffer that stores previously generated images') 41 | parser.add_argument('--lr_policy', type=str, default='linear', help='learning rate policy. [linear | step | plateau | cosine]') 42 | parser.add_argument('--lr_decay_iters', type=int, default=50, help='multiply by a gamma every lr_decay_iters iterations') 43 | 44 | self.isTrain = True 45 | return parser 46 | -------------------------------------------------------------------------------- /requirement.txt: -------------------------------------------------------------------------------- 1 | opencv-python==4.5.2.54 2 | numpy==1.21.2 3 | torch 4 | torchvision 5 | Cython==0.29.22 6 | matplotlib==3.5.1 7 | pandas==1.3.3 8 | Pillow==8.4.0 9 | scikit-image==0.16.2 10 | tqdm==4.64.1 11 | -------------------------------------------------------------------------------- /test.py: -------------------------------------------------------------------------------- 1 | import os 2 | import sys 3 | import cv2 4 | import time 5 | import pprint 6 | import numpy as np 7 | from datetime import datetime 8 | from options.test_options import TestOptions 9 | from models import create_model 10 | from data import create_dataset 11 | 12 | ###usage: python -u test.py --image_size=192 --model perspnet --dataset_mode arkit --csv_path_test 'test.csv' 13 | if __name__ == '__main__': 14 | print('【process_id】', os.getpid()) 15 | print('【command】python -u ' + ' '.join(sys.argv) + '\n') 16 | 17 | opt = TestOptions().parse() # get test options 18 | opt.serial_batches = True 19 | opt.num_thread = 4 20 | dataset = create_dataset(opt) 21 | num_batch = len(dataset.dataloader) 22 | print("The number of testing images = %d" % len(dataset)) 23 | 24 | model = create_model(opt) 25 | t_start = time.time() 26 | for i, data in enumerate(dataset): 27 | print('[%s][progresss]%d %d'%(datetime.now().isoformat(sep=' '), i, num_batch)) 28 | if i==0: 29 | model.data_dependent_initialize(data) 30 | model.setup(opt) 31 | model.parallelize() 32 | if model.eval: 33 | model.eval() 34 | model.init_evaluation() 35 | model.set_input(data) 36 | model.inference_curr_batch() 37 | if i==-1: 38 | break 39 | metric = dataset.dataset.compute_metrics(model.inference_data) 40 | print('\n [metric] ') 41 | pprint.pprint(metric) 42 | print('Done in %.3f sec' %(time.time()-t_start)) 43 | 44 | 45 | 46 | 47 | 48 | -------------------------------------------------------------------------------- /train.py: -------------------------------------------------------------------------------- 1 | import matplotlib 2 | matplotlib.use('Agg') 3 | import matplotlib.pyplot as plt 4 | 5 | import os 6 | import sys 7 | import pdb 8 | import cv2 9 | import time 10 | import copy 11 | import numpy as np 12 | import torch 13 | from datetime import datetime 14 | from options.train_options import TrainOptions 15 | from data import create_dataset 16 | from models import create_model 17 | from util.visualizer import Visualizer 18 | 19 | 20 | if __name__ == '__main__': 21 | print('【process_id】', os.getpid()) 22 | print('【command】python -u ' + ' '.join(sys.argv) + '\n') 23 | 24 | train_opt = TrainOptions().parse() # get training options 25 | train_dataset = create_dataset(train_opt) # create a dataset given opt.dataset_mode and other options 26 | train_dataset_size = len(train_dataset) # get the number of images in the dataset. 27 | 28 | 29 | val_opt = copy.deepcopy(train_opt) 30 | val_opt.isTrain = False 31 | val_opt.batch_size = 64 32 | val_opt.num_thread = 6 33 | val_opt.serial_batches = True 34 | val_dataset = create_dataset(val_opt) 35 | 36 | print('The number of training images = %d' % train_dataset_size) 37 | print('The number of val images = %d\n' % len(val_dataset)) 38 | 39 | 40 | model = create_model(train_opt) # create a model given opt.model and other options 41 | visualizer = Visualizer(train_opt) # create a visualizer that display/save images and plots 42 | train_opt.visualizer = visualizer 43 | total_iters = 0 # the total number of training iterations 44 | 45 | optimize_time = 0.1 46 | 47 | times = [] 48 | for epoch in range(train_opt.epoch_count, train_opt.n_epochs + train_opt.n_epochs_decay + 1): # outer loop for different epochs; we save the model by , + 49 | epoch_start_time = time.time() # timer for entire epoch 50 | iter_data_time = time.time() # timer for data loading per iteration 51 | epoch_iter = 0 # the number of training iterations in current epoch, reset to 0 every epoch 52 | visualizer.reset() # reset the visualizer: make sure it saves the results to HTML at least once every epoch 53 | 54 | 55 | #train_dataset.set_epoch(epoch) 56 | for i, data in enumerate(train_dataset): # inner loop within one epoch 57 | iter_start_time = time.time() # timer for computation per iteration 58 | if total_iters % train_opt.print_freq == 0: 59 | t_data = iter_start_time - iter_data_time 60 | 61 | batch_size = data['img'].size(0) 62 | total_iters += batch_size 63 | epoch_iter += batch_size 64 | if len(train_opt.gpu_ids) > 0: 65 | torch.cuda.synchronize() 66 | 67 | optimize_start_time = time.time() 68 | if epoch ==2 and i==0: 69 | train_opt.epoch_count=2 70 | train_opt.lr=0.0001 71 | train_opt.continue_train=True 72 | model = create_model(train_opt) # create a model given opt.model and other options 73 | 74 | if epoch == train_opt.epoch_count and i == 0: 75 | model.data_dependent_initialize(data) 76 | model.setup(train_opt) # regular setup: load and print networks; create schedulers 77 | model.parallelize() 78 | 79 | 80 | model.set_input(data) # unpack data from dataset and apply preprocessing 81 | model.optimize_parameters() # calculate loss functions, get gradients, update network weights 82 | 83 | 84 | if len(train_opt.gpu_ids) > 0: 85 | torch.cuda.synchronize() 86 | optimize_time = (time.time() - optimize_start_time) / batch_size * 0.005 + 0.995 * optimize_time 87 | 88 | if total_iters % train_opt.display_freq == 0: # display images on visdom and save images to a HTML file 89 | save_result = total_iters % train_opt.update_html_freq == 0 90 | model.compute_visuals() 91 | visualizer.display_current_results(model.get_current_visuals(), epoch, save_result) 92 | 93 | if total_iters % train_opt.print_freq == 0: # print training losses and save logging information to the disk 94 | losses = model.get_current_losses() 95 | visualizer.print_current_losses(epoch, epoch_iter, losses, optimize_time, t_data) 96 | if train_opt.display_id is None or train_opt.display_id > 0: 97 | visualizer.plot_current_losses(epoch, float(epoch_iter) / dataset_size, losses) 98 | 99 | if total_iters % train_opt.save_latest_freq == 0: # cache our latest model every iterations 100 | print('saving the latest model (epoch %d, total_iters %d)' % (epoch, total_iters)) 101 | print(train_opt.name) # it's useful to occasionally show the experiment name on console 102 | save_suffix = 'iter_%d' % total_iters if train_opt.save_by_iter else 'latest' 103 | model.save_networks(save_suffix) 104 | 105 | iter_data_time = time.time() 106 | 107 | if epoch % train_opt.save_epoch_freq == 0: # cache our model every epochs 108 | print('saving the model at the end of epoch %d, iters %d' % (epoch, total_iters)) 109 | model.save_networks('latest') 110 | model.save_networks(epoch) 111 | 112 | print('End of epoch %d / %d \t Time Taken: %d sec' % (epoch, train_opt.n_epochs + train_opt.n_epochs_decay, time.time() - epoch_start_time)) 113 | model.update_learning_rate() # update learning rates at the end of every epoch. 114 | 115 | 116 | if epoch % 1 == 0: 117 | print('===================== evaluation epoch %d begin =====================' % epoch) 118 | t_start = time.time() 119 | model.init_evaluation() 120 | num_batch = len(val_dataset.dataloader) 121 | 122 | for j, data in enumerate(val_dataset): 123 | if (j + 1) % 50 == 0: 124 | print('[%s]【progress】%d/%d' % (datetime.now().isoformat(sep=' '), j + 1, num_batch)) 125 | model.set_input(data) 126 | model.inference_curr_batch() 127 | 128 | metrics = val_dataset.dataset.compute_metrics(model.inference_data) 129 | print(metrics) 130 | print('Done in %.2f sec' % (time.time() - t_start)) 131 | print('===================== evaluation epoch %d end =====================\n' % epoch) 132 | 133 | 134 | 135 | 136 | 137 | -------------------------------------------------------------------------------- /util/__init__.py: -------------------------------------------------------------------------------- 1 | """This package includes a miscellaneous collection of useful helper functions.""" 2 | from util import * 3 | -------------------------------------------------------------------------------- /util/get_data.py: -------------------------------------------------------------------------------- 1 | from __future__ import print_function 2 | import os 3 | import tarfile 4 | import requests 5 | from warnings import warn 6 | from zipfile import ZipFile 7 | from bs4 import BeautifulSoup 8 | from os.path import abspath, isdir, join, basename 9 | 10 | 11 | class GetData(object): 12 | """A Python script for downloading CycleGAN or pix2pix datasets. 13 | 14 | Parameters: 15 | technique (str) -- One of: 'cyclegan' or 'pix2pix'. 16 | verbose (bool) -- If True, print additional information. 17 | 18 | Examples: 19 | >>> from util.get_data import GetData 20 | >>> gd = GetData(technique='cyclegan') 21 | >>> new_data_path = gd.get(save_path='./datasets') # options will be displayed. 22 | 23 | Alternatively, You can use bash scripts: 'scripts/download_pix2pix_model.sh' 24 | and 'scripts/download_cyclegan_model.sh'. 25 | """ 26 | 27 | def __init__(self, technique='cyclegan', verbose=True): 28 | url_dict = { 29 | 'pix2pix': 'http://efrosgans.eecs.berkeley.edu/pix2pix/datasets/', 30 | 'cyclegan': 'https://people.eecs.berkeley.edu/~taesung_park/CycleGAN/datasets' 31 | } 32 | self.url = url_dict.get(technique.lower()) 33 | self._verbose = verbose 34 | 35 | def _print(self, text): 36 | if self._verbose: 37 | print(text) 38 | 39 | @staticmethod 40 | def _get_options(r): 41 | soup = BeautifulSoup(r.text, 'lxml') 42 | options = [h.text for h in soup.find_all('a', href=True) 43 | if h.text.endswith(('.zip', 'tar.gz'))] 44 | return options 45 | 46 | def _present_options(self): 47 | r = requests.get(self.url) 48 | options = self._get_options(r) 49 | print('Options:\n') 50 | for i, o in enumerate(options): 51 | print("{0}: {1}".format(i, o)) 52 | choice = input("\nPlease enter the number of the " 53 | "dataset above you wish to download:") 54 | return options[int(choice)] 55 | 56 | def _download_data(self, dataset_url, save_path): 57 | if not isdir(save_path): 58 | os.makedirs(save_path) 59 | 60 | base = basename(dataset_url) 61 | temp_save_path = join(save_path, base) 62 | 63 | with open(temp_save_path, "wb") as f: 64 | r = requests.get(dataset_url) 65 | f.write(r.content) 66 | 67 | if base.endswith('.tar.gz'): 68 | obj = tarfile.open(temp_save_path) 69 | elif base.endswith('.zip'): 70 | obj = ZipFile(temp_save_path, 'r') 71 | else: 72 | raise ValueError("Unknown File Type: {0}.".format(base)) 73 | 74 | self._print("Unpacking Data...") 75 | obj.extractall(save_path) 76 | obj.close() 77 | os.remove(temp_save_path) 78 | 79 | def get(self, save_path, dataset=None): 80 | """ 81 | 82 | Download a dataset. 83 | 84 | Parameters: 85 | save_path (str) -- A directory to save the data to. 86 | dataset (str) -- (optional). A specific dataset to download. 87 | Note: this must include the file extension. 88 | If None, options will be presented for you 89 | to choose from. 90 | 91 | Returns: 92 | save_path_full (str) -- the absolute path to the downloaded data. 93 | 94 | """ 95 | if dataset is None: 96 | selected_dataset = self._present_options() 97 | else: 98 | selected_dataset = dataset 99 | 100 | save_path_full = join(save_path, selected_dataset.split('.')[0]) 101 | 102 | if isdir(save_path_full): 103 | warn("\n'{0}' already exists. Voiding Download.".format( 104 | save_path_full)) 105 | else: 106 | self._print('Downloading Data...') 107 | url = "{0}/{1}".format(self.url, selected_dataset) 108 | self._download_data(url, save_path=save_path) 109 | 110 | return abspath(save_path_full) 111 | -------------------------------------------------------------------------------- /util/html.py: -------------------------------------------------------------------------------- 1 | import dominate 2 | from dominate.tags import meta, h3, table, tr, td, p, a, img, br 3 | import os 4 | 5 | 6 | class HTML: 7 | """This HTML class allows us to save images and write texts into a single HTML file. 8 | 9 | It consists of functions such as (add a text header to the HTML file), 10 | (add a row of images to the HTML file), and (save the HTML to the disk). 11 | It is based on Python library 'dominate', a Python library for creating and manipulating HTML documents using a DOM API. 12 | """ 13 | 14 | def __init__(self, web_dir, title, refresh=0): 15 | """Initialize the HTML classes 16 | 17 | Parameters: 18 | web_dir (str) -- a directory that stores the webpage. HTML file will be created at /index.html; images will be saved at 0: 32 | with self.doc.head: 33 | meta(http_equiv="refresh", content=str(refresh)) 34 | 35 | def get_image_dir(self): 36 | """Return the directory that stores images""" 37 | return self.img_dir 38 | 39 | def add_header(self, text): 40 | """Insert a header to the HTML file 41 | 42 | Parameters: 43 | text (str) -- the header text 44 | """ 45 | with self.doc: 46 | h3(text) 47 | 48 | def add_images(self, ims, txts, links, width=400): 49 | """add images to the HTML file 50 | 51 | Parameters: 52 | ims (str list) -- a list of image paths 53 | txts (str list) -- a list of image names shown on the website 54 | links (str list) -- a list of hyperref links; when you click an image, it will redirect you to a new page 55 | """ 56 | self.t = table(border=1, style="table-layout: fixed;") # Insert a table 57 | self.doc.add(self.t) 58 | with self.t: 59 | with tr(): 60 | for im, txt, link in zip(ims, txts, links): 61 | with td(style="word-wrap: break-word;", halign="center", valign="top"): 62 | with p(): 63 | with a(href=os.path.join('images', link)): 64 | img(style="width:%dpx" % width, src=os.path.join('images', im)) 65 | br() 66 | p(txt) 67 | 68 | def save(self): 69 | """save the current content to the HMTL file""" 70 | html_file = '%s/index.html' % self.web_dir 71 | f = open(html_file, 'wt') 72 | f.write(self.doc.render()) 73 | f.close() 74 | 75 | 76 | if __name__ == '__main__': # we show an example usage here. 77 | html = HTML('web/', 'test_html') 78 | html.add_header('hello world') 79 | 80 | ims, txts, links = [], [], [] 81 | for n in range(4): 82 | ims.append('image_%d.png' % n) 83 | txts.append('text_%d' % n) 84 | links.append('image_%d.png' % n) 85 | html.add_images(ims, txts, links) 86 | html.save() 87 | -------------------------------------------------------------------------------- /util/image_pool.py: -------------------------------------------------------------------------------- 1 | import random 2 | import torch 3 | 4 | 5 | class ImagePool(): 6 | """This class implements an image buffer that stores previously generated images. 7 | 8 | This buffer enables us to update discriminators using a history of generated images 9 | rather than the ones produced by the latest generators. 10 | """ 11 | 12 | def __init__(self, pool_size): 13 | """Initialize the ImagePool class 14 | 15 | Parameters: 16 | pool_size (int) -- the size of image buffer, if pool_size=0, no buffer will be created 17 | """ 18 | self.pool_size = pool_size 19 | if self.pool_size > 0: # create an empty pool 20 | self.num_imgs = 0 21 | self.images = [] 22 | 23 | def query(self, images): 24 | """Return an image from the pool. 25 | 26 | Parameters: 27 | images: the latest generated images from the generator 28 | 29 | Returns images from the buffer. 30 | 31 | By 50/100, the buffer will return input images. 32 | By 50/100, the buffer will return images previously stored in the buffer, 33 | and insert the current images to the buffer. 34 | """ 35 | if self.pool_size == 0: # if the buffer size is 0, do nothing 36 | return images 37 | return_images = [] 38 | for image in images: 39 | image = torch.unsqueeze(image.data, 0) 40 | if self.num_imgs < self.pool_size: # if the buffer is not full; keep inserting current images to the buffer 41 | self.num_imgs = self.num_imgs + 1 42 | self.images.append(image) 43 | return_images.append(image) 44 | else: 45 | p = random.uniform(0, 1) 46 | if p > 0.5: # by 50% chance, the buffer will return a previously stored image, and insert the current image into the buffer 47 | random_id = random.randint(0, self.pool_size - 1) # randint is inclusive 48 | tmp = self.images[random_id].clone() 49 | self.images[random_id] = image 50 | return_images.append(tmp) 51 | else: # by another 50% chance, the buffer will return the current image 52 | return_images.append(image) 53 | return_images = torch.cat(return_images, 0) # collect all the images and return 54 | return return_images 55 | -------------------------------------------------------------------------------- /util/prior_box.py: -------------------------------------------------------------------------------- 1 | import torch 2 | from itertools import product as product 3 | import numpy as np 4 | from math import ceil 5 | 6 | 7 | class PriorBox(object): 8 | def __init__(self, cfg, image_size=None, phase='train'): 9 | super(PriorBox, self).__init__() 10 | self.min_sizes = cfg['min_sizes'] 11 | self.steps = cfg['steps'] 12 | self.clip = cfg['clip'] 13 | self.image_size = image_size 14 | self.feature_maps = [[ceil(self.image_size[0]/step), ceil(self.image_size[1]/step)] for step in self.steps] 15 | self.name = "s" 16 | 17 | def forward(self): 18 | anchors = [] 19 | for k, f in enumerate(self.feature_maps): 20 | min_sizes = self.min_sizes[k] 21 | for i, j in product(range(f[0]), range(f[1])): 22 | for min_size in min_sizes: 23 | s_kx = min_size / self.image_size[1] 24 | s_ky = min_size / self.image_size[0] 25 | dense_cx = [x * self.steps[k] / self.image_size[1] for x in [j + 0.5]] 26 | dense_cy = [y * self.steps[k] / self.image_size[0] for y in [i + 0.5]] 27 | for cy, cx in product(dense_cy, dense_cx): 28 | anchors += [cx, cy, s_kx, s_ky] 29 | 30 | # back to torch land 31 | output = torch.Tensor(anchors).view(-1, 4) 32 | if self.clip: 33 | output.clamp_(max=1, min=0) 34 | return output 35 | -------------------------------------------------------------------------------- /util/py_cpu_nms.py: -------------------------------------------------------------------------------- 1 | # -------------------------------------------------------- 2 | # Fast R-CNN 3 | # Copyright (c) 2015 Microsoft 4 | # Licensed under The MIT License [see LICENSE for details] 5 | # Written by Ross Girshick 6 | # -------------------------------------------------------- 7 | 8 | import numpy as np 9 | 10 | def py_cpu_nms(dets, thresh): 11 | """Pure Python NMS baseline.""" 12 | x1 = dets[:, 0] 13 | y1 = dets[:, 1] 14 | x2 = dets[:, 2] 15 | y2 = dets[:, 3] 16 | scores = dets[:, 4] 17 | 18 | areas = (x2 - x1 + 1) * (y2 - y1 + 1) 19 | order = scores.argsort()[::-1] 20 | 21 | keep = [] 22 | while order.size > 0: 23 | i = order[0] 24 | keep.append(i) 25 | xx1 = np.maximum(x1[i], x1[order[1:]]) 26 | yy1 = np.maximum(y1[i], y1[order[1:]]) 27 | xx2 = np.minimum(x2[i], x2[order[1:]]) 28 | yy2 = np.minimum(y2[i], y2[order[1:]]) 29 | 30 | w = np.maximum(0.0, xx2 - xx1 + 1) 31 | h = np.maximum(0.0, yy2 - yy1 + 1) 32 | inter = w * h 33 | ovr = inter / (areas[i] + areas[order[1:]] - inter) 34 | 35 | inds = np.where(ovr <= thresh)[0] 36 | order = order[inds + 1] 37 | 38 | return keep 39 | -------------------------------------------------------------------------------- /util/renderer.py: -------------------------------------------------------------------------------- 1 | import pdb 2 | import cv2 3 | import socket 4 | import numpy as np 5 | from pathlib import Path 6 | 7 | from lib.Sim3DR import RenderPipeline 8 | ip = socket.gethostbyname(socket.gethostname()) 9 | 10 | def _to_ctype(arr): 11 | if not arr.flags.c_contiguous: 12 | return arr.copy(order="C") 13 | return arr 14 | 15 | cfg = { 16 | "intensity_ambient": 0.3, 17 | "color_ambient": (1, 1, 1), 18 | "intensity_directional": 0.6, 19 | "color_directional": (1, 1, 1), 20 | "intensity_specular": 0.1, 21 | "specular_exp": 5, 22 | "light_pos": (0, 0, 5), 23 | "view_pos": (0, 0, 5), 24 | } 25 | render_app = RenderPipeline(**cfg) 26 | 27 | 28 | class Renderer(object): 29 | 30 | def __init__(self, img_size=800, alpha=0.7): 31 | npy_path = './npy/tris_2500x3_202110.npy' 32 | tris = np.load(npy_path)[:2304] 33 | tris[:, [0, 1]] = tris[:, [1, 0]] 34 | self.tris = np.ascontiguousarray(tris.astype(np.int32)) # (2304, 3) 35 | 36 | self.img_size = img_size 37 | self.alpha = alpha 38 | 39 | self.M_proj = np.array([ 40 | [1.574437, 0, 0, 0], 41 | [ 0, 1.574437, 0, 0], 42 | [ 0, 0, -0.99999976, -1], 43 | [ 0, 0, -0.0009999998, 0] 44 | ]) 45 | self.M1 = np.array([ 46 | [img_size/2, 0, 0, 0], 47 | [ 0, img_size/2, 0, 0], 48 | [ 0, 0, 1, 0], 49 | [img_size/2, img_size/2, 0, 1] 50 | ]) 51 | 52 | self.f = self.M_proj[0, 0] * img_size / 2 53 | 54 | 55 | 56 | 57 | 58 | 59 | def __call__(self, verts3d, R_t, overlap=None): 60 | ones = np.ones([verts3d.shape[0], 1]) 61 | verts_homo = np.concatenate([verts3d, ones], axis=1) 62 | 63 | verts = verts_homo @ R_t @ self.M_proj @ self.M1 64 | w_ = verts[:, [3]] 65 | verts = verts / w_ 66 | 67 | # points2d，x right，y down 68 | points2d = verts[:, :2] 69 | points2d[:, 1] = self.img_size - points2d[:, 1] 70 | 71 | 72 | verts_temp = np.concatenate([points2d, w_], axis=1) 73 | 74 | tz = R_t[3, 2] 75 | scale = self.f / tz 76 | verts_temp[:, 2] *= scale 77 | 78 | verts_temp = _to_ctype(verts_temp.astype(np.float32)) 79 | 80 | if overlap is None: 81 | overlap = np.zeros([self.img_size, self.img_size, 3], dtype=np.uint8) 82 | 83 | overlap_copy = overlap.copy() 84 | overlap = render_app(verts_temp, self.tris, overlap) # overlap上没有透明度 85 | 86 | img_render = cv2.addWeighted(overlap_copy, 1 - self.alpha, overlap, self.alpha, 0) 87 | return img_render 88 | 89 | 90 | 91 | 92 | -------------------------------------------------------------------------------- /util/util.py: -------------------------------------------------------------------------------- 1 | """This module contains simple helper functions """ 2 | from __future__ import print_function 3 | import torch 4 | import numpy as np 5 | from PIL import Image 6 | import os 7 | import importlib 8 | import argparse 9 | from argparse import Namespace 10 | import torchvision 11 | 12 | 13 | def str2bool(v): 14 | if isinstance(v, bool): 15 | return v 16 | if v.lower() in ('yes', 'true', 't', 'y', '1'): 17 | return True 18 | elif v.lower() in ('no', 'false', 'f', 'n', '0'): 19 | return False 20 | else: 21 | raise argparse.ArgumentTypeError('Boolean value expected.') 22 | 23 | 24 | def copyconf(default_opt, **kwargs): 25 | conf = Namespace(**vars(default_opt)) 26 | for key in kwargs: 27 | setattr(conf, key, kwargs[key]) 28 | return conf 29 | 30 | 31 | def find_class_in_module(target_cls_name, module): 32 | target_cls_name = target_cls_name.replace('_', '').lower() 33 | clslib = importlib.import_module(module) 34 | cls = None 35 | for name, clsobj in clslib.__dict__.items(): 36 | if name.lower() == target_cls_name: 37 | cls = clsobj 38 | 39 | assert cls is not None, "In %s, there should be a class whose name matches %s in lowercase without underscore(_)" % (module, target_cls_name) 40 | 41 | return cls 42 | 43 | 44 | def tensor2im(input_image, imtype=np.uint8): 45 | """"Converts a Tensor array into a numpy image array. 46 | 47 | Parameters: 48 | input_image (tensor) -- the input image tensor array 49 | imtype (type) -- the desired type of the converted numpy array 50 | """ 51 | if not isinstance(input_image, np.ndarray): 52 | if isinstance(input_image, torch.Tensor): # get the data from a variable 53 | image_tensor = input_image.data 54 | else: 55 | return input_image 56 | image_numpy = image_tensor[0].clamp(-1.0, 1.0).cpu().float().numpy() # convert it into a numpy array 57 | if image_numpy.shape[0] == 1: # grayscale to RGB 58 | image_numpy = np.tile(image_numpy, (3, 1, 1)) 59 | image_numpy = (np.transpose(image_numpy, (1, 2, 0)) + 1) / 2.0 * 255.0 # post-processing: tranpose and scaling 60 | else: # if it is a numpy array, do nothing 61 | image_numpy = input_image 62 | return image_numpy.astype(imtype) 63 | 64 | 65 | def diagnose_network(net, name='network'): 66 | """Calculate and print the mean of average absolute(gradients) 67 | 68 | Parameters: 69 | net (torch network) -- Torch network 70 | name (str) -- the name of the network 71 | """ 72 | mean = 0.0 73 | count = 0 74 | for param in net.parameters(): 75 | if param.grad is not None: 76 | mean += torch.mean(torch.abs(param.grad.data)) 77 | count += 1 78 | if count > 0: 79 | mean = mean / count 80 | print(name) 81 | print(mean) 82 | 83 | 84 | def save_image(image_numpy, image_path, aspect_ratio=1.0): 85 | """Save a numpy image to the disk 86 | 87 | Parameters: 88 | image_numpy (numpy array) -- input numpy array 89 | image_path (str) -- the path of the image 90 | """ 91 | 92 | image_pil = Image.fromarray(image_numpy) 93 | h, w, _ = image_numpy.shape 94 | 95 | if aspect_ratio is None: 96 | pass 97 | elif aspect_ratio > 1.0: 98 | image_pil = image_pil.resize((h, int(w * aspect_ratio)), Image.BICUBIC) 99 | elif aspect_ratio < 1.0: 100 | image_pil = image_pil.resize((int(h / aspect_ratio), w), Image.BICUBIC) 101 | image_pil.save(image_path) 102 | 103 | 104 | def print_numpy(x, val=True, shp=False): 105 | """Print the mean, min, max, median, std, and size of a numpy array 106 | 107 | Parameters: 108 | val (bool) -- if print the values of the numpy array 109 | shp (bool) -- if print the shape of the numpy array 110 | """ 111 | x = x.astype(np.float64) 112 | if shp: 113 | print('shape,', x.shape) 114 | if val: 115 | x = x.flatten() 116 | print('mean = %3.3f, min = %3.3f, max = %3.3f, median = %3.3f, std=%3.3f' % ( 117 | np.mean(x), np.min(x), np.max(x), np.median(x), np.std(x))) 118 | 119 | 120 | def mkdirs(paths): 121 | """create empty directories if they don't exist 122 | 123 | Parameters: 124 | paths (str list) -- a list of directory paths 125 | """ 126 | if isinstance(paths, list) and not isinstance(paths, str): 127 | for path in paths: 128 | mkdir(path) 129 | else: 130 | mkdir(paths) 131 | 132 | 133 | def mkdir(path): 134 | """create a single empty directory if it didn't exist 135 | 136 | Parameters: 137 | path (str) -- a single directory path 138 | """ 139 | if not os.path.exists(path): 140 | os.makedirs(path) 141 | 142 | 143 | def correct_resize_label(t, size): 144 | device = t.device 145 | t = t.detach().cpu() 146 | resized = [] 147 | for i in range(t.size(0)): 148 | one_t = t[i, :1] 149 | one_np = np.transpose(one_t.numpy().astype(np.uint8), (1, 2, 0)) 150 | one_np = one_np[:, :, 0] 151 | one_image = Image.fromarray(one_np).resize(size, Image.NEAREST) 152 | resized_t = torch.from_numpy(np.array(one_image)).long() 153 | resized.append(resized_t) 154 | return torch.stack(resized, dim=0).to(device) 155 | 156 | 157 | def correct_resize(t, size, mode=Image.BICUBIC): 158 | device = t.device 159 | t = t.detach().cpu() 160 | resized = [] 161 | for i in range(t.size(0)): 162 | one_t = t[i:i + 1] 163 | one_image = Image.fromarray(tensor2im(one_t)).resize(size, Image.BICUBIC) 164 | resized_t = torchvision.transforms.functional.to_tensor(one_image) * 2 - 1.0 165 | resized.append(resized_t) 166 | return torch.stack(resized, dim=0).to(device) 167 | 168 | 169 | 170 | 171 | def landmarks106to68(lms): 172 | idx = list(range(0, 33, 2)) + list(range(33, 64)) + list(range(84, 104)) 173 | lms_68 = lms[idx] 174 | factor = 0.6 175 | lms_68[37] = lms[53] * factor + lms[72] * (1 - factor) 176 | lms_68[38] = lms[54] * factor + lms[72] * (1 - factor) 177 | lms_68[40] = lms[56] * factor + lms[73] * (1 - factor) 178 | lms_68[41] = lms[57] * factor + lms[73] * (1 - factor) 179 | 180 | lms_68[43] = lms[59] * factor + lms[75] * (1 - factor) 181 | lms_68[44] = lms[60] * factor + lms[75] * (1 - factor) 182 | lms_68[46] = lms[62] * factor + lms[76] * (1 - factor) 183 | lms_68[47] = lms[63] * factor + lms[76] * (1 - factor) 184 | return lms_68 185 | 186 | 187 | 188 | 189 | --------------------------------------------------------------------------------