├── .gitignore ├── 3dfit ├── README.md ├── bodyfit.py └── render.py ├── LICENSE.txt ├── README.md ├── config.py ├── direct3d ├── bodyfit.py ├── conversion_tests.py ├── conversions.pyx ├── fit_forest.py ├── run_partforest_training.sh └── tools │ └── create_dataset.py ├── models ├── 3D │ ├── .AppleDouble │ │ └── .Parent │ ├── gmm_best_06.pkl │ ├── gmm_best_08.pkl │ ├── head_vid_heva.pkl │ ├── mask_filled.png │ ├── mask_filled_uniform.png │ ├── regressors_locked_normalized_hybrid.npz │ ├── template-bodyparts.ply │ ├── template.ply │ └── template_blender_low2.ply ├── pose │ ├── deepercut.prototxt │ ├── landmarks.pkl │ └── testpy_val_91_500_pkg.prototxt └── segmentation │ └── testpy_test_31_500_pkg_dorder.prototxt ├── patches └── deeplab.txt ├── pose ├── evaluate_pose.py ├── pose.py ├── run.sh ├── store_pose_results.py ├── sub.sed ├── tools │ ├── all_stats.txt │ └── create_dataset.py └── training │ └── config │ └── pose │ ├── n_classes.txt │ ├── solver.prototxt │ ├── solver2.prototxt │ ├── target_person_size.txt │ ├── testpy.prototxt │ └── train.prototxt ├── requirements.txt ├── segmentation ├── evaluate_segmentation.py ├── run.sh ├── segmentation.py ├── store_segmentation_results.py ├── sub.sed ├── tools │ ├── create_dataset.py │ └── render_orig_colors.py └── training │ └── config │ └── segmentation │ ├── n_classes.txt │ ├── solver.prototxt │ ├── solver2.prototxt │ ├── target_person_size.txt │ ├── testpy.prototxt │ └── train.prototxt ├── setup.py ├── tests ├── __init__.py └── tools.py └── up_tools ├── __init__.py ├── bake_vertex_colors.py ├── camera.py ├── capsule_ch.py ├── capsule_man.py ├── max_mixture_prior.py ├── mesh.py ├── model.py ├── render_segmented_views.py ├── robustifiers.py ├── sphere_collisions.py └── visualization.py /.gitignore: -------------------------------------------------------------------------------- 1 | *.pyc 2 | *.so 3 | *.egg-info 4 | *.z 5 | build 6 | segmentation/training/features 7 | segmentation/training/features2 8 | segmentation/training/states 9 | segmentation/training/list 10 | segmentation/training/log 11 | segmentation/training/model 12 | segmentation/training/states 13 | pose/training/features 14 | pose/training/features2 15 | pose/training/list 16 | pose/training/log 17 | pose/training/model 18 | pose/training/states 19 | *.caffemodel 20 | up_tools/sample2meshdist.cpp 21 | up_tools/fast_derivatives/smpl_derivatives.cpp 22 | models/pose/deepercut.caffemodel 23 | .#* 24 | applications 25 | 26 | models/pose/train2_iter_500000.caffemodel 27 | .*_cache 28 | .DS_Store 29 | .python-version 30 | -------------------------------------------------------------------------------- /3dfit/README.md: -------------------------------------------------------------------------------- 1 | # 3D Fitting to 14 or 91 Keypoints 2 | 3 | This subfolder contains scripts to fit a 3D body to 14 or 91 keypoints with the 4 | SMPLify fitting objective and render it. The silhouette fitting objective is 5 | currently not included due to licensing issues. The code is based on the 6 | [SMPLify code](http://smplify.is.tuebingen.mpg.de) and some of its library files 7 | in the `up_tools` folder as well as the model 8 | `models/3D/basicModel_neutral_lbs_10_207_0_v1.0.0.pkl` fall under the SMPLify 9 | license. If you find this code useful for your research, please consider also 10 | citing 11 | 12 | ``` 13 | @inproceedings{Bogo:ECCV:2016, 14 | title = {Keep it {SMPL}: Automatic Estimation of {3D} Human Pose and Shape from a Single Image}, 15 | author = {Bogo, Federica and Kanazawa, Angjoo and Lassner, Christoph and Gehler, Peter and Romero, Javier and Black, Michael J.}, 16 | booktitle = {Computer Vision -- ECCV 2016}, 17 | series = {Lecture Notes in Computer Science}, 18 | publisher = {Springer International Publishing}, 19 | month = oct, 20 | year = {2016} 21 | } 22 | ``` 23 | -------------------------------------------------------------------------------- /3dfit/render.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python 2 | """Render a model.""" 3 | import cPickle as pickle 4 | import logging as _logging 5 | import os as _os 6 | import os.path as _path 7 | # pylint: disable=invalid-name 8 | import sys as _sys 9 | from config import SMPL_FP 10 | from copy import copy 11 | 12 | import click as _click 13 | import cv2 14 | import numpy as np 15 | from opendr.camera import ProjectPoints 16 | from opendr.lighting import LambertianPointLight 17 | from opendr.renderer import ColoredRenderer, TexturedRenderer 18 | 19 | from up_tools.camera import rotateY 20 | from up_tools.mesh import Mesh 21 | 22 | _sys.path.insert(0, _path.join(_path.dirname(__file__), '..')) 23 | _sys.path.insert(0, SMPL_FP) 24 | try: 25 | from smpl.serialization import load_model as _load_model 26 | except: 27 | from smpl_webuser.serialization import load_model as _load_model 28 | 29 | _LOGGER = _logging.getLogger(__name__) 30 | _TEMPLATE_MESH = Mesh(filename=_os.path.join(_os.path.dirname(__file__), 31 | '..', 'models', '3D', 'template.ply')) 32 | _TEMPLATE_MESH_SEGMENTED = Mesh(filename=_os.path.join(_os.path.dirname(__file__), 33 | '..', 'models', '3D', 'template-bodyparts.ply')) 34 | _COLORS = { 35 | 'pink': [.6, .6, .8], 36 | 'cyan': [.7, .75, .5], 37 | 'yellow': [.5, .7, .75], 38 | 'grey': [.7, .7, .7], 39 | } 40 | 41 | 42 | def _create_renderer( # pylint: disable=too-many-arguments 43 | w=640, 44 | h=480, 45 | rt=np.zeros(3), 46 | t=np.zeros(3), 47 | f=None, 48 | c=None, 49 | k=None, 50 | near=1., 51 | far=10., 52 | texture=None): 53 | """Create a renderer for the specified parameters.""" 54 | f = np.array([w, w]) / 2. if f is None else f 55 | c = np.array([w, h]) / 2. if c is None else c 56 | k = np.zeros(5) if k is None else k 57 | 58 | if texture is not None: 59 | rn = TexturedRenderer() 60 | else: 61 | rn = ColoredRenderer() 62 | 63 | rn.camera = ProjectPoints(rt=rt, t=t, f=f, c=c, k=k) 64 | rn.frustum = {'near':near, 'far':far, 'height':h, 'width':w} 65 | if texture is not None: 66 | rn.texture_image = np.asarray(cv2.imread(texture), np.float64)/255. 67 | return rn 68 | 69 | 70 | def _stack_with(rn, mesh, texture): 71 | if texture is not None: 72 | if not hasattr(mesh, 'ft'): 73 | mesh.ft = mesh.f 74 | mesh.vt = mesh.v[:, :2] 75 | rn.ft = np.vstack((rn.ft, mesh.ft+len(rn.vt))) 76 | rn.vt = np.vstack((rn.vt, mesh.vt)) 77 | rn.f = np.vstack((rn.f, mesh.f+len(rn.v))) 78 | rn.v = np.vstack((rn.v, mesh.v)) 79 | rn.vc = np.vstack((rn.vc, mesh.vc)) 80 | 81 | 82 | def _simple_renderer(rn, meshes, yrot=0, texture=None): 83 | mesh = meshes[0] 84 | if texture is not None: 85 | if not hasattr(mesh, 'ft'): 86 | mesh.ft = copy(mesh.f) 87 | vt = copy(mesh.v[:, :2]) 88 | vt -= np.min(vt, axis=0).reshape((1, -1)) 89 | vt /= np.max(vt, axis=0).reshape((1, -1)) 90 | mesh.vt = vt 91 | mesh.texture_filepath = rn.texture_image 92 | 93 | # Set camers parameters 94 | if texture is not None: 95 | rn.set(v=mesh.v, f=mesh.f, vc=mesh.vc, ft=mesh.ft, vt=mesh.vt, bgcolor=np.ones(3)) 96 | else: 97 | rn.set(v=mesh.v, f=mesh.f, vc=mesh.vc, bgcolor=np.ones(3)) 98 | 99 | for next_mesh in meshes[1:]: 100 | _stack_with(rn, next_mesh, texture) 101 | 102 | # Construct Back Light (on back right corner) 103 | albedo = rn.vc 104 | 105 | rn.vc = LambertianPointLight( 106 | f=rn.f, 107 | v=rn.v, 108 | num_verts=len(rn.v), 109 | light_pos=rotateY(np.array([-200, -100, -100]), yrot), 110 | vc=albedo, 111 | light_color=np.array([1, 1, 1])) 112 | 113 | # Construct Left Light 114 | rn.vc += LambertianPointLight( 115 | f=rn.f, 116 | v=rn.v, 117 | num_verts=len(rn.v), 118 | light_pos=rotateY(np.array([800, 10, 300]), yrot), 119 | vc=albedo, 120 | light_color=np.array([1, 1, 1])) 121 | 122 | # Construct Right Light 123 | rn.vc += LambertianPointLight( 124 | f=rn.f, 125 | v=rn.v, 126 | num_verts=len(rn.v), 127 | light_pos=rotateY(np.array([-500, 500, 1000]), yrot), 128 | vc=albedo, 129 | light_color=np.array([.7, .7, .7])) 130 | 131 | return rn.r 132 | 133 | 134 | # pylint: disable=too-many-locals 135 | def render(model, image, cam, steps, segmented=False, scale=1.): 136 | """Render a sequence of views from a fitted body model.""" 137 | assert steps >= 1 138 | if segmented: 139 | texture = None 140 | mesh = copy(_TEMPLATE_MESH_SEGMENTED) 141 | else: 142 | texture = None 143 | mesh = copy(_TEMPLATE_MESH) 144 | mesh.vc = _COLORS['pink'] 145 | #cmesh = Mesh(_os.path.join(_os.path.dirname(__file__), 146 | # 'template-bodyparts-corrected-labeled-split5.ply')) 147 | #mesh.vc = cmesh.vc.copy() 148 | # render ply 149 | model.betas[:len(cam['betas'])] = cam['betas'] 150 | model.pose[:] = cam['pose'] 151 | model.trans[:] = cam['trans'] 152 | 153 | mesh.v = model.r 154 | w, h = (image.shape[1], image.shape[0]) 155 | dist = np.abs(cam['t'][2] - np.mean(mesh.v, axis=0)[2]) 156 | rn = _create_renderer(w=int(w * scale), 157 | h=int(h * scale), 158 | near=4., 159 | far=30.+dist, 160 | rt=np.array(cam['rt']), 161 | t=np.array(cam['t']), 162 | f=np.array([cam['f'], cam['f']]) * scale, 163 | # c=np.array(cam['cam_c']), 164 | texture=texture) 165 | light_yrot = np.radians(120) 166 | base_mesh = copy(mesh) 167 | renderings = [] 168 | for angle in np.linspace(0., 2. * np.pi, num=steps, endpoint=False): 169 | mesh.v = rotateY(base_mesh.v, angle) 170 | imtmp = _simple_renderer(rn=rn, 171 | meshes=[mesh], 172 | yrot=light_yrot, 173 | texture=texture) 174 | if segmented: 175 | imtmp = imtmp[:, :, ::-1] 176 | renderings.append(imtmp * 255.) 177 | return renderings 178 | 179 | 180 | @_click.command() 181 | @_click.argument('filename', type=_click.Path(exists=True, readable=True)) 182 | @_click.option('--segmented', 183 | type=_click.BOOL, 184 | is_flag=True, 185 | help='If set, use a segmented mesh.') 186 | @_click.option('--steps', 187 | type=_click.INT, 188 | help='The number of rotated images to render. Default: 1.', 189 | default=1) 190 | @_click.option("--scale", 191 | type=_click.FLOAT, 192 | help="Render the results at this scale.", 193 | default=1.) 194 | def cli(filename, segmented=False, steps=1, scale=1.): 195 | """Render a 3D model for an estimated body fit. Provide the image (!!) filename.""" 196 | model = { 197 | 'neutral': _load_model( 198 | _os.path.join(_os.path.dirname(__file__), 199 | '..', 'models', '3D', 200 | 'basicModel_neutral_lbs_10_207_0_v1.0.0.pkl')) 201 | } 202 | fileending = filename[-4:] 203 | pkl_fname = filename.replace(fileending, '.pkl') 204 | if not _path.exists(pkl_fname): 205 | pkl_fname = filename[:-10] + '_body.pkl' 206 | if not _path.exists(pkl_fname): 207 | pkl_fname = filename + '_body.pkl' 208 | _LOGGER.info("Rendering 3D model for image %s and parameters %s...", 209 | filename, pkl_fname) 210 | assert _os.path.exists(pkl_fname), ( 211 | 'Stored body fit does not exist for {}: {}!'.format( 212 | filename, pkl_fname)) 213 | 214 | image = cv2.imread(filename) 215 | with open(pkl_fname) as f: 216 | cam = pickle.load(f) 217 | renderings = render(model['neutral'], image, cam, steps, segmented, scale) 218 | for ridx, rim in enumerate(renderings): 219 | if segmented: 220 | out_fname = filename + '_body_segmented_%d.png' % (ridx) 221 | else: 222 | out_fname = filename + '_body_%d.png' % (ridx) 223 | cv2.imwrite(out_fname, rim) 224 | 225 | 226 | if __name__ == '__main__': 227 | _logging.basicConfig(level=_logging.INFO) 228 | _logging.getLogger("opendr.lighting").setLevel(_logging.FATAL) # pylint: disable=no-attribute 229 | cli() # pylint: disable=no-value-for-parameter 230 | -------------------------------------------------------------------------------- /LICENSE.txt: -------------------------------------------------------------------------------- 1 | Copyright (C) 2017 University of Tübingen, Christoph Lassner. 2 | 3 | This code is available under the Creative Commons Attribution-NonCommercial 4.0 4 | International Public License (https://creativecommons.org/licenses/by-nc/4.0/) 5 | with the exception of the following files in the folder up_tools: 6 | 7 | * capsule_ch.py, 8 | * capsule_man.py, 9 | * max_mixture_prior.py, 10 | * robustifiers.py, 11 | * sphere_collisions.py, 12 | 13 | which fall under the SMPLify license conditions available here: 14 | http://smplify.is.tuebingen.mpg.de/downloads . 15 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # Unite the People code repository 2 | 3 | Requirements: 4 | 5 | * OpenCV (on Ubuntu, e.g., install libopencv-dev and python-opencv). 6 | * SMPL (download at http://smpl.is.tue.mpg.de/downloads) and unzip to a 7 | place of your choice. 8 | * OpenDR (just run `pip install opendr`, unfortunately can't be done 9 | automatically with the setuptools requirements. 10 | * If you want to train a segmentation model, Deeplab V2 11 | (https://bitbucket.org/aquariusjay/deeplab-public-ver2) with a minimal patch 12 | applied that can be found in the subdirectory `patches`, to enable on the fly 13 | mirroring of the segmented images. Since I didn't use the MATLAB interface and 14 | did not care about fixing related errors, I just deleted 15 | `src/caffe/layers/mat_{read,write}_layer.cpp` as well as 16 | `src/caffe/util/matio_io.cpp` and built with `-DWITH_matlab=Off`. 17 | * If you want to train a pose model, the Deepercut caffe 18 | (https://github.com/eldar/deepcut-cnn). 19 | * If you want to get deepercut-cnn predictions, download the deepercut 20 | .caffemodel file and place it in `models/pose/deepercut.caffemodel`. 21 | * Edit the file `config.py` to set up the paths. 22 | * Register on https://smpl.is.tue.mpg.de/ to obtain a SMPL license and place 23 | the model file at `models/3D/basicModel_neutral_lbs_10_207_0_v1.0.0.pkl`. 24 | 25 | The rest of the requirements is then automatically installed when running: 26 | 27 | ``` 28 | python setup.py develop 29 | ``` 30 | 31 | ## Folder structure 32 | 33 | For each of the tasks we described, there is one subfolder with the related 34 | executables. All files that are being used for training or testing models are 35 | executable and provide a full synopsis when run with the `--help` option. In the 36 | respective `tools` subfolder for each task, there is a `create_dataset.py` 37 | script to summarize the data in the proper formats. This must be usually run 38 | before the training script. The `models` folder contains pretrained models and 39 | infos, `patches` a patch for deeplab caffe, `tests` some Python tests and 40 | `up_tools` some Python tools that are shared between modalities. 41 | 42 | There is a Docker image available that has been created by TheWebMonks here (not 43 | affiliated with the authors): https://github.com/TheWebMonks/demo-2d3d . 44 | 45 | ### Bodyfit 46 | 47 | The adjusted SMPLify code to fit bodies to 91 keypoints is located in the folder 48 | `3dfit`. It can be used for 14 or 91 keypoints. Use the script `3dfit/render.py` 49 | to render a fitted body. 50 | 51 | ### Direct 3D fitting using regression forests 52 | 53 | The relevant files are in the folder `direct3d`. Run 54 | `run_partforest_training.sh` to train all regressors. After that, you can use 55 | `bodyfit.py` to get predictions from estimated keypoints of the 91 keypoint pose 56 | predictor. 57 | 58 | ### 91 keypoint pose prediction 59 | 60 | The `pose` folder containes infrastructure for 91 keypoint pose prediction. Use 61 | the script `pose/tools/create_dataset.py` with a dataset name of your choice and 62 | a target person size of 500 pixels to create the pose data from UP-3D, 63 | alternatively download it from our [website](http://up.is.tuebingen.mpg.de). 64 | 65 | Configure a model by creating the model configuration folder 66 | `pose/training/config/modelname` by cloning the `pose` model. Then you can run 67 | `run.sh {train,test,evaluate,trfull,tefull,evfull} modelname` to run training, 68 | testing or evaluation on either the reduced training set with the held-out 69 | validation set as test data or the full training set and real test data. We 70 | initialized our training from the original Resnet models 71 | (https://github.com/KaimingHe/deep-residual-networks). You can do so by 72 | downloading the model and saving it as 73 | `pose/training/config/modelname/init.caffemodel`. 74 | 75 | 76 | The `pose.py` script will produce a pose prediction for an image. It assumes 77 | that a model with name `pose` has been trained (or downloaded). We normalize the 78 | training images w.r.t. person size, that's why the model works best for images 79 | with a rough person height of 500 pixels. Multiple people are not taken into 80 | account; for every joint the `arg max` position is used over the full image. 81 | 82 | ### 31 part segmentation 83 | 84 | The folder setup is just as for the keypoint estimation: use 85 | `segmentation/tools/create_dataset.py` to create a segmentation dataset from the 86 | UP-3D data or download it (again, we used target person size 500). Then use 87 | `run.sh {train,test,evaluate,trfull,tefull,evfull} modelname` as described above 88 | to create your models. The `segmentation.py` script can be used to get 89 | segmentation results for the model named `segmentation` from and image. We 90 | initialized our models from the Deeplab trained models available 91 | [here](http://liangchiehchen.com/projects/DeepLabv2_resnet.html). Move the 92 | model file to `segmentation/training/modelname/init.caffemodel`. 93 | 94 | ## Website, citation, license 95 | 96 | You can find more information on the [website](http://up.is.tuebingen.mpg.de). 97 | If you use this code for your research, please consider citing us: 98 | 99 | ``` 100 | @inproceedings{Lassner:UP:2017, 101 | title = {Unite the People: Closing the Loop Between 3D and 2D Human Representations}, 102 | author = {Lassner, Christoph and Romero, Javier and Kiefel, Martin and Bogo, Federica and Black, Michael J. and Gehler, Peter V.}, 103 | booktitle = {Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition}, 104 | month = jul, 105 | year = {2017}, 106 | url = {http://up.is.tuebingen.mpg.de}, 107 | month_numeric = {7} 108 | } 109 | ``` 110 | 111 | License: [Creative Commons Non-Commercial 4.0](https://creativecommons.org/licenses/by-nc/4.0/). 112 | 113 | The code for 3D fitting is based on the [SMPLify](http://smplify.is.tue.mpg.de) 114 | code. Parts of the files in the folder `up_tools` (`capsule_ch.py`, 115 | `capsule_man.py`, `max_mixture_prior.py`, `robustifiers.py`, 116 | `sphere_collisions.py`) as well as the model 117 | `models/3D/basicModel_neutral_lbs_10_207_0_v1.0.0.pkl` fall under the SMPLify 118 | license conditions. 119 | -------------------------------------------------------------------------------- /config.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python2 2 | """Configuration values for the project.""" 3 | import os.path as path 4 | import click 5 | 6 | 7 | ############################ EDIT HERE ######################################## 8 | SMPL_FP = path.expanduser("~/smpl") 9 | DEEPLAB_BUILD_FP = path.expanduser("~/git/deeplab-public-ver2/build") 10 | DEEPERCUT_CNN_BUILD_FP = path.expanduser("~/git/deepcut-cnn/build_cluster") 11 | UP3D_FP = path.expanduser("~/datasets/up-3d") 12 | SEG_DATA_FP = path.expanduser("~/datasets/seg_prepared") 13 | POSE_DATA_FP = path.expanduser("~/datasets/pose_prepared") 14 | DIRECT3D_DATA_FP = path.expanduser("~/datasets/2dto3d_prepared") 15 | 16 | 17 | ############################################################################### 18 | # Infrastructure. Don't edit. # 19 | ############################################################################### 20 | 21 | @click.command() 22 | @click.argument('key', type=click.STRING) 23 | def cli(key): 24 | """Print a config value to STDOUT.""" 25 | if key in globals().keys(): 26 | print globals()[key] 27 | else: 28 | raise Exception("Requested configuration value not available! " 29 | "Available keys: " + 30 | str([kval for kval in globals().keys() if kval.isupper()]) + 31 | ".") 32 | 33 | 34 | if __name__ == '__main__': 35 | cli() # pylint: disable=no-value-for-parameter 36 | 37 | -------------------------------------------------------------------------------- /direct3d/conversion_tests.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python2 2 | """Tests for the conversion code.""" 3 | import unittest 4 | 5 | import numpy as np 6 | import pyximport; pyximport.install() # pylint: disable=multiple-statements 7 | import conversions # pylint: disable=import-error 8 | 9 | 10 | class ConversionTests(unittest.TestCase): 11 | 12 | """Tests for the conversions.""" 13 | 14 | def _create_test_axis_angle_matrix(self): # pylint: disable=no-self-use 15 | """Create a test matrix.""" 16 | angles = np.random.normal(size=(1000, 3)).astype('float32') 17 | for angle_idx, angle_w in enumerate(angles): 18 | angle_w = angles[angle_idx] 19 | angle_wnorm = np.linalg.norm(angle_w) 20 | if angle_wnorm > np.pi: 21 | angles[angle_idx] = angle_w / angle_wnorm * np.fmod(angle_wnorm, np.pi) 22 | return angles 23 | 24 | def test_axis_angle_versor(self): 25 | """Test cycling from axis angle through versor back.""" 26 | np.random.seed(1) 27 | angles = self._create_test_axis_angle_matrix() 28 | vers = conversions.axis_angle_to_versor_ipr(angles.copy()) 29 | angles_res = conversions.versor_to_axis_angle_ipr(vers) 30 | ''' 31 | for idx in range(1000): 32 | if not np.allclose(angles[idx], angles_res[idx]): 33 | print idx, angles[idx], angles_res[idx] 34 | ''' 35 | self.assertTrue(np.allclose(angles, angles_res, atol=1e-5, rtol=1e-4)) 36 | 37 | def test_matrix_to_versor(self): 38 | """Test matrix to versor conversion.""" 39 | np.random.seed(1) 40 | angles = self._create_test_axis_angle_matrix() 41 | matvals = np.array([-0.40967655, 0.07308591, -0.90929848, 0.07226962, 42 | -0.99105203, -0.1122174, -0.90936369, -0.11168748, 43 | 0.40072888], dtype=np.float32) 44 | vers = conversions.matrix_to_versor(matvals.reshape((1, 9))) 45 | aangle = conversions.versor_to_axis_angle_ipr(vers) 46 | 47 | def test_matrix_cycle(self): 48 | """Test cycling from axis angle through matrix through versor back.""" 49 | np.random.seed(1) 50 | angles = self._create_test_axis_angle_matrix() 51 | mat = conversions.axis_angle_to_matrix(angles) 52 | vers = conversions.matrix_to_versor(mat) 53 | angles_res = conversions.versor_to_axis_angle_ipr(vers) 54 | ''' 55 | for idx in range(1000): 56 | if not np.allclose(angles[idx], angles_res[idx]): 57 | print idx 58 | ''' 59 | self.assertTrue(np.allclose(angles, angles_res)) 60 | 61 | def test_matrix_projection(self): 62 | """Test the manifold projection.""" 63 | np.random.seed(1) 64 | angles = self._create_test_axis_angle_matrix() 65 | mat = conversions.axis_angle_to_matrix(angles) 66 | n_failed = 0 67 | for matrix_id, matrix in enumerate(mat): # pylint: disable=unused-variable 68 | self.assertTrue(np.allclose(np.linalg.det(matrix.reshape((3, 3))), 69 | 1.)) 70 | """ 71 | if not np.allclose(np.dot(matrix.reshape((3, 3)), 72 | matrix.reshape((3, 3)).T), 73 | np.eye(3), atol=1e-6): 74 | print np.dot(matrix.reshape((3, 3)), matrix.reshape((3, 3)).T), matrix_id 75 | """ 76 | self.assertTrue(np.allclose(np.dot(matrix.reshape((3, 3)), 77 | matrix.reshape((3, 3)).T), 78 | np.eye(3), atol=1e-6)) 79 | # Perturb. 80 | offs = np.random.normal() 81 | idx = np.random.randint(low=0, high=9) 82 | matrix[idx] += offs 83 | # Recover. 84 | try: 85 | rec, dist = conversions.project_to_rot(matrix) # pylint: disable=unused-variable 86 | self.assertTrue(np.allclose(np.linalg.det(rec.reshape((3, 3))), 87 | 1., atol=1e-5)) 88 | """ 89 | if not np.allclose(np.dot(rec.reshape((3, 3)), 90 | rec.reshape((3, 3)).T), 91 | np.eye(3), atol=1e-6): 92 | print np.dot(rec.reshape((3, 3)), 93 | rec.reshape((3, 3)).T), matrix_id 94 | """ 95 | self.assertTrue(np.allclose(np.dot(rec.reshape((3, 3)), 96 | rec.reshape((3, 3)).T), 97 | np.eye(3), atol=1e-6)) 98 | except: # pylint: disable=bare-except 99 | n_failed += 1 100 | self.assertEqual(n_failed, 56) 101 | # print 'could not recover matrix in %d cases.' % (n_failed) 102 | 103 | 104 | class RotationForestTestCase(unittest.TestCase): 105 | 106 | """Test the rotation forests.""" 107 | 108 | def _create_test_axis_angle_matrix(self, n_angles=1): # pylint: disable=no-self-use 109 | """Create a test matrix.""" 110 | angles = np.random.normal(size=(1000, 3*n_angles)).astype('float32') 111 | for angle_idx, angle_w in enumerate(angles): 112 | for line_idx in range((angles.shape[1] // 3)): 113 | angle_w = angles[angle_idx, line_idx * 3] 114 | angle_wnorm = np.linalg.norm(angle_w) 115 | if angle_wnorm > np.pi: 116 | angles[angle_idx, line_idx * 3:(line_idx+1) * 3] = \ 117 | angle_w / angle_wnorm * np.fmod(angle_wnorm, np.pi) 118 | return angles 119 | 120 | def test_basic(self): 121 | """Test the basic functionality.""" 122 | #np.seterr(all='raise') 123 | angles = self._create_test_axis_angle_matrix() 124 | mat = conversions.axis_angle_to_matrix(angles) 125 | dta = np.random.uniform(size=(1000, 5)) 126 | from rotation_forest import RotationForest 127 | rf = RotationForest(n_jobs=2) # pylint: disable=invalid-name 128 | rf.fit(dta, mat) 129 | res = rf.predict(dta) 130 | import cPickle as pickle 131 | pstr = pickle.dumps(rf) 132 | nf = pickle.loads(pstr) # pylint: disable=invalid-name 133 | self.assertTrue(np.all(nf.predict(dta) == res)) 134 | 135 | def test_multitarget(self): 136 | """Test multiple prediction target functionality.""" 137 | #np.seterr(all='raise') 138 | angles = self._create_test_axis_angle_matrix(n_angles=2) 139 | mat = conversions.axis_angle_to_matrix(angles) 140 | dta = np.random.uniform(size=(1000, 5)) 141 | from rotation_forest import RotationForest 142 | rf = RotationForest(n_jobs=2) # pylint: disable=invalid-name 143 | rf.fit(dta, mat) 144 | res = rf.predict(dta) 145 | import cPickle as pickle 146 | pstr = pickle.dumps(rf) 147 | nf = pickle.loads(pstr) # pylint: disable=invalid-name 148 | self.assertTrue(np.all(nf.predict(dta) == res)) 149 | 150 | 151 | def test_opencv_rodrigues(self): 152 | """Test the manifold projection.""" 153 | import cv2 154 | np.random.seed(1) 155 | angles = self._create_test_axis_angle_matrix() 156 | mat = conversions.axis_angle_to_matrix(angles) 157 | n_failed = 0 158 | for matrix_id, matrix in enumerate(mat): # pylint: disable=unused-variable 159 | self.assertTrue(np.allclose(np.linalg.det(matrix.reshape((3, 3))), 160 | 1.)) 161 | """ 162 | if not np.allclose(np.dot(matrix.reshape((3, 3)), 163 | matrix.reshape((3, 3)).T), 164 | np.eye(3), atol=1e-6): 165 | print np.dot(matrix.reshape((3, 3)), matrix.reshape((3, 3)).T), matrix_id 166 | """ 167 | self.assertTrue(np.allclose(np.dot(matrix.reshape((3, 3)), 168 | matrix.reshape((3, 3)).T), 169 | np.eye(3), atol=1e-6)) 170 | # Perturb. 171 | offs = np.random.normal() 172 | idx = np.random.randint(low=0, high=9) 173 | matrix[idx] += offs 174 | # Recover. 175 | rec = cv2.Rodrigues(cv2.Rodrigues(matrix.reshape((3, 3)))[0])[0] # pylint: disable=unused-variable 176 | self.assertTrue(np.allclose(np.linalg.det(rec.reshape((3, 3))), 177 | 1., atol=1e-5)) 178 | """ 179 | if not np.allclose(np.dot(rec.reshape((3, 3)), 180 | rec.reshape((3, 3)).T), 181 | np.eye(3), atol=1e-6): 182 | print np.dot(rec.reshape((3, 3)), 183 | rec.reshape((3, 3)).T), matrix_id 184 | """ 185 | self.assertTrue(np.allclose(np.dot(rec.reshape((3, 3)), 186 | rec.reshape((3, 3)).T), 187 | np.eye(3), atol=1e-6)) 188 | # print 'could not recover matrix in %d cases.' % (n_failed) 189 | 190 | 191 | if __name__ == '__main__': 192 | unittest.main() 193 | -------------------------------------------------------------------------------- /direct3d/conversions.pyx: -------------------------------------------------------------------------------- 1 | """Angle representation conversion functions in Cython. 2 | 3 | Useful reference: 4 | http://www.euclideanspace.com/maths/geometry/rotations/conversions/angleToEuler/ 5 | """ 6 | # cython: infer_types=True 7 | # cython: boundscheck=True 8 | # cython: wraparound=False 9 | import numpy as np 10 | from scipy.linalg import polar 11 | cimport numpy as np 12 | try: 13 | import cv2 14 | CV2_AVAILABLE = True 15 | except ImportError: 16 | CV2_AVAILABLE = False 17 | 18 | 19 | cpdef axis_angle_to_matrix(np.ndarray[np.float32_t, ndim=2] anglearr): 20 | assert anglearr.shape[1] % 3 == 0 21 | cdef int sample_idx, rot_idx 22 | cdef np.float32_t angle 23 | cdef np.ndarray[np.float32_t, ndim=1] rep 24 | cdef np.float32_t x, y, z, s, c, t 25 | cdef np.float32_t tmp1, tmp2 26 | cdef np.ndarray[np.float32_t, ndim=2] retarr 27 | cdef np.float32_t[:,:] workarr 28 | retarr = np.zeros((anglearr.shape[0], anglearr.shape[1] * 3), dtype=np.float32) 29 | for sample_idx in range(anglearr.shape[0]): 30 | for rot_idx in range(0, anglearr.shape[1], 3): 31 | rep = anglearr[sample_idx, rot_idx:rot_idx + 3] 32 | workarr = retarr[sample_idx, rot_idx // 3 * 9:(rot_idx // 3 + 1) * 9].reshape((3, 3)) 33 | angle = np.linalg.norm(rep) 34 | if angle == 0.: 35 | x, y, z = 0., 0., 0. 36 | else: 37 | x, y, z = rep / angle 38 | s = np.sin(angle) 39 | c = np.cos(angle) 40 | t = 1. - c 41 | # Diagonal. 42 | workarr[0, 0] = c + x*x*t 43 | workarr[1, 1] = c + y*y*t 44 | workarr[2, 2] = c + z*z*t 45 | tmp1 = x*y*t 46 | tmp2 = z*s 47 | # Out of diagonal. 48 | workarr[1, 0] = tmp1 + tmp2 49 | workarr[0, 1] = tmp1 - tmp2 50 | tmp1 = x*z*t 51 | tmp2 = y*s 52 | workarr[2, 0] = tmp1 - tmp2 53 | workarr[0, 2] = tmp1 + tmp2 54 | tmp1 = y*z*t 55 | tmp2 = x*s 56 | workarr[2, 1] = tmp1 + tmp2 57 | workarr[1, 2] = tmp1 - tmp2 58 | return retarr 59 | 60 | 61 | cpdef matrix_to_versor(np.ndarray[np.float32_t, ndim=2] anglearr): 62 | assert anglearr.shape[1] % 9 == 0 63 | cdef np.float32_t tr 64 | cdef int sample_idx, rot_idx 65 | cdef np.float32_t S, w 66 | cdef np.ndarray[np.float32_t, ndim=2] retarr, rep 67 | cdef np.ndarray[np.float32_t, ndim=1] workarr 68 | retarr = np.zeros((anglearr.shape[0], anglearr.shape[1] // 3), dtype=np.float32) 69 | for sample_idx in range(anglearr.shape[0]): 70 | for rot_idx in range(0, anglearr.shape[1], 9): 71 | rep = anglearr[sample_idx, rot_idx:rot_idx + 9].reshape((3, 3)) 72 | workarr = retarr[sample_idx, rot_idx // 9 * 3:(rot_idx // 9 + 1) * 3] 73 | # Determine the trace. 74 | tr = rep[0, 0] + rep[1, 1] + rep[2, 2] 75 | if tr > 0: 76 | S = np.sqrt(tr + 1.) * 2 # S=4*q 77 | w = 0.25 * S 78 | workarr[0] = (rep[2, 1] - rep[1, 2]) / S 79 | workarr[1] = (rep[0, 2] - rep[2, 0]) / S 80 | workarr[2] = (rep[1, 0] - rep[0, 1]) / S 81 | elif rep[0, 0] > rep[1, 1] and rep[0, 0] > rep[2, 2]: 82 | S = np.sqrt(1. + rep[0, 0] - rep[1, 1] - rep[2, 2]) * 2. # S=4*qx 83 | w = (rep[2, 1] - rep[1, 2]) / S 84 | workarr[0] = 0.25 * S 85 | workarr[1] = (rep[0, 1] + rep[1, 0]) / S 86 | workarr[2] = (rep[0, 2] + rep[2, 0]) / S 87 | elif rep[1, 1] > rep[2, 2]: 88 | S = np.sqrt(1. + rep[1, 1] - rep[0, 0] - rep[2, 2]) * 2. # S=4*qy 89 | w = (rep[0, 2] - rep[2, 0]) / S 90 | workarr[0] = (rep[0, 1] + rep[1, 0]) / S 91 | workarr[1] = 0.25 * S 92 | workarr[2] = (rep[1, 2] + rep[2, 1]) / S 93 | else: 94 | S = np.sqrt(1. + rep[2, 2] - rep[0, 0] - rep[1, 1]) * 2. # S=4*qz 95 | w = (rep[1, 0] - rep[0, 1]) / S 96 | workarr[0] = (rep[0, 2] + rep[2, 0]) / S 97 | workarr[1] = (rep[1, 2] + rep[2, 1]) / S 98 | workarr[2] = 0.25 * S 99 | if w < 0: 100 | workarr *= -1. 101 | if np.abs(w) >= 1. or not np.allclose(np.sqrt(np.sum(np.square(workarr)) + np.square(w)), 1.): 102 | print 'strange mat to vers result for:', rep 103 | assert np.allclose(np.sqrt(np.sum(np.square(workarr)) + np.square(w)), 1.) 104 | assert np.sum(np.square(workarr)) <= 1. or np.allclose(np.sum(np.square(workarr)), 1.) 105 | return retarr 106 | 107 | 108 | cpdef versor_to_axis_angle_ipr(np.ndarray[np.float32_t, ndim=2] anglearr): 109 | assert anglearr.shape[1] % 3 == 0 110 | cdef np.ndarray[np.float32_t, ndim=1] vec, recvec 111 | cdef np.float32_t comp_s, recrot, vecsqsum 112 | cdef int sample_idx, rot_idx 113 | for sample_idx in range(anglearr.shape[0]): 114 | for rot_idx in range(0, anglearr.shape[1], 3): 115 | vec = anglearr[sample_idx, rot_idx:rot_idx + 3] 116 | vecsqsum = np.sum(np.square(vec)) 117 | assert vecsqsum <= 1. or np.allclose(vecsqsum, 1.) 118 | # Sidestep numerics. 119 | vecsqsum = np.clip(vecsqsum, 0., 1.) 120 | w = np.sqrt(1. - vecsqsum) 121 | #print 'w recovered:', w 122 | angle = 2. * np.arccos(w) 123 | #print 'angle recovered:', angle 124 | s = np.sqrt(1. - w ** 2) 125 | if s < 0.001: 126 | # if s is close to zero then the direction of the axis is not important. 127 | vec[:] = 1., 0., 0. 128 | else: 129 | vec[:] = vec / s * angle 130 | """ 131 | comp_s = np.sqrt(1. - (np.sum(np.square(vec)))) 132 | if comp_s < 0.001: 133 | # if s is close to zero then the direction of the axis is not important. 134 | vec[:] = 1., 0., 0. 135 | else: 136 | recrot = 2. * np.arctan2(np.linalg.norm(vec), comp_s) 137 | recvec = vec / np.sin(recrot / 2.) 138 | vec[:] = recvec / np.linalg.norm(recvec) * recrot 139 | """ 140 | return anglearr 141 | 142 | 143 | cpdef axis_angle_to_versor_ipr(np.ndarray[np.float32_t, ndim=2] anglearr): 144 | assert anglearr.shape[1] % 3 == 0 145 | cdef np.ndarray[np.float32_t, ndim=1] vec, quat, rep 146 | cdef np.float32_t angle 147 | cdef int sample_idx, rot_idx 148 | quat = np.zeros((4,), dtype='float32') 149 | for sample_idx in range(anglearr.shape[0]): 150 | for rot_idx in range(0, anglearr.shape[1], 3): 151 | rep = anglearr[sample_idx, rot_idx:rot_idx + 3] 152 | angle = np.linalg.norm(rep) 153 | #print 'angle orig:', angle 154 | vec = rep / angle 155 | #print 'vec orig:', vec 156 | quat[0] = np.cos(angle / 2.) 157 | quat[1:] = vec * np.sin(angle / 2.) 158 | if quat[0] < 0.: 159 | quat = -quat 160 | rep[:] = quat[1:] 161 | #print 'w orig:', quat[0] 162 | return anglearr 163 | 164 | 165 | cpdef axis_angle_to_euler_angle_ipr(np.ndarray[np.float32_t, ndim=2] anglearr): 166 | assert anglearr.shape[1] % 3 == 0 167 | cdef int startpos 168 | cdef np.float32_t angle 169 | cdef np.ndarray[np.float32_t, ndim=1] rep 170 | cdef np.float32_t x, y, z, s, c, t 171 | cdef np.float32_t heading, attitude, bank 172 | for startpos in range(0, anglearr.size, 3): 173 | rep = anglearr.flat[startpos:startpos + 3] 174 | angle = np.linalg.norm(rep) 175 | x, y, z = rep / angle 176 | s = np.sin(angle) 177 | c = np.cos(angle) 178 | t = 1. - c 179 | if ((x*y*t + z*s) > 0.998): # north pole singularity detected 180 | heading = 2. * np.arctan2(x * np.sin(angle / 2.), 181 | np.cos(angle / 2.)) 182 | attitude = np.pi / 2. 183 | bank = 0. 184 | elif ((x*y*t + z*s) < -0.998): # south pole singularity detected 185 | heading = -2. * np.arctan2(x * np.sin(angle / 2.), 186 | np.cos(angle / 2.)) 187 | attitude = -np.pi / 2. 188 | bank = 0 189 | else: 190 | heading = np.arctan2(y * np.sin(angle) - x * z * (1. - np.cos(angle)), 191 | 1. - (y**2 + z**2) * (1. - np.cos(angle))) 192 | attitude = np.arcsin(x * y * (1. - np.cos(angle)) + z * np.sin(angle)) 193 | bank = np.arctan2(x * np.sin(angle)- y * z * (1. - np.cos(angle)), 194 | 1. - (x**2 + z**2) * (1. - np.cos(angle))) 195 | anglearr.flat[startpos:startpos + 3] = heading, attitude, bank 196 | return anglearr 197 | 198 | 199 | cpdef euler_angle_to_axis_angle_ipr(np.ndarray[np.float32_t, ndim=2] anglearr): 200 | assert anglearr.shape[1] % 3 == 0 201 | cdef int startpos 202 | cdef np.float32_t heading, attitude, bank 203 | cdef np.float32_t c1, c2, c3, s1, s2, s3, c1c2, s1s2, w 204 | cdef np.float32_t x, y, z, angle, norm 205 | for startpos in range(0, anglearr.size, 3): 206 | heading, attitude, bank = anglearr.flat[startpos:startpos + 3] 207 | c1 = np.cos(heading / 2.) 208 | s1 = np.sin(heading / 2.) 209 | c2 = np.cos(attitude / 2.) 210 | s2 = np.sin(attitude / 2.) 211 | c3 = np.cos(bank / 2.) 212 | s3 = np.sin(bank / 2.) 213 | c1c2 = c1 * c2 214 | s1s2 = s1 * s2 215 | w = c1c2 * c3 - s1s2 * s3 216 | x = c1c2 * s3 + s1s2 * c3 217 | y = s1 * c2 * c3 + c1 * s2 * s3 218 | z = c1 * s2 * c3 - s1 * c2 * s3 219 | angle = 2 * np.arccos(w) 220 | norm = x*x+y*y+z*z 221 | if (norm < 0.001): 222 | x = 1. 223 | y, z = 0., 0. 224 | else: 225 | norm = np.sqrt(norm) 226 | x /= norm 227 | y /= norm 228 | z /= norm 229 | anglearr.flat[startpos:startpos + 3] = x * angle, y * angle, z * angle 230 | return anglearr 231 | 232 | 233 | cpdef abs_to_dist(np.ndarray[np.float32_t, ndim=2] absarr, 234 | np.ndarray[np.int_t, ndim=2] combinations): 235 | assert combinations.shape[1] == 2 236 | cdef int sample_idx, feat_comb_idx 237 | cdef np.ndarray[np.float32_t, ndim=2] landmark_pos 238 | cdef np.ndarray[np.int_t, ndim=1] feat_comb 239 | cdef np.ndarray[np.float32_t, ndim=2] retarr 240 | retarr = np.zeros_like(absarr) 241 | for sample_idx in range(absarr.shape[0]): 242 | landmark_pos = absarr[sample_idx].reshape((2, 91)) 243 | for feat_comb_idx, feat_comb in enumerate(combinations): 244 | retarr[sample_idx, 245 | feat_comb_idx * 2:feat_comb_idx*2+2] = \ 246 | landmark_pos[:, feat_comb[0]] - landmark_pos[:, feat_comb[1]] 247 | return retarr 248 | 249 | 250 | class ProjectionError(Exception): 251 | pass 252 | 253 | 254 | cpdef project_to_rot(np.ndarray[np.float32_t, ndim=1] anglearr): 255 | assert anglearr.shape[0] == 9 256 | cdef np.ndarray[np.float32_t, ndim=2] ret_arr 257 | #cdef np.ndarray[np.float32_t, ndim=1] x, y, z 258 | cdef np.float32_t error 259 | ret_arr = polar(anglearr.reshape((3, 3)))[0] 260 | if np.linalg.det(ret_arr) < 0.: 261 | raise ProjectionError("Could not recover rotation, but only a reflection!") 262 | return ret_arr, np.linalg.norm(anglearr.reshape((3, 3)) - ret_arr) 263 | """ 264 | x, y, z = anglearr.copy() 265 | # Determine the error. 266 | error = np.dot(x, y) 267 | # Resolve it. 268 | anglearr[0] = x - (error / 2.) * y 269 | anglearr[1] = y - (error / 2.) * x 270 | anglearr[2] = np.cross(anglearr[0], anglearr[1]) 271 | # Normalize. 272 | #anglearr[0] /= np.linalg.norm(anglearr[0]) 273 | #anglearr[1] /= np.linalg.norm(anglearr[1]) 274 | #anglearr[2] /= np.linalg.norm(anglearr[2]) 275 | #anglearr[0] = 0.5 * (3. - np.dot(anglearr[0], anglearr[0])) * anglearr[0] 276 | #anglearr[1] = 0.5 * (3. - np.dot(anglearr[1], anglearr[1])) * anglearr[1] 277 | #anglearr[2] = 0.5 * (3. - np.dot(anglearr[2], anglearr[2])) * anglearr[2] 278 | return anglearr, error 279 | """ 280 | 281 | cpdef project_to_rot_nofailure_ipr(np.ndarray[np.float32_t, ndim=2] anglearr): 282 | assert anglearr.shape[1] % 9 == 0 283 | cdef int startpos 284 | for startpos in range(0, anglearr.size, 9): 285 | anglearr.flat[startpos:startpos+9] = cv2.Rodrigues( 286 | cv2.Rodrigues(anglearr.flat[startpos:startpos+9].reshape((3, 3)))[0])[0].reshape((9,)) 287 | return anglearr 288 | 289 | 290 | cpdef versor_mean(versors): 291 | assert len(versors) > 0 292 | cdef np.ndarray[np.float32_t, ndim=1] res, firstvers, fullvers 293 | cdef np.ndarray[np.float32_t, ndim=2] workvers 294 | cpdef np.float32_t dnvers, w, w_res, vecsqsum 295 | dnvers = 1. / float(len(versors)) 296 | 297 | vecsqsum = np.sum(np.square(versors[0])) 298 | assert vecsqsum <= 1. or np.allclose(vecsqsum, 1.) 299 | # Sidestep numerics. 300 | vecsqsum = np.clip(vecsqsum, 0., 1.) 301 | w = np.sqrt(1. - vecsqsum) 302 | 303 | firstvers = np.zeros((4,), dtype=np.float32) 304 | firstvers[0] = w 305 | firstvers[1:] = versors[0] 306 | res = np.atleast_2d(versors[0])[0] * dnvers 307 | w_res = w * dnvers 308 | fullvers = np.zeros((4,), dtype=np.float32) 309 | for workvers in versors: 310 | vecsqsum = np.sum(np.square(workvers)) 311 | assert vecsqsum <= 1. or np.allclose(vecsqsum, 1.) 312 | # Sidestep numerics. 313 | vecsqsum = np.clip(vecsqsum, 0., 1.) 314 | w = np.sqrt(1. - vecsqsum) 315 | fullvers[0] = w 316 | fullvers[1:] = workvers 317 | if np.dot(fullvers, firstvers) < 0.: 318 | fullvers *= -1. 319 | res += np.atleast_2d(workvers)[0] * dnvers 320 | w_res += w * dnvers 321 | # Normalize. 322 | fullvers[0] = w_res 323 | fullvers[1:] = res 324 | fullvers /= np.linalg.norm(fullvers) 325 | if fullvers[0] < 0.: 326 | fullvers *= -1. 327 | return np.atleast_2d(fullvers[1:]) 328 | -------------------------------------------------------------------------------- /direct3d/fit_forest.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python2 2 | """Forest training.""" 3 | # pylint: disable=wrong-import-order, redefined-outer-name, line-too-long, invalid-name 4 | import os 5 | import sys 6 | import os.path as path 7 | import h5py 8 | import numpy as np 9 | import logging 10 | import click 11 | import joblib 12 | import pyximport; pyximport.install() # pylint: disable=multiple-statements 13 | from conversions import ( # pylint: disable=import-error 14 | axis_angle_to_matrix) 15 | from clustertools.config import available_cpu_count 16 | from clustertools.log import LOGFORMAT 17 | 18 | sys.path.insert(0, path.join(path.dirname(__file__), 19 | '..')) 20 | 21 | try: 22 | from up_tools.model import landmarks_91 23 | from config import DIRECT3D_DATA_FP 24 | except: 25 | print("Make sure this script is in the up/direct3d/ folder.") 26 | 27 | LOGGER = logging.getLogger(__name__) 28 | OUT_DIR = path.join(path.dirname(__file__), 29 | '..', 'models', '2dto3d', 30 | 'separate_regressors') 31 | if not path.exists(OUT_DIR): 32 | os.makedirs(OUT_DIR) 33 | # Inputs: 34 | # * 2x91 absolute x, y coordinates 35 | # Prediction targets: 36 | # * betas: 10 37 | # * pose: 72 38 | # * trans: 3 39 | # * rt: 3 40 | # * t: 3 41 | # * f: 1 42 | 43 | torso_ids = [landmarks_91.rshoulder, 44 | landmarks_91.rshoulder_back, 45 | landmarks_91.rshoulder_front, 46 | landmarks_91.rshoulder_back, 47 | 48 | landmarks_91.lshoulder, 49 | landmarks_91.lshoulder_back, 50 | landmarks_91.lshoulder_front, 51 | landmarks_91.lshoulder_top, 52 | 53 | landmarks_91.shoulderblade_center, 54 | landmarks_91.solar_plexus, 55 | landmarks_91.lpapilla, 56 | landmarks_91.rpapilla, 57 | 58 | landmarks_91.belly_button, 59 | landmarks_91.rwaist, 60 | landmarks_91.lwaist, 61 | landmarks_91.waist_back, 62 | 63 | landmarks_91.lhip, 64 | landmarks_91.lhip_back, 65 | landmarks_91.lhip_front, 66 | landmarks_91.lhip_outer, 67 | landmarks_91.rhip, 68 | landmarks_91.rhip_back, 69 | landmarks_91.rhip_front, 70 | landmarks_91.rhip_outer] 71 | 72 | head_ids = [landmarks_91.head_top, 73 | landmarks_91.head_back, 74 | landmarks_91.nose, 75 | landmarks_91.lear, 76 | landmarks_91.rear] 77 | 78 | larm_ids = [landmarks_91.luarm_inner, 79 | landmarks_91.luarm_outer, 80 | landmarks_91.lelbow, 81 | landmarks_91.lelbow_bottom, 82 | landmarks_91.lelbow_inner, 83 | landmarks_91.lelbow_outer, 84 | landmarks_91.lelbow_top, 85 | landmarks_91.llarm_lower, 86 | landmarks_91.llarm_upper, 87 | landmarks_91.lwrist] 88 | 89 | rarm_ids = [landmarks_91.ruarm_inner, 90 | landmarks_91.ruarm_outer, 91 | landmarks_91.relbow, 92 | landmarks_91.relbow_bottom, 93 | landmarks_91.relbow_inner, 94 | landmarks_91.relbow_outer, 95 | landmarks_91.relbow_top, 96 | landmarks_91.rlarm_lower, 97 | landmarks_91.rlarm_upper, 98 | landmarks_91.rwrist] 99 | 100 | lleg_ids = [landmarks_91.luleg_back, 101 | landmarks_91.luleg_front, 102 | landmarks_91.luleg_outer, 103 | landmarks_91.luleg_inner, 104 | landmarks_91.lknee, 105 | landmarks_91.llleg_back, 106 | landmarks_91.llleg_front, 107 | landmarks_91.llleg_outer, 108 | landmarks_91.llleg_inner, 109 | landmarks_91.lankle, 110 | landmarks_91.lheel, 111 | landmarks_91.lankle_inner, 112 | landmarks_91.lankle_outer, 113 | landmarks_91.lbigtoe] 114 | 115 | rleg_ids = [landmarks_91.ruleg_back, 116 | landmarks_91.ruleg_front, 117 | landmarks_91.ruleg_outer, 118 | landmarks_91.ruleg_inner, 119 | landmarks_91.rknee, 120 | landmarks_91.rlleg_back, 121 | landmarks_91.rlleg_front, 122 | landmarks_91.rlleg_outer, 123 | landmarks_91.rlleg_inner, 124 | landmarks_91.rankle, 125 | landmarks_91.rheel, 126 | landmarks_91.rankle_inner, 127 | landmarks_91.rankle_outer, 128 | landmarks_91.rbigtoe] 129 | 130 | def create_featseltuple(ids): 131 | """Add the flat y coordinate to an index list.""" 132 | newlist = [] 133 | for part_id in ids: 134 | newlist.extend([part_id, part_id + 91]) 135 | return tuple(sorted(newlist)) 136 | 137 | lmset_to_use = { 138 | (0, 10): torso_ids, # shape 139 | (10, 13): torso_ids, # root joint 140 | (13, 16): torso_ids + lleg_ids, # luleg 141 | (16, 19): torso_ids + rleg_ids, # ruleg 142 | (19, 22): torso_ids, # spine 143 | (22, 25): lleg_ids, # llleg 144 | (25, 28): rleg_ids, # rlleg 145 | (28, 31): torso_ids, # spine1 146 | (31, 34): lleg_ids, # lfoot 147 | (34, 37): rleg_ids, # rfoot 148 | (37, 40): torso_ids, # spine2 149 | (40, 43): lleg_ids, # ltoes 150 | (43, 46): rleg_ids, # rtoes 151 | (46, 49): torso_ids + head_ids, # neck 152 | (49, 52): torso_ids + larm_ids, # lshoulder 153 | (52, 55): torso_ids + rarm_ids, # rshoulder 154 | (55, 58): torso_ids + head_ids, # head 155 | (58, 61): torso_ids + larm_ids, # luarm 156 | (61, 64): torso_ids + rarm_ids, # ruarm 157 | (64, 67): larm_ids, # llarm 158 | (67, 70): rarm_ids, # rlarm 159 | (70, 73): larm_ids, # lhand 160 | (73, 76): rarm_ids, # rhand 161 | (76, 79): larm_ids, # lfingers 162 | (79, 82): rarm_ids, # rfingers 163 | (82, 85): torso_ids # depth 164 | } 165 | 166 | 167 | 168 | def normalize_axis_angle(anglevec): 169 | """Normalize angle periodicity.""" 170 | assert len(anglevec) % 3 == 0 171 | for startpos in range(0, len(anglevec), 3): 172 | rep = anglevec[startpos:startpos + 3] 173 | angle = np.linalg.norm(rep) 174 | angle_norm = np.fmod(angle, np.pi) 175 | anglevec[startpos:startpos + 3] = rep / angle * angle_norm 176 | 177 | def preprocess(dta_arr): 178 | """Make the coordinates relative to mean position, apply the modulo operator to pose.""" 179 | for dta_idx in range(dta_arr.shape[0]): 180 | pose = dta_arr[dta_idx, :2*91].reshape((2, 91)) 181 | mean = np.mean(pose, axis=1) 182 | dta_arr[dta_idx, :2*91] = (pose.T - mean + 513. / 2.).T.flat 183 | normalize_axis_angle(dta_arr[dta_idx, 2*91+10:2*91+10+72]) 184 | 185 | def get_data(prefix, part_rangestart, finalize, debug_run): # pylint: disable=too-many-branches 186 | """Get the data.""" 187 | rangestart = part_rangestart 188 | rangeend = 10 if part_rangestart == 0 else part_rangestart + 3 189 | train_f = h5py.File(path.join( 190 | DIRECT3D_DATA_FP, 191 | '91', '500', prefix, 'train.hdf5')) 192 | train_dset = train_f['2dto3d'] 193 | if debug_run: 194 | train_dta = np.array(train_dset[:10000]) 195 | else: 196 | train_dta = np.array(train_dset) 197 | preprocess(train_dta) 198 | #add_noise(train_dta) 199 | val_f = h5py.File(path.join( 200 | DIRECT3D_DATA_FP, 201 | '91', '500', prefix, 'val.hdf5')) 202 | val_dset = val_f['2dto3d'] 203 | if debug_run: 204 | val_dta = np.array(val_dset[:10]) 205 | else: 206 | val_dta = np.array(val_dset) 207 | preprocess(val_dta) 208 | if finalize: 209 | train_dta = np.vstack((train_dta, val_dta)) 210 | val_f = h5py.File(path.join( 211 | DIRECT3D_DATA_FP, 212 | '91', '500', prefix, 'test.hdf5')) 213 | val_dset = val_f['2dto3d'] 214 | if debug_run: 215 | val_dta = np.array(val_dset[:10]) 216 | else: 217 | val_dta = np.array(val_dset) 218 | preprocess(val_dta) 219 | train_annot = train_dta[:, 182+rangestart:182+rangeend] 220 | val_annot = val_dta[:, 182+rangestart:182+rangeend] 221 | rel_ids = create_featseltuple(lmset_to_use[(rangestart, rangeend)]) 222 | train_dta = train_dta[:, rel_ids] 223 | val_dta = val_dta[:, rel_ids] 224 | if rangestart > 0 and rangestart < 82: 225 | train_annot = axis_angle_to_matrix(train_annot) 226 | val_annot = axis_angle_to_matrix(val_annot) 227 | return train_dta, train_annot, val_dta, val_annot 228 | 229 | def sqdiff(rnge, val_dta, val_results, addoffs=0): 230 | """Error measure robust to angle orientations and mirroring.""" 231 | orig_ids = tuple(np.array(rnge) + 182 + addoffs) 232 | val_ids = tuple(np.array(rnge)) 233 | assert len(orig_ids) == len(val_ids) 234 | assert len(orig_ids) % 3 == 0 235 | diffs = [] 236 | for sample_idx in range(val_dta.shape[0]): 237 | for rot_idx in range(0, len(orig_ids), 3): 238 | plaindiff = np.linalg.norm(val_dta[sample_idx, orig_ids[rot_idx:rot_idx+3]] - 239 | val_results[sample_idx, val_ids[rot_idx:rot_idx+3]]) 240 | mirrdiff = np.linalg.norm(val_dta[sample_idx, orig_ids[rot_idx:rot_idx+3]] + 241 | val_results[sample_idx, val_ids[rot_idx:rot_idx+3]]) 242 | diffs.append(min(plaindiff, mirrdiff)) 243 | return np.mean(diffs) 244 | 245 | @click.command() 246 | @click.argument('train_prefix', type=click.STRING) 247 | @click.argument('part_rangestart', type=click.INT) 248 | @click.option('--finalize', type=click.BOOL, default=False, is_flag=True, 249 | help='Train on train+val, test on test.') 250 | @click.option('--debug_run', type=click.BOOL, default=False, is_flag=True, 251 | help='Use only a small fraction of data for testing.') 252 | def cli(train_prefix, part_rangestart, # pylint: disable=too-many-branches, too-many-locals, too-many-statements, too-many-arguments 253 | finalize=False, debug_run=False): 254 | """Run a RotatingTree experiment.""" 255 | rangestart = part_rangestart 256 | pref = 'forest' 257 | pref += '_' + str(part_rangestart) 258 | if finalize: 259 | pref += '_final' 260 | if debug_run: 261 | pref += '_debug' 262 | out_fp = path.join(OUT_DIR, pref + '.z') 263 | LOGGER.info("Running for configuration `%s`.", out_fp) 264 | LOGGER.info("Loading data...") 265 | train_dta, train_annot, val_dta, val_annot = get_data( # pylint: disable=unused-variable 266 | train_prefix, part_rangestart, finalize, debug_run) 267 | # Checks. 268 | if rangestart > 0 and rangestart < 82: 269 | # Rotation matrices. 270 | assert train_annot.max() <= 1. 271 | assert train_annot.min() >= -1. 272 | assert val_annot.max() <= 1. 273 | assert val_annot.min() >= -1. 274 | import sklearn.ensemble 275 | rf = sklearn.ensemble.RandomForestRegressor(n_jobs=available_cpu_count()) 276 | LOGGER.info("Fitting...") 277 | rf.fit(train_dta, train_annot) 278 | LOGGER.info("Writing results...") 279 | joblib.dump(rf, out_fp, compress=True) 280 | LOGGER.info("Done.") 281 | 282 | 283 | if __name__ == '__main__': 284 | logging.basicConfig(level=logging.INFO, format=LOGFORMAT) 285 | cli() # pylint: disable=no-value-for-parameter 286 | -------------------------------------------------------------------------------- /direct3d/run_partforest_training.sh: -------------------------------------------------------------------------------- 1 | #!/usr/bin/zsh 2 | 3 | trap "exit" INT 4 | 5 | for rangestart in 0 $(seq 10 3 82); do 6 | echo Running training for ${rangestart} 7 | ./fit_forest.py $1 ${rangestart} --finalize 8 | done 9 | -------------------------------------------------------------------------------- /models/3D/.AppleDouble/.Parent: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/classner/up/7ec281c62afae764fa56a649f41a02487ea18c58/models/3D/.AppleDouble/.Parent -------------------------------------------------------------------------------- /models/3D/head_vid_heva.pkl: -------------------------------------------------------------------------------- 1 | (cnumpy.core.multiarray 2 | scalar 3 | p1 4 | (cnumpy 5 | dtype 6 | p2 7 | (S'i8' 8 | I0 9 | I1 10 | tRp3 11 | (I3 12 | S'<' 13 | NNNI-1 14 | I-1 15 | I0 16 | tbS'\x9b\x01\x00\x00\x00\x00\x00\x00' 17 | tRp4 18 | tp5 19 | . -------------------------------------------------------------------------------- /models/3D/mask_filled.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/classner/up/7ec281c62afae764fa56a649f41a02487ea18c58/models/3D/mask_filled.png -------------------------------------------------------------------------------- /models/3D/mask_filled_uniform.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/classner/up/7ec281c62afae764fa56a649f41a02487ea18c58/models/3D/mask_filled_uniform.png -------------------------------------------------------------------------------- /models/3D/regressors_locked_normalized_hybrid.npz: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/classner/up/7ec281c62afae764fa56a649f41a02487ea18c58/models/3D/regressors_locked_normalized_hybrid.npz -------------------------------------------------------------------------------- /models/3D/template-bodyparts.ply: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/classner/up/7ec281c62afae764fa56a649f41a02487ea18c58/models/3D/template-bodyparts.ply -------------------------------------------------------------------------------- /models/3D/template.ply: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/classner/up/7ec281c62afae764fa56a649f41a02487ea18c58/models/3D/template.ply -------------------------------------------------------------------------------- /models/pose/landmarks.pkl: -------------------------------------------------------------------------------- 1 | ccollections 2 | OrderedDict 3 | p1 4 | ((lp2 5 | (lp3 6 | S'llleg_outer' 7 | p4 8 | aI1097 9 | aa(lp5 10 | S'lhip_outer' 11 | p6 12 | aI1454 13 | aa(lp7 14 | S'rmcorner' 15 | p8 16 | aI3841 17 | aa(lp9 18 | S'lknee_inner' 19 | p10 20 | aI1016 21 | aa(lp11 22 | S'ruarm_outer' 23 | p12 24 | aI4862 25 | aa(lp13 26 | S'lankle_outer' 27 | p14 28 | aI3328 29 | aa(lp15 30 | S'rknee_outer' 31 | p16 32 | aI4495 33 | aa(lp17 34 | S'lshoulder_back' 35 | p18 36 | aI734 37 | aa(lp19 38 | S'relbow_inner' 39 | p20 40 | aI5121 41 | aa(lp21 42 | S'rpapilla' 43 | p22 44 | aI6489 45 | aa(lp23 46 | S'lelbow_top' 47 | p24 48 | aI1618 49 | aa(lp25 50 | S'rhip_outer' 51 | p26 52 | aI4920 53 | aa(lp27 54 | S'rbigtoe' 55 | p28 56 | aI6618 57 | aa(lp29 58 | S'lknee_back' 59 | p30 60 | aI1002 61 | aa(lp31 62 | S'lknee_outer' 63 | p32 64 | aI1010 65 | aa(lp33 66 | S'waist_back' 67 | p34 68 | aI3021 69 | aa(lp35 70 | S'rknee_inner' 71 | p36 72 | aI4500 73 | aa(lp37 74 | S'ruleg_front' 75 | p38 76 | aI4431 77 | aa(lp39 78 | S'luarm_inner' 79 | p40 80 | aI627 81 | aa(lp41 82 | S'ruleg_outer' 83 | p42 84 | aI4459 85 | aa(lp43 86 | S'rankle_outer' 87 | p44 88 | aI6728 89 | aa(lp45 90 | S'rhip_back' 91 | p46 92 | aI6539 93 | aa(lp47 94 | S'luleg_inner' 95 | p48 96 | aI949 97 | aa(lp49 98 | S'lknee_front' 99 | p50 100 | aI1043 101 | aa(lp51 102 | S'lbigtoe' 103 | p52 104 | aI3217 105 | aa(lp53 106 | S'belly_button' 107 | p54 108 | aI3500 109 | aa(lp55 110 | S'lpapilla' 111 | p56 112 | aI3042 113 | aa(lp57 114 | S'rlleg_back' 115 | p58 116 | aI4937 117 | aa(lp59 118 | S'ruleg_inner' 119 | p60 120 | aI4435 121 | aa(lp61 122 | S'ruarm_inner' 123 | p62 124 | aI6281 125 | aa(lp63 126 | S'lear' 127 | p64 128 | aI516 129 | aa(lp65 130 | S'lshoulder_front' 131 | p66 132 | aI1873 133 | aa(lp67 134 | S'throat' 135 | p68 136 | aI3068 137 | aa(lp69 138 | S'llleg_front' 139 | p70 140 | aI1103 141 | aa(lp71 142 | S'lwaist' 143 | p72 144 | aI676 145 | aa(lp73 146 | S'luleg_front' 147 | p74 148 | aI946 149 | aa(lp75 150 | S'rknee_front' 151 | p76 152 | aI4529 153 | aa(lp77 154 | S'luleg_outer' 155 | p78 156 | aI973 157 | aa(lp79 158 | S'rlarm_upper' 159 | p80 160 | aI5016 161 | aa(lp81 162 | S'nose' 163 | p82 164 | aI331 165 | aa(lp83 166 | S'rshoulder_back' 167 | p84 168 | aI4222 169 | aa(lp85 170 | S'rshoulder_top' 171 | p86 172 | aI6470 173 | aa(lp87 174 | S'rlleg_inner' 175 | p88 176 | aI4574 177 | aa(lp89 178 | S'lmcorner' 179 | p90 180 | aI337 181 | aa(lp91 182 | S'head_top' 183 | p92 184 | aI411 185 | aa(lp93 186 | S'lhip_back' 187 | p94 188 | aI3117 189 | aa(lp95 190 | S'lelbow_outer' 191 | p96 192 | aI1703 193 | aa(lp97 194 | S'rhip_front' 195 | p98 196 | aI4399 197 | aa(lp99 198 | S'llarm_upper' 199 | p100 200 | aI1971 201 | aa(lp101 202 | S'shoulderblade_center' 203 | p102 204 | aI3028 205 | aa(lp103 206 | S'ruleg_back' 207 | p104 208 | aI4933 209 | aa(lp105 210 | S'lelbow_inner' 211 | p106 212 | aI1650 213 | aa(lp107 214 | S'rknee_back' 215 | p108 216 | aI4486 217 | aa(lp109 218 | S'rshoulder_front' 219 | p110 220 | aI4798 221 | aa(lp111 222 | S'rheel' 223 | p112 224 | aI6786 225 | aa(lp113 226 | S'luarm_outer' 227 | p114 228 | aI1389 229 | aa(lp115 230 | S'relbow_outer' 231 | p116 232 | aI5129 233 | aa(lp117 234 | S'llarm_lower' 235 | p118 236 | aI1953 237 | aa(lp119 238 | S'rear' 239 | p120 240 | aI3998 241 | aa(lp121 242 | S'lhip_front' 243 | p122 244 | aI915 245 | aa(lp123 246 | S'reye' 247 | p124 248 | aI6262 249 | aa(lp125 250 | S'luleg_back' 251 | p126 252 | aI1459 253 | aa(lp127 254 | S'neck' 255 | p128 256 | aI829 257 | aa(lp129 258 | S'rankle_inner' 259 | p130 260 | aI6832 261 | aa(lp131 262 | S'relbow_bottom' 263 | p132 264 | aI5131 265 | aa(lp133 266 | S'lankle_inner' 267 | p134 268 | aI3432 269 | aa(lp135 270 | S'head_back' 271 | p136 272 | aI457 273 | aa(lp137 274 | S'leye' 275 | p138 276 | aI2802 277 | aa(lp139 278 | S'rlleg_outer' 279 | p140 280 | aI4581 281 | aa(lp141 282 | S'rlleg_front' 283 | p142 284 | aI4589 285 | aa(lp143 286 | S'relbow_top' 287 | p144 288 | aI5135 289 | aa(lp145 290 | S'lelbow_bottom' 291 | p146 292 | aI1724 293 | aa(lp147 294 | S'llleg_back' 295 | p148 296 | aI1464 297 | aa(lp149 298 | S'solar_plexus' 299 | p150 300 | aI1329 301 | aa(lp151 302 | S'lheel' 303 | p152 304 | aI3387 305 | aa(lp153 306 | S'lshoulder_top' 307 | p154 308 | aI1822 309 | aa(lp155 310 | S'llleg_inner' 311 | p156 312 | aI1088 313 | aa(lp157 314 | S'rwaist' 315 | p158 316 | aI4317 317 | aa(lp159 318 | S'rlarm_lower' 319 | p160 320 | aI5021 321 | aatRp161 322 | . -------------------------------------------------------------------------------- /patches/deeplab.txt: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/classner/up/7ec281c62afae764fa56a649f41a02487ea18c58/patches/deeplab.txt -------------------------------------------------------------------------------- /pose/evaluate_pose.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python2 2 | """Get evaluation results for stored landmarks.""" 3 | # pylint: disable=invalid-name, wrong-import-order 4 | from __future__ import print_function 5 | import matplotlib 6 | matplotlib.use("Agg") 7 | import matplotlib.pyplot as plt 8 | from collections import OrderedDict 9 | from os import path 10 | import logging 11 | 12 | import numpy as np 13 | import click 14 | 15 | from clustertools.config import available_cpu_count # pylint: disable=unused-import 16 | 17 | 18 | LOGGER = logging.getLogger(__name__) 19 | 20 | 21 | def _getDistPCK(pred, gt, norm_lm_ids): 22 | """ 23 | Calculate the pck distance for all given poses. 24 | 25 | norm_lm_ids: Use the distance of these landmarks for normalization. Usually 26 | lshoulder and rhip. 27 | """ 28 | assert pred.ndim == gt.ndim 29 | assert pred.ndim == 3 30 | assert pred.shape[0] >= 2 31 | assert np.all(pred.shape[1:] == gt.shape[1:]) 32 | dist = np.empty((1, pred.shape[1], pred.shape[2])) 33 | 34 | for imgidx in range(pred.shape[2]): 35 | # torso diameter 36 | refDist = np.linalg.norm(gt[:2, norm_lm_ids[0], imgidx] - 37 | gt[:2, norm_lm_ids[1], imgidx]) 38 | # distance to gt joints 39 | dist[0, :, imgidx] =\ 40 | np.sqrt( 41 | np.sum( 42 | np.power(pred[:2, :, imgidx] - 43 | gt[:2, :, imgidx], 2), 44 | axis=0)) / refDist 45 | return dist 46 | 47 | def _computePCK(dist, rnge, mask): 48 | """Compute PCK values for given joint distances and a range.""" 49 | pck = np.zeros((len(rnge), dist.shape[1] + 1)) 50 | for joint_idx in range(dist.shape[1]): 51 | # compute PCK for each threshold 52 | evaluation_basis = dist[0, 53 | joint_idx, 54 | np.where(mask[joint_idx, :] > 0)] 55 | for k, rngval in enumerate(rnge): 56 | pck[k, joint_idx] = 100. *\ 57 | np.mean(evaluation_basis.flat <= rngval) 58 | # compute average PCK 59 | for k in range(len(rnge)): 60 | pck[k, -1] = np.mean(pck[k, :-1]) 61 | return pck 62 | 63 | def _area_under_curve(xpts, ypts): 64 | """Calculate the AUC.""" 65 | a = np.min(xpts) 66 | b = np.max(xpts) 67 | # remove duplicate points 68 | _, I = np.unique(xpts, return_index=True) # pylint: disable=W0632 69 | xpts = xpts[I] 70 | ypts = ypts[I] 71 | assert np.all(np.diff(xpts) > 0) 72 | if len(xpts) < 2: 73 | return np.NAN 74 | from scipy import integrate 75 | myfun = lambda x: np.interp(x, xpts, ypts) 76 | auc = integrate.quad(myfun, a, b)[0] 77 | return auc 78 | 79 | def PCK(poses, # pylint: disable=too-many-arguments 80 | annotations, 81 | norm_lm_ids, 82 | plot=False, 83 | rnge_max=0.2, 84 | print_res=False, 85 | using_joint_index=-1): 86 | r""" 87 | Implementation of the PCK measure. 88 | 89 | As defined in Sapp&Taskar, CVPR 2013. 90 | Torso height: ||left_shoulder - right hip||. 91 | Validated to give the same results as Pishchulin et al. 92 | 93 | Parameters 94 | ========== 95 | 96 | :param poses: np.ndarray((M>2, L, N)). 97 | M are the coordinates, L joints, N is 98 | the number of samples. 99 | 100 | :param annotations: np.ndarray((O>2, L, N)). 101 | The annotated poses. The coordinate order must match the pose coordinate 102 | order. 103 | 104 | :param norm_lm_ids: 2-tuple(int). 105 | The indices of the two landmarks to use for normalization. 106 | 107 | :param plot: bool. 108 | Whether to directly show a plot of the results. 109 | 110 | :param rnge_max: float. 111 | Up to which point to calculate the AUC. 112 | 113 | :param print_res: bool. 114 | Whether to print a summary of results. 115 | 116 | :param using_joint_index: int. 117 | If > 1, specifies a column in the pose array, which indicates binary 118 | whether to take the given joint into account or not. 119 | """ 120 | assert using_joint_index == -1 or using_joint_index > 1 121 | assert len(norm_lm_ids) == 2 122 | rnge = np.arange(0., rnge_max + 0.001, 0.01) 123 | dist = _getDistPCK(poses, annotations, norm_lm_ids) 124 | # compute PCK 125 | if using_joint_index > 1: 126 | mask = poses[using_joint_index, :, :] > 0 127 | else: 128 | mask = np.ones((poses.shape[1], poses.shape[2])) 129 | pck = _computePCK(dist, rnge, mask) 130 | auc = _area_under_curve(rnge / rnge.max(), pck[:, -1]) 131 | if plot: 132 | plt.plot(rnge, 133 | pck[:, -1], 134 | label='PCK', 135 | linewidth=2) 136 | plt.xlim(0., 0.2) 137 | plt.xticks(np.arange(0, rnge_max + 0.01, 0.02)) 138 | plt.yticks(np.arange(0, 101., 10.)) 139 | plt.ylabel('Detection rate, %') 140 | plt.xlabel('Normalized distance') 141 | plt.grid() 142 | legend = plt.legend(loc=4) 143 | legend.get_frame().set_facecolor('white') 144 | plt.show() 145 | # plot(range,pck(:,end),'color',p.colorName, 146 | #'LineStyle','-','LineWidth',3); 147 | if print_res: 148 | # pylint: disable=superfluous-parens 149 | print("AUC: {}.".format(auc)) 150 | print("@0.2: {}.".format(pck[np.argmax(rnge > 0.2) - 1, -1])) 151 | return rnge, pck, auc 152 | 153 | 154 | @click.command() 155 | @click.argument("image_list_file", type=click.Path(dir_okay=False, readable=True)) 156 | @click.argument("scale_fp", type=click.Path(dir_okay=False, readable=True)) 157 | @click.argument("result_label_folder", type=click.Path(file_okay=False, readable=True)) 158 | @click.argument("n_labels", type=click.INT) 159 | def main(image_list_file, # pylint: disable=too-many-locals, too-many-statements, too-many-branches, too-many-arguments 160 | scale_fp, 161 | result_label_folder, 162 | n_labels): 163 | """Perform the evaluation for previously written result landmarks.""" 164 | LOGGER.info("Evaluating landmarks in folder `%s`.", result_label_folder) 165 | LOGGER.info("Reading image information...") 166 | with open(image_list_file, 'r') as inf: 167 | image_list_lines = inf.readlines() 168 | with open(scale_fp, 'r') as inf: 169 | scale_lines = inf.readlines() 170 | all_scales = dict((line.split(" ")[0].strip(), float(line.split(" ")[1].strip())) 171 | for line in scale_lines) 172 | lm_annots = OrderedDict() 173 | read_lms = None 174 | for line in image_list_lines: 175 | if line.startswith('#'): 176 | if read_lms is not None: 177 | lm_annots[imname] = read_lms[:] # pylint: disable=used-before-assignment, unsubscriptable-object 178 | image_started = True 179 | read_lms = [] 180 | elif image_started: 181 | imname = line.strip() 182 | image_started = False 183 | size_spec = 0 184 | elif size_spec < 3: 185 | size_spec += 1 186 | continue 187 | elif size_spec == 3: 188 | num_lms = int(line.strip()) 189 | assert num_lms == n_labels 190 | size_spec += 1 191 | else: 192 | read_lms.append((int(line.split(" ")[1].strip()), 193 | int(line.split(" ")[2].strip()))) 194 | scales = [all_scales[path.basename(imname)] for imname in lm_annots.keys()] 195 | annots = np.array(lm_annots.values()).transpose((2, 1, 0)).astype('float') 196 | for annot_idx in range(annots.shape[2]): 197 | annots[:2, :, annot_idx] /= scales[annot_idx] 198 | LOGGER.info("Loading results...") 199 | lm_positions = [] 200 | for imgname, scale in zip(lm_annots.keys(), scales): 201 | result_file = path.join(result_label_folder, 202 | path.basename(imgname) + '.npy') 203 | lm_positions.append(np.load(result_file) / scale) 204 | lm_positions = np.array(lm_positions).transpose((1, 2, 0)) 205 | LOGGER.info("Evaluating...") 206 | if lm_positions.shape[1] == 91: 207 | from model import landmarks_91 208 | rnge, pck, auc = PCK(lm_positions, annots, 209 | [landmarks_91.lshoulder, # pylint: disable=no-member 210 | landmarks_91.rhip], # pylint: disable=no-member 211 | print_res=False, 212 | plot=False) 213 | else: 214 | # Assume LSP model. 215 | rnge, pck, auc = PCK(lm_positions, annots, 216 | (9, 2), 217 | print_res=False, 218 | plot=False) 219 | # Create the plot. 220 | plt.figure(figsize=(7, 7)) 221 | plt.plot(rnge, 222 | pck[:, -1], 223 | label='PCK', 224 | linewidth=2) 225 | plt.xlim(0., 0.2) 226 | plt.xticks(np.arange(0, 0.2 + 0.01, 0.02)) 227 | plt.yticks(np.arange(0, 101., 10.)) 228 | plt.ylabel('Detection rate, %') 229 | plt.xlabel('Normalized distance') 230 | plt.grid() 231 | legend = plt.legend(loc=4) 232 | legend.get_frame().set_facecolor('white') 233 | plt.savefig(path.join(result_label_folder, 'pck.png')) 234 | # plot(range,pck(:,end),'color',p.colorName, 235 | #'LineStyle','-','LineWidth',3); 236 | LOGGER.info("AUC: %f.", auc) 237 | LOGGER.info("@0.2: %f.", pck[np.argmax(rnge > 0.2) - 1, -1]) 238 | LOGGER.info("Per-part information (PCK@0.2):") 239 | if lm_positions.shape[1] == 91: 240 | for lmid, lmname in landmarks_91.reverse_mapping.items(): # pylint: disable=no-member 241 | LOGGER.info("%s %f", lmname, pck[np.argmax(rnge > 0.2) - 1, lmid]) 242 | else: 243 | from PoseKit.model import joints_lsp 244 | valsatp2 = [] 245 | for lmid, lmname in joints_lsp.reverse_mapping.items(): # pylint: disable=no-member 246 | LOGGER.info("%s %f", lmname, pck[np.argmax(rnge > 0.2) - 1, lmid]) 247 | valsatp2.append(pck[np.argmax(rnge > 0.2) - 1, lmid]) 248 | LOGGER.info("PCK@0.2 wo neck and head: %f.", np.mean(valsatp2[:-2])) 249 | LOGGER.info("Done.") 250 | 251 | 252 | if __name__ == '__main__': 253 | logging.basicConfig(level=logging.INFO) 254 | main() # pylint: disable=no-value-for-parameter 255 | -------------------------------------------------------------------------------- /pose/run.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | 3 | trap "exit" INT 4 | 5 | CAFFE_BUILD_FOLDER=$(../config.py DEEPERCUT_CNN_BUILD_FP) 6 | CAFFE_BIN=${CAFFE_DIR}/${CAFFE_BUILD_FOLDER}/install/bin/caffe 7 | CAFFE_PYTHON=${CAFFE_DIR}/${CAFFE_BUILD_FOLDER}/install/python 8 | EXP=training 9 | if [ "${EXP}" = "training" ]; then 10 | echo "Running pose" 11 | DATA_ROOT=$(../config.py POSE_DATA_FP) 12 | else 13 | echo "Wrong exp name" 14 | exit 1 15 | fi 16 | 17 | ## Specify which model to train 18 | ########### segnpose ################ 19 | NET_ID=$2 20 | TRAIN_SET_SUFFIX=$3 21 | DEV_ID=0 22 | RUN_TRAIN=`grep -q train <<<$1 && echo 1 || echo 0` 23 | RUN_TEST=`grep -q test <<<$1 && echo 1 || echo 0` 24 | RUN_EVALUATION=`grep -q evaluate <<<$1 && echo 1 || echo 0` 25 | RUN_TRAIN2=`grep -q trfull <<<$1 && echo 1 || echo 0` 26 | RUN_TEST2=`grep -q tefull <<<$1 && echo 1 || echo 0` 27 | RUN_EVALUATION2=`grep -q evfull <<<$1 && echo 1 || echo 0` 28 | RUN_TEST_PLAIN=`grep -q teplain <<<$1 && echo 1 || echo 0` 29 | RUN_EVALUATION_PLAIN=`grep -q evplain <<<$1 && echo 1 || echo 0` 30 | 31 | ##### 32 | ## Create dirs 33 | CONFIG_DIR=${EXP}/config/${NET_ID} 34 | MODEL_DIR=${EXP}/model/${NET_ID} 35 | mkdir -p ${MODEL_DIR} 36 | LOG_DIR=${EXP}/log/${NET_ID} 37 | mkdir -p ${LOG_DIR} 38 | export GLOG_log_dir=${LOG_DIR} 39 | export STAT_DUMMY_FP=$(realpath tools) 40 | 41 | ## Training #1 (on train_aug) 42 | if [ ${RUN_TRAIN} -eq 1 ]; then 43 | export GLOG_minloglevel=0 44 | LIST_DIR=${EXP}/list 45 | NUM_LABELS=`cat ${EXP}/config/${NET_ID}/n_classes.txt` 46 | ((NUM_COORDS=NUM_LABELS*2)) 47 | TARGET_PERSON_SIZE=`cat ${EXP}/config/${NET_ID}/target_person_size.txt` 48 | TRAIN_SET=train_${NUM_LABELS}_${TARGET_PERSON_SIZE}_${TRAIN_SET_SUFFIX} 49 | MODEL=${EXP}/config/${NET_ID}/init.caffemodel 50 | echo Training net ${EXP}/${NET_ID} 51 | for pname in train solver; do 52 | sed "$(eval echo $(cat sub.sed))" \ 53 | ${CONFIG_DIR}/${pname}.prototxt > ${CONFIG_DIR}/${pname}_${TRAIN_SET}.prototxt 54 | done 55 | CMD="${CAFFE_BIN} train \ 56 | --solver=${CONFIG_DIR}/solver_${TRAIN_SET}.prototxt \ 57 | --gpu=${DEV_ID}" 58 | if [ -f ${MODEL} ]; then 59 | CMD="${CMD} --weights=${MODEL}" 60 | fi 61 | echo Running ${CMD} && ${CMD} 62 | fi 63 | 64 | ## Test #1 specification (on val or test) 65 | if [ ${RUN_TEST} -eq 1 ]; then 66 | export GLOG_minloglevel=1 67 | for TEST_SET in val; do 68 | NUM_LABELS=`cat ${EXP}/config/${NET_ID}/n_classes.txt` 69 | ((NUM_COORDS=NUM_LABELS*2)) 70 | TARGET_PERSON_SIZE=`cat ${EXP}/config/${NET_ID}/target_person_size.txt` 71 | TEST_SET=${TEST_SET}_${NUM_LABELS}_${TARGET_PERSON_SIZE}_${TRAIN_SET_SUFFIX} 72 | TEST_ITER=`cat ${EXP}/list/${TEST_SET}.txt | wc -l` 73 | MODEL=${EXP}/model/${NET_ID}/test.caffemodel 74 | if [ ! -f ${MODEL} ]; then 75 | MODEL=`ls -t ${EXP}/model/${NET_ID}/train_iter_*.caffemodel | head -n 1` 76 | fi 77 | echo Testing net ${EXP}/${NET_ID} 78 | FEATURE_DIR=${EXP}/features/${NET_ID} 79 | mkdir -p ${FEATURE_DIR}/${TEST_SET}/seg_score 80 | sed "$(eval echo $(cat sub.sed))" \ 81 | ${CONFIG_DIR}/testpy.prototxt > ${CONFIG_DIR}/testpy_${TEST_SET}.prototxt 82 | CMD="./store_pose_results.py \ 83 | ${CONFIG_DIR}/testpy_${TEST_SET}.prototxt \ 84 | ${MODEL} \ 85 | ${EXP}/list/${TEST_SET}.txt \ 86 | ${DATA_ROOT} \ 87 | ${FEATURE_DIR}/${TEST_SET}/seg_score \ 88 | ${CAFFE_PYTHON}" 89 | echo Running ${CMD} && ${CMD} 90 | done 91 | fi 92 | 93 | ## Evaluation #1 specification (on val or test) 94 | if [ ${RUN_EVALUATION} -eq 1 ]; then 95 | export GLOG_minloglevel=1 96 | for TEST_SET in val; do 97 | echo Evaluating net ${EXP}/${NET_ID} 98 | FEATURE_DIR=${EXP}/features/${NET_ID} 99 | NUM_LABELS=`cat ${EXP}/config/${NET_ID}/n_classes.txt` 100 | ((NUM_COORDS=NUM_LABELS*2)) # Add background. 101 | TARGET_PERSON_SIZE=`cat ${EXP}/config/${NET_ID}/target_person_size.txt` 102 | TEST_SET=${TEST_SET}_${NUM_LABELS}_${TARGET_PERSON_SIZE}_${TRAIN_SET_SUFFIX} 103 | CMD="./evaluate_pose.py \ 104 | ${EXP}/list/${TEST_SET}.txt \ 105 | ${EXP}/list/scale_${NUM_LABELS}_${TARGET_PERSON_SIZE}_${TRAIN_SET_SUFFIX}.txt \ 106 | ${FEATURE_DIR}/${TEST_SET}/seg_score \ 107 | ${NUM_LABELS}" 108 | echo Running ${CMD} && ${CMD} 109 | done 110 | fi 111 | 112 | 113 | ## Training #2 (finetune on trainval_aug) 114 | if [ ${RUN_TRAIN2} -eq 1 ]; then 115 | export GLOG_minloglevel=0 116 | LIST_DIR=${EXP}/list 117 | NUM_LABELS=`cat ${EXP}/config/${NET_ID}/n_classes.txt` 118 | ((NUM_COORDS=NUM_LABELS*2)) 119 | TARGET_PERSON_SIZE=`cat ${EXP}/config/${NET_ID}/target_person_size.txt` 120 | TRAIN_SET=trainval_${NUM_LABELS}_${TARGET_PERSON_SIZE}_${TRAIN_SET_SUFFIX} 121 | MODEL=${EXP}/config/${NET_ID}/init2.caffemodel 122 | if [ ! -f ${MODEL} ]; then 123 | MODEL=`ls -t ${EXP}/model/${NET_ID}/train_iter_*.caffemodel | head -n 1` 124 | fi 125 | echo Training2 net ${EXP}/${NET_ID} 126 | for pname in train solver2; do 127 | sed "$(eval echo $(cat sub.sed))" \ 128 | ${CONFIG_DIR}/${pname}.prototxt > ${CONFIG_DIR}/${pname}_${TRAIN_SET}.prototxt 129 | done 130 | CMD="${CAFFE_BIN} train \ 131 | --solver=${CONFIG_DIR}/solver2_${TRAIN_SET}.prototxt \ 132 | --weights=${MODEL} \ 133 | --gpu=${DEV_ID}" 134 | echo Running ${CMD} && ${CMD} 135 | fi 136 | 137 | ## Test #2 on official test set 138 | if [ ${RUN_TEST2} -eq 1 ]; then 139 | export GLOG_minloglevel=1 140 | for TEST_SET in test; do 141 | NUM_LABELS=`cat ${EXP}/config/${NET_ID}/n_classes.txt` 142 | ((NUM_COORDS=NUM_LABELS*2)) 143 | if [ -f ${EXP}/config/${NET_ID}/ev_classes.txt ]; then 144 | NUM_EV_LABELS=`cat ${EXP}/config/${NET_ID}/ev_classes.txt` 145 | else 146 | NUM_EV_LABELS=${NUM_LABELS} 147 | fi 148 | TARGET_PERSON_SIZE=`cat ${EXP}/config/${NET_ID}/target_person_size.txt` 149 | TEST_SET=${TEST_SET}_${NUM_EV_LABELS}_${TARGET_PERSON_SIZE}_${TRAIN_SET_SUFFIX} 150 | TEST_ITER=`cat ${EXP}/list/${TEST_SET}.txt | wc -l` 151 | MODEL=${EXP}/model/${NET_ID}/test2.caffemodel 152 | if [ ! -f ${MODEL} ]; then 153 | MODEL=`ls -t ${EXP}/model/${NET_ID}/train2_iter_*.caffemodel | head -n 1` 154 | fi 155 | echo Testing2 net ${EXP}/${NET_ID} 156 | FEATURE_DIR=${EXP}/features2/${NET_ID} 157 | mkdir -p ${FEATURE_DIR}/${TEST_SET}/seg_score 158 | sed "$(eval echo $(cat sub.sed))" \ 159 | ${CONFIG_DIR}/testpy.prototxt > ${CONFIG_DIR}/testpy_${TEST_SET}.prototxt 160 | CMD="./store_pose_results.py \ 161 | ${CONFIG_DIR}/testpy_${TEST_SET}.prototxt \ 162 | ${MODEL} \ 163 | ${EXP}/list/${TEST_SET}.txt \ 164 | ${DATA_ROOT} \ 165 | ${FEATURE_DIR}/${TEST_SET}/seg_score \ 166 | ${CAFFE_PYTHON}" 167 | echo Running ${CMD} && ${CMD} 168 | done 169 | fi 170 | 171 | ## Evaluation #2 specification (on val or test) 172 | if [ ${RUN_EVALUATION2} -eq 1 ]; then 173 | export GLOG_minloglevel=1 174 | for TEST_SET in test; do 175 | echo Evaluating net ${EXP}/${NET_ID} 176 | FEATURE_DIR=${EXP}/features2/${NET_ID} 177 | NUM_LABELS=`cat ${EXP}/config/${NET_ID}/n_classes.txt` 178 | ((NUM_COORDS=NUM_LABELS*2)) 179 | if [ -f ${EXP}/config/${NET_ID}/ev_classes.txt ]; then 180 | NUM_EV_LABELS=`cat ${EXP}/config/${NET_ID}/ev_classes.txt` 181 | else 182 | NUM_EV_LABELS=${NUM_LABELS} 183 | fi 184 | TARGET_PERSON_SIZE=`cat ${EXP}/config/${NET_ID}/target_person_size.txt` 185 | TEST_SET=${TEST_SET}_${NUM_EV_LABELS}_${TARGET_PERSON_SIZE}_${TRAIN_SET_SUFFIX} 186 | CMD="./evaluate_pose.py \ 187 | ${EXP}/list/${TEST_SET}.txt \ 188 | ${EXP}/list/scale_${NUM_EV_LABELS}_${TARGET_PERSON_SIZE}_${TRAIN_SET_SUFFIX}.txt \ 189 | ${FEATURE_DIR}/${TEST_SET}/seg_score 190 | ${NUM_EV_LABELS}" 191 | echo Running ${CMD} && ${CMD} 192 | done 193 | fi 194 | 195 | # Test on plain datasets. 196 | if [ ${RUN_TEST_PLAIN} -eq 1 ]; then 197 | export GLOG_minloglevel=1 198 | for TEST_SET in test; do 199 | NUM_LABELS=`cat ${EXP}/config/${NET_ID}/n_classes.txt` 200 | ((NUM_COORDS=NUM_LABELS*2)) 201 | TARGET_PERSON_SIZE=`cat ${EXP}/config/${NET_ID}/target_person_size.txt` 202 | TEST_SET=${TEST_SET}_14_${TARGET_PERSON_SIZE}_${TRAIN_SET_SUFFIX} 203 | TEST_ITER=`cat ${EXP}/list/${TEST_SET}.txt | wc -l` 204 | MODEL=${EXP}/model/${NET_ID}/test2.caffemodel 205 | if [ ! -f ${MODEL} ]; then 206 | MODEL=`ls -t ${EXP}/model/${NET_ID}/train2_iter_*.caffemodel | head -n 1` 207 | fi 208 | echo Testing2 net ${EXP}/${NET_ID} 209 | FEATURE_DIR=${EXP}/features2/${NET_ID} 210 | mkdir -p ${FEATURE_DIR}/${TEST_SET}/seg_score 211 | sed "$(eval echo $(cat sub.sed))" \ 212 | ${CONFIG_DIR}/testpy.prototxt > ${CONFIG_DIR}/testpy_${TEST_SET}.prototxt 213 | CMD="./store_pose_results.py \ 214 | ${CONFIG_DIR}/testpy_${TEST_SET}.prototxt \ 215 | ${MODEL} \ 216 | ${EXP}/list/${TEST_SET}.txt \ 217 | ${FEATURE_DIR}/${TEST_SET}/seg_score \ 218 | ${CAFFE_PYTHON}" 219 | echo Running ${CMD} && ${CMD} 220 | done 221 | fi 222 | 223 | ## Evaluation #2 specification (on val or test) 224 | if [ ${RUN_EVALUATION_PLAIN} -eq 1 ]; then 225 | export GLOG_minloglevel=1 226 | for TEST_SET in test; do 227 | echo Evaluating net ${EXP}/${NET_ID} 228 | FEATURE_DIR=${EXP}/features2/${NET_ID} 229 | NUM_LABELS=`cat ${EXP}/config/${NET_ID}/n_classes.txt` 230 | ((NUM_COORDS=NUM_LABELS*2)) 231 | TARGET_PERSON_SIZE=`cat ${EXP}/config/${NET_ID}/target_person_size.txt` 232 | TEST_SET=${TEST_SET}_14_${TARGET_PERSON_SIZE}_${TRAIN_SET_SUFFIX} 233 | CMD="./evaluate_pose.py \ 234 | ${EXP}/list/${TEST_SET}.txt \ 235 | ${EXP}/list/scale_${NUM_LABELS}_${TARGET_PERSON_SIZE}_${TRAIN_SET_SUFFIX}.txt \ 236 | ${FEATURE_DIR}/${TEST_SET}/seg_score 237 | ${NUM_LABELS}" 238 | echo Running ${CMD} && ${CMD} 239 | done 240 | fi 241 | 242 | 243 | -------------------------------------------------------------------------------- /pose/store_pose_results.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python2 2 | """Write the pose results to disk for further evaluation.""" 3 | # pylint: disable=invalid-name 4 | from os import path 5 | import sys 6 | import logging 7 | from PIL import Image 8 | 9 | import numpy as np 10 | import scipy.misc 11 | import scipy.misc as sm 12 | import click 13 | import tqdm 14 | 15 | from up_tools.model import ( 16 | connections_landmarks_91, dots_landmarks_91, 17 | reduction_91tolsp) 18 | from up_tools.visualization import visualize_pose 19 | from up_tools.model import connections_lsp 20 | LOGGER = logging.getLogger(__name__) 21 | 22 | # Image mean to use. 23 | _MEAN = np.array([104., 117., 123.]) 24 | # Scale factor for the CNN offset predictions. 25 | _LOCREF_SCALE_MUL = np.sqrt(53.) 26 | # Maximum size of one tile to process (to limit the required GPU memory). 27 | _MAX_SIZE = 700 28 | 29 | _STRIDE = 8. 30 | 31 | # CNN model store. 32 | _MODEL = None 33 | 34 | 35 | def estimate_pose(image, model_def, model_bin, caffe, scales=None): # pylint: disable=too-many-locals 36 | """ 37 | Get the estimated pose for an image. 38 | 39 | Uses the CNN pose estimator from "Deepcut: Joint Subset Partition and 40 | Labeling for Multi Person Pose Estimation" (Pishchulin et al., 2016), 41 | "DeeperCut: A Deeper, Stronger, and Faster Multi-Person Pose Estimation 42 | Model" (Insafutdinov et al., 2016). 43 | 44 | 45 | Parameters 46 | ========== 47 | 48 | :param image: np.array(3D). 49 | The image in height X width X BGR format. 50 | 51 | :param scales: list(float) or None. 52 | Scales on which to apply the estimator. The pose with the best confidence 53 | on average will be returned. 54 | 55 | Returns 56 | ======= 57 | 58 | :param pose: np.array(2D). 59 | The pose in 5x14 layout. The first axis is along per-joint information, 60 | the second the joints. Information is: 61 | 1. position x, 62 | 2. position y, 63 | 3. CNN confidence, 64 | 4. CNN offset vector x, 65 | 5. CNN offset vector y. 66 | """ 67 | global _MODEL # pylint: disable=global-statement 68 | if scales is None: 69 | scales = [1.] 70 | if _MODEL is None: 71 | LOGGER.info("Loading pose model...") 72 | _MODEL = caffe.Net(model_def, model_bin, caffe.TEST) 73 | LOGGER.info("Done!") 74 | LOGGER.debug("Processing image...") 75 | im_orig = image.copy() 76 | LOGGER.debug("Image shape: %s.", im_orig.shape) 77 | best_pose = None 78 | highest_confidence = 0. 79 | for scale_factor in scales: 80 | LOGGER.debug("Scale %f...", scale_factor) 81 | image = im_orig.copy() 82 | # Create input for multiples of net input/output changes. 83 | im_bg_width = int(np.ceil( 84 | float(image.shape[1]) * scale_factor / _STRIDE) * _STRIDE) 85 | im_bg_height = int(np.ceil( 86 | float(image.shape[0]) * scale_factor / _STRIDE) * _STRIDE) 87 | pad_size = 64 88 | im_bot_pixels = image[-1:, :, :] 89 | im_bot = np.tile(im_bot_pixels, (pad_size, 1, 1)) 90 | image = np.vstack((image, im_bot)) 91 | im_right_pixels = image[:, -1:, :] 92 | im_right = np.tile(im_right_pixels, (1, pad_size, 1)) 93 | image = np.hstack((image, im_right)) 94 | image = scipy.misc.imresize(image, scale_factor, interp='bilinear') 95 | image = image.astype('float32') - _MEAN 96 | 97 | net_input = np.zeros((im_bg_height, im_bg_width, 3), dtype='float32') 98 | net_input[:min(net_input.shape[0], image.shape[0]), 99 | :min(net_input.shape[1], image.shape[1]), :] =\ 100 | image[:min(net_input.shape[0], image.shape[0]), 101 | :min(net_input.shape[1], image.shape[1]), :] 102 | 103 | LOGGER.debug("Input shape: %d, %d.", 104 | net_input.shape[0], net_input.shape[1]) 105 | unary_maps, locreg_pred = _process_image_tiled(_MODEL, net_input, _STRIDE) 106 | 107 | """ 108 | import matplotlib.pyplot as plt 109 | plt.figure() 110 | for map_idx in range(unary_maps.shape[2]): 111 | plt.imshow(unary_maps[:, :, map_idx], interpolation='none') 112 | plt.imsave('map_%d.png' % map_idx, 113 | unary_maps[:, :, map_idx]) 114 | plt.show() 115 | """ 116 | 117 | pose = _pose_from_mats(unary_maps, locreg_pred, scale=scale_factor) 118 | 119 | minconf = np.min(pose[2, :]) 120 | if minconf > highest_confidence: 121 | LOGGER.debug("New best scale detected: %f (scale), " + 122 | "%f (min confidence).", scale_factor, minconf) 123 | highest_confidence = minconf 124 | best_pose = pose 125 | LOGGER.debug("Pose estimated.") 126 | return best_pose 127 | 128 | 129 | def _pose_from_mats(scoremat, offmat, scale): 130 | """Combine scoremat and offsets to the final pose.""" 131 | pose = [] 132 | for joint_idx in range(scoremat.shape[2]): 133 | maxloc = np.unravel_index(np.argmax(scoremat[:, :, joint_idx]), 134 | scoremat[:, :, joint_idx].shape) 135 | offset = np.array(offmat[maxloc][joint_idx])[::-1] 136 | pos_f8 = (np.array(maxloc).astype('float') * _STRIDE + 0.5 * _STRIDE + 137 | offset * _LOCREF_SCALE_MUL) 138 | pose.append(np.hstack((pos_f8[::-1] / scale, 139 | [scoremat[maxloc][joint_idx]], 140 | offset * _LOCREF_SCALE_MUL / scale))) 141 | return np.array(pose).T 142 | 143 | 144 | def _get_num_tiles(length, max_size, rf): 145 | """Get the number of tiles required to cover the entire image.""" 146 | if length <= max_size: 147 | return 1 148 | k = 0 149 | while True: 150 | new_size = (max_size - rf) * 2 + (max_size - 2*rf) * k 151 | if new_size > length: 152 | break 153 | k += 1 154 | return 2 + k 155 | 156 | 157 | # pylint: disable=too-many-locals 158 | def _process_image_tiled(model, net_input, stride): 159 | """Get the CNN results for the tiled image.""" 160 | rf = 224 # Standard receptive field size. 161 | cut_off = rf / stride 162 | 163 | num_tiles_x = _get_num_tiles(net_input.shape[1], _MAX_SIZE, rf) 164 | num_tiles_y = _get_num_tiles(net_input.shape[0], _MAX_SIZE, rf) 165 | if num_tiles_x > 1 or num_tiles_y > 1: 166 | LOGGER.info("Tiling the image into %d, %d (w, h) tiles...", 167 | num_tiles_x, num_tiles_y) 168 | 169 | scoremaps = [] 170 | locreg_pred = [] 171 | for j in range(num_tiles_y): 172 | start_y = j * (_MAX_SIZE - 2*rf) 173 | if j == num_tiles_y: 174 | end_y = net_input.shape[0] 175 | else: 176 | end_y = start_y + _MAX_SIZE 177 | scoremaps_line = [] 178 | locreg_pred_line = [] 179 | for i in range(num_tiles_x): 180 | start_x = i * (_MAX_SIZE - 2*rf) 181 | if i == num_tiles_x: 182 | end_x = net_input.shape[1] 183 | else: 184 | end_x = start_x + _MAX_SIZE 185 | input_tile = net_input[start_y:end_y, start_x:end_x, :] 186 | LOGGER.debug("Tile info: %d, %d, %d, %d.", 187 | start_y, end_y, start_x, end_x) 188 | scoremaps_tile, locreg_pred_tile = _cnn_process_image(model, 189 | input_tile) 190 | LOGGER.debug("Tile out shape: %s, %s.", 191 | str(scoremaps_tile.shape), 192 | str(locreg_pred_tile.shape)) 193 | scoremaps_tile = _cutoff_tile(scoremaps_tile, 194 | num_tiles_x, i, cut_off, True) 195 | locreg_pred_tile = _cutoff_tile(locreg_pred_tile, 196 | num_tiles_x, i, cut_off, True) 197 | scoremaps_tile = _cutoff_tile(scoremaps_tile, 198 | num_tiles_y, j, cut_off, False) 199 | locreg_pred_tile = _cutoff_tile(locreg_pred_tile, 200 | num_tiles_y, j, cut_off, False) 201 | LOGGER.debug("Cutoff tile out shape: %s, %s.", 202 | str(scoremaps_tile.shape), str(locreg_pred_tile.shape)) 203 | scoremaps_line.append(scoremaps_tile) 204 | locreg_pred_line.append(locreg_pred_tile) 205 | scoremaps_line = np.concatenate(scoremaps_line, axis=1) 206 | locreg_pred_line = np.concatenate(locreg_pred_line, axis=1) 207 | scoremaps_line = _cutoff_tile(scoremaps_line, 208 | num_tiles_y, j, cut_off, False) 209 | locreg_pred_line = _cutoff_tile(locreg_pred_line, 210 | num_tiles_y, j, cut_off, False) 211 | LOGGER.debug("Line tile out shape: %s, %s.", 212 | str(scoremaps_line.shape), str(locreg_pred_line.shape)) 213 | scoremaps.append(scoremaps_line) 214 | locreg_pred.append(locreg_pred_line) 215 | scoremaps = np.concatenate(scoremaps, axis=0) 216 | locreg_pred = np.concatenate(locreg_pred, axis=0) 217 | LOGGER.debug("Final tiled shape: %s, %s.", 218 | str(scoremaps.shape), str(locreg_pred.shape)) 219 | return scoremaps[:, :, 0, :], locreg_pred.transpose((0, 1, 3, 2)) 220 | 221 | 222 | def _cnn_process_image(model, net_input): 223 | """Get the CNN results for a fully prepared image.""" 224 | net_input = net_input.transpose((2, 0, 1)) 225 | model.blobs['data'].reshape(1, 3, net_input.shape[1], net_input.shape[2]) 226 | model.blobs['data'].data[0, ...] = net_input[...] 227 | model.forward() 228 | 229 | out_value = model.blobs['prob'].data 230 | feat_prob = out_value.copy().transpose((2, 3, 0, 1)) 231 | n_joints = feat_prob.shape[3] 232 | out_value = model.blobs['loc_pred'].data 233 | out_value = out_value.reshape((n_joints, 2, 234 | out_value.shape[2], 235 | out_value.shape[3])) 236 | locreg_pred = out_value.transpose((2, 3, 1, 0)) 237 | return feat_prob, locreg_pred 238 | 239 | 240 | def _cutoff_tile(sm, num_tiles, idx, cut_off, is_x): 241 | """Cut the valid parts of the CNN predictions for a tile.""" 242 | if is_x: 243 | sm = sm.transpose((1, 0, 2, 3)) 244 | if num_tiles == 1: 245 | pass 246 | elif idx == 0: 247 | sm = sm[:-cut_off, ...] 248 | elif idx == num_tiles-1: 249 | sm = sm[cut_off:, ...] 250 | else: 251 | sm = sm[cut_off:-cut_off, ...] 252 | if is_x: 253 | sm = sm.transpose((1, 0, 2, 3)) 254 | return sm 255 | 256 | 257 | @click.command() 258 | @click.argument("caffe_prototxt", type=click.Path(dir_okay=False, readable=True)) 259 | @click.argument("caffe_model", type=click.Path(dir_okay=False, readable=True)) 260 | @click.argument("image_list_file", type=click.Path(dir_okay=False, readable=True)) 261 | @click.argument("dset_root", type=click.Path(file_okay=False, readable=True)) 262 | @click.argument("output_folder", type=click.Path(file_okay=False, writable=True)) 263 | @click.argument("caffe_install_path", type=click.Path(file_okay=False, readable=True)) 264 | def main(caffe_prototxt, # pylint: disable=too-many-arguments, too-many-locals, too-many-statements 265 | caffe_model, 266 | image_list_file, 267 | dset_root, 268 | output_folder, 269 | caffe_install_path): 270 | """Store and visualize the pose results for a model.""" 271 | LOGGER.info("Storing pose results to folder `%s`.", output_folder) 272 | LOGGER.info("Using caffe from `%s`.", caffe_install_path) 273 | sys.path.insert(0, path.join(caffe_install_path)) 274 | import caffe # pylint: disable=import-error 275 | caffe.set_mode_gpu() 276 | with open(image_list_file, 'r') as inf: 277 | image_list = inf.readlines() 278 | image_list = [path.join(dset_root, line[1:]) 279 | for line in image_list if line.startswith('/')] 280 | n_landmarks = int(image_list_file[image_list_file.find('_')+1: 281 | image_list_file.find('_')+3]) 282 | for imgnames in tqdm.tqdm(image_list): 283 | imgname = imgnames.split(" ")[0].strip() 284 | LOGGER.debug("Processing `%s`...", imgname) 285 | im = caffe.io.load_image(imgname) # pylint: disable=invalid-name 286 | # caffe.io loads as RGB, and in range [0., 1.]. 287 | im = (im * 255.)[:, :, ::-1].astype('uint8') 288 | landmark_locs = estimate_pose(im, caffe_prototxt, caffe_model, caffe) 289 | if landmark_locs.shape[1] == 91 and n_landmarks == 14: 290 | # Extract LSP joints. 291 | landmark_locs = landmark_locs[:, reduction_91tolsp] 292 | np.save(path.join(output_folder, 293 | path.basename(imgname) + '.npy'), 294 | landmark_locs) 295 | vis_im = visualize_pose(im[:, :, ::-1], 296 | landmark_locs, scale=1., dash_length=5) 297 | sm.imsave(path.join(output_folder, 298 | path.basename(imgname) + '.npy.vis.png'), 299 | vis_im) 300 | 301 | 302 | if __name__ == '__main__': 303 | logging.basicConfig(level=logging.INFO) 304 | main() # pylint: disable=no-value-for-parameter 305 | -------------------------------------------------------------------------------- /pose/sub.sed: -------------------------------------------------------------------------------- 1 | 's,${STAT_DUMMY_FP},'"${STAT_DUMMY_FP}"',g;s,${NUM_COORDS},'"${NUM_COORDS}"',g;s,${DATA_ROOT},'"${DATA_ROOT}"',g;s,${DATA_ROOT1},'"${DATA_ROOT1}"',g;s,${EXP},'"${EXP}"',g;s,${EXP1},'"${EXP1}"',g;s,${TRAIN_SET},'"${TRAIN_SET}"',g;s,${TRAIN_SET1},'"${TRAIN_SET1}"',g;s,${TRAIN_SET1_WEAK},'"${TRAIN_SET1_WEAK}"',g;s,${TRAIN_SET1_STRONG},'"${TRAIN_SET1_STRONG}"',g;s,${TEST_SET},'"${TEST_SET}"',g;s,${NET_ID},'"${NET_ID}"',g;s,${FEATURE_DIR},'"${FEATURE_DIR}"',g;s,${NUM_LABELS},'"${NUM_LABELS}"',g;s,${NUM_LABELS1},'"${NUM_LABELS1}"',g;s,${NUM_LABELS_UNION},'"${NUM_LABELS_UNION}"',g;s,${BG_BIAS},'"${BG_BIAS}"',g;s,${FG_BIAS},'"${FG_BIAS}"',g;s,${TRAIN_SET_STRONG},'"${TRAIN_SET_STRONG}"',g;s,${TRAIN_SET_WEAK},'"${TRAIN_SET_WEAK}"',g;s,${BATCH_SIZE},'"${BATCH_SIZE}"',g;s,${TEST_SET_PREFIX},'"${TEST_SET_PREFIX}"',g;s,${TRAIN_STEP},'"${TRAIN_STEP}"',g' 2 | -------------------------------------------------------------------------------- /pose/tools/all_stats.txt: -------------------------------------------------------------------------------- 1 | # 2 | graph 3 | 182 2 4 | 1.000000 2.000000 5 | 1.000000 3.000000 6 | 1.000000 4.000000 7 | 1.000000 5.000000 8 | 1.000000 6.000000 9 | 1.000000 7.000000 10 | 1.000000 8.000000 11 | 1.000000 9.000000 12 | 1.000000 10.000000 13 | 1.000000 11.000000 14 | 1.000000 12.000000 15 | 1.000000 13.000000 16 | 1.000000 14.000000 17 | 2.000000 3.000000 18 | 2.000000 4.000000 19 | 2.000000 5.000000 20 | 2.000000 6.000000 21 | 2.000000 7.000000 22 | 2.000000 8.000000 23 | 2.000000 9.000000 24 | 2.000000 10.000000 25 | 2.000000 11.000000 26 | 2.000000 12.000000 27 | 2.000000 13.000000 28 | 2.000000 14.000000 29 | 3.000000 4.000000 30 | 3.000000 5.000000 31 | 3.000000 6.000000 32 | 3.000000 7.000000 33 | 3.000000 8.000000 34 | 3.000000 9.000000 35 | 3.000000 10.000000 36 | 3.000000 11.000000 37 | 3.000000 12.000000 38 | 3.000000 13.000000 39 | 3.000000 14.000000 40 | 4.000000 5.000000 41 | 4.000000 6.000000 42 | 4.000000 7.000000 43 | 4.000000 8.000000 44 | 4.000000 9.000000 45 | 4.000000 10.000000 46 | 4.000000 11.000000 47 | 4.000000 12.000000 48 | 4.000000 13.000000 49 | 4.000000 14.000000 50 | 5.000000 6.000000 51 | 5.000000 7.000000 52 | 5.000000 8.000000 53 | 5.000000 9.000000 54 | 5.000000 10.000000 55 | 5.000000 11.000000 56 | 5.000000 12.000000 57 | 5.000000 13.000000 58 | 5.000000 14.000000 59 | 6.000000 7.000000 60 | 6.000000 8.000000 61 | 6.000000 9.000000 62 | 6.000000 10.000000 63 | 6.000000 11.000000 64 | 6.000000 12.000000 65 | 6.000000 13.000000 66 | 6.000000 14.000000 67 | 7.000000 8.000000 68 | 7.000000 9.000000 69 | 7.000000 10.000000 70 | 7.000000 11.000000 71 | 7.000000 12.000000 72 | 7.000000 13.000000 73 | 7.000000 14.000000 74 | 8.000000 9.000000 75 | 8.000000 10.000000 76 | 8.000000 11.000000 77 | 8.000000 12.000000 78 | 8.000000 13.000000 79 | 8.000000 14.000000 80 | 9.000000 10.000000 81 | 9.000000 11.000000 82 | 9.000000 12.000000 83 | 9.000000 13.000000 84 | 9.000000 14.000000 85 | 10.000000 11.000000 86 | 10.000000 12.000000 87 | 10.000000 13.000000 88 | 10.000000 14.000000 89 | 11.000000 12.000000 90 | 11.000000 13.000000 91 | 11.000000 14.000000 92 | 12.000000 13.000000 93 | 12.000000 14.000000 94 | 13.000000 14.000000 95 | 2.000000 1.000000 96 | 3.000000 1.000000 97 | 4.000000 1.000000 98 | 5.000000 1.000000 99 | 6.000000 1.000000 100 | 7.000000 1.000000 101 | 8.000000 1.000000 102 | 9.000000 1.000000 103 | 10.000000 1.000000 104 | 11.000000 1.000000 105 | 12.000000 1.000000 106 | 13.000000 1.000000 107 | 14.000000 1.000000 108 | 3.000000 2.000000 109 | 4.000000 2.000000 110 | 5.000000 2.000000 111 | 6.000000 2.000000 112 | 7.000000 2.000000 113 | 8.000000 2.000000 114 | 9.000000 2.000000 115 | 10.000000 2.000000 116 | 11.000000 2.000000 117 | 12.000000 2.000000 118 | 13.000000 2.000000 119 | 14.000000 2.000000 120 | 4.000000 3.000000 121 | 5.000000 3.000000 122 | 6.000000 3.000000 123 | 7.000000 3.000000 124 | 8.000000 3.000000 125 | 9.000000 3.000000 126 | 10.000000 3.000000 127 | 11.000000 3.000000 128 | 12.000000 3.000000 129 | 13.000000 3.000000 130 | 14.000000 3.000000 131 | 5.000000 4.000000 132 | 6.000000 4.000000 133 | 7.000000 4.000000 134 | 8.000000 4.000000 135 | 9.000000 4.000000 136 | 10.000000 4.000000 137 | 11.000000 4.000000 138 | 12.000000 4.000000 139 | 13.000000 4.000000 140 | 14.000000 4.000000 141 | 6.000000 5.000000 142 | 7.000000 5.000000 143 | 8.000000 5.000000 144 | 9.000000 5.000000 145 | 10.000000 5.000000 146 | 11.000000 5.000000 147 | 12.000000 5.000000 148 | 13.000000 5.000000 149 | 14.000000 5.000000 150 | 7.000000 6.000000 151 | 8.000000 6.000000 152 | 9.000000 6.000000 153 | 10.000000 6.000000 154 | 11.000000 6.000000 155 | 12.000000 6.000000 156 | 13.000000 6.000000 157 | 14.000000 6.000000 158 | 8.000000 7.000000 159 | 9.000000 7.000000 160 | 10.000000 7.000000 161 | 11.000000 7.000000 162 | 12.000000 7.000000 163 | 13.000000 7.000000 164 | 14.000000 7.000000 165 | 9.000000 8.000000 166 | 10.000000 8.000000 167 | 11.000000 8.000000 168 | 12.000000 8.000000 169 | 13.000000 8.000000 170 | 14.000000 8.000000 171 | 10.000000 9.000000 172 | 11.000000 9.000000 173 | 12.000000 9.000000 174 | 13.000000 9.000000 175 | 14.000000 9.000000 176 | 11.000000 10.000000 177 | 12.000000 10.000000 178 | 13.000000 10.000000 179 | 14.000000 10.000000 180 | 12.000000 11.000000 181 | 13.000000 11.000000 182 | 14.000000 11.000000 183 | 13.000000 12.000000 184 | 14.000000 12.000000 185 | 14.000000 13.000000 186 | # 187 | means 188 | 182 2 189 | -0.596845 -59.773762 190 | 1.112818 -117.470432 191 | 16.016898 -117.196001 192 | 18.529572 -59.390058 193 | 18.500304 0.244287 194 | -5.923113 -141.039923 195 | -8.726263 -158.931410 196 | -3.845356 -199.031166 197 | 19.642209 -199.486912 198 | 24.495313 -158.681344 199 | 21.385920 -139.426824 200 | 8.220737 -208.741467 201 | 7.872185 -253.828332 202 | 1.633174 -59.030643 203 | 17.029461 -58.888015 204 | 19.546379 -0.150480 205 | 18.807119 59.624722 206 | -4.400504 -84.450352 207 | -7.818063 -101.651938 208 | -2.994367 -142.339518 209 | 21.316724 -142.765904 210 | 26.365762 -101.368875 211 | 23.144902 -82.830286 212 | 9.544185 -152.215656 213 | 9.658370 -197.668168 214 | 16.659667 0.074488 215 | 18.148599 58.779021 216 | 17.273294 117.474079 217 | -5.790904 -28.603901 218 | -10.029987 -44.534014 219 | -4.866420 -86.148420 220 | 21.622010 -86.571236 221 | 27.032585 -44.167698 222 | 23.503280 -26.821164 223 | 8.781065 -96.459942 224 | 9.198072 -142.613958 225 | 2.711417 58.756738 226 | 2.345221 117.437363 227 | -22.511794 -28.691243 228 | -26.779033 -44.557264 229 | -21.563860 -86.172727 230 | 4.963388 -86.644332 231 | 10.368280 -44.269552 232 | 6.819239 -26.950606 233 | -7.897649 -96.503420 234 | -7.471827 -142.658821 235 | -0.445034 59.943425 236 | -24.107125 -84.203540 237 | -27.670481 -101.341659 238 | -22.754776 -142.063080 239 | 1.642298 -142.610577 240 | 6.786021 -101.210470 241 | 3.630139 -82.659050 242 | -10.164536 -152.012501 243 | -9.937193 -197.440564 244 | -23.655438 -140.914503 245 | -26.817087 -158.710981 246 | -21.906856 -198.888329 247 | 1.738062 -199.622165 248 | 6.674466 -158.919308 249 | 3.765318 -139.542230 250 | -9.700556 -208.716707 251 | -9.703093 -253.749297 252 | -4.439832 -15.497981 253 | 0.804316 -57.107521 254 | 27.717572 -57.474980 255 | 33.281145 -15.075975 256 | 29.533130 1.852609 257 | 14.661406 -67.428782 258 | 15.041111 -113.704108 259 | 5.255915 -41.579967 260 | 32.108831 -41.981623 261 | 37.676739 0.415842 262 | 33.991304 17.334034 263 | 19.077185 -51.924133 264 | 19.464553 -98.191686 265 | 26.851041 -0.419367 266 | 32.387510 41.962568 267 | 28.750543 58.910973 268 | 13.827514 -10.365580 269 | 14.233335 -56.641585 270 | 5.492311 42.394972 271 | 1.771324 59.398506 272 | -13.021406 -9.946049 273 | -12.613133 -56.221677 274 | -3.716947 16.998029 275 | -18.540351 -52.315746 276 | -18.142881 -98.588556 277 | -14.859561 -69.283178 278 | -14.420617 -115.592773 279 | 0.407077 -46.274953 280 | 0.596845 59.773762 281 | -1.112818 117.470432 282 | -16.016898 117.196001 283 | -18.529572 59.390058 284 | -18.500304 -0.244287 285 | 5.923113 141.039923 286 | 8.726263 158.931410 287 | 3.845356 199.031166 288 | -19.642209 199.486912 289 | -24.495313 158.681344 290 | -21.385920 139.426824 291 | -8.220737 208.741467 292 | -7.872185 253.828332 293 | -1.633174 59.030643 294 | -17.029461 58.888015 295 | -19.546379 0.150480 296 | -18.807119 -59.624722 297 | 4.400504 84.450352 298 | 7.818063 101.651938 299 | 2.994367 142.339518 300 | -21.316724 142.765904 301 | -26.365762 101.368875 302 | -23.144902 82.830286 303 | -9.544185 152.215656 304 | -9.658370 197.668168 305 | -16.659667 -0.074488 306 | -18.148599 -58.779021 307 | -17.273294 -117.474079 308 | 5.790904 28.603901 309 | 10.029987 44.534014 310 | 4.866420 86.148420 311 | -21.622010 86.571236 312 | -27.032585 44.167698 313 | -23.503280 26.821164 314 | -8.781065 96.459942 315 | -9.198072 142.613958 316 | -2.711417 -58.756738 317 | -2.345221 -117.437363 318 | 22.511794 28.691243 319 | 26.779033 44.557264 320 | 21.563860 86.172727 321 | -4.963388 86.644332 322 | -10.368280 44.269552 323 | -6.819239 26.950606 324 | 7.897649 96.503420 325 | 7.471827 142.658821 326 | 0.445034 -59.943425 327 | 24.107125 84.203540 328 | 27.670481 101.341659 329 | 22.754776 142.063080 330 | -1.642298 142.610577 331 | -6.786021 101.210470 332 | -3.630139 82.659050 333 | 10.164536 152.012501 334 | 9.937193 197.440564 335 | 23.655438 140.914503 336 | 26.817087 158.710981 337 | 21.906856 198.888329 338 | -1.738062 199.622165 339 | -6.674466 158.919308 340 | -3.765318 139.542230 341 | 9.700556 208.716707 342 | 9.703093 253.749297 343 | 4.439832 15.497981 344 | -0.804316 57.107521 345 | -27.717572 57.474980 346 | -33.281145 15.075975 347 | -29.533130 -1.852609 348 | -14.661406 67.428782 349 | -15.041111 113.704108 350 | -5.255915 41.579967 351 | -32.108831 41.981623 352 | -37.676739 -0.415842 353 | -33.991304 -17.334034 354 | -19.077185 51.924133 355 | -19.464553 98.191686 356 | -26.851041 0.419367 357 | -32.387510 -41.962568 358 | -28.750543 -58.910973 359 | -13.827514 10.365580 360 | -14.233335 56.641585 361 | -5.492311 -42.394972 362 | -1.771324 -59.398506 363 | 13.021406 9.946049 364 | 12.613133 56.221677 365 | 3.716947 -16.998029 366 | 18.540351 52.315746 367 | 18.142881 98.588556 368 | 14.859561 69.283178 369 | 14.420617 115.592773 370 | -0.407077 46.274953 371 | # 372 | std_devs 373 | 182 2 374 | 32.051950 33.818434 375 | 48.779257 59.142557 376 | 60.862712 59.806249 377 | 62.880406 39.985839 378 | 64.326598 32.853912 379 | 75.676988 86.449371 380 | 68.521536 79.449583 381 | 66.849179 81.626041 382 | 79.341679 82.288071 383 | 85.587036 80.234667 384 | 89.724683 85.997029 385 | 70.322892 86.224146 386 | 85.242267 97.121723 387 | 35.866370 37.420458 388 | 51.411395 38.622123 389 | 52.928866 24.244678 390 | 62.644948 41.202537 391 | 58.447866 68.998384 392 | 51.537221 59.119966 393 | 48.614493 59.764199 394 | 65.905382 60.646905 395 | 74.066562 60.250605 396 | 76.129804 68.865127 397 | 52.943721 64.496932 398 | 66.759756 75.740508 399 | 32.135644 10.270822 400 | 51.078909 39.556511 401 | 60.944887 60.779886 402 | 58.851393 55.136783 403 | 44.300666 38.241330 404 | 35.154945 34.379439 405 | 52.011931 35.549636 406 | 63.267739 39.546021 407 | 71.465162 54.960734 408 | 38.661183 39.674491 409 | 59.290249 53.556659 410 | 35.181350 37.686751 411 | 48.530922 59.772016 412 | 73.313200 56.197276 413 | 64.215292 40.072233 414 | 51.975496 35.735912 415 | 35.127562 34.627540 416 | 43.844167 37.863858 417 | 57.653753 53.841986 418 | 38.865993 39.918443 419 | 59.555588 53.784007 420 | 31.602156 34.588807 421 | 77.424142 70.762681 422 | 74.371348 61.337865 423 | 65.432234 61.059446 424 | 47.960537 60.085339 425 | 50.934889 58.931535 426 | 57.650103 67.722889 427 | 52.550498 64.925787 428 | 66.554817 76.140279 429 | 90.404046 88.281407 430 | 86.089607 81.423695 431 | 79.335685 82.780287 432 | 66.341874 82.122438 433 | 67.802548 79.560055 434 | 74.751377 85.555991 435 | 70.209847 86.728923 436 | 85.112237 97.701778 437 | 29.554973 29.721583 438 | 43.027837 44.888210 439 | 72.948304 48.892374 440 | 81.787961 45.228812 441 | 79.077383 44.459847 442 | 53.275982 48.133831 443 | 55.210664 52.182259 444 | 25.747935 24.215435 445 | 67.323145 31.304057 446 | 78.827130 32.380424 447 | 80.428612 44.097068 448 | 44.358790 30.418020 449 | 53.343895 38.624712 450 | 51.330181 16.116188 451 | 66.275290 31.040904 452 | 70.870715 47.624213 453 | 26.055384 12.909241 454 | 38.707068 25.673379 455 | 25.557412 23.940686 456 | 42.401853 43.490125 457 | 26.665069 12.782185 458 | 39.685631 25.577513 459 | 29.860362 29.030683 460 | 44.082640 30.181023 461 | 53.803870 38.697362 462 | 52.096603 46.903246 463 | 54.950107 51.316071 464 | 25.934435 19.009251 465 | 32.051950 33.818434 466 | 48.779257 59.142557 467 | 60.862712 59.806249 468 | 62.880406 39.985839 469 | 64.326598 32.853912 470 | 75.676988 86.449371 471 | 68.521536 79.449583 472 | 66.849179 81.626041 473 | 79.341679 82.288071 474 | 85.587036 80.234667 475 | 89.724683 85.997029 476 | 70.322892 86.224146 477 | 85.242267 97.121723 478 | 35.866370 37.420458 479 | 51.411395 38.622123 480 | 52.928866 24.244678 481 | 62.644948 41.202537 482 | 58.447866 68.998384 483 | 51.537221 59.119966 484 | 48.614493 59.764199 485 | 65.905382 60.646905 486 | 74.066562 60.250605 487 | 76.129804 68.865127 488 | 52.943721 64.496932 489 | 66.759756 75.740508 490 | 32.135644 10.270822 491 | 51.078909 39.556511 492 | 60.944887 60.779886 493 | 58.851393 55.136783 494 | 44.300666 38.241330 495 | 35.154945 34.379439 496 | 52.011931 35.549636 497 | 63.267739 39.546021 498 | 71.465162 54.960734 499 | 38.661183 39.674491 500 | 59.290249 53.556659 501 | 35.181350 37.686751 502 | 48.530922 59.772016 503 | 73.313200 56.197276 504 | 64.215292 40.072233 505 | 51.975496 35.735912 506 | 35.127562 34.627540 507 | 43.844167 37.863858 508 | 57.653753 53.841986 509 | 38.865993 39.918443 510 | 59.555588 53.784007 511 | 31.602156 34.588807 512 | 77.424142 70.762681 513 | 74.371348 61.337865 514 | 65.432234 61.059446 515 | 47.960537 60.085339 516 | 50.934889 58.931535 517 | 57.650103 67.722889 518 | 52.550498 64.925787 519 | 66.554817 76.140279 520 | 90.404046 88.281407 521 | 86.089607 81.423695 522 | 79.335685 82.780287 523 | 66.341874 82.122438 524 | 67.802548 79.560055 525 | 74.751377 85.555991 526 | 70.209847 86.728923 527 | 85.112237 97.701778 528 | 29.554973 29.721583 529 | 43.027837 44.888210 530 | 72.948304 48.892374 531 | 81.787961 45.228812 532 | 79.077383 44.459847 533 | 53.275982 48.133831 534 | 55.210664 52.182259 535 | 25.747935 24.215435 536 | 67.323145 31.304057 537 | 78.827130 32.380424 538 | 80.428612 44.097068 539 | 44.358790 30.418020 540 | 53.343895 38.624712 541 | 51.330181 16.116188 542 | 66.275290 31.040904 543 | 70.870715 47.624213 544 | 26.055384 12.909241 545 | 38.707068 25.673379 546 | 25.557412 23.940686 547 | 42.401853 43.490125 548 | 26.665069 12.782185 549 | 39.685631 25.577513 550 | 29.860362 29.030683 551 | 44.082640 30.181023 552 | 53.803870 38.697362 553 | 52.096603 46.903246 554 | 54.950107 51.316071 555 | 25.934435 19.009251 556 | -------------------------------------------------------------------------------- /pose/tools/create_dataset.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python2 2 | """Create segmentation datasets from select SMPL fits.""" 3 | # pylint: disable=invalid-name 4 | import os 5 | import os.path as path 6 | import sys 7 | import logging 8 | import csv 9 | from collections import OrderedDict 10 | from copy import copy as _copy 11 | import cPickle as pickle 12 | 13 | import numpy as np 14 | import scipy 15 | import scipy.io as sio 16 | import click 17 | import opendr.camera as _odr_c 18 | import tqdm 19 | 20 | from clustertools.log import LOGFORMAT 21 | 22 | from up_tools.model import (robust_person_size, rlswap_lsp, connections_lsp, 23 | rlswap_landmarks_91, get_crop, landmark_mesh_91) 24 | import up_tools.visualization as vs 25 | from up_tools.render_segmented_views import render_body_impl # pylint: disable=unused-import 26 | from up_tools.mesh import Mesh 27 | try: 28 | # Robustify against setup. 29 | from smpl.serialization import load_model 30 | except ImportError: 31 | # pylint: disable=import-error 32 | try: 33 | from psbody.smpl.serialization import load_model 34 | except ImportError: 35 | from smpl_webuser.serialization import load_model 36 | sys.path.insert(0, path.join(path.dirname(__file__), '..', '..')) 37 | from config import POSE_DATA_FP, UP3D_FP 38 | 39 | 40 | LOGGER = logging.getLogger(__name__) 41 | DSET_ROOT_FP = POSE_DATA_FP 42 | MODEL_NEUTRAL_PATH = path.join( 43 | path.dirname(__file__), '..', '..', 'models', '3D', 44 | 'basicModel_neutral_lbs_10_207_0_v1.0.0.pkl') 45 | MODEL_NEUTRAL = load_model(MODEL_NEUTRAL_PATH) 46 | _TEMPLATE_MESH = Mesh(filename=path.join( 47 | path.dirname(__file__), '..', '..', 48 | 'models', '3D', 49 | 'template.ply')) 50 | 51 | if not path.exists(DSET_ROOT_FP): 52 | os.mkdir(DSET_ROOT_FP) 53 | 54 | 55 | 56 | def get_landmark_positions(stored_parameter_fp, resolution, landmarks): 57 | """Get landmark positions for a given image.""" 58 | with open(stored_parameter_fp, 'rb') as inf: 59 | stored_parameters = pickle.load(inf) 60 | # Pose the model. 61 | model = MODEL_NEUTRAL 62 | model.betas[:len(stored_parameters['betas'])] = stored_parameters['betas'] 63 | model.pose[:] = stored_parameters['pose'] 64 | model.trans[:] = stored_parameters['trans'] 65 | mesh = _copy(_TEMPLATE_MESH) 66 | mesh.v = model.r 67 | mesh_points = mesh.v[tuple(landmarks.values()),] 68 | J_onbetas = model.J_regressor.dot(model.r) 69 | skeleton_points = J_onbetas[(8, 5, 2, 1, 4, 7, 21, 19, 17, 16, 18, 20),] 70 | # Do the projection. 71 | camera = _odr_c.ProjectPoints( 72 | rt=stored_parameters['rt'], 73 | t=stored_parameters['t'], 74 | f=(stored_parameters['f'], stored_parameters['f']), 75 | c=np.array(resolution) / 2., 76 | k=np.zeros(5)) 77 | camera.v = np.vstack((skeleton_points, mesh_points)) 78 | return camera.r.T.copy() 79 | 80 | 81 | def add_dataset(dset_fp, dset_fromroot, list_ids, up3d_fp, # pylint: disable=too-many-locals, too-many-arguments, too-many-statements, too-many-branches 82 | train_list_f, val_list_f, train_val_list_f, test_list_f, scale_f, 83 | train_spec, val_spec, test_spec, 84 | target_person_size, landmarks, train_crop, test_crop, running_idx, 85 | only_missing=False, with_rlswap=True, write_gtjoints_as_lm=False, 86 | human_annotations=False): 87 | """Add a dataset to the collection.""" 88 | test_ids = [int(id_[1:6]) for id_ in test_spec] 89 | train_ids = [int(id_[1:6]) for id_ in train_spec] 90 | val_ids = [int(id_[1:6]) for id_ in val_spec] 91 | LOGGER.info("Split: %d train, %d val, %d test.", 92 | len(train_ids), len(val_ids), len(test_ids)) 93 | LOGGER.info("Writing dataset...") 94 | for im_idx in tqdm.tqdm(train_ids + val_ids + test_ids): 95 | image = scipy.misc.imread(path.join(up3d_fp, '%05d_image.png' % (im_idx))) 96 | with open(path.join(up3d_fp, '%05d_fit_crop_info.txt' % (im_idx)), 'r') as inf: 97 | cropinfo = [int(val) for val in inf.readline().strip().split()] 98 | assert image.ndim == 3 99 | out_exists = (path.exists(path.join(dset_fp, '%05d_image.png' % (running_idx))) and 100 | path.exists(path.join(dset_fp, '%05d_ann_vis.png' % (running_idx)))) 101 | if with_rlswap and im_idx not in test_ids: 102 | out_exists = out_exists and ( 103 | path.exists(path.join(dset_fp, '%05d_image.png' % (running_idx + 1))) and 104 | path.exists(path.join(dset_fp, '%05d_ann_vis.png' % (running_idx + 1)))) 105 | if not (only_missing and out_exists or write_gtjoints_as_lm): 106 | if human_annotations: 107 | landmark_pos = np.load(path.join(up3d_fp, '%05d_joints.npy' % (im_idx))) 108 | else: 109 | landmark_pos = get_landmark_positions(path.join(up3d_fp, '%05d_body.pkl' % (im_idx)), 110 | (cropinfo[1], cropinfo[0]), 111 | landmarks) 112 | fac_y = cropinfo[0] / float(cropinfo[3] - cropinfo[2]) 113 | fac_x = cropinfo[1] / float(cropinfo[5] - cropinfo[4]) 114 | landmark_pos[:2, :] /= np.mean([fac_x, fac_y]) 115 | landmark_pos[0, :] += cropinfo[4] 116 | landmark_pos[1, :] += cropinfo[2] 117 | joints = np.load(path.join(up3d_fp, '%05d_joints.npy' % (im_idx))) 118 | joints = np.vstack((joints, np.all(joints > 0, axis=0)[None, :])) 119 | person_size = robust_person_size(joints) 120 | norm_factor = float(target_person_size) / person_size 121 | joints[:2, :] *= norm_factor 122 | if not (only_missing and out_exists or write_gtjoints_as_lm): 123 | landmark_pos[:2, :] *= norm_factor 124 | if write_gtjoints_as_lm: 125 | landmark_pos = joints.copy() 126 | image = scipy.misc.imresize(image, norm_factor, interp='bilinear') 127 | if im_idx in test_ids: 128 | crop = test_crop 129 | else: 130 | crop = train_crop 131 | if image.shape[0] > crop or image.shape[1] > crop: 132 | LOGGER.debug("Image (original %d, here %d) too large (%s)! Cropping...", 133 | im_idx, running_idx, str(image.shape[:2])) 134 | person_center = np.mean(joints[:2, joints[2, :] == 1], axis=1) 135 | crop_y, crop_x = get_crop(image, person_center, crop) 136 | image = image[crop_y[0]:crop_y[1], 137 | crop_x[0]:crop_x[1], :] 138 | landmark_pos[0, :] -= crop_x[0] 139 | landmark_pos[1, :] -= crop_y[0] 140 | assert image.shape[0] == crop or image.shape[1] == crop, ( 141 | "Error cropping image (original %d, here %d)!" % (im_idx, 142 | running_idx)) 143 | assert image.shape[0] <= crop and image.shape[1] <= crop and image.shape[2] == 3, ( 144 | "Wrong image shape (original %d, here %d)!" % (im_idx, running_idx)) 145 | vis_im = vs.visualize_pose(image, landmark_pos, scale=1.) 146 | if not (only_missing and out_exists): 147 | scipy.misc.imsave(path.join(dset_fp, '%05d_image.png' % (running_idx)), image) 148 | scipy.misc.imsave(path.join(dset_fp, '%05d_ann_vis.png' % (running_idx)), vis_im) 149 | if with_rlswap and im_idx not in test_ids: 150 | if landmark_pos.shape[1] == 14: 151 | landmark_pos_swapped = landmark_pos[:, rlswap_lsp] 152 | else: 153 | landmark_pos_swapped = landmark_pos[:, rlswap_landmarks_91] 154 | landmark_pos_swapped[0, :] = image.shape[1] - landmark_pos_swapped[0, :] 155 | image_swapped = image[:, ::-1, :] 156 | # Use core visualization for 14 joints. 157 | vis_im_swapped = vs.visualize_pose(image_swapped, 158 | landmark_pos_swapped, 159 | scale=1) 160 | if not (only_missing and out_exists): 161 | scipy.misc.imsave(path.join(dset_fp, '%05d_image.png' % (running_idx + 1)), 162 | image_swapped) 163 | scipy.misc.imsave(path.join(dset_fp, '%05d_ann_vis.png' % (running_idx + 1)), 164 | vis_im_swapped) 165 | list_fs = [] 166 | list_id_ids = [] 167 | if im_idx in train_ids: 168 | list_fs.append(train_val_list_f) 169 | list_id_ids.append(2) 170 | list_fs.append(train_list_f) 171 | list_id_ids.append(0) 172 | elif im_idx in val_ids: 173 | list_fs.append(train_val_list_f) 174 | list_id_ids.append(2) 175 | list_fs.append(val_list_f) 176 | list_id_ids.append(1) 177 | elif im_idx in test_ids: 178 | list_fs.append(test_list_f) 179 | list_id_ids.append(3) 180 | for list_f, list_id_idx in zip(list_fs, list_id_ids): 181 | # pylint: disable=bad-continuation 182 | list_f.write( 183 | """# %d 184 | %s 185 | 3 186 | %d 187 | %d 188 | %d 189 | """ % ( 190 | list_ids[list_id_idx], 191 | path.join('/' + dset_fromroot, '%05d_image.png' % (running_idx)), 192 | image.shape[0], 193 | image.shape[1], 194 | landmark_pos.shape[1])) 195 | for landmark_idx, landmark_point in enumerate(landmark_pos.T): 196 | list_f.write("%d %d %d\n" % (landmark_idx + 1, 197 | int(landmark_point[0]), 198 | int(landmark_point[1]))) 199 | list_f.flush() 200 | list_ids[list_id_idx] += 1 201 | scale_f.write("%05d_image.png %f\n" % (running_idx, norm_factor)) 202 | scale_f.flush() 203 | running_idx += 1 204 | if with_rlswap and im_idx not in test_ids: 205 | for list_f, list_id_idx in zip(list_fs, list_id_ids): 206 | # pylint: disable=bad-continuation 207 | list_f.write( 208 | """# %d 209 | %s 210 | 3 211 | %d 212 | %d 213 | %d 214 | """ % ( 215 | list_ids[list_id_idx], 216 | path.join('/' + dset_fromroot, '%05d_image.png' % (running_idx)), 217 | image.shape[0], 218 | image.shape[1], 219 | landmark_pos.shape[1])) 220 | for landmark_idx, landmark_point in enumerate(landmark_pos_swapped.T): 221 | list_f.write("%d %d %d\n" % (landmark_idx + 1, 222 | int(landmark_point[0]), 223 | int(landmark_point[1]))) 224 | list_f.flush() 225 | list_ids[list_id_idx] += 1 226 | scale_f.write("%05d_image.png %f\n" % (running_idx, norm_factor)) 227 | scale_f.flush() 228 | running_idx += 1 229 | return running_idx 230 | 231 | 232 | @click.command() 233 | @click.argument("suffix", type=click.STRING) 234 | @click.argument("target_person_size", type=click.INT) 235 | @click.option("--crop", type=click.INT, default=513, # Used to be 513. 236 | help="Crop size for the images.") 237 | @click.option("--test_crop", type=click.INT, default=513, # Used to be 513. 238 | help="Crop size for the images.") 239 | @click.option("--only_missing", type=click.BOOL, default=False, is_flag=True, 240 | help="Only rewrite missing images.") 241 | @click.option("--noswap", type=click.BOOL, default=False, is_flag=True, 242 | help="Do not produce side-swapped samples.") 243 | @click.option("--core_joints", type=click.BOOL, default=False, is_flag=True, 244 | help="Use only the 14 usual joints projected from SMPL.") 245 | @click.option("--human_annotations", type=click.BOOL, default=False, is_flag=True, 246 | help="Use the human annotations (FashionPose should be excluded!).") 247 | @click.option("--up3d_fp", type=click.Path(file_okay=False, readable=True), 248 | default=UP3D_FP, 249 | help="Path to the UP3D folder that you want to use.") 250 | def cli(suffix, target_person_size, crop=513, test_crop=513, # pylint: disable=too-many-locals, too-many-arguments 251 | only_missing=False, noswap=False, core_joints=False, 252 | human_annotations=False, up3d_fp=UP3D_FP): 253 | """Create segmentation datasets from select SMPL fits.""" 254 | np.random.seed(1) 255 | with_rlswap = not noswap 256 | if human_annotations: 257 | assert core_joints 258 | if test_crop < target_person_size or crop < target_person_size: 259 | LOGGER.critical("Too small crop size!") 260 | raise Exception("Too small crop size!") 261 | landmark_mapping = landmark_mesh_91 262 | if core_joints: 263 | LOGGER.info("Using the core joints.") 264 | # Order is important here! This way, we maintain LSP compatibility. 265 | landmark_mapping = OrderedDict([('neck', landmark_mapping['neck']), 266 | ('head_top', landmark_mapping['head_top']),]) 267 | n_landmarks = len(landmark_mapping) + 12 268 | LOGGER.info("Creating pose dataset with %d landmarks with target " 269 | "person size %f and suffix `%s`.", 270 | n_landmarks, target_person_size, suffix) 271 | assert ' ' not in suffix 272 | dset_fromroot = path.join(str(n_landmarks), str(target_person_size), suffix) 273 | dset_fp = path.join(DSET_ROOT_FP, dset_fromroot) 274 | if path.exists(dset_fp): 275 | if not click.confirm("Dataset folder exists: `%s`! Continue?" % (dset_fp)): 276 | return 277 | else: 278 | os.makedirs(dset_fp) 279 | LOGGER.info("Creating list files...") 280 | list_fp = path.join(path.dirname(__file__), '..', 'training', 'list') 281 | if not path.exists(list_fp): 282 | os.makedirs(list_fp) 283 | train_list_f = open(path.join(list_fp, 'train_%d_%s_%s.txt' % ( 284 | n_landmarks, target_person_size, suffix)), 'w') 285 | val_list_f = open(path.join(list_fp, 'val_%d_%s_%s.txt' % ( 286 | n_landmarks, target_person_size, suffix)), 'w') 287 | train_val_list_f = open(path.join(list_fp, 'trainval_%d_%s_%s.txt' % ( 288 | n_landmarks, target_person_size, suffix)), 'w') 289 | test_list_f = open(path.join(list_fp, 'test_%d_%s_%s.txt' % ( 290 | n_landmarks, target_person_size, suffix)), 'w') 291 | scale_f = open(path.join(list_fp, 'scale_%d_%s_%s.txt' % ( 292 | n_landmarks, target_person_size, suffix)), 'w') 293 | with open(path.join(up3d_fp, 'train.txt'), 'r') as f: 294 | train_spec = [line.strip() for line in f.readlines()] 295 | with open(path.join(up3d_fp, 'val.txt'), 'r') as f: 296 | val_spec = [line.strip() for line in f.readlines()] 297 | with open(path.join(up3d_fp, 'test.txt'), 'r') as f: 298 | test_spec = [line.strip() for line in f.readlines()] 299 | 300 | LOGGER.info("Processing...") 301 | list_ids = np.zeros((4,), dtype='int') 302 | add_dataset( 303 | dset_fp, 304 | dset_fromroot, 305 | list_ids, 306 | up3d_fp, 307 | train_list_f, 308 | val_list_f, 309 | train_val_list_f, 310 | test_list_f, 311 | scale_f, 312 | train_spec, 313 | val_spec, 314 | test_spec, 315 | target_person_size, landmark_mapping, 316 | crop, test_crop, 0, 317 | only_missing=only_missing, 318 | with_rlswap=with_rlswap, 319 | human_annotations=human_annotations) 320 | train_list_f.close() 321 | val_list_f.close() 322 | train_val_list_f.close() 323 | test_list_f.close() 324 | scale_f.close() 325 | LOGGER.info("Done.") 326 | 327 | 328 | if __name__ == '__main__': 329 | logging.basicConfig(level=logging.INFO, format=LOGFORMAT) 330 | logging.getLogger("opendr.lighting").setLevel(logging.WARN) 331 | cli() # pylint: disable=no-value-for-parameter 332 | 333 | -------------------------------------------------------------------------------- /pose/training/config/pose/n_classes.txt: -------------------------------------------------------------------------------- 1 | 91 2 | -------------------------------------------------------------------------------- /pose/training/config/pose/solver.prototxt: -------------------------------------------------------------------------------- 1 | train_net: "${EXP}/config/${NET_ID}/train_${TRAIN_SET}.prototxt" 2 | 3 | lr_policy: "multistep" 4 | multistep_lr: 0.0001 5 | stepvalue: 10000 6 | multistep_lr: 0.002 7 | stepvalue: 430000 8 | multistep_lr: 0.0002 9 | stepvalue: 730000 10 | multistep_lr: 0.0001 11 | 12 | display: 20 13 | max_iter: 1030000 14 | momentum: 0.9 15 | weight_decay: 0.0001 16 | 17 | snapshot: 5000 18 | snapshot_prefix: "${EXP}/model/${NET_ID}/train" 19 | solver_mode: GPU 20 | -------------------------------------------------------------------------------- /pose/training/config/pose/solver2.prototxt: -------------------------------------------------------------------------------- 1 | train_net: "${EXP}/config/${NET_ID}/train_${TRAIN_SET}.prototxt" 2 | 3 | lr_policy: "multistep" 4 | multistep_lr: 0.0001 5 | stepvalue: 10000 6 | multistep_lr: 0.002 7 | stepvalue: 430000 8 | multistep_lr: 0.0002 9 | stepvalue: 730000 10 | multistep_lr: 0.0001 11 | 12 | display: 20 13 | max_iter: 500000 14 | momentum: 0.9 15 | weight_decay: 0.0001 16 | 17 | snapshot: 50000 18 | snapshot_prefix: "${EXP}/model/${NET_ID}/train2" 19 | solver_mode: GPU 20 | -------------------------------------------------------------------------------- /pose/training/config/pose/target_person_size.txt: -------------------------------------------------------------------------------- 1 | 500 2 | -------------------------------------------------------------------------------- /requirements.txt: -------------------------------------------------------------------------------- 1 | numpy 2 | scipy 3 | sklearn 4 | joblib 5 | tqdm 6 | click 7 | opendr 8 | cython 9 | chumpy 10 | pillow 11 | fasteners 12 | h5py 13 | pymp-pypi 14 | meshio 15 | git+https://github.com/classner/clustertools.git#egg=clustertools 16 | -------------------------------------------------------------------------------- /segmentation/evaluate_segmentation.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python2 2 | """Get evaluation results for stored scoremaps.""" 3 | # pylint: disable=invalid-name 4 | from __future__ import print_function 5 | from os import path 6 | import logging 7 | 8 | import numpy as np 9 | from PIL import Image 10 | import scipy.ndimage 11 | import click 12 | import tqdm 13 | import pymp 14 | from clustertools.config import available_cpu_count # pylint: disable=unused-import 15 | from clustertools.log import LOGFORMAT 16 | 17 | from up_tools.model import regions, six_region_groups 18 | 19 | 20 | 21 | LOGGER = logging.getLogger(__name__) 22 | 23 | # pylint: disable=no-member 24 | VOC_REGION_GROUPS = [ 25 | [regions.ruhead, regions.rlhead, regions.luhead, regions.llhead], 26 | [regions.neck, regions.rubody, regions.rlbody, regions.lubody, 27 | regions.llbody, regions.rshoulder, regions.lshoulder], 28 | # Lower arms. 29 | [regions.rlarm, regions.rwrist, regions.rhand, 30 | regions.llarm, regions.lwrist, regions.lhand], 31 | [regions.ruarm, regions.luarm, regions.relbow, regions.lelbow], 32 | # Lower legs. 33 | [regions.rlleg, regions.rankle, regions.rfoot, 34 | regions.llleg, regions.lankle, regions.lfoot], 35 | [regions.ruleg, regions.luleg, regions.rknee, regions.lknee], 36 | ] 37 | 38 | 39 | # pylint: enable=no-member 40 | @click.command() 41 | @click.argument("image_list_file", type=click.Path(dir_okay=False, readable=True)) 42 | @click.argument("data_folder", type=click.Path(file_okay=False, readable=True)) 43 | @click.argument("result_label_folder", type=click.Path(file_okay=False, readable=True)) 44 | @click.argument("n_labels", type=click.INT) 45 | @click.option("--as_nolr", type=click.BOOL, is_flag=True, default=False, 46 | help="Evaluate 6 class body part segmentation without lr.") 47 | @click.option("--ev_31_as_6", type=click.BOOL, is_flag=True, default=False, 48 | help="Evaluate 31 region predictions as 6 regions (n_labels must be 6).") 49 | def main(image_list_file, # pylint: disable=too-many-locals, too-many-statements, too-many-branches, too-many-arguments 50 | data_folder, 51 | result_label_folder, 52 | n_labels, 53 | as_nolr=False, 54 | ev_31_as_6=False): 55 | """Perform the evaluation for previously written results scoremaps.""" 56 | LOGGER.info("Evaluating segmentation in folder `%s`.", result_label_folder) 57 | if 'voc' in image_list_file: 58 | voc_mode = True 59 | ev_31_as_6 = True 60 | n_labels = 7 61 | LOGGER.info("Using VOC part segmentation style.") 62 | else: 63 | voc_mode = False 64 | if ev_31_as_6: 65 | assert n_labels == 7 66 | if as_nolr: 67 | assert n_labels == 7 68 | classIOUs = np.zeros((n_labels,)) 69 | overallIOU = 0. 70 | overallAccuracy = 0. 71 | # Parallel stats. 72 | TP = pymp.shared.array((n_labels,), dtype='float32') 73 | FP = pymp.shared.array((n_labels,), dtype='float32') 74 | FN = pymp.shared.array((n_labels,), dtype='float32') 75 | imgTP = pymp.shared.array((1,), dtype='float32') 76 | imgPixels = pymp.shared.array((1,), dtype='float32') 77 | stat_lock = pymp.shared.lock() 78 | warned = False 79 | with open(image_list_file, 'r') as inf: 80 | image_list = inf.readlines() 81 | with pymp.Parallel(available_cpu_count(), if_=False and available_cpu_count() > 1) as p: 82 | for imgnames in p.iterate(tqdm.tqdm(image_list), element_timeout=20): # pylint: disable=too-many-nested-blocks 83 | imgname = imgnames.split(" ")[0].strip()[1:] 84 | old_applied_scale = float(imgnames.split(" ")[2].strip()) 85 | gtname = imgnames.split(" ")[1].strip()[1:] 86 | gt_file = path.join(data_folder, gtname) 87 | gtLabels = np.array(Image.open(gt_file)) 88 | if gtLabels.ndim == 3: 89 | if not warned: 90 | LOGGER.warn("Three-layer ground truth detected. Using first.") 91 | warned = True 92 | gtLabels = gtLabels[:, :, 0] 93 | gtLabels = scipy.misc.imresize(gtLabels, 1. / old_applied_scale, 94 | interp='nearest') 95 | if as_nolr: 96 | gtLabels[gtLabels == 4] = 3 97 | gtLabels[gtLabels == 6] = 5 98 | LOGGER.debug("Evaluating `%s`...", imgname) 99 | result_file = path.join(result_label_folder, 100 | path.basename(imgname) + '.npy') 101 | result_probs = np.load(result_file) 102 | result_probs = np.array( 103 | [scipy.misc.imresize(layer, 1. / old_applied_scale, interp='bilinear', mode='F') 104 | for layer in result_probs]) 105 | if result_probs.shape[0] > n_labels and not ( 106 | ev_31_as_6 and result_probs.shape[0] == 32): 107 | LOGGER.warn('Result has invalid labels: %s!', 108 | str(result_probs.shape[0])) 109 | continue 110 | if result_probs.min() < 0 or result_probs.max() > 1.: 111 | LOGGER.warn('Invalid result probabilities: min `%f`, max `%f`!', 112 | result_probs.min(), result_probs.max()) 113 | continue 114 | else: 115 | MAP = np.argmax(result_probs, axis=0) 116 | if as_nolr: 117 | MAP[MAP == 4] = 3 118 | MAP[MAP == 6] = 5 119 | if ev_31_as_6: 120 | if voc_mode: 121 | groups_to_use = VOC_REGION_GROUPS 122 | else: 123 | groups_to_use = six_region_groups 124 | for classID in range(1, 32): 125 | new_id = -1 126 | for group_idx, group in enumerate(groups_to_use): 127 | for grelem in group: 128 | if regions.reverse_mapping.keys().index(grelem) == classID - 1: # pylint: disable=no-member 129 | new_id = group_idx + 1 130 | assert new_id > 0 131 | if not voc_mode: 132 | gtLabels[gtLabels == classID] = new_id 133 | MAP[MAP == classID] = new_id 134 | for classID in range(n_labels): 135 | classGT = np.equal(gtLabels, classID) 136 | classResult = np.equal(MAP, classID) 137 | classResult[np.equal(gtLabels, 255)] = 0 138 | with stat_lock: 139 | TP[classID] = TP[classID] + np.count_nonzero(classGT & classResult) 140 | FP[classID] = FP[classID] + np.count_nonzero(classResult & ~classGT) 141 | FN[classID] = FN[classID] + np.count_nonzero(~classResult & classGT) 142 | imgResult = MAP 143 | imgGT = gtLabels 144 | imgResult[np.equal(MAP, 255)] = 0 145 | imgGT[np.equal(gtLabels, 255)] = 0 146 | with stat_lock: 147 | imgTP[0] += np.count_nonzero(np.equal(imgGT, imgResult)) 148 | imgPixels[0] += np.size(imgGT) 149 | for classID in range(0, n_labels): 150 | classIOUs[classID] = TP[classID] / (TP[classID] + FP[classID] + FN[classID]) 151 | if as_nolr: 152 | classIOUs = classIOUs[(0, 1, 2, 3, 5),] 153 | overallIOU = np.mean(classIOUs) 154 | overallAccuracy = imgTP[0] / imgPixels[0] 155 | if n_labels == 32: 156 | region_names = ['background'] + regions.reverse_mapping.values()[:-1] # pylint: disable=no-member 157 | LOGGER.info("Class IOUs:") 158 | for region_name, class_iou in zip(region_names, classIOUs): 159 | LOGGER.info('%s: %f', region_name, class_iou) 160 | else: 161 | LOGGER.info('Class IOUs: %s.', str(classIOUs)) 162 | LOGGER.info('Overall IOU: %s.', str(overallIOU)) 163 | LOGGER.info('Overall Accuracy: %s.', str(overallAccuracy)) 164 | 165 | 166 | if __name__ == '__main__': 167 | logging.basicConfig(level=logging.INFO, format=LOGFORMAT) 168 | main() # pylint: disable=no-value-for-parameter 169 | -------------------------------------------------------------------------------- /segmentation/run.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | 3 | trap "exit" INT 4 | 5 | CAFFE_BUILD_FP=$(../config.py DEEPLAB_BUILD_FP) 6 | CAFFE_BIN=${CAFFE_BUILD_FP}/install/bin/caffe 7 | CAFFE_PYTHON=${CAFFE_BUILD_FP}/install/python 8 | EXP=training 9 | if [ "${EXP}" = "training" ]; then 10 | DATA_ROOT=$(../config.py SEG_DATA_FP) 11 | else 12 | echo "Wrong exp name" 13 | exit 1 14 | fi 15 | 16 | ## Specify which model to train 17 | ########### training (standard training) ################ 18 | NET_ID=$2 19 | TRAIN_SET_SUFFIX=$3 20 | RESUME=$(grep -q resume <<<$4 && echo 1 || echo 0) 21 | DEV_ID=0 22 | RUN_TRAIN=`grep -q train <<<$1 && echo 1 || echo 0` 23 | RUN_TEST=`grep -q test <<<$1 && echo 1 || echo 0` 24 | RUN_EVALUATION=`grep -q evaluate <<<$1 && echo 1 || echo 0` 25 | RUN_TRAIN2=`grep -q trfull <<<$1 && echo 1 || echo 0` 26 | RUN_TEST2=`grep -q tefull <<<$1 && echo 1 || echo 0` 27 | RUN_EVALUATION2=`grep -q evfull <<<$1 && echo 1 || echo 0` 28 | ##### 29 | ## Create dirs 30 | CONFIG_DIR=${EXP}/config/${NET_ID} 31 | MODEL_DIR=${EXP}/model/${NET_ID} 32 | mkdir -p ${MODEL_DIR} 33 | LOG_DIR=${EXP}/log/${NET_ID} 34 | mkdir -p ${LOG_DIR} 35 | export GLOG_log_dir=${LOG_DIR} 36 | 37 | ## Training #1 (on train_aug) 38 | if [ ${RUN_TRAIN} -eq 1 ]; then 39 | export GLOG_minloglevel=0 40 | LIST_DIR=${EXP}/list 41 | NUM_CLASSES=`cat ${EXP}/config/${NET_ID}/n_classes.txt` 42 | ((NUM_LABELS=NUM_CLASSES+1)) # Add background. 43 | TARGET_PERSON_SIZE=`cat ${EXP}/config/${NET_ID}/target_person_size.txt` 44 | TRAIN_SET=train_${NUM_CLASSES}_${TARGET_PERSON_SIZE}_${TRAIN_SET_SUFFIX} 45 | if [ $RESUME -gt 0 ]; then 46 | MODEL=`ls -t ${EXP}/model/${NET_ID}/train_iter_*.caffemodel | head -n 1` 47 | SSTATE=`ls -t ${EXP}/model/${NET_ID}/train_iter_*.solverstate | head -n 1` 48 | if [ ! -f ${SSTATE} ]; then 49 | echo "Solverstate not found: ${SSTATE}!" 50 | exit 1 51 | fi 52 | else 53 | MODEL=${EXP}/config/${NET_ID}/init.caffemodel 54 | fi 55 | if [ ! -f ${MODEL} ]; then 56 | echo "Initialization model file not found: ${MODEL}!" 57 | exit 1 58 | fi 59 | echo Training net ${EXP}/${NET_ID} 60 | for pname in train solver; do 61 | sed "$(eval echo $(cat sub.sed))" \ 62 | ${CONFIG_DIR}/${pname}.prototxt > ${CONFIG_DIR}/${pname}_${TRAIN_SET}.prototxt 63 | done 64 | CMD="${CAFFE_BIN} train \ 65 | --solver=${CONFIG_DIR}/solver_${TRAIN_SET}.prototxt \ 66 | --gpu=${DEV_ID}" 67 | if [ -f ${MODEL} -a ${RESUME} -le 0 ]; then 68 | CMD="${CMD} --weights=${MODEL}" 69 | fi 70 | if [ ${RESUME} -gt 0 ]; then 71 | CMD="${CMD} --snapshot=${SSTATE}" 72 | fi 73 | echo Running ${CMD} && ${CMD} 74 | fi 75 | 76 | ## Test #1 specification (on val or test) 77 | if [ ${RUN_TEST} -eq 1 ]; then 78 | export GLOG_minloglevel=1 79 | for TEST_SET in val; do 80 | NUM_CLASSES=`cat ${EXP}/config/${NET_ID}/n_classes.txt` 81 | ((NUM_LABELS=NUM_CLASSES+1)) # Add background. 82 | TARGET_PERSON_SIZE=`cat ${EXP}/config/${NET_ID}/target_person_size.txt` 83 | TEST_SET=${TEST_SET}_${NUM_CLASSES}_${TARGET_PERSON_SIZE}_${TRAIN_SET_SUFFIX} 84 | TEST_ITER=`cat ${EXP}/list/${TEST_SET}.txt | wc -l` 85 | MODEL=${EXP}/model/${NET_ID}/test.caffemodel 86 | if [ ! -f ${MODEL} ]; then 87 | MODEL=`ls -t ${EXP}/model/${NET_ID}/train_iter_*.caffemodel | head -n 1` 88 | fi 89 | if [ ! -f ${MODEL} ]; then 90 | echo "Test model file not found: ${MODEL}!" 91 | exit 1 92 | fi 93 | echo Testing net ${EXP}/${NET_ID} 94 | FEATURE_DIR=${EXP}/features/${NET_ID} 95 | mkdir -p ${FEATURE_DIR}/${TEST_SET}/seg_score 96 | sed "$(eval echo $(cat sub.sed))" \ 97 | ${CONFIG_DIR}/testpy.prototxt > ${CONFIG_DIR}/testpy_${TEST_SET}.prototxt 98 | CMD="./store_segmentation_results.py \ 99 | ${CONFIG_DIR}/testpy_${TEST_SET}.prototxt \ 100 | ${MODEL} \ 101 | ${DATA_ROOT} \ 102 | ${EXP}/list/${TEST_SET}.txt \ 103 | ${FEATURE_DIR}/${TEST_SET}/seg_score \ 104 | --caffe_install_path ${CAFFE_PYTHON} \ 105 | --n_labels ${NUM_LABELS}" 106 | echo Running ${CMD} && ${CMD} 107 | done 108 | fi 109 | 110 | ## Evaluation #1 specification (on val or test) 111 | if [ ${RUN_EVALUATION} -eq 1 ]; then 112 | export GLOG_minloglevel=1 113 | for TEST_SET in val; do 114 | echo Evaluating net ${EXP}/${NET_ID} 115 | FEATURE_DIR=${EXP}/features/${NET_ID} 116 | NUM_CLASSES=`cat ${EXP}/config/${NET_ID}/n_classes.txt` 117 | ((NUM_LABELS=NUM_CLASSES+1)) # Add background. 118 | TARGET_PERSON_SIZE=`cat ${EXP}/config/${NET_ID}/target_person_size.txt` 119 | TEST_SET=${TEST_SET}_${NUM_CLASSES}_${TARGET_PERSON_SIZE}_${TRAIN_SET_SUFFIX} 120 | CMD="./evaluate_segmentation.py \ 121 | ${EXP}/list/${TEST_SET}.txt \ 122 | ${DATA_ROOT} \ 123 | ${FEATURE_DIR}/${TEST_SET}/seg_score \ 124 | ${NUM_LABELS}" 125 | echo Running ${CMD} && ${CMD} 126 | done 127 | fi 128 | 129 | ################################################################################ 130 | ## Training #2 (finetune on trainval_aug) 131 | if [ ${RUN_TRAIN2} -eq 1 ]; then 132 | export GLOG_minloglevel=0 133 | LIST_DIR=${EXP}/list 134 | NUM_CLASSES=`cat ${EXP}/config/${NET_ID}/n_classes.txt` 135 | ((NUM_LABELS=NUM_CLASSES+1)) 136 | TARGET_PERSON_SIZE=`cat ${EXP}/config/${NET_ID}/target_person_size.txt` 137 | TRAIN_SET=trainval_${NUM_CLASSES}_${TARGET_PERSON_SIZE}_${TRAIN_SET_SUFFIX} 138 | if [ $RESUME -gt 0 ]; then 139 | MODEL=`ls -t ${EXP}/model/${NET_ID}/train2_iter_*.caffemodel | head -n 1` 140 | SSTATE=`ls -t ${EXP}/model/${NET_ID}/train2_iter_*.solverstate | head -n 1` 141 | if [ ! -f ${SSTATE} ]; then 142 | echo "Solverstate not found: ${SSTATE}!" 143 | exit 1 144 | fi 145 | else 146 | MODEL=${EXP}/config/${NET_ID}/init2.caffemodel 147 | fi 148 | if [ ! -f ${MODEL} ]; then 149 | echo "WARNING! No init2 model file found at: ${MODEL}. Using latest..." 150 | MODEL=`ls -t ${EXP}/model/${NET_ID}/train_iter_*.caffemodel | head -n 1` 151 | fi 152 | echo Training2 net ${EXP}/${NET_ID} 153 | for pname in train solver2; do 154 | sed "$(eval echo $(cat sub.sed))" \ 155 | ${CONFIG_DIR}/${pname}.prototxt > ${CONFIG_DIR}/${pname}_${TRAIN_SET}.prototxt 156 | done 157 | CMD="${CAFFE_BIN} train \ 158 | --solver=${CONFIG_DIR}/solver2_${TRAIN_SET}.prototxt \ 159 | --gpu=${DEV_ID}" 160 | if [ $RESUME -le 0 ]; then 161 | CMD="${CMD} --weights=${MODEL}" 162 | fi 163 | if [ $RESUME -gt 0 ]; then 164 | CMD="${CMD} --snapshot=${SSTATE}" 165 | fi 166 | echo Running ${CMD} && ${CMD} 167 | fi 168 | 169 | ## Test #2 on official test set 170 | if [ ${RUN_TEST2} -eq 1 ]; then 171 | export GLOG_minloglevel=1 172 | for TEST_SET in test; do 173 | NUM_CLASSES=`cat ${EXP}/config/${NET_ID}/n_classes.txt` 174 | ((NUM_LABELS=NUM_CLASSES+1)) # Add background. 175 | TARGET_PERSON_SIZE=`cat ${EXP}/config/${NET_ID}/target_person_size.txt` 176 | TEST_SET=${TEST_SET}_${NUM_CLASSES}_${TARGET_PERSON_SIZE}_${TRAIN_SET_SUFFIX} 177 | TEST_ITER=`cat ${EXP}/list/${TEST_SET}.txt | wc -l` 178 | MODEL=${EXP}/model/${NET_ID}/test2.caffemodel 179 | if [ ! -f ${MODEL} ]; then 180 | MODEL=`ls -t ${EXP}/model/${NET_ID}/train2_iter_*.caffemodel | head -n 1` 181 | fi 182 | if [ ! -f ${MODEL} ]; then 183 | echo "Test model file not found: ${MODEL}!" 184 | exit 1 185 | fi 186 | echo Testing2 net ${EXP}/${NET_ID} 187 | FEATURE_DIR=${EXP}/features2/${NET_ID} 188 | mkdir -p ${FEATURE_DIR}/${TEST_SET}/seg_score 189 | sed "$(eval echo $(cat sub.sed))" \ 190 | ${CONFIG_DIR}/testpy.prototxt > ${CONFIG_DIR}/testpy_${TEST_SET}.prototxt 191 | CMD="./store_segmentation_results.py \ 192 | ${CONFIG_DIR}/testpy_${TEST_SET}.prototxt \ 193 | ${MODEL} \ 194 | ${DATA_ROOT} \ 195 | ${EXP}/list/${TEST_SET}.txt \ 196 | ${FEATURE_DIR}/${TEST_SET}/seg_score \ 197 | --caffe_install_path ${CAFFE_PYTHON} \ 198 | --n_labels ${NUM_LABELS}" 199 | echo Running ${CMD} && ${CMD} 200 | done 201 | fi 202 | 203 | ## Evaluation #2 specification (on val or test) 204 | if [ ${RUN_EVALUATION2} -eq 1 ]; then 205 | export GLOG_minloglevel=1 206 | for TEST_SET in test; do 207 | echo Evaluating net ${EXP}/${NET_ID} 208 | FEATURE_DIR=${EXP}/features2/${NET_ID} 209 | NUM_CLASSES=`cat ${EXP}/config/${NET_ID}/n_classes.txt` 210 | ((NUM_LABELS=NUM_CLASSES+1)) # Add background. 211 | TARGET_PERSON_SIZE=`cat ${EXP}/config/${NET_ID}/target_person_size.txt` 212 | TEST_SET=${TEST_SET}_${NUM_CLASSES}_${TARGET_PERSON_SIZE}_${TRAIN_SET_SUFFIX} 213 | CMD="./evaluate_segmentation.py \ 214 | ${EXP}/list/${TEST_SET}.txt \ 215 | ${DATA_ROOT} \ 216 | ${FEATURE_DIR}/${TEST_SET}/seg_score 217 | ${NUM_LABELS}" 218 | echo Running ${CMD} && ${CMD} 219 | done 220 | fi 221 | -------------------------------------------------------------------------------- /segmentation/store_segmentation_results.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python2 2 | """Write the segmentation results to disk for further evaluation.""" 3 | from os import path 4 | import sys 5 | import logging 6 | from PIL import Image 7 | sys.path.insert(0, path.abspath(path.join(path.dirname(__file__), '..'))) 8 | 9 | import numpy as np 10 | import click 11 | import tqdm 12 | 13 | from clustertools.visualization import apply_colormap 14 | from clustertools.log import LOGFORMAT 15 | 16 | from config import DEEPLAB_BUILD_FP # pylint: disable=import-error 17 | 18 | LOGGER = logging.getLogger(__name__) 19 | 20 | 21 | @click.command() 22 | @click.argument("caffe_prototxt", type=click.Path(dir_okay=False, readable=True)) 23 | @click.argument("caffe_model", type=click.Path(dir_okay=False, readable=True)) 24 | @click.argument("image_folder", type=click.Path(file_okay=False, readable=True)) 25 | @click.argument("image_list_file", type=click.Path(dir_okay=False, readable=True)) 26 | @click.argument("output_folder", type=click.Path(file_okay=False, writable=True)) 27 | @click.option("--caffe_install_path", 28 | type=click.Path(file_okay=False, readable=True), 29 | default=path.join(DEEPLAB_BUILD_FP, 'install/bin/python')) 30 | @click.option("--n_labels", type=click.INT, default=31) 31 | def main(caffe_prototxt, # pylint: disable=too-many-arguments, too-many-locals, too-many-statements 32 | caffe_model, 33 | image_folder, 34 | image_list_file, 35 | output_folder, 36 | caffe_install_path, 37 | n_labels): 38 | """Store and visualize the segmentation results for a model.""" 39 | LOGGER.info("Storing segmentation results to folder `%s`.", 40 | output_folder) 41 | LOGGER.info("Using caffe from `%s`.", caffe_install_path) 42 | sys.path.insert(0, path.join(caffe_install_path)) 43 | import caffe # pylint: disable=import-error 44 | mean_red = 122.675 45 | mean_green = 116.669 46 | mean_blue = 104.008 47 | 48 | # Configure preprocessing 49 | caffe.set_mode_gpu() 50 | net_full_conv = caffe.Net(caffe_prototxt, caffe_model, caffe.TEST) 51 | net_input_blob = net_full_conv.inputs[0] 52 | transformer = caffe.io.Transformer({ 53 | net_full_conv.inputs[0]: net_full_conv.blobs[net_full_conv.inputs[0]].data.shape}) 54 | transformer.set_transpose(net_input_blob, (2, 0, 1)) 55 | transformer.set_channel_swap(net_input_blob, (2, 1, 0)) 56 | transformer.set_raw_scale(net_input_blob, 255.0) 57 | net_inp_height, net_inp_width = net_full_conv.blobs[net_input_blob].data.shape[2:4] 58 | # Create and configure the mean image. The transformer applies channel-swap 59 | # first, so we have BGR order for the mean image. 60 | mean_image = np.zeros((3, net_inp_height, net_inp_width), dtype='float32') 61 | mean_image[0, :, :] = mean_blue 62 | mean_image[1, :, :] = mean_green 63 | mean_image[2, :, :] = mean_red 64 | transformer.set_mean(net_input_blob, mean_image) 65 | with open(image_list_file, 'r') as inf: 66 | image_list = inf.readlines() 67 | for imgnames in tqdm.tqdm(image_list): 68 | imgname = imgnames.split(" ")[0][1:].strip() 69 | LOGGER.debug("Processing `%s`...", imgname) 70 | image_filename = path.join(image_folder, imgname) 71 | # caffe.io loads as RGB, and in range [0., 1.]. 72 | im = caffe.io.load_image(image_filename) # pylint: disable=invalid-name 73 | height, width = im.shape[:2] 74 | # Pad values. 75 | pad_width = net_inp_width - width 76 | pad_height = net_inp_height - height 77 | im = np.lib.pad(im, # pylint: disable=invalid-name 78 | ((0, pad_height), 79 | (0, pad_width), 80 | (0, 0)), 81 | 'constant', 82 | constant_values=-5) 83 | assert im.shape[0] == net_inp_height 84 | assert im.shape[1] == net_inp_width 85 | R = im[:, :, 0] # pylint: disable=invalid-name 86 | G = im[:, :, 1] # pylint: disable=invalid-name 87 | B = im[:, :, 2] # pylint: disable=invalid-name 88 | # Will be multiplied by 255 by the transformer. 89 | R[R == -5] = mean_red / 255. 90 | G[G == -5] = mean_green / 255. 91 | B[B == -5] = mean_blue / 255. 92 | im[:, :, 0] = R 93 | im[:, :, 1] = G 94 | im[:, :, 2] = B 95 | out = net_full_conv.forward_all( 96 | data=np.asarray([transformer.preprocess(net_input_blob, im)])) 97 | pmap = out['prob'][0] 98 | assert pmap.min() >= 0. and pmap.max() <= 1., ( 99 | "Invalid probability value in result map!") 100 | prob_map = pmap[:, :height, :width] 101 | np.save(path.join(output_folder, 102 | path.basename(imgname) + '.npy'), 103 | prob_map) 104 | maxed_map = np.argmax(prob_map, axis=0) 105 | vis_image = Image.fromarray(apply_colormap(maxed_map, vmax=n_labels-1)) 106 | vis_image.save(path.join(output_folder, 107 | path.basename(imgname) + '.npy.vis.png')) 108 | raw_image = Image.fromarray(maxed_map.astype('uint8')) 109 | raw_image.save(path.join(output_folder, 110 | path.basename(imgname) + '_segmentation.png')) 111 | 112 | 113 | if __name__ == '__main__': 114 | logging.basicConfig(level=logging.INFO, format=LOGFORMAT) 115 | main() # pylint: disable=no-value-for-parameter 116 | -------------------------------------------------------------------------------- /segmentation/sub.sed: -------------------------------------------------------------------------------- 1 | 's,${DATA_ROOT},'"${DATA_ROOT}"',g;s,${DATA_ROOT1},'"${DATA_ROOT1}"',g;s,${EXP},'"${EXP}"',g;s,${EXP1},'"${EXP1}"',g;s,${TRAIN_SET},'"${TRAIN_SET}"',g;s,${TRAIN_SET1},'"${TRAIN_SET1}"',g;s,${TRAIN_SET1_WEAK},'"${TRAIN_SET1_WEAK}"',g;s,${TRAIN_SET1_STRONG},'"${TRAIN_SET1_STRONG}"',g;s,${TEST_SET},'"${TEST_SET}"',g;s,${NET_ID},'"${NET_ID}"',g;s,${FEATURE_DIR},'"${FEATURE_DIR}"',g;s,${NUM_LABELS},'"${NUM_LABELS}"',g;s,${NUM_LABELS1},'"${NUM_LABELS1}"',g;s,${NUM_LABELS_UNION},'"${NUM_LABELS_UNION}"',g;s,${BG_BIAS},'"${BG_BIAS}"',g;s,${FG_BIAS},'"${FG_BIAS}"',g;s,${TRAIN_SET_STRONG},'"${TRAIN_SET_STRONG}"',g;s,${TRAIN_SET_WEAK},'"${TRAIN_SET_WEAK}"',g;s,${BATCH_SIZE},'"${BATCH_SIZE}"',g;s,${TEST_SET_PREFIX},'"${TEST_SET_PREFIX}"',g;s,${TRAIN_STEP},'"${TRAIN_STEP}"',g' 2 | -------------------------------------------------------------------------------- /segmentation/tools/create_dataset.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python2 2 | """Create segmentation datasets from select SMPL fits.""" 3 | import os 4 | import os.path as path 5 | import sys 6 | import logging 7 | 8 | import numpy as np 9 | import scipy 10 | import click 11 | import tqdm 12 | 13 | from clustertools.log import LOGFORMAT 14 | from clustertools.visualization import apply_colormap 15 | from up_tools.model import (robust_person_size, six_region_groups, 16 | regions_to_classes, get_crop) 17 | 18 | from up_tools.render_segmented_views import render_body_impl 19 | sys.path.insert(0, path.join(path.dirname(__file__), '..', '..')) 20 | from config import SEG_DATA_FP, UP3D_FP 21 | 22 | 23 | LOGGER = logging.getLogger(__name__) 24 | DSET_ROOT_FP = SEG_DATA_FP 25 | 26 | if not path.exists(DSET_ROOT_FP): 27 | os.mkdir(DSET_ROOT_FP) 28 | 29 | 30 | def uncrop(annot, fullimsize, cropinfo): 31 | if annot.ndim == 2: 32 | res = np.zeros((fullimsize[0], fullimsize[1]), dtype='uint8') 33 | else: 34 | res = np.ones((fullimsize[0], fullimsize[1], 3), dtype='uint8') * 255 35 | res[cropinfo[2]:cropinfo[3], 36 | cropinfo[4]:cropinfo[5]] = scipy.misc.imresize( 37 | annot, 38 | (cropinfo[3] - cropinfo[2], 39 | cropinfo[5] - cropinfo[4]), 40 | interp='nearest') 41 | return res 42 | 43 | 44 | def add_dataset(dset_fp, dset_rel_fp, up3d_fp, # pylint: disable=too-many-locals, too-many-arguments, too-many-statements, too-many-branches 45 | train_list_f, val_list_f, test_list_f, 46 | train_spec, val_spec, test_spec, 47 | target_person_size, partspec, crop, running_idx, 48 | only_missing=False): 49 | """Add a dataset to the collection.""" 50 | test_ids = [int(id_[1:6]) for id_ in test_spec] 51 | train_ids = [int(id_[1:6]) for id_ in train_spec] 52 | val_ids = [int(id_[1:6]) for id_ in val_spec] 53 | LOGGER.info("Split: %d train, %d val, %d test.", 54 | len(train_ids), len(val_ids), len(test_ids)) 55 | LOGGER.info("Writing dataset...") 56 | for im_idx in tqdm.tqdm(train_ids + val_ids + test_ids): 57 | image = scipy.misc.imread(path.join(up3d_fp, '%05d_image.png' % (im_idx))) 58 | with open(path.join(up3d_fp, '%05d_fit_crop_info.txt' % (im_idx)), 'r') as inf: 59 | cropinfo = [int(val) for val in inf.readline().strip().split()] 60 | assert image.ndim == 3 61 | out_exists = (path.exists(path.join(dset_fp, '%05d_image.png' % (running_idx))) and 62 | path.exists(path.join(dset_fp, '%05d_ann.png' % (running_idx))) and 63 | path.exists(path.join(dset_fp, '%05d_ann_vis.png' % (running_idx))) and 64 | path.exists(path.join(dset_fp, '%05d_render.png' % (running_idx))) and 65 | path.exists(path.join(dset_fp, '%05d_render_light.png' % (running_idx)))) 66 | if not (only_missing and out_exists): 67 | rendering = uncrop(render_body_impl(path.join(up3d_fp, '%05d_body.pkl' % (im_idx)), 68 | resolution=(cropinfo[1], 69 | cropinfo[0]), 70 | quiet=True, 71 | use_light=False)[0], 72 | image.shape[:2], 73 | cropinfo) 74 | rendering_l = uncrop(render_body_impl(path.join(up3d_fp, '%05d_body.pkl' % (im_idx)), 75 | resolution=(cropinfo[1], 76 | cropinfo[0]), 77 | quiet=True, 78 | use_light=True)[0], 79 | image.shape[:2], 80 | cropinfo) 81 | joints = np.load(path.join(up3d_fp, '%05d_joints.npy' % (im_idx))) 82 | joints = np.vstack((joints, np.all(joints > 0, axis=0)[None, :])) 83 | person_size = robust_person_size(joints) 84 | norm_factor = float(target_person_size) / person_size 85 | if not (only_missing and out_exists): 86 | image = scipy.misc.imresize(image, norm_factor, interp='bilinear') 87 | rendering = scipy.misc.imresize(rendering, norm_factor, interp='nearest') 88 | rendering_l = scipy.misc.imresize(rendering_l, norm_factor, interp='bilinear') 89 | if image.shape[0] > crop or image.shape[1] > crop: 90 | LOGGER.debug("Image (original %d, here %d) too large (%s)! Cropping...", 91 | im_idx, running_idx, str(image.shape[:2])) 92 | person_center = np.mean(joints[:2, joints[2, :] == 1], axis=1) * norm_factor 93 | crop_y, crop_x = get_crop(image, person_center, crop) 94 | image = image[crop_y[0]:crop_y[1], 95 | crop_x[0]:crop_x[1], :] 96 | rendering = rendering[crop_y[0]:crop_y[1], 97 | crop_x[0]:crop_x[1], :] 98 | rendering_l = rendering_l[crop_y[0]:crop_y[1], 99 | crop_x[0]:crop_x[1], :] 100 | assert image.shape[0] == crop or image.shape[1] == crop, ( 101 | "Error cropping image (original %d, here %d)!" % (im_idx, 102 | running_idx)) 103 | assert image.shape[0] <= crop and image.shape[1] <= crop and image.shape[2] == 3, ( 104 | "Wrong image shape (original %d, here %d)!" % (im_idx, running_idx)) 105 | class_groups = six_region_groups if partspec == '6' else None 106 | annotation = regions_to_classes(rendering, class_groups, warn_id=str(im_idx)) 107 | if partspec == '1': 108 | annotation = (annotation > 0).astype('uint8') 109 | assert np.max(annotation) <= int(partspec), ( 110 | "Wrong annotation value (original %d, here %d): %s!" % ( 111 | im_idx, running_idx, str(np.unique(annotation)))) 112 | if running_idx == 0: 113 | assert np.max(annotation) == int(partspec), ( 114 | "Probably an error in the number of parts!") 115 | scipy.misc.imsave(path.join(dset_fp, '%05d_image.png' % (running_idx)), image) 116 | scipy.misc.imsave(path.join(dset_fp, '%05d_ann.png' % (running_idx)), annotation) 117 | scipy.misc.imsave(path.join(dset_fp, '%05d_ann_vis.png' % (running_idx)), 118 | apply_colormap(annotation, vmax=int(partspec))) 119 | scipy.misc.imsave(path.join(dset_fp, '%05d_render.png' % (running_idx)), rendering) 120 | scipy.misc.imsave(path.join(dset_fp, '%05d_render_light.png' % (running_idx)), rendering_l) # pylint: disable=line-too-long 121 | if im_idx in train_ids: 122 | list_f = train_list_f 123 | elif im_idx in val_ids: 124 | list_f = val_list_f 125 | elif im_idx in test_ids: 126 | list_f = test_list_f 127 | list_f.write("/%s/%05d_image.png /%s/%05d_ann.png %f\n" % ( 128 | dset_rel_fp, running_idx, dset_rel_fp, running_idx, norm_factor)) 129 | list_f.flush() 130 | running_idx += 1 131 | return running_idx 132 | 133 | 134 | @click.command() 135 | @click.argument("suffix", type=click.STRING) 136 | @click.argument("partspec", type=click.Choice(['1', '6', '31'])) 137 | @click.argument("target_person_size", type=click.INT) 138 | @click.option("--crop", type=click.INT, default=513, 139 | help="Crop size for the images.") 140 | @click.option("--only_missing", type=click.BOOL, default=False, is_flag=True, 141 | help="Only rewrite missing images.") 142 | @click.option("--up3d_fp", type=click.Path(file_okay=False, readable=True), 143 | default=UP3D_FP, 144 | help="Path to the UP3D folder that you want to use.") 145 | def cli(suffix, partspec, target_person_size, crop=513, only_missing=False, up3d_fp=UP3D_FP): # pylint: disable=too-many-locals, too-many-arguments 146 | """Create segmentation datasets from select SMPL fits.""" 147 | np.random.seed(1) 148 | LOGGER.info("Creating segmentation dataset for %s classes with target " 149 | "person size %f and suffix `%s`.", 150 | partspec, target_person_size, suffix) 151 | assert ' ' not in suffix 152 | dset_fromroot = path.join(partspec, str(target_person_size), suffix) 153 | dset_fp = path.join(DSET_ROOT_FP, dset_fromroot) 154 | if path.exists(dset_fp): 155 | if not only_missing: 156 | if not click.confirm("Dataset folder exists: `%s`! Continue?" % (dset_fp)): 157 | return 158 | else: 159 | os.makedirs(dset_fp) 160 | LOGGER.info("Creating list files...") 161 | list_fp = path.join(path.dirname(__file__), '..', 'training', 'list') 162 | if not path.exists(list_fp): 163 | os.makedirs(list_fp) 164 | train_list_f = open(path.join(list_fp, 'train_%s_%d_%s.txt' % ( 165 | partspec, target_person_size, suffix)), 'w') 166 | val_list_f = open(path.join(list_fp, 'val_%s_%d_%s.txt' % ( 167 | partspec, target_person_size, suffix)), 'w') 168 | test_list_f = open(path.join(list_fp, 'test_%s_%d_%s.txt' % ( 169 | partspec, target_person_size, suffix)), 'w') 170 | with open(path.join(up3d_fp, 'train.txt'), 'r') as f: 171 | train_spec = [line.strip() for line in f.readlines()] 172 | with open(path.join(up3d_fp, 'val.txt'), 'r') as f: 173 | val_spec = [line.strip() for line in f.readlines()] 174 | with open(path.join(up3d_fp, 'test.txt'), 'r') as f: 175 | test_spec = [line.strip() for line in f.readlines()] 176 | 177 | LOGGER.info("Processing...") 178 | add_dataset( 179 | dset_fp, 180 | dset_fromroot, 181 | up3d_fp, 182 | train_list_f, val_list_f, test_list_f, 183 | train_spec, val_spec, test_spec, 184 | target_person_size, partspec, 185 | crop, 0, 186 | only_missing=only_missing) 187 | train_list_f.close() 188 | val_list_f.close() 189 | test_list_f.close() 190 | LOGGER.info("Creating trainval file...") 191 | trainval_list_f = open(path.join(list_fp, 'trainval_%s_%d_%s.txt' % ( 192 | partspec, target_person_size, suffix)), 'w') 193 | train_list_f = open(path.join(list_fp, 'train_%s_%d_%s.txt' % ( 194 | partspec, target_person_size, suffix)), 'r') 195 | val_list_f = open(path.join(list_fp, 'val_%s_%d_%s.txt' % ( 196 | partspec, target_person_size, suffix)), 'r') 197 | for line in train_list_f: 198 | trainval_list_f.write(line) 199 | for line in val_list_f: 200 | trainval_list_f.write(line) 201 | trainval_list_f.close() 202 | train_list_f.close() 203 | val_list_f.close() 204 | LOGGER.info("Done.") 205 | 206 | 207 | if __name__ == '__main__': 208 | logging.basicConfig(level=logging.INFO, format=LOGFORMAT) 209 | logging.getLogger("opendr.lighting").setLevel(logging.WARN) 210 | cli() # pylint: disable=no-value-for-parameter 211 | -------------------------------------------------------------------------------- /segmentation/tools/render_orig_colors.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python2 2 | """Create a rendering from a segmentation in the original 31 colors.""" 3 | import os.path as path 4 | import logging 5 | from glob import glob 6 | from scipy.misc import imread, imsave 7 | import numpy as np 8 | import click 9 | import up_tools.model as mdl 10 | from clustertools.log import LOGFORMAT 11 | import tqdm 12 | 13 | 14 | LOGGER = logging.getLogger(__name__) 15 | 16 | 17 | @click.command() 18 | @click.argument('seg_fp', type=click.Path(exists=True, readable=True, 19 | writable=True, file_okay=False)) 20 | @click.option('--image_fp', type=click.Path(exists=True, readable=True), 21 | help='Look for images here.', default=None) 22 | def cli(seg_fp, image_fp=None): 23 | """Create a rendering from a segmentation in the original 31 colors.""" 24 | if image_fp is None: 25 | image_fp = seg_fp 26 | LOGGER.info("Applying original colormap on segmentations in folder `%s`.", 27 | seg_fp) 28 | for segim_fp in tqdm.tqdm(glob(path.join(seg_fp, '*_segmentation.png'))): 29 | seg = imread(segim_fp) 30 | im = imread(path.join(image_fp, # pylint: disable=invalid-name 31 | path.basename(segim_fp)[:-len('_segmentation.png')])) 32 | assert np.all(seg.shape[:2] == im.shape[:2]) 33 | # Map back the colors. 34 | seg_orig = np.ones_like(im) * 255 35 | for y_idx in range(seg.shape[0]): 36 | for x_idx in range(seg.shape[1]): 37 | if seg[y_idx, x_idx] > 0: 38 | seg_orig[y_idx, x_idx, :] = mdl.regions.reverse_mapping.keys()[ # pylint: disable=no-member 39 | seg[y_idx, x_idx] - 1] 40 | # Save. 41 | imsave(segim_fp + '_orig_vis.png', seg_orig) 42 | # Blend. 43 | imsave(segim_fp + '_orig_blend_vis.png', 44 | (seg_orig.astype('float32') * 0.5 + 45 | im.astype('float32') * 0.5).astype('uint8')) 46 | # Blend only on foreground. 47 | blend_fg = im.copy() 48 | fg_regions = np.dstack([(seg != 0)[:, :, None] for _ in range(3)]) 49 | blend_fg[fg_regions] = ((seg_orig[fg_regions].astype('float32') * 0.5 + 50 | im[fg_regions].astype('float32') * 0.5).astype('uint8')) 51 | imsave(segim_fp + '_orig_blend_fg_vis.png', blend_fg) 52 | 53 | 54 | if __name__ == '__main__': 55 | logging.basicConfig(level=logging.INFO, format=LOGFORMAT) 56 | cli() # pylint: disable=no-value-for-parameter 57 | -------------------------------------------------------------------------------- /segmentation/training/config/segmentation/n_classes.txt: -------------------------------------------------------------------------------- 1 | 31 2 | -------------------------------------------------------------------------------- /segmentation/training/config/segmentation/solver.prototxt: -------------------------------------------------------------------------------- 1 | train_net: "${EXP}/config/${NET_ID}/train_${TRAIN_SET}.prototxt" 2 | 3 | iter_size: 10 4 | lr_policy: "poly" 5 | power: 0.9 6 | base_lr: 2.5e-4 7 | 8 | average_loss: 20 9 | display: 20 10 | max_iter: 30000 11 | momentum: 0.9 12 | weight_decay: 0.0005 13 | 14 | snapshot: 2000 15 | snapshot_prefix: "${EXP}/model/${NET_ID}/train" 16 | solver_mode: GPU -------------------------------------------------------------------------------- /segmentation/training/config/segmentation/solver2.prototxt: -------------------------------------------------------------------------------- 1 | train_net: "${EXP}/config/${NET_ID}/train_${TRAIN_SET}.prototxt" 2 | 3 | iter_size: 10 4 | lr_policy: "poly" 5 | power: 0.9 6 | base_lr: 2.5e-4 7 | 8 | average_loss: 20 9 | display: 20 10 | max_iter: 30000 11 | momentum: 0.9 12 | weight_decay: 0.0005 13 | 14 | snapshot: 2000 15 | snapshot_prefix: "${EXP}/model/${NET_ID}/train2" 16 | solver_mode: GPU -------------------------------------------------------------------------------- /segmentation/training/config/segmentation/target_person_size.txt: -------------------------------------------------------------------------------- 1 | 500 2 | -------------------------------------------------------------------------------- /setup.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python2 2 | # -*- coding: utf-8 -*- 3 | """ 4 | The setup script for the entire project. 5 | @author: Christoph Lassner 6 | """ 7 | from setuptools import setup 8 | from pip.req import parse_requirements 9 | 10 | VERSION = '1.0' 11 | REQS = [str(ir.req) for ir in parse_requirements('requirements.txt', 12 | session='tmp')] 13 | 14 | setup( 15 | name='up_tools', 16 | author='Christoph Lassner', 17 | author_email='mail@christophlassner.de', 18 | packages=['up_tools'], 19 | test_suite='tests', 20 | dependency_links=['http://github.com/classner/clustertools/tarball/master#egg=clustertools'], 21 | include_package_data=True, 22 | install_requires=REQS, 23 | version=VERSION, 24 | license='Creative Commons Non-Commercial 4.0', 25 | ) 26 | -------------------------------------------------------------------------------- /tests/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/classner/up/7ec281c62afae764fa56a649f41a02487ea18c58/tests/__init__.py -------------------------------------------------------------------------------- /tests/tools.py: -------------------------------------------------------------------------------- 1 | """Tool unittests.""" 2 | # pylint: disable=invalid-name 3 | import os 4 | import os.path as path 5 | import unittest 6 | import numpy as np 7 | 8 | 9 | class MeshTest(unittest.TestCase): 10 | 11 | """Test Mesh serialization.""" 12 | 13 | def test_constructor(self): 14 | """Test object construction.""" 15 | import up_tools.mesh as upm 16 | # From file. 17 | tm = upm.Mesh(filename=path.join(path.dirname(__file__), 18 | '..', 19 | 'models', 20 | '3D', 21 | 'template-bodyparts.ply')) 22 | # From data. 23 | v = np.zeros((3, 3)) 24 | vc = np.zeros((3, 3)) 25 | f = np.zeros((1, 3)) 26 | tm = upm.Mesh(v=v, vc=vc, f=f) 27 | self.assertTrue(np.all(v == tm.v)) 28 | self.assertTrue(np.all(vc == tm.vc)) 29 | self.assertTrue(np.all(f == tm.f)) 30 | 31 | def test_serialization(self): 32 | """Test serialization.""" 33 | import up_tools.mesh as upm 34 | tm = upm.Mesh(filename=path.join(path.dirname(__file__), 35 | '..', 36 | 'models', 37 | '3D', 38 | 'template-bodyparts.ply')) 39 | self.assertTrue(tm.v is not None) 40 | self.assertTrue(tm.vc is not None) 41 | self.assertTrue(tm.f is not None) 42 | tm.write_ply('test_out.ply') 43 | rm = upm.Mesh(filename='test_out.ply') 44 | self.assertTrue(np.all(tm.v == rm.v)) 45 | self.assertTrue(np.all(tm.vc == rm.vc)) 46 | self.assertTrue(np.all(tm.f == rm.f)) 47 | os.remove('test_out.ply') 48 | 49 | 50 | if __name__ == '__main__': 51 | unittest.main() 52 | -------------------------------------------------------------------------------- /up_tools/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/classner/up/7ec281c62afae764fa56a649f41a02487ea18c58/up_tools/__init__.py -------------------------------------------------------------------------------- /up_tools/bake_vertex_colors.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python2 2 | """Bake the vertex colors by superposition.""" 3 | from copy import copy 4 | import numpy as np 5 | import click 6 | 7 | import opendr.renderer as _odr_r 8 | import opendr.camera as _odr_c 9 | from up_tools.mesh import Mesh 10 | 11 | 12 | def bake_vertex_colors(inmesh): 13 | """Bake the vertex colors by superposition.""" 14 | faces = np.arange(inmesh.f.size).reshape(-1, 3) 15 | vertices = np.empty((len(inmesh.f)*3, 3)) 16 | vc = np.zeros_like(vertices) # pylint: disable=invalid-name 17 | tmpar = np.ascontiguousarray(inmesh.vc).view( 18 | np.dtype((np.void, inmesh.vc.dtype.itemsize * inmesh.vc.shape[1]))) 19 | _, unique_idx = np.unique(tmpar, return_index=True) 20 | unique_clrs = inmesh.vc[unique_idx] 21 | for iface, face in enumerate(inmesh.f): 22 | vertices[iface*3+0] = inmesh.v[face[0]] 23 | vertices[iface*3+1] = inmesh.v[face[1]] 24 | vertices[iface*3+2] = inmesh.v[face[2]] 25 | low_idx = np.argmin([np.linalg.norm(inmesh.vc[face[0]]), 26 | np.linalg.norm(inmesh.vc[face[1]]), 27 | np.linalg.norm(inmesh.vc[face[2]])]) 28 | vc[iface*3+0] = inmesh.vc[face[low_idx]] 29 | vc[iface*3+1] = inmesh.vc[face[low_idx]] 30 | vc[iface*3+2] = inmesh.vc[face[low_idx]] 31 | tmpar = np.ascontiguousarray(vc).view(np.dtype((np.void, vc.dtype.itemsize * vc.shape[1]))) 32 | _, unique_idx = np.unique(tmpar, return_index=True) 33 | unique_clrs_after = vc[unique_idx] 34 | for clr in unique_clrs: 35 | assert clr in unique_clrs_after 36 | for clr in unique_clrs_after: 37 | assert clr in unique_clrs 38 | outmesh = Mesh(v=vertices, f=faces, vc=vc) 39 | return outmesh 40 | 41 | 42 | def get_face_for_pixel(shape, mesh, model, configuration, coords): # pylint: disable=too-many-locals 43 | """Get the face index or -1 for the mesh in the given conf at coords.""" 44 | assert len(shape) == 2, str(shape) 45 | assert np.all(coords >= 0), str(coords) 46 | assert coords.ndim == 2, str(coords.ndim) 47 | for coord in coords.T: 48 | assert coord[0] < shape[1], "%s, %s" % (str(coord), str(shape)) 49 | assert coord[1] < shape[0], "%s, %s" % (str(coord), str(shape)) 50 | mesh = copy(mesh) 51 | # Setup the model. 52 | model.betas[:len(configuration['betas'])] = configuration['betas'] 53 | model.pose[:] = configuration['pose'] 54 | model.trans[:] = configuration['trans'] 55 | mesh.v = model.r 56 | inmesh = mesh 57 | # Assign a different color for each face. 58 | faces = np.arange(inmesh.f.size).reshape(-1, 3) 59 | vertices = np.empty((len(inmesh.f)*3, 3)) 60 | vc = np.zeros_like(vertices) # pylint: disable=invalid-name 61 | for iface, face in enumerate(inmesh.f): 62 | vertices[iface*3+0] = inmesh.v[face[0]] 63 | vertices[iface*3+1] = inmesh.v[face[1]] 64 | vertices[iface*3+2] = inmesh.v[face[2]] 65 | vc[iface*3+0] = (float(iface % 255) / 255., float(iface / 255) / 255., 0.) 66 | vc[iface*3+1] = (float(iface % 255) / 255., float(iface / 255) / 255., 0.) 67 | vc[iface*3+2] = (float(iface % 255) / 255., float(iface / 255) / 255., 0.) 68 | fcmesh = Mesh(v=vertices, f=faces, vc=vc) 69 | # Render the mesh. 70 | dist = np.abs(configuration['t'][2] - np.mean(fcmesh.v, axis=0)[2]) 71 | rn = _odr_r.ColoredRenderer() # pylint: disable=redefined-variable-type, invalid-name 72 | rn.camera = _odr_c.ProjectPoints( 73 | rt=configuration['rt'], 74 | t=configuration['t'], 75 | f=np.array([configuration['f'], configuration['f']]), 76 | c=np.array([shape[1], shape[0]]) / 2., 77 | k=np.zeros(5)) 78 | rn.frustum = {'near': 1., 'far': dist + 20., 'height': shape[0], 'width': shape[1]} 79 | rn.set(v=fcmesh.v, f=fcmesh.f, vc=fcmesh.vc, bgcolor=np.ones(3)) 80 | rendered = rn.r 81 | results = [-1 for _ in range(len(coords.T))] 82 | for coord_idx, coord in enumerate(coords.T): 83 | # Find the face or background. 84 | loc_color = (rendered[int(coord[1]), int(coord[0])] * 255.).astype('uint8') 85 | if np.all(loc_color == 255): 86 | continue 87 | else: 88 | assert loc_color[2] == 0, str(loc_color) 89 | face_idx = loc_color[1] * 255 + loc_color[0] 90 | assert face_idx >= 0 and face_idx < len(mesh.f) 91 | results[coord_idx] = face_idx 92 | return results 93 | 94 | 95 | @click.command() 96 | @click.argument("inmesh_fp", type=click.Path(exists=True, dir_okay=False, readable=True)) 97 | @click.argument("outmesh_fp", type=click.Path(dir_okay=False, writable=True)) 98 | def cli(inmesh_fp, outmesh_fp): 99 | """Bake the vertex colors by superposition.""" 100 | inmesh = Mesh(filename=inmesh_fp) 101 | outmesh = bake_vertex_colors(inmesh) 102 | outmesh.write_ply(outmesh_fp) 103 | 104 | 105 | if __name__ == '__main__': 106 | cli() # pylint: disable=no-value-for-parameter 107 | -------------------------------------------------------------------------------- /up_tools/camera.py: -------------------------------------------------------------------------------- 1 | """Camera tools.""" 2 | # pylint: disable=invalid-name, bad-whitespace 3 | import numpy as np 4 | 5 | 6 | def rotateY(points, angle): 7 | """Rotate all points in a 2D array around the y axis.""" 8 | ry = np.array([ 9 | [np.cos(angle), 0., np.sin(angle)], 10 | [0., 1., 0. ], 11 | [-np.sin(angle), 0., np.cos(angle)] 12 | ]) 13 | return np.dot(points, ry) 14 | 15 | def rotateX( points, angle ): 16 | """Rotate all points in a 2D array around the x axis.""" 17 | rx = np.array([ 18 | [1., 0., 0. ], 19 | [0., np.cos(angle), -np.sin(angle)], 20 | [0., np.sin(angle), np.cos(angle) ] 21 | ]) 22 | return np.dot(points, rx) 23 | 24 | def rotateZ( points, angle ): 25 | """Rotate all points in a 2D array around the z axis.""" 26 | rz = np.array([ 27 | [np.cos(angle), -np.sin(angle), 0. ], 28 | [np.sin(angle), np.cos(angle), 0. ], 29 | [0., 0., 1. ] 30 | ]) 31 | return np.dot(points, rz) 32 | -------------------------------------------------------------------------------- /up_tools/capsule_ch.py: -------------------------------------------------------------------------------- 1 | """ 2 | Copyright 2016 Max Planck Society, Federica Bogo, Angjoo Kanazawa. All 3 | rights reserved. 4 | 5 | This software is provided for research purposes only. 6 | 7 | By using this software you agree to the terms of the SMPLify license here: 8 | 9 | http://smplify.is.tue.mpg.de/license 10 | """ 11 | import numpy as np 12 | import chumpy as ch 13 | from opendr.geometry import Rodrigues 14 | 15 | 16 | cap_f = np.asarray([[ 0, 7, 6], [ 1, 7, 9], [ 0, 6, 11], [ 0, 11, 13], 17 | [ 0, 13, 10], [ 1, 9, 16], [ 2, 8, 18], [ 3, 12, 20], 18 | [ 4, 14, 22], [ 5, 15, 24], [ 1, 16, 19], [ 2, 18, 21], 19 | [ 3, 20, 23], [ 4, 22, 25], [ 5, 24, 17], [16, 17, 26], 20 | [22, 23, 32], [48, 18, 28], [49, 20, 30], [24, 25, 34], 21 | [25, 22, 50], [28, 19, 47], [30, 21, 48], [32, 23, 49], 22 | [17, 24, 51], [26, 17, 51], [34, 25, 50], [23, 20, 49], 23 | [21, 18, 48], [19, 16, 47], [51, 24, 34], [24, 15, 25], 24 | [15, 4, 25], [50, 22, 32], [22, 14, 23], [14, 3, 23], 25 | [20, 21, 30], [20, 12, 21], [12, 2, 21], [18, 19, 28], 26 | [18, 8, 19], [ 8, 1, 19], [47, 16, 26], [16, 9, 17], 27 | [ 9, 5, 17], [10, 15, 5], [10, 13, 15], [13, 4, 15], 28 | [13, 14, 4], [13, 11, 14], [11, 3, 14], [11, 12, 3], 29 | [11, 6, 12], [ 6, 2, 12], [ 9, 10, 5], [ 9, 7, 10], 30 | [ 7, 0, 10], [ 6, 8, 2], [ 6, 7, 8], [ 7, 1, 8], 31 | [29, 36, 41], [31, 37, 44], [33, 38, 45], [35, 39, 46], 32 | [27, 40, 42], [42, 46, 43], [42, 40, 46], [40, 35, 46], 33 | [46, 45, 43], [46, 39, 45], [39, 33, 45], [45, 44, 43], 34 | [45, 38, 44], [38, 31, 44], [44, 41, 43], [44, 37, 41], 35 | [37, 29, 41], [41, 42, 43], [41, 36, 42], [36, 27, 42], 36 | [26, 40, 27], [26, 51, 40], [51, 35, 40], [34, 39, 35], 37 | [34, 50, 39], [50, 33, 39], [32, 38, 33], [32, 49, 38], 38 | [49, 31, 38], [30, 37, 31], [30, 48, 37], [48, 29, 37], 39 | [28, 36, 29], [28, 47, 36], [47, 27, 36], [51, 34, 35], 40 | [50, 32, 33], [49, 30, 31], [48, 28, 29], [47, 26, 27]]) 41 | 42 | elev = np.asarray([ 0. , 0.5535673 , 1.01721871, 0. , -1.01721871, 43 | -0.5535673 , 0.52359324, 0.31415301, 0.94246863, 0. , 44 | -0.31415301, 0. , 0.52359547, -0.52359324, -0.52359547, 45 | -0.94246863, 0.31415501, -0.31415501, 1.57079633, 0.94247719, 46 | 0.31415501, 0.94247719, -0.94247719, -0.31415501, -0.94247719, 47 | -1.57079633, -0.31415624, 0. , 0.94248124, 1.01722122, 48 | 0.94247396, 0.55356579, -0.31415377, -0.55356579, -1.57079233, 49 | -1.01722122, 0.52359706, 0.94246791, 0. , -0.94246791, 50 | -0.52359706, 0.52359371, 0. , 0. , 0.31415246, 51 | -0.31415246, -0.52359371, 0.31415624, 1.57079233, 0.31415377, 52 | -0.94247396, -0.94248124]) 53 | 54 | az = np.asarray([-1.57079633, -0.55358064, -2.12435586, -2.67794236, -2.12435586, 55 | -0.55358064, -1.7595018 , -1.10715248, -1.10714872, -0.55357999, 56 | -1.10715248, -2.12436911, -2.48922865, -1.7595018 , -2.48922865, 57 | -1.10714872, 0. , 0. , 0. , 0. , 58 | 3.14159265, 3.14159265, 3.14159265, 3.14159265, 0. , 59 | 0. , 0. , 0.46365119, 0. , 1.01724226, 60 | 3.14159265, 2.58801549, 3.14159265, 2.58801549, 3.14159265, 61 | 1.01724226, 0.6523668 , 2.03445078, 2.58801476, 2.03445078, 62 | 0.6523668 , 1.38209652, 1.01722642, 1.57080033, 2.03444394, 63 | 2.03444394, 1.38209652, 0. , 3.14159265, 3.14159265, 64 | 3.14159265, 0. ]) 65 | 66 | v = np.vstack([np.cos(az)*np.cos(elev),np.sin(az)*np.cos(elev),np.sin(elev)]).T 67 | 68 | class Capsule(object): 69 | 70 | def __init__(self, t, rod, rad, length): 71 | self.t = t # translation of the axis 72 | self.rod = rod # rotation of the axis in Rodrigues form 73 | self.rad = rad # radious of the capsule 74 | self.length = length # length of the axis 75 | axis0 = ch.vstack([0, ch.abs(self.length), 0]) 76 | self.axis = ch.vstack((t.T, (t + Rodrigues(rod).dot(axis0)).T)) 77 | v0 = ch.hstack([v[:26].T*rad, (v[26:].T*rad)+ axis0]) 78 | self.v = ((t + Rodrigues(rod).dot(v0)).T) 79 | self.set_sphere_centers() 80 | #self._dirty_vars = set() 81 | 82 | 83 | def set_sphere_centers(self, floor=False): 84 | if floor: 85 | n_spheres = int(np.floor(self.length/(2*self.rad) - 1)) 86 | else: 87 | n_spheres = int(np.ceil(self.length/(2*self.rad) - 1)) 88 | 89 | centers = [self.axis[0].r, self.axis[1].r] 90 | 91 | if n_spheres >= 1: 92 | step = self.length.r/(n_spheres+1) 93 | for i in xrange(n_spheres): 94 | centers.append(self.axis[0].r + (self.axis[1].r-self.axis[0].r)*step*(i+1)/self.length.r) 95 | 96 | self.centers = centers 97 | #return capsule.centers 98 | -------------------------------------------------------------------------------- /up_tools/capsule_man.py: -------------------------------------------------------------------------------- 1 | """ 2 | Copyright 2016 Max Planck Society, Federica Bogo, Angjoo Kanazawa. All 3 | rights reserved. 4 | 5 | This software is provided for research purposes only. 6 | 7 | By using this software you agree to the terms of the SMPLify license here: 8 | 9 | http://smplify.is.tue.mpg.de/license 10 | """ 11 | import numpy as np 12 | import chumpy as ch 13 | import cPickle as pickle 14 | 15 | from up_tools.mesh import Mesh 16 | from os.path import exists 17 | 18 | try: 19 | # Robustify against setup. 20 | from smpl.lbs import verts_core 21 | from smpl.serialization import load_model as _load_model 22 | except ImportError: 23 | # pylint: disable=import-error 24 | try: 25 | from psbody.smpl.lbs import verts_core 26 | from psbody.smpl.serialization import load_model as _load_model 27 | except: 28 | from smpl_webuser.lbs import verts_core 29 | from smpl_webuser.serialization import load_model as _load_model 30 | import scipy.sparse as sp 31 | from capsule_ch import Capsule, cap_f 32 | 33 | joint2name = ['pelvis', 'leftThigh', 'rightThigh', 'spine', 'leftCalf', 34 | 'rightCalf', 'spine1', 'leftFoot', 'rightFoot', 'spine2', 35 | 'neck', 'leftShoulder', 'rightShoulder', 36 | 'head', 'leftUpperArm', 'rightUpperArm', 'leftForeArm', 37 | 'rightForeArm', 'leftHand', 'rightHand'] 38 | 39 | # the orientation of each capsule 40 | rots0 = ch.asarray([[0, 0, np.pi/2], 41 | [0, 0, np.pi], 42 | [0, 0, np.pi], 43 | [0, 0, np.pi/2], 44 | [0, 0, np.pi], 45 | [0, 0, np.pi], 46 | [0, 0, np.pi/2], 47 | [np.pi/2, 0, 0], 48 | [np.pi/2, 0, 0], 49 | [0, 0, np.pi/2], 50 | [0, 0, 0], 51 | [0, 0, -np.pi/2], 52 | [0, 0, np.pi/2], 53 | [0, 0, 0], 54 | [0, 0, -np.pi/2], 55 | [0, 0, np.pi/2], 56 | [0, 0, -np.pi/2], 57 | [0, 0, np.pi/2], 58 | [0, 0, -np.pi/2], 59 | [0, 0, np.pi/2]]) 60 | 61 | # groups hands and fingers, feet and toes 62 | mujoco2segm = [[0], # hip 0 63 | [1], # leftThigh 1 64 | [2], # rightThigh 2 65 | [3], # spine 3 66 | [4], # leftCalf 4 67 | [5], # rightCalf 5 68 | [6], # spine1 6 69 | [7,10], # leftFoot + leftToes 7 70 | [8,11], # rightFoot + rightToes 8 71 | [9], # spine2 9 72 | [12], # neck 10 73 | [13], # leftShoulder 11 74 | [14], # rightShoulder 12 75 | [15], # head 13 76 | [16], # leftUpperArm 14 77 | [17], # rightUpperArm 15 78 | [18], # leftForeArm 16 79 | [19], # rightForeArm 17 80 | [20,22], # leftHand + leftFingers 18 81 | [21,23]] # rightHand + rightFingers 19 82 | 83 | collisions = [[0, 16], # hip and leftForeArm 84 | [0, 17], # hip and rightForeArm 85 | [0, 18], # hip and leftHand 86 | [0, 19], # hip and rightHand 87 | #[1, 2], # leftThigh and rightThigh 88 | [3, 16], # spine and leftForeArm 89 | [3, 17], # spine and rightForeArm 90 | #[3, 14], # spine and leftUpperArm 91 | #[3, 15], # spine and rightUpperArm 92 | [3, 18], # spine and leftHand 93 | [3, 19], # spine and rightHand 94 | [4, 5], # leftCalf and rightCalf 95 | [6, 16], # spine1 and leftForeArm 96 | [6, 17], # spine1 and rightForeArm 97 | #[6, 14], # spine1 and leftUpperArm 98 | #[6, 15], # spine1 and rightUpperArm 99 | [6, 18], # spine1 and leftHand 100 | [6, 19], # spine1 and rightHand 101 | [7, 5], # leftFoot and rightCalf 102 | [8, 7], # rightFoot and leftFoot 103 | [8, 4], # rightFoot and leftCalf 104 | [9, 16], # spine2 and leftForeArm 105 | [9, 17], # spine2 and rightForeArm 106 | #[9, 14], # spine2 and leftUpperArm 107 | #[9, 15], # spine2 and rightUpperArm 108 | [9, 18], # spine2 and leftHand 109 | [9, 19], # spine2 and rightHand 110 | [11, 16], # leftShoulder and leftForeArm 111 | [12, 17], # rightShoulder and rightForeArm 112 | [18, 19], # leftHand and rightHand 113 | ] 114 | 115 | 116 | def get_capsules(model, wrt_betas=None, length_regs=None, rad_regs=None): 117 | from opendr.geometry import Rodrigues 118 | 119 | n_shape_dofs = model.betas.r.size 120 | 121 | segm = np.argmax(model.weights_prior, axis=1) 122 | J_off = ch.zeros((len(joint2name), 3)) 123 | rots = rots0.copy() 124 | mujoco_t_mid = [0, 3, 6, 9] 125 | 126 | if wrt_betas is not None: 127 | # we must have the regressors 128 | assert(length_regs is not None and rad_regs is not None) 129 | # betas must be a chumpy object 130 | assert(hasattr(wrt_betas, 'dterms')) 131 | # WHY IS THIS 301 x 20 (and not 300)? 132 | pad = ch.concatenate((wrt_betas, ch.zeros(n_shape_dofs-len(wrt_betas)), ch.ones(1))) 133 | #lengths = ch.ch.MatVecMult(length_regs.T, pad) 134 | #rads = ch.ch.MatVecMult(rad_regs.T, pad) 135 | lengths = pad.dot(length_regs) 136 | rads = pad.dot(rad_regs) 137 | else: 138 | lengths = ch.ones(len(joint2name)) 139 | rads = ch.ones(len(joint2name)) 140 | 141 | betas = wrt_betas if wrt_betas is not None else model.betas 142 | n_betas = len(betas) 143 | 144 | # the joint regressors are the original, pre-optimized ones (middle of the part frontier) 145 | myJ_regressor = model.J_regressor_prior# sp.csc_matrix((data, (rows, cols))) 146 | myJ0 = ch.vstack((ch.ch.MatVecMult(myJ_regressor, model.v_template[:,0] + model.shapedirs[:,:,:n_betas].dot(betas)[:,0]), 147 | ch.ch.MatVecMult(myJ_regressor, model.v_template[:,1] + model.shapedirs[:,:,:n_betas].dot(betas)[:,1]), 148 | ch.ch.MatVecMult(myJ_regressor, model.v_template[:,2] + model.shapedirs[:,:,:n_betas].dot(betas)[:,2]))).T 149 | # with small adjustments for hips, spine and feet 150 | myJ = ch.vstack([ch.concatenate([myJ0[0,0], (.6*myJ0[0,1]+.2*myJ0[1,1]+.2*myJ0[2,1]), myJ0[9,2]]), 151 | ch.vstack([myJ0[i] for i in range(1,7)]), 152 | ch.concatenate([myJ0[7,0], (1.1*myJ0[7,1] - .1*myJ0[4,1]), myJ0[7,2]]), 153 | ch.concatenate([myJ0[8,0], (1.1*myJ0[8,1] - .1*myJ0[5,1]), myJ0[8,2]]), 154 | ch.concatenate([myJ0[9,0], myJ0[9,1], (.2*myJ0[9,2]+.8*myJ0[12,2])]), 155 | #ch.concatenate([myJ0[9,0], .9*myJ0[9,1] + .1*myJ0[12,1], myJ0[9,2]]), 156 | ch.vstack([myJ0[i] for i in range(10,24)])]) 157 | 158 | capsules = [] 159 | # create one capsule per mujoco joint. Assign some parameters (irrelevant now with the regressors) 160 | for ijoint, segms in enumerate(mujoco2segm): 161 | if wrt_betas is None: 162 | vidxs = np.asarray([segm == k for k in segms]).any(axis=0) 163 | verts = model.v_template[vidxs].r 164 | dims = (verts.max(axis=0) - verts.min(axis=0)) 165 | rads[ijoint] = .5*((dims[(np.argmax(dims) + 1)%3] + dims[(np.argmax(dims) + 2)%3])/4.) 166 | lengths[ijoint] = max(dims) - 2.*rads[ijoint].r 167 | # the core joints are different, since the capsule is not in the joint but in the middle 168 | if ijoint in mujoco_t_mid: 169 | len_offset = ch.vstack([ch.zeros(1), ch.abs(lengths[ijoint])/2., ch.zeros(1)]).reshape(3,1) 170 | caps = Capsule((J_off[ijoint] + myJ[mujoco2segm[ijoint][0]]).reshape(3,1) - 171 | #Rodrigues(rots[ijoint])[0].dot(len_offset), 172 | Rodrigues(rots[ijoint]).dot(len_offset), 173 | rots[ijoint], rads[ijoint], lengths[ijoint]) 174 | else: 175 | caps = Capsule((J_off[ijoint] + myJ[mujoco2segm[ijoint][0]]).reshape(3,1), rots[ijoint], rads[ijoint], lengths[ijoint]) 176 | caps.id = ijoint 177 | capsules.append(caps) 178 | return capsules 179 | 180 | 181 | def show_capsules(capsules, mv): 182 | from body.mesh import Mesh 183 | 184 | capsules_v = [cps.v.r for cps in capsules] 185 | vs = np.vstack(capsules_v) 186 | fs = np.vstack([cap_f+(i*len(capsules[0].v)) for i in range(len(capsules_v))]) 187 | mcap = Mesh(v=vs, f=fs) 188 | mv.set_static_meshes([mcap]) 189 | 190 | 191 | def set_sphere_centers(capsule, floor=True): 192 | if np.isnan(capsule.length.r)[0] or np.isnan(capsule.rad.r)[0]: 193 | n_spheres = 0 194 | else: 195 | if floor: 196 | n_spheres = int(np.floor(capsule.length.r/(2*capsule.rad.r) - 1)) 197 | else: 198 | n_spheres = int(np.ceil(capsule.length.r/(2*capsule.rad.r) - 1)) 199 | n_spheres = min(n_spheres, 100) 200 | # remove "redundant" spheres for right and left thigh... 201 | if capsule.id == 1 or capsule.id == 2: 202 | centers = [capsule.axis[1].r] 203 | # ... and right and left upper arm 204 | elif capsule.id == 14 or capsule.id == 15: 205 | if n_spheres >= 1: 206 | centers = [] #[capsule.axis[1].r] 207 | else: 208 | centers = [capsule.axis[1].r] 209 | else: 210 | centers = [capsule.axis[0].r, capsule.axis[1].r] 211 | 212 | if n_spheres >= 1: 213 | step = capsule.length.r/(n_spheres+1) 214 | for i in xrange(n_spheres): 215 | centers.append(capsule.axis[0].r + (capsule.axis[1].r-capsule.axis[0].r)*step*(i+1)/capsule.length.r) 216 | 217 | capsule.centers = centers 218 | return capsule.centers 219 | 220 | 221 | def capsule_dist(capsule0, capsule1, alpha=.3, increase_hand=True): 222 | range0 = range(capsule0.center_id, capsule0.center_id + len(capsule0.centers)) 223 | range1 = range(capsule1.center_id, capsule1.center_id + len(capsule1.centers)) 224 | 225 | #cnt0 = ch.hstack([[cid]*len(range1) for cid in range0]) 226 | cnt0 = ch.concatenate([[cid]*len(range1) for cid in range0]) 227 | #cnt1 = ch.hstack([range1]*len(range0)) 228 | cnt1 = ch.concatenate([range1]*len(range0)) 229 | # dst = capsule0.rad.r + capsule1.rad.r 230 | if increase_hand: 231 | if capsule0.id == 18 or capsule0.id == 19 or capsule1.id == 18 or capsule1.id == 19: 232 | dst = (alpha * 1.2 * capsule0.rad.r)**2 + (alpha * 1.2 * capsule1.rad.r)**2 233 | else: 234 | dst = (alpha * capsule0.rad.r)**2 + (alpha * capsule1.rad.r)**2 235 | else: 236 | dst = (alpha * capsule0.rad.r)**2 + (alpha * capsule1.rad.r)**2 237 | radiuss = np.hstack([dst]*len(cnt0)).squeeze() 238 | return (cnt0, cnt1, radiuss) 239 | 240 | def get_capsule_bweights(vs): 241 | # "blend" weights for the capsule. They are binary 242 | rows = np.arange(vs.shape[0]) 243 | cols = np.tile(np.hstack((range(10), range(12,22))), (52,1)).T.ravel() 244 | data = np.ones(vs.shape[0]) 245 | caps_weights = np.asarray(sp.csc_matrix((data, (rows, cols)), shape=(vs.shape[0], 24)).todense()) 246 | return caps_weights 247 | 248 | def get_sphere_bweights(sph_vs, capsules): 249 | rows = np.arange(sph_vs.shape[0]) 250 | cols = [] 251 | for cps, w in zip(capsules, range(10)+range(12,22)): 252 | cols.append([w]*len(cps.centers)) 253 | cols = np.hstack(cols) 254 | data = np.ones(sph_vs.shape[0]) 255 | sph_weights = np.asarray(sp.csc_matrix((data, (rows, cols)), shape=(sph_vs.shape[0], 24)).todense()) 256 | return sph_weights 257 | -------------------------------------------------------------------------------- /up_tools/max_mixture_prior.py: -------------------------------------------------------------------------------- 1 | """ 2 | Copyright 2016 Max Planck Society, Federica Bogo, Angjoo Kanazawa. All 3 | rights reserved. 4 | 5 | This software is provided for research purposes only. 6 | 7 | By using this software you agree to the terms of the SMPLify license here: 8 | 9 | http://smplify.is.tue.mpg.de/license 10 | """ 11 | # pylint: disable=no-member, invalid-name, attribute-defined-outside-init 12 | import os 13 | import numpy as np 14 | import chumpy as ch 15 | 16 | class MaxMixtureComplete(ch.Ch): 17 | 18 | """Define the MaxMixture class.""" 19 | 20 | dterms = 'x', 'means', 'precs', 'weights' 21 | 22 | def on_changed(self, which): 23 | # setup means, precs and loglikelihood expressions 24 | if 'means' in which or 'precs' in which or 'weights' in which: 25 | # This is just the mahalanobis part. 26 | self.loglikelihoods = [np.sqrt(0.5) * (self.x - m).dot(s) 27 | for m, s in zip(self.means, self.precs)] 28 | 29 | 30 | if 'x' in which: 31 | # start = time.time() 32 | self.min_component_idx = np.argmin([(logl**2).sum().r[0] -np.log(w[0]) 33 | for logl, w in zip(self.loglikelihoods, 34 | self.weights)]) 35 | 36 | def compute_r(self): 37 | # pylint: disable=unsubscriptable-object 38 | min_w = self.weights[self.min_component_idx] 39 | # Add the sqrt(-log(weights)). 40 | return ch.concatenate((self.loglikelihoods[self.min_component_idx].r, 41 | np.sqrt(-np.log(min_w)))) 42 | 43 | 44 | def compute_dr_wrt(self, wrt): 45 | # Returns 69 x 72, when wrt is 69D => return 70x72 with empty last for 46 | # when returning 70D 47 | # Extract the data, rows cols and data, new one with exact same values 48 | # but with size one more rows) 49 | import scipy.sparse as sp 50 | 51 | dr = self.loglikelihoods[self.min_component_idx].dr_wrt(wrt) 52 | if dr is not None: 53 | Is, Js, Vs = sp.find(dr) 54 | dr = sp.csc_matrix((Vs, (Is, Js)), shape=(dr.shape[0]+1, dr.shape[1])) 55 | 56 | return dr 57 | 58 | 59 | # pylint: disable=too-few-public-methods 60 | class MaxMixtureCompleteWrapper(object): 61 | 62 | """Convenience wrapper to match interface spec.""" 63 | 64 | def __init__(self, means, precs, weights, prefix): 65 | self.means = means 66 | self.precs = precs # Already "sqrt"ed 67 | self.weights = weights 68 | self.prefix = prefix 69 | 70 | def __call__(self, x): 71 | return(MaxMixtureComplete(x=x[self.prefix:], means=self.means, 72 | precs=self.precs, weights=self.weights)) 73 | 74 | 75 | 76 | class MaxMixtureCompletePrior(object): 77 | 78 | """Prior density estimation.""" 79 | 80 | def __init__(self, n_gaussians=6, prefix=3): 81 | self.n_gaussians = n_gaussians 82 | self.prefix = prefix 83 | self.prior = self.create_prior_from_cmu() 84 | 85 | def create_prior_from_cmu(self): 86 | """Get the prior from the CMU motion database.""" 87 | from os.path import realpath 88 | import cPickle as pickle 89 | with open(os.path.join(os.path.dirname(__file__), 90 | '..', 'models', '3D', 91 | 'gmm_best_%02d.pkl' % self.n_gaussians)) as f: 92 | gmm = pickle.load(f) 93 | 94 | precs = ch.asarray([np.linalg.inv(cov) for cov in gmm.covars_]) 95 | chols = ch.asarray([np.linalg.cholesky(prec) for prec in precs]) 96 | 97 | # The constant term: 98 | sqrdets = np.array([(np.sqrt(np.linalg.det(c))) for c in gmm.covars_]) 99 | const = (2*np.pi)**(69/2.) 100 | 101 | self.weights = ch.asarray(gmm.weights_ / (const * (sqrdets/sqrdets.min()))) 102 | 103 | return(MaxMixtureCompleteWrapper(means=gmm.means_, precs=chols, 104 | weights=self.weights, prefix=self.prefix)) 105 | 106 | def get_gmm_prior(self): 107 | """Getter implementation.""" 108 | return self.prior 109 | -------------------------------------------------------------------------------- /up_tools/mesh.py: -------------------------------------------------------------------------------- 1 | """Mesh tools.""" 2 | # pylint: disable=invalid-name 3 | import numpy as np 4 | import meshio 5 | 6 | 7 | class Mesh(object): # pylint: disable=too-few-public-methods 8 | 9 | """An easy to use mesh interface.""" 10 | 11 | def __init__(self, filename=None, v=None, vc=None, f=None): 12 | """Construct a mesh either from a file or with the provided data.""" 13 | if filename is not None: 14 | assert v is None and f is None and vc is None 15 | else: 16 | assert v is not None or f is not None 17 | if vc is not None: 18 | assert len(v) == len(vc) 19 | 20 | if filename is None: 21 | self.v = v 22 | self.vc = vc 23 | self.f = f 24 | else: 25 | mesh = meshio.read(filename) 26 | self.v = mesh.points 27 | self.vc = ( 28 | np.column_stack( 29 | [ 30 | mesh.point_data["red"], 31 | mesh.point_data["green"], 32 | mesh.point_data["blue"], 33 | ] 34 | ).astype(np.float) 35 | / 255.0 36 | ) 37 | self.f = mesh.cells["triangle"] 38 | 39 | def write_ply(self, out_name): 40 | """Write to a .ply file.""" 41 | colors = (self.vc * 255).astype("uint8") 42 | meshio.write_points_cells( 43 | out_name, 44 | self.v, 45 | {"triangle": self.f}, 46 | point_data={ 47 | "red": colors[:, 0], 48 | "green": colors[:, 1], 49 | "blue": colors[:, 2], 50 | }, 51 | ) 52 | -------------------------------------------------------------------------------- /up_tools/render_segmented_views.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python 2 | """Data management script for this project.""" 3 | # pylint: disable=invalid-name, wrong-import-order 4 | from __future__ import print_function 5 | 6 | import sys 7 | import logging as _logging 8 | from copy import copy as _copy 9 | from glob import glob as _glob 10 | import cPickle as _pickle 11 | import os as _os 12 | import os.path as _path 13 | 14 | import numpy as _np 15 | import cv2 as _cv2 16 | import click as _click 17 | import pymp as _pymp 18 | 19 | import opendr.renderer as _odr_r 20 | import opendr.camera as _odr_c 21 | import opendr.lighting as _odr_l 22 | 23 | import tqdm 24 | # pylint: disable=no-name-in-module 25 | from up_tools.mesh import Mesh as _Mesh 26 | from up_tools.bake_vertex_colors import bake_vertex_colors 27 | from up_tools.camera import (rotateY as _rotateY) 28 | from clustertools.log import LOGFORMAT 29 | from clustertools.config import available_cpu_count 30 | 31 | sys.path.insert(0, _path.join(_path.dirname(__file__), 32 | '..', '..')) 33 | from config import SMPL_FP 34 | sys.path.insert(0, SMPL_FP) 35 | try: 36 | from smpl.serialization import load_model # pylint: disable=import-error 37 | except ImportError: 38 | try: 39 | from psbody.smpl.serialization import load_model # pylint: disable=import-error 40 | except ImportError: 41 | from smpl_webuser.serialization import load_model 42 | 43 | _LOGGER = _logging.getLogger(__name__) 44 | MODEL_NEUTRAL_PATH = _os.path.join( 45 | _path.dirname(__file__), '..', 'models', '3D', 46 | 'basicModel_neutral_lbs_10_207_0_v1.0.0.pkl') 47 | MODEL_NEUTRAL = load_model(MODEL_NEUTRAL_PATH) 48 | _TEMPLATE_MESH = _Mesh(filename=_os.path.join(_os.path.dirname(__file__), 49 | '..', 'models', '3D', 50 | 'template-bodyparts.ply')) 51 | 52 | def _rodrigues_from_seq(angles_seq): 53 | """Create rodrigues representation of angles.""" 54 | rot = _np.eye(3) 55 | for angle in angles_seq[::-1]: 56 | rot = rot.dot(_cv2.Rodrigues(angle)[0]) 57 | return _cv2.Rodrigues(rot)[0].flatten() 58 | 59 | 60 | def _create_renderer( # pylint: disable=too-many-arguments 61 | w=640, 62 | h=480, 63 | rt=_np.zeros(3), 64 | t=_np.zeros(3), 65 | f=None, 66 | c=None, 67 | k=None, 68 | near=1., 69 | far=10., 70 | texture=None): 71 | """Create a renderer for the specified parameters.""" 72 | f = _np.array([w, w]) / 2. if f is None else f 73 | c = _np.array([w, h]) / 2. if c is None else c 74 | k = _np.zeros(5) if k is None else k 75 | 76 | if texture is not None: 77 | rn = _odr_r.TexturedRenderer() 78 | else: 79 | rn = _odr_r.ColoredRenderer() # pylint: disable=redefined-variable-type 80 | 81 | rn.camera = _odr_c.ProjectPoints(rt=rt, t=t, f=f, c=c, k=k) 82 | rn.frustum = {'near':near, 'far':far, 'height':h, 'width':w} 83 | if texture is not None: 84 | rn.texture_image = _np.asarray(_cv2.imread(texture), _np.float64)/255. 85 | return rn 86 | 87 | 88 | def _stack_with(rn, mesh, texture): 89 | if texture is not None: 90 | if not hasattr(mesh, 'ft'): 91 | mesh.ft = mesh.f 92 | mesh.vt = mesh.v[:, :2] 93 | rn.ft = _np.vstack((rn.ft, mesh.ft+len(rn.vt))) 94 | rn.vt = _np.vstack((rn.vt, mesh.vt)) 95 | rn.f = _np.vstack((rn.f, mesh.f+len(rn.v))) 96 | rn.v = _np.vstack((rn.v, mesh.v)) 97 | rn.vc = _np.vstack((rn.vc, mesh.vc)) 98 | 99 | 100 | def _simple_renderer(rn, meshes, yrot=0, texture=None, use_light=False): 101 | mesh = meshes[0] 102 | if texture is not None: 103 | if not hasattr(mesh, 'ft'): 104 | mesh.ft = _copy(mesh.f) 105 | vt = _copy(mesh.v[:, :2]) 106 | vt -= _np.min(vt, axis=0).reshape((1, -1)) 107 | vt /= _np.max(vt, axis=0).reshape((1, -1)) 108 | mesh.vt = vt 109 | mesh.texture_filepath = rn.texture_image 110 | 111 | # Set camera parameters 112 | if texture is not None: 113 | rn.set(v=mesh.v, f=mesh.f, vc=mesh.vc, ft=mesh.ft, vt=mesh.vt, bgcolor=_np.ones(3)) 114 | else: 115 | rn.set(v=mesh.v, f=mesh.f, vc=mesh.vc, bgcolor=_np.ones(3)) 116 | 117 | for next_mesh in meshes[1:]: 118 | _stack_with(rn, next_mesh, texture) 119 | 120 | # Construct light. 121 | if use_light: 122 | albedo = rn.vc 123 | rn.vc = _odr_l.LambertianPointLight( 124 | f=rn.f, 125 | v=rn.v, 126 | num_verts=len(rn.v), 127 | light_pos=_rotateY(_np.array([-200, -100, -100]), yrot), 128 | vc=albedo, 129 | light_color=_np.array([1, 1, 1])) 130 | # Construct Left Light 131 | rn.vc += _odr_l.LambertianPointLight( 132 | f=rn.f, 133 | v=rn.v, 134 | num_verts=len(rn.v), 135 | light_pos=_rotateY(_np.array([800, 10, 300]), yrot), 136 | vc=albedo, 137 | light_color=_np.array([1, 1, 1])) 138 | 139 | # Construct Right Light 140 | rn.vc += _odr_l.LambertianPointLight( 141 | f=rn.f, 142 | v=rn.v, 143 | num_verts=len(rn.v), 144 | light_pos=_rotateY(_np.array([-500, 500, 1000]), yrot), 145 | vc=albedo, 146 | light_color=_np.array([.7, .7, .7])) 147 | return rn.r 148 | 149 | 150 | # pylint: disable=too-many-locals 151 | def render(model, resolution, cam, steps, segmented=False, use_light=False): # pylint: disable=too-many-arguments 152 | """Render a sequence of views from a fitted body model.""" 153 | assert steps >= 1 154 | if segmented: 155 | texture = _os.path.join(_os.path.dirname(__file__), 156 | '..', 'models', '3D', 'mask_filled.png') 157 | else: 158 | texture = _os.path.join(_os.path.dirname(__file__), 159 | '..', 'models', '3D', 'mask_filled_uniform.png') 160 | mesh = _copy(_TEMPLATE_MESH) 161 | 162 | # render ply 163 | model.betas[:len(cam['betas'])] = cam['betas'] 164 | model.pose[:] = cam['pose'] 165 | model.trans[:] = cam['trans'] 166 | 167 | mesh.v = model.r 168 | w, h = resolution[0], resolution[1] 169 | dist = _np.abs(cam['t'][2] - _np.mean(mesh.v, axis=0)[2]) 170 | rn = _create_renderer(w=w, 171 | h=h, 172 | near=1., 173 | far=20.+dist, 174 | rt=_np.array(cam['rt']), 175 | t=_np.array(cam['t']), 176 | f=_np.array([cam['f'], cam['f']]), 177 | # c=_np.array(cam['cam_c']), 178 | texture=texture) 179 | light_yrot = _np.radians(120) 180 | baked_mesh = bake_vertex_colors(mesh) 181 | base_mesh = _copy(baked_mesh) 182 | mesh.f = base_mesh.f 183 | mesh.vc = base_mesh.vc 184 | renderings = [] 185 | for angle in _np.linspace(0., 2. * (1. - 1. / steps) * _np.pi, steps): 186 | mesh.v = _rotateY(base_mesh.v, angle) 187 | imtmp = _simple_renderer(rn=rn, 188 | meshes=[mesh], 189 | yrot=light_yrot, 190 | texture=texture, 191 | use_light=use_light) 192 | im = _np.zeros(h*w*3).reshape(((h, w, 3))) 193 | im[:h, :w, :] = imtmp*255. 194 | renderings.append(im) 195 | return renderings 196 | 197 | ############################################################################### 198 | # Command-line interface. 199 | ############################################################################### 200 | 201 | 202 | @_click.command() 203 | @_click.argument('input_folder', type=_click.Path(exists=True, file_okay=False, 204 | readable=True)) 205 | @_click.argument('out_folder', type=_click.Path(writable=True)) 206 | @_click.option("--num_shots_per_body", type=_click.INT, default=7, 207 | help="Number of shots to take per body.") 208 | @_click.option("--only_missing", type=_click.BOOL, default=False, is_flag=True, 209 | help="Only run for missing shots.") 210 | @_click.option("--num_threads", type=_click.INT, default=-1, 211 | help="Number of threads to use.") 212 | @_click.option("--use_light", type=_click.BOOL, default=False, is_flag=True, 213 | help="Light the scene for a depth effect.") 214 | @_click.option("--factor", type=_click.FLOAT, default=1., 215 | help="Scaling factor for the rendered human (not the image).") 216 | def sample_shots( # pylint: disable=too-many-arguments 217 | input_folder, 218 | out_folder, 219 | num_shots_per_body=7, 220 | only_missing=False, 221 | num_threads=-1, 222 | use_light=False, 223 | factor=1.): 224 | """Sample body images with visibilities.""" 225 | _LOGGER.info("Sampling 3D body shots.") 226 | if num_threads == -1: 227 | num_threads = available_cpu_count() 228 | else: 229 | assert num_threads > 0 230 | if not _path.exists(out_folder): 231 | _os.mkdir(out_folder) 232 | _np.random.seed(1) 233 | bodies = _glob(_path.join(input_folder, '*.pkl')) 234 | _LOGGER.info("%d bodies detected.", len(bodies)) 235 | with _pymp.Parallel(num_threads, if_=num_threads > 1) as p: 236 | for body_idx in p.iterate(tqdm.tqdm(range(len(bodies)))): 237 | body_filename = bodies[body_idx] 238 | vis_filename = body_filename + '_vis_overlay.png' 239 | vis_filename = body_filename + '_vis.png' 240 | if not _os.path.exists(vis_filename): 241 | vis_filename = body_filename + '_vis_0.png' 242 | if not _os.path.exists(vis_filename): 243 | # Try something else. 244 | vis_filename = body_filename[:-9] 245 | out_names = [_os.path.join(out_folder, 246 | _path.basename(body_filename) + '.' + 247 | str(map_idx) + '.png') 248 | for map_idx in range(num_shots_per_body)] 249 | if only_missing: 250 | all_exist = True 251 | for fname in out_names: 252 | if not _path.exists(fname): 253 | all_exist = False 254 | break 255 | if all_exist: 256 | continue 257 | vis_im = _cv2.imread(vis_filename) 258 | assert vis_im is not None, 'visualization not found: %s' % (vis_filename) 259 | renderings = render_body_impl(body_filename, 260 | [vis_im.shape[1], vis_im.shape[0]], 261 | num_shots_per_body, 262 | quiet=False, 263 | use_light=use_light, 264 | factor=factor) 265 | for map_idx, vmap in enumerate(renderings): 266 | _cv2.imwrite(out_names[map_idx], vmap[:, :, ::-1]) 267 | _LOGGER.info("Done.") 268 | 269 | 270 | def render_body_impl(filename, # pylint: disable=too-many-arguments 271 | resolution=None, 272 | num_steps_around_y=1, 273 | quiet=False, 274 | use_light=False, 275 | factor=1.): 276 | """Create a SMPL rendering.""" 277 | if resolution is None: 278 | resolution = [640, 480] 279 | if not quiet: 280 | _LOGGER.info("Rendering SMPL model from `%s` in resolution %s.", 281 | filename, str(resolution)) 282 | _LOGGER.info(" Using %d steps around y axis.", 283 | num_steps_around_y) 284 | with open(filename, 'rb') as inf: 285 | camera = _pickle.load(inf) 286 | # Render. 287 | renderings = render( 288 | MODEL_NEUTRAL, 289 | (_np.asarray(resolution) * 1. / factor).astype('int'), 290 | camera, 291 | num_steps_around_y, 292 | False, 293 | use_light=use_light) 294 | interp = 'bilinear' if use_light else 'nearest' 295 | import scipy.misc 296 | renderings = [scipy.misc.imresize(renderim, 297 | (resolution[1], 298 | resolution[0]), 299 | interp=interp) 300 | for renderim in renderings] 301 | if not quiet: 302 | _LOGGER.info("Done.") 303 | return renderings 304 | 305 | 306 | if __name__ == '__main__': 307 | _logging.basicConfig(level=_logging.INFO, 308 | format=LOGFORMAT,) 309 | # filename=__file__+'_log.txt') 310 | _logging.getLogger("opendr.lighting").setLevel(_logging.WARN) 311 | sample_shots() # pylint: disable=no-value-for-parameter 312 | -------------------------------------------------------------------------------- /up_tools/robustifiers.py: -------------------------------------------------------------------------------- 1 | """ 2 | Copyright 2016 Max Planck Society, Matthew Loper. All rights reserved. 3 | This software is provided for research purposes only. 4 | By using this software you agree to the terms of the SMPLify license here: 5 | http://smplify.is.tue.mpg.de/license 6 | 7 | This script implements the Geman-McClure robustifier as chumpy object. 8 | """ 9 | 10 | #!/usr/bin/env python 11 | import numpy as np 12 | import scipy 13 | import scipy.sparse as sp 14 | from chumpy import Ch 15 | 16 | __all__ = ['GMOf'] 17 | 18 | 19 | def GMOf(x, sigma): 20 | """Given x and sigma in some units (say mm), 21 | returns robustified values (in same units), 22 | by making use of the Geman-McClure robustifier.""" 23 | 24 | result = SignedSqrt(x=GMOfInternal(x=x, sigma=sigma)) 25 | return result 26 | 27 | 28 | class SignedSqrt(Ch): 29 | dterms = ('x', ) 30 | terms = () 31 | 32 | def compute_r(self): 33 | return np.sqrt(np.abs(self.x.r)) * np.sign(self.x.r) 34 | 35 | def compute_dr_wrt(self, wrt): 36 | if wrt is self.x: 37 | result = (.5 / np.sqrt(np.abs(self.x.r))) 38 | result = np.nan_to_num(result) 39 | result *= (self.x.r != 0).astype(np.uint32) 40 | return sp.spdiags(result.ravel(), [0], self.x.r.size, 41 | self.x.r.size) 42 | 43 | 44 | class GMOfInternal(Ch): 45 | dterms = 'x', 'sigma' 46 | 47 | def on_changed(self, which): 48 | if 'sigma' in which: 49 | assert (self.sigma.r > 0) 50 | 51 | if 'x' in which: 52 | self.squared_input = self.x.r**2. 53 | 54 | def compute_r(self): 55 | return (self.sigma.r**2 * 56 | (self.squared_input / 57 | (self.sigma.r**2 + self.squared_input))) * np.sign(self.x.r) 58 | 59 | def compute_dr_wrt(self, wrt): 60 | if wrt is not self.x and wrt is not self.sigma: 61 | return None 62 | 63 | squared_input = self.squared_input 64 | result = [] 65 | if wrt is self.x: 66 | dx = self.sigma.r**2 / (self.sigma.r**2 + squared_input 67 | ) - self.sigma.r**2 * (squared_input / ( 68 | self.sigma.r**2 + squared_input)**2) 69 | dx = 2 * self.x.r * dx 70 | result.append( 71 | scipy.sparse.spdiags( 72 | (dx * np.sign(self.x.r)).ravel(), [0], 73 | self.x.r.size, 74 | self.x.r.size, 75 | format='csc')) 76 | if wrt is self.sigma: 77 | ds = 2 * self.sigma.r * (squared_input / ( 78 | self.sigma.r**2 + squared_input)) - 2 * self.sigma.r**3 * ( 79 | squared_input / (self.sigma.r**2 + squared_input)**2) 80 | result.append( 81 | scipy.sparse.spdiags( 82 | (ds * np.sign(self.x.r)).ravel(), [0], 83 | self.x.r.size, 84 | self.x.r.size, 85 | format='csc')) 86 | 87 | if len(result) == 1: 88 | return result[0] 89 | else: 90 | return np.sum(result).tocsc() 91 | -------------------------------------------------------------------------------- /up_tools/sphere_collisions.py: -------------------------------------------------------------------------------- 1 | """ 2 | Copyright 2016 Max Planck Society, Federica Bogo, Angjoo Kanazawa. All 3 | rights reserved. 4 | 5 | This software is provided for research purposes only. 6 | 7 | By using this software you agree to the terms of the SMPLify license here: 8 | 9 | http://smplify.is.tue.mpg.de/license 10 | """ 11 | import numpy as np 12 | import chumpy as ch 13 | try: 14 | # Robustify against setup. 15 | from smpl.lbs import verts_core 16 | except ImportError: 17 | # pylint: disable=import-error 18 | try: 19 | from psbody.smpl.lbs import verts_core 20 | except: 21 | from smpl_webuser.lbs import verts_core 22 | 23 | from capsule_man import get_capsules, set_sphere_centers,\ 24 | get_sphere_bweights, collisions, capsule_dist 25 | 26 | 27 | class SphereCollisions(ch.Ch): 28 | dterms = ('pose', 'betas') 29 | terms = ('regs', 'model') 30 | 31 | def update_capsules_and_centers(self): 32 | centers = [set_sphere_centers(capsule) for capsule in self.capsules] 33 | count = 0 34 | for capsule in self.capsules: 35 | capsule.center_id = count 36 | count += len(capsule.centers) 37 | self.sph_vs = ch.vstack(centers) 38 | self.sph_weights = get_sphere_bweights(self.sph_vs, self.capsules) 39 | self.ids0 = [] 40 | self.ids1 = [] 41 | self.radiuss = [] 42 | self.caps_pairs = [] 43 | for collision in collisions: 44 | if hasattr(self, 'no_hands'): 45 | (id0, id1, rd) = capsule_dist(self.capsules[collision[0]], self.capsules[collision[1]], increase_hand=False) 46 | else: 47 | (id0, id1, rd) = capsule_dist(self.capsules[collision[0]], self.capsules[collision[1]]) 48 | self.ids0.append(id0.r) 49 | self.ids1.append(id1.r) 50 | self.radiuss.append(rd) 51 | self.caps_pairs.append(['%02d_%02d' % (collision[0], collision[1])]*len(id0)) 52 | self.ids0 = np.concatenate(self.ids0).astype(int) # numpy? 53 | self.ids1 = np.concatenate(self.ids1).astype(int) 54 | self.radiuss = np.concatenate(self.radiuss) 55 | self.caps_pairs = np.concatenate(self.caps_pairs) 56 | assert(self.caps_pairs.size==self.ids0.size) 57 | assert(self.radiuss.size==self.ids0.size) 58 | 59 | def update_pose(self): 60 | self.sph_v = verts_core(self.pose, 61 | self.sph_vs, 62 | self.model.J, 63 | self.sph_weights, 64 | self.model.kintree_table, 65 | want_Jtr=False)[0] 66 | 67 | def get_objective(self): 68 | return ch.sum((ch.exp(-((ch.sum((self.sph_v[self.ids0]- 69 | self.sph_v[self.ids1])**2, 70 | axis=1))/(self.radiuss))/2.))**2)**.5 71 | 72 | def compute_r(self): 73 | return self.get_objective().r 74 | 75 | 76 | def compute_dr_wrt(self, wrt): 77 | if wrt is self.pose: 78 | return self.get_objective().dr_wrt(wrt) 79 | 80 | def on_changed(self, which): 81 | if 'regs' in which: 82 | self.length_regs = self.regs['betas2lens'] 83 | self.rad_regs = self.regs['betas2rads'] 84 | #self.capsules = get_capsules(self.model) 85 | #self.prev_betas = None 86 | 87 | if 'betas' in which: 88 | if not hasattr(self, 'capsules'): 89 | self.capsules = get_capsules(self.model, 90 | wrt_betas=self.betas, 91 | length_regs=self.length_regs, 92 | rad_regs=self.rad_regs) 93 | self.update_capsules_and_centers() 94 | 95 | if 'pose' in which: 96 | self.update_pose() 97 | -------------------------------------------------------------------------------- /up_tools/visualization.py: -------------------------------------------------------------------------------- 1 | """Visualization tools.""" 2 | # pylint: disable=no-member, invalid-name 3 | import clustertools.visualization as vs 4 | from up_tools.model import (connections_lsp, 5 | connections_landmarks_91, 6 | lm_region_mapping) 7 | 8 | 9 | def visualize_pose(image, # pylint: disable=too-many-arguments, dangerous-default-value 10 | pose, 11 | line_thickness=3, 12 | dash_length=15, 13 | opacity=0.6, 14 | circle_color=(255, 255, 255), 15 | connections=[], 16 | region_mapping=None, 17 | skip_unconnected_joints=True, 18 | scale=1.): 19 | """Visualize a pose.""" 20 | if pose.shape[1] == 91: 21 | this_connections = connections 22 | if this_connections == []: 23 | this_connections = connections_landmarks_91 24 | this_lm_region_mapping = region_mapping 25 | if this_lm_region_mapping is None: 26 | this_lm_region_mapping = lm_region_mapping 27 | else: 28 | this_connections = connections 29 | if this_connections == []: 30 | this_connections = connections_lsp 31 | this_lm_region_mapping = region_mapping 32 | return vs.visualize_pose(image, pose, line_thickness, dash_length, 33 | opacity, circle_color, this_connections, 34 | this_lm_region_mapping, 35 | skip_unconnected_joints, 36 | scale) 37 | --------------------------------------------------------------------------------