├── .gitignore ├── LICENSE.md ├── README.md ├── code ├── datasets.py ├── hopenet.py ├── test_alexnet.py ├── test_hopenet.py ├── test_on_video.py ├── test_on_video_dlib.py ├── test_on_video_dockerface.py ├── test_resnet50_regression.py ├── train_alexnet.py ├── train_hopenet.py ├── train_resnet50_regression.py └── utils.py └── conan-cruise.gif /.gitignore: -------------------------------------------------------------------------------- 1 | *.pyc 2 | *.npy 3 | data/* 4 | output/* 5 | *.jpg 6 | *.png -------------------------------------------------------------------------------- /LICENSE.md: -------------------------------------------------------------------------------- 1 | Nataniel Ruiz, Georgia Institute of Technology 2 | 3 | Apache License 4 | Version 2.0, January 2004 5 | http://www.apache.org/licenses/ 6 | 7 | TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION 8 | 9 | 1. Definitions. 10 | 11 | "License" shall mean the terms and conditions for use, reproduction, 12 | and distribution as defined by Sections 1 through 9 of this document. 13 | 14 | "Licensor" shall mean the copyright owner or entity authorized by 15 | the copyright owner that is granting the License. 16 | 17 | "Legal Entity" shall mean the union of the acting entity and all 18 | other entities that control, are controlled by, or are under common 19 | control with that entity. For the purposes of this definition, 20 | "control" means (i) the power, direct or indirect, to cause the 21 | direction or management of such entity, whether by contract or 22 | otherwise, or (ii) ownership of fifty percent (50%) or more of the 23 | outstanding shares, or (iii) beneficial ownership of such entity. 24 | 25 | "You" (or "Your") shall mean an individual or Legal Entity 26 | exercising permissions granted by this License. 27 | 28 | "Source" form shall mean the preferred form for making modifications, 29 | including but not limited to software source code, documentation 30 | source, and configuration files. 31 | 32 | "Object" form shall mean any form resulting from mechanical 33 | transformation or translation of a Source form, including but 34 | not limited to compiled object code, generated documentation, 35 | and conversions to other media types. 36 | 37 | "Work" shall mean the work of authorship, whether in Source or 38 | Object form, made available under the License, as indicated by a 39 | copyright notice that is included in or attached to the work 40 | (an example is provided in the Appendix below). 41 | 42 | "Derivative Works" shall mean any work, whether in Source or Object 43 | form, that is based on (or derived from) the Work and for which the 44 | editorial revisions, annotations, elaborations, or other modifications 45 | represent, as a whole, an original work of authorship. For the purposes 46 | of this License, Derivative Works shall not include works that remain 47 | separable from, or merely link (or bind by name) to the interfaces of, 48 | the Work and Derivative Works thereof. 49 | 50 | "Contribution" shall mean any work of authorship, including 51 | the original version of the Work and any modifications or additions 52 | to that Work or Derivative Works thereof, that is intentionally 53 | submitted to Licensor for inclusion in the Work by the copyright owner 54 | or by an individual or Legal Entity authorized to submit on behalf of 55 | the copyright owner. For the purposes of this definition, "submitted" 56 | means any form of electronic, verbal, or written communication sent 57 | to the Licensor or its representatives, including but not limited to 58 | communication on electronic mailing lists, source code control systems, 59 | and issue tracking systems that are managed by, or on behalf of, the 60 | Licensor for the purpose of discussing and improving the Work, but 61 | excluding communication that is conspicuously marked or otherwise 62 | designated in writing by the copyright owner as "Not a Contribution." 63 | 64 | "Contributor" shall mean Licensor and any individual or Legal Entity 65 | on behalf of whom a Contribution has been received by Licensor and 66 | subsequently incorporated within the Work. 67 | 68 | 2. Grant of Copyright License. Subject to the terms and conditions of 69 | this License, each Contributor hereby grants to You a perpetual, 70 | worldwide, non-exclusive, no-charge, royalty-free, irrevocable 71 | copyright license to reproduce, prepare Derivative Works of, 72 | publicly display, publicly perform, sublicense, and distribute the 73 | Work and such Derivative Works in Source or Object form. 74 | 75 | 3. Grant of Patent License. Subject to the terms and conditions of 76 | this License, each Contributor hereby grants to You a perpetual, 77 | worldwide, non-exclusive, no-charge, royalty-free, irrevocable 78 | (except as stated in this section) patent license to make, have made, 79 | use, offer to sell, sell, import, and otherwise transfer the Work, 80 | where such license applies only to those patent claims licensable 81 | by such Contributor that are necessarily infringed by their 82 | Contribution(s) alone or by combination of their Contribution(s) 83 | with the Work to which such Contribution(s) was submitted. If You 84 | institute patent litigation against any entity (including a 85 | cross-claim or counterclaim in a lawsuit) alleging that the Work 86 | or a Contribution incorporated within the Work constitutes direct 87 | or contributory patent infringement, then any patent licenses 88 | granted to You under this License for that Work shall terminate 89 | as of the date such litigation is filed. 90 | 91 | 4. Redistribution. You may reproduce and distribute copies of the 92 | Work or Derivative Works thereof in any medium, with or without 93 | modifications, and in Source or Object form, provided that You 94 | meet the following conditions: 95 | 96 | (a) You must give any other recipients of the Work or 97 | Derivative Works a copy of this License; and 98 | 99 | (b) You must cause any modified files to carry prominent notices 100 | stating that You changed the files; and 101 | 102 | (c) You must retain, in the Source form of any Derivative Works 103 | that You distribute, all copyright, patent, trademark, and 104 | attribution notices from the Source form of the Work, 105 | excluding those notices that do not pertain to any part of 106 | the Derivative Works; and 107 | 108 | (d) If the Work includes a "NOTICE" text file as part of its 109 | distribution, then any Derivative Works that You distribute must 110 | include a readable copy of the attribution notices contained 111 | within such NOTICE file, excluding those notices that do not 112 | pertain to any part of the Derivative Works, in at least one 113 | of the following places: within a NOTICE text file distributed 114 | as part of the Derivative Works; within the Source form or 115 | documentation, if provided along with the Derivative Works; or, 116 | within a display generated by the Derivative Works, if and 117 | wherever such third-party notices normally appear. The contents 118 | of the NOTICE file are for informational purposes only and 119 | do not modify the License. You may add Your own attribution 120 | notices within Derivative Works that You distribute, alongside 121 | or as an addendum to the NOTICE text from the Work, provided 122 | that such additional attribution notices cannot be construed 123 | as modifying the License. 124 | 125 | You may add Your own copyright statement to Your modifications and 126 | may provide additional or different license terms and conditions 127 | for use, reproduction, or distribution of Your modifications, or 128 | for any such Derivative Works as a whole, provided Your use, 129 | reproduction, and distribution of the Work otherwise complies with 130 | the conditions stated in this License. 131 | 132 | 5. Submission of Contributions. Unless You explicitly state otherwise, 133 | any Contribution intentionally submitted for inclusion in the Work 134 | by You to the Licensor shall be under the terms and conditions of 135 | this License, without any additional terms or conditions. 136 | Notwithstanding the above, nothing herein shall supersede or modify 137 | the terms of any separate license agreement you may have executed 138 | with Licensor regarding such Contributions. 139 | 140 | 6. Trademarks. This License does not grant permission to use the trade 141 | names, trademarks, service marks, or product names of the Licensor, 142 | except as required for reasonable and customary use in describing the 143 | origin of the Work and reproducing the content of the NOTICE file. 144 | 145 | 7. Disclaimer of Warranty. Unless required by applicable law or 146 | agreed to in writing, Licensor provides the Work (and each 147 | Contributor provides its Contributions) on an "AS IS" BASIS, 148 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or 149 | implied, including, without limitation, any warranties or conditions 150 | of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A 151 | PARTICULAR PURPOSE. You are solely responsible for determining the 152 | appropriateness of using or redistributing the Work and assume any 153 | risks associated with Your exercise of permissions under this License. 154 | 155 | 8. Limitation of Liability. In no event and under no legal theory, 156 | whether in tort (including negligence), contract, or otherwise, 157 | unless required by applicable law (such as deliberate and grossly 158 | negligent acts) or agreed to in writing, shall any Contributor be 159 | liable to You for damages, including any direct, indirect, special, 160 | incidental, or consequential damages of any character arising as a 161 | result of this License or out of the use or inability to use the 162 | Work (including but not limited to damages for loss of goodwill, 163 | work stoppage, computer failure or malfunction, or any and all 164 | other commercial damages or losses), even if such Contributor 165 | has been advised of the possibility of such damages. 166 | 167 | 9. Accepting Warranty or Additional Liability. While redistributing 168 | the Work or Derivative Works thereof, You may choose to offer, 169 | and charge a fee for, acceptance of support, warranty, indemnity, 170 | or other liability obligations and/or rights consistent with this 171 | License. However, in accepting such obligations, You may act only 172 | on Your own behalf and on Your sole responsibility, not on behalf 173 | of any other Contributor, and only if You agree to indemnify, 174 | defend, and hold each Contributor harmless for any liability 175 | incurred by, or claims asserted against, such Contributor by reason 176 | of your accepting any such warranty or additional liability. 177 | 178 | END OF TERMS AND CONDITIONS 179 | 180 | APPENDIX: How to apply the Apache License to your work. 181 | 182 | To apply the Apache License to your work, attach the following 183 | boilerplate notice, with the fields enclosed by brackets "[]" 184 | replaced with your own identifying information. (Don't include 185 | the brackets!) The text should be enclosed in the appropriate 186 | comment syntax for the file format. We also recommend that a 187 | file or class name and description of purpose be included on the 188 | same "printed page" as the copyright notice for easier 189 | identification within third-party archives. 190 | 191 | Copyright [2017] [Nataniel Ruiz] 192 | 193 | Licensed under the Apache License, Version 2.0 (the "License"); 194 | you may not use this file except in compliance with the License. 195 | You may obtain a copy of the License at 196 | 197 | http://www.apache.org/licenses/LICENSE-2.0 198 | 199 | Unless required by applicable law or agreed to in writing, software 200 | distributed under the License is distributed on an "AS IS" BASIS, 201 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 202 | See the License for the specific language governing permissions and 203 | limitations under the License. 204 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # Hopenet # 2 | 3 |
4 |

5 |
6 | 7 | **Hopenet** is an accurate and easy to use head pose estimation network. Models have been trained on the 300W-LP dataset and have been tested on real data with good qualitative performance. 8 | 9 | For details about the method and quantitative results please check the CVPR Workshop [paper](https://arxiv.org/abs/1710.00925). 10 | 11 |
12 |

13 |
14 | 15 | **new** [GoT trailer example video](https://youtu.be/OZdOrSLBQmI) 16 | 17 | **new** [Conan-Cruise-Car example video](https://youtu.be/Bz6eF4Nl1O8) 18 | 19 | 20 | To use please install [PyTorch](http://pytorch.org/) and [OpenCV](https://opencv.org/) (for video) - I believe that's all you need apart from usual libraries such as numpy. You need a GPU to run Hopenet (for now). 21 | 22 | To test on a video using dlib face detections (center of head will be jumpy): 23 | ```bash 24 | python code/test_on_video_dlib.py --snapshot PATH_OF_SNAPSHOT --face_model PATH_OF_DLIB_MODEL --video PATH_OF_VIDEO --output_string STRING_TO_APPEND_TO_OUTPUT --n_frames N_OF_FRAMES_TO_PROCESS --fps FPS_OF_SOURCE_VIDEO 25 | ``` 26 | To test on a video using your own face detections (we recommend using [dockerface](https://github.com/natanielruiz/dockerface), center of head will be smoother): 27 | ```bash 28 | python code/test_on_video_dockerface.py --snapshot PATH_OF_SNAPSHOT --video PATH_OF_VIDEO --bboxes FACE_BOUNDING_BOX_ANNOTATIONS --output_string STRING_TO_APPEND_TO_OUTPUT --n_frames N_OF_FRAMES_TO_PROCESS --fps FPS_OF_SOURCE_VIDEO 29 | ``` 30 | Face bounding box annotations should be in Dockerface format (n_frame x_min y_min x_max y_max confidence). 31 | 32 | Pre-trained models: 33 | 34 | [300W-LP, alpha 1](https://drive.google.com/open?id=1EJPu2sOAwrfuamTitTkw2xJ2ipmMsmD3) 35 | 36 | [300W-LP, alpha 2](https://drive.google.com/open?id=16OZdRULgUpceMKZV6U9PNFiigfjezsCY) 37 | 38 | [300W-LP, alpha 1, robust to image quality](https://drive.google.com/open?id=1m25PrSE7g9D2q2XJVMR6IA7RaCvWSzCR) 39 | 40 | For more information on what alpha stands for please read the paper. First two models are for validating paper results, if used on real data we suggest using the last model as it is more robust to image quality and blur and gives good results on video. 41 | 42 | Please open an issue if you have an problem. 43 | 44 | Some very cool implementations of this work on other platforms by some cool people: 45 | 46 | [Gluon](https://github.com/Cjiangbpcs/gazenet_mxJiang) 47 | 48 | [MXNet](https://github.com/haofanwang/mxnet-Head-Pose) 49 | 50 | [TensorFlow with Keras](https://github.com/Oreobird/tf-keras-deep-head-pose) 51 | 52 | A really cool lightweight version of HopeNet: 53 | 54 | [Deep Head Pose Light](https://github.com/OverEuro/deep-head-pose-lite) 55 | 56 | 57 | If you find Hopenet useful in your research please cite: 58 | 59 | ``` 60 | @InProceedings{Ruiz_2018_CVPR_Workshops, 61 | author = {Ruiz, Nataniel and Chong, Eunji and Rehg, James M.}, 62 | title = {Fine-Grained Head Pose Estimation Without Keypoints}, 63 | booktitle = {The IEEE Conference on Computer Vision and Pattern Recognition (CVPR) Workshops}, 64 | month = {June}, 65 | year = {2018} 66 | } 67 | ``` 68 | 69 | *Nataniel Ruiz*, *Eunji Chong*, *James M. Rehg* 70 | 71 | Georgia Institute of Technology 72 | -------------------------------------------------------------------------------- /code/datasets.py: -------------------------------------------------------------------------------- 1 | import os 2 | import numpy as np 3 | import cv2 4 | import pandas as pd 5 | 6 | import torch 7 | from torch.utils.data.dataset import Dataset 8 | from torchvision import transforms 9 | 10 | from PIL import Image, ImageFilter 11 | 12 | import utils 13 | 14 | def get_list_from_filenames(file_path): 15 | # input: relative path to .txt file with file names 16 | # output: list of relative path names 17 | with open(file_path) as f: 18 | lines = f.read().splitlines() 19 | return lines 20 | 21 | class Synhead(Dataset): 22 | def __init__(self, data_dir, csv_path, transform, test=False): 23 | column_names = ['path', 'bbox_x_min', 'bbox_y_min', 'bbox_x_max', 'bbox_y_max', 'yaw', 'pitch', 'roll'] 24 | tmp_df = pd.read_csv(csv_path, sep=',', names=column_names, index_col=False, encoding="utf-8-sig") 25 | self.data_dir = data_dir 26 | self.transform = transform 27 | self.X_train = tmp_df['path'] 28 | self.y_train = tmp_df[['bbox_x_min', 'bbox_y_min', 'bbox_x_max', 'bbox_y_max', 'yaw', 'pitch', 'roll']] 29 | self.length = len(tmp_df) 30 | self.test = test 31 | 32 | def __getitem__(self, index): 33 | path = os.path.join(self.data_dir, self.X_train.iloc[index]).strip('.jpg') + '.png' 34 | img = Image.open(path) 35 | img = img.convert('RGB') 36 | 37 | x_min, y_min, x_max, y_max, yaw, pitch, roll = self.y_train.iloc[index] 38 | x_min = float(x_min); x_max = float(x_max) 39 | y_min = float(y_min); y_max = float(y_max) 40 | yaw = -float(yaw); pitch = float(pitch); roll = float(roll) 41 | 42 | # k = 0.2 to 0.40 43 | k = np.random.random_sample() * 0.2 + 0.2 44 | x_min -= 0.6 * k * abs(x_max - x_min) 45 | y_min -= 2 * k * abs(y_max - y_min) 46 | x_max += 0.6 * k * abs(x_max - x_min) 47 | y_max += 0.6 * k * abs(y_max - y_min) 48 | 49 | width, height = img.size 50 | # Crop the face 51 | img = img.crop((int(x_min), int(y_min), int(x_max), int(y_max))) 52 | 53 | # Flip? 54 | rnd = np.random.random_sample() 55 | if rnd < 0.5: 56 | yaw = -yaw 57 | roll = -roll 58 | img = img.transpose(Image.FLIP_LEFT_RIGHT) 59 | 60 | # Blur? 61 | rnd = np.random.random_sample() 62 | if rnd < 0.05: 63 | img = img.filter(ImageFilter.BLUR) 64 | 65 | # Bin values 66 | bins = np.array(range(-99, 102, 3)) 67 | binned_pose = np.digitize([yaw, pitch, roll], bins) - 1 68 | 69 | labels = torch.LongTensor(binned_pose) 70 | cont_labels = torch.FloatTensor([yaw, pitch, roll]) 71 | 72 | if self.transform is not None: 73 | img = self.transform(img) 74 | 75 | return img, labels, cont_labels, self.X_train[index] 76 | 77 | def __len__(self): 78 | return self.length 79 | 80 | class Pose_300W_LP(Dataset): 81 | # Head pose from 300W-LP dataset 82 | def __init__(self, data_dir, filename_path, transform, img_ext='.jpg', annot_ext='.mat', image_mode='RGB'): 83 | self.data_dir = data_dir 84 | self.transform = transform 85 | self.img_ext = img_ext 86 | self.annot_ext = annot_ext 87 | 88 | filename_list = get_list_from_filenames(filename_path) 89 | 90 | self.X_train = filename_list 91 | self.y_train = filename_list 92 | self.image_mode = image_mode 93 | self.length = len(filename_list) 94 | 95 | def __getitem__(self, index): 96 | img = Image.open(os.path.join(self.data_dir, self.X_train[index] + self.img_ext)) 97 | img = img.convert(self.image_mode) 98 | mat_path = os.path.join(self.data_dir, self.y_train[index] + self.annot_ext) 99 | 100 | # Crop the face loosely 101 | pt2d = utils.get_pt2d_from_mat(mat_path) 102 | x_min = min(pt2d[0,:]) 103 | y_min = min(pt2d[1,:]) 104 | x_max = max(pt2d[0,:]) 105 | y_max = max(pt2d[1,:]) 106 | 107 | # k = 0.2 to 0.40 108 | k = np.random.random_sample() * 0.2 + 0.2 109 | x_min -= 0.6 * k * abs(x_max - x_min) 110 | y_min -= 2 * k * abs(y_max - y_min) 111 | x_max += 0.6 * k * abs(x_max - x_min) 112 | y_max += 0.6 * k * abs(y_max - y_min) 113 | img = img.crop((int(x_min), int(y_min), int(x_max), int(y_max))) 114 | 115 | # We get the pose in radians 116 | pose = utils.get_ypr_from_mat(mat_path) 117 | # And convert to degrees. 118 | pitch = pose[0] * 180 / np.pi 119 | yaw = pose[1] * 180 / np.pi 120 | roll = pose[2] * 180 / np.pi 121 | 122 | # Flip? 123 | rnd = np.random.random_sample() 124 | if rnd < 0.5: 125 | yaw = -yaw 126 | roll = -roll 127 | img = img.transpose(Image.FLIP_LEFT_RIGHT) 128 | 129 | # Blur? 130 | rnd = np.random.random_sample() 131 | if rnd < 0.05: 132 | img = img.filter(ImageFilter.BLUR) 133 | 134 | # Bin values 135 | bins = np.array(range(-99, 102, 3)) 136 | binned_pose = np.digitize([yaw, pitch, roll], bins) - 1 137 | 138 | # Get target tensors 139 | labels = binned_pose 140 | cont_labels = torch.FloatTensor([yaw, pitch, roll]) 141 | 142 | if self.transform is not None: 143 | img = self.transform(img) 144 | 145 | return img, labels, cont_labels, self.X_train[index] 146 | 147 | def __len__(self): 148 | # 122,450 149 | return self.length 150 | 151 | class Pose_300W_LP_random_ds(Dataset): 152 | # 300W-LP dataset with random downsampling 153 | def __init__(self, data_dir, filename_path, transform, img_ext='.jpg', annot_ext='.mat', image_mode='RGB'): 154 | self.data_dir = data_dir 155 | self.transform = transform 156 | self.img_ext = img_ext 157 | self.annot_ext = annot_ext 158 | 159 | filename_list = get_list_from_filenames(filename_path) 160 | 161 | self.X_train = filename_list 162 | self.y_train = filename_list 163 | self.image_mode = image_mode 164 | self.length = len(filename_list) 165 | 166 | def __getitem__(self, index): 167 | img = Image.open(os.path.join(self.data_dir, self.X_train[index] + self.img_ext)) 168 | img = img.convert(self.image_mode) 169 | mat_path = os.path.join(self.data_dir, self.y_train[index] + self.annot_ext) 170 | 171 | # Crop the face loosely 172 | pt2d = utils.get_pt2d_from_mat(mat_path) 173 | x_min = min(pt2d[0,:]) 174 | y_min = min(pt2d[1,:]) 175 | x_max = max(pt2d[0,:]) 176 | y_max = max(pt2d[1,:]) 177 | 178 | # k = 0.2 to 0.40 179 | k = np.random.random_sample() * 0.2 + 0.2 180 | x_min -= 0.6 * k * abs(x_max - x_min) 181 | y_min -= 2 * k * abs(y_max - y_min) 182 | x_max += 0.6 * k * abs(x_max - x_min) 183 | y_max += 0.6 * k * abs(y_max - y_min) 184 | img = img.crop((int(x_min), int(y_min), int(x_max), int(y_max))) 185 | 186 | # We get the pose in radians 187 | pose = utils.get_ypr_from_mat(mat_path) 188 | pitch = pose[0] * 180 / np.pi 189 | yaw = pose[1] * 180 / np.pi 190 | roll = pose[2] * 180 / np.pi 191 | 192 | ds = 1 + np.random.randint(0,4) * 5 193 | original_size = img.size 194 | img = img.resize((img.size[0] / ds, img.size[1] / ds), resample=Image.NEAREST) 195 | img = img.resize((original_size[0], original_size[1]), resample=Image.NEAREST) 196 | 197 | # Flip? 198 | rnd = np.random.random_sample() 199 | if rnd < 0.5: 200 | yaw = -yaw 201 | roll = -roll 202 | img = img.transpose(Image.FLIP_LEFT_RIGHT) 203 | 204 | # Blur? 205 | rnd = np.random.random_sample() 206 | if rnd < 0.05: 207 | img = img.filter(ImageFilter.BLUR) 208 | 209 | # Bin values 210 | bins = np.array(range(-99, 102, 3)) 211 | binned_pose = np.digitize([yaw, pitch, roll], bins) - 1 212 | 213 | # Get target tensors 214 | labels = binned_pose 215 | cont_labels = torch.FloatTensor([yaw, pitch, roll]) 216 | 217 | if self.transform is not None: 218 | img = self.transform(img) 219 | 220 | return img, labels, cont_labels, self.X_train[index] 221 | 222 | def __len__(self): 223 | # 122,450 224 | return self.length 225 | 226 | class AFLW2000(Dataset): 227 | def __init__(self, data_dir, filename_path, transform, img_ext='.jpg', annot_ext='.mat', image_mode='RGB'): 228 | self.data_dir = data_dir 229 | self.transform = transform 230 | self.img_ext = img_ext 231 | self.annot_ext = annot_ext 232 | 233 | filename_list = get_list_from_filenames(filename_path) 234 | 235 | self.X_train = filename_list 236 | self.y_train = filename_list 237 | self.image_mode = image_mode 238 | self.length = len(filename_list) 239 | 240 | def __getitem__(self, index): 241 | img = Image.open(os.path.join(self.data_dir, self.X_train[index] + self.img_ext)) 242 | img = img.convert(self.image_mode) 243 | mat_path = os.path.join(self.data_dir, self.y_train[index] + self.annot_ext) 244 | 245 | # Crop the face loosely 246 | pt2d = utils.get_pt2d_from_mat(mat_path) 247 | 248 | x_min = min(pt2d[0,:]) 249 | y_min = min(pt2d[1,:]) 250 | x_max = max(pt2d[0,:]) 251 | y_max = max(pt2d[1,:]) 252 | 253 | k = 0.20 254 | x_min -= 2 * k * abs(x_max - x_min) 255 | y_min -= 2 * k * abs(y_max - y_min) 256 | x_max += 2 * k * abs(x_max - x_min) 257 | y_max += 0.6 * k * abs(y_max - y_min) 258 | img = img.crop((int(x_min), int(y_min), int(x_max), int(y_max))) 259 | 260 | # We get the pose in radians 261 | pose = utils.get_ypr_from_mat(mat_path) 262 | # And convert to degrees. 263 | pitch = pose[0] * 180 / np.pi 264 | yaw = pose[1] * 180 / np.pi 265 | roll = pose[2] * 180 / np.pi 266 | # Bin values 267 | bins = np.array(range(-99, 102, 3)) 268 | labels = torch.LongTensor(np.digitize([yaw, pitch, roll], bins) - 1) 269 | cont_labels = torch.FloatTensor([yaw, pitch, roll]) 270 | 271 | if self.transform is not None: 272 | img = self.transform(img) 273 | 274 | return img, labels, cont_labels, self.X_train[index] 275 | 276 | def __len__(self): 277 | # 2,000 278 | return self.length 279 | 280 | class AFLW2000_ds(Dataset): 281 | # AFLW2000 dataset with fixed downsampling 282 | def __init__(self, data_dir, filename_path, transform, img_ext='.jpg', annot_ext='.mat', image_mode='RGB'): 283 | self.data_dir = data_dir 284 | self.transform = transform 285 | self.img_ext = img_ext 286 | self.annot_ext = annot_ext 287 | 288 | filename_list = get_list_from_filenames(filename_path) 289 | 290 | self.X_train = filename_list 291 | self.y_train = filename_list 292 | self.image_mode = image_mode 293 | self.length = len(filename_list) 294 | 295 | def __getitem__(self, index): 296 | img = Image.open(os.path.join(self.data_dir, self.X_train[index] + self.img_ext)) 297 | img = img.convert(self.image_mode) 298 | mat_path = os.path.join(self.data_dir, self.y_train[index] + self.annot_ext) 299 | 300 | # Crop the face loosely 301 | pt2d = utils.get_pt2d_from_mat(mat_path) 302 | x_min = min(pt2d[0,:]) 303 | y_min = min(pt2d[1,:]) 304 | x_max = max(pt2d[0,:]) 305 | y_max = max(pt2d[1,:]) 306 | 307 | k = 0.20 308 | x_min -= 2 * k * abs(x_max - x_min) 309 | y_min -= 2 * k * abs(y_max - y_min) 310 | x_max += 2 * k * abs(x_max - x_min) 311 | y_max += 0.6 * k * abs(y_max - y_min) 312 | img = img.crop((int(x_min), int(y_min), int(x_max), int(y_max))) 313 | 314 | ds = 3 # downsampling factor 315 | original_size = img.size 316 | img = img.resize((img.size[0] / ds, img.size[1] / ds), resample=Image.NEAREST) 317 | img = img.resize((original_size[0], original_size[1]), resample=Image.NEAREST) 318 | 319 | # We get the pose in radians 320 | pose = utils.get_ypr_from_mat(mat_path) 321 | # And convert to degrees. 322 | pitch = pose[0] * 180 / np.pi 323 | yaw = pose[1] * 180 / np.pi 324 | roll = pose[2] * 180 / np.pi 325 | # Bin values 326 | bins = np.array(range(-99, 102, 3)) 327 | labels = torch.LongTensor(np.digitize([yaw, pitch, roll], bins) - 1) 328 | cont_labels = torch.FloatTensor([yaw, pitch, roll]) 329 | 330 | if self.transform is not None: 331 | img = self.transform(img) 332 | 333 | return img, labels, cont_labels, self.X_train[index] 334 | 335 | def __len__(self): 336 | # 2,000 337 | return self.length 338 | 339 | class AFLW_aug(Dataset): 340 | # AFLW dataset with flipping 341 | def __init__(self, data_dir, filename_path, transform, img_ext='.jpg', annot_ext='.txt', image_mode='RGB'): 342 | self.data_dir = data_dir 343 | self.transform = transform 344 | self.img_ext = img_ext 345 | self.annot_ext = annot_ext 346 | 347 | filename_list = get_list_from_filenames(filename_path) 348 | 349 | self.X_train = filename_list 350 | self.y_train = filename_list 351 | self.image_mode = image_mode 352 | self.length = len(filename_list) 353 | 354 | def __getitem__(self, index): 355 | img = Image.open(os.path.join(self.data_dir, self.X_train[index] + self.img_ext)) 356 | img = img.convert(self.image_mode) 357 | txt_path = os.path.join(self.data_dir, self.y_train[index] + self.annot_ext) 358 | 359 | # We get the pose in radians 360 | annot = open(txt_path, 'r') 361 | line = annot.readline().split(' ') 362 | pose = [float(line[1]), float(line[2]), float(line[3])] 363 | # And convert to degrees. 364 | yaw = pose[0] * 180 / np.pi 365 | pitch = pose[1] * 180 / np.pi 366 | roll = pose[2] * 180 / np.pi 367 | # Fix the roll in AFLW 368 | roll *= -1 369 | 370 | # Augment 371 | # Flip? 372 | rnd = np.random.random_sample() 373 | if rnd < 0.5: 374 | yaw = -yaw 375 | roll = -roll 376 | img = img.transpose(Image.FLIP_LEFT_RIGHT) 377 | 378 | # Bin values 379 | bins = np.array(range(-99, 102, 3)) 380 | labels = torch.LongTensor(np.digitize([yaw, pitch, roll], bins) - 1) 381 | cont_labels = torch.FloatTensor([yaw, pitch, roll]) 382 | 383 | if self.transform is not None: 384 | img = self.transform(img) 385 | 386 | return img, labels, cont_labels, self.X_train[index] 387 | 388 | def __len__(self): 389 | # train: 18,863 390 | # test: 1,966 391 | return self.length 392 | 393 | class AFLW(Dataset): 394 | def __init__(self, data_dir, filename_path, transform, img_ext='.jpg', annot_ext='.txt', image_mode='RGB'): 395 | self.data_dir = data_dir 396 | self.transform = transform 397 | self.img_ext = img_ext 398 | self.annot_ext = annot_ext 399 | 400 | filename_list = get_list_from_filenames(filename_path) 401 | 402 | self.X_train = filename_list 403 | self.y_train = filename_list 404 | self.image_mode = image_mode 405 | self.length = len(filename_list) 406 | 407 | def __getitem__(self, index): 408 | img = Image.open(os.path.join(self.data_dir, self.X_train[index] + self.img_ext)) 409 | img = img.convert(self.image_mode) 410 | txt_path = os.path.join(self.data_dir, self.y_train[index] + self.annot_ext) 411 | 412 | # We get the pose in radians 413 | annot = open(txt_path, 'r') 414 | line = annot.readline().split(' ') 415 | pose = [float(line[1]), float(line[2]), float(line[3])] 416 | # And convert to degrees. 417 | yaw = pose[0] * 180 / np.pi 418 | pitch = pose[1] * 180 / np.pi 419 | roll = pose[2] * 180 / np.pi 420 | # Fix the roll in AFLW 421 | roll *= -1 422 | # Bin values 423 | bins = np.array(range(-99, 102, 3)) 424 | labels = torch.LongTensor(np.digitize([yaw, pitch, roll], bins) - 1) 425 | cont_labels = torch.FloatTensor([yaw, pitch, roll]) 426 | 427 | if self.transform is not None: 428 | img = self.transform(img) 429 | 430 | return img, labels, cont_labels, self.X_train[index] 431 | 432 | def __len__(self): 433 | # train: 18,863 434 | # test: 1,966 435 | return self.length 436 | 437 | class AFW(Dataset): 438 | def __init__(self, data_dir, filename_path, transform, img_ext='.jpg', annot_ext='.txt', image_mode='RGB'): 439 | self.data_dir = data_dir 440 | self.transform = transform 441 | self.img_ext = img_ext 442 | self.annot_ext = annot_ext 443 | 444 | filename_list = get_list_from_filenames(filename_path) 445 | 446 | self.X_train = filename_list 447 | self.y_train = filename_list 448 | self.image_mode = image_mode 449 | self.length = len(filename_list) 450 | 451 | def __getitem__(self, index): 452 | txt_path = os.path.join(self.data_dir, self.y_train[index] + self.annot_ext) 453 | img_name = self.X_train[index].split('_')[0] 454 | 455 | img = Image.open(os.path.join(self.data_dir, img_name + self.img_ext)) 456 | img = img.convert(self.image_mode) 457 | txt_path = os.path.join(self.data_dir, self.y_train[index] + self.annot_ext) 458 | 459 | # We get the pose in degrees 460 | annot = open(txt_path, 'r') 461 | line = annot.readline().split(' ') 462 | yaw, pitch, roll = [float(line[1]), float(line[2]), float(line[3])] 463 | 464 | # Crop the face loosely 465 | k = 0.32 466 | x1 = float(line[4]) 467 | y1 = float(line[5]) 468 | x2 = float(line[6]) 469 | y2 = float(line[7]) 470 | x1 -= 0.8 * k * abs(x2 - x1) 471 | y1 -= 2 * k * abs(y2 - y1) 472 | x2 += 0.8 * k * abs(x2 - x1) 473 | y2 += 1 * k * abs(y2 - y1) 474 | 475 | img = img.crop((int(x1), int(y1), int(x2), int(y2))) 476 | 477 | # Bin values 478 | bins = np.array(range(-99, 102, 3)) 479 | labels = torch.LongTensor(np.digitize([yaw, pitch, roll], bins) - 1) 480 | cont_labels = torch.FloatTensor([yaw, pitch, roll]) 481 | 482 | if self.transform is not None: 483 | img = self.transform(img) 484 | 485 | return img, labels, cont_labels, self.X_train[index] 486 | 487 | def __len__(self): 488 | # Around 200 489 | return self.length 490 | 491 | class BIWI(Dataset): 492 | def __init__(self, data_dir, filename_path, transform, img_ext='.png', annot_ext='.txt', image_mode='RGB'): 493 | self.data_dir = data_dir 494 | self.transform = transform 495 | self.img_ext = img_ext 496 | self.annot_ext = annot_ext 497 | 498 | filename_list = get_list_from_filenames(filename_path) 499 | 500 | self.X_train = filename_list 501 | self.y_train = filename_list 502 | self.image_mode = image_mode 503 | self.length = len(filename_list) 504 | 505 | def __getitem__(self, index): 506 | img = Image.open(os.path.join(self.data_dir, self.X_train[index] + '_rgb' + self.img_ext)) 507 | img = img.convert(self.image_mode) 508 | pose_path = os.path.join(self.data_dir, self.y_train[index] + '_pose' + self.annot_ext) 509 | 510 | y_train_list = self.y_train[index].split('/') 511 | bbox_path = os.path.join(self.data_dir, y_train_list[0] + '/dockerface-' + y_train_list[-1] + '_rgb' + self.annot_ext) 512 | 513 | # Load bounding box 514 | bbox = open(bbox_path, 'r') 515 | line = bbox.readline().split(' ') 516 | if len(line) < 4: 517 | x_min, y_min, x_max, y_max = 0, 0, img.size[0], img.size[1] 518 | else: 519 | x_min, y_min, x_max, y_max = [float(line[1]), float(line[2]), float(line[3]), float(line[4])] 520 | bbox.close() 521 | 522 | # Load pose in degrees 523 | pose_annot = open(pose_path, 'r') 524 | R = [] 525 | for line in pose_annot: 526 | line = line.strip('\n').split(' ') 527 | l = [] 528 | if line[0] != '': 529 | for nb in line: 530 | if nb == '': 531 | continue 532 | l.append(float(nb)) 533 | R.append(l) 534 | 535 | R = np.array(R) 536 | T = R[3,:] 537 | R = R[:3,:] 538 | pose_annot.close() 539 | 540 | R = np.transpose(R) 541 | 542 | roll = -np.arctan2(R[1][0], R[0][0]) * 180 / np.pi 543 | yaw = -np.arctan2(-R[2][0], np.sqrt(R[2][1] ** 2 + R[2][2] ** 2)) * 180 / np.pi 544 | pitch = np.arctan2(R[2][1], R[2][2]) * 180 / np.pi 545 | 546 | # Loosely crop face 547 | k = 0.35 548 | x_min -= 0.6 * k * abs(x_max - x_min) 549 | y_min -= k * abs(y_max - y_min) 550 | x_max += 0.6 * k * abs(x_max - x_min) 551 | y_max += 0.6 * k * abs(y_max - y_min) 552 | img = img.crop((int(x_min), int(y_min), int(x_max), int(y_max))) 553 | 554 | # Bin values 555 | bins = np.array(range(-99, 102, 3)) 556 | binned_pose = np.digitize([yaw, pitch, roll], bins) - 1 557 | 558 | labels = torch.LongTensor(binned_pose) 559 | cont_labels = torch.FloatTensor([yaw, pitch, roll]) 560 | 561 | if self.transform is not None: 562 | img = self.transform(img) 563 | 564 | return img, labels, cont_labels, self.X_train[index] 565 | 566 | def __len__(self): 567 | # 15,667 568 | return self.length 569 | -------------------------------------------------------------------------------- /code/hopenet.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | from torch.autograd import Variable 4 | import math 5 | import torch.nn.functional as F 6 | 7 | class Hopenet(nn.Module): 8 | # Hopenet with 3 output layers for yaw, pitch and roll 9 | # Predicts Euler angles by binning and regression with the expected value 10 | def __init__(self, block, layers, num_bins): 11 | self.inplanes = 64 12 | super(Hopenet, self).__init__() 13 | self.conv1 = nn.Conv2d(3, 64, kernel_size=7, stride=2, padding=3, 14 | bias=False) 15 | self.bn1 = nn.BatchNorm2d(64) 16 | self.relu = nn.ReLU(inplace=True) 17 | self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1) 18 | self.layer1 = self._make_layer(block, 64, layers[0]) 19 | self.layer2 = self._make_layer(block, 128, layers[1], stride=2) 20 | self.layer3 = self._make_layer(block, 256, layers[2], stride=2) 21 | self.layer4 = self._make_layer(block, 512, layers[3], stride=2) 22 | self.avgpool = nn.AvgPool2d(7) 23 | self.fc_yaw = nn.Linear(512 * block.expansion, num_bins) 24 | self.fc_pitch = nn.Linear(512 * block.expansion, num_bins) 25 | self.fc_roll = nn.Linear(512 * block.expansion, num_bins) 26 | 27 | # Vestigial layer from previous experiments 28 | self.fc_finetune = nn.Linear(512 * block.expansion + 3, 3) 29 | 30 | for m in self.modules(): 31 | if isinstance(m, nn.Conv2d): 32 | n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels 33 | m.weight.data.normal_(0, math.sqrt(2. / n)) 34 | elif isinstance(m, nn.BatchNorm2d): 35 | m.weight.data.fill_(1) 36 | m.bias.data.zero_() 37 | 38 | def _make_layer(self, block, planes, blocks, stride=1): 39 | downsample = None 40 | if stride != 1 or self.inplanes != planes * block.expansion: 41 | downsample = nn.Sequential( 42 | nn.Conv2d(self.inplanes, planes * block.expansion, 43 | kernel_size=1, stride=stride, bias=False), 44 | nn.BatchNorm2d(planes * block.expansion), 45 | ) 46 | 47 | layers = [] 48 | layers.append(block(self.inplanes, planes, stride, downsample)) 49 | self.inplanes = planes * block.expansion 50 | for i in range(1, blocks): 51 | layers.append(block(self.inplanes, planes)) 52 | 53 | return nn.Sequential(*layers) 54 | 55 | def forward(self, x): 56 | x = self.conv1(x) 57 | x = self.bn1(x) 58 | x = self.relu(x) 59 | x = self.maxpool(x) 60 | 61 | x = self.layer1(x) 62 | x = self.layer2(x) 63 | x = self.layer3(x) 64 | x = self.layer4(x) 65 | 66 | x = self.avgpool(x) 67 | x = x.view(x.size(0), -1) 68 | pre_yaw = self.fc_yaw(x) 69 | pre_pitch = self.fc_pitch(x) 70 | pre_roll = self.fc_roll(x) 71 | 72 | return pre_yaw, pre_pitch, pre_roll 73 | 74 | class ResNet(nn.Module): 75 | # ResNet for regression of 3 Euler angles. 76 | def __init__(self, block, layers, num_classes=1000): 77 | self.inplanes = 64 78 | super(ResNet, self).__init__() 79 | self.conv1 = nn.Conv2d(3, 64, kernel_size=7, stride=2, padding=3, 80 | bias=False) 81 | self.bn1 = nn.BatchNorm2d(64) 82 | self.relu = nn.ReLU(inplace=True) 83 | self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1) 84 | self.layer1 = self._make_layer(block, 64, layers[0]) 85 | self.layer2 = self._make_layer(block, 128, layers[1], stride=2) 86 | self.layer3 = self._make_layer(block, 256, layers[2], stride=2) 87 | self.layer4 = self._make_layer(block, 512, layers[3], stride=2) 88 | self.avgpool = nn.AvgPool2d(7) 89 | self.fc_angles = nn.Linear(512 * block.expansion, num_classes) 90 | 91 | for m in self.modules(): 92 | if isinstance(m, nn.Conv2d): 93 | n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels 94 | m.weight.data.normal_(0, math.sqrt(2. / n)) 95 | elif isinstance(m, nn.BatchNorm2d): 96 | m.weight.data.fill_(1) 97 | m.bias.data.zero_() 98 | 99 | def _make_layer(self, block, planes, blocks, stride=1): 100 | downsample = None 101 | if stride != 1 or self.inplanes != planes * block.expansion: 102 | downsample = nn.Sequential( 103 | nn.Conv2d(self.inplanes, planes * block.expansion, 104 | kernel_size=1, stride=stride, bias=False), 105 | nn.BatchNorm2d(planes * block.expansion), 106 | ) 107 | 108 | layers = [] 109 | layers.append(block(self.inplanes, planes, stride, downsample)) 110 | self.inplanes = planes * block.expansion 111 | for i in range(1, blocks): 112 | layers.append(block(self.inplanes, planes)) 113 | 114 | return nn.Sequential(*layers) 115 | 116 | def forward(self, x): 117 | x = self.conv1(x) 118 | x = self.bn1(x) 119 | x = self.relu(x) 120 | x = self.maxpool(x) 121 | 122 | x = self.layer1(x) 123 | x = self.layer2(x) 124 | x = self.layer3(x) 125 | x = self.layer4(x) 126 | 127 | x = self.avgpool(x) 128 | x = x.view(x.size(0), -1) 129 | x = self.fc_angles(x) 130 | return x 131 | 132 | class AlexNet(nn.Module): 133 | # AlexNet laid out as a Hopenet - classify Euler angles in bins and 134 | # regress the expected value. 135 | def __init__(self, num_bins): 136 | super(AlexNet, self).__init__() 137 | self.features = nn.Sequential( 138 | nn.Conv2d(3, 64, kernel_size=11, stride=4, padding=2), 139 | nn.ReLU(inplace=True), 140 | nn.MaxPool2d(kernel_size=3, stride=2), 141 | nn.Conv2d(64, 192, kernel_size=5, padding=2), 142 | nn.ReLU(inplace=True), 143 | nn.MaxPool2d(kernel_size=3, stride=2), 144 | nn.Conv2d(192, 384, kernel_size=3, padding=1), 145 | nn.ReLU(inplace=True), 146 | nn.Conv2d(384, 256, kernel_size=3, padding=1), 147 | nn.ReLU(inplace=True), 148 | nn.Conv2d(256, 256, kernel_size=3, padding=1), 149 | nn.ReLU(inplace=True), 150 | nn.MaxPool2d(kernel_size=3, stride=2), 151 | ) 152 | self.classifier = nn.Sequential( 153 | nn.Dropout(), 154 | nn.Linear(256 * 6 * 6, 4096), 155 | nn.ReLU(inplace=True), 156 | nn.Dropout(), 157 | nn.Linear(4096, 4096), 158 | nn.ReLU(inplace=True), 159 | ) 160 | self.fc_yaw = nn.Linear(4096, num_bins) 161 | self.fc_pitch = nn.Linear(4096, num_bins) 162 | self.fc_roll = nn.Linear(4096, num_bins) 163 | 164 | def forward(self, x): 165 | x = self.features(x) 166 | x = x.view(x.size(0), 256 * 6 * 6) 167 | x = self.classifier(x) 168 | yaw = self.fc_yaw(x) 169 | pitch = self.fc_pitch(x) 170 | roll = self.fc_roll(x) 171 | return yaw, pitch, roll 172 | -------------------------------------------------------------------------------- /code/test_alexnet.py: -------------------------------------------------------------------------------- 1 | import sys, os, argparse 2 | 3 | import numpy as np 4 | import cv2 5 | import matplotlib.pyplot as plt 6 | 7 | import torch 8 | import torch.nn as nn 9 | from torch.autograd import Variable 10 | from torch.utils.data import DataLoader 11 | from torchvision import transforms 12 | import torch.backends.cudnn as cudnn 13 | import torchvision 14 | import torch.nn.functional as F 15 | 16 | import datasets, hopenet, utils 17 | 18 | def parse_args(): 19 | """Parse input arguments.""" 20 | parser = argparse.ArgumentParser(description='Head pose estimation using the Hopenet network.') 21 | parser.add_argument('--gpu', dest='gpu_id', help='GPU device id to use [0]', 22 | default=0, type=int) 23 | parser.add_argument('--data_dir', dest='data_dir', help='Directory path for data.', 24 | default='', type=str) 25 | parser.add_argument('--filename_list', dest='filename_list', help='Path to text file containing relative paths for every example.', 26 | default='', type=str) 27 | parser.add_argument('--snapshot', dest='snapshot', help='Name of model snapshot.', 28 | default='', type=str) 29 | parser.add_argument('--batch_size', dest='batch_size', help='Batch size.', 30 | default=1, type=int) 31 | parser.add_argument('--save_viz', dest='save_viz', help='Save images with pose cube.', 32 | default=False, type=bool) 33 | parser.add_argument('--dataset', dest='dataset', help='Dataset type.', default='AFLW2000', type=str) 34 | 35 | args = parser.parse_args() 36 | 37 | return args 38 | 39 | if __name__ == '__main__': 40 | args = parse_args() 41 | 42 | cudnn.enabled = True 43 | gpu = args.gpu_id 44 | snapshot_path = args.snapshot 45 | 46 | model = hopenet.AlexNet(66) 47 | 48 | print 'Loading snapshot.' 49 | # Load snapshot 50 | saved_state_dict = torch.load(snapshot_path) 51 | model.load_state_dict(saved_state_dict) 52 | 53 | print 'Loading data.' 54 | 55 | transformations = transforms.Compose([transforms.Scale(224), 56 | transforms.CenterCrop(224), transforms.ToTensor(), 57 | transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])]) 58 | 59 | if args.dataset == 'Pose_300W_LP': 60 | pose_dataset = datasets.Pose_300W_LP(args.data_dir, args.filename_list, transformations) 61 | elif args.dataset == 'Pose_300W_LP_random_ds': 62 | pose_dataset = datasets.Pose_300W_LP_random_ds(args.data_dir, args.filename_list, transformations) 63 | elif args.dataset == 'AFLW2000': 64 | pose_dataset = datasets.AFLW2000(args.data_dir, args.filename_list, transformations) 65 | elif args.dataset == 'AFLW2000_ds': 66 | pose_dataset = datasets.AFLW2000_ds(args.data_dir, args.filename_list, transformations) 67 | elif args.dataset == 'BIWI': 68 | pose_dataset = datasets.BIWI(args.data_dir, args.filename_list, transformations) 69 | elif args.dataset == 'AFLW': 70 | pose_dataset = datasets.AFLW(args.data_dir, args.filename_list, transformations) 71 | elif args.dataset == 'AFLW_aug': 72 | pose_dataset = datasets.AFLW_aug(args.data_dir, args.filename_list, transformations) 73 | elif args.dataset == 'AFW': 74 | pose_dataset = datasets.AFW(args.data_dir, args.filename_list, transformations) 75 | else: 76 | print 'Error: not a valid dataset name' 77 | sys.exit() 78 | test_loader = torch.utils.data.DataLoader(dataset=pose_dataset, 79 | batch_size=args.batch_size, 80 | num_workers=2) 81 | 82 | model.cuda(gpu) 83 | 84 | print 'Ready to test network.' 85 | 86 | # Test the Model 87 | model.eval() # Change model to 'eval' mode (BN uses moving mean/var). 88 | total = 0 89 | 90 | idx_tensor = [idx for idx in xrange(66)] 91 | idx_tensor = torch.FloatTensor(idx_tensor).cuda(gpu) 92 | 93 | yaw_error = .0 94 | pitch_error = .0 95 | roll_error = .0 96 | 97 | l1loss = torch.nn.L1Loss(size_average=False) 98 | 99 | for i, (images, labels, cont_labels, name) in enumerate(test_loader): 100 | images = Variable(images).cuda(gpu) 101 | total += cont_labels.size(0) 102 | label_yaw = cont_labels[:,0].float() 103 | label_pitch = cont_labels[:,1].float() 104 | label_roll = cont_labels[:,2].float() 105 | 106 | yaw, pitch, roll = model(images) 107 | 108 | # Binned predictions 109 | _, yaw_bpred = torch.max(yaw.data, 1) 110 | _, pitch_bpred = torch.max(pitch.data, 1) 111 | _, roll_bpred = torch.max(roll.data, 1) 112 | 113 | # Continuous predictions 114 | yaw_predicted = utils.softmax_temperature(yaw.data, 1) 115 | pitch_predicted = utils.softmax_temperature(pitch.data, 1) 116 | roll_predicted = utils.softmax_temperature(roll.data, 1) 117 | 118 | yaw_predicted = torch.sum(yaw_predicted * idx_tensor, 1).cpu() * 3 - 99 119 | pitch_predicted = torch.sum(pitch_predicted * idx_tensor, 1).cpu() * 3 - 99 120 | roll_predicted = torch.sum(roll_predicted * idx_tensor, 1).cpu() * 3 - 99 121 | 122 | # Mean absolute error 123 | yaw_error += torch.sum(torch.abs(yaw_predicted - label_yaw)) 124 | pitch_error += torch.sum(torch.abs(pitch_predicted - label_pitch)) 125 | roll_error += torch.sum(torch.abs(roll_predicted - label_roll)) 126 | 127 | # Save first image in batch with pose cube or axis. 128 | if args.save_viz: 129 | name = name[0] 130 | if args.dataset == 'BIWI': 131 | cv2_img = cv2.imread(os.path.join(args.data_dir, name + '_rgb.png')) 132 | else: 133 | cv2_img = cv2.imread(os.path.join(args.data_dir, name + '.jpg')) 134 | if args.batch_size == 1: 135 | error_string = 'y %.2f, p %.2f, r %.2f' % (torch.sum(torch.abs(yaw_predicted - label_yaw)), torch.sum(torch.abs(pitch_predicted - label_pitch)), torch.sum(torch.abs(roll_predicted - label_roll))) 136 | cv2.putText(cv2_img, error_string, (30, cv2_img.shape[0]- 30), fontFace=1, fontScale=1, color=(0,0,255), thickness=1) 137 | # utils.plot_pose_cube(cv2_img, yaw_predicted[0], pitch_predicted[0], roll_predicted[0], size=100) 138 | utils.draw_axis(cv2_img, yaw_predicted[0], pitch_predicted[0], roll_predicted[0], tdx = 200, tdy= 200, size=100) 139 | cv2.imwrite(os.path.join('output/images', name + '.jpg'), cv2_img) 140 | 141 | print('Test error in degrees of the model on the ' + str(total) + 142 | ' test images. Yaw: %.4f, Pitch: %.4f, Roll: %.4f' % (yaw_error / total, 143 | pitch_error / total, roll_error / total)) 144 | -------------------------------------------------------------------------------- /code/test_hopenet.py: -------------------------------------------------------------------------------- 1 | import sys, os, argparse 2 | 3 | import numpy as np 4 | import cv2 5 | import matplotlib.pyplot as plt 6 | 7 | import torch 8 | import torch.nn as nn 9 | from torch.autograd import Variable 10 | from torch.utils.data import DataLoader 11 | from torchvision import transforms 12 | import torch.backends.cudnn as cudnn 13 | import torchvision 14 | import torch.nn.functional as F 15 | 16 | import datasets, hopenet, utils 17 | 18 | def parse_args(): 19 | """Parse input arguments.""" 20 | parser = argparse.ArgumentParser(description='Head pose estimation using the Hopenet network.') 21 | parser.add_argument('--gpu', dest='gpu_id', help='GPU device id to use [0]', 22 | default=0, type=int) 23 | parser.add_argument('--data_dir', dest='data_dir', help='Directory path for data.', 24 | default='', type=str) 25 | parser.add_argument('--filename_list', dest='filename_list', help='Path to text file containing relative paths for every example.', 26 | default='', type=str) 27 | parser.add_argument('--snapshot', dest='snapshot', help='Name of model snapshot.', 28 | default='', type=str) 29 | parser.add_argument('--batch_size', dest='batch_size', help='Batch size.', 30 | default=1, type=int) 31 | parser.add_argument('--save_viz', dest='save_viz', help='Save images with pose cube.', 32 | default=False, type=bool) 33 | parser.add_argument('--dataset', dest='dataset', help='Dataset type.', default='AFLW2000', type=str) 34 | 35 | args = parser.parse_args() 36 | 37 | return args 38 | 39 | if __name__ == '__main__': 40 | args = parse_args() 41 | 42 | cudnn.enabled = True 43 | gpu = args.gpu_id 44 | snapshot_path = args.snapshot 45 | 46 | # ResNet50 structure 47 | model = hopenet.Hopenet(torchvision.models.resnet.Bottleneck, [3, 4, 6, 3], 66) 48 | 49 | print 'Loading snapshot.' 50 | # Load snapshot 51 | saved_state_dict = torch.load(snapshot_path) 52 | model.load_state_dict(saved_state_dict) 53 | 54 | print 'Loading data.' 55 | 56 | transformations = transforms.Compose([transforms.Scale(224), 57 | transforms.CenterCrop(224), transforms.ToTensor(), 58 | transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])]) 59 | 60 | if args.dataset == 'Pose_300W_LP': 61 | pose_dataset = datasets.Pose_300W_LP(args.data_dir, args.filename_list, transformations) 62 | elif args.dataset == 'Pose_300W_LP_random_ds': 63 | pose_dataset = datasets.Pose_300W_LP_random_ds(args.data_dir, args.filename_list, transformations) 64 | elif args.dataset == 'AFLW2000': 65 | pose_dataset = datasets.AFLW2000(args.data_dir, args.filename_list, transformations) 66 | elif args.dataset == 'AFLW2000_ds': 67 | pose_dataset = datasets.AFLW2000_ds(args.data_dir, args.filename_list, transformations) 68 | elif args.dataset == 'BIWI': 69 | pose_dataset = datasets.BIWI(args.data_dir, args.filename_list, transformations) 70 | elif args.dataset == 'AFLW': 71 | pose_dataset = datasets.AFLW(args.data_dir, args.filename_list, transformations) 72 | elif args.dataset == 'AFLW_aug': 73 | pose_dataset = datasets.AFLW_aug(args.data_dir, args.filename_list, transformations) 74 | elif args.dataset == 'AFW': 75 | pose_dataset = datasets.AFW(args.data_dir, args.filename_list, transformations) 76 | else: 77 | print 'Error: not a valid dataset name' 78 | sys.exit() 79 | test_loader = torch.utils.data.DataLoader(dataset=pose_dataset, 80 | batch_size=args.batch_size, 81 | num_workers=2) 82 | 83 | model.cuda(gpu) 84 | 85 | print 'Ready to test network.' 86 | 87 | # Test the Model 88 | model.eval() # Change model to 'eval' mode (BN uses moving mean/var). 89 | total = 0 90 | 91 | idx_tensor = [idx for idx in xrange(66)] 92 | idx_tensor = torch.FloatTensor(idx_tensor).cuda(gpu) 93 | 94 | yaw_error = .0 95 | pitch_error = .0 96 | roll_error = .0 97 | 98 | l1loss = torch.nn.L1Loss(size_average=False) 99 | 100 | for i, (images, labels, cont_labels, name) in enumerate(test_loader): 101 | images = Variable(images).cuda(gpu) 102 | total += cont_labels.size(0) 103 | 104 | label_yaw = cont_labels[:,0].float() 105 | label_pitch = cont_labels[:,1].float() 106 | label_roll = cont_labels[:,2].float() 107 | 108 | yaw, pitch, roll = model(images) 109 | 110 | # Binned predictions 111 | _, yaw_bpred = torch.max(yaw.data, 1) 112 | _, pitch_bpred = torch.max(pitch.data, 1) 113 | _, roll_bpred = torch.max(roll.data, 1) 114 | 115 | # Continuous predictions 116 | yaw_predicted = utils.softmax_temperature(yaw.data, 1) 117 | pitch_predicted = utils.softmax_temperature(pitch.data, 1) 118 | roll_predicted = utils.softmax_temperature(roll.data, 1) 119 | 120 | yaw_predicted = torch.sum(yaw_predicted * idx_tensor, 1).cpu() * 3 - 99 121 | pitch_predicted = torch.sum(pitch_predicted * idx_tensor, 1).cpu() * 3 - 99 122 | roll_predicted = torch.sum(roll_predicted * idx_tensor, 1).cpu() * 3 - 99 123 | 124 | # Mean absolute error 125 | yaw_error += torch.sum(torch.abs(yaw_predicted - label_yaw)) 126 | pitch_error += torch.sum(torch.abs(pitch_predicted - label_pitch)) 127 | roll_error += torch.sum(torch.abs(roll_predicted - label_roll)) 128 | 129 | # Save first image in batch with pose cube or axis. 130 | if args.save_viz: 131 | name = name[0] 132 | if args.dataset == 'BIWI': 133 | cv2_img = cv2.imread(os.path.join(args.data_dir, name + '_rgb.png')) 134 | else: 135 | cv2_img = cv2.imread(os.path.join(args.data_dir, name + '.jpg')) 136 | if args.batch_size == 1: 137 | error_string = 'y %.2f, p %.2f, r %.2f' % (torch.sum(torch.abs(yaw_predicted - label_yaw)), torch.sum(torch.abs(pitch_predicted - label_pitch)), torch.sum(torch.abs(roll_predicted - label_roll))) 138 | cv2.putText(cv2_img, error_string, (30, cv2_img.shape[0]- 30), fontFace=1, fontScale=1, color=(0,0,255), thickness=2) 139 | # utils.plot_pose_cube(cv2_img, yaw_predicted[0], pitch_predicted[0], roll_predicted[0], size=100) 140 | utils.draw_axis(cv2_img, yaw_predicted[0], pitch_predicted[0], roll_predicted[0], tdx = 200, tdy= 200, size=100) 141 | cv2.imwrite(os.path.join('output/images', name + '.jpg'), cv2_img) 142 | 143 | print('Test error in degrees of the model on the ' + str(total) + 144 | ' test images. Yaw: %.4f, Pitch: %.4f, Roll: %.4f' % (yaw_error / total, 145 | pitch_error / total, roll_error / total)) 146 | -------------------------------------------------------------------------------- /code/test_on_video.py: -------------------------------------------------------------------------------- 1 | import sys, os, argparse 2 | 3 | import numpy as np 4 | import cv2 5 | import matplotlib.pyplot as plt 6 | 7 | import torch 8 | import torch.nn as nn 9 | from torch.autograd import Variable 10 | from torch.utils.data import DataLoader 11 | from torchvision import transforms 12 | import torch.backends.cudnn as cudnn 13 | import torchvision 14 | import torch.nn.functional as F 15 | from PIL import Image 16 | 17 | import datasets, hopenet, utils 18 | 19 | def parse_args(): 20 | """Parse input arguments.""" 21 | parser = argparse.ArgumentParser(description='Head pose estimation using the Hopenet network.') 22 | parser.add_argument('--gpu', dest='gpu_id', help='GPU device id to use [0]', 23 | default=0, type=int) 24 | parser.add_argument('--snapshot', dest='snapshot', help='Path of model snapshot.', 25 | default='', type=str) 26 | parser.add_argument('--video', dest='video_path', help='Path of video') 27 | parser.add_argument('--bboxes', dest='bboxes', help='Bounding box annotations of frames') 28 | parser.add_argument('--output_string', dest='output_string', help='String appended to output file') 29 | parser.add_argument('--n_frames', dest='n_frames', help='Number of frames', type=int) 30 | parser.add_argument('--fps', dest='fps', help='Frames per second of source video', type=float, default=30.) 31 | args = parser.parse_args() 32 | return args 33 | 34 | if __name__ == '__main__': 35 | args = parse_args() 36 | 37 | cudnn.enabled = True 38 | 39 | batch_size = 1 40 | gpu = args.gpu_id 41 | snapshot_path = args.snapshot 42 | out_dir = 'output/video' 43 | video_path = args.video_path 44 | 45 | if not os.path.exists(out_dir): 46 | os.makedirs(out_dir) 47 | 48 | if not os.path.exists(args.video_path): 49 | sys.exit('Video does not exist') 50 | 51 | # ResNet50 structure 52 | model = hopenet.Hopenet(torchvision.models.resnet.Bottleneck, [3, 4, 6, 3], 66) 53 | 54 | print 'Loading snapshot.' 55 | # Load snapshot 56 | saved_state_dict = torch.load(snapshot_path) 57 | model.load_state_dict(saved_state_dict) 58 | 59 | print 'Loading data.' 60 | 61 | transformations = transforms.Compose([transforms.Scale(224), 62 | transforms.CenterCrop(224), transforms.ToTensor(), 63 | transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])]) 64 | 65 | model.cuda(gpu) 66 | 67 | print 'Ready to test network.' 68 | 69 | # Test the Model 70 | model.eval() # Change model to 'eval' mode (BN uses moving mean/var). 71 | total = 0 72 | 73 | idx_tensor = [idx for idx in xrange(66)] 74 | idx_tensor = torch.FloatTensor(idx_tensor).cuda(gpu) 75 | 76 | video = cv2.VideoCapture(video_path) 77 | 78 | # New cv2 79 | width = int(video.get(cv2.CAP_PROP_FRAME_WIDTH)) # float 80 | height = int(video.get(cv2.CAP_PROP_FRAME_HEIGHT)) # float 81 | 82 | # Define the codec and create VideoWriter object 83 | fourcc = cv2.VideoWriter_fourcc(*'MJPG') 84 | out = cv2.VideoWriter('output/video/output-%s.avi' % args.output_string, fourcc, args.fps, (width, height)) 85 | 86 | # # Old cv2 87 | # width = int(video.get(cv2.cv.CV_CAP_PROP_FRAME_WIDTH)) # float 88 | # height = int(video.get(cv2.cv.CV_CAP_PROP_FRAME_HEIGHT)) # float 89 | # 90 | # # Define the codec and create VideoWriter object 91 | # fourcc = cv2.cv.CV_FOURCC(*'MJPG') 92 | # out = cv2.VideoWriter('output/video/output-%s.avi' % args.output_string, fourcc, 30.0, (width, height)) 93 | 94 | txt_out = open('output/video/output-%s.txt' % args.output_string, 'w') 95 | 96 | frame_num = 1 97 | 98 | with open(args.bboxes, 'r') as f: 99 | bbox_line_list = f.read().splitlines() 100 | 101 | idx = 0 102 | while idx < len(bbox_line_list): 103 | line = bbox_line_list[idx] 104 | line = line.strip('\n') 105 | line = line.split(' ') 106 | det_frame_num = int(line[0]) 107 | 108 | print frame_num 109 | 110 | # Stop at a certain frame number 111 | if frame_num > args.n_frames: 112 | break 113 | 114 | # Save all frames as they are if they don't have bbox annotation. 115 | while frame_num < det_frame_num: 116 | ret, frame = video.read() 117 | if ret == False: 118 | out.release() 119 | video.release() 120 | txt_out.close() 121 | sys.exit(0) 122 | # out.write(frame) 123 | frame_num += 1 124 | 125 | # Start processing frame with bounding box 126 | ret,frame = video.read() 127 | if ret == False: 128 | break 129 | cv2_frame = cv2.cvtColor(frame,cv2.COLOR_BGR2RGB) 130 | 131 | while True: 132 | x_min, y_min, x_max, y_max = int(float(line[1])), int(float(line[2])), int(float(line[3])), int(float(line[4])) 133 | 134 | bbox_width = abs(x_max - x_min) 135 | bbox_height = abs(y_max - y_min) 136 | # x_min -= 3 * bbox_width / 4 137 | # x_max += 3 * bbox_width / 4 138 | # y_min -= 3 * bbox_height / 4 139 | # y_max += bbox_height / 4 140 | x_min -= 50 141 | x_max += 50 142 | y_min -= 50 143 | y_max += 30 144 | x_min = max(x_min, 0) 145 | y_min = max(y_min, 0) 146 | x_max = min(frame.shape[1], x_max) 147 | y_max = min(frame.shape[0], y_max) 148 | # Crop face loosely 149 | img = cv2_frame[y_min:y_max,x_min:x_max] 150 | img = Image.fromarray(img) 151 | 152 | # Transform 153 | img = transformations(img) 154 | img_shape = img.size() 155 | img = img.view(1, img_shape[0], img_shape[1], img_shape[2]) 156 | img = Variable(img).cuda(gpu) 157 | 158 | yaw, pitch, roll = model(img) 159 | 160 | yaw_predicted = F.softmax(yaw) 161 | pitch_predicted = F.softmax(pitch) 162 | roll_predicted = F.softmax(roll) 163 | # Get continuous predictions in degrees. 164 | yaw_predicted = torch.sum(yaw_predicted.data[0] * idx_tensor) * 3 - 99 165 | pitch_predicted = torch.sum(pitch_predicted.data[0] * idx_tensor) * 3 - 99 166 | roll_predicted = torch.sum(roll_predicted.data[0] * idx_tensor) * 3 - 99 167 | 168 | # Print new frame with cube and axis 169 | txt_out.write(str(frame_num) + ' %f %f %f\n' % (yaw_predicted, pitch_predicted, roll_predicted)) 170 | # utils.plot_pose_cube(frame, yaw_predicted, pitch_predicted, roll_predicted, (x_min + x_max) / 2, (y_min + y_max) / 2, size = bbox_width) 171 | utils.draw_axis(frame, yaw_predicted, pitch_predicted, roll_predicted, tdx = (x_min + x_max) / 2, tdy= (y_min + y_max) / 2, size = bbox_height/2) 172 | # Plot expanded bounding box 173 | # cv2.rectangle(frame, (x_min, y_min), (x_max, y_max), (0,255,0), 1) 174 | 175 | # Peek next frame detection 176 | next_frame_num = int(bbox_line_list[idx+1].strip('\n').split(' ')[0]) 177 | # print 'next_frame_num ', next_frame_num 178 | if next_frame_num == det_frame_num: 179 | idx += 1 180 | line = bbox_line_list[idx].strip('\n').split(' ') 181 | det_frame_num = int(line[0]) 182 | else: 183 | break 184 | 185 | idx += 1 186 | out.write(frame) 187 | frame_num += 1 188 | 189 | out.release() 190 | video.release() 191 | txt_out.close() 192 | -------------------------------------------------------------------------------- /code/test_on_video_dlib.py: -------------------------------------------------------------------------------- 1 | import sys, os, argparse 2 | 3 | import numpy as np 4 | import cv2 5 | import matplotlib.pyplot as plt 6 | 7 | import torch 8 | import torch.nn as nn 9 | from torch.autograd import Variable 10 | from torch.utils.data import DataLoader 11 | from torchvision import transforms 12 | import torch.backends.cudnn as cudnn 13 | import torchvision 14 | import torch.nn.functional as F 15 | from PIL import Image 16 | 17 | import datasets, hopenet, utils 18 | 19 | from skimage import io 20 | import dlib 21 | 22 | def parse_args(): 23 | """Parse input arguments.""" 24 | parser = argparse.ArgumentParser(description='Head pose estimation using the Hopenet network.') 25 | parser.add_argument('--gpu', dest='gpu_id', help='GPU device id to use [0]', 26 | default=0, type=int) 27 | parser.add_argument('--snapshot', dest='snapshot', help='Path of model snapshot.', 28 | default='', type=str) 29 | parser.add_argument('--face_model', dest='face_model', help='Path of DLIB face detection model.', 30 | default='', type=str) 31 | parser.add_argument('--video', dest='video_path', help='Path of video') 32 | parser.add_argument('--output_string', dest='output_string', help='String appended to output file') 33 | parser.add_argument('--n_frames', dest='n_frames', help='Number of frames', type=int) 34 | parser.add_argument('--fps', dest='fps', help='Frames per second of source video', type=float, default=30.) 35 | args = parser.parse_args() 36 | return args 37 | 38 | if __name__ == '__main__': 39 | args = parse_args() 40 | 41 | cudnn.enabled = True 42 | 43 | batch_size = 1 44 | gpu = args.gpu_id 45 | snapshot_path = args.snapshot 46 | out_dir = 'output/video' 47 | video_path = args.video_path 48 | 49 | if not os.path.exists(out_dir): 50 | os.makedirs(out_dir) 51 | 52 | if not os.path.exists(args.video_path): 53 | sys.exit('Video does not exist') 54 | 55 | # ResNet50 structure 56 | model = hopenet.Hopenet(torchvision.models.resnet.Bottleneck, [3, 4, 6, 3], 66) 57 | 58 | # Dlib face detection model 59 | cnn_face_detector = dlib.cnn_face_detection_model_v1(args.face_model) 60 | 61 | print 'Loading snapshot.' 62 | # Load snapshot 63 | saved_state_dict = torch.load(snapshot_path) 64 | model.load_state_dict(saved_state_dict) 65 | 66 | print 'Loading data.' 67 | 68 | transformations = transforms.Compose([transforms.Scale(224), 69 | transforms.CenterCrop(224), transforms.ToTensor(), 70 | transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])]) 71 | 72 | model.cuda(gpu) 73 | 74 | print 'Ready to test network.' 75 | 76 | # Test the Model 77 | model.eval() # Change model to 'eval' mode (BN uses moving mean/var). 78 | total = 0 79 | 80 | idx_tensor = [idx for idx in xrange(66)] 81 | idx_tensor = torch.FloatTensor(idx_tensor).cuda(gpu) 82 | 83 | video = cv2.VideoCapture(video_path) 84 | 85 | # New cv2 86 | width = int(video.get(cv2.CAP_PROP_FRAME_WIDTH)) # float 87 | height = int(video.get(cv2.CAP_PROP_FRAME_HEIGHT)) # float 88 | 89 | # Define the codec and create VideoWriter object 90 | fourcc = cv2.VideoWriter_fourcc(*'MJPG') 91 | out = cv2.VideoWriter('output/video/output-%s.avi' % args.output_string, fourcc, args.fps, (width, height)) 92 | 93 | # # Old cv2 94 | # width = int(video.get(cv2.cv.CV_CAP_PROP_FRAME_WIDTH)) # float 95 | # height = int(video.get(cv2.cv.CV_CAP_PROP_FRAME_HEIGHT)) # float 96 | # 97 | # # Define the codec and create VideoWriter object 98 | # fourcc = cv2.cv.CV_FOURCC(*'MJPG') 99 | # out = cv2.VideoWriter('output/video/output-%s.avi' % args.output_string, fourcc, 30.0, (width, height)) 100 | 101 | txt_out = open('output/video/output-%s.txt' % args.output_string, 'w') 102 | 103 | frame_num = 1 104 | 105 | while frame_num <= args.n_frames: 106 | print frame_num 107 | 108 | ret,frame = video.read() 109 | if ret == False: 110 | break 111 | 112 | cv2_frame = cv2.cvtColor(frame,cv2.COLOR_BGR2RGB) 113 | 114 | # Dlib detect 115 | dets = cnn_face_detector(cv2_frame, 1) 116 | 117 | for idx, det in enumerate(dets): 118 | # Get x_min, y_min, x_max, y_max, conf 119 | x_min = det.rect.left() 120 | y_min = det.rect.top() 121 | x_max = det.rect.right() 122 | y_max = det.rect.bottom() 123 | conf = det.confidence 124 | 125 | if conf > 1.0: 126 | bbox_width = abs(x_max - x_min) 127 | bbox_height = abs(y_max - y_min) 128 | x_min -= 2 * bbox_width / 4 129 | x_max += 2 * bbox_width / 4 130 | y_min -= 3 * bbox_height / 4 131 | y_max += bbox_height / 4 132 | x_min = max(x_min, 0); y_min = max(y_min, 0) 133 | x_max = min(frame.shape[1], x_max); y_max = min(frame.shape[0], y_max) 134 | # Crop image 135 | img = cv2_frame[y_min:y_max,x_min:x_max] 136 | img = Image.fromarray(img) 137 | 138 | # Transform 139 | img = transformations(img) 140 | img_shape = img.size() 141 | img = img.view(1, img_shape[0], img_shape[1], img_shape[2]) 142 | img = Variable(img).cuda(gpu) 143 | 144 | yaw, pitch, roll = model(img) 145 | 146 | yaw_predicted = F.softmax(yaw) 147 | pitch_predicted = F.softmax(pitch) 148 | roll_predicted = F.softmax(roll) 149 | # Get continuous predictions in degrees. 150 | yaw_predicted = torch.sum(yaw_predicted.data[0] * idx_tensor) * 3 - 99 151 | pitch_predicted = torch.sum(pitch_predicted.data[0] * idx_tensor) * 3 - 99 152 | roll_predicted = torch.sum(roll_predicted.data[0] * idx_tensor) * 3 - 99 153 | 154 | # Print new frame with cube and axis 155 | txt_out.write(str(frame_num) + ' %f %f %f\n' % (yaw_predicted, pitch_predicted, roll_predicted)) 156 | # utils.plot_pose_cube(frame, yaw_predicted, pitch_predicted, roll_predicted, (x_min + x_max) / 2, (y_min + y_max) / 2, size = bbox_width) 157 | utils.draw_axis(frame, yaw_predicted, pitch_predicted, roll_predicted, tdx = (x_min + x_max) / 2, tdy= (y_min + y_max) / 2, size = bbox_height/2) 158 | # Plot expanded bounding box 159 | # cv2.rectangle(frame, (x_min, y_min), (x_max, y_max), (0,255,0), 1) 160 | 161 | out.write(frame) 162 | frame_num += 1 163 | 164 | out.release() 165 | video.release() 166 | -------------------------------------------------------------------------------- /code/test_on_video_dockerface.py: -------------------------------------------------------------------------------- 1 | import sys, os, argparse 2 | 3 | import numpy as np 4 | import cv2 5 | import matplotlib.pyplot as plt 6 | 7 | import torch 8 | import torch.nn as nn 9 | from torch.autograd import Variable 10 | from torch.utils.data import DataLoader 11 | from torchvision import transforms 12 | import torch.backends.cudnn as cudnn 13 | import torchvision 14 | import torch.nn.functional as F 15 | from PIL import Image 16 | 17 | import datasets, hopenet, utils 18 | 19 | def parse_args(): 20 | """Parse input arguments.""" 21 | parser = argparse.ArgumentParser(description='Head pose estimation using the Hopenet network.') 22 | parser.add_argument('--gpu', dest='gpu_id', help='GPU device id to use [0]', 23 | default=0, type=int) 24 | parser.add_argument('--snapshot', dest='snapshot', help='Path of model snapshot.', 25 | default='', type=str) 26 | parser.add_argument('--video', dest='video_path', help='Path of video') 27 | parser.add_argument('--bboxes', dest='bboxes', help='Bounding box annotations of frames') 28 | parser.add_argument('--output_string', dest='output_string', help='String appended to output file') 29 | parser.add_argument('--n_frames', dest='n_frames', help='Number of frames', type=int) 30 | parser.add_argument('--fps', dest='fps', help='Frames per second of source video', type=float, default=30.) 31 | args = parser.parse_args() 32 | return args 33 | 34 | if __name__ == '__main__': 35 | args = parse_args() 36 | 37 | cudnn.enabled = True 38 | 39 | batch_size = 1 40 | gpu = args.gpu_id 41 | snapshot_path = args.snapshot 42 | out_dir = 'output/video' 43 | video_path = args.video_path 44 | 45 | if not os.path.exists(out_dir): 46 | os.makedirs(out_dir) 47 | 48 | if not os.path.exists(args.video_path): 49 | sys.exit('Video does not exist') 50 | 51 | # ResNet50 structure 52 | model = hopenet.Hopenet(torchvision.models.resnet.Bottleneck, [3, 4, 6, 3], 66) 53 | 54 | print 'Loading snapshot.' 55 | # Load snapshot 56 | saved_state_dict = torch.load(snapshot_path) 57 | model.load_state_dict(saved_state_dict) 58 | 59 | print 'Loading data.' 60 | 61 | transformations = transforms.Compose([transforms.Scale(224), 62 | transforms.CenterCrop(224), transforms.ToTensor(), 63 | transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])]) 64 | 65 | model.cuda(gpu) 66 | 67 | print 'Ready to test network.' 68 | 69 | # Test the Model 70 | model.eval() # Change model to 'eval' mode (BN uses moving mean/var). 71 | total = 0 72 | 73 | idx_tensor = [idx for idx in xrange(66)] 74 | idx_tensor = torch.FloatTensor(idx_tensor).cuda(gpu) 75 | 76 | video = cv2.VideoCapture(video_path) 77 | 78 | # New cv2 79 | width = int(video.get(cv2.CAP_PROP_FRAME_WIDTH)) # float 80 | height = int(video.get(cv2.CAP_PROP_FRAME_HEIGHT)) # float 81 | 82 | # Define the codec and create VideoWriter object 83 | fourcc = cv2.VideoWriter_fourcc(*'MJPG') 84 | out = cv2.VideoWriter('output/video/output-%s.avi' % args.output_string, fourcc, args.fps, (width, height)) 85 | 86 | # # Old cv2 87 | # width = int(video.get(cv2.cv.CV_CAP_PROP_FRAME_WIDTH)) # float 88 | # height = int(video.get(cv2.cv.CV_CAP_PROP_FRAME_HEIGHT)) # float 89 | # 90 | # # Define the codec and create VideoWriter object 91 | # fourcc = cv2.cv.CV_FOURCC(*'MJPG') 92 | # out = cv2.VideoWriter('output/video/output-%s.avi' % args.output_string, fourcc, 30.0, (width, height)) 93 | 94 | txt_out = open('output/video/output-%s.txt' % args.output_string, 'w') 95 | 96 | frame_num = 1 97 | 98 | with open(args.bboxes, 'r') as f: 99 | bbox_line_list = f.read().splitlines() 100 | 101 | idx = 0 102 | while idx < len(bbox_line_list): 103 | line = bbox_line_list[idx] 104 | line = line.strip('\n') 105 | line = line.split(' ') 106 | det_frame_num = int(line[0]) 107 | 108 | print frame_num 109 | 110 | # Stop at a certain frame number 111 | if frame_num > args.n_frames: 112 | break 113 | 114 | # Save all frames as they are if they don't have bbox annotation. 115 | while frame_num < det_frame_num: 116 | ret, frame = video.read() 117 | if ret == False: 118 | out.release() 119 | video.release() 120 | txt_out.close() 121 | sys.exit(0) 122 | out.write(frame) 123 | frame_num += 1 124 | 125 | # Start processing frame with bounding box 126 | ret,frame = video.read() 127 | if ret == False: 128 | break 129 | cv2_frame = cv2.cvtColor(frame,cv2.COLOR_BGR2RGB) 130 | 131 | while True: 132 | x_min, y_min, x_max, y_max, conf = int(float(line[1])), int(float(line[2])), int(float(line[3])), int(float(line[4])), float(line[5]) 133 | 134 | if conf > 0.98: 135 | bbox_width = abs(x_max - x_min) 136 | bbox_height = abs(y_max - y_min) 137 | # x_min -= 3 * bbox_width / 4 138 | # x_max += 3 * bbox_width / 4 139 | # y_min -= 3 * bbox_height / 4 140 | # y_max += bbox_height / 4 141 | x_min -= 50 142 | x_max += 50 143 | y_min -= 50 144 | y_max += 30 145 | x_min = max(x_min, 0) 146 | y_min = max(y_min, 0) 147 | x_max = min(frame.shape[1], x_max) 148 | y_max = min(frame.shape[0], y_max) 149 | # Crop image 150 | img = cv2_frame[y_min:y_max,x_min:x_max] 151 | img = Image.fromarray(img) 152 | 153 | # Transform 154 | img = transformations(img) 155 | img_shape = img.size() 156 | img = img.view(1, img_shape[0], img_shape[1], img_shape[2]) 157 | img = Variable(img).cuda(gpu) 158 | 159 | yaw, pitch, roll = model(img) 160 | 161 | yaw_predicted = F.softmax(yaw) 162 | pitch_predicted = F.softmax(pitch) 163 | roll_predicted = F.softmax(roll) 164 | # Get continuous predictions in degrees. 165 | yaw_predicted = torch.sum(yaw_predicted.data[0] * idx_tensor) * 3 - 99 166 | pitch_predicted = torch.sum(pitch_predicted.data[0] * idx_tensor) * 3 - 99 167 | roll_predicted = torch.sum(roll_predicted.data[0] * idx_tensor) * 3 - 99 168 | 169 | # Print new frame with cube and axis 170 | txt_out.write(str(frame_num) + ' %f %f %f\n' % (yaw_predicted, pitch_predicted, roll_predicted)) 171 | # utils.plot_pose_cube(frame, yaw_predicted, pitch_predicted, roll_predicted, (x_min + x_max) / 2, (y_min + y_max) / 2, size = bbox_width) 172 | utils.draw_axis(frame, yaw_predicted, pitch_predicted, roll_predicted, tdx = (x_min + x_max) / 2, tdy= (y_min + y_max) / 2, size = bbox_height/2) 173 | # Plot expanded bounding box 174 | # cv2.rectangle(frame, (x_min, y_min), (x_max, y_max), (0,255,0), 1) 175 | 176 | # Peek next frame detection 177 | next_frame_num = int(bbox_line_list[idx+1].strip('\n').split(' ')[0]) 178 | # print 'next_frame_num ', next_frame_num 179 | if next_frame_num == det_frame_num: 180 | idx += 1 181 | line = bbox_line_list[idx].strip('\n').split(' ') 182 | det_frame_num = int(line[0]) 183 | else: 184 | break 185 | 186 | idx += 1 187 | out.write(frame) 188 | frame_num += 1 189 | 190 | out.release() 191 | video.release() 192 | txt_out.close() 193 | -------------------------------------------------------------------------------- /code/test_resnet50_regression.py: -------------------------------------------------------------------------------- 1 | import sys, os, argparse 2 | 3 | import numpy as np 4 | import cv2 5 | import matplotlib.pyplot as plt 6 | 7 | import torch 8 | import torch.nn as nn 9 | from torch.autograd import Variable 10 | from torch.utils.data import DataLoader 11 | from torchvision import transforms 12 | import torch.backends.cudnn as cudnn 13 | import torchvision 14 | import torch.nn.functional as F 15 | 16 | import datasets, hopenet, utils 17 | 18 | def parse_args(): 19 | """Parse input arguments.""" 20 | parser = argparse.ArgumentParser(description='Head pose estimation using the Hopenet network.') 21 | parser.add_argument('--gpu', dest='gpu_id', help='GPU device id to use [0]', 22 | default=0, type=int) 23 | parser.add_argument('--data_dir', dest='data_dir', help='Directory path for data.', 24 | default='', type=str) 25 | parser.add_argument('--filename_list', dest='filename_list', help='Path to text file containing relative paths for every example.', 26 | default='', type=str) 27 | parser.add_argument('--snapshot', dest='snapshot', help='Name of model snapshot.', 28 | default='', type=str) 29 | parser.add_argument('--batch_size', dest='batch_size', help='Batch size.', 30 | default=1, type=int) 31 | parser.add_argument('--save_viz', dest='save_viz', help='Save images with pose cube.', 32 | default=False, type=bool) 33 | parser.add_argument('--dataset', dest='dataset', help='Dataset type.', default='AFLW2000', type=str) 34 | 35 | args = parser.parse_args() 36 | 37 | return args 38 | 39 | if __name__ == '__main__': 40 | args = parse_args() 41 | 42 | cudnn.enabled = True 43 | gpu = args.gpu_id 44 | snapshot_path = args.snapshot 45 | 46 | model = hopenet.ResNet(torchvision.models.resnet.Bottleneck, [3, 4, 6, 3], 3) 47 | 48 | print 'Loading snapshot.' 49 | # Load snapshot 50 | saved_state_dict = torch.load(snapshot_path) 51 | model.load_state_dict(saved_state_dict) 52 | 53 | print 'Loading data.' 54 | 55 | transformations = transforms.Compose([transforms.Scale(224), 56 | transforms.CenterCrop(224), transforms.ToTensor(), 57 | transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])]) 58 | 59 | if args.dataset == 'Pose_300W_LP': 60 | pose_dataset = datasets.Pose_300W_LP(args.data_dir, args.filename_list, transformations) 61 | elif args.dataset == 'Pose_300W_LP_random_ds': 62 | pose_dataset = datasets.Pose_300W_LP_random_ds(args.data_dir, args.filename_list, transformations) 63 | elif args.dataset == 'AFLW2000': 64 | pose_dataset = datasets.AFLW2000(args.data_dir, args.filename_list, transformations) 65 | elif args.dataset == 'AFLW2000_ds': 66 | pose_dataset = datasets.AFLW2000_ds(args.data_dir, args.filename_list, transformations) 67 | elif args.dataset == 'BIWI': 68 | pose_dataset = datasets.BIWI(args.data_dir, args.filename_list, transformations) 69 | elif args.dataset == 'AFLW': 70 | pose_dataset = datasets.AFLW(args.data_dir, args.filename_list, transformations) 71 | elif args.dataset == 'AFLW_aug': 72 | pose_dataset = datasets.AFLW_aug(args.data_dir, args.filename_list, transformations) 73 | elif args.dataset == 'AFW': 74 | pose_dataset = datasets.AFW(args.data_dir, args.filename_list, transformations) 75 | else: 76 | print 'Error: not a valid dataset name' 77 | sys.exit() 78 | test_loader = torch.utils.data.DataLoader(dataset=pose_dataset, 79 | batch_size=args.batch_size, 80 | num_workers=2) 81 | 82 | model.cuda(gpu) 83 | 84 | print 'Ready to test network.' 85 | 86 | # Test the Model 87 | model.eval() # Change model to 'eval' mode (BN uses moving mean/var). 88 | total = 0 89 | 90 | yaw_error = .0 91 | pitch_error = .0 92 | roll_error = .0 93 | 94 | l1loss = torch.nn.L1Loss(size_average=False) 95 | 96 | for i, (images, labels, cont_labels, name) in enumerate(test_loader): 97 | images = Variable(images).cuda(gpu) 98 | total += cont_labels.size(0) 99 | label_yaw = cont_labels[:,0].float() 100 | label_pitch = cont_labels[:,1].float() 101 | label_roll = cont_labels[:,2].float() 102 | 103 | angles = model(images) 104 | yaw_predicted = angles[:,0].data.cpu() 105 | pitch_predicted = angles[:,1].data.cpu() 106 | roll_predicted = angles[:,2].data.cpu() 107 | 108 | # Mean absolute error 109 | yaw_error += torch.sum(torch.abs(yaw_predicted - label_yaw)) 110 | pitch_error += torch.sum(torch.abs(pitch_predicted - label_pitch)) 111 | roll_error += torch.sum(torch.abs(roll_predicted - label_roll)) 112 | 113 | # Save first image in batch with pose cube or axis. 114 | if args.save_viz: 115 | name = name[0] 116 | if args.dataset == 'BIWI': 117 | cv2_img = cv2.imread(os.path.join(args.data_dir, name + '_rgb.png')) 118 | else: 119 | cv2_img = cv2.imread(os.path.join(args.data_dir, name + '.jpg')) 120 | if args.batch_size == 1: 121 | error_string = 'y %.2f, p %.2f, r %.2f' % (torch.sum(torch.abs(yaw_predicted - label_yaw)), torch.sum(torch.abs(pitch_predicted - label_pitch)), torch.sum(torch.abs(roll_predicted - label_roll))) 122 | cv2.putText(cv2_img, error_string, (30, cv2_img.shape[0]- 30), fontFace=1, fontScale=1, color=(0,0,255), thickness=1) 123 | # utils.plot_pose_cube(cv2_img, yaw_predicted[0], pitch_predicted[0], roll_predicted[0], size=100) 124 | utils.draw_axis(cv2_img, yaw_predicted[0], pitch_predicted[0], roll_predicted[0], tdx = 200, tdy= 200, size=100) 125 | cv2.imwrite(os.path.join('output/images', name + '.jpg'), cv2_img) 126 | 127 | print('Test error in degrees of the model on the ' + str(total) + 128 | ' test images. Yaw: %.4f, Pitch: %.4f, Roll: %.4f' % (yaw_error / total, 129 | pitch_error / total, roll_error / total)) 130 | -------------------------------------------------------------------------------- /code/train_alexnet.py: -------------------------------------------------------------------------------- 1 | import sys, os, argparse, time 2 | 3 | import numpy as np 4 | import cv2 5 | import matplotlib.pyplot as plt 6 | 7 | import torch 8 | import torch.nn as nn 9 | from torch.autograd import Variable 10 | from torch.utils.data import DataLoader 11 | from torchvision import transforms 12 | import torchvision 13 | import torch.backends.cudnn as cudnn 14 | import torch.nn.functional as F 15 | 16 | import datasets, hopenet 17 | import torch.utils.model_zoo as model_zoo 18 | 19 | model_urls = { 20 | 'alexnet': 'https://download.pytorch.org/models/alexnet-owt-4df8aa71.pth', 21 | } 22 | 23 | def parse_args(): 24 | """Parse input arguments.""" 25 | parser = argparse.ArgumentParser(description='Head pose estimation using the Hopenet network.') 26 | parser.add_argument('--gpu', dest='gpu_id', help='GPU device id to use [0]', 27 | default=0, type=int) 28 | parser.add_argument('--num_epochs', dest='num_epochs', help='Maximum number of training epochs.', 29 | default=5, type=int) 30 | parser.add_argument('--batch_size', dest='batch_size', help='Batch size.', 31 | default=16, type=int) 32 | parser.add_argument('--lr', dest='lr', help='Base learning rate.', 33 | default=0.001, type=float) 34 | parser.add_argument('--data_dir', dest='data_dir', help='Directory path for data.', 35 | default='', type=str) 36 | parser.add_argument('--filename_list', dest='filename_list', help='Path to text file containing relative paths for every example.', 37 | default='', type=str) 38 | parser.add_argument('--output_string', dest='output_string', help='String appended to output snapshots.', default = '', type=str) 39 | parser.add_argument('--alpha', dest='alpha', help='Regression loss coefficient.', 40 | default=0.001, type=float) 41 | parser.add_argument('--dataset', dest='dataset', help='Dataset type.', default='Pose_300W_LP', type=str) 42 | args = parser.parse_args() 43 | return args 44 | 45 | def get_ignored_params(model): 46 | # Generator function that yields ignored params. 47 | b = [model.features[0], model.features[1], model.features[2]] 48 | for i in range(len(b)): 49 | for module_name, module in b[i].named_modules(): 50 | if 'bn' in module_name: 51 | module.eval() 52 | for name, param in module.named_parameters(): 53 | yield param 54 | 55 | def get_non_ignored_params(model): 56 | # Generator function that yields params that will be optimized. 57 | b = [] 58 | for idx in xrange(3, len(model.features)): 59 | b.append(model.features[idx]) 60 | for layer in model.classifier: 61 | b.append(layer) 62 | for i in range(len(b)): 63 | for module_name, module in b[i].named_modules(): 64 | if 'bn' in module_name: 65 | module.eval() 66 | for name, param in module.named_parameters(): 67 | yield param 68 | 69 | def get_fc_params(model): 70 | b = [model.fc_yaw, model.fc_pitch, model.fc_roll] 71 | for i in range(len(b)): 72 | for module_name, module in b[i].named_modules(): 73 | for name, param in module.named_parameters(): 74 | yield param 75 | 76 | def load_filtered_state_dict(model, snapshot): 77 | # By user apaszke from discuss.pytorch.org 78 | model_dict = model.state_dict() 79 | snapshot = {k: v for k, v in snapshot.items() if k in model_dict} 80 | model_dict.update(snapshot) 81 | model.load_state_dict(model_dict) 82 | 83 | if __name__ == '__main__': 84 | args = parse_args() 85 | 86 | cudnn.enabled = True 87 | num_epochs = args.num_epochs 88 | batch_size = args.batch_size 89 | gpu = args.gpu_id 90 | 91 | if not os.path.exists('output/snapshots'): 92 | os.makedirs('output/snapshots') 93 | 94 | model = hopenet.AlexNet(66) 95 | load_filtered_state_dict(model, model_zoo.load_url(model_urls['alexnet'])) 96 | 97 | print 'Loading data.' 98 | 99 | transformations = transforms.Compose([transforms.Scale(240), 100 | transforms.RandomCrop(224), transforms.ToTensor(), 101 | transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])]) 102 | 103 | if args.dataset == 'Pose_300W_LP': 104 | pose_dataset = datasets.Pose_300W_LP(args.data_dir, args.filename_list, transformations) 105 | elif args.dataset == 'Pose_300W_LP_random_ds': 106 | pose_dataset = datasets.Pose_300W_LP_random_ds(args.data_dir, args.filename_list, transformations) 107 | elif args.dataset == 'AFLW2000': 108 | pose_dataset = datasets.AFLW2000(args.data_dir, args.filename_list, transformations) 109 | elif args.dataset == 'BIWI': 110 | pose_dataset = datasets.BIWI(args.data_dir, args.filename_list, transformations) 111 | elif args.dataset == 'AFLW': 112 | pose_dataset = datasets.AFLW(args.data_dir, args.filename_list, transformations) 113 | elif args.dataset == 'AFLW_aug': 114 | pose_dataset = datasets.AFLW_aug(args.data_dir, args.filename_list, transformations) 115 | elif args.dataset == 'AFW': 116 | pose_dataset = datasets.AFW(args.data_dir, args.filename_list, transformations) 117 | else: 118 | print 'Error: not a valid dataset name' 119 | sys.exit() 120 | train_loader = torch.utils.data.DataLoader(dataset=pose_dataset, 121 | batch_size=batch_size, 122 | shuffle=True, 123 | num_workers=2) 124 | 125 | model.cuda(gpu) 126 | softmax = nn.Softmax().cuda(gpu) 127 | criterion = nn.CrossEntropyLoss().cuda(gpu) 128 | reg_criterion = nn.MSELoss().cuda(gpu) 129 | # Regression loss coefficient 130 | alpha = args.alpha 131 | 132 | idx_tensor = [idx for idx in xrange(66)] 133 | idx_tensor = Variable(torch.FloatTensor(idx_tensor)).cuda(gpu) 134 | 135 | optimizer = torch.optim.Adam([{'params': get_ignored_params(model), 'lr': 0}, 136 | {'params': get_non_ignored_params(model), 'lr': args.lr}, 137 | {'params': get_fc_params(model), 'lr': args.lr * 5}], 138 | lr = args.lr) 139 | 140 | print 'Ready to train network.' 141 | for epoch in range(num_epochs): 142 | for i, (images, labels, cont_labels, name) in enumerate(train_loader): 143 | images = Variable(images).cuda(gpu) 144 | 145 | # Binned labels 146 | label_yaw = Variable(labels[:,0]).cuda(gpu) 147 | label_pitch = Variable(labels[:,1]).cuda(gpu) 148 | label_roll = Variable(labels[:,2]).cuda(gpu) 149 | 150 | # Continuous labels 151 | label_yaw_cont = Variable(cont_labels[:,0]).cuda(gpu) 152 | label_pitch_cont = Variable(cont_labels[:,1]).cuda(gpu) 153 | label_roll_cont = Variable(cont_labels[:,2]).cuda(gpu) 154 | 155 | # Forward pass 156 | pre_yaw, pre_pitch, pre_roll = model(images) 157 | 158 | # Cross entropy loss 159 | loss_yaw = criterion(pre_yaw, label_yaw) 160 | loss_pitch = criterion(pre_pitch, label_pitch) 161 | loss_roll = criterion(pre_roll, label_roll) 162 | 163 | # MSE loss 164 | yaw_predicted = softmax(pre_yaw) 165 | pitch_predicted = softmax(pre_pitch) 166 | roll_predicted = softmax(pre_roll) 167 | 168 | yaw_predicted = torch.sum(yaw_predicted * idx_tensor, 1) * 3 - 99 169 | pitch_predicted = torch.sum(pitch_predicted * idx_tensor, 1) * 3 - 99 170 | roll_predicted = torch.sum(roll_predicted * idx_tensor, 1) * 3 - 99 171 | 172 | loss_reg_yaw = reg_criterion(yaw_predicted, label_yaw_cont) 173 | loss_reg_pitch = reg_criterion(pitch_predicted, label_pitch_cont) 174 | loss_reg_roll = reg_criterion(roll_predicted, label_roll_cont) 175 | 176 | # Total loss 177 | loss_yaw += alpha * loss_reg_yaw 178 | loss_pitch += alpha * loss_reg_pitch 179 | loss_roll += alpha * loss_reg_roll 180 | 181 | loss_seq = [loss_yaw, loss_pitch, loss_roll] 182 | grad_seq = [torch.ones(1).cuda(gpu) for _ in range(len(loss_seq))] 183 | torch.autograd.backward(loss_seq, grad_seq) 184 | optimizer.step() 185 | 186 | if (i+1) % 100 == 0: 187 | print ('Epoch [%d/%d], Iter [%d/%d] Losses: Yaw %.4f, Pitch %.4f, Roll %.4f' 188 | %(epoch+1, num_epochs, i+1, len(pose_dataset)//batch_size, loss_yaw.data[0], loss_pitch.data[0], loss_roll.data[0])) 189 | 190 | # Save models at numbered epochs. 191 | if epoch % 1 == 0 and epoch < num_epochs: 192 | print 'Taking snapshot...' 193 | torch.save(model.state_dict(), 194 | 'output/snapshots/' + args.output_string + '_epoch_'+ str(epoch+1) + '.pkl') 195 | -------------------------------------------------------------------------------- /code/train_hopenet.py: -------------------------------------------------------------------------------- 1 | import sys, os, argparse, time 2 | 3 | import numpy as np 4 | import cv2 5 | import matplotlib.pyplot as plt 6 | 7 | import torch 8 | import torch.nn as nn 9 | from torch.autograd import Variable 10 | from torch.utils.data import DataLoader 11 | from torchvision import transforms 12 | import torchvision 13 | import torch.backends.cudnn as cudnn 14 | import torch.nn.functional as F 15 | 16 | import datasets, hopenet 17 | import torch.utils.model_zoo as model_zoo 18 | 19 | def parse_args(): 20 | """Parse input arguments.""" 21 | parser = argparse.ArgumentParser(description='Head pose estimation using the Hopenet network.') 22 | parser.add_argument('--gpu', dest='gpu_id', help='GPU device id to use [0]', 23 | default=0, type=int) 24 | parser.add_argument('--num_epochs', dest='num_epochs', help='Maximum number of training epochs.', 25 | default=5, type=int) 26 | parser.add_argument('--batch_size', dest='batch_size', help='Batch size.', 27 | default=16, type=int) 28 | parser.add_argument('--lr', dest='lr', help='Base learning rate.', 29 | default=0.001, type=float) 30 | parser.add_argument('--dataset', dest='dataset', help='Dataset type.', default='Pose_300W_LP', type=str) 31 | parser.add_argument('--data_dir', dest='data_dir', help='Directory path for data.', 32 | default='', type=str) 33 | parser.add_argument('--filename_list', dest='filename_list', help='Path to text file containing relative paths for every example.', 34 | default='', type=str) 35 | parser.add_argument('--output_string', dest='output_string', help='String appended to output snapshots.', default = '', type=str) 36 | parser.add_argument('--alpha', dest='alpha', help='Regression loss coefficient.', 37 | default=0.001, type=float) 38 | parser.add_argument('--snapshot', dest='snapshot', help='Path of model snapshot.', 39 | default='', type=str) 40 | 41 | args = parser.parse_args() 42 | return args 43 | 44 | def get_ignored_params(model): 45 | # Generator function that yields ignored params. 46 | b = [model.conv1, model.bn1, model.fc_finetune] 47 | for i in range(len(b)): 48 | for module_name, module in b[i].named_modules(): 49 | if 'bn' in module_name: 50 | module.eval() 51 | for name, param in module.named_parameters(): 52 | yield param 53 | 54 | def get_non_ignored_params(model): 55 | # Generator function that yields params that will be optimized. 56 | b = [model.layer1, model.layer2, model.layer3, model.layer4] 57 | for i in range(len(b)): 58 | for module_name, module in b[i].named_modules(): 59 | if 'bn' in module_name: 60 | module.eval() 61 | for name, param in module.named_parameters(): 62 | yield param 63 | 64 | def get_fc_params(model): 65 | # Generator function that yields fc layer params. 66 | b = [model.fc_yaw, model.fc_pitch, model.fc_roll] 67 | for i in range(len(b)): 68 | for module_name, module in b[i].named_modules(): 69 | for name, param in module.named_parameters(): 70 | yield param 71 | 72 | def load_filtered_state_dict(model, snapshot): 73 | # By user apaszke from discuss.pytorch.org 74 | model_dict = model.state_dict() 75 | snapshot = {k: v for k, v in snapshot.items() if k in model_dict} 76 | model_dict.update(snapshot) 77 | model.load_state_dict(model_dict) 78 | 79 | if __name__ == '__main__': 80 | args = parse_args() 81 | 82 | cudnn.enabled = True 83 | num_epochs = args.num_epochs 84 | batch_size = args.batch_size 85 | gpu = args.gpu_id 86 | 87 | if not os.path.exists('output/snapshots'): 88 | os.makedirs('output/snapshots') 89 | 90 | # ResNet50 structure 91 | model = hopenet.Hopenet(torchvision.models.resnet.Bottleneck, [3, 4, 6, 3], 66) 92 | 93 | if args.snapshot == '': 94 | load_filtered_state_dict(model, model_zoo.load_url('https://download.pytorch.org/models/resnet50-19c8e357.pth')) 95 | else: 96 | saved_state_dict = torch.load(args.snapshot) 97 | model.load_state_dict(saved_state_dict) 98 | 99 | print 'Loading data.' 100 | 101 | transformations = transforms.Compose([transforms.Scale(240), 102 | transforms.RandomCrop(224), transforms.ToTensor(), 103 | transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])]) 104 | 105 | if args.dataset == 'Pose_300W_LP': 106 | pose_dataset = datasets.Pose_300W_LP(args.data_dir, args.filename_list, transformations) 107 | elif args.dataset == 'Pose_300W_LP_random_ds': 108 | pose_dataset = datasets.Pose_300W_LP_random_ds(args.data_dir, args.filename_list, transformations) 109 | elif args.dataset == 'Synhead': 110 | pose_dataset = datasets.Synhead(args.data_dir, args.filename_list, transformations) 111 | elif args.dataset == 'AFLW2000': 112 | pose_dataset = datasets.AFLW2000(args.data_dir, args.filename_list, transformations) 113 | elif args.dataset == 'BIWI': 114 | pose_dataset = datasets.BIWI(args.data_dir, args.filename_list, transformations) 115 | elif args.dataset == 'AFLW': 116 | pose_dataset = datasets.AFLW(args.data_dir, args.filename_list, transformations) 117 | elif args.dataset == 'AFLW_aug': 118 | pose_dataset = datasets.AFLW_aug(args.data_dir, args.filename_list, transformations) 119 | elif args.dataset == 'AFW': 120 | pose_dataset = datasets.AFW(args.data_dir, args.filename_list, transformations) 121 | else: 122 | print 'Error: not a valid dataset name' 123 | sys.exit() 124 | 125 | train_loader = torch.utils.data.DataLoader(dataset=pose_dataset, 126 | batch_size=batch_size, 127 | shuffle=True, 128 | num_workers=2) 129 | 130 | model.cuda(gpu) 131 | criterion = nn.CrossEntropyLoss().cuda(gpu) 132 | reg_criterion = nn.MSELoss().cuda(gpu) 133 | # Regression loss coefficient 134 | alpha = args.alpha 135 | 136 | softmax = nn.Softmax().cuda(gpu) 137 | idx_tensor = [idx for idx in xrange(66)] 138 | idx_tensor = Variable(torch.FloatTensor(idx_tensor)).cuda(gpu) 139 | 140 | optimizer = torch.optim.Adam([{'params': get_ignored_params(model), 'lr': 0}, 141 | {'params': get_non_ignored_params(model), 'lr': args.lr}, 142 | {'params': get_fc_params(model), 'lr': args.lr * 5}], 143 | lr = args.lr) 144 | 145 | print 'Ready to train network.' 146 | for epoch in range(num_epochs): 147 | for i, (images, labels, cont_labels, name) in enumerate(train_loader): 148 | images = Variable(images).cuda(gpu) 149 | 150 | # Binned labels 151 | label_yaw = Variable(labels[:,0]).cuda(gpu) 152 | label_pitch = Variable(labels[:,1]).cuda(gpu) 153 | label_roll = Variable(labels[:,2]).cuda(gpu) 154 | 155 | # Continuous labels 156 | label_yaw_cont = Variable(cont_labels[:,0]).cuda(gpu) 157 | label_pitch_cont = Variable(cont_labels[:,1]).cuda(gpu) 158 | label_roll_cont = Variable(cont_labels[:,2]).cuda(gpu) 159 | 160 | # Forward pass 161 | yaw, pitch, roll = model(images) 162 | 163 | # Cross entropy loss 164 | loss_yaw = criterion(yaw, label_yaw) 165 | loss_pitch = criterion(pitch, label_pitch) 166 | loss_roll = criterion(roll, label_roll) 167 | 168 | # MSE loss 169 | yaw_predicted = softmax(yaw) 170 | pitch_predicted = softmax(pitch) 171 | roll_predicted = softmax(roll) 172 | 173 | yaw_predicted = torch.sum(yaw_predicted * idx_tensor, 1) * 3 - 99 174 | pitch_predicted = torch.sum(pitch_predicted * idx_tensor, 1) * 3 - 99 175 | roll_predicted = torch.sum(roll_predicted * idx_tensor, 1) * 3 - 99 176 | 177 | loss_reg_yaw = reg_criterion(yaw_predicted, label_yaw_cont) 178 | loss_reg_pitch = reg_criterion(pitch_predicted, label_pitch_cont) 179 | loss_reg_roll = reg_criterion(roll_predicted, label_roll_cont) 180 | 181 | # Total loss 182 | loss_yaw += alpha * loss_reg_yaw 183 | loss_pitch += alpha * loss_reg_pitch 184 | loss_roll += alpha * loss_reg_roll 185 | 186 | loss_seq = [loss_yaw, loss_pitch, loss_roll] 187 | grad_seq = [torch.ones(1).cuda(gpu) for _ in range(len(loss_seq))] 188 | optimizer.zero_grad() 189 | torch.autograd.backward(loss_seq, grad_seq) 190 | optimizer.step() 191 | 192 | if (i+1) % 100 == 0: 193 | print ('Epoch [%d/%d], Iter [%d/%d] Losses: Yaw %.4f, Pitch %.4f, Roll %.4f' 194 | %(epoch+1, num_epochs, i+1, len(pose_dataset)//batch_size, loss_yaw.data[0], loss_pitch.data[0], loss_roll.data[0])) 195 | 196 | # Save models at numbered epochs. 197 | if epoch % 1 == 0 and epoch < num_epochs: 198 | print 'Taking snapshot...' 199 | torch.save(model.state_dict(), 200 | 'output/snapshots/' + args.output_string + '_epoch_'+ str(epoch+1) + '.pkl') 201 | -------------------------------------------------------------------------------- /code/train_resnet50_regression.py: -------------------------------------------------------------------------------- 1 | import sys, os, argparse, time 2 | 3 | import numpy as np 4 | import cv2 5 | import matplotlib.pyplot as plt 6 | 7 | import torch 8 | import torch.nn as nn 9 | from torch.autograd import Variable 10 | from torch.utils.data import DataLoader 11 | from torchvision import transforms 12 | import torchvision 13 | import torch.backends.cudnn as cudnn 14 | import torch.nn.functional as F 15 | 16 | import datasets, hopenet 17 | import torch.utils.model_zoo as model_zoo 18 | 19 | def parse_args(): 20 | """Parse input arguments.""" 21 | parser = argparse.ArgumentParser(description='Head pose estimation using the Hopenet network.') 22 | parser.add_argument('--gpu', dest='gpu_id', help='GPU device id to use [0]', 23 | default=0, type=int) 24 | parser.add_argument('--num_epochs', dest='num_epochs', help='Maximum number of training epochs.', 25 | default=5, type=int) 26 | parser.add_argument('--batch_size', dest='batch_size', help='Batch size.', 27 | default=16, type=int) 28 | parser.add_argument('--lr', dest='lr', help='Base learning rate.', 29 | default=0.001, type=float) 30 | parser.add_argument('--data_dir', dest='data_dir', help='Directory path for data.', 31 | default='', type=str) 32 | parser.add_argument('--filename_list', dest='filename_list', help='Path to text file containing relative paths for every example.', 33 | default='', type=str) 34 | parser.add_argument('--output_string', dest='output_string', help='String appended to output snapshots.', default = '', type=str) 35 | parser.add_argument('--dataset', dest='dataset', help='Dataset type.', default='Pose_300W_LP', type=str) 36 | 37 | args = parser.parse_args() 38 | return args 39 | 40 | def get_ignored_params(model): 41 | # Generator function that yields ignored params. 42 | b = [model.conv1, model.bn1] 43 | for i in range(len(b)): 44 | for module_name, module in b[i].named_modules(): 45 | if 'bn' in module_name: 46 | module.eval() 47 | for name, param in module.named_parameters(): 48 | yield param 49 | 50 | def get_non_ignored_params(model): 51 | # Generator function that yields params that will be optimized. 52 | b = [model.layer1, model.layer2, model.layer3, model.layer4] 53 | for i in range(len(b)): 54 | for module_name, module in b[i].named_modules(): 55 | if 'bn' in module_name: 56 | module.eval() 57 | for name, param in module.named_parameters(): 58 | yield param 59 | 60 | def get_fc_params(model): 61 | # Generator function that yields fc layer params. 62 | b = [model.fc_angles] 63 | for i in range(len(b)): 64 | for module_name, module in b[i].named_modules(): 65 | for name, param in module.named_parameters(): 66 | yield param 67 | 68 | def load_filtered_state_dict(model, snapshot): 69 | # By user apaszke from discuss.pytorch.org 70 | model_dict = model.state_dict() 71 | snapshot = {k: v for k, v in snapshot.items() if k in model_dict} 72 | model_dict.update(snapshot) 73 | model.load_state_dict(model_dict) 74 | 75 | if __name__ == '__main__': 76 | args = parse_args() 77 | 78 | cudnn.enabled = True 79 | num_epochs = args.num_epochs 80 | batch_size = args.batch_size 81 | gpu = args.gpu_id 82 | 83 | if not os.path.exists('output/snapshots'): 84 | os.makedirs('output/snapshots') 85 | 86 | # ResNet50 87 | model = hopenet.ResNet(torchvision.models.resnet.Bottleneck, [3, 4, 6, 3], 3) 88 | load_filtered_state_dict(model, model_zoo.load_url('https://download.pytorch.org/models/resnet50-19c8e357.pth')) 89 | 90 | print 'Loading data.' 91 | 92 | transformations = transforms.Compose([transforms.Scale(240), 93 | transforms.RandomCrop(224), transforms.ToTensor(), 94 | transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])]) 95 | 96 | if args.dataset == 'Pose_300W_LP': 97 | pose_dataset = datasets.Pose_300W_LP(args.data_dir, args.filename_list, transformations) 98 | elif args.dataset == 'Pose_300W_LP_random_ds': 99 | pose_dataset = datasets.Pose_300W_LP_random_ds(args.data_dir, args.filename_list, transformations) 100 | elif args.dataset == 'AFLW2000': 101 | pose_dataset = datasets.AFLW2000(args.data_dir, args.filename_list, transformations) 102 | elif args.dataset == 'BIWI': 103 | pose_dataset = datasets.BIWI(args.data_dir, args.filename_list, transformations) 104 | elif args.dataset == 'AFLW': 105 | pose_dataset = datasets.AFLW(args.data_dir, args.filename_list, transformations) 106 | elif args.dataset == 'AFLW_aug': 107 | pose_dataset = datasets.AFLW_aug(args.data_dir, args.filename_list, transformations) 108 | elif args.dataset == 'AFW': 109 | pose_dataset = datasets.AFW(args.data_dir, args.filename_list, transformations) 110 | else: 111 | print 'Error: not a valid dataset name' 112 | sys.exit() 113 | train_loader = torch.utils.data.DataLoader(dataset=pose_dataset, 114 | batch_size=batch_size, 115 | shuffle=True, 116 | num_workers=2) 117 | 118 | model.cuda(gpu) 119 | criterion = nn.MSELoss().cuda(gpu) 120 | 121 | optimizer = torch.optim.Adam([{'params': get_ignored_params(model), 'lr': 0}, 122 | {'params': get_non_ignored_params(model), 'lr': args.lr}, 123 | {'params': get_fc_params(model), 'lr': args.lr * 5}], 124 | lr = args.lr) 125 | 126 | print 'Ready to train network.' 127 | print 'First phase of training.' 128 | for epoch in range(num_epochs): 129 | for i, (images, labels, cont_labels, name) in enumerate(train_loader): 130 | images = Variable(images).cuda(gpu) 131 | 132 | label_angles = Variable(cont_labels[:,:3]).cuda(gpu) 133 | angles = model(images) 134 | 135 | loss = criterion(angles, label_angles) 136 | optimizer.zero_grad() 137 | loss.backward() 138 | optimizer.step() 139 | 140 | if (i+1) % 100 == 0: 141 | print ('Epoch [%d/%d], Iter [%d/%d] Loss: %.4f' 142 | %(epoch+1, num_epochs, i+1, len(pose_dataset)//batch_size, loss.data[0])) 143 | 144 | # Save models at numbered epochs. 145 | if epoch % 1 == 0 and epoch < num_epochs: 146 | print 'Taking snapshot...' 147 | torch.save(model.state_dict(), 148 | 'output/snapshots/' + args.output_string + '_epoch_'+ str(epoch+1) + '.pkl') 149 | -------------------------------------------------------------------------------- /code/utils.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import torch 3 | from torch.utils.serialization import load_lua 4 | import os 5 | import scipy.io as sio 6 | import cv2 7 | import math 8 | from math import cos, sin 9 | 10 | def softmax_temperature(tensor, temperature): 11 | result = torch.exp(tensor / temperature) 12 | result = torch.div(result, torch.sum(result, 1).unsqueeze(1).expand_as(result)) 13 | return result 14 | 15 | def get_pose_params_from_mat(mat_path): 16 | # This functions gets the pose parameters from the .mat 17 | # Annotations that come with the Pose_300W_LP dataset. 18 | mat = sio.loadmat(mat_path) 19 | # [pitch yaw roll tdx tdy tdz scale_factor] 20 | pre_pose_params = mat['Pose_Para'][0] 21 | # Get [pitch, yaw, roll, tdx, tdy] 22 | pose_params = pre_pose_params[:5] 23 | return pose_params 24 | 25 | def get_ypr_from_mat(mat_path): 26 | # Get yaw, pitch, roll from .mat annotation. 27 | # They are in radians 28 | mat = sio.loadmat(mat_path) 29 | # [pitch yaw roll tdx tdy tdz scale_factor] 30 | pre_pose_params = mat['Pose_Para'][0] 31 | # Get [pitch, yaw, roll] 32 | pose_params = pre_pose_params[:3] 33 | return pose_params 34 | 35 | def get_pt2d_from_mat(mat_path): 36 | # Get 2D landmarks 37 | mat = sio.loadmat(mat_path) 38 | pt2d = mat['pt2d'] 39 | return pt2d 40 | 41 | def mse_loss(input, target): 42 | return torch.sum(torch.abs(input.data - target.data) ** 2) 43 | 44 | def plot_pose_cube(img, yaw, pitch, roll, tdx=None, tdy=None, size=150.): 45 | # Input is a cv2 image 46 | # pose_params: (pitch, yaw, roll, tdx, tdy) 47 | # Where (tdx, tdy) is the translation of the face. 48 | # For pose we have [pitch yaw roll tdx tdy tdz scale_factor] 49 | 50 | p = pitch * np.pi / 180 51 | y = -(yaw * np.pi / 180) 52 | r = roll * np.pi / 180 53 | if tdx != None and tdy != None: 54 | face_x = tdx - 0.50 * size 55 | face_y = tdy - 0.50 * size 56 | else: 57 | height, width = img.shape[:2] 58 | face_x = width / 2 - 0.5 * size 59 | face_y = height / 2 - 0.5 * size 60 | 61 | x1 = size * (cos(y) * cos(r)) + face_x 62 | y1 = size * (cos(p) * sin(r) + cos(r) * sin(p) * sin(y)) + face_y 63 | x2 = size * (-cos(y) * sin(r)) + face_x 64 | y2 = size * (cos(p) * cos(r) - sin(p) * sin(y) * sin(r)) + face_y 65 | x3 = size * (sin(y)) + face_x 66 | y3 = size * (-cos(y) * sin(p)) + face_y 67 | 68 | # Draw base in red 69 | cv2.line(img, (int(face_x), int(face_y)), (int(x1),int(y1)),(0,0,255),3) 70 | cv2.line(img, (int(face_x), int(face_y)), (int(x2),int(y2)),(0,0,255),3) 71 | cv2.line(img, (int(x2), int(y2)), (int(x2+x1-face_x),int(y2+y1-face_y)),(0,0,255),3) 72 | cv2.line(img, (int(x1), int(y1)), (int(x1+x2-face_x),int(y1+y2-face_y)),(0,0,255),3) 73 | # Draw pillars in blue 74 | cv2.line(img, (int(face_x), int(face_y)), (int(x3),int(y3)),(255,0,0),2) 75 | cv2.line(img, (int(x1), int(y1)), (int(x1+x3-face_x),int(y1+y3-face_y)),(255,0,0),2) 76 | cv2.line(img, (int(x2), int(y2)), (int(x2+x3-face_x),int(y2+y3-face_y)),(255,0,0),2) 77 | cv2.line(img, (int(x2+x1-face_x),int(y2+y1-face_y)), (int(x3+x1+x2-2*face_x),int(y3+y2+y1-2*face_y)),(255,0,0),2) 78 | # Draw top in green 79 | cv2.line(img, (int(x3+x1-face_x),int(y3+y1-face_y)), (int(x3+x1+x2-2*face_x),int(y3+y2+y1-2*face_y)),(0,255,0),2) 80 | cv2.line(img, (int(x2+x3-face_x),int(y2+y3-face_y)), (int(x3+x1+x2-2*face_x),int(y3+y2+y1-2*face_y)),(0,255,0),2) 81 | cv2.line(img, (int(x3), int(y3)), (int(x3+x1-face_x),int(y3+y1-face_y)),(0,255,0),2) 82 | cv2.line(img, (int(x3), int(y3)), (int(x3+x2-face_x),int(y3+y2-face_y)),(0,255,0),2) 83 | 84 | return img 85 | 86 | def draw_axis(img, yaw, pitch, roll, tdx=None, tdy=None, size = 100): 87 | 88 | pitch = pitch * np.pi / 180 89 | yaw = -(yaw * np.pi / 180) 90 | roll = roll * np.pi / 180 91 | 92 | if tdx != None and tdy != None: 93 | tdx = tdx 94 | tdy = tdy 95 | else: 96 | height, width = img.shape[:2] 97 | tdx = width / 2 98 | tdy = height / 2 99 | 100 | # X-Axis pointing to right. drawn in red 101 | x1 = size * (cos(yaw) * cos(roll)) + tdx 102 | y1 = size * (cos(pitch) * sin(roll) + cos(roll) * sin(pitch) * sin(yaw)) + tdy 103 | 104 | # Y-Axis | drawn in green 105 | # v 106 | x2 = size * (-cos(yaw) * sin(roll)) + tdx 107 | y2 = size * (cos(pitch) * cos(roll) - sin(pitch) * sin(yaw) * sin(roll)) + tdy 108 | 109 | # Z-Axis (out of the screen) drawn in blue 110 | x3 = size * (sin(yaw)) + tdx 111 | y3 = size * (-cos(yaw) * sin(pitch)) + tdy 112 | 113 | cv2.line(img, (int(tdx), int(tdy)), (int(x1),int(y1)),(0,0,255),3) 114 | cv2.line(img, (int(tdx), int(tdy)), (int(x2),int(y2)),(0,255,0),3) 115 | cv2.line(img, (int(tdx), int(tdy)), (int(x3),int(y3)),(255,0,0),2) 116 | 117 | return img 118 | -------------------------------------------------------------------------------- /conan-cruise.gif: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/natanielruiz/deep-head-pose/f7bbb9981c2953c2eca67748d6492a64c8243946/conan-cruise.gif --------------------------------------------------------------------------------