[
  {
    "path": ".gitignore",
    "content": "*.pyc\n__pycache__\n.idea/\n*.iml\n**/venv\ndata/test\nexternal/cocoapi\ntensorflow_toolkit/tests/models\ntensorflow_toolkit/**/model"
  },
  {
    "path": ".travis.yml",
    "content": "language: python\nsudo: required\ndist: xenial\n\npython:\n  - \"3.5\"\ncache: pip\n\ninstall:\n  - bash ./init_venv.sh\n\njobs:\n  include:\n  - stage: Tests\n    script:\n      - . venv/bin/activate\n      - python -m unittest\n"
  },
  {
    "path": "LICENSE",
    "content": "                                 Apache License\n                           Version 2.0, January 2004\n                        http://www.apache.org/licenses/\n\n   TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION\n\n   1. Definitions.\n\n      \"License\" shall mean the terms and conditions for use, reproduction,\n      and distribution as defined by Sections 1 through 9 of this document.\n\n      \"Licensor\" shall mean the copyright owner or entity authorized by\n      the copyright owner that is granting the License.\n\n      \"Legal Entity\" shall mean the union of the acting entity and all\n      other entities that control, are controlled by, or are under common\n      control with that entity. For the purposes of this definition,\n      \"control\" means (i) the power, direct or indirect, to cause the\n      direction or management of such entity, whether by contract or\n      otherwise, or (ii) ownership of fifty percent (50%) or more of the\n      outstanding shares, or (iii) beneficial ownership of such entity.\n\n      \"You\" (or \"Your\") shall mean an individual or Legal Entity\n      exercising permissions granted by this License.\n\n      \"Source\" form shall mean the preferred form for making modifications,\n      including but not limited to software source code, documentation\n      source, and configuration files.\n\n      \"Object\" form shall mean any form resulting from mechanical\n      transformation or translation of a Source form, including but\n      not limited to compiled object code, generated documentation,\n      and conversions to other media types.\n\n      \"Work\" shall mean the work of authorship, whether in Source or\n      Object form, made available under the License, as indicated by a\n      copyright notice that is included in or attached to the work\n      (an example is provided in the Appendix below).\n\n      \"Derivative Works\" shall mean any work, whether in Source or Object\n      form, that is based on (or derived from) the Work and for which the\n      editorial revisions, annotations, elaborations, or other modifications\n      represent, as a whole, an original work of authorship. For the purposes\n      of this License, Derivative Works shall not include works that remain\n      separable from, or merely link (or bind by name) to the interfaces of,\n      the Work and Derivative Works thereof.\n\n      \"Contribution\" shall mean any work of authorship, including\n      the original version of the Work and any modifications or additions\n      to that Work or Derivative Works thereof, that is intentionally\n      submitted to Licensor for inclusion in the Work by the copyright owner\n      or by an individual or Legal Entity authorized to submit on behalf of\n      the copyright owner. For the purposes of this definition, \"submitted\"\n      means any form of electronic, verbal, or written communication sent\n      to the Licensor or its representatives, including but not limited to\n      communication on electronic mailing lists, source code control systems,\n      and issue tracking systems that are managed by, or on behalf of, the\n      Licensor for the purpose of discussing and improving the Work, but\n      excluding communication that is conspicuously marked or otherwise\n      designated in writing by the copyright owner as \"Not a Contribution.\"\n\n      \"Contributor\" shall mean Licensor and any individual or Legal Entity\n      on behalf of whom a Contribution has been received by Licensor and\n      subsequently incorporated within the Work.\n\n   2. Grant of Copyright License. Subject to the terms and conditions of\n      this License, each Contributor hereby grants to You a perpetual,\n      worldwide, non-exclusive, no-charge, royalty-free, irrevocable\n      copyright license to reproduce, prepare Derivative Works of,\n      publicly display, publicly perform, sublicense, and distribute the\n      Work and such Derivative Works in Source or Object form.\n\n   3. Grant of Patent License. Subject to the terms and conditions of\n      this License, each Contributor hereby grants to You a perpetual,\n      worldwide, non-exclusive, no-charge, royalty-free, irrevocable\n      (except as stated in this section) patent license to make, have made,\n      use, offer to sell, sell, import, and otherwise transfer the Work,\n      where such license applies only to those patent claims licensable\n      by such Contributor that are necessarily infringed by their\n      Contribution(s) alone or by combination of their Contribution(s)\n      with the Work to which such Contribution(s) was submitted. If You\n      institute patent litigation against any entity (including a\n      cross-claim or counterclaim in a lawsuit) alleging that the Work\n      or a Contribution incorporated within the Work constitutes direct\n      or contributory patent infringement, then any patent licenses\n      granted to You under this License for that Work shall terminate\n      as of the date such litigation is filed.\n\n   4. Redistribution. You may reproduce and distribute copies of the\n      Work or Derivative Works thereof in any medium, with or without\n      modifications, and in Source or Object form, provided that You\n      meet the following conditions:\n\n      (a) You must give any other recipients of the Work or\n          Derivative Works a copy of this License; and\n\n      (b) You must cause any modified files to carry prominent notices\n          stating that You changed the files; and\n\n      (c) You must retain, in the Source form of any Derivative Works\n          that You distribute, all copyright, patent, trademark, and\n          attribution notices from the Source form of the Work,\n          excluding those notices that do not pertain to any part of\n          the Derivative Works; and\n\n      (d) If the Work includes a \"NOTICE\" text file as part of its\n          distribution, then any Derivative Works that You distribute must\n          include a readable copy of the attribution notices contained\n          within such NOTICE file, excluding those notices that do not\n          pertain to any part of the Derivative Works, in at least one\n          of the following places: within a NOTICE text file distributed\n          as part of the Derivative Works; within the Source form or\n          documentation, if provided along with the Derivative Works; or,\n          within a display generated by the Derivative Works, if and\n          wherever such third-party notices normally appear. The contents\n          of the NOTICE file are for informational purposes only and\n          do not modify the License. You may add Your own attribution\n          notices within Derivative Works that You distribute, alongside\n          or as an addendum to the NOTICE text from the Work, provided\n          that such additional attribution notices cannot be construed\n          as modifying the License.\n\n      You may add Your own copyright statement to Your modifications and\n      may provide additional or different license terms and conditions\n      for use, reproduction, or distribution of Your modifications, or\n      for any such Derivative Works as a whole, provided Your use,\n      reproduction, and distribution of the Work otherwise complies with\n      the conditions stated in this License.\n\n   5. Submission of Contributions. Unless You explicitly state otherwise,\n      any Contribution intentionally submitted for inclusion in the Work\n      by You to the Licensor shall be under the terms and conditions of\n      this License, without any additional terms or conditions.\n      Notwithstanding the above, nothing herein shall supersede or modify\n      the terms of any separate license agreement you may have executed\n      with Licensor regarding such Contributions.\n\n   6. Trademarks. This License does not grant permission to use the trade\n      names, trademarks, service marks, or product names of the Licensor,\n      except as required for reasonable and customary use in describing the\n      origin of the Work and reproducing the content of the NOTICE file.\n\n   7. Disclaimer of Warranty. Unless required by applicable law or\n      agreed to in writing, Licensor provides the Work (and each\n      Contributor provides its Contributions) on an \"AS IS\" BASIS,\n      WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or\n      implied, including, without limitation, any warranties or conditions\n      of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A\n      PARTICULAR PURPOSE. You are solely responsible for determining the\n      appropriateness of using or redistributing the Work and assume any\n      risks associated with Your exercise of permissions under this License.\n\n   8. Limitation of Liability. In no event and under no legal theory,\n      whether in tort (including negligence), contract, or otherwise,\n      unless required by applicable law (such as deliberate and grossly\n      negligent acts) or agreed to in writing, shall any Contributor be\n      liable to You for damages, including any direct, indirect, special,\n      incidental, or consequential damages of any character arising as a\n      result of this License or out of the use or inability to use the\n      Work (including but not limited to damages for loss of goodwill,\n      work stoppage, computer failure or malfunction, or any and all\n      other commercial damages or losses), even if such Contributor\n      has been advised of the possibility of such damages.\n\n   9. Accepting Warranty or Additional Liability. While redistributing\n      the Work or Derivative Works thereof, You may choose to offer,\n      and charge a fee for, acceptance of support, warranty, indemnity,\n      or other liability obligations and/or rights consistent with this\n      License. However, in accepting such obligations, You may act only\n      on Your own behalf and on Your sole responsibility, not on behalf\n      of any other Contributor, and only if You agree to indemnify,\n      defend, and hold each Contributor harmless for any liability\n      incurred by, or claims asserted against, such Contributor by reason\n      of your accepting any such warranty or additional liability.\n\n   END OF TERMS AND CONDITIONS\n\n   APPENDIX: How to apply the Apache License to your work.\n\n      To apply the Apache License to your work, attach the following\n      boilerplate notice, with the fields enclosed by brackets \"{}\"\n      replaced with your own identifying information. (Don't include\n      the brackets!)  The text should be enclosed in the appropriate\n      comment syntax for the file format. We also recommend that a\n      file or class name and description of purpose be included on the\n      same \"printed page\" as the copyright notice for easier\n      identification within third-party archives.\n\n   Copyright 2018 algo\n\n   Licensed under the Apache License, Version 2.0 (the \"License\");\n   you may not use this file except in compliance with the License.\n   You may obtain a copy of the License at\n\n       http://www.apache.org/licenses/LICENSE-2.0\n\n   Unless required by applicable law or agreed to in writing, software\n   distributed under the License is distributed on an \"AS IS\" BASIS,\n   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n   See the License for the specific language governing permissions and\n   limitations under the License.\n"
  },
  {
    "path": "README.md",
    "content": "# Face Recognition in PyTorch\n[![License](https://img.shields.io/badge/License-Apache%202.0-blue.svg)](https://opensource.org/licenses/Apache-2.0)\n[![Build Status](https://travis-ci.com/grib0ed0v/face_recognition.pytorch.svg?branch=develop)](https://travis-ci.com/grib0ed0v/face_recognition.pytorch)\n\nBy [Alexey Gruzdev](https://www.linkedin.com/in/alexey-gruzdev-454399128/) and [Vladislav Sovrasov](https://www.linkedin.com/in/%D0%B2%D0%BB%D0%B0%D0%B4%D0%B8%D1%81%D0%BB%D0%B0%D0%B2-%D1%81%D0%BE%D0%B2%D1%80%D0%B0%D1%81%D0%BE%D0%B2-173b23104/)\n\n## Introduction\n\n*A repository for different experimental Face Recognition models such as [CosFace](https://arxiv.org/pdf/1801.09414.pdf), [ArcFace](https://arxiv.org/pdf/1801.07698.pdf), [SphereFace](https://arxiv.org/pdf/1704.08063.pdf), [SV-Softmax](https://arxiv.org/pdf/1812.11317.pdf), etc.*\n\n## Contents\n1. [Installation](#installation)\n2. [Preparation](#preparation)\n3. [Train/Eval](#traineval)\n4. [Models](#models)\n5. [Face Recognition Demo](#demo)\n\n\n## Installation\n1. Create and activate virtual python environment\n```bash\nbash init_venv.sh\n. venv/bin/activate\n```\n\n\n\n\n## Preparation\n\n1. For Face Recognition training you should download [VGGFace2](http://www.robots.ox.ac.uk/~vgg/data/vgg_face2/) data. We will refer to this folder as `$VGGFace2_ROOT`.\n2. For Face Recognition evaluation you need to download [LFW](http://vis-www.cs.umass.edu/lfw/) data and [LFW landmarks](https://github.com/clcarwin/sphereface_pytorch/blob/master/data/lfw_landmark.txt).  Place everything in one folder, which will be `$LFW_ROOT`.\n\n\n\n\n## Train/Eval\n1. Go to `$FR_ROOT` folder\n```bash\ncd $FR_ROOT/\n```\n\n2. To start training FR model:\n```bash\npython train.py --train_data_root $VGGFace2_ROOT/train/ --train_list $VGGFace2_ROOT/meta/train_list.txt\n--train_landmarks  $VGGFace2_ROOT/bb_landmark/ --val_data_root  $LFW_ROOT/lfw/ --val_list $LFW_ROOT/pairs.txt  \n--val_landmarks $LFW_ROOT/lfw_landmark.txt --train_batch_size 200  --snap_prefix mobilenet_256 --lr 0.35\n--embed_size 256 --model mobilenet --device 1\n```\n\n3. To evaluate FR snapshot (let's say we have MobileNet with 256 embedding size trained for 300k):\n```bash\n python evaluate_lfw.py --val_data_root $LFW_ROOT/lfw/ --val_list $LFW_ROOT/pairs.txt\n --val_landmarks $LFW_ROOT/lfw_landmark.txt --snap /path/to/snapshot/mobilenet_256_300000.pt --model mobilenet --embed_size 256\n```\n\n## Configuration files\nBesides passing all the required parameters via command line, the training script allows to read them from a `yaml` configuration file.\nEach line of such file should contain a valid description of one parameter in the `yaml` fromat.\nExample:\n```yml\n#optimizer parameters\nlr: 0.4\ntrain_batch_size: 256\n#loss options\nmargin_type: cos\ns: 30\nm: 0.35\n#model parameters\nmodel: mobilenet\nembed_size: 256\n#misc\nsnap_prefix: MobileFaceNet\ndevices: [0, 1]\n#datasets\ntrain_dataset: vgg\ntrain_data_root: $VGGFace2_ROOT/train/\n#... and so on\n```\nPath to the config file can be passed to the training script via command line. In case if any other arguments were passed before the config, they will be overwritten.\n```bash\npython train.py -m 0.35 @./my_config.yml #here m can be overwritten with the value from my_config.yml\n```\n\n\n\n## Models\n\n1. You can download pretrained model from fileshare as well - https://download.01.org/openvinotoolkit/open_model_zoo/training_toolbox_pytorch/models/fr/Mobilenet_se_focal_121000.pt\n```bash\ncd $FR_ROOT\npython evaluate_lfw.py --val_data_root $LFW_ROOT/lfw/ --val_list $LFW_ROOT/pairs.txt --val_landmarks $LFW_ROOT/lfw_landmark.txt\n--snap /path/to/snapshot/Mobilenet_se_focal_121000.pt --model mobilenet --embed_size 256\n```\n\n2. You should get the following output:\n```\nI1114 09:33:37.846870 10544 evaluate_lfw.py:242] Accuracy/Val_same_accuracy mean: 0.9923\nI1114 09:33:37.847019 10544 evaluate_lfw.py:243] Accuracy/Val_diff_accuracy mean: 0.9970\nI1114 09:33:37.847069 10544 evaluate_lfw.py:244] Accuracy/Val_accuracy mean: 0.9947\nI1114 09:33:37.847179 10544 evaluate_lfw.py:245] Accuracy/Val_accuracy std dev: 0.0035\nI1114 09:33:37.847229 10544 evaluate_lfw.py:246] AUC: 0.9995\nI1114 09:33:37.847305 10544 evaluate_lfw.py:247] Estimated threshold: 0.7241\n```\n\n## Demo\n\n1. For setting up demo, please go to [Face Recognition demo with OpenVINO Toolkit](./demo/README.md)\n"
  },
  {
    "path": "__init__.py",
    "content": ""
  },
  {
    "path": "configs/mobilefacenet_vgg2.yml",
    "content": "#optimizer parameters\nlr: 0.4\ntrain_batch_size: 256\n#loss options\nmargin_type: cos\ns: 30\nm: 0.35\nmining_type: sv\nt: 1.1\n#model parameters\nmodel: mobilenet\nembed_size: 256\n\ntrain_dataset: vgg\nsnap_prefix: MobileFaceNet\ndevices: [0, 1]\n"
  },
  {
    "path": "datasets/__init__.py",
    "content": "from .lfw import LFW\nfrom .vggface2 import VGGFace2\nfrom .ms_celeb1m import MSCeleb1M\nfrom .trillion_pairs import TrillionPairs\nfrom .imdbface import IMDBFace\n\nfrom .celeba import CelebA\nfrom .ndg import NDG\n\n__all__ = [LFW, VGGFace2, MSCeleb1M, TrillionPairs, IMDBFace, CelebA, NDG]\n"
  },
  {
    "path": "datasets/casia.py",
    "content": "\"\"\"\n Copyright (c) 2018 Intel Corporation\n Licensed under the Apache License, Version 2.0 (the \"License\");\n you may not use this file except in compliance with the License.\n You may obtain a copy of the License at\n      http://www.apache.org/licenses/LICENSE-2.0\n Unless required by applicable law or agreed to in writing, software\n distributed under the License is distributed on an \"AS IS\" BASIS,\n WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n See the License for the specific language governing permissions and\n limitations under the License.\n\"\"\"\n\nimport os.path as osp\n\nfrom tqdm import tqdm\nfrom torch.utils.data import Dataset\nimport cv2 as cv\n\nfrom utils.face_align import FivePointsAligner\n\nclass CASIA(Dataset):\n    \"\"\"CASIA Dataset compatible with PyTorch DataLoader.\"\"\"\n    def __init__(self, images_root_path, image_list_path, transform, use_landmarks=True):\n        self.image_list_path = image_list_path\n        self.images_root_path = images_root_path\n        self.identities = {}\n        self.use_landmarks = use_landmarks\n        self.samples_info = self._read_samples_info()\n        self.transform = transform\n\n    def _read_samples_info(self):\n        \"\"\"Reads annotation of the dataset\"\"\"\n        samples = []\n        with open(self.image_list_path, 'r') as f:\n            for line in tqdm(f.readlines(), 'Preparing CASIA dataset'):\n                sample = line.split()\n                sample_id = sample[1]\n                landmarks = [[sample[i], sample[i+1]] for i in range(2, 12, 2)]\n                self.identities[sample_id] = [1]\n                samples.append((osp.join(self.images_root_path, sample[0]), sample_id, landmarks))\n\n        return samples\n\n    def get_num_classes(self):\n        \"\"\"Returns total number of identities\"\"\"\n        return len(self.identities)\n\n    def __len__(self):\n        \"\"\"Returns total number of samples\"\"\"\n        return len(self.samples_info)\n\n    def __getitem__(self, idx):\n        img = cv.imread(self.samples_info[idx][0])\n        if self.use_landmarks:\n            img = FivePointsAligner.align(img, self.samples_info[idx][2],\n                                          d_size=(200, 200), normalized=True, show=False)\n\n        if self.transform:\n            img = self.transform(img)\n        return {'img': img, 'label': int(self.samples_info[idx][1])}\n"
  },
  {
    "path": "datasets/celeba.py",
    "content": "\"\"\"\n Copyright (c) 2018 Intel Corporation\n Licensed under the Apache License, Version 2.0 (the \"License\");\n you may not use this file except in compliance with the License.\n You may obtain a copy of the License at\n      http://www.apache.org/licenses/LICENSE-2.0\n Unless required by applicable law or agreed to in writing, software\n distributed under the License is distributed on an \"AS IS\" BASIS,\n WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n See the License for the specific language governing permissions and\n limitations under the License.\n\"\"\"\n\nimport os.path as osp\n\nimport numpy as np\nfrom tqdm import tqdm\nfrom torch.utils.data import Dataset\nimport cv2 as cv\n\n\nclass CelebA(Dataset):\n    \"\"\"CelebA Dataset compatible with PyTorch DataLoader.\"\"\"\n    def __init__(self, images_root_path, landmarks_folder_path, transform=None, test=False):\n        self.test = test\n        self.have_landmarks = True\n        self.images_root_path = images_root_path\n        bb_file_name = 'list_bbox_celeba.txt'\n        landmarks_file_name = 'list_landmarks_celeba.txt'\n        self.detections_file = open(osp.join(landmarks_folder_path, bb_file_name), 'r')\n        self.landmarks_file = open(osp.join(landmarks_folder_path, landmarks_file_name), 'r')\n        self.samples_info = self._read_samples_info()\n        self.transform = transform\n\n    def _read_samples_info(self):\n        \"\"\"Reads annotation of the dataset\"\"\"\n        samples = []\n\n        detections_file_lines = self.detections_file.readlines()[2:]\n        landmarks_file_lines = self.landmarks_file.readlines()[2:]\n        assert len(detections_file_lines) == len(landmarks_file_lines)\n\n        if self.test:\n            images_range = range(182638, len(landmarks_file_lines))\n        else:\n            images_range = range(182637)\n\n        for i in tqdm(images_range):\n            line = detections_file_lines[i].strip()\n            img_name = line.split(' ')[0]\n            img_path = osp.join(self.images_root_path, img_name)\n\n            bbox = list(filter(bool, line.split(' ')[1:]))\n            bbox = [int(coord) for coord in bbox]\n            if bbox[2] == 0 or bbox[3] == 0:\n                continue\n\n            line_landmarks = landmarks_file_lines[i].strip().split(' ')[1:]\n            landmarks = list(filter(bool, line_landmarks))\n            landmarks = [float(coord) for coord in landmarks]\n            samples.append((img_path, bbox, landmarks))\n\n        return samples\n\n    def __len__(self):\n        \"\"\"Returns total number of samples\"\"\"\n        return len(self.samples_info)\n\n    def __getitem__(self, idx):\n        \"\"\"Returns sample (image, landmarks) by index\"\"\"\n        img = cv.imread(self.samples_info[idx][0], cv.IMREAD_COLOR)\n        bbox = self.samples_info[idx][1]\n        landmarks = self.samples_info[idx][2]\n\n        img = img[bbox[1]:bbox[1] + bbox[3], bbox[0]:bbox[0] + bbox[2]]\n        landmarks = np.array([(float(landmarks[2*i]-bbox[0]) / bbox[2],\n                               float(landmarks[2*i + 1]-bbox[1])/ bbox[3]) \\\n                               for i in range(len(landmarks)//2)]).reshape(-1)\n        data = {'img': img, 'landmarks': landmarks}\n        if self.transform:\n            data = self.transform(data)\n        return data\n"
  },
  {
    "path": "datasets/imdbface.py",
    "content": "\"\"\"\n Copyright (c) 2018 Intel Corporation\n Licensed under the Apache License, Version 2.0 (the \"License\");\n you may not use this file except in compliance with the License.\n You may obtain a copy of the License at\n      http://www.apache.org/licenses/LICENSE-2.0\n Unless required by applicable law or agreed to in writing, software\n distributed under the License is distributed on an \"AS IS\" BASIS,\n WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n See the License for the specific language governing permissions and\n limitations under the License.\n\"\"\"\n\nimport os.path as osp\n\nimport cv2 as cv\nfrom tqdm import tqdm\nfrom torch.utils.data import Dataset\n\nfrom utils.face_align import FivePointsAligner\n\n\nclass IMDBFace(Dataset):\n    \"\"\"IMDBFace Dataset compatible with PyTorch DataLoader.\"\"\"\n    def __init__(self, images_root_path, image_list_path, transform=None):\n        self.image_list_path = image_list_path\n        self.images_root_path = images_root_path\n        self.identities = {}\n\n        assert osp.isfile(image_list_path)\n        self.have_landmarks = True\n\n        self.all_samples_info = self._read_samples_info()\n        self.samples_info = self.all_samples_info\n        self.transform = transform\n\n    def _read_samples_info(self):\n        \"\"\"Reads annotation of the dataset\"\"\"\n        samples = []\n\n        with open(self.image_list_path, 'r') as f:\n            images_file_lines = f.readlines()\n            last_class_id = -1\n\n            for i in tqdm(range(len(images_file_lines))):\n                line = images_file_lines[i]\n                terms = line.split('|')\n                if len(terms) < 3:\n                    continue # FD has failed on this imsage\n                path, landmarks, _ = terms\n                image_id, _ = path.rsplit('/', 1)\n\n                if image_id in self.identities:\n                    self.identities[image_id].append(len(samples))\n                else:\n                    last_class_id += 1\n                    self.identities[image_id] = [len(samples)]\n\n                landmarks = [float(coord) for coord in landmarks.strip().split(' ')]\n                assert len(landmarks) == 10\n                samples.append((osp.join(self.images_root_path, path).strip(), last_class_id, image_id, landmarks))\n\n        return samples\n\n    def get_weights(self):\n        \"\"\"Computes weights of the each identity in dataset according to frequency of it's occurance\"\"\"\n        weights = [0.]*len(self.all_samples_info)\n        for i, sample in enumerate(self.all_samples_info):\n            weights[i] = float(len(self.all_samples_info)) / len(self.identities[sample[2]])\n        return weights\n\n    def get_num_classes(self):\n        \"\"\"Returns total number of identities\"\"\"\n        return len(self.identities)\n\n    def __len__(self):\n        \"\"\"Returns total number of samples\"\"\"\n        return len(self.samples_info)\n\n    def __getitem__(self, idx):\n        \"\"\"Returns sample (image, class id, image id) by index\"\"\"\n        img = cv.imread(self.samples_info[idx][0], cv.IMREAD_COLOR)\n        landmarks = self.samples_info[idx][-1]\n        img = FivePointsAligner.align(img, landmarks, d_size=(200, 200), normalized=True, show=False)\n\n        if self.transform:\n            img = self.transform(img)\n\n        return {'img': img, 'label': self.samples_info[idx][1], 'instance': self.samples_info[idx][2]}\n"
  },
  {
    "path": "datasets/lfw.py",
    "content": "\"\"\"\n Copyright (c) 2018 Intel Corporation\n Licensed under the Apache License, Version 2.0 (the \"License\");\n you may not use this file except in compliance with the License.\n You may obtain a copy of the License at\n      http://www.apache.org/licenses/LICENSE-2.0\n Unless required by applicable law or agreed to in writing, software\n distributed under the License is distributed on an \"AS IS\" BASIS,\n WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n See the License for the specific language governing permissions and\n limitations under the License.\n\"\"\"\n\nimport os.path as osp\nimport cv2 as cv\nimport numpy as np\nfrom torch.utils.data import Dataset\n\nfrom utils.face_align import FivePointsAligner\n\n\nclass LFW(Dataset):\n    \"\"\"LFW Dataset compatible with PyTorch DataLoader.\"\"\"\n    def __init__(self, images_root_path, pairs_path, landmark_file_path='', transform=None):\n        self.pairs_path = pairs_path\n        self.images_root_path = images_root_path\n        self.landmark_file_path = landmark_file_path\n        self.use_landmarks = len(self.landmark_file_path) > 0\n        if self.use_landmarks:\n            self.landmarks = self._read_landmarks()\n        self.pairs = self._read_pairs()\n        self.transform = transform\n\n    def _read_landmarks(self):\n        \"\"\"Reads landmarks of the dataset\"\"\"\n        landmarks = {}\n        with open(self.landmark_file_path, 'r') as f:\n            for line in f.readlines():\n                sp = line.split()\n                key = sp[0][sp[0].rfind('/')+1:]\n                landmarks[key] = [[int(sp[i]), int(sp[i+1])] for i in range(1, 11, 2)]\n\n        return landmarks\n\n    def _read_pairs(self):\n        \"\"\"Reads annotation of the dataset\"\"\"\n        pairs = []\n        with open(self.pairs_path, 'r') as f:\n            for line in f.readlines()[1:]:  # skip header\n                pair = line.strip().split()\n                pairs.append(pair)\n\n        file_ext = 'jpg'\n        lfw_dir = self.images_root_path\n        path_list = []\n\n        for pair in pairs:\n            if len(pair) == 3:\n                path0 = osp.join(lfw_dir, pair[0], pair[0] + '_' + '%04d' % int(pair[1]) + '.' + file_ext)\n                id0 = pair[0]\n                path1 = osp.join(lfw_dir, pair[0], pair[0] + '_' + '%04d' % int(pair[2]) + '.' + file_ext)\n                id1 = pair[0]\n                issame = True\n            elif len(pair) == 4:\n                path0 = osp.join(lfw_dir, pair[0], pair[0] + '_' + '%04d' % int(pair[1]) + '.' + file_ext)\n                id0 = pair[0]\n                path1 = osp.join(lfw_dir, pair[2], pair[2] + '_' + '%04d' % int(pair[3]) + '.' + file_ext)\n                id1 = pair[0]\n                issame = False\n\n            path_list.append((path0, path1, issame, id0, id1))\n\n        return path_list\n\n    def _load_img(self, img_path):\n        \"\"\"Loads an image from dist, then performs face alignment and applies transform\"\"\"\n        img = cv.imread(img_path, cv.IMREAD_COLOR)\n\n        if self.use_landmarks:\n            landmarks = np.array(self.landmarks[img_path[img_path.rfind('/')+1:]]).reshape(-1)\n            img = FivePointsAligner.align(img, landmarks, show=False)\n\n        if self.transform is None:\n            return img\n\n        return self.transform(img)\n\n    def show_item(self, index):\n        \"\"\"Saves a pair with a given index to disk\"\"\"\n        path_1, path_2, _, _, _ = self.pairs[index]\n        img1 = cv.imread(path_1)\n        img2 = cv.imread(path_2)\n        if self.use_landmarks:\n            landmarks1 = np.array(self.landmarks[path_1[path_1.rfind('/')+1:]]).reshape(-1)\n            landmarks2 = np.array(self.landmarks[path_2[path_2.rfind('/')+1:]]).reshape(-1)\n            img1 = FivePointsAligner.align(img1, landmarks1)\n            img2 = FivePointsAligner.align(img2, landmarks2)\n        else:\n            img1 = cv.resize(img1, (400, 400))\n            img2 = cv.resize(img2, (400, 400))\n        cv.imwrite('misclassified_{}.jpg'.format(index), np.hstack([img1, img2]))\n\n    def __getitem__(self, index):\n        \"\"\"Returns a pair of images and similarity flag by index\"\"\"\n        (path_1, path_2, is_same, id0, id1) = self.pairs[index]\n        img1, img2 = self._load_img(path_1), self._load_img(path_2)\n\n        return {'img1': img1, 'img2': img2, 'is_same': is_same, 'id0': id0, 'id1': id1}\n\n    def __len__(self):\n        \"\"\"Returns total number of samples\"\"\"\n        return len(self.pairs)\n"
  },
  {
    "path": "datasets/megaface.py",
    "content": "\"\"\"\n Copyright (c) 2018 Intel Corporation\n Licensed under the Apache License, Version 2.0 (the \"License\");\n you may not use this file except in compliance with the License.\n You may obtain a copy of the License at\n      http://www.apache.org/licenses/LICENSE-2.0\n Unless required by applicable law or agreed to in writing, software\n distributed under the License is distributed on an \"AS IS\" BASIS,\n WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n See the License for the specific language governing permissions and\n limitations under the License.\n\"\"\"\n\nimport numpy as np\nfrom torch.utils.data import Dataset\nimport cv2 as cv\n\nfrom utils.face_align import FivePointsAligner\n\n\nclass MegaFace(Dataset):\n    \"\"\"MegaFace Dataset compatible with PyTorch DataLoader.\"\"\"\n    def __init__(self, images_lsit, transform=None):\n        self.samples_info = images_lsit\n        self.transform = transform\n\n    def __len__(self):\n        \"\"\"Returns total number of samples\"\"\"\n        return len(self.samples_info)\n\n    def __getitem__(self, idx):\n        \"\"\"Returns sample (image, index)\"\"\"\n        img = None\n        try:\n            img = cv.imread(self.samples_info[idx]['path'], cv.IMREAD_COLOR)\n            bbox = self.samples_info[idx]['bbox']\n            landmarks = self.samples_info[idx]['landmarks']\n\n            if bbox is not None or landmarks is not None:\n                if landmarks is not None:\n                    landmarks = np.array(landmarks).reshape(5, -1)\n                    landmarks[:,0] = landmarks[:,0]*bbox[2] + bbox[0]\n                    landmarks[:,1] = landmarks[:,1]*bbox[3] + bbox[1]\n                    img = FivePointsAligner.align(img, landmarks.reshape(-1), d_size=(bbox[2], bbox[3]),\n                                                  normalized=False, show=False)\n                if bbox is not None and landmarks is None:\n                    img = img[bbox[1]:bbox[1] + bbox[3], bbox[0]:bbox[0] + bbox[2]]\n        except BaseException:\n            print('Corrupted image!', self.samples_info[idx])\n            img = np.zeros((128, 128, 3), dtype='uint8')\n\n        if self.transform:\n            img = self.transform(img)\n\n        return {'img': img, 'idx': idx}\n"
  },
  {
    "path": "datasets/ms_celeb1m.py",
    "content": "\"\"\"\n Copyright (c) 2018 Intel Corporation\n Licensed under the Apache License, Version 2.0 (the \"License\");\n you may not use this file except in compliance with the License.\n You may obtain a copy of the License at\n      http://www.apache.org/licenses/LICENSE-2.0\n Unless required by applicable law or agreed to in writing, software\n distributed under the License is distributed on an \"AS IS\" BASIS,\n WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n See the License for the specific language governing permissions and\n limitations under the License.\n\"\"\"\n\nimport os.path as osp\n\nimport cv2 as cv\nfrom tqdm import tqdm\nfrom torch.utils.data import Dataset\n\nfrom utils.face_align import FivePointsAligner\n\n\nclass MSCeleb1M(Dataset):\n    \"\"\"MSCeleb1M Dataset compatible with PyTorch DataLoader.\"\"\"\n    def __init__(self, images_root_path, image_list_path, transform=None):\n        self.image_list_path = image_list_path\n        self.images_root_path = images_root_path\n        self.identities = {}\n\n        assert osp.isfile(image_list_path)\n        self.have_landmarks = True\n\n        self.all_samples_info = self._read_samples_info()\n        self.samples_info = self.all_samples_info\n        self.transform = transform\n\n    def _read_samples_info(self):\n        \"\"\"Reads annotation of the dataset\"\"\"\n        samples = []\n\n        with open(self.image_list_path, 'r') as f:\n            images_file_lines = f.readlines()\n            last_class_id = -1\n\n            for i in tqdm(range(len(images_file_lines))):\n                line = images_file_lines[i]\n                terms = line.split('|')\n                if len(terms) < 3:\n                    continue # FD has failed on this imsage\n                path, landmarks, bbox = terms\n                image_id, _ = path.split('/')\n\n                if image_id in self.identities:\n                    self.identities[image_id].append(len(samples))\n                else:\n                    last_class_id += 1\n                    self.identities[image_id] = [len(samples)]\n\n                bbox = [max(int(coord), 0) for coord in bbox.strip().split(' ')]\n                landmarks = [float(coord) for coord in landmarks.strip().split(' ')]\n                assert len(bbox) == 4\n                assert len(landmarks) == 10\n                samples.append((osp.join(self.images_root_path, path).strip(),\n                                last_class_id, image_id, bbox, landmarks))\n\n        return samples\n\n    def get_weights(self):\n        \"\"\"Computes weights of the each identity in dataset according to frequency of it's occurance\"\"\"\n        weights = [0.]*len(self.all_samples_info)\n        for i, sample in enumerate(self.all_samples_info):\n            weights[i] = float(len(self.all_samples_info)) / len(self.identities[sample[2]])\n        return weights\n\n    def get_num_classes(self):\n        \"\"\"Returns total number of identities\"\"\"\n        return len(self.identities)\n\n    def __len__(self):\n        \"\"\"Returns total number of samples\"\"\"\n        return len(self.samples_info)\n\n    def __getitem__(self, idx):\n        \"\"\"Returns sample (image, class id, image id) by index\"\"\"\n        img = cv.imread(self.samples_info[idx][0], cv.IMREAD_COLOR)\n        bbox = self.samples_info[idx][-2]\n        landmarks = self.samples_info[idx][-1]\n\n        img = img[bbox[1]:bbox[1] + bbox[3], bbox[0]:bbox[0] + bbox[2]]\n        img = FivePointsAligner.align(img, landmarks, d_size=(200, 200), normalized=True, show=False)\n\n        if self.transform:\n            img = self.transform(img)\n\n        return {'img': img, 'label': self.samples_info[idx][1], 'instance': self.samples_info[idx][2]}\n"
  },
  {
    "path": "datasets/ndg.py",
    "content": "\"\"\"\n Copyright (c) 2018 Intel Corporation\n Licensed under the Apache License, Version 2.0 (the \"License\");\n you may not use this file except in compliance with the License.\n You may obtain a copy of the License at\n      http://www.apache.org/licenses/LICENSE-2.0\n Unless required by applicable law or agreed to in writing, software\n distributed under the License is distributed on an \"AS IS\" BASIS,\n WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n See the License for the specific language governing permissions and\n limitations under the License.\n\"\"\"\n\nimport os.path as osp\nimport json\nimport numpy as np\nfrom tqdm import tqdm\nfrom torch.utils.data import Dataset\nimport cv2 as cv\n\n\nclass NDG(Dataset):\n    \"\"\"NDG Dataset compatible with PyTorch DataLoader.\"\"\"\n    def __init__(self, images_root_path, annotation_list, transform=None, test=False):\n        self.test = test\n        self.have_landmarks = True\n        self.images_root_path = images_root_path\n        self.landmarks_file = open(annotation_list, 'r')\n        self.samples_info = self._read_samples_info()\n        self.transform = transform\n\n    def _read_samples_info(self):\n        \"\"\"Reads annotation of the dataset\"\"\"\n        samples = []\n        data = json.load(self.landmarks_file)\n\n        for image_info in tqdm(data):\n            img_name = image_info['path']\n            img_path = osp.join(self.images_root_path, img_name)\n            landmarks = image_info['lm']\n            samples.append((img_path, landmarks))\n\n        return samples\n\n    def __len__(self):\n        \"\"\"Returns total number of samples\"\"\"\n        return len(self.samples_info)\n\n    def __getitem__(self, idx):\n        \"\"\"Returns sample (image, landmarks) by index\"\"\"\n        img = cv.imread(self.samples_info[idx][0], cv.IMREAD_COLOR)\n        landmarks = self.samples_info[idx][1]\n        width, height = img.shape[1], img.shape[0]\n        landmarks = np.array([(float(landmarks[i][0]) / width,\n                               float(landmarks[i][1]) / height) for i in range(len(landmarks))]).reshape(-1)\n        data = {'img': img, 'landmarks': landmarks}\n        if self.transform:\n            data = self.transform(data)\n        return data\n"
  },
  {
    "path": "datasets/trillion_pairs.py",
    "content": "\"\"\"\n Copyright (c) 2018 Intel Corporation\n Licensed under the Apache License, Version 2.0 (the \"License\");\n you may not use this file except in compliance with the License.\n You may obtain a copy of the License at\n      http://www.apache.org/licenses/LICENSE-2.0\n Unless required by applicable law or agreed to in writing, software\n distributed under the License is distributed on an \"AS IS\" BASIS,\n WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n See the License for the specific language governing permissions and\n limitations under the License.\n\"\"\"\n\nimport os.path as osp\n\nimport cv2 as cv\nfrom tqdm import tqdm\nfrom torch.utils.data import Dataset\n\nfrom utils.face_align import FivePointsAligner\n\n\nclass TrillionPairs(Dataset):\n    \"\"\"TrillionPairs Dataset compatible with PyTorch DataLoader. For details visit http://trillionpairs.deepglint.com/data\"\"\"\n    def __init__(self, images_root_path, image_list_path, test_mode=False, transform=None):\n        self.image_list_path = image_list_path\n        self.images_root_path = images_root_path\n        self.test_mode = test_mode\n        self.identities = {}\n\n        assert osp.isfile(image_list_path)\n        self.have_landmarks = True\n\n        self.all_samples_info = self._read_samples_info()\n        self.samples_info = self.all_samples_info\n        self.transform = transform\n\n    def _read_samples_info(self):\n        \"\"\"Reads annotation of the dataset\"\"\"\n        samples = []\n\n        with open(self.image_list_path, 'r') as f:\n            images_file_lines = f.readlines()\n\n            for i in tqdm(range(len(images_file_lines))):\n                line = images_file_lines[i].strip()\n                terms = line.split(' ')\n                path = terms[0]\n                if not self.test_mode:\n                    label = int(terms[1])\n                    landmarks = terms[2:]\n                    if label in self.identities:\n                        self.identities[label].append(len(samples))\n                    else:\n                        self.identities[label] = [len(samples)]\n                else:\n                    label = 0\n                    landmarks = terms[1:]\n\n                landmarks = [float(coord) for coord in landmarks]\n                assert(len(landmarks) == 10)\n                samples.append((osp.join(self.images_root_path, path).strip(),\n                                label, landmarks))\n\n        return samples\n\n    def get_weights(self):\n        \"\"\"Computes weights of the each identity in dataset according to frequency of it's occurance\"\"\"\n        weights = [0.]*len(self.all_samples_info)\n        for i, sample in enumerate(self.all_samples_info):\n            weights[i] = float(len(self.all_samples_info)) / len(self.identities[sample[1]])\n        return weights\n\n    def get_num_classes(self):\n        \"\"\"Returns total number of identities\"\"\"\n        return len(self.identities)\n\n    def __len__(self):\n        \"\"\"Returns total number of samples\"\"\"\n        return len(self.samples_info)\n\n    def __getitem__(self, idx):\n        \"\"\"Returns sample (image, class id, image id) by index\"\"\"\n        img = cv.imread(self.samples_info[idx][0], cv.IMREAD_COLOR)\n        landmarks = self.samples_info[idx][-1]\n\n        img = FivePointsAligner.align(img, landmarks, d_size=(200, 200), normalized=False, show=False)\n\n        if self.transform:\n            img = self.transform(img)\n\n        return {'img': img, 'label': self.samples_info[idx][1], 'idx': idx}\n"
  },
  {
    "path": "datasets/vggface2.py",
    "content": "\"\"\"\n Copyright (c) 2018 Intel Corporation\n Licensed under the Apache License, Version 2.0 (the \"License\");\n you may not use this file except in compliance with the License.\n You may obtain a copy of the License at\n      http://www.apache.org/licenses/LICENSE-2.0\n Unless required by applicable law or agreed to in writing, software\n distributed under the License is distributed on an \"AS IS\" BASIS,\n WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n See the License for the specific language governing permissions and\n limitations under the License.\n\"\"\"\n\nimport os.path as osp\nimport cv2 as cv\nfrom tqdm import tqdm\nimport numpy as np\nfrom torch.utils.data import Dataset\n\nfrom utils.face_align import FivePointsAligner\n\n\nclass VGGFace2(Dataset):\n    \"\"\"VGGFace2 Dataset compatible with PyTorch DataLoader.\"\"\"\n    def __init__(self, images_root_path, image_list_path, landmarks_folder_path='',\n                 transform=None, landmarks_training=False):\n        self.image_list_path = image_list_path\n        self.images_root_path = images_root_path\n        self.identities = {}\n\n        self.landmarks_file = None\n        self.detections_file = None\n        if osp.isdir(landmarks_folder_path):\n            if 'train' in image_list_path:\n                bb_file_name = 'loose_landmark_train.csv'\n                landmarks_file_name = 'loose_bb_train.csv'\n            elif 'test' in image_list_path:\n                bb_file_name = 'loose_landmark_test.csv'\n                landmarks_file_name = 'loose_bb_test.csv'\n            else:\n                bb_file_name = 'loose_landmark_all.csv'\n                landmarks_file_name = 'loose_bb_all.csv'\n            self.landmarks_file = open(osp.join(landmarks_folder_path, bb_file_name), 'r')\n            self.detections_file = open(osp.join(landmarks_folder_path, landmarks_file_name), 'r')\n        self.have_landmarks = not self.landmarks_file is None\n        self.landmarks_training = landmarks_training\n        if self.landmarks_training:\n            assert self.have_landmarks is True\n\n        self.samples_info = self._read_samples_info()\n\n        self.transform = transform\n\n    def _read_samples_info(self):\n        \"\"\"Reads annotation of the dataset\"\"\"\n        samples = []\n\n        with open(self.image_list_path, 'r') as f:\n            last_class_id = -1\n            images_file_lines = f.readlines()\n\n            if self.have_landmarks:\n                detections_file_lines = self.detections_file.readlines()[1:]\n                landmarks_file_lines = self.landmarks_file.readlines()[1:]\n                assert len(detections_file_lines) == len(landmarks_file_lines)\n                assert len(images_file_lines) == len(detections_file_lines)\n\n            for i in tqdm(range(len(images_file_lines))):\n                sample = images_file_lines[i].strip()\n                sample_id = int(sample.split('/')[0][1:])\n                frame_id = int(sample.split('/')[1].split('_')[0])\n                if sample_id in self.identities:\n                    self.identities[sample_id].append(len(samples))\n                else:\n                    last_class_id += 1\n                    self.identities[sample_id] = [len(samples)]\n                if not self.have_landmarks:\n                    samples.append((osp.join(self.images_root_path, sample), last_class_id, frame_id))\n                else:\n                    _, bbox = detections_file_lines[i].split('\",')\n                    bbox = [max(int(coord), 0) for coord in bbox.split(',')]\n                    _, landmarks = landmarks_file_lines[i].split('\",')\n                    landmarks = [float(coord) for coord in landmarks.split(',')]\n                    samples.append((osp.join(self.images_root_path, sample), last_class_id, sample_id, bbox, landmarks))\n\n        return samples\n\n    def get_weights(self):\n        \"\"\"Computes weights of the each identity in dataset according to frequency of it's occurance\"\"\"\n        weights = [0.]*len(self.samples_info)\n        for i, sample in enumerate(self.samples_info):\n            weights[i] = len(self.samples_info) / float(len(self.identities[sample[2]]))\n\n        return weights\n\n    def get_num_classes(self):\n        \"\"\"Returns total number of identities\"\"\"\n        return len(self.identities)\n\n    def __len__(self):\n        \"\"\"Returns total number of samples\"\"\"\n        return len(self.samples_info)\n\n    def __getitem__(self, idx):\n        \"\"\"Returns sample (image, class id, image id) by index\"\"\"\n        img = cv.imread(self.samples_info[idx][0], cv.IMREAD_COLOR)\n        if self.landmarks_training:\n            landmarks = self.samples_info[idx][-1]\n            bbox = self.samples_info[idx][-2]\n            img = img[bbox[1]:bbox[1] + bbox[3], bbox[0]:bbox[0] + bbox[2]]\n            landmarks = [(float(landmarks[2*i]-bbox[0]) / bbox[2],\n                          float(landmarks[2*i + 1]-bbox[1])/ bbox[3]) for i in range(len(landmarks)//2)]\n            data = {'img': img, 'landmarks': np.array(landmarks)}\n            if self.transform:\n                data = self.transform(data)\n            return data\n\n        if self.have_landmarks:\n            landmarks = self.samples_info[idx][-1]\n            img = FivePointsAligner.align(img, landmarks, d_size=(200, 200), normalized=False)\n\n        if self.transform:\n            img = self.transform(img)\n\n        return {'img': img, 'label': self.samples_info[idx][1], 'instance': self.samples_info[idx][2]}\n"
  },
  {
    "path": "demo/README.md",
    "content": "# Face Recognition demo with [OpenVINO™ Toolkit](https://software.intel.com/en-us/openvino-toolkit)\n\n![](./demo.png)\n\n## Demo Preparation\n\n1. Install **OpenVINO Toolkit** - [Linux installation guide](https://software.intel.com/en-us/articles/OpenVINO-Install-Linux)\n\n2. Create virtual python environment:\n```bash\n  mkvirtualenv fr --python=python3\n```\n3. Install dependencies:\n```bash\n  pip install -r requirements.txt\n```\n4. Initialize OpenVINO environment:\n```bash\nsource /opt/intel/computer_vision_sdk/bin/setupvars.sh\n```\n\n## Deep Face Recognition\n1. Set up `PATH_TO_GALLERY` variable to point to folder with gallery images (faces to be recognized):\n```bash\nexport PATH_TO_GALLERY=/path/to/gallery/with/images/\n```\n2. For using OpenVINO pretrained models, please specify `IR_MODELS_ROOT`, otherwise you need to modify running command.\n```bash\nexport IR_MODELS_ROOT=$INTEL_CVSDK_DIR/deployment_tools/intel_models/\n```\n3. If you are running from pure console, you need to specify `PYTHONPATH` variable:\n```bash\nexport PYTHONPATH=`pwd`:$PYTHONPATH\n```\n4. Run Face Recognition demo:\n```bash\npython demo/run_demo.py --path_to_gallery $PATH_TO_GALLERY --cam_id 0 \\\n  --fd_model $IR_MODELS_ROOT/face-detection-retail-0004/FP32/face-detection-retail-0004.xml \\\n  --fr_model $IR_MODELS_ROOT/face-reidentification-retail-0095/FP32/face-reidentification-retail-0095.xml  \\\n  --ld_model $IR_MODELS_ROOT/landmarks-regression-retail-0009/FP32/landmarks-regression-retail-0009.xml \\\n  -l libcpu_extension_avx2.so\n```\n*Note:* `libcpu_extension_avx2.so` is located at the `$INTEL_CVSDK_DIR/inference_engine/lib/<system_name>/intel64/` folder.\nHere the `<system_name>` is a name detected by the OpenVINO. It can be for example `ubuntu_16.04` if you are running the demo under Ubuntu 16.04 system. The folder with CPU extensions is already in `LD_LIBRARY_PATH` after initialization of the OpenVINO environment, that's why it can be omitted in the launch command.\n"
  },
  {
    "path": "demo/run_demo.py",
    "content": "\"\"\"\n Copyright (c) 2018 Intel Corporation\n Licensed under the Apache License, Version 2.0 (the \"License\");\n you may not use this file except in compliance with the License.\n You may obtain a copy of the License at\n      http://www.apache.org/licenses/LICENSE-2.0\n Unless required by applicable law or agreed to in writing, software\n distributed under the License is distributed on an \"AS IS\" BASIS,\n WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n See the License for the specific language governing permissions and\n limitations under the License.\n\"\"\"\n\nimport argparse\nimport os\nimport os.path as osp\n\nimport glog as log\nimport cv2 as cv\nimport numpy as np\nfrom scipy.spatial.distance import cosine\n\nfrom utils import face_align\nfrom utils.ie_tools import load_ie_model\n\n\nclass FaceDetector:\n    \"\"\"Wrapper class for face detector\"\"\"\n    def __init__(self, model_path, conf=.6, device='CPU', ext_path=''):\n        self.net = load_ie_model(model_path, device, None, ext_path)\n        self.confidence = conf\n        self.expand_ratio = (1.1, 1.05)\n\n    def get_detections(self, frame):\n        \"\"\"Returns all detections on frame\"\"\"\n        _, _, h, w = self.net.get_input_shape().shape\n        out = self.net.forward(cv.resize(frame, (w, h)))\n        detections = self.__decode_detections(out, frame.shape)\n        return detections\n\n    def __decode_detections(self, out, frame_shape):\n        \"\"\"Decodes raw SSD output\"\"\"\n        detections = []\n\n        for detection in out[0, 0]:\n            confidence = detection[2]\n            if confidence > self.confidence:\n                left = int(max(detection[3], 0) * frame_shape[1])\n                top = int(max(detection[4], 0) * frame_shape[0])\n                right = int(max(detection[5], 0) * frame_shape[1])\n                bottom = int(max(detection[6], 0) * frame_shape[0])\n                if self.expand_ratio != (1., 1.):\n                    w = (right - left)\n                    h = (bottom - top)\n                    dw = w * (self.expand_ratio[0] - 1.) / 2\n                    dh = h * (self.expand_ratio[1] - 1.) / 2\n                    left = max(int(left - dw), 0)\n                    right = int(right + dw)\n                    top = max(int(top - dh), 0)\n                    bottom = int(bottom + dh)\n\n                detections.append(((left, top, right, bottom), confidence))\n\n        if len(detections) > 1:\n            detections.sort(key=lambda x: x[1], reverse=True)\n\n        return detections\n\n\nclass VectorCNN:\n    \"\"\"Wrapper class for a nework returning a vector\"\"\"\n    def __init__(self, model_path, device='CPU'):\n        self.net = load_ie_model(model_path, device, None)\n\n    def forward(self, batch):\n        \"\"\"Performs forward of the underlying network on a given batch\"\"\"\n        _, _, h, w = self.net.get_input_shape().shape\n        outputs = [self.net.forward(cv.resize(frame, (w, h))) for frame in batch]\n        return outputs\n\n\ndef get_embeddings(frame, detections, face_reid, landmarks_predictor):\n    \"\"\"Get embeddings for all detected faces on the frame\"\"\"\n    rois = []\n    embeddings = []\n    for rect, _ in detections:\n        left, top, right, bottom = rect\n        rois.append(frame[top:bottom, left:right])\n\n    if rois:\n        landmarks = landmarks_predictor.forward(rois)\n        assert len(landmarks) == len(rois)\n\n        for i, _ in enumerate(rois):\n            roi_keypoints = landmarks[i].reshape(-1)\n            rois[i] = face_align.FivePointsAligner.align(rois[i], roi_keypoints,\n                                                         d_size=(rois[i].shape[1], rois[i].shape[0]),\n                                                         normalized=True, show=False)\n        embeddings = face_reid.forward(rois)\n        assert len(rois) == len(embeddings)\n\n    return embeddings\n\n\ndef find_nearest(x, gallery, thr):\n    \"\"\"Finds the nearest to a given embedding in the gallery\"\"\"\n    if gallery:\n        diffs = np.array([cosine(x, y) for y in gallery.values()])\n        min_pos = diffs.argmin()\n        min_dist = diffs[min_pos]\n        if min_dist < thr:\n            return min_pos, list(gallery.keys())[min_pos]\n    return None, None\n\n\ndef match_embeddings(embeds, gallery, thr):\n    \"\"\"Matches input embeddings with ones in the gallery\"\"\"\n    indexes = []\n    for emb in embeds:\n        _, name = find_nearest(emb, gallery, thr)\n        if name is not None:\n            indexes.append(name)\n        else:\n            indexes.append('Unknown')\n\n    return indexes, gallery\n\n\ndef draw_detections(frame, detections, indexes):\n    \"\"\"Draws detections and labels\"\"\"\n    for i, rect in enumerate(detections):\n        left, top, right, bottom = rect[0]\n        cv.rectangle(frame, (left, top), (right, bottom), (0, 255, 0), thickness=2)\n        label = str(indexes[i])\n        label_size, base_line = cv.getTextSize(label, cv.FONT_HERSHEY_SIMPLEX, 1, 1)\n        top = max(top, label_size[1])\n        cv.rectangle(frame, (left, top - label_size[1]), (left + label_size[0], top + base_line),\n                     (255, 255, 255), cv.FILLED)\n        cv.putText(frame, label, (left, top), cv.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 0))\n\n    return frame\n\n\ndef load_gallery(path_to_gallery, face_det, landmarks_detector, face_recognizer):\n    \"\"\"Computes embeddings for gallery\"\"\"\n    gallery = {}\n    files = os.listdir(path_to_gallery)\n    files = [file for file in files if file.endswith('.png') or file.endswith('.jpg')]\n    for file in files:\n        img = cv.imread(osp.join(path_to_gallery, file))\n        detections = face_det.get_detections(img)\n\n        if not detections:\n            detections = [[0, 0, img.shape[0], img.shape[1]], 0]\n            log.warn('Warning: failed to detect face on the image ' + file)\n\n        embed = get_embeddings(img, detections, face_recognizer, landmarks_detector)\n        gallery[file.replace('.png', '').replace('.jpg', '')] = embed[0]\n    return gallery\n\n\ndef run(params, capture, face_det, face_recognizer, landmarks_detector):\n    \"\"\"Starts the face recognition demo\"\"\"\n    win_name = 'Deep Face Recognition'\n    gallery = load_gallery(params.path_to_gallery, face_det, landmarks_detector, face_recognizer)\n\n    while cv.waitKey(1) != 27:\n        has_frame, frame = capture.read()\n        if not has_frame:\n            return\n\n        detections = face_det.get_detections(frame)\n        embeds = get_embeddings(frame, detections, face_recognizer, landmarks_detector)\n        ids, gallery = match_embeddings(embeds, gallery, params.fr_thresh)\n        frame = draw_detections(frame, detections, ids)\n        cv.imshow(win_name, frame)\n\ndef main():\n    \"\"\"Prepares data for the face recognition demo\"\"\"\n    parser = argparse.ArgumentParser(description='Face recognition live demo script')\n    parser.add_argument('--video', type=str, default=None, help='Input video')\n    parser.add_argument('--cam_id', type=int, default=-1, help='Input cam')\n\n    parser.add_argument('--fd_model', type=str, required=True)\n    parser.add_argument('--fd_thresh', type=float, default=0.6, help='Threshold for FD')\n\n    parser.add_argument('--fr_model', type=str, required=True)\n    parser.add_argument('--fr_thresh', type=float, default=0.6, help='Threshold for FR')\n\n    parser.add_argument('--path_to_gallery', type=str, required=True, help='Path to gallery with subjects')\n\n    parser.add_argument('--ld_model', type=str, default='', help='Path to a snapshots with landmarks detection model')\n\n    parser.add_argument('--device', type=str, default='CPU')\n    parser.add_argument('-l', '--cpu_extension',\n                        help='MKLDNN (CPU)-targeted custom layers.Absolute path to a shared library with the kernels '\n                             'impl.', type=str, default=None)\n\n    args = parser.parse_args()\n\n    if args.cam_id >= 0:\n        log.info('Reading from cam {}'.format(args.cam_id))\n        cap = cv.VideoCapture(args.cam_id)\n        cap.set(cv.CAP_PROP_FRAME_WIDTH, 1280)\n        cap.set(cv.CAP_PROP_FRAME_HEIGHT, 720)\n        cap.set(cv.CAP_PROP_FOURCC, cv.VideoWriter_fourcc('M', 'J', 'P', 'G'))\n    else:\n        assert args.video\n        log.info('Reading from {}'.format(args.video))\n        cap = cv.VideoCapture(args.video)\n    assert cap.isOpened()\n\n    face_detector = FaceDetector(args.fd_model, args.fd_thresh, args.device, args.cpu_extension)\n    face_recognizer = VectorCNN(args.fr_model, args.device)\n    landmarks_detector = VectorCNN(args.ld_model, args.device)\n    run(args, cap, face_detector, face_recognizer, landmarks_detector)\n\nif __name__ == '__main__':\n    main()\n"
  },
  {
    "path": "devtools/pylint.rc",
    "content": "[MASTER]\n\n# Specify a configuration file.\n#rcfile=\n\n# Python code to execute, usually for sys.path manipulation such as\n# pygtk.require().\n#init-hook=\n\n# Profiled execution.\nprofile=no\n\n# Add <file or directory> to the black list. It should be a base name, not a\n# path. You may set this option multiple times.\nignore=CVS\n\n# Pickle collected data for later comparisons.\npersistent=yes\n\n# List of plugins (as comma separated values of python modules names) to load,\n# usually to register additional checkers.\nload-plugins=\n\n\n[MESSAGES CONTROL]\n\n# Enable the message, report, category or checker with the given id(s). You can\n# either give multiple identifier separated by comma (,) or put this option\n# multiple time.\n#enable=\n\n# Disable the message, report, category or checker with the given id(s). You\n# can either give multiple identifier separated by comma (,) or put this option\n# multiple time.\ndisable=R0903, W0221\n\n\n[REPORTS]\n\n# Set the output format. Available formats are text, parseable, colorized, msvs\n# (visual studio) and html\noutput-format=text\n\n# Include message's id in output\ninclude-ids=no\n\n# Put messages in a separate file for each module / package specified on the\n# command line instead of printing them on stdout. Reports (if any) will be\n# written in a file name \"pylint_global.[txt|html]\".\nfiles-output=no\n\n# Tells whether to display a full report or only the messages\nreports=yes\n\n# Python expression which should return a note less than 10 (10 is the highest\n# note). You have access to the variables errors warning, statement which\n# respectively contain the number of errors / warnings messages and the total\n# number of statements analyzed. This is used by the global evaluation report\n# (R0004).\nevaluation=10.0 - ((float(5 * error + warning + refactor + convention) / statement) * 10)\n\n# Add a comment according to your evaluation note. This is used by the global\n# evaluation report (R0004).\ncomment=no\n\n\n[VARIABLES]\n\n# Tells whether we should check for unused import in __init__ files.\ninit-import=no\n\n# A regular expression matching names used for dummy variables (i.e. not used).\ndummy-variables-rgx=_|dummy\n\n# List of additional names supposed to be defined in builtins. Remember that\n# you should avoid to define new builtins when possible.\nadditional-builtins=\n\n\n[BASIC]\n\n# Required attributes for module, separated by a comma\nrequired-attributes=\n\n# List of builtins function names that should not be used, separated by a comma\nbad-functions=map,filter,apply,input\n\n# Regular expression which should only match correct module names\nmodule-rgx=(([a-z_][a-z0-9_]*)|([A-Z][a-zA-Z0-9]+))$\n\n# Regular expression which should only match correct module level names\nconst-rgx=(([A-Z_][A-Z0-9_]*)|(__.*__))$\n\n# Regular expression which should only match correct class names\nclass-rgx=[A-Z_][a-zA-Z0-9]+$\n\n# Regular expression which should only match correct function names\nfunction-rgx=[a-z_][a-z0-9_]{2,40}$\n\n# Regular expression which should only match correct method names\nmethod-rgx=[a-z_][a-z0-9_]{2,30}$\n\n# Regular expression which should only match correct instance attribute names\nattr-rgx=[a-z_][a-z0-9_]{0,30}$\n\n# Regular expression which should only match correct argument names\nargument-rgx=[a-z_][a-z0-9_]{0,30}$\n\n# Regular expression which should only match correct variable names\nvariable-rgx=[a-z_][a-z0-9_]{0,30}$\n\n# Regular expression which should only match correct list comprehension /\n# generator expression variable names\ninlinevar-rgx=[A-Za-z_][A-Za-z0-9_]*$\n\n# Good variable names which should always be accepted, separated by a comma\ngood-names=i,j,k,ex,Run,_\n\n# Bad variable names which should always be refused, separated by a comma\nbad-names=foo,bar,baz,toto,tutu,tata\n\n# Regular expression which should only match functions or classes name which do\n# not require a docstring\nno-docstring-rgx=__.*__\n\n\n[MISCELLANEOUS]\n\n# List of note tags to take in consideration, separated by a comma.\nnotes=FIXME,XXX,TODO\n\n\n[FORMAT]\n\n# Maximum number of characters on a single line.\nmax-line-length=120\n\n# Maximum number of lines in a module\nmax-module-lines=1000\n\n# String used as indentation unit. This is usually \" \" (4 spaces) or \"\\t\" (1\n# tab).\nindent-string='    '\nindent-after-paren=4\n\n\n[SIMILARITIES]\n\n# Minimum lines number of a similarity.\nmin-similarity-lines=4\n\n# Ignore comments when computing similarities.\nignore-comments=yes\n\n# Ignore docstrings when computing similarities.\nignore-docstrings=yes\n\n\n[TYPECHECK]\n\n# Tells whether missing members accessed in mixin class should be ignored. A\n# mixin class is detected if its name ends with \"mixin\" (case insensitive).\nignore-mixin-members=yes\n\n# List of classes names for which member attributes should not be checked\n# (useful for classes with attributes dynamically set).\nignored-classes=SQLObject\n\n# When zope mode is activated, add a predefined set of Zope acquired attributes\n# to generated-members.\nzope=no\n\n# List of members which are set dynamically and missed by pylint inference\n# system, and so shouldn't trigger E0201 when accessed.\ngenerated-members=REQUEST,acl_users,aq_parent,torch,cv\n\n\n[DESIGN]\n\n# Maximum number of arguments for function / method\nmax-args=5\n\n# Argument names that match this expression will be ignored. Default to name\n# with leading underscore\nignored-argument-names=_.*\n\n# Maximum number of locals for function / method body\nmax-locals=15\n\n# Maximum number of return / yield for function / method body\nmax-returns=6\n\n# Maximum number of branch for function / method body\nmax-branchs=12\n\n# Maximum number of statements in function / method body\nmax-statements=50\n\n# Maximum number of parents for a class (see R0901).\nmax-parents=7\n\n# Maximum number of attributes for a class (see R0902).\nmax-attributes=7\n\n# Minimum number of public methods for a class (see R0903).\nmin-public-methods=2\n\n# Maximum number of public methods for a class (see R0904).\nmax-public-methods=20\n\n\n[IMPORTS]\n\n# Deprecated modules which should not be used, separated by a comma\ndeprecated-modules=regsub,string,TERMIOS,Bastion,rexec\n\n# Create a graph of every (i.e. internal and external) dependencies in the\n# given file (report RP0402 must not be disabled)\nimport-graph=\n\n# Create a graph of external dependencies in the given file (report RP0402 must\n# not be disabled)\next-import-graph=\n\n# Create a graph of internal dependencies in the given file (report RP0402 must\n# not be disabled)\nint-import-graph=\n\nextension-pkg-whitelist=cv2\n\n[CLASSES]\n\n# List of interface methods to ignore, separated by a comma. This is used for\n# instance to not check methods defines in Zope's Interface base class.\nignore-iface-methods=isImplementedBy,deferred,extends,names,namesAndDescriptions,queryDescriptionFor,getBases,getDescriptionFor,getDoc,getName,getTaggedValue,getTaggedValueTags,isEqualOrExtendedBy,setTaggedValue,isImplementedByInstancesOf,adaptWith,is_implemented_by\n\n# List of method names used to declare (i.e. assign) instance attributes.\ndefining-attr-methods=__init__,__new__,setUp\n"
  },
  {
    "path": "dump_features.py",
    "content": "\"\"\"\n Copyright (c) 2018 Intel Corporation\n Licensed under the Apache License, Version 2.0 (the \"License\");\n you may not use this file except in compliance with the License.\n You may obtain a copy of the License at\n      http://www.apache.org/licenses/LICENSE-2.0\n Unless required by applicable law or agreed to in writing, software\n distributed under the License is distributed on an \"AS IS\" BASIS,\n WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n See the License for the specific language governing permissions and\n limitations under the License.\n\"\"\"\n\nimport sys\nimport argparse\nimport os\nimport os.path as osp\n\nfrom tqdm import tqdm\nimport numpy as np\nimport glog as log\nimport torch\nimport torch.nn.functional as F\nfrom torch.utils.data import DataLoader\nfrom torchvision import transforms as t\n\nfrom scripts.matio import save_mat\nfrom model.common import models_backbones\nfrom datasets.megaface import MegaFace\nfrom datasets.trillion_pairs import TrillionPairs\nfrom utils.utils import load_model_state\nfrom utils.augmentation import ResizeNumpy, NumpyToTensor\n\n\ndef clean_megaface(filenames, features, noises_list_path):\n    \"\"\"Filters megaface from outliers\"\"\"\n    with open(noises_list_path, 'r') as f:\n        noises_list = f.readlines()\n        noises_list = [line.strip() for line in noises_list]\n    clean_features = np.zeros((features.shape[0], features.shape[1] + 1), dtype=np.float32)\n\n    for i, filename in enumerate(tqdm(filenames)):\n        clean_features[i, 0: features.shape[1]] = features[i, :]\n        for line in noises_list:\n            if line in filename:\n                clean_features[i, features.shape[1]] = 100.0\n                break\n\n    return clean_features\n\n\ndef clean_facescrub(filenames, features, noises_list_path):\n    \"\"\"Replaces wrong instances of identities from the Facescrub with the centroids of these identities\"\"\"\n    clean_feature_size = features.shape[1] + 1\n    with open(noises_list_path, 'r') as f:\n        noises_list = f.readlines()\n        noises_list = [osp.splitext(line.strip())[0] for line in noises_list]\n    clean_features = np.zeros((features.shape[0], clean_feature_size), dtype=np.float32)\n\n    centroids = {}\n    for i, filename in enumerate(tqdm(filenames)):\n        clean_features[i, 0: features.shape[1]] = features[i, :]\n        id_name = osp.basename(filename).split('_')[0]\n        if not id_name in centroids:\n            centroids[id_name] = np.zeros(clean_feature_size, dtype=np.float32)\n        centroids[id_name] += clean_features[i, :]\n\n    for i, file_path in enumerate(tqdm(filenames)):\n        filename = osp.basename(file_path)\n        for line in noises_list:\n            if line in filename.replace(' ', '_'):\n                id_name = filename.split('_')[0]\n                clean_features[i, :] = centroids[id_name] + np.random.uniform(-0.001, 0.001, clean_feature_size)\n                clean_features[i, :] /= np.linalg.norm(clean_features[i, :])\n                break\n\n    return clean_features\n\n\n@torch.no_grad()\ndef main(args):\n    input_filenames = []\n    output_filenames = []\n    input_dir = os.path.abspath(args.input_dir)\n    output_dir = os.path.abspath(args.output_dir)\n\n    if not args.trillion_format:\n        log.info('Reading info...')\n        with open(os.path.join(args.input_dir, os.path.basename(args.input_list)), 'r') as f:\n            lines = f.readlines()\n\n            for line in tqdm(lines):\n                info = line.strip().split('|')\n                file = info[0].strip()\n                filename = os.path.join(input_dir, file)\n\n                path, _ = osp.split(filename)\n                out_folder = path.replace(input_dir, output_dir)\n                if not osp.isdir(out_folder):\n                    os.makedirs(out_folder)\n\n                landmarks = None\n                bbox = None\n\n                if len(info) > 2:\n                    landmarks = info[1].strip().split(' ')\n                    landmarks = [float(x) for x in landmarks]\n                    bbox = info[2].strip().split(' ')\n                    bbox = [int(float(x)) for x in bbox]\n                outname = filename.replace(input_dir, output_dir) + args.file_ending\n                input_filenames.append({'path': filename, 'landmarks': landmarks, 'bbox': bbox})\n                output_filenames += [outname]\n\n        nrof_images = len(input_filenames)\n        log.info(\"Total number of images: \", nrof_images)\n        dataset = MegaFace(input_filenames)\n    else:\n        dataset = TrillionPairs(args.input_dir, osp.join(args.input_dir, 'testdata_lmk.txt'), test_mode=True)\n        nrof_images = len(dataset)\n\n    emb_array = np.zeros((nrof_images, args.embedding_size), dtype=np.float32)\n\n    dataset.transform = t.Compose([ResizeNumpy(models_backbones[args.model].get_input_res()),\n                                   NumpyToTensor(switch_rb=True)])\n    val_loader = DataLoader(dataset, batch_size=args.batch_size, num_workers=5, shuffle=False)\n\n    model = models_backbones[args.model](embedding_size=args.embedding_size, feature=True)\n    assert args.snap is not None\n    log.info('Snapshot ' + args.snap + ' ...')\n    log.info('Extracting embeddings ...')\n    model = load_model_state(model, args.snap, args.devices[0], eval_state=True)\n    model = torch.nn.DataParallel(model, device_ids=args.devices, output_device=args.devices[0])\n\n    f_output_filenames = []\n\n    with torch.cuda.device(args.devices[0]):\n        for i, data in enumerate(tqdm(val_loader), 0):\n            idxs, imgs = data['idx'], data['img']\n            batch_embeddings = F.normalize(model(imgs), p=2, dim=1).data.cpu().numpy()\n            batch_embeddings = batch_embeddings.reshape(batch_embeddings.shape[0], -1)\n            path_indices = idxs.data.cpu().numpy()\n\n            start_index = i*args.batch_size\n            end_index = min((i+1)*args.batch_size, nrof_images)\n            assert start_index == path_indices[0]\n            assert end_index == path_indices[-1] + 1\n            assert emb_array[start_index:end_index, :].shape == batch_embeddings.shape\n            emb_array[start_index:end_index, :] = batch_embeddings\n\n            if not args.trillion_format:\n                for index in path_indices:\n                    f_output_filenames.append(output_filenames[index])\n\n    assert len(output_filenames) == len(output_filenames)\n\n    log.info('Extracting features Done.')\n\n    if args.trillion_format:\n        save_mat(args.file_ending, emb_array)\n    else:\n        if 'megaface_noises.txt' in args.noises_list:\n            log.info('Cleaning Megaface features')\n            emb_array = clean_megaface(f_output_filenames, emb_array, args.noises_list)\n        elif 'facescrub_noises.txt' in args.noises_list:\n            log.info('Cleaning Facescrub features')\n            emb_array = clean_facescrub(f_output_filenames, emb_array, args.noises_list)\n        else:\n            log.info('Megaface features are not cleaned up.')\n\n        log.info('Saving features to files...')\n        for i in tqdm(range(len(f_output_filenames))):\n            save_mat(f_output_filenames[i], emb_array[i, :])\n\n\ndef parse_argument(argv):\n    parser = argparse.ArgumentParser(description='Save embeddings to MegaFace features files')\n    parser.add_argument('--model', choices=models_backbones.keys(), type=str, default='rmnet', help='Model type.')\n    parser.add_argument('input_dir', help='Path to MegaFace Features')\n    parser.add_argument('output_dir', help='Path to FaceScrub Features')\n    parser.add_argument('--input_list', default='list.txt', type=str, required=False)\n    parser.add_argument('--batch_size', type=int, default=128)\n    parser.add_argument('--embedding_size', type=int, default=128)\n    parser.add_argument('--devices', type=int, nargs='+', default=[0], help='CUDA devices to use.')\n    parser.add_argument('--snap', type=str, required=True, help='Snapshot to evaluate.')\n    parser.add_argument('--noises_list', type=str, default='', required=False, help='A list of the Megaface or Facescrub noises produced by insightface. \\\n                                                                        See https://github.com/deepinsight/insightface/blob/master/src/megaface/README.md')\n    parser.add_argument('--file_ending', help='Ending appended to original photo files. i.e.\\\n                        11084833664_0.jpg_LBP_100x100.bin => _LBP_100x100.bin', default='_rmnet.bin')\n    parser.add_argument('--trillion_format', action='store_true')\n    return parser.parse_args(argv)\n\nif __name__ == '__main__':\n    main(parse_argument(sys.argv[1:]))\n"
  },
  {
    "path": "evaluate_landmarks.py",
    "content": "\"\"\"\n Copyright (c) 2018 Intel Corporation\n Licensed under the Apache License, Version 2.0 (the \"License\");\n you may not use this file except in compliance with the License.\n You may obtain a copy of the License at\n      http://www.apache.org/licenses/LICENSE-2.0\n Unless required by applicable law or agreed to in writing, software\n distributed under the License is distributed on an \"AS IS\" BASIS,\n WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n See the License for the specific language governing permissions and\n limitations under the License.\n\"\"\"\n\nimport argparse\n\nimport glog as log\nimport torch\nimport torch.backends.cudnn as cudnn\nfrom torch.utils.data import DataLoader\nfrom torchvision.transforms import transforms as t\nfrom tqdm import tqdm\n\nfrom datasets import VGGFace2, CelebA, NDG\n\nfrom model.common import models_landmarks\nfrom utils.landmarks_augmentation import Rescale, ToTensor\nfrom utils.utils import load_model_state\n\n\ndef evaluate(val_loader, model):\n    \"\"\"Calculates average error\"\"\"\n    total_loss = 0.\n    total_pp_error = 0.\n    failures_num = 0\n    items_num = 0\n    for _, data in enumerate(tqdm(val_loader), 0):\n        data, gt_landmarks = data['img'].cuda(), data['landmarks'].cuda()\n        predicted_landmarks = model(data)\n        loss = predicted_landmarks - gt_landmarks\n        items_num += loss.shape[0]\n        n_points = loss.shape[1] // 2\n        per_point_error = loss.data.view(-1, n_points, 2)\n        per_point_error = torch.norm(per_point_error, p=2, dim=2)\n        avg_error = torch.sum(per_point_error, 1) / n_points\n        eyes_dist = torch.norm(gt_landmarks[:, 0:2] - gt_landmarks[:, 2:4], p=2, dim=1).reshape(-1)\n\n        per_point_error = torch.div(per_point_error, eyes_dist.view(-1, 1))\n        total_pp_error += torch.sum(per_point_error, 0)\n\n        avg_error = torch.div(avg_error, eyes_dist)\n        failures_num += torch.nonzero(avg_error > 0.1).shape[0]\n        total_loss += torch.sum(avg_error)\n\n    return total_loss / items_num, (total_pp_error / items_num).data.cpu().numpy(), float(failures_num) / items_num\n\n\ndef start_evaluation(args):\n    \"\"\"Launches the evaluation process\"\"\"\n\n    if args.dataset == 'vgg':\n        dataset = VGGFace2(args.val, args.v_list, args.v_land, landmarks_training=True)\n    elif args.dataset == 'celeb':\n        dataset = CelebA(args.val, args.v_land, test=True)\n    else:\n        dataset = NDG(args.val, args.v_land)\n\n    if dataset.have_landmarks:\n        log.info('Use alignment for the train data')\n        dataset.transform = t.Compose([Rescale((48, 48)), ToTensor(switch_rb=True)])\n    else:\n        exit()\n\n    val_loader = DataLoader(dataset, batch_size=args.val_batch_size, num_workers=4, shuffle=False, pin_memory=True)\n\n    model = models_landmarks['landnet']\n    assert args.snapshot is not None\n    log.info('Testing snapshot ' + args.snapshot + ' ...')\n    model = load_model_state(model, args.snapshot, args.device, eval_state=True)\n    model.eval()\n    cudnn.benchmark = True\n    model = torch.nn.DataParallel(model, device_ids=[args.device], )\n\n    log.info('Face landmarks model:')\n    log.info(model)\n\n    avg_err, per_point_avg_err, failures_rate = evaluate(val_loader, model)\n\n    log.info('Avg RMSE error: {}'.format(avg_err))\n    log.info('Per landmark RMSE error: {}'.format(per_point_avg_err))\n    log.info('Failure rate: {}'.format(failures_rate))\n\n\ndef main():\n    \"\"\"Creates a cl parser\"\"\"\n    parser = argparse.ArgumentParser(description='Evaluation script for landmarks detection network')\n    parser.add_argument('--device', '-d', default=0, type=int)\n    parser.add_argument('--val_data_root', dest='val', required=True, type=str, help='Path to val data.')\n    parser.add_argument('--val_list', dest='v_list', required=False, type=str, help='Path to test data image list.')\n    parser.add_argument('--val_landmarks', dest='v_land', default='', required=False, type=str,\n                        help='Path to landmarks for test images.')\n    parser.add_argument('--val_batch_size', type=int, default=1, help='Validation batch size.')\n    parser.add_argument('--snapshot', type=str, default=None, help='Snapshot to evaluate.')\n    parser.add_argument('--dataset', choices=['vgg', 'celeb', 'ngd'], type=str, default='vgg', help='Dataset.')\n    arguments = parser.parse_args()\n\n    with torch.cuda.device(arguments.device):\n        start_evaluation(arguments)\n\nif __name__ == '__main__':\n    main()\n"
  },
  {
    "path": "evaluate_lfw.py",
    "content": "\"\"\"\n Copyright (c) 2018 Intel Corporation\n Licensed under the Apache License, Version 2.0 (the \"License\");\n you may not use this file except in compliance with the License.\n You may obtain a copy of the License at\n      http://www.apache.org/licenses/LICENSE-2.0\n Unless required by applicable law or agreed to in writing, software\n distributed under the License is distributed on an \"AS IS\" BASIS,\n WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n See the License for the specific language governing permissions and\n limitations under the License.\n\"\"\"\n\nimport argparse\nimport datetime\nfrom functools import partial\n\nimport cv2 as cv\nimport torch\nimport torch.nn.functional as F\nfrom torch.utils.data import DataLoader\nfrom torchvision import transforms as t\n\nfrom scipy.spatial.distance import cosine\nimport glog as log\nfrom tqdm import tqdm\nimport numpy as np\nfrom tensorboardX import SummaryWriter\n\nfrom datasets.lfw import LFW\nfrom utils.utils import load_model_state, get_model_parameters_number, flip_tensor\nfrom utils.augmentation import ResizeNumpy, CenterCropNumpy, NumpyToTensor\nfrom utils.face_align import FivePointsAligner\nfrom model.common import models_backbones\n\n\ndef get_subset(container, subset_bounds):\n    \"\"\"Returns a subset of the given list with respect to the list of bounds\"\"\"\n    subset = []\n    for bound in subset_bounds:\n        subset += container[bound[0]: bound[1]]\n    return subset\n\n\ndef get_roc(scores_with_gt, n_threshs=400):\n    \"\"\"Computes a ROC cureve on the LFW dataset\"\"\"\n    thresholds = np.linspace(0., 4., n_threshs)\n\n    fp_rates = []\n    tp_rates = []\n\n    for threshold in thresholds:\n        fp = 0\n        tp = 0\n        for score_with_gt in scores_with_gt:\n            predict_same = score_with_gt['score'] < threshold\n            actual_same = score_with_gt['is_same']\n\n            if predict_same and actual_same:\n                tp += 1\n            elif predict_same and not actual_same:\n                fp += 1\n\n        fp_rates.append(float(fp) / len(scores_with_gt) * 2)\n        tp_rates.append(float(tp) / len(scores_with_gt) * 2)\n\n    return np.array(fp_rates), np.array(tp_rates)\n\n\ndef get_auc(fprs, tprs):\n    \"\"\"Computes AUC under a ROC curve\"\"\"\n    sorted_fprs, sorted_tprs = zip(*sorted(zip(*(fprs, tprs))))\n    sorted_fprs = list(sorted_fprs)\n    sorted_tprs = list(sorted_tprs)\n    if sorted_fprs[-1] != 1.0:\n        sorted_fprs.append(1.0)\n        sorted_tprs.append(sorted_tprs[-1])\n    return np.trapz(sorted_tprs, sorted_fprs)\n\n\ndef save_roc(fp_rates, tp_rates, fname):\n    assert fp_rates.shape[0] == tp_rates.shape[0]\n    with open(fname + '.txt', 'w') as f:\n        for i in range(fp_rates.shape[0]):\n            f.write('{} {}\\n'.format(fp_rates[i], tp_rates[i]))\n\n\n@torch.no_grad()\ndef compute_embeddings_lfw(args, dataset, model, batch_size, dump_embeddings=False,\n                           pdist=lambda x, y: 1. - F.cosine_similarity(x, y), flipped_embeddings=False):\n    \"\"\"Computes embeddings of all images from the LFW dataset using PyTorch\"\"\"\n    val_loader = DataLoader(dataset, batch_size=batch_size, num_workers=4, shuffle=False)\n    scores_with_gt = []\n    embeddings = []\n    ids = []\n\n    for batch_idx, data in enumerate(tqdm(val_loader, 'Computing embeddings')):\n        images_1 = data['img1']\n        images_2 = data['img2']\n        is_same = data['is_same']\n        if torch.cuda.is_available() and args.devices[0] != -1:\n            images_1 = images_1.cuda()\n            images_2 = images_2.cuda()\n        emb_1 = model(images_1)\n        emb_2 = model(images_2)\n        if flipped_embeddings:\n            images_1_flipped = flip_tensor(images_1, 3)\n            images_2_flipped = flip_tensor(images_2, 3)\n            emb_1_flipped = model(images_1_flipped)\n            emb_2_flipped = model(images_2_flipped)\n            emb_1 = (emb_1 + emb_1_flipped)*.5\n            emb_2 = (emb_2 + emb_2_flipped)*.5\n        scores = pdist(emb_1, emb_2).data.cpu().numpy()\n\n        for i, _ in enumerate(scores):\n            scores_with_gt.append({'score': scores[i], 'is_same': is_same[i], 'idx': batch_idx*batch_size + i})\n\n        if dump_embeddings:\n            id0 = data['id0']\n            id1 = data['id1']\n            ids.append(id0)\n            ids.append(id1)\n            to_dump_1 = emb_1.data.cpu()\n            to_dump_2 = emb_2.data.cpu()\n            embeddings.append(to_dump_1)\n            embeddings.append(to_dump_2)\n\n    if dump_embeddings:\n        total_emb = np.concatenate(embeddings, axis=0)\n        total_ids = np.concatenate(ids, axis=0)\n        log_path = './logs/{:%Y_%m_%d_%H_%M}'.format(datetime.datetime.now())\n        writer = SummaryWriter(log_path)\n        writer.add_embedding(torch.from_numpy(total_emb), total_ids)\n\n    return scores_with_gt\n\n\ndef compute_embeddings_lfw_ie(args, dataset, model, batch_size=1, dump_embeddings=False,\n                              pdist=cosine, flipped_embeddings=False, lm_model=None):\n    \"\"\"Computes embeddings of all images from the LFW dataset using Inference Engine\"\"\"\n    assert batch_size == 1\n    scores_with_gt = []\n\n    for batch_idx, data in enumerate(tqdm(dataset, 'Computing embeddings')):\n        images_1 = data['img1']\n        images_2 = data['img2']\n        if lm_model:\n            lm_input_size = tuple(lm_model.get_input_shape()[2:])\n            landmarks_1 = lm_model.forward(cv.resize(images_1, lm_input_size)).reshape(-1)\n            images_1 = FivePointsAligner.align(images_1, landmarks_1, *images_1.shape[:2], normalize=False, show=False)\n\n            landmarks_2 = lm_model.forward(cv.resize(images_2, lm_input_size)).reshape(-1)\n            images_2 = FivePointsAligner.align(images_2, landmarks_2, *images_2.shape[:2], normalize=False)\n\n        is_same = data['is_same']\n        emb_1 = model.forward(images_1).reshape(-1)\n        emb_2 = model.forward(images_2).reshape(-1)\n        score = pdist(emb_1, emb_2)\n        scores_with_gt.append({'score': score, 'is_same': is_same, 'idx': batch_idx * batch_size})\n\n    return scores_with_gt\n\n\ndef compute_optimal_thresh(scores_with_gt):\n    \"\"\"Computes an optimal threshold for pairwise face verification\"\"\"\n    pos_scores = []\n    neg_scores = []\n    for score_with_gt in scores_with_gt:\n        if score_with_gt['is_same']:\n            pos_scores.append(score_with_gt['score'])\n        else:\n            neg_scores.append(score_with_gt['score'])\n\n    hist_pos, bins = np.histogram(np.array(pos_scores), 60)\n    hist_neg, _ = np.histogram(np.array(neg_scores), bins)\n\n    intersection_bins = []\n\n    for i in range(1, len(hist_neg)):\n        if hist_pos[i - 1] >= hist_neg[i - 1] and 0.05 < hist_pos[i] <= hist_neg[i]:\n            intersection_bins.append(bins[i])\n\n    if not intersection_bins:\n        intersection_bins.append(0.5)\n\n    return np.mean(intersection_bins)\n\n\ndef evaluate(args, dataset, model, compute_embeddings_fun, val_batch_size=16,\n             dump_embeddings=False, roc_fname='', snap_name='', verbose=True, show_failed=False):\n    \"\"\"Computes the LFW score of given model\"\"\"\n    if verbose and isinstance(model, torch.nn.Module):\n        log.info('Face recognition model config:')\n        log.info(model)\n        log.info('Number of parameters: {}'.format(get_model_parameters_number(model)))\n\n    scores_with_gt = compute_embeddings_fun(args, dataset, model, val_batch_size, dump_embeddings)\n    num_pairs = len(scores_with_gt)\n\n    subsets = []\n    for i in range(10):\n        lower_bnd = i * num_pairs // 10\n        upper_bnd = (i + 1) * num_pairs // 10\n        subset_test = [(lower_bnd, upper_bnd)]\n        subset_train = [(0, lower_bnd), (upper_bnd, num_pairs)]\n        subsets.append({'test': subset_test, 'train': subset_train})\n\n    same_scores = []\n    diff_scores = []\n    val_scores = []\n    threshs = []\n    mean_fpr = np.zeros(400)\n    mean_tpr = np.zeros(400)\n    failed_pairs = []\n\n    for subset in tqdm(subsets, '{} evaluation'.format(snap_name), disable=not verbose):\n        train_list = get_subset(scores_with_gt, subset['train'])\n        optimal_thresh = compute_optimal_thresh(train_list)\n        threshs.append(optimal_thresh)\n\n        test_list = get_subset(scores_with_gt, subset['test'])\n        same_correct = 0\n        diff_correct = 0\n        pos_pairs_num = neg_pairs_num = len(test_list) // 2\n\n        for score_with_gt in test_list:\n            if score_with_gt['score'] < optimal_thresh and score_with_gt['is_same']:\n                same_correct += 1\n            elif score_with_gt['score'] >= optimal_thresh and not score_with_gt['is_same']:\n                diff_correct += 1\n\n            if score_with_gt['score'] >= optimal_thresh and score_with_gt['is_same']:\n                failed_pairs.append(score_with_gt['idx'])\n            if score_with_gt['score'] < optimal_thresh and not score_with_gt['is_same']:\n                failed_pairs.append(score_with_gt['idx'])\n\n        same_scores.append(float(same_correct) / pos_pairs_num)\n        diff_scores.append(float(diff_correct) / neg_pairs_num)\n        val_scores.append(0.5*(same_scores[-1] + diff_scores[-1]))\n\n        fprs, tprs = get_roc(test_list, mean_fpr.shape[0])\n        mean_fpr = mean_fpr + fprs\n        mean_tpr = mean_tpr + tprs\n\n    mean_fpr /= 10\n    mean_tpr /= 10\n\n    if roc_fname:\n        save_roc(mean_tpr, mean_fpr, roc_fname)\n\n    same_acc = np.mean(same_scores)\n    diff_acc = np.mean(diff_scores)\n    overall_acc = np.mean(val_scores)\n    auc = get_auc(mean_fpr, mean_tpr)\n\n    if show_failed:\n        log.info('Number of misclassified pairs: {}'.format(len(failed_pairs)))\n        for pair in failed_pairs:\n            dataset.show_item(pair)\n\n    if verbose:\n        log.info('Accuracy/Val_same_accuracy mean: {0:.4f}'.format(same_acc))\n        log.info('Accuracy/Val_diff_accuracy mean: {0:.4f}'.format(diff_acc))\n        log.info('Accuracy/Val_accuracy mean: {0:.4f}'.format(overall_acc))\n        log.info('Accuracy/Val_accuracy std dev: {0:.4f}'.format(np.std(val_scores)))\n        log.info('AUC: {0:.4f}'.format(auc))\n        log.info('Estimated threshold: {0:.4f}'.format(np.mean(threshs)))\n\n    return same_acc, diff_acc, overall_acc, auc\n\n\ndef load_test_dataset(arguments):\n    \"\"\"Loads and configures the LFW dataset\"\"\"\n    input_size = models_backbones[arguments.model].get_input_res()\n    lfw = LFW(arguments.val, arguments.v_list, arguments.v_land)\n    assert lfw.use_landmarks\n    log.info('Using landmarks for the LFW images.')\n    transform = t.Compose([ResizeNumpy(input_size),\n                           NumpyToTensor(switch_rb=True)])\n    lfw.transform = transform\n    return lfw, partial(compute_embeddings_lfw, flipped_embeddings=arguments.flipped_emb)\n\ndef main():\n    parser = argparse.ArgumentParser(description='Evaluation script for Face Recognition in PyTorch')\n    parser.add_argument('--devices', type=int, nargs='+', default=[0], help='CUDA devices to use.')\n    parser.add_argument('--embed_size', type=int, default=128, help='Size of the face embedding.')\n    parser.add_argument('--val_data_root', dest='val', required=True, type=str, help='Path to validation data.')\n    parser.add_argument('--val_list', dest='v_list', required=True, type=str, help='Path to train data image list.')\n    parser.add_argument('--val_landmarks', dest='v_land', default='', required=False, type=str,\n                        help='Path to landmarks for the test images.')\n    parser.add_argument('--val_batch_size', type=int, default=8, help='Validation batch size.')\n    parser.add_argument('--snap', type=str, required=False, help='Snapshot to evaluate.')\n    parser.add_argument('--roc_fname', type=str, default='', help='ROC file.')\n    parser.add_argument('--dump_embeddings', action='store_true', help='Dump embeddings to summary writer.')\n    parser.add_argument('--dist', choices=['l2', 'cos'], type=str, default='cos', help='Distance.')\n    parser.add_argument('--flipped_emb', action='store_true', help='Flipped embedding concatenation trick.')\n    parser.add_argument('--show_failed', action='store_true', help='Show misclassified pairs.')\n    parser.add_argument('--model', choices=models_backbones.keys(), type=str, default='rmnet', help='Model type.')\n    parser.add_argument('--engine', choices=['pt', 'ie'], type=str, default='pt', help='Framework to use for eval.')\n\n    # IE-related options\n    parser.add_argument('--fr_model', type=str, required=False)\n    parser.add_argument('--lm_model', type=str, required=False)\n    parser.add_argument('-pp', '--plugin_dir', type=str, default=None, help='Path to a plugin folder')\n    args = parser.parse_args()\n\n    if args.engine == 'pt':\n        assert args.snap is not None, 'To evaluate PyTorch snapshot, please, specify --snap option.'\n        with torch.cuda.device(args.devices[0]):\n            data, embeddings_fun = load_test_dataset(args)\n            model = models_backbones[args.model](embedding_size=args.embed_size, feature=True)\n            model = load_model_state(model, args.snap, args.devices[0])\n            evaluate(args, data, model, embeddings_fun, args.val_batch_size, args.dump_embeddings,\n                     args.roc_fname, args.snap, True, args.show_failed)\n    else:\n        from utils.ie_tools import load_ie_model\n\n        assert args.fr_model is not None, 'To evaluate IE model, please, specify --fr_model option.'\n        fr_model = load_ie_model(args.fr_model, 'CPU', args.plugin_dir)\n        lm_model = None\n        if args.lm_model:\n            lm_model = load_ie_model(args.lm_model, 'CPU', args.plugin_dir)\n        input_size = tuple(fr_model.get_input_shape()[2:])\n\n        lfw = LFW(args.val, args.v_list, args.v_land)\n        if not lfw.use_landmarks or lm_model:\n            lfw.transform = t.Compose([ResizeNumpy(220), CenterCropNumpy(input_size)])\n            lfw.use_landmarks = False\n        else:\n            log.info('Using landmarks for the LFW images.')\n            lfw.transform = t.Compose([ResizeNumpy(input_size)])\n\n        evaluate(args, lfw, fr_model, partial(compute_embeddings_lfw_ie, lm_model=lm_model), val_batch_size=1,\n                 dump_embeddings=False, roc_fname='', snap_name='', verbose=True, show_failed=False)\n\nif __name__ == '__main__':\n    main()\n"
  },
  {
    "path": "init_venv.sh",
    "content": "#!/usr/bin/env bash\n\nwork_dir=$(realpath \"$(dirname $0)\")\n\ncd ${work_dir}\nif [[ -e venv ]]; then\n  echo \"Please remove a previously virtual environment folder '${work_dir}/venv'.\"\n  exit\nfi\n\n# Create virtual environment\nvirtualenv venv -p python3 --prompt=\"(deep=fr) \"\necho \"export PYTHONPATH=\\$PYTHONPATH:${work_dir}\" >> venv/bin/activate\n. venv/bin/activate\npip install -r ${work_dir}/requirements.txt\n\n\necho\necho \"====================================================\"\necho \"To start to work, you need to activate a virtualenv:\"\necho \"$ . venv/bin/activate\"\necho \"====================================================\"\n"
  },
  {
    "path": "losses/__init__.py",
    "content": ""
  },
  {
    "path": "losses/alignment.py",
    "content": "\"\"\"\n Copyright (c) 2018 Intel Corporation\n Licensed under the Apache License, Version 2.0 (the \"License\");\n you may not use this file except in compliance with the License.\n You may obtain a copy of the License at\n      http://www.apache.org/licenses/LICENSE-2.0\n Unless required by applicable law or agreed to in writing, software\n distributed under the License is distributed on an \"AS IS\" BASIS,\n WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n See the License for the specific language governing permissions and\n limitations under the License.\n\"\"\"\n\nimport math\nimport torch\nimport torch.nn as nn\n\nVALID_CORE_FUNC_TYPES = ['l1', 'l2', 'wing']\n\n\ndef wing_core(abs_x, w, eps):\n    \"\"\"Calculates the wing function from https://arxiv.org/pdf/1711.06753.pdf\"\"\"\n    return w*math.log(1. + abs_x / eps)\n\nclass AlignmentLoss(nn.Module):\n    \"\"\"Regression loss to train landmarks model\"\"\"\n    def __init__(self, loss_type='l2'):\n        super(AlignmentLoss, self).__init__()\n        assert loss_type in VALID_CORE_FUNC_TYPES\n        self.uniform_weights = True\n        self.weights = None\n        self.core_func_type = loss_type\n        self.eps = 0.031\n        self.w = 0.156\n\n    def set_weights(self, weights):\n        \"\"\"Set weights for the each landmark point in loss\"\"\"\n        self.uniform_weights = False\n        self.weights = torch.FloatTensor(weights).cuda()\n\n    def forward(self, input_values, target):\n        bs = input_values.shape[0]\n        loss = input_values - target\n        n_points = loss.shape[1] // 2\n        loss = loss.view(-1, n_points, 2)\n\n        if self.core_func_type == 'l2':\n            loss = torch.norm(loss, p=2, dim=2)\n            loss = loss.pow(2)\n            eyes_dist = (torch.norm(target[:, 0:2] - target[:, 2:4], p=2, dim=1).reshape(-1)).pow_(2)\n        elif self.core_func_type == 'l1':\n            loss = torch.norm(loss, p=1, dim=2)\n            eyes_dist = (torch.norm(target[:, 0:2] - target[:, 2:4], p=1, dim=1).reshape(-1))\n        elif self.core_func_type == 'wing':\n            wing_const = self.w - wing_core(self.w, self.w, self.eps)\n            loss = torch.abs(loss)\n            loss[loss < wing_const] = self.w*torch.log(1. + loss[loss < wing_const] / self.eps)\n            loss[loss >= wing_const] -= wing_const\n            loss = torch.sum(loss, 2)\n            eyes_dist = (torch.norm(target[:, 0:2] - target[:, 2:4], p=1, dim=1).reshape(-1))\n\n        if self.uniform_weights:\n            loss = torch.sum(loss, 1)\n            loss /= n_points\n        else:\n            assert self.weights.shape[0] == loss.shape[1]\n            loss = torch.mul(loss, self.weights)\n            loss = torch.sum(loss, 1)\n\n        loss = torch.div(loss, eyes_dist)\n        loss = torch.sum(loss)\n        return loss / (2.*bs)\n"
  },
  {
    "path": "losses/am_softmax.py",
    "content": "\"\"\"\n Copyright (c) 2018 Intel Corporation\n Licensed under the Apache License, Version 2.0 (the \"License\");\n you may not use this file except in compliance with the License.\n You may obtain a copy of the License at\n      http://www.apache.org/licenses/LICENSE-2.0\n Unless required by applicable law or agreed to in writing, software\n distributed under the License is distributed on an \"AS IS\" BASIS,\n WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n See the License for the specific language governing permissions and\n limitations under the License.\n\"\"\"\n\nimport math\n\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nfrom torch.nn import Parameter\n\n\nclass AngleSimpleLinear(nn.Module):\n    \"\"\"Computes cos of angles between input vectors and weights vectors\"\"\"\n    def __init__(self, in_features, out_features):\n        super(AngleSimpleLinear, self).__init__()\n        self.in_features = in_features\n        self.out_features = out_features\n        self.weight = Parameter(torch.Tensor(in_features, out_features))\n        self.weight.data.uniform_(-1, 1).renorm_(2, 1, 1e-5).mul_(1e5)\n\n    def forward(self, x):\n        cos_theta = F.normalize(x, dim=1).mm(F.normalize(self.weight, dim=0))\n        return cos_theta.clamp(-1, 1)\n\n\ndef focal_loss(input_values, gamma):\n    \"\"\"Computes the focal loss\"\"\"\n    p = torch.exp(-input_values)\n    loss = (1 - p) ** gamma * input_values\n    return loss.mean()\n\n\nclass AMSoftmaxLoss(nn.Module):\n    \"\"\"Computes the AM-Softmax loss with cos or arc margin\"\"\"\n    margin_types = ['cos', 'arc']\n\n    def __init__(self, margin_type='cos', gamma=0., m=0.5, s=30, t=1.):\n        super(AMSoftmaxLoss, self).__init__()\n        assert margin_type in AMSoftmaxLoss.margin_types\n        self.margin_type = margin_type\n        assert gamma >= 0\n        self.gamma = gamma\n        assert m > 0\n        self.m = m\n        assert s > 0\n        self.s = s\n        self.cos_m = math.cos(m)\n        self.sin_m = math.sin(m)\n        self.th = math.cos(math.pi - m)\n        assert t >= 1\n        self.t = t\n\n    def forward(self, cos_theta, target):\n        if self.margin_type == 'cos':\n            phi_theta = cos_theta - self.m\n        else:\n            sine = torch.sqrt(1.0 - torch.pow(cos_theta, 2))\n            phi_theta = cos_theta * self.cos_m - sine * self.sin_m #cos(theta+m)\n            phi_theta = torch.where(cos_theta > self.th, phi_theta, cos_theta - self.sin_m * self.m)\n\n        index = torch.zeros_like(cos_theta, dtype=torch.uint8)\n        index.scatter_(1, target.data.view(-1, 1), 1)\n        output = torch.where(index, phi_theta, cos_theta)\n\n        if self.gamma == 0 and self.t == 1.:\n            return F.cross_entropy(self.s*output, target)\n\n        if self.t > 1:\n            h_theta = self.t - 1 + self.t*cos_theta\n            support_vecs_mask = (1 - index) * \\\n                torch.lt(torch.masked_select(phi_theta, index).view(-1, 1).repeat(1, h_theta.shape[1]) - cos_theta, 0)\n            output = torch.where(support_vecs_mask, h_theta, output)\n            return F.cross_entropy(self.s*output, target)\n\n        return focal_loss(F.cross_entropy(self.s*output, target, reduction='none'), self.gamma)\n"
  },
  {
    "path": "losses/centroid_based.py",
    "content": "\"\"\"\n Copyright (c) 2018 Intel Corporation\n Licensed under the Apache License, Version 2.0 (the \"License\");\n you may not use this file except in compliance with the License.\n You may obtain a copy of the License at\n      http://www.apache.org/licenses/LICENSE-2.0\n Unless required by applicable law or agreed to in writing, software\n distributed under the License is distributed on an \"AS IS\" BASIS,\n WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n See the License for the specific language governing permissions and\n limitations under the License.\n\"\"\"\n\nimport torch.nn as nn\nimport torch\nimport torch.nn.functional as F\nimport numpy as np\n\n\nclass CenterLoss(nn.Module):\n    \"\"\"Implements the Center loss from https://ydwen.github.io/papers/WenECCV16.pdf\"\"\"\n    def __init__(self, num_classes, embed_size, cos_dist=True):\n        super().__init__()\n        self.cos_dist = cos_dist\n        self.num_classes = num_classes\n        self.centers = nn.Parameter(torch.randn(self.num_classes, embed_size).cuda())\n        self.embed_size = embed_size\n        self.mse = nn.MSELoss(reduction='elementwise_mean')\n\n    def get_centers(self):\n        \"\"\"Returns estimated centers\"\"\"\n        return self.centers\n\n    def forward(self, features, labels):\n        features = F.normalize(features)\n        batch_size = labels.size(0)\n        features_dim = features.size(1)\n        assert features_dim == self.embed_size\n\n        if self.cos_dist:\n            self.centers.data = F.normalize(self.centers.data, p=2, dim=1)\n\n        centers_batch = self.centers[labels, :]\n\n        if self.cos_dist:\n            cos_sim = nn.CosineSimilarity()\n            cos_diff = 1. - cos_sim(features, centers_batch)\n            center_loss = torch.sum(cos_diff) / batch_size\n        else:\n            center_loss = self.mse(centers_batch, features)\n\n        return center_loss\n\n\nclass MinimumMargin(nn.Module):\n    \"\"\"Implements the Minimum margin loss from https://arxiv.org/abs/1805.06741\"\"\"\n    def __init__(self, margin=.6):\n        super().__init__()\n        self.margin = margin\n\n    def forward(self, centers, labels):\n        loss_value = 0\n\n        batch_centers = centers[labels, :]\n        labels = labels.cpu().data.numpy()\n\n        all_pairs = labels.reshape([-1, 1]) != labels.reshape([1, -1])\n        valid_pairs = (all_pairs * np.tri(*all_pairs.shape, k=-1, dtype=np.bool)).astype(np.float32)\n        losses = 1. - torch.mm(batch_centers, torch.t(batch_centers)) - self.margin\n\n        valid_pairs *= (losses.cpu().data.numpy() > 0.0)\n        num_valid = float(np.sum(valid_pairs))\n\n        if num_valid > 0:\n            loss_value = torch.sum(losses * torch.from_numpy(valid_pairs).cuda())\n        else:\n            return loss_value\n\n        return loss_value / num_valid\n\n\nclass GlobalPushPlus(nn.Module):\n    \"\"\"Implements the Global Push Plus loss\"\"\"\n    def __init__(self, margin=.6):\n        super().__init__()\n        self.min_margin = 0.15\n        self.max_margin = margin\n        self.num_calls = 0\n\n    def forward(self, features, centers, labels):\n        self.num_calls += 1\n        features = F.normalize(features)\n        loss_value = 0\n        batch_centers = centers[labels, :]\n        labels = labels.cpu().data.numpy()\n        assert len(labels.shape) == 1\n\n        center_ids = np.arange(centers.shape[0], dtype=np.int32)\n        different_class_pairs = labels.reshape([-1, 1]) != center_ids.reshape([1, -1])\n\n        pos_distances = 1.0 - torch.sum(features * batch_centers, dim=1)\n        neg_distances = 1.0 - torch.mm(features, torch.t(centers))\n\n        margin = self.min_margin + float(self.num_calls) / float(40000) * (self.max_margin - self.min_margin)\n        margin = min(margin, self.max_margin)\n\n        losses = margin + pos_distances.view(-1, 1) - neg_distances\n\n        valid_pairs = (different_class_pairs * (losses.cpu().data.numpy() > 0.0)).astype(np.float32)\n        num_valid = float(np.sum(valid_pairs))\n\n        if num_valid > 0:\n            loss_value = torch.sum(losses * torch.from_numpy(valid_pairs).cuda())\n        else:\n            return loss_value\n\n        return loss_value / num_valid\n\n\nclass PushPlusLoss(nn.Module):\n    \"\"\"Implements the Push Plus loss\"\"\"\n    def __init__(self, margin=.7):\n        super().__init__()\n        self.margin = margin\n\n    def forward(self, features, centers, labels):\n        features = F.normalize(features)\n        loss_value = 0\n        batch_centers = centers[labels, :]\n        labels = labels.cpu().data.numpy()\n        assert len(labels.shape) == 1\n\n        all_pairs = labels.reshape([-1, 1]) != labels.reshape([1, -1])\n        pos_distances = 1.0 - torch.sum(features * batch_centers, dim=1)\n        neg_distances = 1.0 - torch.mm(features, torch.t(features))\n\n        losses = self.margin + pos_distances.view(-1, 1) - neg_distances\n        valid_pairs = (all_pairs * (losses.cpu().data.numpy() > 0.0)).astype(np.float32)\n        num_valid = float(np.sum(valid_pairs))\n\n        if num_valid > 0:\n            loss_value = torch.sum(losses * torch.from_numpy(valid_pairs).cuda())\n        else:\n            return loss_value\n\n        return loss_value / num_valid\n\n\nclass PushLoss(nn.Module):\n    \"\"\"Implements the Push loss\"\"\"\n    def __init__(self, soft=True, margin=0.5):\n        super().__init__()\n        self.soft = soft\n        self.margin = margin\n\n    def forward(self, features, labels):\n        features = F.normalize(features)\n        loss_value = 0\n        labels = labels.cpu().data.numpy()\n        assert len(labels.shape) == 1\n\n        all_pairs = labels.reshape([-1, 1]) != labels.reshape([1, -1])\n        valid_pairs = (all_pairs * np.tri(*all_pairs.shape, k=-1, dtype=np.bool)).astype(np.float32)\n\n        if self.soft:\n            losses = torch.log(1. + torch.exp(torch.mm(features, torch.t(features)) - 1))\n            num_valid = float(np.sum(valid_pairs))\n        else:\n            losses = self.margin - (1. - torch.mm(features, torch.t(features)))\n            valid_pairs *= (losses.cpu().data.numpy() > 0.0)\n            num_valid = float(np.sum(valid_pairs))\n\n        if num_valid > 0:\n            loss_value = torch.sum(losses * torch.from_numpy(valid_pairs).cuda())\n        else:\n            return loss_value\n\n        return loss_value / num_valid\n"
  },
  {
    "path": "losses/metric_losses.py",
    "content": "\"\"\"\n Copyright (c) 2018 Intel Corporation\n Licensed under the Apache License, Version 2.0 (the \"License\");\n you may not use this file except in compliance with the License.\n You may obtain a copy of the License at\n      http://www.apache.org/licenses/LICENSE-2.0\n Unless required by applicable law or agreed to in writing, software\n distributed under the License is distributed on an \"AS IS\" BASIS,\n WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n See the License for the specific language governing permissions and\n limitations under the License.\n\"\"\"\n\nimport torch\nfrom losses.centroid_based import CenterLoss, PushLoss, MinimumMargin, PushPlusLoss, GlobalPushPlus\n\n\nclass MetricLosses:\n    \"\"\"Class-aggregator for all metric-learning losses\"\"\"\n    def __init__(self, classes_num, embed_size, writer):\n        self.writer = writer\n        self.center_loss = CenterLoss(classes_num, embed_size, cos_dist=True)\n        self.optimizer_centloss = torch.optim.SGD(self.center_loss.parameters(), lr=0.5)\n        self.center_coeff = 0.0\n\n        self.push_loss = PushLoss(soft=False, margin=0.7)\n        self.push_loss_coeff = 0.0\n\n        self.push_plus_loss = PushPlusLoss(margin=0.7)\n        self.push_plus_loss_coeff = 0.0\n\n        self.glob_push_plus_loss = GlobalPushPlus(margin=0.7)\n        self.glob_push_plus_loss_coeff = 0.0\n\n        self.min_margin_loss = MinimumMargin(margin=.7)\n        self.min_margin_loss_coeff = 0.0\n\n    def __call__(self, features, labels, epoch_num, iteration):\n        log_string = ''\n\n        center_loss_val = 0\n        if self.center_coeff > 0.:\n            center_loss_val = self.center_loss(features, labels)\n            self.writer.add_scalar('Loss/center_loss', center_loss_val, iteration)\n            log_string += ' Center loss: %.4f' % center_loss_val\n\n        push_loss_val = 0\n        if self.push_loss_coeff > 0.0:\n            push_loss_val = self.push_loss(features, labels)\n            self.writer.add_scalar('Loss/push_loss', push_loss_val, iteration)\n            log_string += ' Push loss: %.4f' % push_loss_val\n\n        push_plus_loss_val = 0\n        if self.push_plus_loss_coeff > 0.0 and self.center_coeff > 0.0:\n            push_plus_loss_val = self.push_plus_loss(features, self.center_loss.get_centers(), labels)\n            self.writer.add_scalar('Loss/push_plus_loss', push_plus_loss_val, iteration)\n            log_string += ' Push Plus loss: %.4f' % push_plus_loss_val\n\n        glob_push_plus_loss_val = 0\n        if self.glob_push_plus_loss_coeff > 0.0 and self.center_coeff > 0.0:\n            glob_push_plus_loss_val = self.glob_push_plus_loss(features, self.center_loss.get_centers(), labels)\n            self.writer.add_scalar('Loss/global_push_plus_loss', glob_push_plus_loss_val, iteration)\n            log_string += ' Global Push Plus loss: %.4f' % glob_push_plus_loss_val\n\n        min_margin_loss_val = 0\n        if self.min_margin_loss_coeff > 0.0 and self.center_coeff > 0.0:\n            min_margin_loss_val = self.min_margin_loss(self.center_loss.get_centers(), labels)\n            self.writer.add_scalar('Loss/min_margin_loss', min_margin_loss_val, iteration)\n            log_string += ' Min margin loss: %.4f' % min_margin_loss_val\n\n        loss_value = self.center_coeff * center_loss_val + self.push_loss_coeff * push_loss_val + \\\n                     self.push_plus_loss_coeff * push_plus_loss_val + self.min_margin_loss_coeff * min_margin_loss_val \\\n                     + self.glob_push_plus_loss_coeff * glob_push_plus_loss_val\n\n        if self.min_margin_loss_coeff + self.center_coeff + self.push_loss_coeff + self.push_plus_loss_coeff > 0.:\n            self.writer.add_scalar('Loss/AUX_losses', loss_value, iteration)\n\n        return loss_value, log_string\n\n    def init_iteration(self):\n        \"\"\"Initializes a training iteration\"\"\"\n        if self.center_coeff > 0.:\n            self.optimizer_centloss.zero_grad()\n\n    def end_iteration(self):\n        \"\"\"Finalizes a training iteration\"\"\"\n        if self.center_coeff > 0.:\n            for param in self.center_loss.parameters():\n                param.grad.data *= (1. / self.center_coeff)\n            self.optimizer_centloss.step()\n"
  },
  {
    "path": "losses/regularizer.py",
    "content": "\"\"\"\n Copyright (c) 2018 Intel Corporation\n Licensed under the Apache License, Version 2.0 (the \"License\");\n you may not use this file except in compliance with the License.\n You may obtain a copy of the License at\n      http://www.apache.org/licenses/LICENSE-2.0\n Unless required by applicable law or agreed to in writing, software\n distributed under the License is distributed on an \"AS IS\" BASIS,\n WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n See the License for the specific language governing permissions and\n limitations under the License.\n\"\"\"\n\nimport torch\nimport torch.nn.functional as F\n\n\ndef l2_reg_ortho(mdl):\n    \"\"\"\n        Function used for Orthogonal Regularization.\n    \"\"\"\n    l2_reg = None\n    for w in mdl.parameters():\n        if w.ndimension() < 2:\n            continue\n        else:\n            cols = w[0].numel()\n            w1 = w.view(-1, cols)\n            wt = torch.transpose(w1, 0, 1)\n            m = torch.matmul(wt, w1)\n            ident = torch.eye(cols, cols).cuda()\n\n            w_tmp = (m - ident)\n            height = w_tmp.size(0)\n            u = F.normalize(w_tmp.new_empty(height).normal_(0, 1), dim=0, eps=1e-12)\n            v = F.normalize(torch.matmul(w_tmp.t(), u), dim=0, eps=1e-12)\n            u = F.normalize(torch.matmul(w_tmp, v), dim=0, eps=1e-12)\n            sigma = torch.dot(u, torch.matmul(w_tmp, v))\n\n            if l2_reg is None:\n                l2_reg = (torch.norm(sigma, 2))**2\n            else:\n                l2_reg += (torch.norm(sigma, 2))**2\n    return l2_reg\n\n\nclass ODecayScheduler():\n    \"\"\"Scheduler for the decay of the orthogonal regularizer\"\"\"\n    def __init__(self, schedule, initial_decay, mult_factor):\n        assert len(schedule) > 1\n        self.schedule = schedule\n        self.epoch_num = 0\n        self.mult_factor = mult_factor\n        self.decay = initial_decay\n\n    def step(self):\n        \"\"\"Switches to the next step\"\"\"\n        self.epoch_num += 1\n        if self.epoch_num in self.schedule:\n            self.decay *= self.mult_factor\n        if self.epoch_num == self.schedule[-1]:\n            self.decay = 0.0\n\n    def get_decay(self):\n        \"\"\"Returns the current value of decay according to th schedule\"\"\"\n        return self.decay\n"
  },
  {
    "path": "model/__init__.py",
    "content": ""
  },
  {
    "path": "model/backbones/__init__.py",
    "content": ""
  },
  {
    "path": "model/backbones/resnet.py",
    "content": "\"\"\"\n Copyright (c) 2018 Intel Corporation\n Licensed under the Apache License, Version 2.0 (the \"License\");\n you may not use this file except in compliance with the License.\n You may obtain a copy of the License at\n      http://www.apache.org/licenses/LICENSE-2.0\n Unless required by applicable law or agreed to in writing, software\n distributed under the License is distributed on an \"AS IS\" BASIS,\n WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n See the License for the specific language governing permissions and\n limitations under the License.\n\"\"\"\n\nimport torch.nn as nn\n\nfrom model.blocks.resnet_blocks import Bottleneck, BasicBlock\nfrom model.blocks.shared_blocks import make_activation\n\n\nclass ResNet(nn.Module):\n    def __init__(self, block, layers, num_classes=1000, activation=nn.ReLU):\n        self.inplanes = 64\n        super(ResNet, self).__init__()\n        self.conv1 = nn.Conv2d(3, 64, kernel_size=3, stride=1, padding=1,\n                               bias=False)\n        self.bn1 = nn.BatchNorm2d(64)\n        self.relu = make_activation(nn.ReLU)\n        self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)\n        self.layer1 = self._make_layer(block, 64, layers[0], activation=activation)\n        self.layer2 = self._make_layer(block, 128, layers[1], stride=2, activation=activation)\n        self.layer3 = self._make_layer(block, 256, layers[2], stride=2, activation=activation)\n        self.layer4 = self._make_layer(block, 512, layers[3], stride=2, activation=activation)\n        self.avgpool = nn.Conv2d(512 * block.expansion, 512 * block.expansion, 7,\n                                 groups=512 * block.expansion, bias=False)\n        self.fc = nn.Conv2d(512 * block.expansion, num_classes, 1, stride=1, padding=0, bias=False)\n\n        for m in self.modules():\n            if isinstance(m, nn.Conv2d):\n                nn.init.kaiming_normal_(m.weight, mode='fan_out')\n            elif isinstance(m, nn.BatchNorm2d):\n                nn.init.constant_(m.weight, 1)\n                nn.init.constant_(m.bias, 0)\n\n    def _make_layer(self, block, planes, blocks, stride=1, activation=nn.ReLU):\n        downsample = None\n        if stride != 1 or self.inplanes != planes * block.expansion:\n            downsample = nn.Sequential(\n                nn.Conv2d(self.inplanes, planes * block.expansion,\n                          kernel_size=1, stride=stride, bias=False),\n                nn.BatchNorm2d(planes * block.expansion),\n            )\n\n        layers = []\n        layers.append(block(self.inplanes, planes, stride, downsample, activation=activation))\n        self.inplanes = planes * block.expansion\n        for _ in range(1, blocks):\n            layers.append(block(self.inplanes, planes, activation=activation))\n\n        return nn.Sequential(*layers)\n\n    def forward(self, x):\n        x = self.conv1(x)\n        x = self.bn1(x)\n        x = self.relu(x)\n        x = self.maxpool(x)\n\n        x = self.layer1(x)\n        x = self.layer2(x)\n        x = self.layer3(x)\n        x = self.layer4(x)\n\n        x = self.avgpool(x)\n        x = self.fc(x)\n\n        return x\n\n\ndef resnet50(**kwargs):\n    model = ResNet(Bottleneck, [3, 4, 6, 3], **kwargs)\n    return model\n\n\ndef resnet34(**kwargs):\n    model = ResNet(BasicBlock, [3, 4, 6, 3], **kwargs)\n    return model\n"
  },
  {
    "path": "model/backbones/rmnet.py",
    "content": "\"\"\"\n Copyright (c) 2018 Intel Corporation\n Licensed under the Apache License, Version 2.0 (the \"License\");\n you may not use this file except in compliance with the License.\n You may obtain a copy of the License at\n      http://www.apache.org/licenses/LICENSE-2.0\n Unless required by applicable law or agreed to in writing, software\n distributed under the License is distributed on an \"AS IS\" BASIS,\n WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n See the License for the specific language governing permissions and\n limitations under the License.\n\"\"\"\n\nfrom collections import OrderedDict\n\nimport torch.nn as nn\nfrom ..blocks.rmnet_blocks import RMBlock\n\n\nclass RMNetBody(nn.Module):\n    def __init__(self, block=RMBlock, blocks_per_stage=(None, 4, 8, 10, 11), trunk_width=(32, 32, 64, 128, 256),\n                 bottleneck_width=(None, 8, 16, 32, 64)):\n        super(RMNetBody, self).__init__()\n        assert len(blocks_per_stage) == len(trunk_width) == len(bottleneck_width)\n        self.dim_out = trunk_width[-1]\n\n        stages = [nn.Sequential(OrderedDict([\n            ('data_bn', nn.BatchNorm2d(3)),\n            ('conv1', nn.Conv2d(3, trunk_width[0], kernel_size=3, stride=2, padding=1, bias=False)),\n            ('bn1', nn.BatchNorm2d(trunk_width[0])),\n            ('relu1', nn.ReLU(inplace=True))])), ]\n\n        for i, (blocks_num, w, wb) in enumerate(zip(blocks_per_stage, trunk_width, bottleneck_width)):\n            # Zeroth stage is already added.\n            if i == 0:\n                continue\n            stage = []\n            # Do not downscale input to the first stage.\n            if i > 1:\n                stage.append(block(trunk_width[i - 1], wb, w, downsample=True))\n            for _ in range(blocks_num):\n                stage.append(block(w, wb, w))\n            stages.append(nn.Sequential(*stage))\n\n        self.stages = nn.Sequential(OrderedDict([('stage_{}'.format(i), stage) for i, stage in enumerate(stages)]))\n\n        self.init_weights()\n\n    def init_weights(self):\n        m = self.stages[0][0]  # ['data_bn']\n        nn.init.constant_(m.weight, 1)\n        nn.init.constant_(m.bias, 0)\n        m = self.stages[0][1]  # ['conv1']\n        nn.init.kaiming_normal_(m.weight, mode='fan_out')\n        m = self.stages[0][2]  # ['bn1']\n        nn.init.constant_(m.weight, 1)\n        nn.init.constant_(m.bias, 0)\n        # All other blocks should be initialized internally during instantiation.\n\n    def forward(self, x):\n        return self.stages(x)\n"
  },
  {
    "path": "model/backbones/se_resnet.py",
    "content": "\"\"\"\n Copyright (c) 2018 Intel Corporation\n Licensed under the Apache License, Version 2.0 (the \"License\");\n you may not use this file except in compliance with the License.\n You may obtain a copy of the License at\n      http://www.apache.org/licenses/LICENSE-2.0\n Unless required by applicable law or agreed to in writing, software\n distributed under the License is distributed on an \"AS IS\" BASIS,\n WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n See the License for the specific language governing permissions and\n limitations under the License.\n\"\"\"\n\nimport math\n\nimport torch.nn as nn\n\nfrom model.blocks.se_resnet_blocks import SEBottleneck\n\n\nclass SEResNet(nn.Module):\n    def __init__(self, block, layers, num_classes=1000, activation=nn.ReLU):\n        self.inplanes = 64\n        super(SEResNet, self).__init__()\n        self.conv1 = nn.Conv2d(3, 64, kernel_size=3, stride=1, padding=1,\n                               bias=False)\n        self.bn1 = nn.BatchNorm2d(64)\n        self.relu = nn.ReLU(inplace=True)\n        self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)\n        self.layer1 = self._make_layer(block, 64, layers[0], activation=activation)\n        self.layer2 = self._make_layer(block, 128, layers[1], stride=2, activation=activation)\n        self.layer3 = self._make_layer(block, 256, layers[2], stride=2, activation=activation)\n        self.layer4 = self._make_layer(block, 512, layers[3], stride=2, activation=activation)\n        self.avgpool = nn.Conv2d(512 * block.expansion, 512 * block.expansion, 7,\n                                 groups=512 * block.expansion, bias=False)\n        self.fc = nn.Conv2d(512 * block.expansion, num_classes, 1, stride=1, padding=0, bias=False)\n\n        for m in self.modules():\n            if isinstance(m, nn.Conv2d):\n                n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels\n                m.weight.data.normal_(0, math.sqrt(2. / n))\n            elif isinstance(m, nn.BatchNorm2d):\n                m.weight.data.fill_(1)\n                m.bias.data.zero_()\n\n    def _make_layer(self, block, planes, blocks, stride=1, activation=nn.ReLU):\n        downsample = None\n        if stride != 1 or self.inplanes != planes * block.expansion:\n            downsample = nn.Sequential(\n                nn.Conv2d(self.inplanes, planes * block.expansion,\n                          kernel_size=1, stride=stride, bias=False),\n                nn.BatchNorm2d(planes * block.expansion),\n            )\n\n        layers = []\n        layers.append(block(self.inplanes, planes, stride, downsample, activation=activation))\n        self.inplanes = planes * block.expansion\n        for _ in range(1, blocks):\n            layers.append(block(self.inplanes, planes, activation=activation))\n\n        return nn.Sequential(*layers)\n\n    def forward(self, x):\n        x = self.conv1(x)\n        x = self.bn1(x)\n        x = self.relu(x)\n        x = self.maxpool(x)\n\n        x = self.layer1(x)\n        x = self.layer2(x)\n        x = self.layer3(x)\n        x = self.layer4(x)\n\n        x = self.avgpool(x)\n        x = self.fc(x)\n\n        return x\n\n\ndef se_resnet50(**kwargs):\n    model = SEResNet(SEBottleneck, [3, 4, 6, 3], **kwargs)\n    return model\n\n\ndef se_resnet101(**kwargs):\n    model = SEResNet(SEBottleneck, [3, 4, 23, 3], **kwargs)\n    return model\n\n\ndef se_resnet152(**kwargs):\n    model = SEResNet(SEBottleneck, [3, 8, 36, 3], **kwargs)\n    return model\n"
  },
  {
    "path": "model/backbones/se_resnext.py",
    "content": "\"\"\"\n Copyright (c) 2018 Intel Corporation\n Licensed under the Apache License, Version 2.0 (the \"License\");\n you may not use this file except in compliance with the License.\n You may obtain a copy of the License at\n      http://www.apache.org/licenses/LICENSE-2.0\n Unless required by applicable law or agreed to in writing, software\n distributed under the License is distributed on an \"AS IS\" BASIS,\n WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n See the License for the specific language governing permissions and\n limitations under the License.\n\"\"\"\n\nimport math\nimport torch.nn as nn\n\nfrom model.blocks.se_resnext_blocks import SEBottleneckX\n\n\nclass SEResNeXt(nn.Module):\n\n    def __init__(self, block, layers, cardinality=32, num_classes=1000):\n        super(SEResNeXt, self).__init__()\n        self.cardinality = cardinality\n        self.inplanes = 64\n\n        self.conv1 = nn.Conv2d(3, 64, kernel_size=7, stride=2, padding=3,\n                               bias=False)\n        self.bn1 = nn.BatchNorm2d(64)\n        self.relu = nn.ReLU(inplace=True)\n        self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)\n\n        self.layer1 = self._make_layer(block, 64, layers[0])\n        self.layer2 = self._make_layer(block, 128, layers[1], stride=2)\n        self.layer3 = self._make_layer(block, 256, layers[2], stride=2)\n        self.layer4 = self._make_layer(block, 512, layers[3], stride=2)\n\n        self.avgpool = nn.AdaptiveAvgPool2d(1)\n        self.fc = nn.Linear(512 * block.expansion, num_classes)\n\n        for m in self.modules():\n            if isinstance(m, nn.Conv2d):\n                n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels\n                m.weight.data.normal_(0, math.sqrt(2. / n))\n                if m.bias is not None:\n                    m.bias.data.zero_()\n            elif isinstance(m, nn.BatchNorm2d):\n                m.weight.data.fill_(1)\n                m.bias.data.zero_()\n\n    def _make_layer(self, block, planes, blocks, stride=1):\n        downsample = None\n        if stride != 1 or self.inplanes != planes * block.expansion:\n            downsample = nn.Sequential(\n                nn.Conv2d(self.inplanes, planes * block.expansion,\n                          kernel_size=1, stride=stride, bias=False),\n                nn.BatchNorm2d(planes * block.expansion),\n            )\n\n        layers = []\n        layers.append(block(self.inplanes, planes, self.cardinality, stride, downsample))\n        self.inplanes = planes * block.expansion\n        for _ in range(1, blocks):\n            layers.append(block(self.inplanes, planes, self.cardinality))\n\n        return nn.Sequential(*layers)\n\n    def forward(self, x):\n        x = self.conv1(x)\n        x = self.bn1(x)\n        x = self.relu(x)\n        x = self.maxpool(x)\n\n        x = self.layer1(x)\n        x = self.layer2(x)\n        x = self.layer3(x)\n        x = self.layer4(x)\n\n        x = self.avgpool(x)\n        x = x.view(x.size(0), -1)\n\n        x = self.fc(x)\n\n        return x\n\n\ndef se_resnext50(**kwargs):\n    model = SEResNeXt(SEBottleneckX, [3, 4, 6, 3], **kwargs)\n    return model\n\n\ndef se_resnext101(**kwargs):\n    model = SEResNeXt(SEBottleneckX, [3, 4, 23, 3], **kwargs)\n    return model\n\n\ndef se_resnext152(**kwargs):\n    model = SEResNeXt(SEBottleneckX, [3, 8, 36, 3], **kwargs)\n    return model\n"
  },
  {
    "path": "model/backbones/shufflenet_v2.py",
    "content": "\"\"\"\n Copyright (c) 2018 Intel Corporation\n Licensed under the Apache License, Version 2.0 (the \"License\");\n you may not use this file except in compliance with the License.\n You may obtain a copy of the License at\n      http://www.apache.org/licenses/LICENSE-2.0\n Unless required by applicable law or agreed to in writing, software\n distributed under the License is distributed on an \"AS IS\" BASIS,\n WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n See the License for the specific language governing permissions and\n limitations under the License.\n\"\"\"\n\nimport torch.nn as nn\n\nfrom model.blocks.shufflenet_v2_blocks import ShuffleInvertedResidual, conv_bn, conv_1x1_bn\n\n\nclass ShuffleNetV2Body(nn.Module):\n    def __init__(self, input_size=224, width_mult=1.):\n        super(ShuffleNetV2Body, self).__init__()\n\n        assert input_size % 32 == 0\n\n        self.stage_repeats = [4, 8, 4]\n        if width_mult == 0.5:\n            self.stage_out_channels = [-1, 24, 48, 96, 192, 1024]\n        elif width_mult == 1.0:\n            self.stage_out_channels = [-1, 24, 116, 232, 464, 1024]\n        elif width_mult == 1.5:\n            self.stage_out_channels = [-1, 24, 176, 352, 704, 1024]\n        elif width_mult == 2.0:\n            self.stage_out_channels = [-1, 24, 224, 488, 976, 2048]\n        else:\n            raise ValueError(\"Unsupported width multiplier\")\n\n        # building first layer\n        self.bn_first = nn.BatchNorm2d(3)\n        input_channel = self.stage_out_channels[1]\n        self.conv1 = conv_bn(3, input_channel, 2)\n\n        self.features = []\n\n        # building inverted residual blocks\n        for idxstage in range(len(self.stage_repeats)):\n            numrepeat = self.stage_repeats[idxstage]\n            output_channel = self.stage_out_channels[idxstage+2]\n            for i in range(numrepeat):\n                if i == 0:\n                    self.features.append(ShuffleInvertedResidual(input_channel, output_channel,\n                                                                 2, 2, activation=nn.PReLU))\n                else:\n                    self.features.append(ShuffleInvertedResidual(input_channel, output_channel,\n                                                                 1, 1, activation=nn.PReLU))\n                input_channel = output_channel\n\n        self.features = nn.Sequential(*self.features)\n        self.conv_last = conv_1x1_bn(input_channel, self.stage_out_channels[-1], activation=nn.PReLU)\n        self.init_weights()\n\n    @staticmethod\n    def get_downscale_factor():\n        return 16\n\n    def init_weights(self):\n        m = self.bn_first\n        nn.init.constant_(m.weight, 1)\n        nn.init.constant_(m.bias, 0)\n\n    def get_num_output_channels(self):\n        return self.stage_out_channels[-1]\n\n    def forward(self, x):\n        x = self.conv1(self.bn_first(x))\n        x = self.features(x)\n        x = self.conv_last(x)\n        return x\n"
  },
  {
    "path": "model/blocks/__init__.py",
    "content": ""
  },
  {
    "path": "model/blocks/mobilenet_v2_blocks.py",
    "content": "\"\"\"\n Copyright (c) 2018 Intel Corporation\n Licensed under the Apache License, Version 2.0 (the \"License\");\n you may not use this file except in compliance with the License.\n You may obtain a copy of the License at\n      http://www.apache.org/licenses/LICENSE-2.0\n Unless required by applicable law or agreed to in writing, software\n distributed under the License is distributed on an \"AS IS\" BASIS,\n WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n See the License for the specific language governing permissions and\n limitations under the License.\n\"\"\"\n\nimport torch.nn as nn\n\nfrom model.blocks.shared_blocks import SELayer\n\n\nclass InvertedResidual(nn.Module):\n    \"\"\"Implementation of the modified Inverted residual block\"\"\"\n    def __init__(self, in_channels, out_channels, stride, expand_ratio, outp_size=None):\n        super(InvertedResidual, self).__init__()\n        self.stride = stride\n        assert stride in [1, 2]\n\n        self.use_res_connect = self.stride == 1 and in_channels == out_channels\n\n        self.inv_block = nn.Sequential(\n            nn.Conv2d(in_channels, in_channels * expand_ratio, 1, 1, 0, bias=False),\n            nn.BatchNorm2d(in_channels * expand_ratio),\n            nn.PReLU(),\n\n            nn.Conv2d(in_channels * expand_ratio, in_channels * expand_ratio, 3, stride, 1,\n                      groups=in_channels * expand_ratio, bias=False),\n            nn.BatchNorm2d(in_channels * expand_ratio),\n            nn.PReLU(),\n\n            nn.Conv2d(in_channels * expand_ratio, out_channels, 1, 1, 0, bias=False),\n            nn.BatchNorm2d(out_channels),\n            SELayer(out_channels, 8, nn.PReLU, outp_size)\n        )\n\n    def forward(self, x):\n        if self.use_res_connect:\n            return x + self.inv_block(x)\n\n        return self.inv_block(x)\n"
  },
  {
    "path": "model/blocks/resnet_blocks.py",
    "content": "\"\"\"\n Copyright (c) 2018 Intel Corporation\n Licensed under the Apache License, Version 2.0 (the \"License\");\n you may not use this file except in compliance with the License.\n You may obtain a copy of the License at\n      http://www.apache.org/licenses/LICENSE-2.0\n Unless required by applicable law or agreed to in writing, software\n distributed under the License is distributed on an \"AS IS\" BASIS,\n WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n See the License for the specific language governing permissions and\n limitations under the License.\n\"\"\"\n\nimport torch.nn as nn\n\nfrom model.blocks.shared_blocks import make_activation\n\n\nclass Bottleneck(nn.Module):\n    expansion = 4\n\n    def __init__(self, inplanes, planes, stride=1, downsample=None, activation=nn.ReLU):\n        super(Bottleneck, self).__init__()\n        self.conv1 = nn.Conv2d(inplanes, planes, kernel_size=1, bias=False)\n        self.bn1 = nn.BatchNorm2d(planes)\n        self.act1 = make_activation(activation)\n\n        self.conv2 = nn.Conv2d(planes, planes, kernel_size=3, stride=stride, padding=1, bias=False)\n        self.bn2 = nn.BatchNorm2d(planes)\n        self.act2 = make_activation(activation)\n\n        self.conv3 = nn.Conv2d(planes, planes * self.expansion, kernel_size=1, bias=False)\n        self.bn3 = nn.BatchNorm2d(planes * self.expansion)\n        self.act3 = make_activation(activation)\n\n        self.downsample = downsample\n        self.stride = stride\n\n    def forward(self, x):\n        residual = x\n\n        out = self.conv1(x)\n        out = self.bn1(out)\n        out = self.act1(out)\n\n        out = self.conv2(out)\n        out = self.bn2(out)\n        out = self.act2(out)\n\n        out = self.conv3(out)\n        out = self.bn3(out)\n\n        if self.downsample is not None:\n            residual = self.downsample(x)\n\n        out += residual\n        out = self.act3(out)\n\n        return out\n\n\nclass BasicBlock(nn.Module):\n    expansion = 1\n\n    def __init__(self, inplanes, planes, stride=1, downsample=None, activation=nn.ReLU):\n        super(BasicBlock, self).__init__()\n        self.conv1 = nn.Conv2d(inplanes, planes, kernel_size=3, stride=stride, padding=1, bias=False)\n        self.bn1 = nn.BatchNorm2d(planes)\n        self.relu = make_activation(activation)\n        self.conv2 = nn.Conv2d(inplanes, planes, kernel_size=3, stride=1, padding=1, bias=False)\n        self.bn2 = nn.BatchNorm2d(planes)\n        self.downsample = downsample\n        self.stride = stride\n\n    def forward(self, x):\n        residual = x\n\n        out = self.conv1(x)\n        out = self.bn1(out)\n        out = self.relu(out)\n\n        out = self.conv2(out)\n        out = self.bn2(out)\n\n        if self.downsample is not None:\n            residual = self.downsample(x)\n\n        out += residual\n        out = self.relu(out)\n\n        return out\n"
  },
  {
    "path": "model/blocks/rmnet_blocks.py",
    "content": "\"\"\"\n Copyright (c) 2018 Intel Corporation\n Licensed under the Apache License, Version 2.0 (the \"License\");\n you may not use this file except in compliance with the License.\n You may obtain a copy of the License at\n      http://www.apache.org/licenses/LICENSE-2.0\n Unless required by applicable law or agreed to in writing, software\n distributed under the License is distributed on an \"AS IS\" BASIS,\n WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n See the License for the specific language governing permissions and\n limitations under the License.\n\"\"\"\n\nimport torch.nn as nn\nimport torch.nn.functional as F\n\nfrom model.blocks.shared_blocks import make_activation\n\n\nclass RMBlock(nn.Module):\n    def __init__(self, input_planes, squeeze_planes, output_planes, downsample=False, dropout_ratio=0.1,\n                 activation=nn.ELU):\n        super(RMBlock, self).__init__()\n        self.downsample = downsample\n        self.input_planes = input_planes\n        self.output_planes = output_planes\n\n        self.squeeze_conv = nn.Conv2d(input_planes, squeeze_planes, kernel_size=1, bias=False)\n        self.squeeze_bn = nn.BatchNorm2d(squeeze_planes)\n\n        self.dw_conv = nn.Conv2d(squeeze_planes, squeeze_planes, groups=squeeze_planes, kernel_size=3, padding=1,\n                                 stride=2 if downsample else 1, bias=False)\n        self.dw_bn = nn.BatchNorm2d(squeeze_planes)\n\n        self.expand_conv = nn.Conv2d(squeeze_planes, output_planes, kernel_size=1, bias=False)\n        self.expand_bn = nn.BatchNorm2d(output_planes)\n\n        self.activation = make_activation(activation)\n\n        self.dropout_ratio = dropout_ratio\n\n        if self.downsample:\n            self.skip_conv = nn.Conv2d(input_planes, output_planes, kernel_size=1, bias=False)\n            self.skip_conv_bn = nn.BatchNorm2d(output_planes)\n\n        self.init_weights()\n\n    def init_weights(self):\n        for m in self.children():\n            if isinstance(m, nn.Conv2d):\n                nn.init.kaiming_normal_(m.weight, mode='fan_out')\n            elif isinstance(m, nn.BatchNorm2d):\n                nn.init.constant_(m.weight, 1)\n                nn.init.constant_(m.bias, 0)\n\n    def forward(self, x):\n        residual = x\n        out = self.activation(self.squeeze_bn(self.squeeze_conv(x)))\n        out = self.activation(self.dw_bn(self.dw_conv(out)))\n        out = self.expand_bn(self.expand_conv(out))\n        if self.dropout_ratio > 0:\n            out = F.dropout(out, p=self.dropout_ratio, training=self.training, inplace=True)\n        if self.downsample:\n            residual = F.max_pool2d(x, kernel_size=2, stride=2, padding=0)\n            residual = self.skip_conv(residual)\n            residual = self.skip_conv_bn(residual)\n        out += residual\n        return self.activation(out)\n"
  },
  {
    "path": "model/blocks/se_resnet_blocks.py",
    "content": "\"\"\"\n Copyright (c) 2018 Intel Corporation\n Licensed under the Apache License, Version 2.0 (the \"License\");\n you may not use this file except in compliance with the License.\n You may obtain a copy of the License at\n      http://www.apache.org/licenses/LICENSE-2.0\n Unless required by applicable law or agreed to in writing, software\n distributed under the License is distributed on an \"AS IS\" BASIS,\n WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n See the License for the specific language governing permissions and\n limitations under the License.\n\"\"\"\n\nimport torch.nn as nn\n\nfrom model.blocks.shared_blocks import make_activation\n\n\nclass SEBottleneck(nn.Module):\n    expansion = 4\n\n    def __init__(self, inplanes, planes, stride=1, downsample=None, activation=nn.ReLU):\n        super(SEBottleneck, self).__init__()\n        self.conv1 = nn.Conv2d(inplanes, planes, kernel_size=1, bias=False)\n        self.bn1 = nn.BatchNorm2d(planes)\n\n        self.conv2 = nn.Conv2d(planes, planes, kernel_size=3, stride=stride, padding=1, bias=False)\n        self.bn2 = nn.BatchNorm2d(planes)\n\n        self.conv3 = nn.Conv2d(planes, planes * 4, kernel_size=1, bias=False)\n        self.bn3 = nn.BatchNorm2d(planes * 4)\n\n        self.relu = make_activation(activation)\n\n        # SE\n        self.global_pool = nn.AdaptiveAvgPool2d(1)\n        self.conv_down = nn.Conv2d(planes * 4, planes // 4, kernel_size=1, bias=False)\n        self.conv_up = nn.Conv2d(planes // 4, planes * 4, kernel_size=1, bias=False)\n        self.sig = nn.Sigmoid()\n\n        # Downsample\n        self.downsample = downsample\n        self.stride = stride\n\n    def forward(self, x):\n        residual = x\n\n        out = self.conv1(x)\n        out = self.bn1(out)\n        out = self.relu(out)\n\n        out = self.conv2(out)\n        out = self.bn2(out)\n        out = self.relu(out)\n\n        out = self.conv3(out)\n        out = self.bn3(out)\n\n        out1 = self.global_pool(out)\n        out1 = self.conv_down(out1)\n        out1 = self.relu(out1)\n        out1 = self.conv_up(out1)\n        out1 = self.sig(out1)\n\n        if self.downsample is not None:\n            residual = self.downsample(x)\n\n        res = out1 * out + residual\n        res = self.relu(res)\n\n        return res\n"
  },
  {
    "path": "model/blocks/se_resnext_blocks.py",
    "content": "\"\"\"\n Copyright (c) 2018 Intel Corporation\n Licensed under the Apache License, Version 2.0 (the \"License\");\n you may not use this file except in compliance with the License.\n You may obtain a copy of the License at\n      http://www.apache.org/licenses/LICENSE-2.0\n Unless required by applicable law or agreed to in writing, software\n distributed under the License is distributed on an \"AS IS\" BASIS,\n WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n See the License for the specific language governing permissions and\n limitations under the License.\n\"\"\"\n\nimport torch.nn as nn\n\nfrom model.blocks.shared_blocks import SELayer\n\n\nclass SEBottleneckX(nn.Module):\n    expansion = 4\n\n    def __init__(self, inplanes, planes, cardinality, stride=1, downsample=None):\n        super(SEBottleneckX, self).__init__()\n        self.conv1 = nn.Conv2d(inplanes, planes * 2, kernel_size=1, bias=False)\n        self.bn1 = nn.BatchNorm2d(planes * 2)\n\n        self.conv2 = nn.Conv2d(planes * 2, planes * 2, kernel_size=3, stride=stride,\n                               padding=1, groups=cardinality, bias=False)\n        self.bn2 = nn.BatchNorm2d(planes * 2)\n\n        self.conv3 = nn.Conv2d(planes * 2, planes * 4, kernel_size=1, bias=False)\n        self.bn3 = nn.BatchNorm2d(planes * 4)\n\n        self.selayer = SELayer(planes * 4, 16, nn.ReLU)\n\n        self.relu = nn.ReLU(inplace=True)\n        self.downsample = downsample\n        self.stride = stride\n\n    def forward(self, x):\n        residual = x\n\n        out = self.conv1(x)\n        out = self.bn1(out)\n        out = self.relu(out)\n\n        out = self.conv2(out)\n        out = self.bn2(out)\n        out = self.relu(out)\n\n        out = self.conv3(out)\n        out = self.bn3(out)\n\n        out = self.selayer(out)\n\n        if self.downsample is not None:\n            residual = self.downsample(x)\n\n        out += residual\n        out = self.relu(out)\n\n        return out\n"
  },
  {
    "path": "model/blocks/shared_blocks.py",
    "content": "\"\"\"\n Copyright (c) 2018 Intel Corporation\n Licensed under the Apache License, Version 2.0 (the \"License\");\n you may not use this file except in compliance with the License.\n You may obtain a copy of the License at\n      http://www.apache.org/licenses/LICENSE-2.0\n Unless required by applicable law or agreed to in writing, software\n distributed under the License is distributed on an \"AS IS\" BASIS,\n WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n See the License for the specific language governing permissions and\n limitations under the License.\n\"\"\"\n\nimport torch\nimport torch.nn as nn\n\n\ndef make_activation(activation):\n    \"\"\"Factory for activation functions\"\"\"\n    if activation != nn.PReLU:\n        return activation(inplace=True)\n\n    return activation()\n\n\nclass SELayer(nn.Module):\n    \"\"\"Implementation of the Squeeze-Excitaion layer from https://arxiv.org/abs/1709.01507\"\"\"\n    def __init__(self, inplanes, squeeze_ratio=8, activation=nn.PReLU, size=None):\n        super(SELayer, self).__init__()\n        assert squeeze_ratio >= 1\n        assert inplanes > 0\n        if size is not None:\n            self.global_avgpool = nn.AvgPool2d(size)\n        else:\n            self.global_avgpool = nn.AdaptiveAvgPool2d(1)\n        self.conv1 = nn.Conv2d(inplanes, int(inplanes / squeeze_ratio), kernel_size=1, stride=1)\n        self.conv2 = nn.Conv2d(int(inplanes / squeeze_ratio), inplanes, kernel_size=1, stride=1)\n        self.relu = make_activation(activation)\n        self.sigmoid = nn.Sigmoid()\n\n    def forward(self, x):\n        out = self.global_avgpool(x)\n        out = self.conv1(out)\n        out = self.relu(out)\n        out = self.conv2(out)\n        out = self.sigmoid(out)\n        return x * out\n\n\nclass ScaleFilter(nn.Module):\n    \"\"\"Implementaion of the ScaleFilter regularizer\"\"\"\n    def __init__(self, q):\n        super(ScaleFilter, self).__init__()\n        assert 0 < q < 1\n        self.q = q\n\n    def forward(self, x):\n        if not self.training:\n            return x\n\n        scale_factors = 1. + self.q \\\n                - 2*self.q*torch.rand(x.shape[1], 1, 1, dtype=torch.float32, requires_grad=False).to(x.device)\n        return x * scale_factors\n"
  },
  {
    "path": "model/blocks/shufflenet_v2_blocks.py",
    "content": "\"\"\"\n Copyright (c) 2018 Intel Corporation\n Licensed under the Apache License, Version 2.0 (the \"License\");\n you may not use this file except in compliance with the License.\n You may obtain a copy of the License at\n      http://www.apache.org/licenses/LICENSE-2.0\n Unless required by applicable law or agreed to in writing, software\n distributed under the License is distributed on an \"AS IS\" BASIS,\n WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n See the License for the specific language governing permissions and\n limitations under the License.\n\"\"\"\n\nimport torch\nimport torch.nn as nn\nfrom model.blocks.shared_blocks import make_activation\n\n\ndef conv_bn(inp, oup, stride, activation=nn.ReLU):\n    conv = nn.Sequential(\n        nn.Conv2d(inp, oup, 3, stride, 1, bias=False),\n        nn.BatchNorm2d(oup),\n        make_activation(activation)\n    )\n    nn.init.kaiming_normal_(conv[0].weight, mode='fan_out')\n    return conv\n\n\ndef conv_1x1_bn(inp, oup, activation=nn.ReLU):\n    conv = nn.Sequential(\n        nn.Conv2d(inp, oup, 1, 1, 0, bias=False),\n        nn.BatchNorm2d(oup),\n        make_activation(activation)\n    )\n    nn.init.kaiming_normal_(conv[0].weight, mode='fan_out')\n    return conv\n\n\ndef channel_shuffle(x, groups):\n    batchsize, num_channels, height, width = x.data.size()\n    channels_per_group = num_channels // groups\n    # reshape\n    x = x.view(batchsize, groups, channels_per_group, height, width)\n    x = torch.transpose(x, 1, 2).contiguous()\n    # flatten\n    x = x.view(batchsize, -1, height, width)\n    return x\n\n\nclass ShuffleInvertedResidual(nn.Module):\n    def __init__(self, inp, oup, stride, benchmodel, activation=nn.ReLU):\n        super(ShuffleInvertedResidual, self).__init__()\n        self.benchmodel = benchmodel\n        self.stride = stride\n        assert stride in [1, 2]\n\n        oup_inc = oup//2\n\n        if self.benchmodel == 1:\n            # assert inp == oup_inc\n            self.branch2 = nn.Sequential(\n                # pw\n                nn.Conv2d(oup_inc, oup_inc, 1, 1, 0, bias=False),\n                nn.BatchNorm2d(oup_inc),\n                make_activation(activation),\n                # dw\n                nn.Conv2d(oup_inc, oup_inc, 3, stride, 1, groups=oup_inc, bias=False),\n                nn.BatchNorm2d(oup_inc),\n                # pw-linear\n                nn.Conv2d(oup_inc, oup_inc, 1, 1, 0, bias=False),\n                nn.BatchNorm2d(oup_inc),\n                make_activation(activation),\n            )\n        else:\n            self.branch1 = nn.Sequential(\n                # dw\n                nn.Conv2d(inp, inp, 3, stride, 1, groups=inp, bias=False),\n                nn.BatchNorm2d(inp),\n                # pw-linear\n                nn.Conv2d(inp, oup_inc, 1, 1, 0, bias=False),\n                nn.BatchNorm2d(oup_inc),\n                make_activation(activation),\n            )\n\n            self.branch2 = nn.Sequential(\n                # pw\n                nn.Conv2d(inp, oup_inc, 1, 1, 0, bias=False),\n                nn.BatchNorm2d(oup_inc),\n                make_activation(activation),\n                # dw\n                nn.Conv2d(oup_inc, oup_inc, 3, stride, 1, groups=oup_inc, bias=False),\n                nn.BatchNorm2d(oup_inc),\n                # pw-linear\n                nn.Conv2d(oup_inc, oup_inc, 1, 1, 0, bias=False),\n                nn.BatchNorm2d(oup_inc),\n                make_activation(activation),\n            )\n        self.init_weights()\n\n    @staticmethod\n    def _concat(x, out):\n        # concatenate along channel axis\n        return torch.cat((x, out), 1)\n\n    def init_weights(self):\n        for m in self.children():\n            if isinstance(m, nn.Conv2d):\n                nn.init.kaiming_normal_(m.weight, mode='fan_out')\n            elif isinstance(m, nn.BatchNorm2d):\n                nn.init.constant_(m.weight, 1)\n                nn.init.constant_(m.bias, 0)\n\n    def forward(self, x):\n        if self.benchmodel == 1:\n            x1 = x[:, :(x.shape[1]//2), :, :]\n            x2 = x[:, (x.shape[1]//2):, :, :]\n            out = self._concat(x1, self.branch2(x2))\n        elif self.benchmodel == 2:\n            out = self._concat(self.branch1(x), self.branch2(x))\n\n        return channel_shuffle(out, 2)\n"
  },
  {
    "path": "model/common.py",
    "content": "\"\"\"\n Copyright (c) 2018 Intel Corporation\n Licensed under the Apache License, Version 2.0 (the \"License\");\n you may not use this file except in compliance with the License.\n You may obtain a copy of the License at\n      http://www.apache.org/licenses/LICENSE-2.0\n Unless required by applicable law or agreed to in writing, software\n distributed under the License is distributed on an \"AS IS\" BASIS,\n WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n See the License for the specific language governing permissions and\n limitations under the License.\n\"\"\"\nfrom abc import abstractmethod\nimport torch.nn as nn\n\n\nclass ModelInterface(nn.Module):\n    \"\"\"Abstract class for models\"\"\"\n\n    @abstractmethod\n    def set_dropout_ratio(self, ratio):\n        \"\"\"Sets dropout ratio of the model\"\"\"\n\n    @abstractmethod\n    def get_input_res(self):\n        \"\"\"Returns input resolution\"\"\"\n\n\nfrom .rmnet_angular import RMNetAngular\nfrom .mobilefacenet import MobileFaceNet\nfrom .landnet import LandmarksNet\nfrom .resnet_angular import ResNetAngular\nfrom .se_resnet_angular import SEResNetAngular\nfrom .shufflenet_v2_angular import ShuffleNetV2Angular\n\n\nmodels_backbones = {'rmnet': RMNetAngular, 'mobilenet': MobileFaceNet, 'resnet': ResNetAngular,\n                    'shufflenetv2': ShuffleNetV2Angular, 'se_resnet': SEResNetAngular}\n\nmodels_landmarks = {'landnet': LandmarksNet}\n"
  },
  {
    "path": "model/landnet.py",
    "content": "\"\"\"\n Copyright (c) 2018 Intel Corporation\n Licensed under the Apache License, Version 2.0 (the \"License\");\n you may not use this file except in compliance with the License.\n You may obtain a copy of the License at\n      http://www.apache.org/licenses/LICENSE-2.0\n Unless required by applicable law or agreed to in writing, software\n distributed under the License is distributed on an \"AS IS\" BASIS,\n WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n See the License for the specific language governing permissions and\n limitations under the License.\n\"\"\"\n\nimport torch.nn as nn\n\nfrom .common import ModelInterface\n\n\nclass LandmarksNet(ModelInterface):\n    \"\"\"Facial landmarks localization network\"\"\"\n    def __init__(self):\n        super(LandmarksNet, self).__init__()\n        self.bn_first = nn.BatchNorm2d(3)\n        activation = nn.PReLU\n        self.landnet = nn.Sequential(\n            nn.Conv2d(3, 16, kernel_size=3, padding=1),\n            activation(),\n            nn.MaxPool2d(2, stride=2),\n            nn.BatchNorm2d(16),\n            nn.Conv2d(16, 32, kernel_size=3, padding=1),\n            activation(),\n            nn.MaxPool2d(2, stride=2),\n            nn.BatchNorm2d(32),\n            nn.Conv2d(32, 64, kernel_size=3, padding=1),\n            activation(),\n            nn.MaxPool2d(2, stride=2),\n            nn.BatchNorm2d(64),\n            nn.Conv2d(64, 64, kernel_size=3, padding=1),\n            activation(),\n            nn.BatchNorm2d(64),\n            nn.Conv2d(64, 128, kernel_size=3, padding=1),\n            activation(),\n            nn.BatchNorm2d(128)\n        )\n        # dw pooling\n        self.bottleneck_size = 256\n        self.pool = nn.Sequential(\n            nn.Conv2d(128, 128, kernel_size=6, padding=0, groups=128),\n            activation(),\n            nn.BatchNorm2d(128),\n            nn.Conv2d(128, self.bottleneck_size, kernel_size=1, padding=0),\n            activation(),\n            nn.BatchNorm2d(self.bottleneck_size),\n        )\n        # Regressor for 5 landmarks (10 coordinates)\n        self.fc_loc = nn.Sequential(\n            nn.Conv2d(self.bottleneck_size, 64, kernel_size=1),\n            activation(),\n            nn.Conv2d(64, 10, kernel_size=1),\n            nn.Sigmoid()\n        )\n\n    def forward(self, x):\n        xs = self.landnet(self.bn_first(x))\n        xs = self.pool(xs)\n        xs = self.fc_loc(xs)\n        return xs\n\n    def get_input_res(self):\n        return 48, 48\n\n    def set_dropout_ratio(self, ratio):\n        pass\n"
  },
  {
    "path": "model/mobilefacenet.py",
    "content": "\"\"\"\n Copyright (c) 2018 Intel Corporation\n Licensed under the Apache License, Version 2.0 (the \"License\");\n you may not use this file except in compliance with the License.\n You may obtain a copy of the License at\n      http://www.apache.org/licenses/LICENSE-2.0\n Unless required by applicable law or agreed to in writing, software\n distributed under the License is distributed on an \"AS IS\" BASIS,\n WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n See the License for the specific language governing permissions and\n limitations under the License.\n\"\"\"\n\nimport math\nimport torch.nn as nn\n\nfrom losses.am_softmax import AngleSimpleLinear\nfrom model.blocks.mobilenet_v2_blocks import InvertedResidual\nfrom model.blocks.shared_blocks import make_activation\nfrom .common import ModelInterface\n\n\ndef init_block(in_channels, out_channels, stride, activation=nn.PReLU):\n    \"\"\"Builds the first block of the MobileFaceNet\"\"\"\n    return nn.Sequential(\n        nn.BatchNorm2d(3),\n        nn.Conv2d(in_channels, out_channels, 3, stride, 1, bias=False),\n        nn.BatchNorm2d(out_channels),\n        make_activation(activation)\n    )\n\n\nclass MobileFaceNet(ModelInterface):\n    \"\"\"Implements modified MobileFaceNet from https://arxiv.org/abs/1804.07573\"\"\"\n    def __init__(self, embedding_size=128, num_classes=1, width_multiplier=1., feature=True):\n        super(MobileFaceNet, self).__init__()\n        assert embedding_size > 0\n        assert num_classes > 0\n        assert width_multiplier > 0\n        self.feature = feature\n\n        # Set up of inverted residual blocks\n        inverted_residual_setting = [\n            # t, c, n, s\n            [2, 64, 5, 2],\n            [4, 128, 1, 2],\n            [2, 128, 6, 1],\n            [4, 128, 1, 2],\n            [2, 128, 2, 1]\n        ]\n\n        first_channel_num = 64\n        last_channel_num = 512\n        self.features = [init_block(3, first_channel_num, 2)]\n\n        self.features.append(nn.Conv2d(first_channel_num, first_channel_num, 3, 1, 1,\n                                       groups=first_channel_num, bias=False))\n        self.features.append(nn.BatchNorm2d(64))\n        self.features.append(nn.PReLU())\n\n        # Inverted Residual Blocks\n        in_channel_num = first_channel_num\n        size_h, size_w = MobileFaceNet.get_input_res()\n        size_h, size_w = size_h // 2, size_w // 2\n        for t, c, n, s in inverted_residual_setting:\n            output_channel = int(c * width_multiplier)\n            for i in range(n):\n                if i == 0:\n                    size_h, size_w = size_h // s, size_w // s\n                    self.features.append(InvertedResidual(in_channel_num, output_channel,\n                                                          s, t, outp_size=(size_h, size_w)))\n                else:\n                    self.features.append(InvertedResidual(in_channel_num, output_channel,\n                                                          1, t, outp_size=(size_h, size_w)))\n                in_channel_num = output_channel\n\n        # 1x1 expand block\n        self.features.append(nn.Sequential(nn.Conv2d(in_channel_num, last_channel_num, 1, 1, 0, bias=False),\n                                           nn.BatchNorm2d(last_channel_num),\n                                           nn.PReLU()))\n        self.features = nn.Sequential(*self.features)\n\n        # Depth-wise pooling\n        k_size = (MobileFaceNet.get_input_res()[0] // 16, MobileFaceNet.get_input_res()[1] // 16)\n        self.dw_pool = nn.Conv2d(last_channel_num, last_channel_num, k_size,\n                                 groups=last_channel_num, bias=False)\n        self.dw_bn = nn.BatchNorm2d(last_channel_num)\n        self.conv1_extra = nn.Conv2d(last_channel_num, embedding_size, 1, stride=1, padding=0, bias=False)\n\n        if not self.feature:\n            self.fc_angular = AngleSimpleLinear(embedding_size, num_classes)\n\n        self.init_weights()\n\n    def forward(self, x):\n        x = self.features(x)\n        x = self.dw_bn(self.dw_pool(x))\n        x = self.conv1_extra(x)\n\n        if self.feature or not self.training:\n            return x\n\n        x = x.view(x.size(0), -1)\n        y = self.fc_angular(x)\n\n        return x, y\n\n    @staticmethod\n    def get_input_res():\n        return 128, 128\n\n    def set_dropout_ratio(self, ratio):\n        assert 0 <= ratio < 1.\n\n    def init_weights(self):\n        \"\"\"Initializes weights of the model before training\"\"\"\n        for m in self.modules():\n            if isinstance(m, nn.Conv2d):\n                n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels\n                m.weight.data.normal_(0, math.sqrt(2. / n))\n                if m.bias is not None:\n                    m.bias.data.zero_()\n            elif isinstance(m, nn.BatchNorm2d):\n                m.weight.data.fill_(1)\n                m.bias.data.zero_()\n            elif isinstance(m, nn.Linear):\n                n = m.weight.size(1)\n                m.weight.data.normal_(0, 0.01)\n                m.bias.data.zero_()\n"
  },
  {
    "path": "model/resnet_angular.py",
    "content": "\"\"\"\n Copyright (c) 2018 Intel Corporation\n Licensed under the Apache License, Version 2.0 (the \"License\");\n you may not use this file except in compliance with the License.\n You may obtain a copy of the License at\n      http://www.apache.org/licenses/LICENSE-2.0\n Unless required by applicable law or agreed to in writing, software\n distributed under the License is distributed on an \"AS IS\" BASIS,\n WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n See the License for the specific language governing permissions and\n limitations under the License.\n\"\"\"\n\nimport torch.nn as nn\nfrom losses.am_softmax import AngleSimpleLinear\nfrom model.backbones.resnet import resnet50\nfrom .common import ModelInterface\n\n\nclass ResNetAngular(ModelInterface):\n    \"\"\"Face reid head for the ResNet architecture\"\"\"\n    def __init__(self, embedding_size=128, num_classes=0, feature=True):\n        super(ResNetAngular, self).__init__()\n\n        self.bn_first = nn.BatchNorm2d(3)\n        self.feature = feature\n        self.model = resnet50(num_classes=embedding_size, activation=nn.PReLU)\n        self.embedding_size = embedding_size\n\n        if not self.feature:\n            self.fc_angular = AngleSimpleLinear(self.embedding_size, num_classes)\n\n    def forward(self, x):\n\n        x = self.bn_first(x)\n        x = self.model(x)\n\n        if self.feature or not self.training:\n            return x\n\n        x = x.view(x.size(0), -1)\n        y = self.fc_angular(x)\n\n        return x, y\n\n    @staticmethod\n    def get_input_res():\n        return 112, 112\n\n    def set_dropout_ratio(self, ratio):\n        assert 0 <= ratio < 1.\n"
  },
  {
    "path": "model/rmnet_angular.py",
    "content": "\"\"\"\n Copyright (c) 2018 Intel Corporation\n Licensed under the Apache License, Version 2.0 (the \"License\");\n you may not use this file except in compliance with the License.\n You may obtain a copy of the License at\n      http://www.apache.org/licenses/LICENSE-2.0\n Unless required by applicable law or agreed to in writing, software\n distributed under the License is distributed on an \"AS IS\" BASIS,\n WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n See the License for the specific language governing permissions and\n limitations under the License.\n\"\"\"\n\nimport torch.nn as nn\n\nfrom losses.am_softmax import AngleSimpleLinear\nfrom model.backbones.rmnet import RMNetBody\nfrom model.blocks.rmnet_blocks import RMBlock\nfrom .common import ModelInterface\n\n\nclass RMNetAngular(ModelInterface):\n    \"\"\"Face reid head for the ResMobNet architecture. See https://arxiv.org/pdf/1812.02465.pdf for details\n    about the ResMobNet backbone.\"\"\"\n    def __init__(self, embedding_size, num_classes=0, feature=True, body=RMNetBody):\n        super(RMNetAngular, self).__init__()\n        self.feature = feature\n        self.backbone = body()\n        self.global_pooling = nn.MaxPool2d((8, 8))\n        self.conv1_extra = nn.Conv2d(256, embedding_size, 1, stride=1, padding=0, bias=False)\n        if not feature:\n            self.fc_angular = AngleSimpleLinear(embedding_size, num_classes)\n\n    def forward(self, x):\n        x = self.backbone(x)\n        x = self.global_pooling(x)\n        x = self.conv1_extra(x)\n\n        if self.feature or not self.training:\n            return x\n\n        x = x.view(x.size(0), -1)\n        y = self.fc_angular(x)\n\n        return x, y\n\n    def set_dropout_ratio(self, ratio):\n        assert 0 <= ratio < 1.\n\n        for m in self.backbone.modules():\n            if isinstance(m, RMBlock):\n                m.dropout_ratio = ratio\n\n    @staticmethod\n    def get_input_res():\n        return 128, 128\n"
  },
  {
    "path": "model/se_resnet_angular.py",
    "content": "\"\"\"\n Copyright (c) 2018 Intel Corporation\n Licensed under the Apache License, Version 2.0 (the \"License\");\n you may not use this file except in compliance with the License.\n You may obtain a copy of the License at\n      http://www.apache.org/licenses/LICENSE-2.0\n Unless required by applicable law or agreed to in writing, software\n distributed under the License is distributed on an \"AS IS\" BASIS,\n WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n See the License for the specific language governing permissions and\n limitations under the License.\n\"\"\"\n\nimport torch.nn as nn\n\nfrom losses.am_softmax import AngleSimpleLinear\nfrom model.backbones.se_resnet import se_resnet50\nfrom .common import ModelInterface\n\n\nclass SEResNetAngular(ModelInterface):\n    \"\"\"Face reid head for the SE ResNet architecture\"\"\"\n    def __init__(self, embedding_size=128, num_classes=0, feature=True):\n        super(SEResNetAngular, self).__init__()\n\n        self.bn_first = nn.BatchNorm2d(3)\n        self.feature = feature\n        self.model = se_resnet50(num_classes=embedding_size, activation=nn.PReLU)\n        self.embedding_size = embedding_size\n\n        if not self.feature:\n            self.fc_angular = AngleSimpleLinear(self.embedding_size, num_classes)\n\n    def forward(self, x):\n        x = self.bn_first(x)\n        x = self.model(x)\n\n        if self.feature or not self.training:\n            return x\n\n        x = x.view(x.size(0), -1)\n        y = self.fc_angular(x)\n\n        return x, y\n\n    @staticmethod\n    def get_input_res():\n        return 112, 112\n\n    def set_dropout_ratio(self, ratio):\n        assert 0 <= ratio < 1.\n"
  },
  {
    "path": "model/shufflenet_v2_angular.py",
    "content": "\"\"\"\n Copyright (c) 2018 Intel Corporation\n Licensed under the Apache License, Version 2.0 (the \"License\");\n you may not use this file except in compliance with the License.\n You may obtain a copy of the License at\n      http://www.apache.org/licenses/LICENSE-2.0\n Unless required by applicable law or agreed to in writing, software\n distributed under the License is distributed on an \"AS IS\" BASIS,\n WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n See the License for the specific language governing permissions and\n limitations under the License.\n\"\"\"\n\nimport torch.nn as nn\n\nfrom losses.am_softmax import AngleSimpleLinear\nfrom model.backbones.shufflenet_v2 import ShuffleNetV2Body\nfrom .common import ModelInterface\n\n\nclass ShuffleNetV2Angular(ModelInterface):\n    \"\"\"Face reid head for the ShuffleNetV2 architecture\"\"\"\n    def __init__(self, embedding_size, num_classes=0, feature=True, body=ShuffleNetV2Body, **kwargs):\n        super(ShuffleNetV2Angular, self).__init__()\n        self.feature = feature\n        kwargs['input_size'] = ShuffleNetV2Angular.get_input_res()[0]\n        kwargs['width_mult'] = 1.\n        self.backbone = body(**kwargs)\n        k_size = int(kwargs['input_size'] / self.backbone.get_downscale_factor())\n        self.global_pool = nn.Conv2d(self.backbone.stage_out_channels[-1], self.backbone.stage_out_channels[-1],\n                                     (k_size, k_size), groups=self.backbone.stage_out_channels[-1], bias=False)\n        self.conv1_extra = nn.Conv2d(self.backbone.get_num_output_channels(), embedding_size, 1, padding=0, bias=False)\n        if not feature:\n            self.fc_angular = AngleSimpleLinear(embedding_size, num_classes)\n\n    def forward(self, x):\n        x = self.backbone(x)\n        x = self.global_pool(x)\n        x = self.conv1_extra(x)\n\n        if self.feature or not self.training:\n            return x\n\n        x = x.view(x.size(0), -1)\n        y = self.fc_angular(x)\n\n        return x, y\n\n    def set_dropout_ratio(self, ratio):\n        assert 0 <= ratio < 1.\n\n    @staticmethod\n    def get_input_res():\n        res = 128\n        return res, res\n"
  },
  {
    "path": "requirements.txt",
    "content": "glog==0.3.1\nnumpy==1.15.4\nopencv-python==3.4.4.19\nPillow==5.3.0\nprotobuf==3.6.1\npython-gflags==3.1.2\nscipy==1.1.0\nsix==1.11.0\ntensorboardX==1.4\ntorch==0.4.1\ntorchvision==0.2.1\ntqdm==4.28.1\npyyaml>=3.12\nptflops==0.1\n"
  },
  {
    "path": "scripts/__init__.py",
    "content": ""
  },
  {
    "path": "scripts/accuracy_check.py",
    "content": "\"\"\"\n Copyright (c) 2018 Intel Corporation\n Licensed under the Apache License, Version 2.0 (the \"License\");\n you may not use this file except in compliance with the License.\n You may obtain a copy of the License at\n      http://www.apache.org/licenses/LICENSE-2.0\n Unless required by applicable law or agreed to in writing, software\n distributed under the License is distributed on an \"AS IS\" BASIS,\n WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n See the License for the specific language governing permissions and\n limitations under the License.\n\"\"\"\n\nimport argparse\nimport glog as log\n\nimport numpy as np\nimport torch\nfrom tqdm import tqdm\nimport cv2 as cv\n\nfrom utils.utils import load_model_state\nfrom utils.ie_tools import load_ie_model\nfrom model.common import models_backbones, models_landmarks\n\ndef main():\n    \"\"\"Runs the accuracy check\"\"\"\n    parser = argparse.ArgumentParser(description='Accuracy check script (pt vs caffe)')\n    parser.add_argument('--embed_size', type=int, default=128, help='Size of the face embedding.')\n    parser.add_argument('--snap', type=str, required=True, help='Snapshot to convert.')\n    parser.add_argument('--device', '-d', default=0, type=int, help='Device for model placement.')\n    parser.add_argument('--model', choices=list(models_backbones.keys()) + list(models_landmarks.keys()), type=str,\n                        default='rmnet')\n\n    # IE-related options\n    parser.add_argument('--ie_model', type=str, required=True)\n    parser.add_argument(\"-l\", \"--cpu_extension\",\n                        help=\"MKLDNN (CPU)-targeted custom layers.Absolute path to a shared library with the kernels \"\n                             \"impl.\", type=str, default=None)\n    parser.add_argument(\"-pp\", \"--plugin_dir\", help=\"Path to a plugin folder\", type=str, default=None)\n    parser.add_argument(\"-d_ie\", \"--device_ie\",\n                        help=\"Specify the target device to infer on; CPU, GPU, FPGA or MYRIAD is acceptable. Sample \"\n                             \"will look for a suitable plugin for device specified (CPU by default)\", default=\"CPU\",\n                        type=str)\n\n    args = parser.parse_args()\n\n    max_err = 0.\n    with torch.cuda.device(args.device):\n        if args.model in models_landmarks.keys():\n            pt_model = models_landmarks[args.model]\n        else:\n            pt_model = models_backbones[args.model](embedding_size=args.embed_size, feature=True)\n        pt_model = load_model_state(pt_model, args.snap, args.device)\n\n        ie_model = load_ie_model(args.ie_model, args.device_ie, args.plugin_dir, args.cpu_extension)\n        np.random.seed(0)\n\n        for _ in tqdm(range(100)):\n            input_img = np.random.randint(0, high=255, size=(*pt_model.get_input_res(), 3), dtype=np.uint8)\n            input_bgr = cv.cvtColor(input_img, cv.COLOR_BGR2RGB)\n\n            input_pt = torch.unsqueeze(torch.from_numpy(input_img.transpose(2, 0, 1).astype('float32') / 255.).cuda(),\n                                       dim=0)\n            pt_output = (pt_model(input_pt)).data.cpu().numpy().reshape(1, -1)\n            ie_output = ie_model.forward(input_bgr).reshape(1, -1)\n\n            max_err = max(np.linalg.norm(pt_output - ie_output, np.inf), max_err)\n\n    log.info('Max l_inf error: %e', max_err)\n\nif __name__ == '__main__':\n    main()\n"
  },
  {
    "path": "scripts/align_images.py",
    "content": "\"\"\"\n Copyright (c) 2018 Intel Corporation\n Licensed under the Apache License, Version 2.0 (the \"License\");\n you may not use this file except in compliance with the License.\n You may obtain a copy of the License at\n      http://www.apache.org/licenses/LICENSE-2.0\n Unless required by applicable law or agreed to in writing, software\n distributed under the License is distributed on an \"AS IS\" BASIS,\n WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n See the License for the specific language governing permissions and\n limitations under the License.\n\"\"\"\n\nimport argparse\nimport os\nimport os.path as osp\nimport json\n\nimport cv2 as cv\nimport torch\nfrom tqdm import tqdm\nfrom torchvision.transforms import transforms\n\nfrom model import landnet\nfrom utils import utils\nfrom utils import augmentation\nfrom utils.face_align import FivePointsAligner\n\nclass LandnetPT:\n    \"\"\"Wrapper for landmarks regression model\"\"\"\n    def __init__(self, model):\n        self.net = model\n        self.transformer = transforms.Compose(\n            [augmentation.ResizeNumpy((48, 48)), augmentation.NumpyToTensor(switch_rb=True)])\n\n    def get_landmarks(self, batch):\n        converted_batch = []\n        for item in batch:\n            converted_batch.append(self.transformer(item))\n        pt_blob = torch.stack(converted_batch).cuda()\n        landmarks = self.net(pt_blob)\n        return landmarks.data.cpu().numpy()\n\n\nclass FaceDetector:\n    \"\"\"Wrapper class for face detector\"\"\"\n    def __init__(self, proto, model, conf=.6, expand_ratio=(1.1, 1.05), size=(300, 300)):\n        self.net = cv.dnn.readNetFromCaffe(proto, model)\n        self.net.setPreferableBackend(cv.dnn.DNN_BACKEND_DEFAULT)\n        self.net.setPreferableTarget(cv.dnn.DNN_TARGET_CPU)\n        last_layer_id = self.net.getLayerId(self.net.getLayerNames()[-1])\n        last_layer = self.net.getLayer(last_layer_id)\n        assert last_layer.type == 'DetectionOutput'\n\n        self.confidence = conf\n        self.expand_ratio = expand_ratio\n        self.det_res = size\n\n    def __decode_detections(self, out, frame_shape):\n        \"\"\"Decodes raw SSD output\"\"\"\n        frame_height = frame_shape[0]\n        frame_width = frame_shape[1]\n        detections = []\n\n        for detection in out[0, 0]:\n            confidence = detection[2]\n            if confidence > self.confidence:\n                left = int(max(detection[3], 0) * frame_width)\n                top = int(max(detection[4], 0) * frame_height)\n                right = int(max(detection[5], 0) * frame_width)\n                bottom = int(max(detection[6], 0) * frame_height)\n                if self.expand_ratio != (1., 1.):\n                    w = (right - left)\n                    h = (bottom - top)\n                    dw = w * (self.expand_ratio[0] - 1.) / 2\n                    dh = h * (self.expand_ratio[1] - 1.) / 2\n                    left = max(int(left - dw), 0)\n                    right = int(right + dw)\n                    top = max(int(top - dh), 0)\n                    bottom = int(bottom + dh)\n\n                # classId = int(detection[1]) - 1  # Skip background label\n                detections.append(((left, top, right, bottom), confidence))\n\n        if len(detections) > 1:\n            detections.sort(key=lambda x: x[1], reverse=True)\n\n        return detections\n\n    def get_detections(self, frame):\n        \"\"\"Returns all detections on frame\"\"\"\n        blob = cv.dnn.blobFromImage(frame, 1., (self.det_res[0], self.det_res[1]), crop=False)\n        self.net.setInput(blob)\n        out = self.net.forward()\n        detections = self.__decode_detections(out, frame.shape)\n        return detections\n\n\ndef draw_detections(frame, detections, landmarks):\n    \"\"\"Draw detections and landmarks on a frame\"\"\"\n    for _, rect in enumerate(detections):\n        left, top, right, bottom = rect\n        cv.rectangle(frame, (left, top), (right, bottom), (0, 255, 0), thickness=2)\n        for point in landmarks.reshape(-1, 2):\n            point = (int(left + point[0] * (right - left)), int(top + point[1] * (bottom - top)))\n            cv.circle(frame, point, 5, (255, 0, 0), -1)\n\n    return frame\n\n\ndef run_dumping(images_list, face_det, landmarks_regressor, vis_flag):\n    \"\"\"Dumps detections and landmarks from images\"\"\"\n    detected_num = 0\n    data = []\n    for path in tqdm(images_list, 'Dumping data'):\n        image = cv.imread(path, cv.IMREAD_COLOR)\n        if image is None:\n            continue\n\n        detections = face_det.get_detections(image)\n        landmarks = None\n        if detections:\n            left, top, right, bottom = detections[0][0]\n            roi = image[top:bottom, left:right]\n            landmarks = landmarks_regressor.get_landmarks([roi]).reshape(-1)\n            data.append({'path': path, 'bbox': detections[0][0], 'landmarks': landmarks})\n            detected_num += 1\n            if vis_flag:\n                FivePointsAligner.align(roi, landmarks,\n                                        d_size=(200,200), normalize=False, show=True)\n        else:\n            data.append({'path': path, 'bbox': None, 'landmarks': None})\n\n    print('Detection ratio: ', float(detected_num) / len(data))\n\n    return data\n\n\ndef create_images_list(images_root, imgs_list):\n    input_filenames = []\n    input_dir = os.path.abspath(images_root)\n\n    if imgs_list is None:\n        stop = False\n        for path, _, files in os.walk(input_dir):\n            if stop:\n                break\n            for name in files:\n                if name.lower().endswith('.jpg') or name.lower().endswith('.png') \\\n                        or name.lower().endswith('.jpeg') or name.lower().endswith('.gif') \\\n                        or not '.' in name:\n                    filename = os.path.join(path, name)\n                    input_filenames.append(filename)\n    else:\n        with open(imgs_list) as f:\n            data = json.load(f)\n            for path in data['path']:\n                filename = osp.join(images_root, path)\n                input_filenames.append(filename)\n\n    return input_filenames\n\n\ndef save_data(data, filename, root_dir):\n    print('Saving data...')\n    with open(filename, 'w') as f:\n        for instance in data:\n            line = osp.relpath(instance['path'], start=root_dir) + ' | '\n            if instance['bbox'] is not None:\n                for x in instance['landmarks']:\n                    line += str(x) + ' '\n                line += ' | '\n                left, top, right, bottom = instance['bbox']\n                line += str(left) + ' ' + str(top) + ' ' + str(right - left) + ' ' + str(bottom - top)\n\n            f.write(line.strip() + '\\n')\n\ndef main():\n    parser = argparse.ArgumentParser(description='')\n    parser.add_argument('--images_root', type=str, default=None, required=True)\n    parser.add_argument('--images_list', type=str, default=None, required=False)\n    parser.add_argument('--fd_proto', type=str, default='../demo/face_detector/deploy_fd.prototxt', help='')\n    parser.add_argument('--fd_model', type=str, default='../demo/face_detector/sq_300x300_iter_120000.caffemodel',\n                        help='')\n    parser.add_argument('--fr_thresh', type=float, default=0.1)\n    parser.add_argument('--det_res', type=int, nargs=2, default=[300, 300], help='Detection net input resolution.')\n    parser.add_argument('--landnet_model', type=str)\n    parser.add_argument('--device', type=int, default=0)\n    parser.add_argument('--visualize', action='store_true')\n    args = parser.parse_args()\n\n    face_detector = FaceDetector(args.fd_proto, args.fd_model, conf=args.fr_thresh, size=args.det_res)\n\n    with torch.cuda.device(args.device):\n        landmarks_regressor = utils.load_model_state(landnet.LandmarksNet(), args.landnet_model, args.device)\n        data = run_dumping(create_images_list(args.images_root, args.images_list), face_detector,\n                           LandnetPT(landmarks_regressor), args.visualize)\n        save_data(data, osp.join(args.images_root, 'list.txt'), args.images_root)\n\nif __name__ == '__main__':\n    main()\n"
  },
  {
    "path": "scripts/count_flops.py",
    "content": "\"\"\"\n Copyright (c) 2018 Intel Corporation\n Licensed under the Apache License, Version 2.0 (the \"License\");\n you may not use this file except in compliance with the License.\n You may obtain a copy of the License at\n      http://www.apache.org/licenses/LICENSE-2.0\n Unless required by applicable law or agreed to in writing, software\n distributed under the License is distributed on an \"AS IS\" BASIS,\n WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n See the License for the specific language governing permissions and\n limitations under the License.\n\"\"\"\n\nimport argparse\nimport torch\n\nfrom model.common import models_backbones, models_landmarks\nfrom ptflops import get_model_complexity_info\n\n\ndef main():\n    \"\"\"Runs flops counter\"\"\"\n    parser = argparse.ArgumentParser(description='Evaluation script for Face Recognition in PyTorch')\n    parser.add_argument('--embed_size', type=int, default=128, help='Size of the face embedding.')\n    parser.add_argument('--model', choices=list(models_backbones.keys()) + list(models_landmarks.keys()), type=str,\n                        default='rmnet')\n    args = parser.parse_args()\n\n    with torch.no_grad():\n        if args.model in models_landmarks.keys():\n            model = models_landmarks[args.model]()\n        else:\n            model = models_backbones[args.model](embedding_size=args.embed_size, feature=True)\n\n        flops, params = get_model_complexity_info(model, model.get_input_res(),\n                                                  as_strings=True, print_per_layer_stat=True)\n        print('Flops:  {}'.format(flops))\n        print('Params: {}'.format(params))\n\n\nif __name__ == '__main__':\n    main()\n"
  },
  {
    "path": "scripts/matio.py",
    "content": "# pylint: skip-file\nimport struct\nimport numpy as np\n\ncv_type_to_dtype = {\n\t5 : np.dtype('float32'),\n\t6 : np.dtype('float64')\n}\n\ndtype_to_cv_type = {v : k for k,v in cv_type_to_dtype.items()}\n\ndef write_mat(f, m):\n    \"\"\"Write mat m to file f\"\"\"\n    if len(m.shape) == 1:\n        rows = m.shape[0]\n        cols = 1\n    else:\n        rows, cols = m.shape\n    header = struct.pack('iiii', rows, cols, cols * 4, dtype_to_cv_type[m.dtype])\n    f.write(header)\n    f.write(m.data)\n\n\ndef read_mat(f):\n\t\"\"\"\n\tReads an OpenCV mat from the given file opened in binary mode\n\t\"\"\"\n\trows, cols, stride, type_ = struct.unpack('iiii', f.read(4*4))\n\tmat = np.fromstring(f.read(rows*stride),dtype=cv_type_to_dtype[type_])\n\treturn mat.reshape(rows,cols)\n\ndef read_mkl_vec(f):\n\t\"\"\"\n\tReads an OpenCV mat from the given file opened in binary mode\n\t\"\"\"\n\t# Read past the header information\n\tf.read(4*4)\n\n\tlength, stride, type_ = struct.unpack('iii', f.read(3*4))\n\tmat = np.fromstring(f.read(length*4),dtype=np.float32)\n\treturn mat\n\ndef load_mkl_vec(filename):\n\t\"\"\"\n\tReads a OpenCV Mat from the given filename\n\t\"\"\"\n\treturn read_mkl_vec(open(filename,'rb'))\n\ndef load_mat(filename):\n\t\"\"\"\n\tReads a OpenCV Mat from the given filename\n\t\"\"\"\n\treturn read_mat(open(filename,'rb'))\n\ndef save_mat(filename, m):\n    \"\"\"Saves mat m to the given filename\"\"\"\n    return write_mat(open(filename,'wb'), m)\n\ndef main():\n\tf = open('1_to_0.bin','rb')\n\tvx = read_mat(f)\n\tvy = read_mat(f)\n\nif __name__ == '__main__':\n\tmain()\n"
  },
  {
    "path": "scripts/plot_roc_curves_lfw.py",
    "content": "\"\"\"\n Copyright (c) 2018 Intel Corporation\n Licensed under the Apache License, Version 2.0 (the \"License\");\n you may not use this file except in compliance with the License.\n You may obtain a copy of the License at\n      http://www.apache.org/licenses/LICENSE-2.0\n Unless required by applicable law or agreed to in writing, software\n distributed under the License is distributed on an \"AS IS\" BASIS,\n WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n See the License for the specific language governing permissions and\n limitations under the License.\n\"\"\"\n\nimport argparse\n\nimport matplotlib.pyplot as plt\nfrom evaluate_lfw import get_auc\n\ndef main():\n    parser = argparse.ArgumentParser(description='')\n    parser.add_argument('rocs', metavar='ROCs', type=str, nargs='+',\n                        help='paths to roc curves')\n\n    args = parser.parse_args()\n\n    plt.xlabel(\"False Positive Rate\")\n    plt.ylabel(\"True Positive Rate\")\n    plt.grid(b=True, which='major', color='k', linestyle='-')\n    plt.grid(b=True, which='minor', color='k', linestyle='-', alpha=0.2)\n    plt.minorticks_on()\n\n    for curve_file in args.rocs:\n        fprs = []\n        tprs = []\n        with open(curve_file, 'r') as f:\n            for line in f.readlines():\n                values = line.strip().split()\n                fprs.append(float(values[1]))\n                tprs.append(float(values[0]))\n\n        curve_name = curve_file.split('/')[-1].split('.')[0]\n        plt.plot(fprs, tprs, label=curve_name)\n        plt.legend(loc='best', fontsize=10)\n\n        print('AUC for {}: {}'.format(curve_name, get_auc(fprs, tprs)))\n\n    plt.show()\n\nif __name__ == '__main__':\n    main()\n"
  },
  {
    "path": "scripts/pytorch2onnx.py",
    "content": "\"\"\"\n Copyright (c) 2018 Intel Corporation\n Licensed under the Apache License, Version 2.0 (the \"License\");\n you may not use this file except in compliance with the License.\n You may obtain a copy of the License at\n      http://www.apache.org/licenses/LICENSE-2.0\n Unless required by applicable law or agreed to in writing, software\n distributed under the License is distributed on an \"AS IS\" BASIS,\n WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n See the License for the specific language governing permissions and\n limitations under the License.\n\"\"\"\n\nimport argparse\nimport torch\n\nfrom utils.utils import load_model_state\nfrom model.common import models_backbones, models_landmarks\n\ndef main():\n    parser = argparse.ArgumentParser(description='Conversion script for FR models from PyTorch to ONNX')\n    parser.add_argument('--embed_size', type=int, default=128, help='Size of the face embedding.')\n    parser.add_argument('--snap', type=str, required=True, help='Snapshot to convert.')\n    parser.add_argument('--device', '-d', default=-1, type=int, help='Device for model placement.')\n    parser.add_argument('--output_dir', default='./', type=str, help='Output directory.')\n    parser.add_argument('--model', choices=list(models_backbones.keys()) + list(models_landmarks.keys()),\n                        type=str, default='rmnet')\n\n    args = parser.parse_args()\n\n    if args.model in models_landmarks.keys():\n        model = models_landmarks[args.model]()\n    else:\n        model = models_backbones[args.model](embedding_size=args.embed_size, feature=True)\n\n    model = load_model_state(model, args.snap, args.device, eval_state=True)\n    input_var = torch.rand(1, 3, *model.get_input_res())\n    dump_name = args.snap[args.snap.rfind('/') + 1:-3]\n\n    torch.onnx.export(model, input_var, dump_name + '.onnx', verbose=True, export_params=True)\n\nif __name__ == '__main__':\n    main()\n"
  },
  {
    "path": "tests/__init__.py",
    "content": ""
  },
  {
    "path": "tests/test_alignment.py",
    "content": "\"\"\"\n Copyright (c) 2018 Intel Corporation\n Licensed under the Apache License, Version 2.0 (the \"License\");\n you may not use this file except in compliance with the License.\n You may obtain a copy of the License at\n      http://www.apache.org/licenses/LICENSE-2.0\n Unless required by applicable law or agreed to in writing, software\n distributed under the License is distributed on an \"AS IS\" BASIS,\n WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n See the License for the specific language governing permissions and\n limitations under the License.\n\"\"\"\n\nimport unittest\nimport cv2 as cv\nimport numpy as np\n\nfrom utils.face_align import FivePointsAligner\nfrom utils.landmarks_augmentation import RandomRotate\n\n\nclass FaceAlignmentTests(unittest.TestCase):\n    \"\"\"Tests for alignment methods\"\"\"\n    def test_align_image(self):\n        \"\"\"Synthetic test for alignment function\"\"\"\n        image = np.zeros((128, 128, 3), dtype=np.float32)\n        for point in FivePointsAligner.ref_landmarks:\n            point_scaled = point * [128, 128]\n            cv.circle(image, tuple(point_scaled.astype(np.int)), 5, (255, 255, 255), cv.FILLED)\n\n        transform = RandomRotate(40., p=1.)\n        rotated_data = transform({'img': image, 'landmarks': FivePointsAligner.ref_landmarks})\n        aligned_image = FivePointsAligner.align(rotated_data['img'], \\\n                                                rotated_data['landmarks'].reshape(-1),\n                                                d_size=(128, 128), normalized=True)\n\n        for point in FivePointsAligner.ref_landmarks:\n            point_scaled = (point * [128, 128]).astype(np.int)\n            check_sum = np.mean(aligned_image[point_scaled[1] - 3 : point_scaled[1] + 3,\n                                              point_scaled[0] - 3 : point_scaled[0] + 3])\n            self.assertGreaterEqual(check_sum, 220)\n\nif __name__ == '__main__':\n    unittest.main()\n"
  },
  {
    "path": "tests/test_models.py",
    "content": "\"\"\"\n Copyright (c) 2018 Intel Corporation\n Licensed under the Apache License, Version 2.0 (the \"License\");\n you may not use this file except in compliance with the License.\n You may obtain a copy of the License at\n      http://www.apache.org/licenses/LICENSE-2.0\n Unless required by applicable law or agreed to in writing, software\n distributed under the License is distributed on an \"AS IS\" BASIS,\n WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n See the License for the specific language governing permissions and\n limitations under the License.\n\"\"\"\n\nimport unittest\nimport os\nimport torch\n\nfrom model.common import models_backbones, models_landmarks\nfrom utils.utils import save_model_cpu, load_model_state\n\n\nclass BackbonesTests(unittest.TestCase):\n    \"\"\"Tests for backbones\"\"\"\n    def test_output_shape(self):\n        \"\"\"Checks output shape\"\"\"\n        embed_size = 256\n        for model_type in models_backbones.values():\n            model = model_type(embedding_size=embed_size, feature=True).eval()\n            batch = torch.Tensor(1, 3, *model.get_input_res()).uniform_()\n            output = model(batch)\n            self.assertEqual(list(output.shape), list((1, embed_size, 1, 1)))\n\n    def test_save_load_snap(self):\n        \"\"\"Checks an ability to save and load model correctly\"\"\"\n        embed_size = 256\n        snap_name = os.path.join(os.getcwd(), 'test_snap.pt')\n        for model_type in models_backbones.values():\n            model = model_type(embedding_size=embed_size, feature=True).eval()\n            batch = torch.Tensor(1, 3, *model.get_input_res()).uniform_()\n            output = model(batch)\n            save_model_cpu(model, None, snap_name, 0, write_solverstate=False)\n\n            model_loaded = model_type(embedding_size=embed_size, feature=True)\n            load_model_state(model_loaded, snap_name, -1, eval_state=True)\n\n            output_loaded = model_loaded(batch)\n\n            self.assertEqual(torch.norm(output - output_loaded), 0)\n\n\nclass LandnetTests(unittest.TestCase):\n    \"\"\"Tests for landmark regressor\"\"\"\n    def test_output_shape(self):\n        \"\"\"Checks output shape\"\"\"\n        model = models_landmarks['landnet']().eval()\n        batch = torch.Tensor(1, 3, *model.get_input_res())\n        output = model(batch)\n        self.assertEqual(list(output.shape), list((1, 10, 1, 1)))\n\n    def test_save_load_snap(self):\n        \"\"\"Checks an ability to save and load model correctly\"\"\"\n        snap_name = os.path.join(os.getcwd(), 'test_snap.pt')\n        model = models_landmarks['landnet']().eval()\n        batch = torch.Tensor(1, 3, *model.get_input_res()).uniform_()\n        output = model(batch)\n        save_model_cpu(model, None, snap_name, 0, write_solverstate=False)\n\n        model_loaded = models_landmarks['landnet']()\n        load_model_state(model_loaded, snap_name, -1, eval_state=True)\n\n        output_loaded = model_loaded(batch)\n\n        self.assertEqual(torch.norm(output - output_loaded), 0)\n\nif __name__ == '__main__':\n    unittest.main()\n"
  },
  {
    "path": "tests/test_utils.py",
    "content": "\"\"\"\n Copyright (c) 2018 Intel Corporation\n Licensed under the Apache License, Version 2.0 (the \"License\");\n you may not use this file except in compliance with the License.\n You may obtain a copy of the License at\n      http://www.apache.org/licenses/LICENSE-2.0\n Unless required by applicable law or agreed to in writing, software\n distributed under the License is distributed on an \"AS IS\" BASIS,\n WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n See the License for the specific language governing permissions and\n limitations under the License.\n\"\"\"\n\nimport unittest\n\nimport torch\nfrom utils.utils import get_model_parameters_number\n\n\nclass UtilsTests(unittest.TestCase):\n    \"\"\"Tests for utils\"\"\"\n    def test_parameters_counter(self):\n        \"\"\"Checks output of get_model_parameters_number\"\"\"\n        class ParamsHolder(torch.nn.Module):\n            \"\"\"Dummy parameters holder\"\"\"\n            def __init__(self, n_params):\n                super(ParamsHolder, self).__init__()\n                self.p1 = torch.nn.Parameter(torch.Tensor(n_params // 2))\n                self.p2 = torch.nn.Parameter(torch.Tensor(n_params // 2))\n                self.dummy = -1\n\n        params_num = 1000\n        module = ParamsHolder(params_num)\n        estimated_params = get_model_parameters_number(module, as_string=False)\n        self.assertEqual(estimated_params, params_num)\n\nif __name__ == '__main__':\n    unittest.main()\n"
  },
  {
    "path": "train.py",
    "content": "\"\"\"\n Copyright (c) 2018 Intel Corporation\n Licensed under the Apache License, Version 2.0 (the \"License\");\n you may not use this file except in compliance with the License.\n You may obtain a copy of the License at\n      http://www.apache.org/licenses/LICENSE-2.0\n Unless required by applicable law or agreed to in writing, software\n distributed under the License is distributed on an \"AS IS\" BASIS,\n WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n See the License for the specific language governing permissions and\n limitations under the License.\n\"\"\"\n\nimport argparse\nimport datetime\nimport os.path as osp\nimport os\nfrom pprint import pformat\n\nimport glog as log\nimport torch\nimport torch.optim as optim\nfrom torch.utils.data import DataLoader\nimport torch.backends.cudnn as cudnn\nfrom torchvision import transforms as t\nfrom tensorboardX import SummaryWriter\n\nfrom datasets import LFW, VGGFace2, MSCeleb1M, IMDBFace, TrillionPairs\n\nfrom losses.am_softmax import AMSoftmaxLoss\nfrom losses.metric_losses import MetricLosses\nfrom evaluate_lfw import evaluate, compute_embeddings_lfw\n\nfrom utils.utils import load_model_state, save_model_cpu\nimport utils.augmentation as augm\nfrom utils.parser_yaml import ArgumentParserWithYaml\nfrom model.common import models_backbones\n\ndef train(args):\n    \"\"\"Performs training of a face recognition network\"\"\"\n    input_size = models_backbones[args.model].get_input_res()\n    if args.train_dataset == 'vgg':\n        assert args.t_list\n        dataset = VGGFace2(args.train, args.t_list, args.t_land)\n    elif args.train_dataset == 'imdbface':\n        dataset = IMDBFace(args.train, args.t_list)\n    elif args.train_dataset == 'trp':\n        dataset = TrillionPairs(args.train, args.t_list)\n    else:\n        dataset = MSCeleb1M(args.train, args.t_list)\n\n    if dataset.have_landmarks:\n        log.info('Use alignment for the train data')\n        dataset.transform = t.Compose([augm.HorizontalFlipNumpy(p=.5),\n                                       augm.CutOutWithPrior(p=0.05, max_area=0.1),\n                                       augm.RandomRotationNumpy(10, p=.95),\n                                       augm.ResizeNumpy(input_size),\n                                       augm.BlurNumpy(k=5, p=.2),\n                                       augm.NumpyToTensor(switch_rb=True)])\n    else:\n        dataset.transform = t.Compose([augm.ResizeNumpy(input_size),\n                                       augm.HorizontalFlipNumpy(),\n                                       augm.RandomRotationNumpy(10),\n                                       augm.NumpyToTensor(switch_rb=True)])\n\n    if args.weighted:\n        train_weights = dataset.get_weights()\n        train_weights = torch.DoubleTensor(train_weights)\n        sampler = torch.utils.data.sampler.WeightedRandomSampler(train_weights, len(train_weights))\n        train_loader = torch.utils.data.DataLoader(dataset, batch_size=args.train_batch_size,\n                                                   sampler=sampler, num_workers=3, pin_memory=False)\n    else:\n        train_loader = DataLoader(dataset, batch_size=args.train_batch_size, num_workers=4, shuffle=True)\n\n    lfw = LFW(args.val, args.v_list, args.v_land)\n    if lfw.use_landmarks:\n        log.info('Use alignment for the test data')\n        lfw.transform = t.Compose([augm.ResizeNumpy(input_size),\n                                   augm.NumpyToTensor(switch_rb=True)])\n    else:\n        lfw.transform = t.Compose([augm.ResizeNumpy((160, 160)),\n                                   augm.CenterCropNumpy(input_size),\n                                   augm.NumpyToTensor(switch_rb=True)])\n\n    log_path = './logs/{:%Y_%m_%d_%H_%M}_{}'.format(datetime.datetime.now(), args.snap_prefix)\n    writer = SummaryWriter(log_path)\n\n    if not osp.exists(args.snap_folder):\n        os.mkdir(args.snap_folder)\n\n    model = models_backbones[args.model](embedding_size=args.embed_size,\n                                         num_classes=dataset.get_num_classes(), feature=False)\n    if args.snap_to_resume is not None:\n        log.info('Resuming snapshot ' + args.snap_to_resume + ' ...')\n        model = load_model_state(model, args.snap_to_resume, args.devices[0], eval_state=False)\n        model = torch.nn.DataParallel(model, device_ids=args.devices)\n    else:\n        model = torch.nn.DataParallel(model, device_ids=args.devices, output_device=args.devices[0])\n        model.cuda()\n        model.train()\n        cudnn.benchmark = True\n\n    log.info('Face Recognition model:')\n    log.info(model)\n\n    if args.mining_type == 'focal':\n        softmax_criterion = AMSoftmaxLoss(gamma=args.gamma, m=args.m, margin_type=args.margin_type, s=args.s)\n    else:\n        softmax_criterion = AMSoftmaxLoss(t=args.t, m=0.35, margin_type=args.margin_type, s=args.s)\n    aux_losses = MetricLosses(dataset.get_num_classes(), args.embed_size, writer)\n    optimizer = optim.SGD(model.parameters(), lr=args.lr, momentum=args.momentum, weight_decay=args.weight_decay)\n    scheduler = optim.lr_scheduler.MultiStepLR(optimizer, [3, 6, 9, 13])\n    for epoch_num in range(args.epoch_total_num):\n        scheduler.step()\n        if epoch_num > 6:\n            model.module.set_dropout_ratio(0.)\n        classification_correct = 0\n        classification_total = 0\n\n        for i, data in enumerate(train_loader, 0):\n            iteration = epoch_num * len(train_loader) + i\n\n            if iteration % args.val_step == 0:\n                snapshot_name = osp.join(args.snap_folder, args.snap_prefix + '_{0}.pt'.format(iteration))\n                if iteration > 0:\n                    log.info('Saving Snapshot: ' + snapshot_name)\n                    save_model_cpu(model, optimizer, snapshot_name, epoch_num)\n\n                log.info('Evaluating Snapshot: ' + snapshot_name)\n                model.eval()\n                same_acc, diff_acc, all_acc, auc = evaluate(args, lfw, model, compute_embeddings_lfw,\n                                                            args.val_batch_size, verbose=False)\n\n                model.train()\n\n                log.info('Validation accuracy: {0:.4f}, {1:.4f}'.format(same_acc, diff_acc))\n                log.info('Validation accuracy mean: {0:.4f}'.format(all_acc))\n                log.info('Validation AUC: {0:.4f}'.format(auc))\n                writer.add_scalar('Accuracy/Val_same_accuracy', same_acc, iteration)\n                writer.add_scalar('Accuracy/Val_diff_accuracy', diff_acc, iteration)\n                writer.add_scalar('Accuracy/Val_accuracy', all_acc, iteration)\n                writer.add_scalar('Accuracy/AUC', auc, iteration)\n\n            data, label = data['img'], data['label'].cuda()\n            features, sm_outputs = model(data)\n\n            optimizer.zero_grad()\n            aux_losses.init_iteration()\n            aux_loss, aux_log = aux_losses(features, label, epoch_num, iteration)\n            loss_sm = softmax_criterion(sm_outputs, label)\n            loss = loss_sm + aux_loss\n            loss.backward()\n            aux_losses.end_iteration()\n            optimizer.step()\n\n            _, predicted = torch.max(sm_outputs.data, 1)\n            classification_total += int(label.size(0))\n            classification_correct += int(torch.sum(predicted.eq(label)))\n            train_acc = float(classification_correct) / classification_total\n\n            if i % 10 == 0:\n                log.info('Iteration %d, Softmax loss: %.4f, Total loss: %.4f' % (iteration, loss_sm, loss) + aux_log)\n                log.info('Learning rate: %f' % scheduler.get_lr()[0])\n                writer.add_scalar('Loss/train_loss', loss, iteration)\n                writer.add_scalar('Loss/softmax_loss', loss_sm, iteration)\n                writer.add_scalar('Learning_rate', scheduler.get_lr()[0], iteration)\n                writer.add_scalar('Accuracy/classification', train_acc, iteration)\n\n\ndef main():\n    \"\"\"Creates a command line parser and starts training\"\"\"\n    parser = ArgumentParserWithYaml(description='Training Face Recognition in PyTorch',\n                                    fromfile_prefix_chars='@',\n                                    epilog=\"Please, note that you can parse parameters from a yaml file if \\\n                                    you add @<path_to_yaml_file> to command line\")\n\n    #datasets configuration\n    parser.add_argument('--train_dataset', choices=['vgg', 'ms1m', 'trp', 'imdbface'],\n                        type=str, default='vgg', help='Name of the train dataset.')\n    parser.add_argument('--train_data_root', dest='train', required=True, type=str, help='Path to train data.')\n    parser.add_argument('--train_list', dest='t_list', required=False, type=str, help='Path to train data image list.')\n    parser.add_argument('--train_landmarks', default='', dest='t_land', required=False, type=str,\n                        help='Path to landmarks for the train images.')\n\n    parser.add_argument('--val_data_root', dest='val', required=True, type=str, help='Path to val data.')\n    parser.add_argument('--val_step', type=int, default=1000, help='Evaluate model each val_step during each epoch.')\n    parser.add_argument('--val_list', dest='v_list', required=True, type=str, help='Path to test data image list.')\n    parser.add_argument('--val_landmarks', dest='v_land', default='', required=False, type=str,\n                        help='Path to landmarks for test images.')\n\n    #model configuration\n    parser.add_argument('--model', choices=models_backbones.keys(), type=str, default='mobilenet', help='Model type.')\n    parser.add_argument('--embed_size', type=int, default=256, help='Size of the face embedding.')\n\n    #optimizer configuration\n    parser.add_argument('--train_batch_size', type=int, default=170, help='Train batch size.')\n    parser.add_argument('--epoch_total_num', type=int, default=30, help='Number of epochs to train.')\n    parser.add_argument('--lr', type=float, default=0.4, help='Learning rate.')\n    parser.add_argument('--momentum', type=float, default=0.9, help='Momentum.')\n    parser.add_argument('--weight_decay', type=float, default=0.0001, help='Weight decay.')\n\n    #loss configuration\n    parser.add_argument('--mining_type', choices=['focal', 'sv'],\n                        type=str, default='sv', help='Hard mining method in loss.')\n    parser.add_argument('--t', type=float, default=1.1, help='t in support vector softmax. See https://arxiv.org/abs/1812.11317 for details')\n    parser.add_argument('--gamma', type=float, default=2., help='Gamma in focal loss. See https://arxiv.org/abs/1708.02002 for details')\n    parser.add_argument('--m', type=float, default=0.35, help='Margin size for AMSoftmax.')\n    parser.add_argument('--s', type=float, default=30., help='Scale for AMSoftmax.')\n    parser.add_argument('--margin_type', choices=['cos', 'arc'],\n                        type=str, default='cos', help='Margin type for AMSoftmax loss.')\n\n    #other parameters\n    parser.add_argument('--devices', type=int, nargs='+', default=[0], help='CUDA devices to use.')\n    parser.add_argument('--val_batch_size', type=int, default=20, help='Validation batch size.')\n    parser.add_argument('--snap_folder', type=str, default='./snapshots/', help='Folder to save snapshots.')\n    parser.add_argument('--snap_prefix', type=str, default='FaceReidNet', help='Prefix for snapshots.')\n    parser.add_argument('--snap_to_resume', type=str, default=None, help='Snapshot to resume.')\n    parser.add_argument('--weighted', action='store_true')\n\n    args = parser.parse_args()\n    log.info('Arguments:\\n' + pformat(args.__dict__))\n\n    with torch.cuda.device(args.devices[0]):\n        train(args)\n\nif __name__ == '__main__':\n    main()\n"
  },
  {
    "path": "train_landmarks.py",
    "content": "\"\"\"\n Copyright (c) 2018 Intel Corporation\n Licensed under the Apache License, Version 2.0 (the \"License\");\n you may not use this file except in compliance with the License.\n You may obtain a copy of the License at\n      http://www.apache.org/licenses/LICENSE-2.0\n Unless required by applicable law or agreed to in writing, software\n distributed under the License is distributed on an \"AS IS\" BASIS,\n WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n See the License for the specific language governing permissions and\n limitations under the License.\n\"\"\"\n\nimport argparse\nimport datetime\nimport os.path as osp\n\nimport numpy as np\nimport glog as log\nfrom tensorboardX import SummaryWriter\nimport torch\nimport torch.backends.cudnn as cudnn\nimport torch.optim as optim\nfrom torch.utils.data import DataLoader\nfrom torchvision.transforms import transforms\n\nfrom datasets import VGGFace2, CelebA, NDG\n\nfrom model.common import models_landmarks\nfrom utils import landmarks_augmentation\nfrom utils.utils import save_model_cpu, load_model_state\nfrom losses.alignment import AlignmentLoss\nfrom evaluate_landmarks import evaluate\n\n\ndef train(args):\n    \"\"\"Launches training of landmark regression model\"\"\"\n    if args.dataset == 'vgg':\n        drops_schedule = [1, 6, 9, 13]\n        dataset = VGGFace2(args.train, args.t_list, args.t_land, landmarks_training=True)\n    elif args.dataset == 'celeba':\n        drops_schedule = [10, 20]\n        dataset = CelebA(args.train, args.t_land)\n    else:\n        drops_schedule = [90, 140, 200]\n        dataset = NDG(args.train, args.t_land)\n\n    if dataset.have_landmarks:\n        log.info('Use alignment for the train data')\n        dataset.transform = transforms.Compose([landmarks_augmentation.Rescale((56, 56)),\n                                                landmarks_augmentation.Blur(k=3, p=.2),\n                                                landmarks_augmentation.HorizontalFlip(p=.5),\n                                                landmarks_augmentation.RandomRotate(50),\n                                                landmarks_augmentation.RandomScale(.8, .9, p=.4),\n                                                landmarks_augmentation.RandomCrop(48),\n                                                landmarks_augmentation.ToTensor(switch_rb=True)])\n    else:\n        log.info('Error: training dataset has no landmarks data')\n        exit()\n\n    train_loader = DataLoader(dataset, batch_size=args.train_batch_size, num_workers=4, shuffle=True)\n    writer = SummaryWriter('./logs_landm/{:%Y_%m_%d_%H_%M}_'.format(datetime.datetime.now()) + args.snap_prefix)\n    model = models_landmarks['landnet']\n\n    if args.snap_to_resume is not None:\n        log.info('Resuming snapshot ' + args.snap_to_resume + ' ...')\n        model = load_model_state(model, args.snap_to_resume, args.device, eval_state=False)\n        model = torch.nn.DataParallel(model, device_ids=[args.device])\n    else:\n        model = torch.nn.DataParallel(model, device_ids=[args.device])\n        model.cuda()\n        model.train()\n        cudnn.enabled = True\n        cudnn.benchmark = True\n\n    log.info('Face landmarks model:')\n    log.info(model)\n\n    criterion = AlignmentLoss('wing')\n    optimizer = optim.SGD(model.parameters(), lr=args.lr, momentum=args.momentum, weight_decay=args.weight_decay)\n    scheduler = optim.lr_scheduler.MultiStepLR(optimizer, drops_schedule)\n    for epoch_num in range(args.epoch_total_num):\n        scheduler.step()\n        if epoch_num > 5:\n            model.module.set_dropout_ratio(0.)\n        for i, data in enumerate(train_loader, 0):\n            iteration = epoch_num * len(train_loader) + i\n\n            data, gt_landmarks = data['img'].cuda(), data['landmarks'].cuda()\n            predicted_landmarks = model(data)\n\n            optimizer.zero_grad()\n            loss = criterion(predicted_landmarks, gt_landmarks)\n            loss.backward()\n            optimizer.step()\n\n            if i % 10 == 0:\n                log.info('Iteration %d, Loss: %.4f' % (iteration, loss))\n                log.info('Learning rate: %f' % scheduler.get_lr()[0])\n                writer.add_scalar('Loss/train_loss', loss.item(), iteration)\n                writer.add_scalar('Learning_rate', scheduler.get_lr()[0], iteration)\n\n            if iteration % args.val_step == 0:\n                snapshot_name = osp.join(args.snap_folder, args.snap_prefix + '_{0}.pt'.format(iteration))\n                log.info('Saving Snapshot: ' + snapshot_name)\n                save_model_cpu(model, optimizer, snapshot_name, epoch_num)\n\n                model.eval()\n                log.info('Evaluating Snapshot: ' + snapshot_name)\n                avg_err, per_point_avg_err, failures_rate = evaluate(train_loader, model)\n                weights = per_point_avg_err / np.sum(per_point_avg_err)\n                criterion.set_weights(weights)\n                log.info(str(weights))\n                log.info('Avg train error: {}'.format(avg_err))\n                log.info('Train failure rate: {}'.format(failures_rate))\n                writer.add_scalar('Quality/Avg_error', avg_err, iteration)\n                writer.add_scalar('Quality/Failure_rate', failures_rate, iteration)\n                model.train()\n\ndef main():\n    \"\"\"Creates a command line parser\"\"\"\n    parser = argparse.ArgumentParser(description='Training Landmarks detector in PyTorch')\n    parser.add_argument('--train_data_root', dest='train', required=True, type=str, help='Path to train data.')\n    parser.add_argument('--train_list', dest='t_list', required=False, type=str, help='Path to train data image list.')\n    parser.add_argument('--train_landmarks', default='', dest='t_land', required=False, type=str,\n                        help='Path to landmarks for the train images.')\n    parser.add_argument('--train_batch_size', type=int, default=170, help='Train batch size.')\n    parser.add_argument('--epoch_total_num', type=int, default=30, help='Number of epochs to train.')\n    parser.add_argument('--lr', type=float, default=0.4, help='Learning rate.')\n    parser.add_argument('--momentum', type=float, default=0.9, help='Momentum.')\n    parser.add_argument('--val_step', type=int, default=2000, help='Evaluate model each val_step during each epoch.')\n    parser.add_argument('--weight_decay', type=float, default=0.0001, help='Weight decay.')\n    parser.add_argument('--device', '-d', default=0, type=int)\n    parser.add_argument('--snap_folder', type=str, default='./snapshots/', help='Folder to save snapshots.')\n    parser.add_argument('--snap_prefix', type=str, default='LandmarksNet', help='Prefix for snapshots.')\n    parser.add_argument('--snap_to_resume', type=str, default=None, help='Snapshot to resume.')\n    parser.add_argument('--dataset', choices=['vgg', 'celeb', 'ngd'], type=str, default='vgg', help='Dataset.')\n    arguments = parser.parse_args()\n\n    with torch.cuda.device(arguments.device):\n        train(arguments)\n\nif __name__ == '__main__':\n    main()\n"
  },
  {
    "path": "utils/__init__.py",
    "content": ""
  },
  {
    "path": "utils/augmentation.py",
    "content": "\"\"\"\n Copyright (c) 2018 Intel Corporation\n Licensed under the Apache License, Version 2.0 (the \"License\");\n you may not use this file except in compliance with the License.\n You may obtain a copy of the License at\n      http://www.apache.org/licenses/LICENSE-2.0\n Unless required by applicable law or agreed to in writing, software\n distributed under the License is distributed on an \"AS IS\" BASIS,\n WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n See the License for the specific language governing permissions and\n limitations under the License.\n\"\"\"\n\nimport math\nimport torch\nimport numpy as np\nimport cv2 as cv\n\ntry:\n    from .face_align import FivePointsAligner\nexcept (ImportError, SystemError) as exp:\n    from face_align import FivePointsAligner\n\n\nclass HorizontalFlipNumpy:\n    \"\"\"Horizontal flip augmentation with probability p\"\"\"\n    def __init__(self, p=.5):\n        assert 0 <= p <= 1.\n        self.p = p\n\n    def __call__(self, img):\n        if float(torch.FloatTensor(1).uniform_()) < self.p:\n            return cv.flip(img, 1)\n        return img\n\n\nclass ShowTransform:\n    \"\"\"Show image using opencv\"\"\"\n    def __call__(self, sample):\n        img = np.array(sample)\n        cv.imshow('image', img)\n        cv.waitKey()\n        return sample\n\n\nclass NumpyToTensor:\n    \"\"\"Converts a numpy array to torch.Tensor with optionally swapping R and B channels\"\"\"\n    def __init__(self, switch_rb=False):\n        self.switch_rb = switch_rb\n\n    def __call__(self, image):\n        # swap color axis because\n        # numpy image: H x W x C\n        # torch image: C X H X W\n        if self.switch_rb:\n            image = cv.cvtColor(image, cv.COLOR_RGB2BGR)\n        image = image.transpose((2, 0, 1))\n        return torch.from_numpy(image).type(torch.FloatTensor) / 255.\n\n\nclass RandomShiftNumpy:\n    \"\"\"Shifts an image by a randomly generated offset along x and y axes\"\"\"\n    def __init__(self, max_rel_shift, p=.5):\n        self.p = p\n        self.max_rel_shift = max_rel_shift\n\n    def __call__(self, image):\n        if float(torch.FloatTensor(1).uniform_()) < self.p:\n            rel_shift = 2 * (torch.FloatTensor(1).uniform_() - .5) * self.max_rel_shift\n            h, w = image.shape[:2]\n            shift_w = w * rel_shift\n            shift_h = h * rel_shift\n            transl_mat = np.array([[1., 0., shift_w], [0., 1., shift_h]])\n            image = cv.warpAffine(image, transl_mat, (w, h))\n\n        return image\n\n\nclass RandomRotationNumpy:\n    \"\"\"Rotates an image around it's center by a randomly generated angle\"\"\"\n    def __init__(self, max_angle, p=.5):\n        self.max_angle = max_angle\n        self.p = p\n\n    def __call__(self, image):\n        if float(torch.FloatTensor(1).uniform_()) < self.p:\n            angle = 2 * (torch.FloatTensor(1).uniform_() - .5) * self.max_angle\n            h, w = image.shape[:2]\n            rot_mat = cv.getRotationMatrix2D((w * 0.5, h * 0.5), angle, 1.)\n            image = cv.warpAffine(image, rot_mat, (w, h), flags=cv.INTER_LANCZOS4)\n\n        return image\n\n\nclass ResizeNumpy:\n    \"\"\"Resizes an image in numpy format\"\"\"\n    def __init__(self, output_size):\n        assert isinstance(output_size, (int, tuple))\n        self.output_size = output_size\n\n    def __call__(self, image):\n        h, w = image.shape[:2]\n        if isinstance(self.output_size, int):\n            if h > w:\n                new_h, new_w = self.output_size * h / w, self.output_size\n            else:\n                new_h, new_w = self.output_size, self.output_size * w / h\n        else:\n            new_h, new_w = self.output_size\n\n        new_h, new_w = int(new_h), int(new_w)\n        img = cv.resize(image, (new_h, new_w))\n        return img\n\n\nclass CenterCropNumpy:\n    \"\"\"Performs a center crop of an images\"\"\"\n    def __init__(self, output_size):\n        assert isinstance(output_size, (int, tuple))\n        self.output_size = output_size\n\n    def __call__(self, image):\n        h, w = image.shape[:2]\n        if isinstance(self.output_size, int):\n            new_h, new_w = self.output_size, self.output_size\n        else:\n            new_h, new_w = self.output_size\n\n        s_h = int(h / 2 - new_h / 2)\n        s_w = int(w / 2 - new_w / 2)\n        image = image[s_h: s_h + new_h, s_w: s_w + new_w]\n        return image\n\n\nclass BlurNumpy:\n    \"\"\"Blurs an image with the given sigma and probability\"\"\"\n    def __init__(self, p, k):\n        self.p = p\n        assert k % 2 == 1\n        self.k = k\n\n    def __call__(self, img):\n        if float(torch.FloatTensor(1).uniform_()) < self.p:\n            img = cv.blur(img, (self.k, self.k))\n        return img\n\n\nclass CutOutWithPrior:\n    \"\"\"Cuts rectangular patches from an image around pre-defined landmark locations\"\"\"\n    def __init__(self, p, max_area):\n        self.p = p\n        self.max_area = max_area\n\n    # use after resize transform\n    def __call__(self, img):\n        height, width = img.shape[:2]\n        keypoints_ref = np.zeros((5, 2), dtype=np.float32)\n        keypoints_ref[:, 0] = FivePointsAligner.ref_landmarks[:, 0] * width\n        keypoints_ref[:, 1] = FivePointsAligner.ref_landmarks[:, 1] * height\n\n        if float(torch.FloatTensor(1).uniform_()) < self.p:\n            erase_num = torch.LongTensor(1).random_(1, 4)\n            erase_ratio = torch.FloatTensor(1).uniform_(self.max_area / 2, self.max_area)\n            erase_h = math.sqrt(erase_ratio) / float(erase_num) * height\n            erase_w = math.sqrt(erase_ratio) / float(erase_num) * width\n\n            erased_idx = []\n            for _ in range(erase_num):\n                erase_pos = int(torch.LongTensor(1).random_(0, 5))\n                while erase_pos in erased_idx:\n                    erase_pos = int(torch.LongTensor(1).random_(0, 5))\n\n                left_corner = (\n                    int(keypoints_ref[erase_pos][0] - erase_h / 2), int(keypoints_ref[erase_pos][1] - erase_w / 2))\n                right_corner = (\n                    int(keypoints_ref[erase_pos][0] + erase_h / 2), int(keypoints_ref[erase_pos][1] + erase_w / 2))\n\n                cv.rectangle(img, tuple(left_corner), tuple(right_corner), (0, 0, 0), thickness=-1)\n                erased_idx.append(erase_pos)\n\n        return img\n"
  },
  {
    "path": "utils/face_align.py",
    "content": "\"\"\"\n Copyright (c) 2018 Intel Corporation\n Licensed under the Apache License, Version 2.0 (the \"License\");\n you may not use this file except in compliance with the License.\n You may obtain a copy of the License at\n      http://www.apache.org/licenses/LICENSE-2.0\n Unless required by applicable law or agreed to in writing, software\n distributed under the License is distributed on an \"AS IS\" BASIS,\n WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n See the License for the specific language governing permissions and\n limitations under the License.\n\"\"\"\n\nimport cv2 as cv\nimport numpy as np\n\n\nclass FivePointsAligner():\n    \"\"\"This class performs face alignmet by five reference points\"\"\"\n    ref_landmarks = np.array([30.2946 / 96, 51.6963 / 112,\n                              65.5318 / 96, 51.5014 / 112,\n                              48.0252 / 96, 71.7366 / 112,\n                              33.5493 / 96, 92.3655 / 112,\n                              62.7299 / 96, 92.2041 / 112], dtype=np.float64).reshape(5, 2)\n    @staticmethod\n    def align(img, landmarks, d_size=(400, 400), normalized=False, show=False):\n        \"\"\"Transforms given image in such a way that landmarks are located near ref_landmarks after transformation\"\"\"\n        assert len(landmarks) == 10\n        assert isinstance(img, np.ndarray)\n        landmarks = np.array(landmarks).reshape(5, 2)\n        dw, dh = d_size\n\n        keypoints = landmarks.copy().astype(np.float64)\n        if normalized:\n            keypoints[:, 0] *= img.shape[1]\n            keypoints[:, 1] *= img.shape[0]\n\n        keypoints_ref = np.zeros((5, 2), dtype=np.float64)\n        keypoints_ref[:, 0] = FivePointsAligner.ref_landmarks[:, 0] * dw\n        keypoints_ref[:, 1] = FivePointsAligner.ref_landmarks[:, 1] * dh\n\n        transform_matrix = transformation_from_points(keypoints_ref, keypoints)\n        output_im = cv.warpAffine(img, transform_matrix, d_size, flags=cv.WARP_INVERSE_MAP)\n\n        if show:\n            tmp_output = output_im.copy()\n            for point in keypoints_ref:\n                cv.circle(tmp_output, (int(point[0]), int(point[1])), 5, (255, 0, 0), -1)\n            for point in keypoints:\n                cv.circle(img, (int(point[0]), int(point[1])), 5, (255, 0, 0), -1)\n            img = cv.resize(img, d_size)\n            cv.imshow('source/warped', np.hstack((img, tmp_output)))\n            cv.waitKey()\n\n        return output_im\n\n\ndef transformation_from_points(points1, points2):\n    \"\"\"Builds an affine transformation matrix form points1 to points2\"\"\"\n    points1 = points1.astype(np.float64)\n    points2 = points2.astype(np.float64)\n\n    c1 = np.mean(points1, axis=0)\n    c2 = np.mean(points2, axis=0)\n    points1 -= c1\n    points2 -= c2\n\n    s1 = np.std(points1)\n    s2 = np.std(points2)\n    points1 /= s1\n    points2 /= s2\n\n    u, _, vt = np.linalg.svd(np.matmul(points1.T, points2))\n    r = np.matmul(u, vt).T\n\n    return np.hstack(((s2 / s1) * r, (c2.T - (s2 / s1) * np.matmul(r, c1.T)).reshape(2, -1)))\n"
  },
  {
    "path": "utils/ie_tools.py",
    "content": "\"\"\"\n Copyright (c) 2018 Intel Corporation\n Licensed under the Apache License, Version 2.0 (the \"License\");\n you may not use this file except in compliance with the License.\n You may obtain a copy of the License at\n      http://www.apache.org/licenses/LICENSE-2.0\n Unless required by applicable law or agreed to in writing, software\n distributed under the License is distributed on an \"AS IS\" BASIS,\n WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n See the License for the specific language governing permissions and\n limitations under the License.\n\"\"\"\n\nimport sys\nimport os\n\nimport glog as log\nimport numpy as np\nfrom openvino.inference_engine import IENetwork, IEPlugin # pylint: disable=import-error,E0611\n\nclass IEModel:\n    \"\"\"Class for inference of models in the Inference Engine format\"\"\"\n    def __init__(self, exec_net, inputs_info, input_key, output_key):\n        self.net = exec_net\n        self.inputs_info = inputs_info\n        self.input_key = input_key\n        self.output_key = output_key\n\n    def forward(self, img):\n        \"\"\"Performs forward pass of the wrapped IE model\"\"\"\n        res = self.net.infer(inputs={self.input_key: np.expand_dims(img.transpose(2, 0, 1), axis=0)})\n        return np.copy(res[self.output_key])\n\n    def get_input_shape(self):\n        \"\"\"Returns an input shape of the wrapped IE model\"\"\"\n        return self.inputs_info[self.input_key]\n\n\ndef load_ie_model(model_xml, device, plugin_dir, cpu_extension=''):\n    \"\"\"Loads a model in the Inference Engine format\"\"\"\n    model_bin = os.path.splitext(model_xml)[0] + \".bin\"\n    # Plugin initialization for specified device and load extensions library if specified\n    plugin = IEPlugin(device=device, plugin_dirs=plugin_dir)\n    if cpu_extension and 'CPU' in device:\n        plugin.add_cpu_extension(cpu_extension)\n    # Read IR\n    log.info(\"Loading network files:\\n\\t%s\\n\\t%s\", model_xml, model_bin)\n    net = IENetwork(model=model_xml, weights=model_bin)\n\n    if \"CPU\" in plugin.device:\n        supported_layers = plugin.get_supported_layers(net)\n        not_supported_layers = [l for l in net.layers.keys() if l not in supported_layers]\n        if not_supported_layers:\n            log.error(\"Following layers are not supported by the plugin for specified device %s:\\n %s\",\n                      plugin.device, ', '.join(not_supported_layers))\n            log.error(\"Please try to specify cpu extensions library path in sample's command line parameters using -l \"\n                      \"or --cpu_extension command line argument\")\n            sys.exit(1)\n\n    assert len(net.inputs.keys()) == 1, \"Checker supports only single input topologies\"\n    assert len(net.outputs) == 1, \"Checker supports only single output topologies\"\n\n    log.info(\"Preparing input blobs\")\n    input_blob = next(iter(net.inputs))\n    out_blob = next(iter(net.outputs))\n    net.batch_size = 1\n\n    # Loading model to the plugin\n    log.info(\"Loading model to the plugin\")\n    exec_net = plugin.load(network=net)\n    model = IEModel(exec_net, net.inputs, input_blob, out_blob)\n    del net\n    return model\n"
  },
  {
    "path": "utils/landmarks_augmentation.py",
    "content": "\"\"\"\n Copyright (c) 2018 Intel Corporation\n Licensed under the Apache License, Version 2.0 (the \"License\");\n you may not use this file except in compliance with the License.\n You may obtain a copy of the License at\n      http://www.apache.org/licenses/LICENSE-2.0\n Unless required by applicable law or agreed to in writing, software\n distributed under the License is distributed on an \"AS IS\" BASIS,\n WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n See the License for the specific language governing permissions and\n limitations under the License.\n\"\"\"\n\nimport cv2 as cv\nimport numpy as np\nimport torch\n\n\nclass Rescale:\n    \"\"\"Resizes an image and corresponding landmarks\"\"\"\n    def __init__(self, output_size):\n        assert isinstance(output_size, (int, tuple))\n        self.output_size = output_size\n\n    def __call__(self, sample):\n        image, landmarks = sample['img'], sample['landmarks']\n\n        h, w = image.shape[:2]\n        if isinstance(self.output_size, int):\n            if w > h:\n                new_h, new_w = self.output_size, self.output_size * w / h\n            else:\n                new_h, new_w = self.output_size * h / w, self.output_size\n        else:\n            new_h, new_w = self.output_size\n        new_h, new_w = int(new_h), int(new_w)\n        img = cv.resize(image, (new_h, new_w))\n        return {'img': img, 'landmarks': landmarks}\n\n\nclass RandomCrop:\n    \"\"\"Makes a random crop from the source image with corresponding transformation of landmarks\"\"\"\n    def __init__(self, output_size):\n        assert isinstance(output_size, (int, tuple))\n        if isinstance(output_size, int):\n            self.output_size = (output_size, output_size)\n        else:\n            assert len(output_size) == 2\n            self.output_size = output_size\n\n    def __call__(self, sample):\n        image, landmarks = sample['img'], sample['landmarks'].reshape(-1, 2)\n\n        h, w = image.shape[:2]\n        new_h, new_w = self.output_size\n\n        top = np.random.randint(0, h - new_h)\n        left = np.random.randint(0, w - new_w)\n\n        image = image[top: top + new_h,\n                      left: left + new_w]\n\n        landmarks = landmarks - [left / float(w), top / float(h)]\n        for point in landmarks:\n            point[0] *= float(h) / new_h\n            point[1] *= float(w) / new_w\n\n        return {'img': image, 'landmarks': landmarks}\n\n\nclass HorizontalFlip:\n    \"\"\"Flips an input image and landmarks horizontally with a given probability\"\"\"\n    def __init__(self, p=.5):\n        self.p = p\n\n    def __call__(self, sample):\n        image, landmarks = sample['img'], sample['landmarks'].reshape(-1, 2)\n\n        if float(torch.FloatTensor(1).uniform_()) < self.p:\n            image = cv.flip(image, 1)\n            landmarks = landmarks.reshape(5, 2)\n            landmarks[:, 0] = 1. - landmarks[:, 0]\n            tmp = np.copy(landmarks[0])\n            landmarks[0] = landmarks[1]\n            landmarks[1] = tmp\n\n            tmp = np.copy(landmarks[3])\n            landmarks[3] = landmarks[4]\n            landmarks[4] = tmp\n\n        return {'img': image, 'landmarks': landmarks}\n\n\nclass Blur:\n    \"\"\"Blurs an image with the given sigma and probability\"\"\"\n    def __init__(self, p, k):\n        self.p = p\n        assert k % 2 == 1\n        self.k = k\n\n    def __call__(self, sample):\n        image, landmarks = sample['img'], sample['landmarks']\n\n        if float(torch.FloatTensor(1).uniform_()) < self.p:\n            image = cv.blur(image, (self.k, self.k))\n\n        return {'img': image, 'landmarks': landmarks}\n\n\nclass Show:\n    \"\"\"Show image using opencv\"\"\"\n    def __call__(self, sample):\n        image, landmarks = sample['img'].copy(), sample['landmarks'].reshape(-1, 2)\n        h, w = image.shape[:2]\n        for point in landmarks:\n            cv.circle(image, (int(point[0]*w), int(point[1]*h)), 3, (255, 0, 0), -1)\n        cv.imshow('image', image)\n        cv.waitKey()\n        return sample\n\n\nclass RandomRotate:\n    \"\"\"\n        Rotates an image around it's center by a randomly generated angle.\n        Also performs the same transformation with landmark points.\n    \"\"\"\n    def __init__(self, max_angle, p=.5):\n        self.max_angle = max_angle\n        self.p = p\n\n    def __call__(self, sample):\n        image, landmarks = sample['img'], sample['landmarks']\n\n        if float(torch.FloatTensor(1).uniform_()) < self.p:\n            angle = 2*(torch.FloatTensor(1).uniform_() - .5)*self.max_angle\n            h, w = image.shape[:2]\n            rot_mat = cv.getRotationMatrix2D((w*0.5, h*0.5), angle, 1.)\n            image = cv.warpAffine(image, rot_mat, (w, h), flags=cv.INTER_LANCZOS4)\n            rot_mat_l = cv.getRotationMatrix2D((0.5, 0.5), angle, 1.)\n            landmarks = cv.transform(landmarks.reshape(1, 5, 2), rot_mat_l).reshape(5, 2)\n\n        return {'img': image, 'landmarks': landmarks}\n\n\nclass ToTensor:\n    \"\"\"Convert ndarrays in sample to Tensors.\"\"\"\n    def __init__(self, switch_rb=False):\n        self.switch_rb = switch_rb\n\n    def __call__(self, sample):\n        image, landmarks = sample['img'], sample['landmarks']\n        # swap color axis because\n        # numpy image: H x W x C\n        # torch image: C X H X W\n        if self.switch_rb:\n            image = cv.cvtColor(image, cv.COLOR_RGB2BGR)\n        image = image.transpose((2, 0, 1))\n        return {'img': torch.from_numpy(image).type(torch.FloatTensor) / 255,\n                'landmarks': torch.from_numpy(landmarks).type(torch.FloatTensor).view(-1, 1, 1)}\n\n\nclass RandomScale:\n    \"\"\"Performs uniform scale with a random magnitude\"\"\"\n    def __init__(self, max_scale, min_scale, p=.5):\n        self.max_scale = max_scale\n        self.min_scale = min_scale\n        self.p = p\n\n    def __call__(self, sample):\n        image, landmarks = sample['img'], sample['landmarks']\n\n        if float(torch.FloatTensor(1).uniform_()) < self.p:\n            scale = self.min_scale + torch.FloatTensor(1).uniform_()*(self.max_scale - self.min_scale)\n            h, w = image.shape[:2]\n            rot_mat = cv.getRotationMatrix2D((w*0.5, h*0.5), 0, scale)\n            image = cv.warpAffine(image, rot_mat, (w, h), flags=cv.INTER_LANCZOS4)\n            rot_mat_l = cv.getRotationMatrix2D((0.5, 0.5), 0, scale)\n            landmarks = cv.transform(landmarks.reshape(1, 5, 2), rot_mat_l).reshape(5, 2)\n\n        return {'img': image, 'landmarks': landmarks}\n"
  },
  {
    "path": "utils/parser_yaml.py",
    "content": "\"\"\"\n Copyright (c) 2018 Intel Corporation\n Licensed under the Apache License, Version 2.0 (the \"License\");\n you may not use this file except in compliance with the License.\n You may obtain a copy of the License at\n      http://www.apache.org/licenses/LICENSE-2.0\n Unless required by applicable law or agreed to in writing, software\n distributed under the License is distributed on an \"AS IS\" BASIS,\n WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n See the License for the specific language governing permissions and\n limitations under the License.\n\"\"\"\n\nfrom argparse import ArgumentParser\nimport yaml\n\nclass ArgumentParserWithYaml(ArgumentParser):\n    \"\"\"\n    Attention, this will work with simple yaml files only, and if there is no action=store_false\n    \"\"\"\n    @staticmethod\n    def _check_arg_line_repr_None(arg_line, k, v):\n        \"\"\" The method is required, since by default python prints None value as None, whereas yaml waiths for null \"\"\"\n        s = arg_line.strip()\n        prefixes = [k, \"'\" + k + \"'\", '\"' + k + '\"']\n        is_ok = False\n        for prefix in prefixes:\n            if s.startswith(prefix):\n                s = s[len(prefix):]\n                is_ok = True\n                break\n        if not is_ok:\n            raise RuntimeError(\"Unknown prefix in line '{}', k = '{}', v = '{}'\".format(arg_line, k, v))\n        s = s.strip()\n        assert s.startswith(':'), \"Bad format of line '{}', k = '{}', v = '{}'\".format(arg_line, k, v)\n        s = s[1:]\n        s = s.strip()\n        #print(\"arg line '{}' repr None = {}, s = '{}'\".format(arg_line, s == \"None\", s))\n\n        return s == \"None\" #note that 'None' will be a string, whereas just None will be None\n\n    def convert_arg_line_to_args(self, arg_line):\n        arg_line = arg_line.strip()\n        if not arg_line:\n            return []\n        if arg_line.endswith(','):\n            arg_line = arg_line[:-1]\n\n        data = yaml.load(arg_line)\n        if data is None:\n            return []\n        assert type(data) is dict\n        assert len(data) == 1\n\n        res = []\n        for k, v in data.items():\n            if v == 'None': # default value is None -- skipping\n                if self._check_arg_line_repr_None(arg_line, k, v): #additional check that somebody passed string \"None\"\n                    continue\n                else:\n                    print(\"WARNING: DURING PARSING ARGUMENTS FILE: possible error in the argument line '{}' -- probably None value is missed\".format(arg_line))\n\n            if type(v) is list:\n                res.append('--' + str(k))\n                [res.append(str(item)) for item in v]\n                continue\n\n            if type(v) is bool: # special case, action=store_true, do not use store_false!\n                if v:\n                    res.append('--' + str(k))\n                continue\n\n            # attention, there may be small issue with converting float -> string -> float -> string\n            res.extend(['--' + str(k), str(v)])\n\n        return res\n"
  },
  {
    "path": "utils/utils.py",
    "content": "\"\"\"\n Copyright (c) 2018 Intel Corporation\n Licensed under the Apache License, Version 2.0 (the \"License\");\n you may not use this file except in compliance with the License.\n You may obtain a copy of the License at\n      http://www.apache.org/licenses/LICENSE-2.0\n Unless required by applicable law or agreed to in writing, software\n distributed under the License is distributed on an \"AS IS\" BASIS,\n WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n See the License for the specific language governing permissions and\n limitations under the License.\n\"\"\"\n\nfrom collections import OrderedDict\n\nimport torch\nimport torch.backends.cudnn as cudnn\n\n\ndef save_model_cpu(net, optim, ckpt_fname, epoch, write_solverstate=False):\n    \"\"\"Saves model weights and optimizer state (optionally) to a file\"\"\"\n    state_dict = net.state_dict()\n    for key in state_dict.keys():\n        state_dict[key] = state_dict[key].cpu()\n    snapshot_dict = {\n        'epoch': epoch,\n        'state_dict': state_dict}\n\n    if write_solverstate:\n        snapshot_dict['optimizer'] = optim\n\n    torch.save(snapshot_dict, ckpt_fname)\n\n\ndef get_model_parameters_number(model, as_string=True):\n    \"\"\"Returns a total number of trainable parameters in a specified model\"\"\"\n    params_num = sum(p.numel() for p in model.parameters() if p.requires_grad)\n    if not as_string:\n        return params_num\n\n    if params_num // 10 ** 6 > 0:\n        flops_str = str(round(params_num / 10. ** 6, 2)) + 'M'\n    elif params_num // 10 ** 3 > 0:\n        flops_str = str(round(params_num / 10. ** 3, 2)) + 'k'\n    else:\n        flops_str = str(params_num)\n    return flops_str\n\n\ndef load_model_state(model, snap, device_id, eval_state=True):\n    \"\"\"Loads model weight from a file produced by save_model_cpu\"\"\"\n    if device_id != -1:\n        location = 'cuda:' + str(device_id)\n    else:\n        location = 'cpu'\n    state_dict = torch.load(snap, map_location=location)['state_dict']\n\n    new_state_dict = OrderedDict()\n    for k, v in state_dict.items():\n        head = k[:7]\n        if head == 'module.':\n            name = k[7:]  # remove `module.`\n        else:\n            name = k\n        new_state_dict[name] = v\n\n    model.load_state_dict(new_state_dict, strict=False)\n\n    if device_id != -1:\n        model.cuda(device_id)\n        cudnn.benchmark = True\n\n    if eval_state:\n        model.eval()\n    else:\n        model.train()\n\n    return model\n\n\ndef flip_tensor(x, dim):\n    \"\"\"Flips a tensor along the specified axis\"\"\"\n    xsize = x.size()\n    dim = x.dim() + dim if dim < 0 else dim\n    x = x.view(-1, *xsize[dim:])\n    x = x.view(x.size(0), x.size(1), -1)[:, getattr(torch.arange(x.size(1) - 1, -1, -1),\n                                                    ('cpu', 'cuda')[x.is_cuda])().long(), :]\n    return x.view(xsize)\n"
  }
]