[
  {
    "path": ".gitignore",
    "content": "#/*\n**/__pycache__\n.spyproject/\n/libs/annotator/fitted\n/libs/annotator/smplify/models\n/model\n/examples/imgs\n/examples/example_annot.npy\n/examples/example_model.th\n/examples/stats.npy\n/examples/h36m2Dpose/final_state.pth\n*.yml\n*.log\n*.ini\n*.bak\n*.pth\n*.csv\n*.pyc\n\n\n"
  },
  {
    "path": "LICENSE",
    "content": "MIT License\n\nCopyright (c) 2020 Shichao Li\n\nPermission is hereby granted, free of charge, to any person obtaining a copy\nof this software and associated documentation files (the \"Software\"), to deal\nin the Software without restriction, including without limitation the rights\nto use, copy, modify, merge, publish, distribute, sublicense, and/or sell\ncopies of the Software, and to permit persons to whom the Software is\nfurnished to do so, subject to the following conditions:\n\nThe above copyright notice and this permission notice shall be included in all\ncopies or substantial portions of the Software.\n\nTHE SOFTWARE IS PROVIDED \"AS IS\", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR\nIMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,\nFITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE\nAUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER\nLIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,\nOUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE\nSOFTWARE.\n"
  },
  {
    "path": "README.md",
    "content": "[![PWC](https://img.shields.io/endpoint.svg?url=https://paperswithcode.com/badge/cascaded-deep-monocular-3d-human-pose-1/weakly-supervised-3d-human-pose-estimation-on)](https://paperswithcode.com/sota/weakly-supervised-3d-human-pose-estimation-on?p=cascaded-deep-monocular-3d-human-pose-1)\n[![PWC](https://img.shields.io/endpoint.svg?url=https://paperswithcode.com/badge/cascaded-deep-monocular-3d-human-pose-1/monocular-3d-human-pose-estimation-on-human3)](https://paperswithcode.com/sota/monocular-3d-human-pose-estimation-on-human3?p=cascaded-deep-monocular-3d-human-pose-1)\n[![PWC](https://img.shields.io/endpoint.svg?url=https://paperswithcode.com/badge/cascaded-deep-monocular-3d-human-pose-1/3d-human-pose-estimation-on-human36m)](https://paperswithcode.com/sota/3d-human-pose-estimation-on-human36m?p=cascaded-deep-monocular-3d-human-pose-1)\n# EvoSkeleton\nThis is the project website containing relevant files for the CVPR 2020 paper \"Cascaded deep monocular 3D human pose estimation with evolutionary training data\". The usage and instructions are organized into several parts serving distinct purposes. Please visit the corresponding sub-page for details. For Q&A, go to [discussions](https://github.com/Nicholasli1995/EvoSkeleton/discussions). If you believe there is a technical problem, submit to [issues](https://github.com/Nicholasli1995/EvoSkeleton/issues). \n\nNews:\n\n(2021-04-08): Release v-1.0. The support for pre-trained models is strengthened. More details have been added to the supplementary material.\n  \n## Cascaded 2D-to-3D Lifting\n[This sub-page](https://github.com/Nicholasli1995/EvoSkeleton/blob/master/docs/TRAINING.md) details how to train a cascaded model to lift 2D key-points to 3D skeletons on H36M.\n\nIf you do not want to prepare synthetic data and train the model by yourself, you can access an examplar pre-trained model [here](https://drive.google.com/file/d/158oCTK-9Y8Bl9qxidoHcXfqfeeA7qT93/view?usp=sharing) and follow the instructions in the [document](https://github.com/Nicholasli1995/EvoSkeleton/blob/master/docs/TRAINING.md). This model can be used for in-the-wild inference as well as reproducing the results on MPI-INF-3DHP. The evaluation metric for MPI-INF-3DHP can be accessed [here](https://github.com/chenxuluo/OriNet-demo/tree/master/src/test_util).\n<p align=\"center\">\n  <img src=\"https://github.com/Nicholasli1995/EvoSkeleton/blob/master/imgs/architecture.jpg\"/>\n</p>\n\nPerformance on H36M ([Link to pre-trained models](https://github.com/Nicholasli1995/EvoSkeleton/blob/master/docs/Zoo.md))\n| Protocol \\#1| Avg.|Dir. | Disc| Eat| Greet| Phone| Photo | Pose | Purch.| Sit| SitD.| Smoke| Wait| WalkD.| Walk | WalkT.| \n|-------------------------------------------------------------|------------------|------------------|---------------|------------------|---------------|---------------|------|---------------|------------------|------------------|---------------|---------------|---------------|---------------|---------------|---------------|\n| [Martinez](https://github.com/una-dinosauria/3d-pose-baseline) et al. (ICCV'17)   |62.9| 51.8 | 56.2| 58.1| 59.0   | 69.5 | 78.4| 55.2 | 58.1  | 74.0 | 94.6| 62.3 | 59.1  | 65.1 | 49.5 | 52.4  |\n| Ours (S15678)                                        |**49.7**|**45.6**|**44.6**|**49.3**|**49.3**|**52.5**|**58.5**|**46.4**|**44.3**|**53.8**|**67.5**|**49.4**|**46.1**|**52.5**|**41.4**|**44.4**|\n\n| Protocol \\#2| Avg.|Dir. | Disc| Eat| Greet| Phone| Photo | Pose | Purch.| Sit| SitD.| Smoke| Wait| WalkD.| Walk | WalkT.| \n|-------------------------------------------------------------|------------------|------------------|---------------|------------------|---------------|---------------|------|---------------|------------------|------------------|---------------|---------------|---------------|---------------|---------------|---------------|\n| [Martinez](https://github.com/una-dinosauria/3d-pose-baseline) et al. (ICCV'17)   |47.7| 39.5 | 43.2 | 46.4 | 47.0 | 51.0| 56.0 | 41.4 | 40.6 | 56.5 | 69.4 | 49.2 | 45.0  | 49.5 | 38.0  | 43.1  |\n| Ours (S15678)                                        |**37.7** |**34.2**|**34.6**|**37.3**|**39.3**|**38.5**|**45.6**|**34.5**|**32.7**|**40.5**|**51.3**|**37.7**|**35.4**|**39.9**|**29.9**|**34.5**|\n\n## Hierarchical Human Representation and Data Synthesis\n[This sub-page](https://github.com/Nicholasli1995/EvoSkeleton/blob/master/docs/HHR.md) gives instructions on how to use the 3D skeleton model and how the evolution algorithm can be used to discover novel data.\n<p align=\"center\">\n  <img src=\"https://github.com/Nicholasli1995/EvoSkeleton/blob/master/imgs/hierarchical.jpg\"  width=\"394\" height=\"243\" />\n</p>\n\n## 2D Human Pose Estimation on H3.6M\n\n[This page](https://github.com/Nicholasli1995/EvoSkeleton/blob/master/docs/2DHPE.md) shows how to perform 2D human pose estimation on Human 3.6M dataset with the pre-trained high-resolution heatmap regression model. The highly accurate 2D joint predictions may benefit your 3D human pose estimation project.\n<p align=\"center\">\n  <img src=\"https://github.com/Nicholasli1995/EvoSkeleton/blob/master/imgs/h36m2dpose2.png\" width=\"789\" height=\"208\"/>\n</p>\n\n| Method                    | Parameters| FLOPs|Average Joint Localization Error (pixels) |\n| ------------------------- | ---------------| --------------| --------------| \n| CPN (CVPR' 18)            | -|-| 5.4           |\n| Ours (HRN + U + S)           |63.6M| 32.9G           | **4.4**        |\n\n\n\n## Dataset: Unconstrained 3D Pose in the Wild\n[This sub-page](https://github.com/Nicholasli1995/EvoSkeleton/blob/master/docs/DATASET.md) describs the newly collected dataset Unconstrained 3D Human Pose in the Wild (U3DPW) and gives instructions on how to download it.\n<p align=\"center\">\n  <img src=\"https://github.com/Nicholasli1995/EvoSkeleton/blob/master/imgs/U3DPW.png\"/>\n</p>\n\n## Interactive Annotation Tool\n[This sub-page](https://github.com/Nicholasli1995/EvoSkeleton/blob/master/docs/ANNOTATOR.md) provides usage of an annotation tool that can be used to label 2D and 3D skeleton for an input image. U3DPW was obtained created with this tool and this tool may help increasing the scale of 3D annotation for in-the-wild images.\n<p align=\"center\">\n  <img src=\"https://github.com/Nicholasli1995/EvoSkeleton/blob/master/imgs/tool.gif\" width=\"531\" height=\"291\"/>\n</p>\n\n## Environment\n- Python 3.6\n- Numpy 1.16\n- PyTorch 1.0.1\n- CUDA 9\n\nFor a complete list of other python packages, please refer to [spec-list.txt](https://github.com/Nicholasli1995/EvoSkeleton/blob/master/spec-list.txt). The recommended environment manager is Anaconda, which can create an environment using the provided spec-list. Certain tool in this project may need other specified environment, which is detailed in its corresponding page.\n\n## License\nA MIT license is used for this repository. However, certain third-party dataset (Human 3.6M) and tool (SMPLify) are subject to their respective licenses and may not grant commercial use.\n\n## Citation\nPlease star this repository and cite the following paper in your publications if it helps your research:\n\n    @InProceedings{Li_2020_CVPR,\n    author = {Li, Shichao and Ke, Lei and Pratama, Kevin and Tai, Yu-Wing and Tang, Chi-Keung and Cheng, Kwang-Ting},\n    title = {Cascaded Deep Monocular 3D Human Pose Estimation With Evolutionary Training Data},\n    booktitle = {The IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR)},\n    month = {June},\n    year = {2020}\n    }\n    \nLink to the paper:\n[Cascaded Deep Monocular 3D Human Pose Estimation With Evolutionary Training Data](https://arxiv.org/abs/2006.07778)\n\nLink to the oral presentation video:\n[Youtube](https://www.youtube.com/watch?v=erYymlWw2bo)\n"
  },
  {
    "path": "data/.gitignore",
    "content": "# Ignore everything in this directory\n*\n# Except this file\n!.gitignore\n"
  },
  {
    "path": "docs/2DHPE.md",
    "content": "The pre-trained model can be downloaded from [here](https://drive.google.com/file/d/1NjQFCz0GdS7oIdYrK5ouxEI07wYYh4r8/view?usp=sharing) and placed under \"${EvoSkeleton}/examples/h36m2Dpose\" folder. Your directory should look like this:\n\n   ```\n   ${EvoSkeleton}\n   ├── examples\n      ├── h36m2Dpose\n          ├── cropped (prepared testing images from Human 3.6M)\n          ├── cfgs.yaml (configuration file)\n          ├── final_state.pth (pre-trained high-resolution heatmap regression model)\n   ```\n   \nThen run h36m2Dpose.py at ${EvoSkeleton}/examples\n```bash\npython h36m2Dpose.py\n```\n\nYou should expect to see results like [this](https://github.com/Nicholasli1995/EvoSkeleton/blob/master/imgs/h36m2dpose2.png). I only uploaded a few example images since I cannot upload the whole video due to the license requirement. For your own images, you should crop the humans and prepare your data accordingly.\n<p align=\"center\">\n  <img src=\"https://github.com/Nicholasli1995/EvoSkeleton/blob/master/imgs/h36m2dpose.png\" width=\"924\" height=\"506\"/>\n</p>\n"
  },
  {
    "path": "docs/ANNOTATOR.md",
    "content": "The annotator is composed of three parts:\n1. 2D annotation: interactively annotate 2D key-points for RGB images\n2. 3D parameter fitting: obtain coarse 3D skeleton fitting results based on SMPLify.\n3. 3D annotation: interactively modify 3D parameters.\n\n## 2D Keypoints Annotation\nUsers can annotate 2D Keypoints of images by running the script `annotate_2d.py` under ${EvoSkeleton}/tools. \n```bash\npython annotate_2d.py -d DATASET_PATH\n```\nDATASET_PATH is the path to the folder containing images.\n<p align=\"center\">\n  <img src=\"https://github.com/Nicholasli1995/EvoSkeleton/blob/master/imgs/annotator_2d.gif\"/>\n</p>\n\nUsers can annotate 2D Keypoints in the following order by clicking on the image:\nRight Ankle, Right Knee, Right Hip, Left Hip, Left Knee, Left Ankle, Right Wrist, Right Elbow, Right Shoulder, \nLeft Shoulder, Left Elbow, Left Wrist, Neck, Head top, Spine, Thorax, Nose\n\nOther keyborad short-cuts are:\n\nPress Q to exit the tool.\n\nPress N to go to the next image.\n\nPress Z to save the annotation.\n\nPress C to erase all of the assigned keypoints from the image and start over.\n\n## Coarse 3D Keypoints Estimation\nManually annotating 3D skeleton from scratch is time-consuming, thus we use a tool to obtain an initial 3D pose estimation given 2D annotation. Any method that outputs 3D pose inference given 2D key-points can be employed.\n\nHere we use SMPLify to estimate coarse 3D skeleton. You need to set up a Python 2.7 environment where a [spec-list](https://github.com/Nicholasli1995/EvoSkeleton/blob/master/libs/annotator/smpl-spec-list.txt) can be used as your reference if you use Anaconda. Then you need to install chumpy and opendr using pip:\n```bash\npip install chumpy\npip install opendr\n```\nAfter setting up the environment, you need to download the SMPL model files [here](https://drive.google.com/drive/folders/12qJQP-h4E43FkgE74tybQUjeP_pnqAor?usp=sharing) and organize your project files as follows:\n  ```\n   ${EvoSkeleton}\n   ├── libs\n      ├── annotator\n          ├── smplify\n              ├── models\n                  ├── basicModel_neutral_lbs_10_207_0_v1.0.0.pkl \n                  ├── gmm_08.pkl \n      ├── fit_3d.py               \n                  \n   ```\nThen one can run fit_3d under ${EvoSkeleton}/libs/annotator/fit_3d.py to fit the SMPL model\n```bash\npython fit_3d.py -dataset_dir DATASET_PATH -model_dir MODEL_PATH\n```\nDATASET_PATH is the path to the folder containing the annotated 2D key-point file \"annotation.npy\".\nMODEL_PATH is the path to the used SMPL model (for example, basicModel_neutral_lbs_10_207_0_v1.0.0.pkl). There are other available models depending on the gender of the subject.\nThe fitting process can be shown during running and the file annotation.npy will be updated with 3D parameters.\n<p align=\"center\">\n  <img src=\"https://github.com/Nicholasli1995/EvoSkeleton/blob/master/imgs/fitted.png\"/>\n</p>\n\n## 3D Skeleton Interactive Annotation \nAfter you have obtained fitted parameters for your dataset, you can modify the 3D parameters interactively with this tool. If you just want to try this tool without finishing the above steps, you can play with the pre-fitted parameters [here](https://drive.google.com/file/d/1OJokg844KDpRG3YQsZNpXwFXlVF8iOH5/view?usp=sharing) for U3DPW images. To start the tool, go to ${EvoSkeleton}/tools and run\n```bash\npython annotate_3D.py -dataset_dir DATASET_PATH\n```\nDATASET_PATH is the path to the folder containing the fitted parameters \"fitted.npy\". The tool will select one unannotated image, display it and start the interactive process. You can modify the global orientation as well as the limb orientation of the 3D skeleton. A 2D image with projected key-points will be plotted on-line so that you can check if your annotation is reasonable or not. \n\nSome important keyborad short-cuts are:\n\nPress number keys (2-9) to select which bone vector to rotate.\n\nPress 0 so that the 3D skeleton will rotate as a whole.\n\nPress arrow keys (up and down) to rotate the bone vector.\n\nPress \"m\" to save an updated annotation file.\n\nOther keyboard inputs are detailed in annotate_3D.py\n\n<p align=\"center\">\n  <img src=\"https://github.com/Nicholasli1995/EvoSkeleton/blob/master/imgs/tool.gif\"/>\n</p>\n"
  },
  {
    "path": "docs/DATASET.md",
    "content": "## Download\nYou can access the dataset [here](https://drive.google.com/file/d/1JRJuL69J0drZOAUT8VDK5ywmxt-Gm7s-/view?usp=sharing).\n\n## Folder Structure\n- imgs - Contains the collected images\n- annotation.npy - Contains the pose annotation\n\n   ```\n   ${U3DPW}\n   ├── imgs\n   ├── annotation.npy\n   ```\n## Annotation\nThe annotation file is a Python dictionary that has the following format:\np2d is a numpy array of shape (num_keypoints, 2) that stores the image coordinates of the 2D key-points. Each row in the array stores (x, y) coordinate of the corresponding key-point. These key-points are re-annotated with a style similar to that of Human 3.6M, and can be accessed through key 'h36m'.\nlsp is a boolean flag that indicates whether the image is collected from [Leeds Sport Pose dataset](https://sam.johnson.io/research/lsp.html) or not.\n\n```\n{\n'image_name1':{'p2d':array1, 'lsp':True/False, 'h36m':array2},\n'image_name2':{'p2d':array3, 'lsp':True/False, 'h36m':array4},\n...\n}\n```\n\n## Key-point Semantics\nThe name of the Human 3.6M style key-points are:\n\n| Index | Keypoint |\n|---|-------------|\n| 0 | Hip Center  | \n| 1 | Right Hip   | \n| 2 | Right Knee  |\n| 3 | Right Ankle |\n| 4 | Left Hip    |\n| 5 | Left Knee   |\n| 6 | Left Ankle  | \n| 7 | Spine       |\n| 8 | Thorax      |\n| 9 | Neck        |  \n| 10 | Head Top   |\n| 11 | Left SHoulder |\n| 12 | Left Elbow |\n| 13 | Left Wrist |\n| 14 | Right Shoulder |\n| 15 | Right Elbow|\n| 16 | Right Wrist| \n"
  },
  {
    "path": "docs/HHR.md",
    "content": "## Data Preparation\nPlease prepare data as instructed in the model training [sub-page](https://github.com/Nicholasli1995/EvoSkeleton/blob/master/docs/TRAINING.md).\nThe training data need to downloaded from [here](https://drive.google.com/drive/folders/1zyW8ryGXLq4bumWnVGUROpDNdubWUExg?usp=sharing) and placed under \"${EvoSkeleton}/data\" folder:\n   ```\n   ${EvoSkeleton}\n   ├── data\n      ├── human3.6M\n          ├── your downloaded files\n   ```\n## Model Preparation\nDuring data space exploration, a function that evaluates the validity of 3D skeletons is used. This function is parametrized with a model propsosed by Ijaz Akhter in CVPR 2015. You need to download the \"constraints\" folder from [here](https://drive.google.com/drive/folders/1MUcR9oBNUpTAJ7YUWdyVLKCQW874FszI?usp=sharing) which contains the model parameters and place them under \"${EvoSkeleton}/resources\" folder:\n   ```\n   ${EvoSkeleton}\n   ├── recources\n      ├── constraints\n   ```\n## Dataset Evolution\nTo evolve from a population of 3D skeleton (default to Human 3.6M data), go to \"${EvoSkeleton}/tools\" folder and run\n```bash\npython evolve.py -generate True\n```\n### Controling the Initial Population\nTo reproduce the experiments in different settings, you need to specify the choice of initial population.\nFor weakly-supervised experiments, you should only start with subject 1 (S1) data (a subset of H36M training data) as follows\n```bash\npython evolve.py -generate True -WS True -SS \"S1\"\n```\nYou can even start with extremly scarce data (e.g., 1 percent of S1 data) as follows\n```bash\npython evolve.py -generate True -WS True -SS \"0.01S1\"\n```\n\nAfter finished dataset evolution, you can use the saved file for [training](https://github.com/Nicholasli1995/EvoSkeleton/blob/master/docs/TRAINING.md) to see how dataset evolution might help improve model generalization especially when the initial population is scarce.\n\n## Reference\n\n      @inproceedings{akhter2015pose,\n        title={Pose-conditioned joint angle limits for 3D human pose reconstruction},\n        author={Akhter, Ijaz and Black, Michael J},\n        booktitle={Proceedings of the IEEE conference on computer vision and pattern recognition},\n        pages={1446--1455},\n        year={2015}\n      }\n"
  },
  {
    "path": "docs/TRAINING.md",
    "content": "## Data Preparation\nSimilar to other repositories ([SimpleBaseline](https://github.com/una-dinosauria/3d-pose-baseline), [TemporalConvolution](https://github.com/facebookresearch/VideoPose3D)) on training 2D-to-3D  networks, we provide pre-processed 2D detections, camera parameters and 3D poses for training. The 2D detections are produced by our modified high-resolution model, while the camera parameters and the 3D poses are taken from [SimpleBaseline](https://github.com/una-dinosauria/3d-pose-baseline).\n\nThe training data need to downloaded from [here](https://drive.google.com/drive/folders/1zyW8ryGXLq4bumWnVGUROpDNdubWUExg?usp=sharing) and placed under \"${EvoSkeleton}/data\" folder. Your directory should look like this:\n   ```\n   ${EvoSkeleton}\n   ├── data\n      ├── human3.6M\n          ├── cameras.npy (Camera parameters provided by Human 3.6M)\n          ├── threeDPose_train.npy (3D skeletons from Human 3.6M training split)\n          ├── threeDPose_test.npy (3D skeletons from Human 3.6M test split)\n          ├── twoDPose_HRN_train.npy (2D key-point detections obtained from the heatmap regression model for the training split)\n          ├── twoDPose_HRN_test.npy (2D key-point detections obtained from the heatmap regression model for the test split)\n   ```\n   \n## Weakly-Supervised Experiments on Human 3.6M Dataset\nTo compare with other weakly-supervised methods, only a subset of training data (e.g., subject 1 data) is used to simulate an environment where data is scarce. \nTo perform training, go to ./tools and run\n```bash\npython 2Dto3Dnet.py -train True -num_stages 2 -ws True -ws_name \"S1\"\n```\nThis command performs training on synthetic 2D key-points to remove the influence of the heatmap regression model, whose results correspond to P1* in the performance table. S1 stands for subject 1 data. \"num_stages\" specify the number of deep learners used in the cascade.\nTo train on real detections obtained by the high-resolution heatmap regression model, run\n```bash\npython 2Dto3Dnet.py -train True -num_stages 2 -ws True -ws_name \"S1\" -twoD_source \"HRN\"\n```\nTo train on evolved dataset, you need to specify the path to the evolved data as\n```bash\npython 2Dto3Dnet.py -train True -num_stages 2 -ws True -ws_name \"S1\" -twoD_source \"HRN/synthetic\" -evolved_path \"YourDataPath\"\n```\nSee [this page](https://github.com/Nicholasli1995/EvoSkeleton/blob/master/docs/HHR.md) on how to evolve a dataset.\n\n\nAfter data augmentation using the evolved data, we noticed the model generalization improves significantly despite the initial population size is small. Other methods utilize multi-view or temporal consistency instead of data augmentation to supervise deep models when data is scarce. Compared to them, we achieve state-of-the-art performance by synthesizing new data to supervise the deep model. P1 and P2 refers to the two protocols used for calculating average MPJPE over all 15 actions in H36M. \n\n| Method                    | Avg. MPJPE (P1) |  Avg. MPJPE (P2) |\n| ------------------------- | --------------- |  --------------- |\n| Rhodin et al. (CVPR' 18)  | -               |  64.6            |\n| Kocabas et al. (CVPR' 19) | 65.3            |  57.2            |\n| Pavllo et al. (CVPR' 19)  | 64.7            |  -               |\n| Li et al. (ICCV' 19)      | 88.8            |  66.5            |\n| Ours                      | **60.8**        |  **46.2**        |\n\n## Fully-Supervised Experiments on Human 3.6M Dataset\nTo train on real detections obtained by the high-resolution heatmap regression model, run\n```bash\npython 2Dto3Dnet.py -train True -num_stages 2 -twoD_source \"HRN\"\n```\nTo train on evolved dataset, you need to specify the path to the evolved data as\n```bash\npython 2Dto3Dnet.py -train True -num_stages 3 -num_blocks 3 -twoD_source \"HRN/synthetic\" -evolved_path \"YourDataPath\"\n```\nHere we increase model capacity with \"-num_stages 3 -num_blocks 3\" since the training data size is much larger (if you evolve enough generations).\nWhile the improvement using data evolution is less obvious in fully-supervised setting compared with weakly-supervised setting, our cascaded model still achieved competitive performance compared with other 2D-to-3D lifting models.\n\n| Method                     | Avg. MPJPE (P1) |  Avg. MPJPE (P2) |\n| -------------------------- | --------------- |  --------------- |\n| [Martinez et al.](https://github.com/una-dinosauria/3d-pose-baseline) (ICCV' 17) | 62.9            |  47.7            |\n| Yang et al. (CVPR' 18)     | 58.6            |  **37.7**        |\n| Zhao et al. (CVPR' 19)     | 57.6            |  -               |\n| Sharma et al. (CVPR' 19)   | 58.0            |  40.9            |\n| Moon et al. (ICCV' 19)     | 54.4            |  -               |\n| Ours                       | **50.9**        |  38.0            |\n\n## Inference Example\nIf you only want to use a pre-trained model to conduct inference on in-the-wild images (skipping data synthesis and model training), you can download the sample images and a pre-trained checkpoint [here](https://drive.google.com/file/d/158oCTK-9Y8Bl9qxidoHcXfqfeeA7qT93/view?usp=sharing). Un-zip the downloaded file to \"${EvoSkeleton}/examples\" folder and your directory should look like this:\n   ```\n   ${EvoSkeleton}\n   ├── examples\n      ├── imgs (sample images)\n      ├── example_annot.npy (2D key-points for the samples)\n      ├── example_model.th (pre-trained model)\n      ├── stats.npy (model statistics)\n      ├── inference.py\n   ``` \nThen you can run the following command at \"${EvoSkeleton}/examples\" to perform inference\n```bash\npython inference.py\n```\n\n<p align=\"center\">\n  <img src=\"https://github.com/Nicholasli1995/EvoSkeleton/blob/master/imgs/example.png\"/>\n</p>\n"
  },
  {
    "path": "docs/Zoo.md",
    "content": "This page presents model performance on H36M under various settings. Pre-trained models and instructions for reproduction can also be found.\n\n## Fully-supervised Setting (S15678)\n\n[Download our pre-trained model](https://drive.google.com/drive/folders/1IRKUWrnheD03Dj30LLGlh_LLT1CK6dr5?usp=sharing)\n\n[Download our pre-evolved data](https://drive.google.com/drive/folders/1FKFkmTJQcEdrCvZOSc8cF5OTTjFvyLav?usp=sharing)\n\nInference command:\n```bash\npython 2Dto3Dnet.py -evaluate True -twoD_source \"HRN\" -ckpt_dir \"YourMODELPath\"\n```\nTraining command ([Docs](https://github.com/Nicholasli1995/EvoSkeleton/blob/master/docs/TRAINING.md)):\n```bash\npython 2Dto3Dnet.py -train True -num_stages 3 -num_blocks 3 -twoD_source \"HRN\" -evolved_path \"YourDataPath\"\n```\nData synthesis command ([Docs](https://github.com/Nicholasli1995/EvoSkeleton/blob/master/docs/HHR.md)):\n```bash\npython evolve.py -SS \"S15678\" -T 1.5 -SD \"YourDataPath\" -generate True\n```\nMPJPE (P1) for each action under fully-supervised setting is shown in the table below.\n| Protocol \\#1                                      | Dir.             | Disc             | Eat           | Greet            | Phone         | Photo         | Pose | Purch.        | Sit              | SitD.            | Smoke         | Wait          | WalkD.        | Walk          | WalkT.        | Avg.          |\n|-------------------------------------------------------------|------------------|------------------|---------------|------------------|---------------|---------------|------|---------------|------------------|------------------|---------------|---------------|---------------|---------------|---------------|---------------|\n| [Martinez](https://github.com/una-dinosauria/3d-pose-baseline) et al. (ICCV'17)   | 51.8             | 56.2             | 58.1          | 59.0             | 69.5          | 78.4          | 55.2 | 58.1          | 74.0             | 94.6             | 62.3          | 59.1          | 65.1          | 49.5          | 52.4          | 62.9          |\n| [Fang](https://arxiv.org/abs/1710.06513) et al. (AAAI'18)        | 50.1             | 54.3             | 57.0          | 57.1             | 66.6          | 73.3          | 53.4 | 55.7          | 72.8             | 88.6             | 60.3          | 57.7          | 62.7          | 47.5          | 50.6          | 60.4          |\n| [Yang](https://arxiv.org/abs/1803.09722) et al. (CVPR'18)               | 51.5             | 58.9             | 50.4          | 57.0             | 62.1          | 65.4          | 49.8 | 52.7          | 69.2             | 85.2             | 57.4          | 58.4          | 43.6          | 60.1          | 47.7          | 58.6          |\n| [Pavlakos](https://github.com/geopavlakos/ordinal-pose3d) et al. (CVPR'18)  | 48.5             | 54.4             | 54.4          | 52.0             | 59.4          | 65.3          | 49.9 | 52.9          | 65.8             | 71.1             | 56.6          | 52.9          | 60.9          | 44.7          | 47.8          | 56.2          |\n| [Lee](https://openaccess.thecvf.com/content_ECCV_2018/papers/Kyoungoh_Lee_Propagating_LSTM_3D_ECCV_2018_paper.pdf) et al. (ECCV'18)        | 40.2    | 49.2             | 47.8          | 52.6             | 50.1 | 75.0          | 50.2 | 43.0          | 55.8             | 73.9             | 54.1          | 55.6          | 58.2          | 43.3          | 43.3 | 52.8          |\n| [Zhao](https://arxiv.org/abs/1904.03345) et al. (CVPR'19)           | 47.3             | 60.7             | 51.4          | 60.5             | 61.1          | 49.9 | 47.3 | 68.1          | 86.2             | 55.0             | 67.8          | 61.0          | 42.1 | 60.6          | 45.3          | 57.6          |\n| [Sharma](https://arxiv.org/abs/1904.01324) et al. (ICCV'19)    | 48.6             | 54.5             | 54.2          | 55.7             | 62.6          | 72.0          | 50.5 | 54.3          | 70.0             | 78.3             | 58.1          | 55.4          | 61.4          | 45.2          | 49.7          | 58.0          |\n| [Moon](https://github.com/mks0601/3DMPPE_POSENET_RELEASE) et al. (ICCV'19)    | 51.5             | 56.8             | 51.2          | 52.2             | 55.2          | 47.7 | 50.9 | 63.3          | 69.9             | 54.2    | 57.4          | 50.4          | 42.5          | 57.5          | 47.7          | 54.4          |\n| [Liu](https://www.ecva.net/papers/eccv_2020/papers_ECCV/papers/123550324.pdf) et al. (ECCV'20)      | 46.3             | 52.2             | 47.3 | 50.7             | 55.5          | 67.1          | 49.2 | 46.0          | 60.4             | 71.1             | 51.5          | 50.1          | 54.5          | 40.3 | 43.7          | 52.4          |\n| Ours (S15678)                                        |45.6|44.6|49.3|49.3|52.5|58.5|46.4|44.3|53.8|67.5|49.4|46.1|52.5|41.4|44.4| 49.7 |\n\nMPJPE (P2) for each action under fully-supervised setting is shown in the table below.\n\n| Protocol \\#2                                       | Dir.             | Disc             | Eat           | Greet            | Phone         | Photo         | Pose | Purch.        | Sit              | SitD.            | Smoke         | Wait          | WalkD.        | Walk          | WalkT.        | Avg.          |\n|-------------------------------------------------------------|------------------|------------------|---------------|------------------|---------------|---------------|------|---------------|------------------|------------------|---------------|---------------|---------------|---------------|---------------|---------------|\n| [Martinez](https://github.com/una-dinosauria/3d-pose-baseline) et al. (ICCV'17)   | 39.5             | 43.2             | 46.4          | 47.0             | 51.0          | 56.0          | 41.4 | 40.6          | 56.5             | 69.4             | 49.2          | 45.0          | 49.5          | 38.0          | 43.1          | 47.7          |\n| [Fang](https://arxiv.org/abs/1710.06513) et al. (AAAI'18)        | 38.2             | 41.7             | 43.7          | 44.9             | 48.5          | 55.3          | 40.2 | 38.2          | 54.5             | 64.4             | 47.2          | 44.3          | 47.3          | 36.7          | 41.7          | 45.7          |\n| [Pavlakos](https://github.com/geopavlakos/ordinal-pose3d) et al. (CVPR'18)  | 34.7             | 39.8             | 41.8          | 38.6             | 42.5          | 47.5          | 38.0 | 36.6          | 50.7             | 56.8             | 42.6          | 39.6          | 43.9          | 32.1          | 36.5          | 41.8          |\n| [Yang](https://arxiv.org/abs/1803.09722) et al. (CVPR'18)               | 26.9             | 30.9             | 36.3          | 39.9             | 43.9          | 47.4          | 28.8 | 29.4          | 36.9             | 58.4             | 41.5          | 30.5          | 29.5          | 42.5          | 32.2          | 37.7          |\n| [Sharma](https://arxiv.org/abs/1904.01324) et al. (ICCV'19)    | 35.3             | 35.9             | 45.8          | 42.0             | 40.9          | 52.6          | 36.9 | 35.8          | 43.5             | 51.9             | 44.3          | 38.8          | 45.5          | 29.4          | 34.3          | 40.9          |\n| [Cai](https://openaccess.thecvf.com/content_ICCV_2019/papers/Cai_Exploiting_Spatial-Temporal_Relationships_for_3D_Pose_Estimation_via_Graph_Convolutional_ICCV_2019_paper.pdf) et al. (ICCV'19)     | 35.7             | 37.8             | 36.9          | 40.7             | 39.6          | 45.2          | 37.4 | 34.5          | 46.9             | 50.1 | 40.5          | 36.1          | 41.0          | 29.6          | 33.2          | 39.0          |\n| [Liu](https://www.ecva.net/papers/eccv_2020/papers_ECCV/papers/123550324.pdf) et al. (ECCV'20)    | 35.9             | 40.0             | 38.0          | 41.5             | 42.5          | 51.4          | 37.8 | 36.0          | 48.6             | 56.6             | 41.8          | 38.3          | 42.7          | 31.7          | 36.2          | 41.2          |\n| Ours (S15678)                                        |34.2|34.6|37.3|39.3|38.5|45.6|34.5|32.7|40.5|51.3|37.7|35.4|39.9|29.9|34.5| 37.7 |\n\n## Weakly-supervised Setting (S1)\n\n[Download our pre-trained model](https://drive.google.com/drive/folders/1PZoiizPKeoFTsvnFKIxaRDNbyb0Csx50?usp=sharing)\n\n[Download our pre-evolved data](https://drive.google.com/drive/folders/1nTW2CCCT_sbJ1CejhuiQLTgDDU5sJjZj?usp=sharing)\n\nInference command:\n```bash\npython 2Dto3Dnet.py -evaluate True -twoD_source \"HRN\" -ckpt_dir \"YourMODELPath\" \n```\nTraining command ([Docs](https://github.com/Nicholasli1995/EvoSkeleton/blob/master/docs/TRAINING.md)):\n```bash\npython 2Dto3Dnet.py -train True -num_stages 2 -ws True -ws_name \"S1\" -twoD_source \"HRN\" -evolved_path \"YourDataPath\"\n```\nData synthesis command ([Docs](https://github.com/Nicholasli1995/EvoSkeleton/blob/master/docs/HHR.md)):\n```bash\npython evolve.py -generate True -WS True -SS \"S1\"\n```\nMPJPE (P1) for each action under weakly-supervised setting is shown in the table below.\n| Protocol \\#1                                           | Dir.             | Disc             | Eat           | Greet            | Phone         | Photo            | Pose | Purch.        | Sit              | SitD. | Smoke         | Wait          | WalkD.        | Walk          | WalkT.        | Avg.          |\n|--------------------------------------------------------|------------------|------------------|---------------|------------------|---------------|------------------|------|---------------|------------------|-------|---------------|---------------|---------------|---------------|---------------|---------------|\n| [Kocabas](https://arxiv.org/abs/1903.02330) et al. (CVPR'19)  | -                | -                | -             | -                | -             | -                | -    | -             | -                | -     | -             | -             | -             | -             | -             | 65.3          |\n| [Pavllo](https://arxiv.org/abs/1811.11742) et al. (CVPR'19)  | -                | -                | -             | -                | -             | -                | -    | -             | -                | -     | -             | -             | -             | -             | -             | 64.7          |\n| [Li](https://openaccess.thecvf.com/content_ICCV_2019/papers/Li_On_Boosting_Single-Frame_3D_Human_Pose_Estimation_via_Monocular_Videos_ICCV_2019_paper.pdf) et al. (ICCV'19)          | 70.4             | 83.6             | 76.6          | 78.0             | 85.4          | 106.1   | 72.2 | 103.0         | 115.8            | 165.0 | 82.4          | 74.3          | 94.6          | 60.1          | 70.6          | 88.8          |\n| Ours (S1)                                       | 52.8             | 56.6 | 54.0 | 57.5 | 62.8 | 72.0 | 55.0 | 61.3 | 65.8 | 80.7  | 58.9 | 56.7 | 69.7 | 51.6 | 57.2 | 60.8 |\n\nMPJPE (P2) for each action under fully-supervised setting is shown in the table below.\n\n| Protocol \\#2                                           | Dir.             | Disc             | Eat           | Greet            | Phone         | Photo            | Pose | Purch.        | Sit              | SitD. | Smoke         | Wait          | WalkD.        | Walk          | WalkT.        | Avg.          |\n|--------------------------------------------------------|------------------|------------------|---------------|------------------|---------------|------------------|------|---------------|------------------|-------|---------------|---------------|---------------|---------------|---------------|---------------|\n| [Rhodin](https://arxiv.org/abs/1803.04775) et al. (CVPR'18)  | -                | -                | -             | -                | -             | -                | -    | -             | -                | -     | -             | -             | -             | -             | -             | 64.6          |\n| [Kocabas](https://arxiv.org/abs/1903.02330) et al. (CVPR'19)  | -                | -                | -             | -                | -             | -                | -    | -             | -                | -     | -             | -             | -             | -             | -             | 57.2          |\n| [Li](https://openaccess.thecvf.com/content_ICCV_2019/papers/Li_On_Boosting_Single-Frame_3D_Human_Pose_Estimation_via_Monocular_Videos_ICCV_2019_paper.pdf) et al. (ICCV'19)          | -                | -                | -             | -                | -             | -                | -    | -             | -                | -     | -             | -             | -             | -             | -             | 66.5          |\n| Ours (S1)                                       | 40.2             | 43.4 | 41.9| 46.1 | 48.2 | 55.1 | 42.8 | 42.6 | 49.6 | 61.1  | 44.5 | 43.2 | 51.5 | 38.1 | 44.4 | 46.2 |\n\n"
  },
  {
    "path": "examples/h36m2Dpose/cfgs.yaml",
    "content": "CUDNN:\n  BENCHMARK: true\n  DETERMINISTIC: false\n  ENABLED: true\nGPUS: (0,)\n\nMODEL:\n  NAME: pose_hrnet\n  NUM_JOINTS: 17\n  TARGET_TYPE: 'coordinate'\n  IMAGE_SIZE:\n  - 288\n  - 384\n  HEATMAP_SIZE:\n  - 288\n  - 384\n  EXTRA:\n    PRETRAINED_LAYERS:\n    - 'conv1'\n    - 'bn1'\n    - 'conv2'\n    - 'bn2'\n    - 'layer1'\n    - 'transition1'\n    - 'stage2'\n    - 'transition2'\n    - 'stage3'\n    - 'transition3'\n    - 'stage4'\n    FINAL_CONV_KERNEL: 1\n    STAGE2:\n      NUM_MODULES: 1\n      NUM_BRANCHES: 2\n      BLOCK: BASIC\n      NUM_BLOCKS:\n      - 4\n      - 4\n      NUM_CHANNELS:\n      - 48\n      - 96\n      FUSE_METHOD: SUM\n    STAGE3:\n      NUM_MODULES: 4\n      NUM_BRANCHES: 3\n      BLOCK: BASIC\n      NUM_BLOCKS:\n      - 4\n      - 4\n      - 4\n      NUM_CHANNELS:\n      - 48\n      - 96\n      - 192\n      FUSE_METHOD: SUM\n    STAGE4:\n      NUM_MODULES: 3\n      NUM_BRANCHES: 4\n      BLOCK: BASIC\n      NUM_BLOCKS:\n      - 4\n      - 4\n      - 4\n      - 4\n      NUM_CHANNELS:\n      - 48\n      - 96\n      - 192\n      - 384\n      FUSE_METHOD: SUM\n"
  },
  {
    "path": "examples/h36m2Dpose.py",
    "content": "\"\"\"\nExamplar code showing how to use pre-trained heatmap regression model H() to \nperform 2D pose estimation on Human 3.6M images. \n\"\"\"\n\nimport sys\nsys.path.append(\"../\")\n\nfrom libs.hhr.config import cfg\nfrom libs.hhr.config import update_config\nfrom libs.hhr.utils.utils import get_model_summary\nfrom libs.model.pose_hrnet import get_pose_net\nfrom libs.hhr.utils.transforms import get_affine_transform\nfrom libs.hhr.core.loss import get_max_preds_soft_pt\n\nimport torch\nimport torch.nn.parallel\nimport torch.backends.cudnn as cudnn\nimport torchvision.transforms as transforms\n\nimport argparse\nimport os\nimport logging\nimport cv2\nimport numpy as np\nimport matplotlib.pyplot as plt\n\npose_connection = np.array([[0,1], [1,2], [2,3], [0,4], [4,5], [5,6], [0,7], [7,8],\n                            [8,9], [9,10], [8,11], [11,12], [12,13], [8,14], [14,15],\n                            [15,16]], dtype=np.int\n                           )\nI = pose_connection[:, 0]\nJ = pose_connection[:, 1]\npose_color = np.array([[255, 0, 0], [255, 85, 0], [255, 170, 0], [255, 255, 0], \n                       [170, 255, 0], [85, 255, 0], [0, 255, 0], [0, 255, 85], \n                       [0, 255, 170], [0, 255, 255], [0, 170, 255], [0, 85, 255], \n                       [0, 0, 255], [85, 0, 255], [170, 0, 255], [255, 0, 255]\n                       ])/255.\nre_order = [3, 12, 14, 16, 11, 13, 15, 1, 2, 0, 4, 5, 7, 9, 6, 8, 10]\n\ndef show2Dpose(vals, ax, lcolor=\"#3498db\", rcolor=\"#e74c3c\", add_labels=False):\n    for i in np.arange( len(I) ):\n        x, y = [np.array([vals[I[i], j], vals[J[i], j]]) for j in range(2)]\n        ax.plot(x, y, c=pose_color[i])\n    \ndef parse_args():\n    parser = argparse.ArgumentParser(description='2D pose estimation example')\n    parser.add_argument('--cfg',\n                        help='configuration file',\n                        default='./h36m2Dpose/cfgs.yaml',\n                        type=str)\n    parser.add_argument('--data_path',\n                        help='path to pre-processed testing images',\n                        default='./h36m2Dpose/cropped',\n                        type=str)\n    args = parser.parse_args()\n    return args\n\n\ndef xywh2cs(x, y, w, h, aspect_ratio=0.75):\n    center = np.zeros((2), dtype=np.float32)\n    center[0] = x + w * 0.5\n    center[1] = y + h * 0.5\n\n    if w > aspect_ratio * h:\n        h = w * 1.0 / aspect_ratio\n    elif w < aspect_ratio * h:\n        w = h * aspect_ratio\n    scale = np.array([w * 1.0 / 200, h * 1.0 / 200], dtype=np.float32)\n\n    return center, scale\n    \ndef gather_inputs(args, logger, image_size = (288, 384)):\n    root = args.data_path\n    img_names = os.listdir(root)\n    normalize = transforms.Normalize(\n        mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]\n    )    \n    transform = transforms.Compose([\n                transforms.ToTensor(),\n                normalize,\n                ])    \n    inputs = []\n    # these testing images were cropped from videos of subject 9 and 11\n    for name in img_names:\n        pass\n        image_file = os.path.join(root, name)\n        data_numpy = cv2.imread(image_file, 1 | 128)        \n        data_numpy = cv2.cvtColor(data_numpy, cv2.COLOR_BGR2RGB)\n        if data_numpy is None:\n            logger.error('=> fail to read {}'.format(image_file))\n            raise ValueError('Fail to read {}'.format(image_file))\n\n        c, s = xywh2cs(0, 0, data_numpy.shape[1], data_numpy.shape[0])\n        r = 0    \n        trans = get_affine_transform(c, s, r, image_size)\n        input = cv2.warpAffine(\n            data_numpy,\n            trans,\n            (image_size[0], image_size[1]),\n            flags=cv2.INTER_LINEAR)\n        inputs.append(transform(input).unsqueeze(0))\n    return torch.cat(inputs)\n\ndef unnormalize(tensor):\n    img = tensor.data.cpu().numpy()\n    mean = np.array([0.485, 0.456, 0.406]).reshape(3, 1, 1)\n    std = np.array([0.229, 0.224, 0.225]).reshape(3, 1, 1)\n    img = np.transpose((((img * std) + mean) * 255).astype(np.uint8), (1, 2, 0))\n    return img\n\ndef visualize(inputs, model):\n    output = model(inputs)\n    pred, max_vals = get_max_preds_soft_pt(output)\n    pred = pred.data.cpu().numpy()\n    pred = pred[:, re_order, :]\n    plt.figure()\n    for i in range(len(pred)):\n        ax = plt.subplot(1, len(pred), i+1)\n        ax.imshow(unnormalize(inputs[i]))\n        ax.plot(pred[i][:, 0], pred[i][:, 1], 'ro')\n        show2Dpose(pred[i], ax)\n    return\n\ndef main():\n    args = parse_args()\n    update_config(cfg, args)\n\n    logging.basicConfig(format='%(asctime)-15s %(message)s')\n    logger = logging.getLogger()\n    logger.setLevel(logging.INFO)\n\n    # cudnn related setting\n    cudnn.benchmark = cfg.CUDNN.BENCHMARK\n    torch.backends.cudnn.deterministic = cfg.CUDNN.DETERMINISTIC\n    torch.backends.cudnn.enabled = cfg.CUDNN.ENABLED\n\n    model = get_pose_net(cfg, is_train=False)\n    \n    # load the pre-trained weights\n    if not os.path.exists('./h36m2Dpose/final_state.pth'):\n        logger.info('Please download the pre-trained model first.')\n        return\n    checkpoint = torch.load('./h36m2Dpose/final_state.pth')\n    model.load_state_dict(checkpoint)\n    \n    dump_input = torch.rand(\n        (1, 3, cfg.MODEL.IMAGE_SIZE[1], cfg.MODEL.IMAGE_SIZE[0])\n    )\n    logger.info(get_model_summary(model, dump_input))\n    \n    # modify the configuration file for multiple GPUs\n    model = torch.nn.DataParallel(model, device_ids=cfg.GPUS).cuda() \n    \n    inputs = gather_inputs(args, logger)\n    visualize(inputs, model)\n\nif __name__ == '__main__':\n    main()"
  },
  {
    "path": "examples/inference.py",
    "content": "\"\"\"\nAm examplar script showing inference on the newly collected images in U3DPW. \n\"\"\"\n\nimport sys\nsys.path.append(\"../\")\n\nimport libs.model.model as libm\nfrom libs.dataset.h36m.data_utils import unNormalizeData\n\nimport torch\nimport numpy as np\nimport imageio\nimport matplotlib.pyplot as plt\n\nnum_joints = 16\ngt_3d = False  \npose_connection = [[0,1], [1,2], [2,3], [0,4], [4,5], [5,6], [0,7], [7,8],\n                   [8,9], [9,10], [8,11], [11,12], [12,13], [8, 14], [14, 15], [15,16]]\n# 16 out of 17 key-points are used as inputs in this examplar model\nre_order_indices= [0, 1, 2, 3, 4, 5, 6, 7, 8, 10, 11, 12, 13, 14, 15, 16]\n# paths\ndata_dic_path = './example_annot.npy'     \nmodel_path = './example_model.th'\nstats = np.load('./stats.npy', allow_pickle=True).item()\ndim_used_2d = stats['dim_use_2d']\nmean_2d = stats['mean_2d']\nstd_2d = stats['std_2d'] \n# load the checkpoint and statistics\nckpt = torch.load(model_path)\ndata_dic = np.load(data_dic_path, allow_pickle=True).item()\n# initialize the model\ncascade = libm.get_cascade()\ninput_size = 32\noutput_size = 48\nfor stage_id in range(2):\n    # initialize a single deep learner\n    stage_model = libm.get_model(stage_id + 1,\n                                 refine_3d=False,\n                                 norm_twoD=False, \n                                 num_blocks=2,\n                                 input_size=input_size,\n                                 output_size=output_size,\n                                 linear_size=1024,\n                                 dropout=0.5,\n                                 leaky=False)\n    cascade.append(stage_model)\ncascade.load_state_dict(ckpt)\ncascade.eval()\n# process and show total_to_show examples\ncount = 0\ntotal_to_show = 10\n\ndef draw_skeleton(ax, skeleton, gt=False, add_index=True):\n    for segment_idx in range(len(pose_connection)):\n        point1_idx = pose_connection[segment_idx][0]\n        point2_idx = pose_connection[segment_idx][1]\n        point1 = skeleton[point1_idx]\n        point2 = skeleton[point2_idx]\n        color = 'k' if gt else 'r'\n        plt.plot([int(point1[0]),int(point2[0])], \n                 [int(point1[1]),int(point2[1])], \n                 c=color, \n                 linewidth=2)\n    if add_index:\n        for (idx, re_order_idx) in enumerate(re_order_indices):\n            plt.text(skeleton[re_order_idx][0], \n                     skeleton[re_order_idx][1],\n                     str(idx+1), \n                     color='b'\n                     )\n    return\n\ndef normalize(skeleton, re_order=None):\n    norm_skel = skeleton.copy()\n    if re_order is not None:\n        norm_skel = norm_skel[re_order].reshape(32)\n    norm_skel = norm_skel.reshape(16, 2)\n    mean_x = np.mean(norm_skel[:,0])\n    std_x = np.std(norm_skel[:,0])\n    mean_y = np.mean(norm_skel[:,1])\n    std_y = np.std(norm_skel[:,1])\n    denominator = (0.5*(std_x + std_y))\n    norm_skel[:,0] = (norm_skel[:,0] - mean_x)/denominator\n    norm_skel[:,1] = (norm_skel[:,1] - mean_y)/denominator\n    norm_skel = norm_skel.reshape(32)         \n    return norm_skel\n\ndef get_pred(cascade, data):\n    \"\"\"\n    Get prediction from a cascaded model\n    \"\"\"\n    # forward pass to get prediction for the first stage\n    num_stages = len(cascade)\n    # for legacy code that does not have the num_blocks attribute\n    for i in range(len(cascade)):\n        cascade[i].num_blocks = len(cascade[i].res_blocks)\n    prediction = cascade[0](data)\n    # prediction for later stages\n    for stage_idx in range(1, num_stages):\n        prediction += cascade[stage_idx](data)\n    return prediction\n\ndef show3Dpose(channels, \n               ax, \n               lcolor=\"#3498db\", \n               rcolor=\"#e74c3c\", \n               add_labels=True,\n               gt=False,\n               pred=False\n               ):\n    vals = np.reshape( channels, (32, -1) )\n    I   = np.array([1,2,3,1,7,8,1, 13,14,15,14,18,19,14,26,27])-1 # start points\n    J   = np.array([2,3,4,7,8,9,13,14,15,16,18,19,20,26,27,28])-1 # end points\n    LR  = np.array([1,1,1,0,0,0,0, 0, 0, 0, 0, 0, 0, 1, 1, 1], dtype=bool)\n    # Make connection matrix\n    for i in np.arange( len(I) ):\n        x, y, z = [np.array( [vals[I[i], j], vals[J[i], j]] ) for j in range(3)]\n        if gt or pred:\n            color = 'k' if gt else 'r'\n            ax.plot(x,y, z,  lw=2, c=color)        \n        else:\n            ax.plot(x,y, z,  lw=2, c=lcolor if LR[i] else rcolor)\n    RADIUS = 750 # space around the subject\n    xroot, yroot, zroot = vals[0,0], vals[0,1], vals[0,2]\n    ax.set_xlim3d([-RADIUS+xroot, RADIUS+xroot])\n    ax.set_zlim3d([-RADIUS+zroot, RADIUS+zroot])\n    ax.set_ylim3d([-RADIUS+yroot, RADIUS+yroot])\n    if add_labels:\n        ax.set_xlabel(\"x\")\n        ax.set_ylabel(\"z\")\n        ax.set_zlabel(\"y\")\n    ax.set_aspect('equal')\n    # Get rid of the panes (actually, make them white)\n    white = (1.0, 1.0, 1.0, 0.0)\n    ax.w_xaxis.set_pane_color(white)\n    ax.w_yaxis.set_pane_color(white)\n    # Get rid of the lines in 3d\n    ax.w_xaxis.line.set_color(white)\n    ax.w_yaxis.line.set_color(white)\n    ax.w_zaxis.line.set_color(white)\n    ax.invert_zaxis()\n    return\n\ndef re_order(skeleton):\n    skeleton = skeleton.copy().reshape(-1,3)\n    # permute the order of x,y,z axis\n    skeleton[:,[0,1,2]] = skeleton[:, [0,2,1]]    \n    return skeleton.reshape(96)\n    \ndef plot_3d_ax(ax, \n               elev, \n               azim, \n               pred, \n               title=None\n               ):\n    ax.view_init(elev=elev, azim=azim)\n    show3Dpose(re_order(pred), ax)     \n    plt.title(title)    \n    return\n\ndef adjust_figure(left = 0, \n                  right = 1, \n                  bottom = 0.01, \n                  top = 0.95,\n                  wspace = 0, \n                  hspace = 0.4\n                  ):  \n    plt.subplots_adjust(left, bottom, right, top, wspace, hspace)\n    return\n\nfor image_name in data_dic.keys():\n    image_path = './imgs/' + image_name\n    img = imageio.imread(image_path)\n    f = plt.figure(figsize=(9, 3))\n    ax1 = plt.subplot(131)\n    ax1.imshow(img)\n    plt.title('Input image')\n    ax2 = plt.subplot(132)\n    plt.title('2D key-point inputs: {:d}*2'.format(num_joints))\n    ax2.set_aspect('equal')\n    ax2.invert_yaxis()\n    skeleton_pred = None\n    skeleton_2d = data_dic[image_name]['p2d']\n    # The order for the 2D keypoints is:\n    # 'Hip', 'RHip', 'RKnee', 'RFoot', 'LHip', 'LKnee', 'LFoot', 'Spine', \n    # 'Thorax', 'Neck/Nose', 'Head', 'LShoulder', 'LElbow', 'LWrist', 'RShoulder'\n    # 'RElbow', 'RWrist'\n    draw_skeleton(ax2, skeleton_2d, gt=True)\n    plt.plot(skeleton_2d[:,0], skeleton_2d[:,1], 'ro', 2)       \n    # Nose was not used for this examplar model\n    norm_ske_gt = normalize(skeleton_2d, re_order_indices).reshape(1,-1)\n    pred = get_pred(cascade, torch.from_numpy(norm_ske_gt.astype(np.float32)))      \n    pred = unNormalizeData(pred.data.numpy(),\n                           stats['mean_3d'],\n                           stats['std_3d'],\n                           stats['dim_ignore_3d']\n                           )      \n    ax3 = plt.subplot(133, projection='3d')\n    plot_3d_ax(ax=ax3, \n               pred=pred, \n               elev=10., \n               azim=-90,\n               title='3D prediction'\n               )    \n    adjust_figure(left = 0.05, \n                  right = 0.95, \n                  bottom = 0.08, \n                  top = 0.92,\n                  wspace = 0.3, \n                  hspace = 0.3\n                  )       \n    count += 1       \n    if count >= total_to_show:\n        break"
  },
  {
    "path": "libs/__init__.py",
    "content": "#!/usr/bin/env python3\n# -*- coding: utf-8 -*-\n\"\"\"\nEmpty file.\n\"\"\"\n\n\n"
  },
  {
    "path": "libs/annotator/__init__.py",
    "content": "#!/usr/bin/env python3\n# -*- coding: utf-8 -*-\n\"\"\"\nEmpty file.\n\"\"\"\n#import libs.dataset.h36m\n\n"
  },
  {
    "path": "libs/annotator/angle.py",
    "content": "import numpy as np \nimport scipy.io as sio \nimport matplotlib.pyplot as plt \nfrom mpl_toolkits.mplot3d import Axes3D \n\nbone_name = {\n 1: 'thorax to head top',\n 2: 'left shoulder to left elbow',\n 3: 'left elbow to left wrist',\n 4: 'right shoulder to right elbow',\n 5: 'right elbow to right wrist',\n 6: 'left hip to left knee',\n 7: 'left knee to left ankle',\n 8: 'right hip to right knee',\n 9: 'right knee to right ankle'\n}      \n\nstatic_pose_path = \"/media/nicholas/Database/Github/EvoSkeleton/resources/constraints/staticPose.npy\"\nstatic_pose = np.load(static_pose_path, allow_pickle=True).item()\ndi = static_pose['di']\na = static_pose['a'].reshape(3)\n\ndi_indices = {2:5, 4:2, 6:13, 8:9}\nnt_parent_indices = [13, 17, 18, 25, 26, 6, 7, 1, 2]\nnt_child_indices = [15, 18, 19, 26, 27, 7, 8, 2, 3]\n\ndef gram_schmidt_columns(X):\n    B = np.zeros(X.shape)\n    B[:, 0] = (1/np.linalg.norm(X[:, 0]))*X[:, 0]\n    for i in range(1, 3):\n        v = X[:, i]\n        U = B[:, 0:i] # subspace basis which has already been orthonormalized\n        pc = U.T @ v # orthogonal projection coefficients of v onto U\n        p = U@pc\n        v = v - p\n        if np.linalg.norm(v) < 2e-16:\n            # vectors are not linearly independent!\n            raise ValueError\n        else:\n            v = normalize(v)\n            B[:, i] = v\n    return B\n\ndef normalize(vector):\n    return vector/np.linalg.norm(vector)\n\ndef get_basis1(skeleton):\n    \"\"\"\n    get local coordinate system\n    \"\"\"\n    # compute the vector from the left shoulder to the right shoulder\n    left_shoulder = skeleton[17]\n    right_shoulder = skeleton[25]\n    v1 = normalize(right_shoulder - left_shoulder)    \n    # compute the backbone vector from the thorax to the spine \n    thorax = skeleton[13]\n    spine = skeleton[12]\n    v2 = normalize(spine - thorax)\n    # v3 is the cross product of v1 and v2 (front-facing vector for upper-body)\n    v3 = normalize(np.cross(v1, v2))    \n    return v1, v2, v3\n\ndef get_normal(x1, a, x):\n    nth = 1e-4\n    # x and a are parallel\n    if np.linalg.norm(x - a) < nth or np.linalg.norm(x + a) < nth:   \n        n = np.cross(x, x1)\n        flag = True\n    else:\n        n = np.cross(a, x)\n        flag = False\n    return normalize(n), flag\n\ndef to_spherical(xyz):\n    \"\"\"\n    convert from Cartisian coordinate to spherical coordinate\n    theta: [-pi, pi]\n    phi: [-pi/2, pi/2]\n    note that xyz should be float number\n    \"\"\"\n    # return in r, phi, and theta (elevation angle from z axis down)\n    return_value = np.zeros(xyz.shape, dtype=xyz.dtype)\n    xy = xyz[:,0]**2 + xyz[:,1]**2\n    return_value[:,0] = np.sqrt(xy + xyz[:,2]**2) # r      \n    return_value[:,1] = np.arctan2(xyz[:,1], xyz[:,0]) # theta\n    return_value[:,2] = np.arctan2(xyz[:,2], np.sqrt(xy)) # phi\n    \n    return return_value\n\ndef to_xyz(rthetaphi):\n    \"\"\"\n    convert from spherical coordinate to Cartisian coordinate\n    theta: [0, 2*pi] or [-pi, pi]\n    phi: [-pi/2, pi/2]\n    \"\"\"\n    return_value = np.zeros(rthetaphi.shape, dtype=rthetaphi.dtype)\n    sintheta = np.sin(rthetaphi[:,1])\n    costheta = np.cos(rthetaphi[:,1])\n    sinphi = np.sin(rthetaphi[:,2])\n    cosphi = np.cos(rthetaphi[:,2])\n    return_value[:,0] = rthetaphi[:,0]*costheta*cosphi # x\n    return_value[:,1] = rthetaphi[:,0]*sintheta*cosphi # y\n    return_value[:,2] = rthetaphi[:,0]*sinphi #z\n    return return_value\n"
  },
  {
    "path": "libs/annotator/fit_3d.py",
    "content": "\"\"\"\nObtain 3D skeleton with 2D key-points as inputs using SMPLify\nPlease run this script in Python 2 environment for now.\nTODO: transfer this tool to Python 3.\n\"\"\"\n\nfrom smpl_webuser.serialization import load_model \nfrom smplify.fit_3d import run_single_fit\n# you can use multi-processing to fit the images in parallel\n#from multiprocessing import Pool \n\nimport matplotlib.pyplot as plt \nimport numpy as np \nimport argparse\nimport cv2\nimport os \n\n# whether to use the regularization terms\nsph_regs = None\n# number of blend shapes to use\nn_betas = 1\n# focal length of the camera\nflength = 5000\npix_thsh = 25\nscale_factor = 1\n# viewpoints for rendering\ndo_degrees = [0.]\n\ndef main(opt):\n    model = load_model(opt.model_dir)\n    annotation_path = os.path.join(opt.dataset_dir, 'annotation.npy')\n    assert os.path.exists(annotation_path), \"Please prepare the 2D annotation first.\"\n    annotation = np.load(annotation_path, allow_pickle=True).item()\n    if opt.save_image and not os.path.exists(os.path.join(opt.dataset_dir, 'fitted')):\n        os.makedirs(os.path.join(opt.dataset_dir, 'fitted'))\n    for (image_name, annots) in annotation.items():\n        assert 'p2d' in annots, \"The image must be annotated with 2D key-points\"\n        joints_2d = annots['p2d']\n        # use 14 key-points for model fitting\n        # one may adjust this number by considering different 2D-3D corespondence\n        if joints_2d.shape[0] > 14:\n            joints_2d = joints_2d[:14, :]        \n        # Prepare fitting parameters\n        filling_list = range(12) + [13]\n        conf = np.zeros(joints_2d.shape[0])\n        conf[filling_list] = 1.0\n\n        img = cv2.imread(os.path.join(opt.dataset_dir, image_name))\n        \n        # Run single fit\n        params, vis = run_single_fit(\n            img,\n            joints_2d, \n            conf, \n            model, \n            regs=sph_regs,\n            n_betas=n_betas,\n            flength=flength,\n            pix_thsh=pix_thsh,\n            scale_factor=scale_factor,\n            viz=opt.viz,\n            do_degrees=do_degrees\n        )\n\n        # Show result then close after 1 second\n        f = plt.figure(figsize=(6, 3))\n        plt.subplot(121)\n        plt.imshow(cv2.cvtColor(img, cv2.COLOR_BGR2RGB))\n\n        for di, deg in enumerate(do_degrees):\n            plt.subplot(122)\n            plt.cla()\n            plt.imshow(vis[di])\n            plt.draw()\n            plt.pause(1)\n\n        # record the params\n        annotation[image_name]['fitting_params'] = params\n        # save fitted image\n        if opt.save_image:\n            image_name = image_name.replace(\".jpg\", \".png\")\n            img_save_path = os.path.join(opt.dataset_dir, 'fitted', image_name)\n            f.savefig(img_save_path)   \n            plt.close(f)\n            print(\"fitted image saved at \", img_save_path)\n            \n    # save the annotation dictionary with fitted parameters\n    np.save(os.path.join(opt.dataset_dir, \"fitted.npy\"), annotation)\n    print('3D prameters saved at ' + os.path.join(opt.dataset_dir, \"fitted.npy\"))\n\nif __name__ == '__main__':\n    parser = argparse.ArgumentParser(description='2D Annotation')\n    parser.add_argument('-dataset_dir', type=str)\n    parser.add_argument('-model_dir', type=str)\n    parser.add_argument('-save_image',default=True, type=bool)\n    # visualize intermeadiate results\n    parser.add_argument('-viz',default=False, type=bool)\n    opt = parser.parse_args() \n    main(opt)"
  },
  {
    "path": "libs/annotator/smpl-spec-list.txt",
    "content": "# This file may be used to create an environment using:\n# $ conda create --name <env> --file <this file>\n# platform: linux-64\n@EXPLICIT\nhttps://repo.anaconda.com/pkgs/main/linux-64/_libgcc_mutex-0.1-main.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/blas-1.0-mkl.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/ca-certificates-2020.6.24-0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/intel-openmp-2020.1-217.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/libgfortran-ng-7.3.0-hdf63c60_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/libstdcxx-ng-9.1.0-hdf63c60_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/pandoc-2.9.2.1-0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/libgcc-ng-9.1.0-hdf63c60_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/mkl-2020.1-217.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/expat-2.2.9-he6710b0_2.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/icu-58.2-he6710b0_3.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/jpeg-9b-h024ee3a_2.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/libffi-3.3-he6710b0_1.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/libsodium-1.0.18-h7b6447c_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/libspatialindex-1.9.3-he6710b0_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/libuuid-1.0.3-h1bed415_2.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/libxcb-1.14-h7b6447c_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/ncurses-6.2-he6710b0_1.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/openssl-1.1.1g-h7b6447c_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/pcre-8.44-he6710b0_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/pixman-0.40.0-h7b6447c_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/xz-5.2.5-h7b6447c_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/yaml-0.1.7-had09818_2.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/zlib-1.2.11-h7b6447c_3.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/glib-2.65.0-h3eb4bd4_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/free/linux-64/hdf5-1.8.17-2.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/libedit-3.1.20191231-h7b6447c_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/libpng-1.6.37-hbc83047_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/libxml2-2.9.10-he19cac6_1.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/readline-8.0-h7b6447c_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/tk-8.6.10-hbc83047_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/zeromq-4.3.2-he6710b0_2.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/zstd-1.3.7-h0b5b093_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/dbus-1.13.16-hb2f20db_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/freetype-2.10.2-h5ab3b9f_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/gstreamer-1.14.0-hb31296c_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/libtiff-4.0.10-h2733197_2.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/sqlite-3.32.3-h62c20be_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/fontconfig-2.13.0-h9420a91_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/gst-plugins-base-1.14.0-hbbd80ab_1.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/python-2.7.18-h15b4118_1.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/noarch/alabaster-0.7.12-py_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/argh-0.26.2-py27_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/asn1crypto-1.3.0-py27_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/noarch/atomicwrites-1.4.0-py_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/noarch/attrs-19.3.0-py_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/noarch/backports-1.0-py_2.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/noarch/backports_abc-0.5-py_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/cairo-1.14.12-h8948797_3.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/certifi-2019.11.28-py27_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/chardet-3.0.4-py27_1003.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/noarch/cloudpickle-1.2.2-py_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/noarch/contextlib2-0.6.0.post1-py_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/noarch/decorator-4.4.2-py_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/noarch/defusedxml-0.6.0-py_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/noarch/diff-match-patch-20181111-py_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/docutils-0.15.2-py27_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/enum34-1.1.6-py27_1.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/functools32-3.2.3.2-py27_1.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/future-0.18.2-py27_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/futures-3.3.0-py27_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/noarch/idna-2.10-py_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/noarch/imagesize-1.2.0-py_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/noarch/ipaddress-1.0.23-py_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/ipython_genutils-0.2.0-py27_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/lazy-object-proxy-1.4.3-py27h7b6447c_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/markupsafe-1.1.1-py27h7b6447c_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/mccabe-0.6.1-py27_1.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/mistune-0.8.4-py27h7b6447c_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/pandocfilters-1.4.2-py27_1.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/noarch/parso-0.5.2-py_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/noarch/pathtools-0.1.2-py_1.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/psutil-5.6.7-py27h7b6447c_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/ptyprocess-0.6.0-py27_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/pycodestyle-2.5.0-py27_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/noarch/pycparser-2.20-py_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/pyflakes-2.1.1-py27_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/noarch/pyparsing-2.4.7-py_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/pysocks-1.7.1-py27_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/noarch/pytz-2020.1-py_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/noarch/pyxdg-0.26-py_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/pyyaml-5.2-py27h7b6447c_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/pyzmq-18.1.0-py27he6710b0_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/noarch/qdarkstyle-2.8.1-py_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/qt-5.9.7-h5867ecd_1.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/noarch/qtpy-1.9.0-py_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/noarch/rope-0.17.0-py_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/rtree-0.8.3-py27_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/scandir-1.10.0-py27h7b6447c_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/noarch/selectors2-2.0.1-py_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/simplegeneric-0.8.1-py27_2.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/sip-4.19.8-py27hf484d3e_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/noarch/six-1.15.0-py_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/noarch/snowballstemmer-2.0.0-py_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/noarch/sortedcontainers-2.2.2-py_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/sphinxcontrib-1.0-py27_1.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/noarch/testpath-0.4.4-py_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/typing-3.7.4.1-py27_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/ujson-1.35-py27h14c3975_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/noarch/wcwidth-0.2.5-py_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/webencodings-0.5.1-py27_1.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/wrapt-1.11.2-py27h7b6447c_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/noarch/yapf-0.30.0-py_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/noarch/autopep8-1.4.4-py_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/noarch/babel-2.8.0-py_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/backports.shutil_get_terminal_size-1.0.0-py27_2.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/cffi-1.14.0-py27he30daa8_1.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/configparser-4.0.2-py27_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/free/linux-64/harfbuzz-0.9.39-1.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/noarch/intervaltree-3.0.2-py_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/jedi-0.14.1-py27_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/mkl-service-2.3.0-py27he904b0f_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/more-itertools-5.0.0-py27_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/noarch/packaging-20.4-py_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/pathlib2-2.3.5-py27_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/pexpect-4.7.0-py27_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/pyqt-5.9.2-py27h05f1152_2.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/pyrsistent-0.15.6-py27h7b6447c_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/noarch/python-dateutil-2.8.1-py_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/noarch/python-jsonrpc-server-0.3.4-py_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/noarch/qtawesome-0.7.2-py_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/setuptools-44.0.0-py27_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/singledispatch-3.4.0.3-py27_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/noarch/sphinxcontrib-websupport-1.2.3-py_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/traitlets-4.3.3-py27_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/watchdog-0.9.0-py27_1.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/wurlitzer-2.0.0-py27_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/noarch/backports.functools_lru_cache-1.6.1-py_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/noarch/bleach-3.1.5-py_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/cryptography-2.8-py27h1ba5d50_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/entrypoints-0.3-py27_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/noarch/jinja2-2.11.2-py_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/jupyter_core-4.6.1-py27_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/numpy-base-1.16.6-py27hde5b4d6_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/pickleshare-0.7.5-py27_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/pydocstyle-3.0.0-py27_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/noarch/pygments-2.5.2-py_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/tornado-5.1.1-py27h7b6447c_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/wheel-0.33.6-py27_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/noarch/zipp-0.6.0-py_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/astroid-1.6.5-py27_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/flake8-3.7.9-py27_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/importlib_metadata-1.3.0-py27_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/isort-4.3.21-py27_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/jupyter_client-5.3.4-py27_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/pip-19.3.1-py27_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/prompt_toolkit-1.0.15-py27_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/pyopenssl-19.1.0-py27_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/ipython-5.8.0-py27_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/jsonschema-3.2.0-py27_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/pluggy-0.13.1-py27_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/pylint-1.9.2-py27_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/urllib3-1.25.7-py27_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/ipykernel-4.10.0-py27_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/nbformat-4.4.0-py27_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/python-language-server-0.31.2-py27_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/noarch/requests-2.24.0-py_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/nbconvert-5.6.1-py27_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/noarch/qtconsole-4.7.5-py_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/sphinx-1.8.5-py27_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/spyder-kernels-1.8.1-py27_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/noarch/numpydoc-1.0.0-py_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/spyder-4.0.1-py27_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/mkl_fft-1.0.15-py27ha843d7b_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/mkl_random-1.1.0-py27hd6b4f25_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/numpy-1.16.6-py27hbc911f0_0.tar.bz2\nhttps://conda.anaconda.org/menpo/linux-64/opencv3-3.1.0-py27_0.tar.bz2\n"
  },
  {
    "path": "libs/annotator/smpl_webuser/LICENSE.txt",
    "content": "Please read carefully the following terms and conditions and any accompanying documentation before you download and/or use the SMPL body model and software, (the \"Model\"), including 3D meshes, blend weights blend shapes, textures, software, scripts, and animations. By downloading and/or using the Model, you acknowledge that you have read these terms and conditions, understand them, and agree to be bound by them. If you do not agree with these terms and conditions, you must not download and/or use the Model.\n\nOwnership\nThe Model has been developed at the Max Planck Institute for Intelligent Systems (hereinafter \"MPI\") and is owned by and proprietary material of the Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (hereinafter “MPG”; MPI and MPG hereinafter collectively “Max-Planck”).\n\nLicense Grant\nMax-Planck grants you a non-exclusive, non-transferable, free of charge right:\n\nTo download the Model and use it on computers owned, leased or otherwise controlled by you and/or your organisation;\nTo use the Model for the sole purpose of performing non-commercial scientific research, non-commercial education, or non-commercial artistic projects.\nAny other use, in particular any use for commercial purposes, is prohibited. This includes, without limitation, incorporation in a commercial product, use in a commercial service, as training data for a commercial product, for commercial ergonomic analysis (e.g. product design, architectural design, etc.), or production of other artifacts for commercial purposes including, for example, web services, movies, television programs, mobile applications, or video games. The Model may not be used for pornographic purposes or to generate pornographic material whether commercial or not. This license also prohibits the use of the Model to train methods/algorithms/neural networks/etc. for commercial use of any kind. The Model may not be reproduced, modified and/or made available in any form to any third party without Max-Planck’s prior written permission. By downloading the Model, you agree not to reverse engineer it.\n\nDisclaimer of Representations and Warranties\nYou expressly acknowledge and agree that the Model results from basic research, is provided “AS IS”, may contain errors, and that any use of the Model is at your sole risk. MAX-PLANCK MAKES NO REPRESENTATIONS OR WARRANTIES OF ANY KIND CONCERNING THE MODEL, NEITHER EXPRESS NOR IMPLIED, AND THE ABSENCE OF ANY LEGAL OR ACTUAL DEFECTS, WHETHER DISCOVERABLE OR NOT. Specifically, and not to limit the foregoing, Max-Planck makes no representations or warranties (i) regarding the merchantability or fitness for a particular purpose of the Model, (ii) that the use of the Model will not infringe any patents, copyrights or other intellectual property rights of a third party, and (iii) that the use of the Model will not cause any damage of any kind to you or a third party.\n\nLimitation of Liability\nUnder no circumstances shall Max-Planck be liable for any incidental, special, indirect or consequential damages arising out of or relating to this license, including but not limited to, any lost profits, business interruption, loss of programs or other data, or all other commercial damages or losses, even if advised of the possibility thereof.\n\nNo Maintenance Services\nYou understand and agree that Max-Planck is under no obligation to provide either maintenance services, update services, notices of latent defects, or corrections of defects with regard to the Model. Max-Planck nevertheless reserves the right to update, modify, or discontinue the Model at any time.\n\nPublication with SMPL\nYou agree to cite the most recent paper describing the model as specified on the download website. This website lists the most up to date bibliographic information on the about page.\n\nMedia projects with SMPL\nWhen using SMPL in a media project please give credit to Max Planck Institute for Intelligent Systems. For example: SMPL was used for character animation courtesy of the Max Planck Institute for Intelligent Systems.\nCommercial licensing opportunities\nFor commercial use in the fields of medicine, psychology, and biomechanics, please contact ps-license@tue.mpg.de.\nFor commercial use in all other fields please contact Body Labs Inc at smpl@bodylabs.com"
  },
  {
    "path": "libs/annotator/smpl_webuser/README.txt",
    "content": "License:\n========\nTo learn about SMPL, please visit our website: http://smpl.is.tue.mpg\nYou can find the SMPL paper at: http://files.is.tue.mpg.de/black/papers/SMPL2015.pdf\n\nVisit our downloads page to download some sample animation files (FBX), and python code:\nhttp://smpl.is.tue.mpg/downloads\n\nFor comments or questions, please email us at: smpl@tuebingen.mpg.de\n\n\nSystem Requirements:\n====================\nOperating system: OSX, Linux\n\nPython Dependencies:\n- Numpy & Scipy  [http://www.scipy.org/scipylib/download.html]\n- Chumpy \t\t [https://github.com/mattloper/chumpy]\n- OpenCV \t\t [http://opencv.org/downloads.html] \n\n\nGetting Started:\n================\n\n1. Extract the Code:\n--------------------\nExtract the 'smpl.zip' file to your home directory (or any other location you wish)\n\n\n2. Set the PYTHONPATH:\n----------------------\nWe need to update the PYTHONPATH environment variable so that the system knows how to find the SMPL code. Add the following lines to your ~/.bash_profile file (create it if it doesn't exist; Linux users might have ~/.bashrc file instead), replacing ~/smpl with the location where you extracted the smpl.zip file:\n\n\tSMPL_LOCATION=~/smpl\n\texport PYTHONPATH=$PYTHONPATH:$SMPL_LOCATION\n\n\nOpen a new terminal window to check if the python path has been updated by typing the following:\n>  echo $PYTHONPATH\n\n\n3. Run the Hello World scripts:\n-------------------------------\nIn the new Terminal window, navigate to the smpl/smpl_webuser/hello_world directory. You can run the hello world scripts now by typing the following:\n\n> python hello_smpl.py\n\nOR \n\n> python render_smpl.py\n\n\n\nNote:\nBoth of these scripts will require the dependencies listed above. The scripts are provided as a sample to help you get started. \n\n"
  },
  {
    "path": "libs/annotator/smpl_webuser/__init__.py",
    "content": "'''\nCopyright 2015 Matthew Loper, Naureen Mahmood and the Max Planck Gesellschaft.  All rights reserved.\nThis software is provided for research purposes only.\nBy using this software you agree to the terms of the SMPL Model license here http://smpl.is.tue.mpg.de/license\n\nMore information about SMPL is available here http://smpl.is.tue.mpg.\nFor comments or questions, please email us at: smpl@tuebingen.mpg.de\n\n\nAbout this file:\n================\nThis is an initialization file to help python look for submodules in this directory.\n'''"
  },
  {
    "path": "libs/annotator/smpl_webuser/hello_world/hello_smpl.py",
    "content": "'''\nCopyright 2015 Matthew Loper, Naureen Mahmood and the Max Planck Gesellschaft.  All rights reserved.\nThis software is provided for research purposes only.\nBy using this software you agree to the terms of the SMPL Model license here http://smpl.is.tue.mpg.de/license\n\nMore information about SMPL is available here http://smpl.is.tue.mpg.\nFor comments or questions, please email us at: smpl@tuebingen.mpg.de\n\n\nPlease Note:\n============\nThis is a demo version of the script for driving the SMPL model with python.\nWe would be happy to receive comments, help and suggestions on improving this code \nand in making it available on more platforms. \n\n\nSystem Requirements:\n====================\nOperating system: OSX, Linux\n\nPython Dependencies:\n- Numpy & Scipy  [http://www.scipy.org/scipylib/download.html]\n- Chumpy [https://github.com/mattloper/chumpy]\n\n\nAbout the Script:\n=================\nThis script demonstrates a few basic functions to help users get started with using \nthe SMPL model. The code shows how to:\n  - Load the SMPL model\n  - Edit pose & shape parameters of the model to create a new body in a new pose\n  - Save the resulting body as a mesh in .OBJ format\n\n\nRunning the Hello World code:\n=============================\nInside Terminal, navigate to the smpl/webuser/hello_world directory. You can run \nthe hello world script now by typing the following:\n>\tpython hello_smpl.py\n\n'''\n\nfrom smpl_webuser.serialization import load_model\nimport numpy as np\n\n## Load SMPL model (here we load the female model)\n## Make sure path is correct\nm = load_model( '../../models/basicModel_f_lbs_10_207_0_v1.0.0.pkl' )\n\n## Assign random pose and shape parameters\nm.pose[:] = np.random.rand(m.pose.size) * .2\nm.betas[:] = np.random.rand(m.betas.size) * .03\n\n## Write to an .obj file\noutmesh_path = './hello_smpl.obj'\nwith open( outmesh_path, 'w') as fp:\n    for v in m.r:\n        fp.write( 'v %f %f %f\\n' % ( v[0], v[1], v[2]) )\n\n    for f in m.f+1: # Faces are 1-based, not 0-based in obj files\n        fp.write( 'f %d %d %d\\n' %  (f[0], f[1], f[2]) )\n\n## Print message\nprint '..Output mesh saved to: ', outmesh_path \n"
  },
  {
    "path": "libs/annotator/smpl_webuser/hello_world/render_smpl.py",
    "content": "'''\nCopyright 2015 Matthew Loper, Naureen Mahmood and the Max Planck Gesellschaft.  All rights reserved.\nThis software is provided for research purposes only.\nBy using this software you agree to the terms of the SMPL Model license here http://smpl.is.tue.mpg.de/license\n\nMore information about SMPL is available here http://smpl.is.tue.mpg.\nFor comments or questions, please email us at: smpl@tuebingen.mpg.de\n\n\nPlease Note:\n============\nThis is a demo version of the script for driving the SMPL model with python.\nWe would be happy to receive comments, help and suggestions on improving this code \nand in making it available on more platforms. \n\n\nSystem Requirements:\n====================\nOperating system: OSX, Linux\n\nPython Dependencies:\n- Numpy & Scipy  [http://www.scipy.org/scipylib/download.html]\n- Chumpy [https://github.com/mattloper/chumpy]\n- OpenCV [http://opencv.org/downloads.html] \n  --> (alternatively: matplotlib [http://matplotlib.org/downloads.html])\n\n\nAbout the Script:\n=================\nThis script demonstrates loading the smpl model and rendering it using OpenDR \nto render and OpenCV to display (or alternatively matplotlib can also be used\nfor display, as shown in commented code below). \n\nThis code shows how to:\n  - Load the SMPL model\n  - Edit pose & shape parameters of the model to create a new body in a new pose\n  - Create an OpenDR scene (with a basic renderer, camera & light)\n  - Render the scene using OpenCV / matplotlib\n\n\nRunning the Hello World code:\n=============================\nInside Terminal, navigate to the smpl/webuser/hello_world directory. You can run \nthe hello world script now by typing the following:\n>\tpython render_smpl.py\n\n\n'''\n\nimport numpy as np\nfrom opendr.renderer import ColoredRenderer\nfrom opendr.lighting import LambertianPointLight\nfrom opendr.camera import ProjectPoints\nfrom smpl_webuser.serialization import load_model\n\n## Load SMPL model (here we load the female model)\nm = load_model('../../models/basicModel_f_lbs_10_207_0_v1.0.0.pkl')\n\n## Assign random pose and shape parameters\nm.pose[:] = np.random.rand(m.pose.size) * .2\nm.betas[:] = np.random.rand(m.betas.size) * .03\nm.pose[0] = np.pi\n\n## Create OpenDR renderer\nrn = ColoredRenderer()\n\n## Assign attributes to renderer\nw, h = (640, 480)\n\nrn.camera = ProjectPoints(v=m, rt=np.zeros(3), t=np.array([0, 0, 2.]), f=np.array([w,w])/2., c=np.array([w,h])/2., k=np.zeros(5))\nrn.frustum = {'near': 1., 'far': 10., 'width': w, 'height': h}\nrn.set(v=m, f=m.f, bgcolor=np.zeros(3))\n\n## Construct point light source\nrn.vc = LambertianPointLight(\n    f=m.f,\n    v=rn.v,\n    num_verts=len(m),\n    light_pos=np.array([-1000,-1000,-2000]),\n    vc=np.ones_like(m)*.9,\n    light_color=np.array([1., 1., 1.]))\n\n\n## Show it using OpenCV\nimport cv2\ncv2.imshow('render_SMPL', rn.r)\nprint ('..Print any key while on the display window')\ncv2.waitKey(0)\ncv2.destroyAllWindows()\n\n\n## Could also use matplotlib to display\n# import matplotlib.pyplot as plt\n# plt.ion()\n# plt.imshow(rn.r)\n# plt.show()\n# import pdb; pdb.set_trace()"
  },
  {
    "path": "libs/annotator/smpl_webuser/lbs.py",
    "content": "'''\nCopyright 2015 Matthew Loper, Naureen Mahmood and the Max Planck Gesellschaft.  All rights reserved.\nThis software is provided for research purposes only.\nBy using this software you agree to the terms of the SMPL Model license here http://smpl.is.tue.mpg.de/license\n\nMore information about SMPL is available here http://smpl.is.tue.mpg.\nFor comments or questions, please email us at: smpl@tuebingen.mpg.de\n\n\nAbout this file:\n================\nThis file defines linear blend skinning for the SMPL loader which \ndefines the effect of bones and blendshapes on the vertices of the template mesh.\n\nModules included:\n- global_rigid_transformation: \n  computes global rotation & translation of the model\n- verts_core: [overloaded function inherited from verts.verts_core]\n  computes the blending of joint-influences for each vertex based on type of skinning\n\n'''\n\nfrom smpl_webuser.posemapper import posemap\nimport chumpy\nimport numpy as np\n\ndef global_rigid_transformation(pose, J, kintree_table, xp):\n    results = {}\n    pose = pose.reshape((-1,3))\n    id_to_col = {kintree_table[1,i] : i for i in range(kintree_table.shape[1])}\n    parent = {i : id_to_col[kintree_table[0,i]] for i in range(1, kintree_table.shape[1])}\n\n    if xp == chumpy:\n        from smpl_webuser.posemapper import Rodrigues\n        rodrigues = lambda x : Rodrigues(x)\n    else:\n        import cv2\n        rodrigues = lambda x : cv2.Rodrigues(x)[0]\n\n    with_zeros = lambda x : xp.vstack((x, xp.array([[0.0, 0.0, 0.0, 1.0]])))\n    results[0] = with_zeros(xp.hstack((rodrigues(pose[0,:]), J[0,:].reshape((3,1)))))        \n        \n    for i in range(1, kintree_table.shape[1]):\n        results[i] = results[parent[i]].dot(with_zeros(xp.hstack((\n            rodrigues(pose[i,:]),\n            ((J[i,:] - J[parent[i],:]).reshape((3,1)))\n            ))))\n\n    pack = lambda x : xp.hstack([np.zeros((4, 3)), x.reshape((4,1))])\n    \n    results = [results[i] for i in sorted(results.keys())]\n    results_global = results\n\n    if True:\n        results2 = [results[i] - (pack(\n            results[i].dot(xp.concatenate( ( (J[i,:]), 0 ) )))\n            ) for i in range(len(results))]\n        results = results2\n    result = xp.dstack(results)\n    return result, results_global\n\n\ndef verts_core(pose, v, J, weights, kintree_table, want_Jtr=False, xp=chumpy):\n    A, A_global = global_rigid_transformation(pose, J, kintree_table, xp)\n    T = A.dot(weights.T)\n\n    rest_shape_h = xp.vstack((v.T, np.ones((1, v.shape[0]))))\n        \n    v =(T[:,0,:] * rest_shape_h[0, :].reshape((1, -1)) + \n        T[:,1,:] * rest_shape_h[1, :].reshape((1, -1)) + \n        T[:,2,:] * rest_shape_h[2, :].reshape((1, -1)) + \n        T[:,3,:] * rest_shape_h[3, :].reshape((1, -1))).T\n\n    v = v[:,:3] \n    \n    if not want_Jtr:\n        return v\n    Jtr = xp.vstack([g[:3,3] for g in A_global])\n    return (v, Jtr)\n    \n"
  },
  {
    "path": "libs/annotator/smpl_webuser/posemapper.py",
    "content": "'''\nCopyright 2015 Matthew Loper, Naureen Mahmood and the Max Planck Gesellschaft.  All rights reserved.\nThis software is provided for research purposes only.\nBy using this software you agree to the terms of the SMPL Model license here http://smpl.is.tue.mpg.de/license\n\nMore information about SMPL is available here http://smpl.is.tue.mpg.\nFor comments or questions, please email us at: smpl@tuebingen.mpg.de\n\n\nAbout this file:\n================\nThis module defines the mapping of joint-angles to pose-blendshapes. \n\nModules included:\n- posemap:\n  computes the joint-to-pose blend shape mapping given a mapping type as input\n\n'''\n\nimport chumpy as ch\nimport numpy as np\nimport cv2\n\n\nclass Rodrigues(ch.Ch):\n    dterms = 'rt'\n    \n    def compute_r(self):\n        return cv2.Rodrigues(self.rt.r)[0]\n    \n    def compute_dr_wrt(self, wrt):\n        if wrt is self.rt:\n            return cv2.Rodrigues(self.rt.r)[1].T\n\n\ndef lrotmin(p): \n    if isinstance(p, np.ndarray):\n        p = p.ravel()[3:]\n        return np.concatenate([(cv2.Rodrigues(np.array(pp))[0]-np.eye(3)).ravel() for pp in p.reshape((-1,3))]).ravel()        \n    if p.ndim != 2 or p.shape[1] != 3:\n        p = p.reshape((-1,3))\n    p = p[1:]\n    return ch.concatenate([(Rodrigues(pp)-ch.eye(3)).ravel() for pp in p]).ravel()\n\ndef posemap(s):\n    if s == 'lrotmin':\n        return lrotmin\n    else:\n        raise Exception('Unknown posemapping: %s' % (str(s),))\n"
  },
  {
    "path": "libs/annotator/smpl_webuser/serialization.py",
    "content": "'''\nCopyright 2015 Matthew Loper, Naureen Mahmood and the Max Planck Gesellschaft.  All rights reserved.\nThis software is provided for research purposes only.\nBy using this software you agree to the terms of the SMPL Model license here http://smpl.is.tue.mpg.de/license\n\nMore information about SMPL is available here http://smpl.is.tue.mpg.\nFor comments or questions, please email us at: smpl@tuebingen.mpg.de\n\n\nAbout this file:\n================\nThis file defines the serialization functions of the SMPL model. \n\nModules included:\n- save_model:\n  saves the SMPL model to a given file location as a .pkl file\n- load_model:\n  loads the SMPL model from a given file location (i.e. a .pkl file location), \n  or a dictionary object.\n\n'''\n\n__all__ = ['load_model', 'save_model']\n\nimport numpy as np\nimport pickle\nimport chumpy as ch\nfrom chumpy.ch import MatVecMult\nfrom smpl_webuser.posemapper import posemap\nfrom smpl_webuser.verts import verts_core\n    \ndef save_model(model, fname):\n    m0 = model\n    trainer_dict = {'v_template': np.asarray(m0.v_template),'J': np.asarray(m0.J),'weights': np.asarray(m0.weights),'kintree_table': m0.kintree_table,'f': m0.f, 'bs_type': m0.bs_type, 'posedirs': np.asarray(m0.posedirs)}    \n    if hasattr(model, 'J_regressor'):\n        trainer_dict['J_regressor'] = m0.J_regressor\n    if hasattr(model, 'J_regressor_prior'):\n        trainer_dict['J_regressor_prior'] = m0.J_regressor_prior\n    if hasattr(model, 'weights_prior'):\n        trainer_dict['weights_prior'] = m0.weights_prior\n    if hasattr(model, 'shapedirs'):\n        trainer_dict['shapedirs'] = m0.shapedirs\n    if hasattr(model, 'vert_sym_idxs'):\n        trainer_dict['vert_sym_idxs'] = m0.vert_sym_idxs\n    if hasattr(model, 'bs_style'):\n        trainer_dict['bs_style'] = model.bs_style\n    else:\n        trainer_dict['bs_style'] = 'lbs'\n    pickle.dump(trainer_dict, open(fname, 'w'), -1)\n\n\ndef backwards_compatibility_replacements(dd):\n\n    # replacements\n    if 'default_v' in dd:\n        dd['v_template'] = dd['default_v']\n        del dd['default_v']\n    if 'template_v' in dd:\n        dd['v_template'] = dd['template_v']\n        del dd['template_v']\n    if 'joint_regressor' in dd:\n        dd['J_regressor'] = dd['joint_regressor']\n        del dd['joint_regressor']\n    if 'blendshapes' in dd:\n        dd['posedirs'] = dd['blendshapes']\n        del dd['blendshapes']\n    if 'J' not in dd:\n        dd['J'] = dd['joints']\n        del dd['joints']\n\n    # defaults\n    if 'bs_style' not in dd:\n        dd['bs_style'] = 'lbs'\n\n\n\ndef ready_arguments(fname_or_dict):\n\n    if not isinstance(fname_or_dict, dict):\n        dd = pickle.load(open(fname_or_dict, 'rb'))\n        # dd = pickle.load(open(fname_or_dict, 'rb'), encoding='latin1')\n    else:\n        dd = fname_or_dict\n        \n    backwards_compatibility_replacements(dd)\n        \n    want_shapemodel = 'shapedirs' in dd\n    nposeparms = dd['kintree_table'].shape[1]*3\n\n    if 'trans' not in dd:\n        dd['trans'] = np.zeros(3)\n    if 'pose' not in dd:\n        dd['pose'] = np.zeros(nposeparms)\n    if 'shapedirs' in dd and 'betas' not in dd:\n        dd['betas'] = np.zeros(dd['shapedirs'].shape[-1])\n\n    for s in ['v_template', 'weights', 'posedirs', 'pose', 'trans', 'shapedirs', 'betas', 'J']:\n        if (s in dd) and not hasattr(dd[s], 'dterms'):\n            dd[s] = ch.array(dd[s])\n\n    if want_shapemodel:\n        dd['v_shaped'] = dd['shapedirs'].dot(dd['betas'])+dd['v_template']\n        v_shaped = dd['v_shaped']\n        J_tmpx = MatVecMult(dd['J_regressor'], v_shaped[:,0])        \n        J_tmpy = MatVecMult(dd['J_regressor'], v_shaped[:,1])        \n        J_tmpz = MatVecMult(dd['J_regressor'], v_shaped[:,2])        \n        dd['J'] = ch.vstack((J_tmpx, J_tmpy, J_tmpz)).T    \n        dd['v_posed'] = v_shaped + dd['posedirs'].dot(posemap(dd['bs_type'])(dd['pose']))\n    else:    \n        dd['v_posed'] = dd['v_template'] + dd['posedirs'].dot(posemap(dd['bs_type'])(dd['pose']))\n            \n    return dd\n\n\n\ndef load_model(fname_or_dict):\n    dd = ready_arguments(fname_or_dict)\n    \n    args = {\n        'pose': dd['pose'],\n        'v': dd['v_posed'],\n        'J': dd['J'],\n        'weights': dd['weights'],\n        'kintree_table': dd['kintree_table'],\n        'xp': ch,\n        'want_Jtr': True,\n        'bs_style': dd['bs_style']\n    }\n    \n    result, Jtr = verts_core(**args)\n    result = result + dd['trans'].reshape((1,3))\n    result.J_transformed = Jtr + dd['trans'].reshape((1,3))\n\n    for k, v in dd.items():\n        setattr(result, k, v)\n        \n    return result\n\n"
  },
  {
    "path": "libs/annotator/smpl_webuser/verts.py",
    "content": "'''\nCopyright 2015 Matthew Loper, Naureen Mahmood and the Max Planck Gesellschaft.  All rights reserved.\nThis software is provided for research purposes only.\nBy using this software you agree to the terms of the SMPL Model license here http://smpl.is.tue.mpg.de/license\n\nMore information about SMPL is available here http://smpl.is.tue.mpg.\nFor comments or questions, please email us at: smpl@tuebingen.mpg.de\n\n\nAbout this file:\n================\nThis file defines the basic skinning modules for the SMPL loader which \ndefines the effect of bones and blendshapes on the vertices of the template mesh.\n\nModules included:\n- verts_decorated: \n  creates an instance of the SMPL model which inherits model attributes from another \n  SMPL model.\n- verts_core: [overloaded function inherited by lbs.verts_core]\n  computes the blending of joint-influences for each vertex based on type of skinning\n\n'''\n\nimport chumpy\nimport smpl_webuser.lbs as lbs\nfrom smpl_webuser.posemapper import posemap\nimport scipy.sparse as sp\nfrom chumpy.ch import MatVecMult\n\ndef ischumpy(x): return hasattr(x, 'dterms')\n\ndef verts_decorated(trans, pose, \n    v_template, J, weights, kintree_table, bs_style, f,\n    bs_type=None, posedirs=None, betas=None, shapedirs=None, want_Jtr=False):\n\n    for which in [trans, pose, v_template, weights, posedirs, betas, shapedirs]:\n        if which is not None:\n            assert ischumpy(which)\n\n    v = v_template\n\n    if shapedirs is not None:\n        if betas is None:\n            betas = chumpy.zeros(shapedirs.shape[-1])\n        v_shaped = v + shapedirs.dot(betas)\n    else:\n        v_shaped = v\n        \n    if posedirs is not None:\n        v_posed = v_shaped + posedirs.dot(posemap(bs_type)(pose))\n    else:\n        v_posed = v_shaped\n        \n    v = v_posed\n        \n    if sp.issparse(J):\n        regressor = J\n        J_tmpx = MatVecMult(regressor, v_shaped[:,0])        \n        J_tmpy = MatVecMult(regressor, v_shaped[:,1])        \n        J_tmpz = MatVecMult(regressor, v_shaped[:,2])        \n        J = chumpy.vstack((J_tmpx, J_tmpy, J_tmpz)).T            \n    else:    \n        assert(ischumpy(J))\n        \n    assert(bs_style=='lbs')\n    result, Jtr = lbs.verts_core(pose, v, J, weights, kintree_table, want_Jtr=True, xp=chumpy)\n     \n    tr = trans.reshape((1,3))\n    result = result + tr\n    Jtr = Jtr + tr\n\n    result.trans = trans\n    result.f = f\n    result.pose = pose\n    result.v_template = v_template\n    result.J = J\n    result.weights = weights\n    result.kintree_table = kintree_table\n    result.bs_style = bs_style\n    result.bs_type =bs_type\n    if posedirs is not None:\n        result.posedirs = posedirs\n        result.v_posed = v_posed\n    if shapedirs is not None:\n        result.shapedirs = shapedirs\n        result.betas = betas\n        result.v_shaped = v_shaped\n    if want_Jtr:\n        result.J_transformed = Jtr\n    return result\n\ndef verts_core(pose, v, J, weights, kintree_table, bs_style, want_Jtr=False, xp=chumpy):\n    \n    if xp == chumpy:\n        assert(hasattr(pose, 'dterms'))\n        assert(hasattr(v, 'dterms'))\n        assert(hasattr(J, 'dterms'))\n        assert(hasattr(weights, 'dterms'))\n     \n    assert(bs_style=='lbs')\n    result = lbs.verts_core(pose, v, J, weights, kintree_table, want_Jtr, xp)\n\n    return result\n"
  },
  {
    "path": "libs/annotator/smplify/__init__.py",
    "content": ""
  },
  {
    "path": "libs/annotator/smplify/fit_3d.py",
    "content": "\"\"\"\nCopyright 2016 Max Planck Society, Federica Bogo, Angjoo Kanazawa. All rights reserved.\nThis software is provided for research purposes only.\nBy using this software you agree to the terms of the SMPLify license here:\n     http://smplify.is.tue.mpg.de/license\n\nAbout this Script:\n============\nThis is a demo version of the algorithm implemented in the paper,\nwhich fits the SMPL body model to the image given the joint detections.\nThe code is organized to be run on the LSP dataset.\nSee README to see how to download images and the detected joints.\n\"\"\"\n\nfrom os.path import join, exists, abspath, dirname\nfrom os import makedirs\nimport logging\nimport pickle\nfrom time import time\nfrom glob import glob\nimport argparse\n\nimport cv2\nimport numpy as np\nimport chumpy as ch\n\nfrom opendr.camera import ProjectPoints\n\nfrom .lib.robustifiers import GMOf\nfrom .lib.sphere_collisions import SphereCollisions\nfrom .lib.max_mixture_prior import MaxMixtureCompletePrior\n\nfrom .render_model import render_model\n\nfrom smpl_webuser.serialization import load_model\nfrom smpl_webuser.lbs import global_rigid_transformation\nfrom smpl_webuser.verts import verts_decorated\n\n_LOGGER = logging.getLogger(__name__)\n\n# Mapping from LSP joints to SMPL joints.\n# 0 Right ankle  8\n# 1 Right knee   5\n# 2 Right hip    2\n# 3 Left hip     1\n# 4 Left knee    4\n# 5 Left ankle   7\n# 6 Right wrist  21\n# 7 Right elbow  19\n# 8 Right shoulder 17\n# 9 Left shoulder  16\n# 10 Left elbow    18\n# 11 Left wrist    20\n# 12 Neck           -\n# 13 Head top       added\n\n\n# --------------------Camera estimation --------------------\ndef guess_init(model, focal_length, j2d, init_pose):\n    \"\"\"Initialize the camera translation via triangle similarity, by using the torso joints        .\n    :param model: SMPL model\n    :param focal_length: camera focal length (kept fixed)\n    :param j2d: 14x2 array of CNN joints\n    :param init_pose: 72D vector of pose parameters used for initialization (kept fixed)\n    :returns: 3D vector corresponding to the estimated camera translation\n    \"\"\"\n    cids = np.arange(0, 12)\n    # map from LSP to SMPL joints\n    j2d_here = j2d[cids]\n    smpl_ids = [8, 5, 2, 1, 4, 7, 21, 19, 17, 16, 18, 20]\n\n    opt_pose = ch.array(init_pose)\n    (_, A_global) = global_rigid_transformation(\n        opt_pose, model.J, model.kintree_table, xp=ch)\n    Jtr = ch.vstack([g[:3, 3] for g in A_global])\n    Jtr = Jtr[smpl_ids].r\n\n    # 9 is L shoulder, 3 is L hip\n    # 8 is R shoulder, 2 is R hip\n    diff3d = np.array([Jtr[9] - Jtr[3], Jtr[8] - Jtr[2]])\n    mean_height3d = np.mean(np.sqrt(np.sum(diff3d**2, axis=1)))\n\n    diff2d = np.array([j2d_here[9] - j2d_here[3], j2d_here[8] - j2d_here[2]])\n    mean_height2d = np.mean(np.sqrt(np.sum(diff2d**2, axis=1)))\n\n    est_d = focal_length * (mean_height3d / mean_height2d)\n    # just set the z value\n    init_t = np.array([0., 0., est_d])\n    return init_t\n\n\ndef initialize_camera(model,\n                      j2d,\n                      img,\n                      init_pose,\n                      flength=5000.,\n                      pix_thsh=25.,\n                      viz=False):\n    \"\"\"Initialize camera translation and body orientation\n    :param model: SMPL model\n    :param j2d: 14x2 array of CNN joints\n    :param img: h x w x 3 image \n    :param init_pose: 72D vector of pose parameters used for initialization\n    :param flength: camera focal length (kept fixed)\n    :param pix_thsh: threshold (in pixel), if the distance between shoulder joints in 2D\n                     is lower than pix_thsh, the body orientation as ambiguous (so a fit is run on both\n                     the estimated one and its flip)\n    :param viz: boolean, if True enables visualization during optimization\n    :returns: a tuple containing the estimated camera,\n              a boolean deciding if both the optimized body orientation and its flip should be considered,\n              3D vector for the body orientation\n    \"\"\"\n    # optimize camera translation and body orientation based on torso joints\n    # LSP torso ids:\n    # 2=right hip, 3=left hip, 8=right shoulder, 9=left shoulder\n    torso_cids = [2, 3, 8, 9]\n    # corresponding SMPL torso ids\n    torso_smpl_ids = [2, 1, 17, 16]\n\n    center = np.array([img.shape[1] / 2, img.shape[0] / 2])\n\n    # initialize camera rotation\n    rt = ch.zeros(3)\n    # initialize camera translation\n    _LOGGER.info('initializing translation via similar triangles')\n    init_t = guess_init(model, flength, j2d, init_pose)\n    t = ch.array(init_t)\n\n    # check how close the shoulder joints are\n    try_both_orient = np.linalg.norm(j2d[8] - j2d[9]) < pix_thsh\n\n    opt_pose = ch.array(init_pose)\n    (_, A_global) = global_rigid_transformation(\n        opt_pose, model.J, model.kintree_table, xp=ch)\n    Jtr = ch.vstack([g[:3, 3] for g in A_global])\n\n    # initialize the camera\n    cam = ProjectPoints(\n        f=np.array([flength, flength]), rt=rt, t=t, k=np.zeros(5), c=center)\n\n    # we are going to project the SMPL joints\n    cam.v = Jtr\n\n    if viz:\n        viz_img = img.copy()\n\n        # draw the target (CNN) joints\n        for coord in np.around(j2d).astype(int):\n            if (coord[0] < img.shape[1] and coord[0] >= 0 and\n                    coord[1] < img.shape[0] and coord[1] >= 0):\n                cv2.circle(viz_img, tuple(coord), 3, [0, 255, 0])\n\n        import matplotlib.pyplot as plt\n        plt.ion()\n\n        # draw optimized joints at each iteration\n        def on_step(_):\n            \"\"\"Draw a visualization.\"\"\"\n            plt.figure(1, figsize=(5, 5))\n            plt.subplot(1, 1, 1)\n            viz_img = img.copy()\n            for coord in np.around(cam.r[torso_smpl_ids]).astype(int):\n                if (coord[0] < viz_img.shape[1] and coord[0] >= 0 and\n                        coord[1] < viz_img.shape[0] and coord[1] >= 0):\n                    cv2.circle(viz_img, tuple(coord), 3, [0, 0, 255])\n            plt.imshow(viz_img[:, :, ::-1])\n            plt.draw()\n            plt.show()\n            plt.pause(1e-3)\n    else:\n        on_step = None\n    # optimize for camera translation and body orientation\n    free_variables = [cam.t, opt_pose[:3]]\n    ch.minimize(\n        # data term defined over torso joints...\n        {'cam': j2d[torso_cids] - cam[torso_smpl_ids],\n         # ...plus a regularizer for the camera translation\n         'cam_t': 1e2 * (cam.t[2] - init_t[2])},\n        x0=free_variables,\n        method='dogleg',\n        callback=on_step,\n        options={'maxiter': 100,\n                 'e_3': .0001,\n                 # disp set to 1 enables verbose output from the optimizer\n                 'disp': 0})\n    if viz:\n        plt.ioff()\n    \n    return (cam, try_both_orient, opt_pose[:3].r)\n\n\n# --------------------Core optimization --------------------\ndef optimize_on_joints(j2d,\n                       model,\n                       cam,\n                       img,\n                       prior,\n                       try_both_orient,\n                       body_orient,\n                       n_betas=10,\n                       regs=None,\n                       conf=None,\n                       viz=False):\n    \"\"\"Fit the model to the given set of joints, given the estimated camera\n    :param j2d: 14x2 array of CNN joints\n    :param model: SMPL model\n    :param cam: estimated camera\n    :param img: h x w x 3 image \n    :param prior: mixture of gaussians pose prior\n    :param try_both_orient: boolean, if True both body_orient and its flip are considered for the fit\n    :param body_orient: 3D vector, initialization for the body orientation\n    :param n_betas: number of shape coefficients considered during optimization\n    :param regs: regressors for capsules' axis and radius, if not None enables the interpenetration error term\n    :param conf: 14D vector storing the confidence values from the CNN\n    :param viz: boolean, if True enables visualization during optimization\n    :returns: a tuple containing the optimized model, its joints projected on image space, the camera translation\n    \"\"\"\n    t0 = time()\n    # define the mapping LSP joints -> SMPL joints\n    # cids are joints ids for LSP:\n    cids = range(12) + [13]\n    # joint ids for SMPL\n    # SMPL does not have a joint for head, instead we use a vertex for the head\n    # and append it later.\n    smpl_ids = [8, 5, 2, 1, 4, 7, 21, 19, 17, 16, 18, 20]\n\n    # the vertex id for the joint corresponding to the head\n    head_id = 411\n\n    # weights assigned to each joint during optimization;\n    # the definition of hips in SMPL and LSP is significantly different so set\n    # their weights to zero\n    base_weights = np.array(\n        [1, 1, 0, 0, 1, 1, 1, 1, 1, 1, 1, 1, 1], dtype=np.float64)\n\n    if try_both_orient:\n        flipped_orient = cv2.Rodrigues(body_orient)[0].dot(\n            cv2.Rodrigues(np.array([0., np.pi, 0]))[0])\n        flipped_orient = cv2.Rodrigues(flipped_orient)[0].ravel()\n        orientations = [body_orient, flipped_orient]\n    else:\n        orientations = [body_orient]\n\n    if try_both_orient:\n        # store here the final error for both orientations,\n        # and pick the orientation resulting in the lowest error\n        errors = []\n\n    svs = []\n    cams = []\n    for o_id, orient in enumerate(orientations):\n        # initialize the shape to the mean shape in the SMPL training set\n        betas = ch.zeros(n_betas)\n\n        # initialize the pose by using the optimized body orientation and the\n        # pose prior\n        init_pose = np.hstack((orient, prior.weights.dot(prior.means)))\n\n        # instantiate the model:\n        # verts_decorated allows us to define how many\n        # shape coefficients (directions) we want to consider (here, n_betas)\n        sv = verts_decorated(\n            trans=ch.zeros(3),\n            pose=ch.array(init_pose),\n            v_template=model.v_template,\n            J=model.J_regressor,\n            betas=betas,\n            shapedirs=model.shapedirs[:, :, :n_betas],\n            weights=model.weights,\n            kintree_table=model.kintree_table,\n            bs_style=model.bs_style,\n            f=model.f,\n            bs_type=model.bs_type,\n            posedirs=model.posedirs)\n\n        # make the SMPL joints depend on betas\n        Jdirs = np.dstack([model.J_regressor.dot(model.shapedirs[:, :, i])\n                           for i in range(len(betas))])\n        J_onbetas = ch.array(Jdirs).dot(betas) + model.J_regressor.dot(\n            model.v_template.r)\n\n        # get joint positions as a function of model pose, betas and trans\n        (_, A_global) = global_rigid_transformation(\n            sv.pose, J_onbetas, model.kintree_table, xp=ch)\n        Jtr = ch.vstack([g[:3, 3] for g in A_global]) + sv.trans\n\n        # add the head joint, corresponding to a vertex...\n        Jtr = ch.vstack((Jtr, sv[head_id]))\n\n        # ... and add the joint id to the list\n        if o_id == 0:\n            smpl_ids.append(len(Jtr) - 1)\n\n        # update the weights using confidence values\n        weights = base_weights * conf[\n            cids] if conf is not None else base_weights\n\n        # project SMPL joints on the image plane using the estimated camera\n        cam.v = Jtr\n\n        # data term: distance between observed and estimated joints in 2D\n        obj_j2d = lambda w, sigma: (\n            w * weights.reshape((-1, 1)) * GMOf((j2d[cids] - cam[smpl_ids]), sigma))\n\n        # mixture of gaussians pose prior\n        pprior = lambda w: w * prior(sv.pose)\n        # joint angles pose prior, defined over a subset of pose parameters:\n        # 55: left elbow,  90deg bend at -np.pi/2\n        # 58: right elbow, 90deg bend at np.pi/2\n        # 12: left knee,   90deg bend at np.pi/2\n        # 15: right knee,  90deg bend at np.pi/2\n        alpha = 10\n        my_exp = lambda x: alpha * ch.exp(x)\n        obj_angle = lambda w: w * ch.concatenate([my_exp(sv.pose[55]), my_exp(-sv.pose[\n                                                 58]), my_exp(-sv.pose[12]), my_exp(-sv.pose[15])])\n\n        if viz:\n            import matplotlib.pyplot as plt\n            plt.ion()\n\n            def on_step(_):\n                \"\"\"Create visualization.\"\"\"\n                plt.figure(1, figsize=(10, 10))\n                plt.subplot(1, 2, 1)\n                # show optimized joints in 2D\n                tmp_img = img.copy()\n                for coord, target_coord in zip(\n                        np.around(cam.r[smpl_ids]).astype(int),\n                        np.around(j2d[cids]).astype(int)):\n                    if (coord[0] < tmp_img.shape[1] and coord[0] >= 0 and\n                            coord[1] < tmp_img.shape[0] and coord[1] >= 0):\n                        cv2.circle(tmp_img, tuple(coord), 3, [0, 0, 255])\n                    if (target_coord[0] < tmp_img.shape[1] and\n                            target_coord[0] >= 0 and\n                            target_coord[1] < tmp_img.shape[0] and\n                            target_coord[1] >= 0):\n                        cv2.circle(tmp_img, tuple(target_coord), 3,\n                                   [0, 255, 0])\n                plt.imshow(tmp_img[:, :, ::-1])\n                plt.draw()\n                plt.show()\n                plt.pause(1e-2)\n\n            on_step(_)\n        else:\n            on_step = None\n\n        if regs is not None:\n            # interpenetration term\n            sp = SphereCollisions(\n                pose=sv.pose, betas=sv.betas, model=model, regs=regs)\n            sp.no_hands = True\n        # weight configuration used in the paper, with joints + confidence values from the CNN\n        # (all the weights used in the code were obtained via grid search, see the paper for more details)\n        # the first list contains the weights for the pose priors,\n        # the second list contains the weights for the shape prior\n        opt_weights = zip([4.04 * 1e2, 4.04 * 1e2, 57.4, 4.78],\n                          [1e2, 5 * 1e1, 1e1, .5 * 1e1])\n\n        # run the optimization in 4 stages, progressively decreasing the\n        # weights for the priors\n        for stage, (w, wbetas) in enumerate(opt_weights):\n            _LOGGER.info('stage %01d', stage)\n            objs = {}\n\n            objs['j2d'] = obj_j2d(1., 100)\n\n            objs['pose'] = pprior(w)\n\n            objs['pose_exp'] = obj_angle(0.317 * w)\n\n            objs['betas'] = wbetas * betas\n\n            if regs is not None:\n                objs['sph_coll'] = 1e3 * sp\n\n            ch.minimize(\n                objs,\n                x0=[sv.betas, sv.pose],\n                method='dogleg',\n                callback=on_step,\n                options={'maxiter': 100,\n                         'e_3': .0001,\n                         'disp': 0})\n\n        t1 = time()\n        _LOGGER.info('elapsed %.05f', (t1 - t0))\n        if try_both_orient:\n            errors.append((objs['j2d'].r**2).sum())\n        svs.append(sv)\n        cams.append(cam)\n\n    if try_both_orient and errors[0] > errors[1]:\n        choose_id = 1\n    else:\n        choose_id = 0\n    if viz:\n        plt.ioff()\n    return (svs[choose_id], cams[choose_id].r, cams[choose_id].t.r, Jtr)\n\n\ndef run_single_fit(img,\n                   j2d,\n                   conf,\n                   model,\n                   regs=None,\n                   n_betas=10,\n                   flength=5000.,\n                   pix_thsh=25.,\n                   scale_factor=1,\n                   viz=False,\n                   do_degrees=None):\n    \"\"\"Run the fit for one specific image.\n    :param img: h x w x 3 image \n    :param j2d: 14x2 array of CNN joints\n    :param conf: 14D vector storing the confidence values from the CNN\n    :param model: SMPL model\n    :param regs: regressors for capsules' axis and radius, if not None enables the interpenetration error term\n    :param n_betas: number of shape coefficients considered during optimization\n    :param flength: camera focal length (kept fixed during optimization)\n    :param pix_thsh: threshold (in pixel), if the distance between shoulder joints in 2D\n                     is lower than pix_thsh, the body orientation as ambiguous (so a fit is run on both\n                     the estimated one and its flip)\n    :param scale_factor: int, rescale the image (for LSP, slightly greater images -- 2x -- help obtain better fits)\n    :param viz: boolean, if True enables visualization during optimization\n    :param do_degrees: list of degrees in azimuth to render the final fit when saving results\n    :returns: a tuple containing camera/model parameters and images with rendered fits\n    \"\"\"\n    if do_degrees is None:\n        do_degrees = []\n\n    # create the pose prior (GMM over CMU)\n    prior = MaxMixtureCompletePrior(n_gaussians=8).get_gmm_prior()\n    # get the mean pose as our initial pose\n    init_pose = np.hstack((np.zeros(3), prior.weights.dot(prior.means)))\n\n    if scale_factor != 1:\n        img = cv2.resize(img, (img.shape[1] * scale_factor,\n                               img.shape[0] * scale_factor))\n        j2d[:, 0] *= scale_factor\n        j2d[:, 1] *= scale_factor\n\n    # estimate the camera parameters\n    (cam, try_both_orient, body_orient) = initialize_camera(\n        model,\n        j2d,\n        img,\n        init_pose,\n        flength=flength,\n        pix_thsh=pix_thsh,\n        viz=viz)\n\n    # fit\n    (sv, opt_j2d, t, v) = optimize_on_joints(\n        j2d,\n        model,\n        cam,\n        img,\n        prior,\n        try_both_orient,\n        body_orient,\n        n_betas=n_betas,\n        conf=conf,\n        viz=viz,\n        regs=regs, )\n\n    h = img.shape[0]\n    w = img.shape[1]\n    dist = np.abs(cam.t.r[2] - np.mean(sv.r, axis=0)[2])\n\n    images = []\n    orig_v = sv.r\n\n    for deg in do_degrees:\n        if deg != 0:\n            aroundy = cv2.Rodrigues(np.array([0, np.radians(deg), 0]))[0]\n            center = orig_v.mean(axis=0)\n            new_v = np.dot((orig_v - center), aroundy)\n            verts = new_v + center\n        else:\n            verts = orig_v\n        \n        # now render\n        im = (render_model(\n            verts, model.f, w, h, cam, far=20 + dist) * 255.).astype('uint8')\n        images.append(im)\n\n    # return fit parameters\n    # .r converts a chumpy array into numpy array\n    params = {'cam_t': cam.t.r,\n              'f': cam.f.r,\n              'v': v.r,\n              'pose': sv.pose.r,\n              'betas': sv.betas.r}\n\n    return params, images\n\n\ndef main(base_dir,\n         out_dir,\n         use_interpenetration=True,\n         n_betas=10,\n         flength=5000.,\n         pix_thsh=25.,\n         use_neutral=False,\n         viz=True):\n    \"\"\"Set up paths to image and joint data, saves results.\n    :param base_dir: folder containing LSP images and data\n    :param out_dir: output folder\n    :param use_interpenetration: boolean, if True enables the interpenetration term\n    :param n_betas: number of shape coefficients considered during optimization\n    :param flength: camera focal length (an estimate)\n    :param pix_thsh: threshold (in pixel), if the distance between shoulder joints in 2D\n                     is lower than pix_thsh, the body orientation as ambiguous (so a fit is run on both\n                     the estimated one and its flip)\n    :param use_neutral: boolean, if True enables uses the neutral gender SMPL model\n    :param viz: boolean, if True enables visualization during optimization\n    \"\"\"\n\n    img_dir = join(abspath(base_dir), 'images/lsp')\n    data_dir = join(abspath(base_dir), 'results/lsp')\n\n    if not exists(out_dir):\n        makedirs(out_dir)\n\n    # Render degrees: List of degrees in azimuth to render the final fit.\n    # Note that rendering many views can take a while.\n    do_degrees = [0.]\n\n    sph_regs = None\n    if not use_neutral:\n        _LOGGER.info(\"Reading genders...\")\n        # File storing information about gender in LSP\n        with open(join(data_dir, 'lsp_gender.csv')) as f:\n            genders = f.readlines()\n        model_female = load_model(MODEL_FEMALE_PATH)\n        model_male = load_model(MODEL_MALE_PATH)\n        if use_interpenetration:\n            sph_regs_male = np.load(SPH_REGS_MALE_PATH)\n            sph_regs_female = np.load(SPH_REGS_FEMALE_PATH)\n    else:\n        gender = 'neutral'\n        model = load_model(MODEL_NEUTRAL_PATH)\n        if use_interpenetration:\n            sph_regs = np.load(SPH_REGS_NEUTRAL_PATH)\n\n    # Load joints\n    est = np.load(join(data_dir, 'est_joints.npz'))['est_joints']\n\n    # Load images\n    img_paths = sorted(glob(join(img_dir, '*[0-9].jpg')))\n    for ind, img_path in enumerate(img_paths):\n        out_path = '%s/%04d.pkl' % (out_dir, ind)\n        if not exists(out_path):\n            _LOGGER.info('Fitting 3D body on `%s` (saving to `%s`).', img_path,\n                         out_path)\n            img = cv2.imread(img_path)\n            if img.ndim == 2:\n                _LOGGER.warn(\"The image is grayscale!\")\n                img = np.dstack((img, img, img))\n\n            # [x-y, keypoints, idx]\n            joints = est[:2, :, ind].T\n            conf = est[2, :, ind]\n\n            if not use_neutral:\n                gender = 'male' if int(genders[ind]) == 0 else 'female'\n                if gender == 'female':\n                    model = model_female\n                    if use_interpenetration:\n                        sph_regs = sph_regs_female\n                elif gender == 'male':\n                    model = model_male\n                    if use_interpenetration:\n                        sph_regs = sph_regs_male\n\n            params, vis = run_single_fit(\n                img,\n                joints,\n                conf,\n                model,\n                regs=sph_regs,\n                n_betas=n_betas,\n                flength=flength,\n                pix_thsh=pix_thsh,\n                scale_factor=2,\n                viz=viz,\n                do_degrees=do_degrees)\n            if viz:\n                import matplotlib.pyplot as plt\n                plt.ion()\n                plt.show()\n                plt.subplot(121)\n                plt.imshow(img[:, :, ::-1])\n                if do_degrees is not None:\n                    print(do_degrees)\n                    for di, deg in enumerate(do_degrees):\n                        plt.subplot(122)\n                        plt.cla()\n                        plt.imshow(vis[di])\n                        plt.draw()\n                        plt.title('%d deg' % deg)\n                        plt.pause(1)\n                raw_input('Press any key to continue...')\n\n            with open(out_path, 'w') as outf:\n                pickle.dump(params, outf)\n\n            print(do_degrees)\n\n            # This only saves the first rendering.\n            if do_degrees is not None:\n                cv2.imwrite(out_path.replace('.pkl', '.png'), vis[0])\n\n\nif __name__ == '__main__':\n    logging.basicConfig(level=logging.INFO)\n\n    parser = argparse.ArgumentParser(description='run SMPLify on LSP dataset')\n    parser.add_argument(\n        'base_dir',\n        default='/scratch1/projects/smplify_public/',\n        nargs='?',\n        help=\"Directory that contains images/lsp and results/lps , i.e.\"\n        \"the directory you untared smplify_code.tar.gz\")\n    parser.add_argument(\n        '--out_dir',\n        default='/home/kp-kepra/Github/smplify_public',\n        type=str,\n        help='Where results will be saved, default is /tmp/smplify_lsp')\n    parser.add_argument(\n        '--no_interpenetration',\n        default=False,\n        action='store_true',\n        help=\"Using this flag removes the interpenetration term, which speeds\"\n        \"up optimization at the expense of possible interpenetration.\")\n    parser.add_argument(\n        '--gender_neutral',\n        default=False,\n        action='store_true',\n        help=\"Using this flag always uses the neutral SMPL model, otherwise \"\n        \"gender specified SMPL models are used.\")\n    parser.add_argument(\n        '--n_betas',\n        default=10,\n        type=int,\n        help=\"Specify the number of shape coefficients to use.\")\n    parser.add_argument(\n        '--flength',\n        default=5000,\n        type=float,\n        help=\"Specify value of focal length.\")\n    parser.add_argument(\n        '--side_view_thsh',\n        default=25,\n        type=float,\n        help=\"This is thresholding value that determines whether the human is captured in a side view. If the pixel distance between the shoulders is less than this value, two initializations of SMPL fits are tried.\")\n    parser.add_argument(\n        '--viz',\n        default=False,\n        action='store_true',\n        help=\"Turns on visualization of intermediate optimization steps \"\n        \"and final results.\")\n    args = parser.parse_args()\n\n    use_interpenetration = not args.no_interpenetration\n    if not use_interpenetration:\n        _LOGGER.info('Not using interpenetration term.')\n    if args.gender_neutral:\n        _LOGGER.info('Using gender neutral model.')\n\n    # Set up paths & load models.\n    # Assumes 'models' in the 'code/' directory where this file is in.\n    MODEL_DIR = join(abspath(dirname(__file__)), 'models')\n    # Model paths:\n    MODEL_NEUTRAL_PATH = join(\n        MODEL_DIR, 'basicModel_neutral_lbs_10_207_0_v1.0.0.pkl')\n    MODEL_FEMALE_PATH = join(\n        MODEL_DIR, 'basicModel_f_lbs_10_207_0_v1.0.0.pkl')\n    MODEL_MALE_PATH = join(MODEL_DIR,\n                           'basicmodel_m_lbs_10_207_0_v1.0.0.pkl')\n\n    if use_interpenetration:\n        # paths to the npz files storing the regressors for capsules\n        SPH_REGS_NEUTRAL_PATH = join(MODEL_DIR,\n                                     'regressors_locked_normalized_hybrid.npz')\n        SPH_REGS_FEMALE_PATH = join(MODEL_DIR,\n                                    'regressors_locked_normalized_female.npz')\n        SPH_REGS_MALE_PATH = join(MODEL_DIR,\n                                  'regressors_locked_normalized_male.npz')\n\n    main(args.base_dir, args.out_dir, use_interpenetration, args.n_betas,\n         args.flength, args.side_view_thsh, args.gender_neutral, args.viz)\n"
  },
  {
    "path": "libs/annotator/smplify/lib/__init__.py",
    "content": ""
  },
  {
    "path": "libs/annotator/smplify/lib/capsule_body.py",
    "content": "\"\"\"\nCopyright 2016 Max Planck Society, Federica Bogo, Angjoo Kanazawa. All rights reserved.\nThis software is provided for research purposes only.\nBy using this software you agree to the terms of the SMPLify license here:\n     http://smplify.is.tue.mpg.de/license\n\nThis script implements an approximation of the body by means of capsules (20 in total).\nCapsules can be further simplified into spheres (with centers along the capsule axis and\nradius corresponding to the capsule radius) to efficiently compute an interpenetration error term\n(as in sphere_collisions.py).\n\"\"\"\n\nimport numpy as np\nimport chumpy as ch\nimport scipy.sparse as sp\n\nfrom .capsule_ch import Capsule\n\njoint2name = ['pelvis', 'leftThigh', 'rightThigh', 'spine', 'leftCalf',\n              'rightCalf', 'spine1', 'leftFoot', 'rightFoot', 'spine2', 'neck',\n              'leftShoulder', 'rightShoulder', 'head', 'leftUpperArm',\n              'rightUpperArm', 'leftForeArm', 'rightForeArm', 'leftHand',\n              'rightHand']\n\n# the orientation of each capsule\nrots0 = ch.asarray(\n    [[0, 0, np.pi / 2], [0, 0, np.pi], [0, 0, np.pi], [0, 0, np.pi / 2],\n     [0, 0, np.pi], [0, 0, np.pi], [0, 0, np.pi / 2], [np.pi / 2, 0, 0],\n     [np.pi / 2, 0, 0], [0, 0, np.pi / 2], [0, 0, 0], [0, 0, -np.pi / 2],\n     [0, 0, np.pi / 2], [0, 0, 0], [0, 0, -np.pi / 2], [0, 0, np.pi / 2],\n     [0, 0, -np.pi / 2], [0, 0, np.pi / 2], [0, 0, -np.pi / 2],\n     [0, 0, np.pi / 2]])\n\n# groups hands and fingers, feet and toes\n# each comment line provides the body part corresonding to the capsule\n# and the corresponding id\nmujoco2segm = [[0],  # hip 0\n               [1],  # leftThigh 1\n               [2],  # rightThigh 2\n               [3],  # spine 3\n               [4],  # leftCalf 4\n               [5],  # rightCalf 5\n               [6],  # spine1 6\n               [7, 10],  # leftFoot + leftToes 7\n               [8, 11],  # rightFoot + rightToes 8\n               [9],  # spine2 9\n               [12],  # neck 10\n               [13],  # leftShoulder 11\n               [14],  # rightShoulder 12\n               [15],  # head 13\n               [16],  # leftUpperArm 14\n               [17],  # rightUpperArm 15\n               [18],  # leftForeArm 16\n               [19],  # rightForeArm 17\n               [20, 22],  # leftHand + leftFingers 18\n               [21, 23]]  # rightHand + rightFingers 19\n\n# sets pairs of ids, corresponding to capsules that should not\n# penetrate each other\ncollisions = [\n    [0, 16],  # hip and leftForeArm\n    [0, 17],  # hip and rightForeArm\n    [0, 18],  # hip and leftHand\n    [0, 19],  # hip and rightHand\n    [3, 16],  # spine and leftForeArm\n    [3, 17],  # spine and rightForeArm\n    [3, 18],  # spine and leftHand\n    [3, 19],  # spine and rightHand\n    [4, 5],  # leftCalf and rightCalf\n    [6, 16],  # spine1 and leftForeArm\n    [6, 17],  # spine1 and rightForeArm\n    [6, 18],  # spine1 and leftHand\n    [6, 19],  # spine1 and rightHand\n    [7, 5],  # leftFoot and rightCalf\n    [8, 7],  # rightFoot and leftFoot\n    [8, 4],  # rightFoot and leftCalf\n    [9, 16],  # spine2 and leftForeArm\n    [9, 17],  # spine2 and rightForeArm\n    [9, 18],  # spine2 and leftHand\n    [9, 19],  # spine2 and rightHand\n    [11, 16],  # leftShoulder and leftForeArm\n    [12, 17],  # rightShoulder and rightForeArm\n    [18, 19],  # leftHand and rightHand\n]\n\n\ndef get_capsules(model, wrt_betas=None, length_regs=None, rad_regs=None):\n    from opendr.geometry import Rodrigues\n    if length_regs is not None:\n        n_shape_dofs = length_regs.shape[0] - 1\n    else:\n        n_shape_dofs = model.betas.r.size\n    segm = np.argmax(model.weights_prior, axis=1)\n    J_off = ch.zeros((len(joint2name), 3))\n    rots = rots0.copy()\n    mujoco_t_mid = [0, 3, 6, 9]\n    if wrt_betas is not None:\n        # if we want to differentiate wrt betas (shape), we must have the\n        # regressors...\n        assert (length_regs is not None and rad_regs is not None)\n        # ... and betas must be a chumpy object\n        assert (hasattr(wrt_betas, 'dterms'))\n        pad = ch.concatenate(\n            (wrt_betas, ch.zeros(n_shape_dofs - len(wrt_betas)), ch.ones(1)))\n        lengths = pad.dot(length_regs)\n        rads = pad.dot(rad_regs)\n    else:\n        lengths = ch.ones(len(joint2name))\n        rads = ch.ones(len(joint2name))\n    betas = wrt_betas if wrt_betas is not None else model.betas\n    n_betas = len(betas)\n    # the joint regressors are the original, pre-optimized ones\n    # (middle of the part frontier)\n    myJ_regressor = model.J_regressor_prior\n    myJ0 = ch.vstack(\n        (ch.ch.MatVecMult(myJ_regressor, model.v_template[:, 0] +\n                          model.shapedirs[:, :, :n_betas].dot(betas)[:, 0]),\n         ch.ch.MatVecMult(myJ_regressor, model.v_template[:, 1] +\n                          model.shapedirs[:, :, :n_betas].dot(betas)[:, 1]),\n         ch.ch.MatVecMult(myJ_regressor, model.v_template[:, 2] +\n                          model.shapedirs[:, :, :n_betas].dot(betas)[:, 2]))).T\n    # with small adjustments for hips, spine and feet\n    myJ = ch.vstack(\n        [ch.concatenate([myJ0[0, 0], (\n            .6 * myJ0[0, 1] + .2 * myJ0[1, 1] + .2 * myJ0[2, 1]), myJ0[9, 2]]),\n         ch.vstack([myJ0[i] for i in range(1, 7)]), ch.concatenate(\n             [myJ0[7, 0], (1.1 * myJ0[7, 1] - .1 * myJ0[4, 1]), myJ0[7, 2]]),\n         ch.concatenate(\n             [myJ0[8, 0], (1.1 * myJ0[8, 1] - .1 * myJ0[5, 1]), myJ0[8, 2]]),\n         ch.concatenate(\n             [myJ0[9, 0], myJ0[9, 1], (.2 * myJ0[9, 2] + .8 * myJ0[12, 2])]),\n         ch.vstack([myJ0[i] for i in range(10, 24)])])\n    capsules = []\n    # create one capsule per mujoco joint\n    for ijoint, segms in enumerate(mujoco2segm):\n        if wrt_betas is None:\n            vidxs = np.asarray([segm == k for k in segms]).any(axis=0)\n            verts = model.v_template[vidxs].r\n            dims = (verts.max(axis=0) - verts.min(axis=0))\n            rads[ijoint] = .5 * ((dims[(np.argmax(dims) + 1) % 3] + dims[(\n                np.argmax(dims) + 2) % 3]) / 4.)\n            lengths[ijoint] = max(dims) - 2. * rads[ijoint].r\n        # the core joints are different, since the capsule is not in the joint\n        # but in the middle\n        if ijoint in mujoco_t_mid:\n            len_offset = ch.vstack([ch.zeros(1), ch.abs(lengths[ijoint]) / 2.,\n                                    ch.zeros(1)]).reshape(3, 1)\n            caps = Capsule(\n                (J_off[ijoint] + myJ[mujoco2segm[ijoint][0]]).reshape(\n                    3, 1) - Rodrigues(rots[ijoint]).dot(len_offset),\n                rots[ijoint], rads[ijoint], lengths[ijoint])\n        else:\n            caps = Capsule(\n                (J_off[ijoint] + myJ[mujoco2segm[ijoint][0]]).reshape(3, 1),\n                rots[ijoint], rads[ijoint], lengths[ijoint])\n        caps.id = ijoint\n        capsules.append(caps)\n    return capsules\n\n\ndef set_sphere_centers(capsule, floor=True):\n    if floor:\n        n_spheres = int(np.floor(capsule.length.r / (2 * capsule.rad.r) - 1))\n    else:\n        n_spheres = int(np.ceil(capsule.length.r / (2 * capsule.rad.r) - 1))\n\n    # remove \"redundant\" spheres for right and left thigh...\n    if capsule.id == 1 or capsule.id == 2:\n        centers = [capsule.axis[1].r]\n    # ... and right and left upper arm\n    elif capsule.id == 14 or capsule.id == 15:\n        if n_spheres >= 1:\n            centers = []\n        else:\n            centers = [capsule.axis[1].r]\n    else:\n        centers = [capsule.axis[0].r, capsule.axis[1].r]\n\n    if n_spheres >= 1:\n        step = capsule.length.r / (n_spheres + 1)\n        for i in xrange(n_spheres):\n            centers.append(capsule.axis[0].r + (capsule.axis[\n                1].r - capsule.axis[0].r) * step * (i + 1) / capsule.length.r)\n\n    capsule.centers = centers\n    return capsule.centers\n\n\ndef capsule_dist(capsule0, capsule1, alpha=.3, increase_hand=True):\n    range0 = range(capsule0.center_id,\n                   capsule0.center_id + len(capsule0.centers))\n    range1 = range(capsule1.center_id,\n                   capsule1.center_id + len(capsule1.centers))\n    cnt0 = ch.concatenate([[cid] * len(range1) for cid in range0])\n    cnt1 = ch.concatenate([range1] * len(range0))\n    if increase_hand:\n        if (capsule0.id == 18) or (capsule0.id == 19) or (\n                capsule1.id == 18) or (capsule1.id == 19):\n            dst = (alpha * 1.2 * capsule0.rad.r)**2 + (alpha * 1.2 *\n                                                       capsule1.rad.r)**2\n        else:\n            dst = (alpha * capsule0.rad.r)**2 + (alpha * capsule1.rad.r)**2\n    else:\n        dst = (alpha * capsule0.rad.r)**2 + (alpha * capsule1.rad.r)**2\n    radiuss = np.hstack([dst] * len(cnt0)).squeeze()\n    return (cnt0, cnt1, radiuss)\n\n\ndef get_capsule_bweights(vs):\n    # \"blend\" weights for the capsule. They are binary\n    rows = np.arange(vs.shape[0])\n    cols = np.tile(np.hstack((range(10), range(12, 22))), (52, 1)).T.ravel()\n    data = np.ones(vs.shape[0])\n    caps_weights = np.asarray(\n        sp.csc_matrix(\n            (data, (rows, cols)), shape=(vs.shape[0], 24)).todense())\n    return caps_weights\n\n\ndef get_sphere_bweights(sph_vs, capsules):\n    rows = np.arange(sph_vs.shape[0])\n    cols = []\n    for cps, w in zip(capsules, range(10) + range(12, 22)):\n        cols.append([w] * len(cps.centers))\n    cols = np.hstack(cols)\n    data = np.ones(sph_vs.shape[0])\n    sph_weights = np.asarray(\n        sp.csc_matrix(\n            (data, (rows, cols)), shape=(sph_vs.shape[0], 24)).todense())\n    return sph_weights\n"
  },
  {
    "path": "libs/annotator/smplify/lib/capsule_ch.py",
    "content": "\"\"\"\nCopyright 2016 Max Planck Society, Federica Bogo, Angjoo Kanazawa. All rights reserved.\nThis software is provided for research purposes only.\nBy using this software you agree to the terms of the SMPLify license here:\n     http://smplify.is.tue.mpg.de/license\n\nThis script implements a Capsule object, used in the body approximation implemented\nin capsule_body.py. Capsule sizes depend on body shape (and are differentiable with respect to it).\nCapsules are the basis to compute an approximation based on spheres, used to compute efficiently\nthe interpenetration error term in sphere_collisions.py.\n\"\"\"\n\nimport numpy as np\nimport chumpy as ch\nfrom opendr.geometry import Rodrigues\n\n# faces for the capsules. Useful only for visualization purposes\ncap_f = np.asarray(\n    [[0, 7, 6], [1, 7, 9], [0, 6, 11], [0, 11, 13], [0, 13, 10], [1, 9, 16],\n     [2, 8, 18], [3, 12, 20], [4, 14, 22], [5, 15, 24], [1, 16, 19],\n     [2, 18, 21], [3, 20, 23], [4, 22, 25], [5, 24, 17], [16, 17, 26],\n     [22, 23, 32], [48, 18, 28], [49, 20, 30], [24, 25, 34], [25, 22, 50],\n     [28, 19, 47], [30, 21, 48], [32, 23, 49], [17, 24, 51], [26, 17, 51],\n     [34, 25, 50], [23, 20, 49], [21, 18, 48], [19, 16, 47], [51, 24, 34],\n     [24, 15, 25], [15, 4, 25], [50, 22, 32], [22, 14, 23], [14, 3, 23],\n     [20, 21, 30], [20, 12, 21], [12, 2, 21], [18, 19, 28], [18, 8, 19],\n     [8, 1, 19], [47, 16, 26], [16, 9, 17], [9, 5, 17], [10, 15, 5],\n     [10, 13, 15], [13, 4, 15], [13, 14, 4], [13, 11, 14], [11, 3, 14],\n     [11, 12, 3], [11, 6, 12], [6, 2, 12], [9, 10, 5], [9, 7, 10], [7, 0, 10],\n     [6, 8, 2], [6, 7, 8], [7, 1, 8], [29, 36, 41], [31, 37, 44], [33, 38, 45],\n     [35, 39, 46], [27, 40, 42], [42, 46, 43], [42, 40, 46], [40, 35, 46],\n     [46, 45, 43], [46, 39, 45], [39, 33, 45], [45, 44, 43], [45, 38, 44],\n     [38, 31, 44], [44, 41, 43], [44, 37, 41], [37, 29, 41], [41, 42, 43],\n     [41, 36, 42], [36, 27, 42], [26, 40, 27], [26, 51, 40], [51, 35, 40],\n     [34, 39, 35], [34, 50, 39], [50, 33, 39], [32, 38, 33], [32, 49, 38],\n     [49, 31, 38], [30, 37, 31], [30, 48, 37], [48, 29, 37], [28, 36, 29],\n     [28, 47, 36], [47, 27, 36], [51, 34, 35], [50, 32, 33], [49, 30, 31],\n     [48, 28, 29], [47, 26, 27]])\n\nelev = np.asarray(\n    [0., 0.5535673, 1.01721871, 0., -1.01721871, -0.5535673, 0.52359324,\n     0.31415301, 0.94246863, 0., -0.31415301, 0., 0.52359547, -0.52359324,\n     -0.52359547, -0.94246863, 0.31415501, -0.31415501, 1.57079633, 0.94247719,\n     0.31415501, 0.94247719, -0.94247719, -0.31415501, -0.94247719,\n     -1.57079633, -0.31415624, 0., 0.94248124, 1.01722122, 0.94247396,\n     0.55356579, -0.31415377, -0.55356579, -1.57079233, -1.01722122,\n     0.52359706, 0.94246791, 0., -0.94246791, -0.52359706, 0.52359371, 0., 0.,\n     0.31415246, -0.31415246, -0.52359371, 0.31415624, 1.57079233, 0.31415377,\n     -0.94247396, -0.94248124])\n\naz = np.asarray(\n    [-1.57079633, -0.55358064, -2.12435586, -2.67794236, -2.12435586,\n     -0.55358064, -1.7595018, -1.10715248, -1.10714872, -0.55357999,\n     -1.10715248, -2.12436911, -2.48922865, -1.7595018, -2.48922865,\n     -1.10714872, 0., 0., 0., 0., 3.14159265, 3.14159265, 3.14159265,\n     3.14159265, 0., 0., 0., 0.46365119, 0., 1.01724226, 3.14159265,\n     2.58801549, 3.14159265, 2.58801549, 3.14159265, 1.01724226, 0.6523668,\n     2.03445078, 2.58801476, 2.03445078, 0.6523668, 1.38209652, 1.01722642,\n     1.57080033, 2.03444394, 2.03444394, 1.38209652, 0., 3.14159265,\n     3.14159265, 3.14159265, 0.])\n\n# vertices for the capsules\nv = np.vstack(\n    [np.cos(az) * np.cos(elev), np.sin(az) * np.cos(elev), np.sin(elev)]).T\n\n\nclass Capsule(object):\n\n    def __init__(self, t, rod, rad, length):\n        assert (hasattr(t, 'dterms'))\n        # the translation should be a chumpy object (differentiable wrt shape)\n        self.t = t  # translation of the axis\n        self.rod = rod  # rotation of the axis in Rodrigues form\n        # the radius should be a chumpy object (differentiable wrt shape)\n        assert (hasattr(rad, 'dterms'))\n        self.rad = rad  # radius of the capsule\n        # the length should be a chumpy object (differentiable wrt shape)\n        assert (hasattr(length, 'dterms'))\n        self.length = length  # length of the axis\n        axis0 = ch.vstack([0, ch.abs(self.length), 0])\n        self.axis = ch.vstack((t.T, (t + Rodrigues(rod).dot(axis0)).T))\n        v0 = ch.hstack([v[:26].T * rad, (v[26:].T * rad) + axis0])\n        self.v = ((t + Rodrigues(rod).dot(v0)).T)\n        self.set_sphere_centers()\n\n    def set_sphere_centers(self, floor=False):\n        # sphere centers are evenly spaced along the capsule axis length\n        if floor:\n            n_spheres = int(np.floor(self.length / (2 * self.rad) - 1))\n        else:\n            n_spheres = int(np.ceil(self.length / (2 * self.rad) - 1))\n\n        centers = [self.axis[0].r, self.axis[1].r]\n\n        if n_spheres >= 1:\n            step = self.length.r / (n_spheres + 1)\n            for i in xrange(n_spheres):\n                centers.append(self.axis[0].r + (self.axis[1].r - self.axis[\n                    0].r) * step * (i + 1) / self.length.r)\n\n        self.centers = centers\n"
  },
  {
    "path": "libs/annotator/smplify/lib/max_mixture_prior.py",
    "content": "\"\"\"\nCopyright 2016 Max Planck Society, Federica Bogo, Angjoo Kanazawa. All rights reserved.\nThis software is provided for research purposes only.\nBy using this software you agree to the terms of the SMPLify license here:\n     http://smplify.is.tue.mpg.de/license\n\nThis script implements the pose prior based on a mixture of Gaussians.\nTo simplify the log-likelihood computation, the sum in the mixture of Gaussians\nis approximated by a max operator (see the paper for more details).\n\"\"\"\n\nimport os\nimport numpy as np\nimport chumpy as ch\n\n\nclass MaxMixtureComplete(ch.Ch):\n    \"\"\"Define the MaxMixture class.\"\"\"\n    # x is the input vector we want to evaluate the prior on;\n    # means, precs and weights are the parameters of the mixture\n    dterms = 'x'\n    terms = 'means', 'precs', 'weights'\n\n    def on_changed(self, which):\n        # on_changed is called before any call to r or dr_wrt,\n        # therefore it can be used also for initialization\n        # setup means, precs and loglikelihood expressions\n        if 'means' in which or 'precs' in which or 'weights' in which:\n            # This is just the mahalanobis part.\n            self.loglikelihoods = [np.sqrt(0.5) * (self.x - m).dot(s)\n                                   for m, s in zip(self.means, self.precs)]\n\n        if 'x' in which:\n            self.min_component_idx = np.argmin(\n                [(logl**2).sum().r[0] - np.log(w[0])\n                 for logl, w in zip(self.loglikelihoods, self.weights)])\n\n    def compute_r(self):\n        min_w = self.weights[self.min_component_idx]\n        # Add the sqrt(-log(weights))\n        return ch.concatenate((self.loglikelihoods[self.min_component_idx].r,\n                               np.sqrt(-np.log(min_w))))\n\n    def compute_dr_wrt(self, wrt):\n        # the call to dr_wrt returns a jacobian 69 x 72,\n        # when wrt has 72 elements (pose vector)\n        # here we intercept the call and return a 70 x 72 matrix,\n        # with an additional row of zeroes (these are the jacobian\n        # entries corresponding to sqrt(-log(weights))\n        import scipy.sparse as sp\n\n        dr = self.loglikelihoods[self.min_component_idx].dr_wrt(wrt)\n        if dr is not None:\n            # extract rows, cols and data, and return a new matrix with\n            # the same values but 1 additional row\n            Is, Js, Vs = sp.find(dr)\n            dr = sp.csc_matrix(\n                (Vs, (Is, Js)), shape=(dr.shape[0] + 1, dr.shape[1]))\n\n        return dr\n\n\nclass MaxMixtureCompleteWrapper(object):\n    \"\"\"Convenience wrapper to match interface spec.\"\"\"\n\n    def __init__(self, means, precs, weights, prefix):\n        self.means = means\n        self.precs = precs  # Already \"sqrt\"ed\n        self.weights = weights\n        self.prefix = prefix\n\n    def __call__(self, x):\n        # wrapping since __call__ couldn't be defined directly for a chumpy\n        # object\n        return (MaxMixtureComplete(\n            x=x[self.prefix:],\n            means=self.means,\n            precs=self.precs,\n            weights=self.weights))\n\n\nclass MaxMixtureCompletePrior(object):\n    \"\"\"Prior density estimation.\"\"\"\n\n    def __init__(self, n_gaussians=8, prefix=3):\n        self.n_gaussians = n_gaussians\n        self.prefix = prefix\n        self.prior = self.create_prior_from_cmu()\n\n    def create_prior_from_cmu(self):\n        \"\"\"Load the gmm from the CMU motion database.\"\"\"\n        from os.path import dirname\n        import cPickle as pickle\n        with open(\n                os.path.join(\n                    dirname(dirname(__file__)), 'models', 'gmm_%02d.pkl' %\n                    self.n_gaussians)) as f:\n            gmm = pickle.load(f)\n\n        precs = ch.asarray([np.linalg.inv(cov) for cov in gmm['covars']])\n        chols = ch.asarray([np.linalg.cholesky(prec) for prec in precs])\n\n        # The constant term:\n        sqrdets = np.array([(np.sqrt(np.linalg.det(c)))\n                            for c in gmm['covars']])\n        const = (2 * np.pi)**(69 / 2.)\n\n        self.weights = ch.asarray(gmm['weights'] / (const *\n                                                    (sqrdets / sqrdets.min())))\n\n        return (MaxMixtureCompleteWrapper(\n            means=gmm['means'],\n            precs=chols,\n            weights=self.weights,\n            prefix=self.prefix))\n\n    def get_gmm_prior(self):\n        \"\"\"Getter implementation.\"\"\"\n        return self.prior\n"
  },
  {
    "path": "libs/annotator/smplify/lib/robustifiers.py",
    "content": "\"\"\"\nCopyright 2016 Max Planck Society, Matthew Loper. All rights reserved.\nThis software is provided for research purposes only.\nBy using this software you agree to the terms of the SMPLify license here:\n     http://smplify.is.tue.mpg.de/license\n\nThis script implements the Geman-McClure robustifier as chumpy object.\n\"\"\"\n\n#!/usr/bin/env python\nimport numpy as np\nimport scipy\nimport scipy.sparse as sp\nfrom chumpy import Ch\n\n__all__ = ['GMOf']\n\n\ndef GMOf(x, sigma):\n    \"\"\"Given x and sigma in some units (say mm),\n    returns robustified values (in same units),\n    by making use of the Geman-McClure robustifier.\"\"\"\n\n    result = SignedSqrt(x=GMOfInternal(x=x, sigma=sigma))\n    return result\n\n\nclass SignedSqrt(Ch):\n    dterms = ('x', )\n    terms = ()\n\n    def compute_r(self):\n        return np.sqrt(np.abs(self.x.r)) * np.sign(self.x.r)\n\n    def compute_dr_wrt(self, wrt):\n        if wrt is self.x:\n            result = (.5 / np.sqrt(np.abs(self.x.r)))\n            result = np.nan_to_num(result)\n            result *= (self.x.r != 0).astype(np.uint32)\n            return sp.spdiags(result.ravel(), [0], self.x.r.size,\n                              self.x.r.size)\n\n\nclass GMOfInternal(Ch):\n    dterms = 'x', 'sigma'\n\n    def on_changed(self, which):\n        if 'sigma' in which:\n            assert (self.sigma.r > 0)\n\n        if 'x' in which:\n            self.squared_input = self.x.r**2.\n\n    def compute_r(self):\n        return (self.sigma.r**2 *\n                (self.squared_input /\n                 (self.sigma.r**2 + self.squared_input))) * np.sign(self.x.r)\n\n    def compute_dr_wrt(self, wrt):\n        if wrt is not self.x and wrt is not self.sigma:\n            return None\n\n        squared_input = self.squared_input\n        result = []\n        if wrt is self.x:\n            dx = self.sigma.r**2 / (self.sigma.r**2 + squared_input\n                                    ) - self.sigma.r**2 * (squared_input / (\n                                        self.sigma.r**2 + squared_input)**2)\n            dx = 2 * self.x.r * dx\n            result.append(\n                scipy.sparse.spdiags(\n                    (dx * np.sign(self.x.r)).ravel(), [0],\n                    self.x.r.size,\n                    self.x.r.size,\n                    format='csc'))\n        if wrt is self.sigma:\n            ds = 2 * self.sigma.r * (squared_input / (\n                self.sigma.r**2 + squared_input)) - 2 * self.sigma.r**3 * (\n                    squared_input / (self.sigma.r**2 + squared_input)**2)\n            result.append(\n                scipy.sparse.spdiags(\n                    (ds * np.sign(self.x.r)).ravel(), [0],\n                    self.x.r.size,\n                    self.x.r.size,\n                    format='csc'))\n\n        if len(result) == 1:\n            return result[0]\n        else:\n            return np.sum(result).tocsc()\n"
  },
  {
    "path": "libs/annotator/smplify/lib/sphere_collisions.py",
    "content": "\"\"\"\nCopyright 2016 Max Planck Society, Federica Bogo, Angjoo Kanazawa. All rights reserved.\nThis software is provided for research purposes only.\nBy using this software you agree to the terms of the SMPLify license here:\n     http://smplify.is.tue.mpg.de/license\n\nThis script implements the interpenetration error term. It uses the capsule\napproximation provided in capsule_body.py, and is differentiable with respect\nto body shape and pose.\n\"\"\"\n\nimport numpy as np\nimport chumpy as ch\nfrom smpl_webuser.lbs import verts_core\n\nfrom .capsule_body import get_capsules, set_sphere_centers,\\\n    get_sphere_bweights, collisions,\\\n    capsule_dist\n\n\nclass SphereCollisions(ch.Ch):\n    dterms = ('pose', 'betas')\n    terms = ('regs', 'model')\n\n    def update_capsules_and_centers(self):\n        centers = [set_sphere_centers(capsule) for capsule in self.capsules]\n        count = 0\n        for capsule in self.capsules:\n            capsule.center_id = count\n            count += len(capsule.centers)\n        self.sph_vs = ch.vstack(centers)\n        self.sph_weights = get_sphere_bweights(self.sph_vs, self.capsules)\n        self.ids0 = []\n        self.ids1 = []\n        self.radiuss = []\n        self.caps_pairs = []\n        for collision in collisions:\n            if hasattr(self, 'no_hands'):\n                (id0, id1, rd) = capsule_dist(\n                    self.capsules[collision[0]],\n                    self.capsules[collision[1]],\n                    increase_hand=False)\n            else:\n                (id0, id1, rd) = capsule_dist(self.capsules[collision[0]],\n                                              self.capsules[collision[1]])\n            self.ids0.append(id0.r)\n            self.ids1.append(id1.r)\n            self.radiuss.append(rd)\n            self.caps_pairs.append(['%02d_%02d' % (collision[0], collision[1])]\n                                   * len(id0))\n\n        self.ids0 = np.concatenate(self.ids0).astype(int)\n        self.ids1 = np.concatenate(self.ids1).astype(int)\n        self.radiuss = np.concatenate(self.radiuss)\n        self.caps_pairs = np.concatenate(self.caps_pairs)\n\n        assert (self.caps_pairs.size == self.ids0.size)\n        assert (self.radiuss.size == self.ids0.size)\n\n    def update_pose(self):\n        self.sph_v = verts_core(\n            self.pose,\n            self.sph_vs,\n            self.model.J,\n            self.sph_weights,\n            self.model.kintree_table,\n            want_Jtr=False)\n\n    def get_objective(self):\n        return ch.sum((ch.exp(\n            -((ch.sum((self.sph_v[self.ids0] - self.sph_v[self.ids1])**2,\n                      axis=1)) / (self.radiuss)) / 2.))**2)**.5\n\n    def compute_r(self):\n        return self.get_objective().r\n\n    def compute_dr_wrt(self, wrt):\n        # we consider derivatives only with respect to pose here,\n        # to avoid bias towards thin body shapes\n        if wrt is self.pose:\n            return self.get_objective().dr_wrt(wrt)\n\n    def on_changed(self, which):\n        if 'regs' in which:\n            self.length_regs = self.regs['betas2lens']\n            self.rad_regs = self.regs['betas2rads']\n\n        if 'betas' in which:\n            if not hasattr(self, 'capsules'):\n                self.capsules = get_capsules(\n                    self.model,\n                    wrt_betas=self.betas,\n                    length_regs=self.length_regs,\n                    rad_regs=self.rad_regs)\n\n            self.update_capsules_and_centers()\n\n        if 'pose' in which:\n            self.update_pose()\n"
  },
  {
    "path": "libs/annotator/smplify/render_model.py",
    "content": "\"\"\"\nCopyright 2016 Max Planck Society, Federica Bogo, Angjoo Kanazawa. All rights reserved.\nThis software is provided for research purposes only.\nBy using this software you agree to the terms of the SMPLify license here:\n     http://smplify.is.tue.mpg.de/license\n\nUtility script for rendering the SMPL model using OpenDR.\n\"\"\"\n\nimport numpy as np\nfrom opendr.camera import ProjectPoints\nfrom opendr.renderer import ColoredRenderer\nfrom opendr.lighting import LambertianPointLight\nimport cv2\n\ncolors = {\n    'pink': [.7, .7, .9],\n    'neutral': [.9, .9, .8],\n    'capsule': [.7, .75, .5],\n    'yellow': [.5, .7, .75],\n}\n\n\ndef _create_renderer(w=640,\n                     h=480,\n                     rt=np.zeros(3),\n                     t=np.zeros(3),\n                     f=None,\n                     c=None,\n                     k=None,\n                     near=.5,\n                     far=10.):\n\n    f = np.array([w, w]) / 2. if f is None else f\n    c = np.array([w, h]) / 2. if c is None else c\n    k = np.zeros(5) if k is None else k\n\n    rn = ColoredRenderer()\n\n    rn.camera = ProjectPoints(rt=rt, t=t, f=f, c=c, k=k)\n    rn.frustum = {'near': near, 'far': far, 'height': h, 'width': w}\n    return rn\n\n\ndef _rotateY(points, angle):\n    \"\"\"Rotate the points by a specified angle.\"\"\"\n    ry = np.array([\n        [np.cos(angle), 0., np.sin(angle)], [0., 1., 0.],\n        [-np.sin(angle), 0., np.cos(angle)]\n    ])\n    return np.dot(points, ry)\n\n\ndef simple_renderer(rn, verts, faces, yrot=np.radians(120)):\n\n    # Rendered model color\n    color = colors['pink']\n\n    rn.set(v=verts, f=faces, vc=color, bgcolor=np.ones(3))\n\n    albedo = rn.vc\n\n    # Construct Back Light (on back right corner)\n    rn.vc = LambertianPointLight(\n        f=rn.f,\n        v=rn.v,\n        num_verts=len(rn.v),\n        light_pos=_rotateY(np.array([-200, -100, -100]), yrot),\n        vc=albedo,\n        light_color=np.array([1, 1, 1]))\n\n    # Construct Left Light\n    rn.vc += LambertianPointLight(\n        f=rn.f,\n        v=rn.v,\n        num_verts=len(rn.v),\n        light_pos=_rotateY(np.array([800, 10, 300]), yrot),\n        vc=albedo,\n        light_color=np.array([1, 1, 1]))\n\n    # Construct Right Light\n    rn.vc += LambertianPointLight(\n        f=rn.f,\n        v=rn.v,\n        num_verts=len(rn.v),\n        light_pos=_rotateY(np.array([-500, 500, 1000]), yrot),\n        vc=albedo,\n        light_color=np.array([.7, .7, .7]))\n\n    return rn.r\n\n\ndef get_alpha(imtmp, bgval=1.):\n    h, w = imtmp.shape[:2]\n    alpha = (~np.all(imtmp == bgval, axis=2)).astype(imtmp.dtype)\n\n    b_channel, g_channel, r_channel = cv2.split(imtmp)\n\n    im_RGBA = cv2.merge(\n        (b_channel, g_channel, r_channel, alpha.astype(imtmp.dtype)))\n    return im_RGBA\n\n\ndef render_model(verts, faces, w, h, cam, near=0.5, far=25, img=None):\n    rn = _create_renderer(\n        w=w, h=h, near=near, far=far, rt=cam.rt, t=cam.t, f=cam.f, c=cam.c)\n    # Uses img as background, otherwise white background.\n    if img is not None:\n        rn.background_image = img / 255. if img.max() > 1 else img\n\n    imtmp = simple_renderer(rn, verts, faces)\n\n    # If white bg, make transparent.\n    if img is None:\n        imtmp = get_alpha(imtmp)\n\n    return imtmp\n"
  },
  {
    "path": "libs/dataset/__init__.py",
    "content": "#!/usr/bin/env python3\n# -*- coding: utf-8 -*-\n\"\"\"\nEmpty file.\n\"\"\"\n#import libs.dataset.h36m\n\n"
  },
  {
    "path": "libs/dataset/h36m/__init__.py",
    "content": "#!/usr/bin/env python3\n# -*- coding: utf-8 -*-\n\n"
  },
  {
    "path": "libs/dataset/h36m/cameras.py",
    "content": "\"\"\"\nUtilities to deal with the cameras of human3.6m.\nReference: https://github.com/una-dinosauria/3d-pose-baseline \n\"\"\"\nimport numpy as np\n# import h5py\n\ndef project_point_radial( P, R, T, f, c, k, p ):\n    \"\"\"\n    Project points from 3d to 2d using camera parameters\n    including radial and tangential distortion\n    \n    Args\n        P: Nx3 points in world coordinates\n        R: 3x3 Camera rotation matrix\n        T: 3x1 Camera translation parameters\n        f: (scalar) Camera focal length\n        c: 2x1 Camera center\n        k: 3x1 Camera radial distortion coefficients\n        p: 2x1 Camera tangential distortion coefficients\n    Returns\n        Proj: Nx2 points in pixel space\n        D: 1xN depth of each point in camera space\n        radial: 1xN radial distortion per point\n        tan: 1xN tangential distortion per point\n        r2: 1xN squared radius of the projected points before distortion\n    \"\"\"\n    # P is a matrix of 3-dimensional points\n    assert len(P.shape) == 2\n    assert P.shape[1] == 3\n    \n    N = P.shape[0]\n    X = R.dot( P.T - T ) # rotate and translate\n    XX = X[:2,:] / X[2,:]\n    r2 = XX[0,:]**2 + XX[1,:]**2\n    \n    radial = 1 + np.einsum( 'ij,ij->j', np.tile(k,(1, N)), \n                           np.array([r2, r2**2, r2**3])\n                           );\n    tan = p[0]*XX[1,:] + p[1]*XX[0,:]\n    \n    XXX = (XX * np.tile(radial + tan,(2,1)) \n          + np.outer(np.array([p[1], p[0]]).reshape(-1), r2))\n    \n    Proj = (f * XXX) + c\n    Proj = Proj.T\n    \n    D = X[2,]\n    return Proj, D, radial, tan, r2\n\ndef world_to_camera_frame(P, R, T):\n    \"\"\"\n    Convert points from world to camera coordinates\n    \n    Args\n        P: Nx3 3d points in world coordinates\n        R: 3x3 Camera rotation matrix\n        T: 3x1 Camera translation parameters\n    Returns\n        X_cam: Nx3 3d points in camera coordinates\n    \"\"\"\n\n    assert len(P.shape) == 2\n    assert P.shape[1] == 3\n    \n    X_cam = R.dot( P.T - T ) # rotate and translate\n    \n    return X_cam.T\n\ndef camera_to_world_frame(P, R, T):\n    \"\"\"\n    Inverse of world_to_camera_frame\n    \n    Args\n        P: Nx3 points in camera coordinates\n        R: 3x3 Camera rotation matrix\n        T: 3x1 Camera translation parameters\n    Returns\n        X_cam: Nx3 points in world coordinates\n    \"\"\"\n\n    assert len(P.shape) == 2\n    assert P.shape[1] == 3\n    \n    X_cam = R.T.dot( P.T ) + T # rotate and translate\n    \n    return X_cam.T\n\ndef load_camera_params( hf, path ):\n    \"\"\"\n    Load h36m camera parameters\n    \n    Args\n        hf: hdf5 open file with h36m cameras data\n        path: path or key inside hf to the camera we are interested in\n    Returns\n        R: 3x3 Camera rotation matrix\n        T: 3x1 Camera translation parameters\n        f: (scalar) Camera focal length\n        c: 2x1 Camera center\n        k: 3x1 Camera radial distortion coefficients\n        p: 2x1 Camera tangential distortion coefficients\n        name: String with camera id\n    \"\"\"\n\n    R = hf[ path.format('R') ][:]\n    R = R.T\n    \n    T = hf[ path.format('T') ][:]\n    f = hf[ path.format('f') ][:]\n    c = hf[ path.format('c') ][:]\n    k = hf[ path.format('k') ][:]\n    p = hf[ path.format('p') ][:]\n    \n    name = hf[ path.format('Name') ][:]\n    name = \"\".join( [chr(item) for item in name] )\n    \n    return R, T, f, c, k, p, name\n\n# def load_cameras( bpath='cameras.h5', subjects=[1,5,6,7,8,9,11] ):\n#     \"\"\"\n#     Loads the cameras of h36m\n\n#     Args\n#         bpath: path to hdf5 file with h36m camera data\n#         subjects: List of ints representing the subject IDs for which cameras \n#         are requested\n#     Returns\n#         rcams: dictionary of 4 tuples per subject ID containing its camera \n#         parameters for the 4 h36m cams\n#     \"\"\"\n#     rcams = {}\n    \n#     with h5py.File(bpath,'r') as hf:\n#         for s in subjects:\n#             for c in range(4): # There are 4 cameras in human3.6m\n#                 string = 'subject%d/camera%d/{0}' % (s,c+1)\n#                 rcams[(s, c+1)] = load_camera_params(hf, string)\n    \n#     return rcams\n"
  },
  {
    "path": "libs/dataset/h36m/data_utils.py",
    "content": "\"\"\"\nUtility functions for dealing with Human3.6M dataset.\nSome functions are adapted from https://github.com/una-dinosauria/3d-pose-baseline \n\"\"\"\n\nimport os\nimport numpy as np\nimport copy\nimport logging\nimport matplotlib.pyplot as plt\nimport torch\nfrom mpl_toolkits.mplot3d import Axes3D\n\nimport libs.dataset.h36m.cameras as cameras\nimport libs.dataset.h36m.pth_dataset as dataset\n# Human3.6m IDs for training and testing\nTRAIN_SUBJECTS = [1, 5, 6, 7, 8]\nTEST_SUBJECTS  = [9, 11]\n\n# Use camera coordinate system \ncamera_frame = True\n\n# Joint names in H3.6M -- data has 32 joints, but only 17 that move; \n# these are the indices.\nH36M_NAMES = ['']*32\nH36M_NAMES[0]  = 'Hip'\nH36M_NAMES[1]  = 'RHip'\nH36M_NAMES[2]  = 'RKnee'\nH36M_NAMES[3]  = 'RFoot'\nH36M_NAMES[6]  = 'LHip'\nH36M_NAMES[7]  = 'LKnee'\nH36M_NAMES[8]  = 'LFoot'\nH36M_NAMES[12] = 'Spine'\nH36M_NAMES[13] = 'Thorax'\nH36M_NAMES[14] = 'Neck/Nose'\nH36M_NAMES[15] = 'Head'\nH36M_NAMES[17] = 'LShoulder'\nH36M_NAMES[18] = 'LElbow'\nH36M_NAMES[19] = 'LWrist'\nH36M_NAMES[25] = 'RShoulder'\nH36M_NAMES[26] = 'RElbow'\nH36M_NAMES[27] = 'RWrist'\n\nparent_indices = np.array([0, 1, 2, 0, 6, 7, 0, 12, 13, 14, 13, 17, 18, 13, 25, 26])\nchildren_indices = np.array([1, 2, 3, 6, 7, 8, 12, 13, 14, 15, 17, 18, 19, 25, 26, 27])\n\n# Stacked Hourglass produces 16 joints. These are the names.\nSH_NAMES = ['']*16\nSH_NAMES[0]  = 'RFoot'\nSH_NAMES[1]  = 'RKnee'\nSH_NAMES[2]  = 'RHip'\nSH_NAMES[3]  = 'LHip'\nSH_NAMES[4]  = 'LKnee'\nSH_NAMES[5]  = 'LFoot'\nSH_NAMES[6]  = 'Hip'\nSH_NAMES[7]  = 'Spine'\nSH_NAMES[8]  = 'Thorax'\nSH_NAMES[9]  = 'Head'\nSH_NAMES[10] = 'RWrist'\nSH_NAMES[11] = 'RElbow'\nSH_NAMES[12] = 'RShoulder'\nSH_NAMES[13] = 'LShoulder'\nSH_NAMES[14] = 'LElbow'\nSH_NAMES[15] = 'LWrist'\n\n# The .h5 suffix in pose sequence name is just inherited from the original \n# naming convention. The '.sh' suffix means stacked hourglass key-point detector \n# used in previous works. Here we just use '.sh' to represent key-points obtained\n# from any heat-map regression model. We used high-resolution net instead of \n# stacked-hourglass model.  \n\ndef load_ckpt(opt):\n    cascade = torch.load(os.path.join(opt.ckpt_dir, 'model.th'))\n    stats = np.load(os.path.join(opt.ckpt_dir, 'stats.npy'), allow_pickle=True).item()\n    if opt.cuda:\n        cascade.cuda()\n    return cascade, stats\n\ndef list_remove(list_a, list_b):\n    \"\"\"\n    Fine all elements of a list A that does not exist in list B.\n    \n    Args\n      list_a: list A\n      list_b: list B\n    Returns\n      list_c: result  \n    \"\"\"\n    list_c = []\n    for item in list_a:\n        if item not in list_b:\n            list_c.append(item)\n    return list_c\n\ndef add_virtual_cams(cams, visualize=False):\n    \"\"\"\n    Deprecated. Add virtual cameras.\n    \"\"\"\n    # add more cameras to the scene\n    #R, T, f, c, k, p, name = cams[ (1,1) ]\n    # plot the position of human subjects\n    old_cam_num = 4\n    def add_coordinate_system(ax, origin, system, length=300, new=False):\n        # draw a coordinate system at a specified origin\n        origin = origin.reshape(3, 1)\n        start_points = np.repeat(origin, 3, axis=1)\n        # system: [v1, v2, v3] \n        end_points = start_points + system*length\n        color = ['g', 'y', 'k'] # color for v1, v2 and v3\n        if new:\n            color = ['b', 'r', 'g']\n        def get_args(start_points, end_points):\n            x = [start_points[0], end_points[0]]\n            y = [start_points[1], end_points[1]]\n            z = [start_points[2], end_points[2]]\n            return x, y, z\n        for i in range(3):\n            x, y, z = get_args(start_points[:,i], end_points[:,i])\n            ax.plot(x, y, z,  lw=2, c=color[i])\n        return\n    \n    def get_new_camera(system, center, rotation = [0,0,90.]):\n        from scipy.spatial.transform import Rotation as Rotation\n        center = center.reshape(3, 1)\n        start_points = np.repeat(center, 3, axis=1)\n        end_points = start_points + system\n        r = Rotation.from_euler('xyz', rotation, degrees=True)\n        start_points_new = r.as_dcm() @ start_points  \n        end_points_new = r.as_dcm() @ end_points\n        new_system = [(end_points_new[:,i] - start_points_new[:,i]).reshape(3,1) for i in range(3)]\n        new_system = np.hstack(new_system)\n        return new_system, start_points_new[:,0]    \n\n    \n    # the new cameras are added by rotating one existing camera\n    # TODO: more rotations\n    new_cams = cams.copy()\n    for key in cams.keys():\n        subject, camera_idx = key\n        if camera_idx != 1: # only rotate the first camera\n            continue\n        R, T, f, c, k, p, name = cams[key]\n        angles = [80., 130., 270., 320.]\n        for angle_idx in range(len(angles)):\n            angle = angles[angle_idx]\n            new_R, new_T = get_new_camera(R.T, T, [0., 0., angle])\n            new_cams[(subject, old_cam_num + angle_idx + 1)]\\\n            = (new_R.T, new_T.reshape(3,1), f, c, k, p, name+'new'+str(angle_idx+1))\n    # visualize cameras used\n    if visualize:\n        train_set_3d = np.load('../data/human3.6M/h36m/numpy/threeDPose_train.npy').item()\n        test_set_3d = np.load('../data/human3.6M/h36m/numpy/threeDPose_test.npy').item()\n        hips_train = np.vstack(list(train_set_3d.values()))\n        hips_test = np.vstack(list(test_set_3d.values()))        \n        ax = plt.subplot(111, projection='3d')\n        chosen = np.random.choice(len(hips_train), 1000, replace=False)\n        chosen_hips = hips_train[chosen, :3]\n        ax.plot(chosen_hips[:,0], chosen_hips[:,1], chosen_hips[:,2], 'bo')\n        chosen = np.random.choice(len(hips_test), 1000, replace=False)\n        chosen_hips = hips_test[chosen, :3]\n        ax.plot(chosen_hips[:,0], chosen_hips[:,1], chosen_hips[:,2], 'ro')        \n        ax.set_xlabel(\"x\");ax.set_ylabel(\"y\");ax.set_zlabel(\"z\")\n        plt.title('Blue dots: Hip positions in the h36m training set. \\\n                  Red dots: testing set. \\\n                  Old camera coordinates: x-green, y-yellow, z-black \\\n                  New camera coordinates: x-blue, y-red, z-green')\n        plt.pause(0.1)\n        for key in new_cams.keys():\n            R, T, f, c, k, p, name = new_cams[key]\n            # R gives camera basis vectors row-by-row, T gives camera center\n            if 'new' in name:\n                new = True\n            else:\n                new = False\n            add_coordinate_system(ax, T, R.T, new=new)\n        RADIUS = 3000 # space around the subject\n        xroot, yroot, zroot = 0., 0., 500.\n        ax.set_xlim3d([-RADIUS+xroot, RADIUS+xroot])\n        ax.set_zlim3d([-RADIUS+zroot, RADIUS+zroot])\n        ax.set_ylim3d([-RADIUS+yroot, RADIUS+yroot])\n        ax.set_aspect(\"equal\")\n    return new_cams\n\ndef down_sample_training_data(train_dict, opt):\n    \"\"\"\n    Down-sample the training data.\n\n    Args\n      train_dict: python dictionary contraining the training data\n      opt: experiment options\n    Returns\n      train_dict/sampled_dict: a dictionary containing a subset of training data\n    \"\"\"\n    if opt.ws_name in ['S1', 'S15', 'S156']:\n        sub_list = [int(opt.ws_name[i]) for i in range(1, len(opt.ws_name))]\n        keys_to_delete = []\n        for key in train_dict.keys():\n            if key[0] not in sub_list:\n                keys_to_delete.append(key)\n        for key in keys_to_delete:    \n            del train_dict[key]\n        return train_dict\n    elif opt.ws_name in ['0.001S1','0.01S1', '0.05S1', '0.1S1', '0.5S1']:\n        ratio = float(opt.ws_name.split('S')[0])\n        # randomly sample a portion of 3D data\n        sampled_dict = {}\n        for key in train_dict.keys():\n            if key[0] != 1:\n                continue\n            total = len(train_dict[key])\n            sampled_num = int(ratio*total)\n            chosen_indices = np.random.choice(total, sampled_num, replace=False)\n            sampled_dict[key] = train_dict[key][chosen_indices].copy()\n        return sampled_dict\n    else:\n        raise ValueError('Unknown experiment setting.')\n    \ndef get_train_dict_3d(opt):\n    \"\"\"\n    Get the training 3d skeletons as a Python dictionary.\n\n    Args\n      opt: experiment options\n    Returns\n      train_dict_3d: a dictionary containing training 3d poses\n    \"\"\"\n    if not opt.train:\n        return None\n    dict_path = os.path.join(opt.data_dir, 'threeDPose_train.npy')\n    #=========================================================================#\n    # For real 2D detections, the down-sampling and data augmentation \n    # are done later in get_train_dict_2d\n    if opt.twoD_source != 'synthetic':\n        train_dict_3d = np.load(dict_path, allow_pickle=True).item()\n        return train_dict_3d\n    #=========================================================================#\n    # For synthetic 2D detections (For Protocol P1*), the down-sampling is \n    # performed here and the data augmentation is assumed to be already done\n    if opt.evolved_path is not None:\n        # the data is pre-augmented\n        train_dict_3d = np.load(opt.evolved_path, allow_pickle=True).item()\n    elif opt.ws:\n        # raw training data from Human 3.6M (S15678) \n        # Down-sample the raw data to simulate an environment with scarce \n        # training data, which is used in weakly-supervised experiments        \n        train_dict_3d = np.load(dict_path, allow_pickle=True).item()\n        train_dict_3d = down_sample_training_data(train_dict_3d, opt)\n    else:\n        # raw training data from Human 3.6M (S15678) \n        train_dict_3d = np.load(dict_path, allow_pickle=True).item()\n    return train_dict_3d\n\ndef get_test_dict_3d(opt):\n    \"\"\"\n    Get the testing 3d skeletons as a Python dictionary.\n\n    Args\n      opt: experiment options\n    Returns\n      test_dict_3d: a dictionary containing testing 3d poses\n    \"\"\"       \n    if opt.test_source == 'h36m':     \n        # for h36m\n        dict_path = os.path.join(opt.data_dir, 'threeDPose_test.npy')\n        test_dict_3d  = np.load(dict_path, allow_pickle=True).item()   \n    else:\n        raise NotImplementedError    \n    return test_dict_3d\n\ndef get_dict_2d(train_dict_3d, test_dict_3d, rcams, ncams, opt):\n    \"\"\"\n    Prepare 2D training and testing data as Python dictionaries.\n\n    Args\n      train_dict_3d: dictionary containing training 3d poses\n      test_dict_3d: dictionary containing testing 3d poses\n      rcams: camera parameters\n      ncams: number of camera to use\n      opt: experiment options\n    Returns\n      train_dict_2d: a dictionary containing training 2d poses\n      test_dict_2d: a dictionary containing testing 2d poses\n      train_dict_3d: the dictionary containing training 3d poses, which may be\n      updated\n    \"\"\"     \n    if opt.twoD_source == 'synthetic':\n        # project the 3D key-points to 2D ones\n        # This type of key-points is used to validate the performance of \n        # 2D-to-3D networks and the noise of 2D key-point detector is ignored.\n        # In fact, these 2D key-points are used as ground-truth to train the \n        # first stage of TAG-net.\n        if opt.virtual_cams:\n            ncams *= 2       \n        if opt.train:\n            train_dict_2d = project_to_cameras(train_dict_3d, rcams, ncams=ncams)\n        else:\n            train_dict_2d = None\n        test_dict_2d  = project_to_cameras(test_dict_3d, rcams, ncams=ncams)\n    elif opt.twoD_source == 'HRN':\n        # The 2D key-point detections obtained by the heatmap regression model.\n        # The model uses high-resolution net as backbone and pixel-shuffle super-resolution\n        # to regress high-resolution heatmaps.\n        if opt.train:\n            train_dict_2d = np.load(os.path.join(opt.data_dir, 'twoDPose_HRN_train.npy'), allow_pickle=True).item()           \n        else:\n            train_dict_2d = None\n        test_dict_2d = np.load(os.path.join(opt.data_dir, 'twoDPose_HRN_test.npy'), allow_pickle=True).item()\n        \n        def delete(dic, actions):\n            keys_to_delete = []\n            for key in dic.keys():\n                sub, act, name = key\n                if act not in actions:\n                    keys_to_delete.append(key)\n            for key in keys_to_delete:\n                del dic[key]\n            return dic\n\n        def replace(dic, temp):\n            for key in dic.keys():\n                sub, act, name = key\n                temp_key = (sub, act, name[:-3])\n                synthetic = temp[temp_key]\n                assert len(dic[key]) == len(synthetic)\n                indices = np.random.choice(len(synthetic), int(0.5*len(synthetic)), replace=False)\n                dic[key][indices] = synthetic[indices].copy()\n            return dic\n\n#        # weakly-supervised experiment\n        def remove_keys(dic, name_list):\n            keys_to_delete = []\n            for key in dic.keys():\n                if key[0] not in name_list:\n                    keys_to_delete.append(key)\n            for key in keys_to_delete:    \n                del dic[key]\n            return dic\n        \n        # down-sample the data for weakly-supervised experiment\n        if opt.ws and opt.ws_name in ['S1', 'S15', 'S156']:\n            sub_list = [int(opt.ws_name[i]) for i in range(1, len(opt.ws_name))]\n            remove_keys(train_dict_3d, sub_list)\n            if train_dict_2d is not None:\n                remove_keys(train_dict_2d, sub_list)\n        # data augmentation with evolved data\n        if opt.evolved_path is not None:\n            evolved_dict_3d = np.load(opt.evolved_path, allow_pickle=True).item()\n            evolved_dict_2d = project_to_cameras(evolved_dict_3d, rcams, ncams=ncams)\n            # combine the synthetic 2D-3D pair with the real 2D-3D pair\n            train_dict_3d = {**train_dict_3d, **evolved_dict_3d}\n            train_dict_2d = {**train_dict_2d, **evolved_dict_2d} \n    return train_dict_2d, test_dict_2d, train_dict_3d        \n\ndef prepare_data_dict(rcams, \n                      opt,\n                      ncams=4,\n                      predict_14=False, \n                      use_nose=True\n                      ):\n    \"\"\"\n    Prepare 2D and 3D data as Python dictionaries.\n    \n    Args\n      rcams: camera parameters\n      opt: experiment options\n      ncams: number of camera to use\n      predict_14: whether to predict 14 joints or not\n      use_nose: whether to use nose joint or not\n    Returns\n      data_dic: a dictionary containing training and testing data\n      data_stats: statistics computed from training data \n    \"\"\"\n    assert opt.twoD_source in ['synthetic', 'HRN'], 'Unknown 2D key-point type.'\n    data_dic = {}\n    # get 3D skeleton data\n    train_dict_3d = get_train_dict_3d(opt)\n    test_dict_3d = get_test_dict_3d(opt)\n    # get 2D key-point data\n    train_dict_2d, test_dict_2d, train_dict_3d = get_dict_2d(train_dict_3d, \n                                                             test_dict_3d, \n                                                             rcams, \n                                                             ncams,\n                                                             opt\n                                                             )\n    # compute normalization statistics and normalize the 2D data\n    if opt.train:\n        complete_train_2d = copy.deepcopy(np.vstack(list(train_dict_2d.values())))\n        data_mean_2d, data_std_2d, dim_to_ignore_2d, dim_to_use_2d = \\\n        normalization_stats(complete_train_2d, \n                            dim=2, \n                            norm_twoD=opt.norm_twoD, \n                            use_nose=use_nose\n                            )\n        \n        data_dic['train_set_2d'] = normalize_data(train_dict_2d,\n                                                  data_mean_2d,\n                                                  data_std_2d, \n                                                  dim_to_use_2d, \n                                                  norm_single=opt.norm_single\n                                                  )\n    else:\n        _, data_stats = load_ckpt(opt)\n        data_mean_2d, data_std_2d = data_stats['mean_2d'], data_stats['std_2d']\n        dim_to_use_2d = data_stats['dim_use_2d']\n    data_dic['test_set_2d']  = normalize_data(test_dict_2d,\n                                              data_mean_2d,\n                                              data_std_2d, \n                                              dim_to_use_2d, \n                                              norm_single=opt.norm_single\n                                              )\n    # The 3D joint position is represented in the world coordinate,\n    # which is converted to camera coordinate system as the regression target\n    if opt.train:\n        train_dict_3d = transform_world_to_camera(train_dict_3d, rcams, ncams=ncams)\n        # apply 3d post-processing (centering around root)\n        train_dict_3d, train_root_positions = postprocess_3d(train_dict_3d)\n    test_dict_3d  = transform_world_to_camera(test_dict_3d, rcams, ncams=ncams)\n    test_dict_3d,  test_root_positions  = postprocess_3d(test_dict_3d)\n    if opt.train:\n        # compute normalization statistics and normalize the 3D data\n        complete_train_3d = copy.deepcopy(np.vstack(list(train_dict_3d.values())))\n        data_mean_3d, data_std_3d, dim_to_ignore_3d, dim_to_use_3d =\\\n        normalization_stats(complete_train_3d, dim=3, predict_14=predict_14)\n        data_dic['train_set_3d'] = normalize_data(train_dict_3d,\n                                                  data_mean_3d,\n                                                  data_std_3d, \n                                                  dim_to_use_3d\n                                                  )\n        # some joints are not used during training\n        dim_use_2d = list_remove([i for i in range(len(data_mean_2d))], \n                                  list(dim_to_ignore_2d))\n        dim_use_3d = list_remove([i for i in range(len(data_mean_3d))], \n                                  list(dim_to_ignore_3d))\n        # assemble a dictionary for data statistics\n        data_stats = {'mean_2d':data_mean_2d, \n                      'std_2d':data_std_2d,\n                      'mean_3d':data_mean_3d, \n                      'std_3d':data_std_3d,\n                      'dim_ignore_2d':dim_to_ignore_2d, \n                      'dim_ignore_3d':dim_to_ignore_3d,\n                      'dim_use_2d':dim_use_2d,\n                      'dim_use_3d':dim_use_3d\n                      }         \n    else:\n        data_mean_3d, data_std_3d = data_stats['mean_3d'], data_stats['std_3d']\n        dim_to_use_3d = data_stats['dim_use_3d']            \n    data_dic['test_set_3d']  = normalize_data(test_dict_3d,  \n                                              data_mean_3d,\n                                              data_std_3d, \n                                              dim_to_use_3d\n                                              )    \n   \n    return data_dic, data_stats\n\ndef select_action(dic_2d, dic_3d, action, twoD_source):\n    \"\"\"\n    Construct sub-dictionaries by specifying which action to use\n    \n    Args\n        dic_2d: dictionary containing 2d poses\n        dic_3d: dictionary containing 3d poses\n        action: the action to use\n        twoD_source: how the key-points are generated (synthetic or real)\n    Returns\n        dic_2d_action: sub-dictionary containing 2d poses for the specified action\n        dic_3d_action: sub-dictionary containing 3d poses for the specified action \n    \"\"\"    \n    dic_2d_action = {}\n    dic_3d_action = {}\n    for key in dic_2d.keys():\n        if key[1] == action:\n            dic_2d_action[key] = dic_2d[key].copy()\n            if twoD_source == 'synthetic':\n                key3d = key\n            else:\n                key3d = (key[0], key[1], key[2][:-3])\n            dic_3d_action[key3d] = dic_3d[key3d].copy()    \n    return dic_2d_action, dic_3d_action\n\ndef split_action(dic_2d, dic_3d, actions, camera_frame, opt, input_size, output_size):\n    \"\"\"\n    Generate a list of datasets for each action.\n    \n    Args\n        dic_2d: dictionary containing 2d poses\n        dic_3d: dictionary containing 3d poses\n        actions: list of defined actions\n        camera_frame: use camera coordinate system\n        opt: experiment options\n        input_size: input vector length\n        output_size: output vector length\n    Returns\n        action_dataset_list: a list of datasets where each element correspond\n        to one action   \n    \"\"\"\n    action_dataset_list = []\n    for act_id in range(len(actions)):\n        action = actions[act_id]\n        dic_2d_action, dic_3d_action = select_action(dic_2d, dic_3d, action, opt.twoD_source)\n        eval_input, eval_output = get_all_data(dic_2d_action, \n                                               dic_3d_action,\n                                               camera_frame, \n                                               norm_twoD=opt.norm_twoD,\n                                               input_size=input_size,\n                                               output_size=output_size)\n        action_dataset = dataset.PoseDataset(eval_input, \n                                             eval_output, \n                                             'eval', \n                                             action_name=action,\n                                             refine_3d=opt.refine_3d)\n        action_dataset_list.append(action_dataset)\n    return action_dataset_list\n\ndef normalization_stats(complete_data, \n                        dim, \n                        predict_14=False, \n                        norm_twoD=False, \n                        use_nose=False\n                        ):\n    \"\"\"\n    Computes normalization statistics: mean and stdev, dimensions used and ignored\n    \n    Args\n        complete_data: nxd np array with poses\n        dim. integer={2,3} dimensionality of the data\n        predict_14. boolean. Whether to use only 14 joints\n        use_nose: whether to use nose or not\n    Returns\n        data_mean: np vector with the mean of the data\n        data_std: np vector with the standard deviation of the data\n        dimensions_to_ignore: list of dimensions not used in the model\n        dimensions_to_use: list of dimensions used in the model\n    \"\"\"\n    if not dim in [2,3]:\n        raise(ValueError, 'dim must be 2 or 3')\n    data_mean = np.mean(complete_data, axis=0)\n    data_std  =  np.std(complete_data, axis=0)\n    # Encodes which 17 (or 14) 2d-3d pairs we are predicting\n    dimensions_to_ignore = []\n    if dim == 2:\n        if not use_nose:\n            dimensions_to_use = np.where(np.array([x != '' and x != 'Neck/Nose' for x in H36M_NAMES]))[0]\n        else:\n            dimensions_to_use = np.where(np.array([x != '' for x in H36M_NAMES]))[0]        \n        if norm_twoD:\n            dimensions_to_use = np.delete(dimensions_to_use, 0)\n        dimensions_to_use = np.sort(np.hstack((dimensions_to_use*2, \n                                               dimensions_to_use*2+1)))\n        dimensions_to_ignore = np.delete(np.arange(len(H36M_NAMES)*2), \n                                         dimensions_to_use)\n    else: \n        dimensions_to_use = np.where(np.array([x != '' for x in H36M_NAMES]))[0]\n        # hip is deleted\n        # spine and neck are also deleted if predict_14 \n        dimensions_to_use = np.delete(dimensions_to_use, [0,7,9] if predict_14 else 0)\n        dimensions_to_use = np.sort(np.hstack((dimensions_to_use*3,\n                                               dimensions_to_use*3+1,\n                                               dimensions_to_use*3+2)))\n        dimensions_to_ignore = np.delete(np.arange(len(H36M_NAMES)*3), \n                                         dimensions_to_use)\n    return data_mean, data_std, dimensions_to_ignore, dimensions_to_use\n\ndef transform_world_to_camera(poses_set, cams, ncams=4):\n    \"\"\"\n    Transform 3d poses from world coordinate to camera coordinate system\n    Args\n      poses_set: dictionary with 3d poses\n      cams: dictionary with cameras\n      ncams: number of cameras per subject\n    Return:\n      t3d_camera: dictionary with 3d poses in camera coordinate\n    \"\"\"\n    t3d_camera = {}\n    for t3dk in sorted(poses_set.keys()):\n      subj, action, seqname = t3dk\n      t3d_world = poses_set[t3dk]\n      for c in range(ncams):\n        R, T, f, c, k, p, name = cams[(subj, c+1)]\n        camera_coord = cameras.world_to_camera_frame(np.reshape(t3d_world, [-1, 3]), R, T)\n        camera_coord = np.reshape(camera_coord, [-1, len(H36M_NAMES)*3])\n        sname = seqname[:-3]+\".\"+name+\".h5\" # e.g.: Waiting 1.58860488.h5\n        t3d_camera[(subj, action, sname)] = camera_coord\n    return t3d_camera\n\ndef normalize_data(data, data_mean, data_std, dim_to_use, norm_single=False):\n    \"\"\"\n    Normalizes a dictionary of poses\n    \n    Args\n        data: dictionary where values are\n        data_mean: np vector with the mean of the data\n        data_std: np vector with the standard deviation of the data\n        dim_to_use: list of dimensions to keep in the data\n        norm_single: whether to perform normalization independently for each \n        sample\n    Returns\n        data_out: dictionary with same keys as data, but values have been normalized\n    \"\"\"\n    data_out = {}\n    for key in data.keys():\n        data[ key ] = data[ key ][ :, dim_to_use ]\n        if norm_single:\n            # does not use statistics over the whole dataset\n            temp = data[key]\n            temp = temp.reshape(len(temp), -1, 2)\n            mean_x = np.mean(temp[:,:,0], axis=1).reshape(len(temp), 1)\n            std_x = np.std(temp[:,:,0], axis=1)\n            mean_y = np.mean(temp[:,:,1], axis=1).reshape(len(temp), 1)\n            std_y = np.std(temp[:,:,1], axis=1)\n            denominator = (0.5*(std_x + std_y)).reshape(len(std_x), 1)\n            temp[:,:,0] = (temp[:,:,0] - mean_x)/denominator\n            temp[:,:,1] = (temp[:,:,1] - mean_y)/denominator\n            data_out[key] = temp.reshape(len(temp), -1)\n        else:\n            mu = data_mean[dim_to_use]\n            stddev = data_std[dim_to_use]\n            data_out[ key ] = np.divide( (data[key] - mu), stddev )    \n    return data_out\n\ndef unNormalizeData(normalized_data, data_mean, data_std, dimensions_to_ignore):\n    \"\"\"\n    Un-normalizes a matrix whose mean has been substracted and that has been \n    divided by standard deviation. Some dimensions might also be missing.\n    \n    Args\n    normalized_data: nxd matrix to unnormalize\n    data_mean: np vector with the mean of the data\n    data_std: np vector with the standard deviation of the data\n    dimensions_to_ignore: list of dimensions that were removed from the original data\n    Returns\n    orig_data: the unnormalized data\n    \"\"\"\n    T = normalized_data.shape[0] # batch size\n    D = data_mean.shape[0] # dimensionality\n    orig_data = np.zeros((T, D), dtype=np.float32)\n    dimensions_to_use = np.array([dim for dim in range(D)\n                                    if dim not in dimensions_to_ignore])\n    orig_data[:, dimensions_to_use] = normalized_data\n    # multiply times stdev and add the mean\n    stdMat = data_std.reshape((1, D))\n    stdMat = np.repeat(stdMat, T, axis=0)\n    meanMat = data_mean.reshape((1, D))\n    meanMat = np.repeat(meanMat, T, axis=0)\n    orig_data = np.multiply(orig_data, stdMat) + meanMat\n    return orig_data\n\ndef define_actions(action):\n    \"\"\"\n    Given an action string, returns a list of corresponding actions.\n    \n    Args\n        action: String. either \"all\" or one of the h36m actions\n    Returns\n        actions: List of strings. Actions to use.\n    Raises\n        ValueError: if the action is not a valid action in Human 3.6M\n    \"\"\"\n    actions = [\"Directions\",\n               \"Discussion\",\n               \"Eating\",\n               \"Greeting\",\n               \"Phoning\",\n               \"Photo\",\n               \"Posing\",\n               \"Purchases\",\n               \"Sitting\",\n               \"SittingDown\",\n               \"Smoking\",\n               \"Waiting\",\n               \"WalkDog\",\n               \"Walking\",\n               \"WalkTogether\"\n               ]\n    \n    if action == \"All\" or action == \"all\":\n        return actions\n    \n    if not action in actions:\n        raise( ValueError, \"Unrecognized action: %s\" % action )\n    \n    return [action]\n\ndef project_to_cameras(poses_set, cams, ncams=4):\n    \"\"\"\n    Project 3d poses using camera parameters\n    \n    Args\n        poses_set: dictionary containing 3d poses\n        cams: dictionary containing camera parameters\n        ncams: number of cameras per subject\n    Returns\n        t2d: dictionary with 2d poses\n    \"\"\"\n    t2d = {}\n    for t3dk in sorted(poses_set.keys()):\n        subj, a, seqname = t3dk\n        t3d = poses_set[t3dk]\n        for cam in range(ncams):\n            R, T, f, c, k, p, name = cams[(subj, cam+1)]\n            pts2d, _, _, _, _ = cameras.project_point_radial(np.reshape(t3d, [-1, 3]), R, T, f, c, k, p)        \n            pts2d = np.reshape(pts2d, [-1, len(H36M_NAMES)*2])\n            sname = seqname[:-3] + \".\" + name + \".h5\" # e.g.: Waiting 1.58860488.h5\n            t2d[ (subj, a, sname) ] = pts2d\n    return t2d\n\ndef postprocess_3d(poses_set):\n    \"\"\"\n    Center 3d points around root\n    \n    Args\n        poses_set: dictionary with 3d data\n    Returns\n        poses_set: dictionary with 3d data centred around root (center hip) joint\n        root_positions: dictionary with the original 3d position of each pose\n    \"\"\"\n    root_positions = {}\n    for k in poses_set.keys():\n        # Keep track of the global position\n        root_positions[k] = copy.deepcopy(poses_set[k][:,:3])\n        # Remove the root from the 3d position\n        poses = poses_set[k]\n        poses = poses - np.tile( poses[:,:3], [1, len(H36M_NAMES)] )\n        poses_set[k] = poses\n    return poses_set, root_positions\n\ndef postprocess_2d(poses_set):\n    \"\"\"\n    Center 2d points around root\n    \n    Args\n        poses_set: dictionary with 2d data\n    Returns\n        poses_set: dictionary with 2d data centred around root (center hip) joint\n        root_positions: dictionary with the original 2d position of each pose\n    \"\"\"\n    root_positions = {}\n    for k in poses_set.keys():\n    # Keep track of the global position\n        root_positions[k] = copy.deepcopy(poses_set[k][:,:2])\n        # Remove the root from the 3d position\n        poses = poses_set[k]\n        poses = poses - np.tile( poses[:,:2], [1, len(H36M_NAMES)] )\n        poses_set[k] = poses\n    return poses_set, root_positions\n\ndef get_all_data(data_x, \n                 data_y, \n                 camera_frame, \n                 norm_twoD=False, \n                 input_size=32,\n                 output_size=48\n                 ):\n    \"\"\"\n    Obtain numpy arrays for network inputs/outputs\n    \n    Args\n      data_x: dictionary with 2d inputs\n      data_y: dictionary with 3d expected outputs\n      camera_frame: whether the 3d data is in camera coordinates\n      input_size: input vector length for each sample\n      output_size: output vector length for each sample\n    Returns\n      encoder_inputs: numpy array for the input data \n      decoder_outputs: numpy array for the output data\n    \"\"\"\n    if norm_twoD:\n        input_size -= 2\n    # Figure out how many frames we have\n    n = 0\n    for key2d in data_x.keys():\n      n2d, _ = data_x[ key2d ].shape\n      n = n + n2d\n\n    encoder_inputs  = np.zeros((n, input_size), dtype=np.float32)\n    decoder_outputs = np.zeros((n, output_size), dtype=np.float32)\n\n    # Put all the data into big arrays\n    idx = 0\n    for key2d in data_x.keys():\n      (subj, b, fname) = key2d\n      # keys should be the same if 3d is in camera coordinates\n      key3d = key2d if (camera_frame) else (subj, b, '{0}.h5'.format(fname.split('.')[0]))\n      # '-sh' suffix means detected key-points are used \n      key3d = (subj, b, fname[:-3]) if fname.endswith('-sh') and camera_frame else key3d\n\n      n2d, _ = data_x[ key2d ].shape\n      encoder_inputs[idx:idx+n2d, :]  = data_x[ key2d ]\n      decoder_outputs[idx:idx+n2d, :] = data_y[ key3d ]\n      idx = idx + n2d\n\n    return encoder_inputs, decoder_outputs\n\ndef prepare_dataset(opt):\n    \"\"\"\n    Prepare PyTorch dataset objects used for training 2D-to-3D deep network\n    \n    Args\n        opt: experiment options\n    Returns\n        train_dataset: training dataset as PyTorch dataset object\n        eval_dataset: evaluation dataset as PyTorch dataset object \n        data_stats: dataset statistics computed from the training dataset\n        action_eval_list: a list of evaluation dataset objects where each \n        corresponds to one action\n    \"\"\"\n    # get relevant paths\n    data_dir =  opt.data_dir\n    cameras_path = os.path.join(data_dir, 'cameras.npy')\n    # By default, all actions are used\n    actions = define_actions(opt.actions)\n    # load camera parameters to project 3D skeleton\n    rcams = np.load(cameras_path, allow_pickle=True).item()    \n    # produce more camera views by adding virtual cameras if needed\n    if opt.virtual_cams:\n        rcams = add_virtual_cams(rcams)\n    # first prepare Python dictionary containing 2D and 3D data\n    data_dic, data_stats = prepare_data_dict(rcams, \n                                             opt,\n                                             predict_14=False\n                                             )\n    input_size = len(data_stats['dim_use_2d'])\n    output_size = len(data_stats['dim_use_3d'])\n    if opt.train:\n        # convert Python dictionary to numpy array\n        train_input, train_output = get_all_data(data_dic['train_set_2d'],\n                                                 data_dic['train_set_3d'], \n                                                 camera_frame,\n                                                 norm_twoD=opt.norm_twoD,\n                                                 input_size=input_size,\n                                                 output_size=output_size\n                                                 )\n        # The Numpy arrays are finally used to initialize the dataset objects\n        train_dataset = dataset.PoseDataset(train_input, \n                                            train_output, \n                                            'train',\n                                            refine_3d = opt.refine_3d\n                                            )\n    else:\n        train_dataset = None\n            \n    eval_input, eval_output = get_all_data(data_dic['test_set_2d'],\n                                           data_dic['test_set_3d'], \n                                           camera_frame,\n                                           norm_twoD=opt.norm_twoD,\n                                           input_size=input_size,\n                                           output_size=output_size\n                                           )\n\n    eval_dataset = dataset.PoseDataset(eval_input, \n                                       eval_output, \n                                       'eval',\n                                       refine_3d = opt.refine_3d\n                                       )\n    # Create a list of dataset objects for action-wise evaluation\n    action_eval_list = split_action(data_dic['test_set_2d'],\n                                    data_dic['test_set_3d'],\n                                    actions, \n                                    camera_frame,\n                                    opt,\n                                    input_size=input_size,\n                                    output_size=output_size\n                                    )\n    \n    return train_dataset, eval_dataset, data_stats, action_eval_list"
  },
  {
    "path": "libs/dataset/h36m/h36m_pose.py",
    "content": "#!/usr/bin/env python3\n# -*- coding: utf-8 -*-\n# ------------------------------------------------------------------------------\n# Copyright (c) Microsoft\n# Licensed under the MIT License.\n# Written by Bin Xiao (Bin.Xiao@microsoft.com)\n# Modified by Shichao Li (nicholas.li@connect.ust.hk)\n# ------------------------------------------------------------------------------\nimport numpy as np\nimport copy\nimport torch\nimport cv2\nimport random\n\nfrom libs.dataset.h36m.pose_dataset import JointsDataset\nfrom libs.hhr.utils.transforms import get_affine_transform\nfrom libs.hhr.utils.transforms import affine_transform\nfrom libs.hhr.utils.transforms import fliplr_joints\n\nimport logging\n\nlogger = logging.getLogger(__name__)\n\n\n## Human 3.6M dataset class\nclass H36MDataset(JointsDataset):\n    '''\n    COCO annotation:\n    \"keypoints\": {\n        0: \"nose\",\n        1: \"left_eye\",\n        2: \"right_eye\",\n        3: \"left_ear\",\n        4: \"right_ear\",\n        5: \"left_shoulder\",\n        6: \"right_shoulder\",\n        7: \"left_elbow\",\n        8: \"right_elbow\",\n        9: \"left_wrist\",\n        10: \"right_wrist\",\n        11: \"left_hip\",\n        12: \"right_hip\",\n        13: \"left_knee\",\n        14: \"right_knee\",\n        15: \"left_ankle\",\n        16: \"right_ankle\"\n    },\n\t\"skeleton\": [\n        [16,14],[14,12],[17,15],[15,13],[12,13],[6,12],[7,13], [6,7],[6,8],\n        [7,9],[8,10],[9,11],[2,3],[1,2],[1,3],[2,4],[3,5],[4,6],[5,7]]\n    H36M annotation:\n        H36M_NAMES[0]  = 'Hip'\n        H36M_NAMES[1]  = 'RHip'\n        H36M_NAMES[2]  = 'RKnee'\n        H36M_NAMES[3]  = 'RFoot'\n        H36M_NAMES[4]  = 'LHip'\n        H36M_NAMES[5]  = 'LKnee'\n        H36M_NAMES[6]  = 'LFoot'\n        H36M_NAMES[7] = 'Spine'\n        H36M_NAMES[8] = 'Thorax'\n        H36M_NAMES[9] = 'Neck/Nose'\n        H36M_NAMES[10] = 'Head'\n        H36M_NAMES[11] = 'LShoulder'\n        H36M_NAMES[12] = 'LElbow'\n        H36M_NAMES[13] = 'LWrist'\n        H36M_NAMES[14] = 'RShoulder'\n        H36M_NAMES[15] = 'RElbow'\n        H36M_NAMES[16] = 'RWrist'   \n    \"skeleton\": [\n        [0,1], [1,2], [2,3], [0,4], [4,5], [5,6], [0,7], [7,8],\n        [8,9], [9,10], [8,11], [11,12], [12,13], [8,14], [14,15], [15,16]] \n    permutation from H36M to COCO:\n    [9, 7, 8, 0, 10, 11, 14, 12, 15, 13, 16, 4, 1, 5, 2, 6, 3]\n    permutation to get back:\n        \n    '''\n\n    def __init__(self, cfg, is_train, annot_path, transform=None):\n        super().__init__(cfg, is_train, transform)\n        self.nms_thre = cfg.TEST.NMS_THRE\n        self.image_thre = cfg.TEST.IMAGE_THRE\n        self.soft_nms = cfg.TEST.SOFT_NMS\n        self.oks_thre = cfg.TEST.OKS_THRE\n        self.in_vis_thre = cfg.TEST.IN_VIS_THRE\n        self.image_width = cfg.MODEL.IMAGE_SIZE[0]\n        self.image_height = cfg.MODEL.IMAGE_SIZE[1]\n        self.loss_type = cfg.MODEL.TARGET_TYPE\n        self.aspect_ratio = self.image_width * 1.0 / self.image_height\n        self.pixel_std = 200\n        # path to pre-processed annotation\n        self.annot_path = annot_path\n        self.num_joints = 17\n        self.flip_pairs = [[5, 6], [7, 8], [9, 10], [11, 12], [13, 14], [15, 16]]\n        self.parent_ids = None\n        self.upper_body_ids = (0, 1, 2, 4, 5, 6, 7, 8, 9, 10)\n        self.lower_body_ids = (3, 11, 12, 13, 14, 15, 16)\n        self.joints_weight = np.ones((self.num_joints,1), np.float32)\n        self.joints_weight[[7,8,13,14]] = 1.2\n        self.joints_weight[[9,10,15,16]] = 1.5\n        \n        ## permute joint order for fine-tuning purpose\n        self.fine_tune_re_order = [9, 7, 8, 0, 10, 11, 14, 12, 15, 13, 16, 4, 1, 5, 2, 6, 3]\n        \n        self.ratio = float(cfg.MODEL.IMAGE_SIZE[0]/cfg.MODEL.HEATMAP_SIZE[0])\n\n        self.db = self._get_db()\n        logging.info('=> total annotation for images: {}'.format(len(self.db)))\n        \n        if is_train and cfg.DATASET.SELECT_DATA:\n            self.db = self.select_data(self.db)\n            \n        logging.info('=> load {} samples'.format(len(self.db)))\n\n    def _get_db(self):\n        gt_db = np.load(self.annot_path)\n        # permute joints\n        for record in gt_db:\n            record['p_2d'] = record['p_2d'][self.fine_tune_re_order, :]\n        return gt_db\n    \n    def get_weights(self):\n        weights = []\n        for sample_idx in range(len(self.db)):\n            path = self.db[sample_idx]['path']\n            if 'S6' in path or 'S8' in path:\n                weights.append(1.5)\n            else:\n                weights.append(1.0)\n        return weights\n    \n    def _box2cs(self, box):\n        x, y, w, h = box[:4]\n        return self._xywh2cs(x, y, w, h)\n\n    def _xywh2cs(self, x, y, w, h):\n        center = np.zeros((2), dtype=np.float32)\n        center[0] = x + w * 0.5\n        center[1] = y + h * 0.5\n\n        if w > self.aspect_ratio * h:\n            h = w * 1.0 / self.aspect_ratio\n        elif w < self.aspect_ratio * h:\n            w = h * self.aspect_ratio\n        scale = np.array(\n            [w * 1.0 / self.pixel_std, h * 1.0 / self.pixel_std],\n            dtype=np.float32)\n\n        return center, scale\n    \n    def __getitem__(self, idx):\n        db_rec = copy.deepcopy(self.db[idx])\n        image_file = db_rec['path']\n\n\n#        data_numpy = cv2.imread(\n#            image_file, cv2.IMREAD_COLOR | cv2.IMREAD_IGNORE_ORIENTATION\n#        )\n\n        # opencv 3\n        data_numpy = cv2.imread(\n            image_file, 1 | 128\n        )        \n\n        if self.color_rgb:\n            data_numpy = cv2.cvtColor(data_numpy, cv2.COLOR_BGR2RGB)\n\n        if data_numpy is None:\n            logger.error('=> fail to read {}'.format(image_file))\n            raise ValueError('Fail to read {}'.format(image_file))\n\n        joints = db_rec['p_2d']\n        joints_original = joints.copy()\n        joints_vis = np.ones(joints.shape, dtype=np.float32)\n        c, s = self._xywh2cs(0, 0, data_numpy.shape[1], data_numpy.shape[0])\n        score = 1\n        r = 0\n\n        if self.is_train:\n            # do not do half body transform since there is not so much occlusion\n            if (np.sum(joints_vis[:, 0]) > self.num_joints_half_body\n                and np.random.rand() < self.prob_half_body):\n                c_half_body, s_half_body = self.half_body_transform(\n                    joints, joints_vis\n                )\n\n                if c_half_body is not None and s_half_body is not None:\n                    c, s = c_half_body, s_half_body\n\n            sf = self.scale_factor\n            rf = self.rotation_factor\n            s = s * np.clip(np.random.randn()*sf + 1, 1 - sf, 1 + sf)\n            r = np.clip(np.random.randn()*rf, -rf*2, rf*2) \\\n                if random.random() <= 0.3 else 0\n            \n            if self.flip and random.random() <= 0.5:\n                data_numpy = data_numpy[:, ::-1, :]\n                joints, joints_vis = fliplr_joints(\n                    joints, joints_vis, data_numpy.shape[1], self.flip_pairs)\n                c[0] = data_numpy.shape[1] - c[0] - 1\n\n        trans = get_affine_transform(c, s, r, self.image_size)\n        input = cv2.warpAffine(\n            data_numpy,\n            trans,\n            (int(self.image_size[0]), int(self.image_size[1])),\n            flags=cv2.INTER_LINEAR)\n\n        if self.transform:\n            input = self.transform(input)\n\n        for i in range(self.num_joints):\n            if joints_vis[i, 0] > 0.0:\n                joints[i, 0:2] = affine_transform(joints[i, 0:2], trans)\n                # set joints to in-visible if they are out-side of the image\n                if joints[i, 0] >= self.image_width or joints[i, 1] >= self.image_height:\n                    joints_vis[i, 0] = 0.0\n        \n        \n        target, target_weight = self.generate_target(joints, joints_vis)\n\n        target = torch.from_numpy(target)\n        target_weight = torch.from_numpy(target_weight)\n\n        meta = {\n            'image': image_file,\n            'joints': joints,\n            'joints_vis': joints_vis,\n            'j_original':joints_original, #original coordinates\n            'center': c,\n            'scale': s,\n            'rotation': r,\n            'score': score,\n            'trans':trans,\n            'bbox':[0, 0, data_numpy.shape[1], data_numpy.shape[0]]\n        }\n\n        return input, target, target_weight, meta    \n    \n    def generate_target(self, joints, joints_vis):\n        '''\n        :param joints:  [num_joints, 3]\n        :param joints_vis: [num_joints, 3]\n        :return: target, target_weight(1: visible, 0: invisible)\n        '''\n        target_weight = np.ones((self.num_joints, 1), dtype=np.float32)\n        target_weight[:, 0] = joints_vis[:, 0]\n\n        assert self.target_type in ['gaussian', 'coordinate'], \\\n            'Unsupported target type'\n\n        if self.target_type == 'gaussian':\n            target = np.zeros((self.num_joints,\n                               self.heatmap_size[1],\n                               self.heatmap_size[0]),\n                              dtype=np.float32)\n\n            tmp_size = self.sigma * 3\n\n            for joint_id in range(self.num_joints):\n                feat_stride = self.image_size / self.heatmap_size\n                mu_x = int(joints[joint_id][0] / feat_stride[0] + 0.5)\n                mu_y = int(joints[joint_id][1] / feat_stride[1] + 0.5)\n                # Check that any part of the gaussian is in-bounds\n                ul = [int(mu_x - tmp_size), int(mu_y - tmp_size)]\n                br = [int(mu_x + tmp_size + 1), int(mu_y + tmp_size + 1)]\n                if ul[0] >= self.heatmap_size[0] or ul[1] >= self.heatmap_size[1] \\\n                        or br[0] < 0 or br[1] < 0:\n                    # If not, just return the image as is\n                    target_weight[joint_id] = 0\n                    continue\n\n                # # Generate gaussian\n                size = 2 * tmp_size + 1\n                x = np.arange(0, size, 1, np.float32)\n                y = x[:, np.newaxis]\n                x0 = y0 = size // 2\n                # The gaussian is not normalized, we want the center value to equal 1\n                g = np.exp(- ((x - x0) ** 2 + (y - y0) ** 2) / (2 * self.sigma ** 2))\n\n                # Usable gaussian range\n                g_x = max(0, -ul[0]), min(br[0], self.heatmap_size[0]) - ul[0]\n                g_y = max(0, -ul[1]), min(br[1], self.heatmap_size[1]) - ul[1]\n                # Image range\n                img_x = max(0, ul[0]), min(br[0], self.heatmap_size[0])\n                img_y = max(0, ul[1]), min(br[1], self.heatmap_size[1])\n\n                v = target_weight[joint_id]\n                if v > 0.5:\n                    target[joint_id][img_y[0]:img_y[1], img_x[0]:img_x[1]] = \\\n                        g[g_y[0]:g_y[1], g_x[0]:g_x[1]]\n        elif self.target_type == 'coordinate':\n            target = joints/self.ratio  \n        if self.use_different_joints_weight:\n            target_weight = np.multiply(target_weight, self.joints_weight)\n\n        return target, target_weight    "
  },
  {
    "path": "libs/dataset/h36m/pose_dataset.py",
    "content": "# ------------------------------------------------------------------------------\n# Copyright (c) Microsoft\n# Licensed under the MIT License.\n# Written by Bin Xiao (Bin.Xiao@microsoft.com)\n# Modified by Shichao Li (nicholas.li@connect.ust.hk)\n# ------------------------------------------------------------------------------\n\nfrom __future__ import absolute_import\nfrom __future__ import division\nfrom __future__ import print_function\n\nimport copy\nimport logging\nimport random\n\nimport cv2\nimport numpy as np\nimport torch\nfrom torch.utils.data import Dataset\n\nfrom libs.hhr.utils.transforms import get_affine_transform\nfrom libs.hhr.utils.transforms import affine_transform\nfrom libs.hhr.utils.transforms import fliplr_joints\n\n\nlogger = logging.getLogger(__name__)\n\n\nclass JointsDataset(Dataset):\n    def __init__(self, cfg, is_train, root=None, image_set=None, transform=None):\n        self.num_joints = 0\n        self.pixel_std = 200\n        self.flip_pairs = []\n        self.parent_ids = []\n\n        self.is_train = is_train\n        self.root = root\n        self.image_set = image_set\n\n        self.output_path = cfg.OUTPUT_DIR\n        self.data_format = cfg.DATASET.DATA_FORMAT\n\n        self.scale_factor = cfg.DATASET.SCALE_FACTOR\n        self.rotation_factor = cfg.DATASET.ROT_FACTOR\n        self.flip = cfg.DATASET.FLIP\n        self.num_joints_half_body = cfg.DATASET.NUM_JOINTS_HALF_BODY\n        self.prob_half_body = cfg.DATASET.PROB_HALF_BODY\n        self.color_rgb = cfg.DATASET.COLOR_RGB\n\n        self.target_type = cfg.MODEL.TARGET_TYPE\n        self.image_size = np.array(cfg.MODEL.IMAGE_SIZE)\n        self.heatmap_size = np.array(cfg.MODEL.HEATMAP_SIZE)\n        self.sigma = cfg.MODEL.SIGMA\n        self.use_different_joints_weight = cfg.LOSS.USE_DIFFERENT_JOINTS_WEIGHT\n        self.joints_weight = 1\n\n        self.transform = transform\n        self.db = []\n\n    def _get_db(self):\n        raise NotImplementedError\n\n    def evaluate(self, cfg, preds, output_dir, *args, **kwargs):\n        raise NotImplementedError\n\n    def half_body_transform(self, joints, joints_vis):\n        upper_joints = []\n        lower_joints = []\n        for joint_id in range(self.num_joints):\n            if joints_vis[joint_id][0] > 0:\n                if joint_id in self.upper_body_ids:\n                    upper_joints.append(joints[joint_id])\n                else:\n                    lower_joints.append(joints[joint_id])\n\n        if np.random.randn() < 0.5 and len(upper_joints) > 2:\n            selected_joints = upper_joints\n        else:\n            selected_joints = lower_joints \\\n                if len(lower_joints) > 2 else upper_joints\n\n        if len(selected_joints) < 2:\n            return None, None\n\n        selected_joints = np.array(selected_joints, dtype=np.float32)\n        center = selected_joints.mean(axis=0)[:2]\n\n        left_top = np.amin(selected_joints, axis=0)\n        right_bottom = np.amax(selected_joints, axis=0)\n\n        w = right_bottom[0] - left_top[0]\n        h = right_bottom[1] - left_top[1]\n\n        if w > self.aspect_ratio * h:\n            h = w * 1.0 / self.aspect_ratio\n        elif w < self.aspect_ratio * h:\n            w = h * self.aspect_ratio\n\n        scale = np.array(\n            [\n                w * 1.0 / self.pixel_std,\n                h * 1.0 / self.pixel_std\n            ],\n            dtype=np.float32\n        )\n\n        scale = scale * 1.5\n\n        return center, scale\n\n    def __len__(self,):\n        return len(self.db)\n\n    def __getitem__(self, idx):\n        db_rec = copy.deepcopy(self.db[idx])\n\n        image_file = db_rec['image']\n        filename = db_rec['filename'] if 'filename' in db_rec else ''\n        imgnum = db_rec['imgnum'] if 'imgnum' in db_rec else ''\n\n        if self.data_format == 'zip':\n            from utils import zipreader\n            data_numpy = zipreader.imread(\n                image_file, cv2.IMREAD_COLOR | cv2.IMREAD_IGNORE_ORIENTATION\n            )\n        else:\n            data_numpy = cv2.imread(\n                image_file, cv2.IMREAD_COLOR | cv2.IMREAD_IGNORE_ORIENTATION\n            )\n\n        if self.color_rgb:\n            data_numpy = cv2.cvtColor(data_numpy, cv2.COLOR_BGR2RGB)\n\n        if data_numpy is None:\n            logger.error('=> fail to read {}'.format(image_file))\n            raise ValueError('Fail to read {}'.format(image_file))\n\n        joints = db_rec['joints_3d']\n        joints_vis = db_rec['joints_3d_vis']\n\n        c = db_rec['center']\n        s = db_rec['scale']\n        score = db_rec['score'] if 'score' in db_rec else 1\n        r = 0\n\n        if self.is_train:\n            if (np.sum(joints_vis[:, 0]) > self.num_joints_half_body\n                and np.random.rand() < self.prob_half_body):\n                c_half_body, s_half_body = self.half_body_transform(\n                    joints, joints_vis\n                )\n\n                if c_half_body is not None and s_half_body is not None:\n                    c, s = c_half_body, s_half_body\n\n            sf = self.scale_factor\n            rf = self.rotation_factor\n            s = s * np.clip(np.random.randn()*sf + 1, 1 - sf, 1 + sf)\n            r = np.clip(np.random.randn()*rf, -rf*2, rf*2) \\\n                if random.random() <= 0.6 else 0\n\n            if self.flip and random.random() <= 0.5:\n                data_numpy = data_numpy[:, ::-1, :]\n                joints, joints_vis = fliplr_joints(\n                    joints, joints_vis, data_numpy.shape[1], self.flip_pairs)\n                c[0] = data_numpy.shape[1] - c[0] - 1\n\n        trans = get_affine_transform(c, s, r, self.image_size)\n        input = cv2.warpAffine(\n            data_numpy,\n            trans,\n            (int(self.image_size[0]), int(self.image_size[1])),\n            flags=cv2.INTER_LINEAR)\n\n        if self.transform:\n            input = self.transform(input)\n\n        for i in range(self.num_joints):\n            if joints_vis[i, 0] > 0.0:\n                joints[i, 0:2] = affine_transform(joints[i, 0:2], trans)\n\n        target, target_weight = self.generate_target(joints, joints_vis)\n\n        target = torch.from_numpy(target)\n        target_weight = torch.from_numpy(target_weight)\n\n        meta = {\n            'image': image_file,\n            'filename': filename,\n            'imgnum': imgnum,\n            'joints': joints,\n            'joints_vis': joints_vis,\n            'center': c,\n            'scale': s,\n            'rotation': r,\n            'score': score\n        }\n\n        return input, target, target_weight, meta\n\n    def select_data(self, db):\n        db_selected = []\n        for rec in db:\n            num_vis = 0\n            joints_x = 0.0\n            joints_y = 0.0\n            for joint, joint_vis in zip(\n                    rec['joints_3d'], rec['joints_3d_vis']):\n                if joint_vis[0] <= 0:\n                    continue\n                num_vis += 1\n\n                joints_x += joint[0]\n                joints_y += joint[1]\n            if num_vis == 0:\n                continue\n\n            joints_x, joints_y = joints_x / num_vis, joints_y / num_vis\n\n            area = rec['scale'][0] * rec['scale'][1] * (self.pixel_std**2)\n            joints_center = np.array([joints_x, joints_y])\n            bbox_center = np.array(rec['center'])\n            diff_norm2 = np.linalg.norm((joints_center-bbox_center), 2)\n            ks = np.exp(-1.0*(diff_norm2**2) / ((0.2)**2*2.0*area))\n\n            metric = (0.2 / 16) * num_vis + 0.45 - 0.2 / 16\n            if ks > metric:\n                db_selected.append(rec)\n\n        logger.info('=> num db: {}'.format(len(db)))\n        logger.info('=> num selected db: {}'.format(len(db_selected)))\n        return db_selected\n\n    def generate_target(self, joints, joints_vis):\n        '''\n        :param joints:  [num_joints, 3]\n        :param joints_vis: [num_joints, 3]\n        :return: target, target_weight(1: visible, 0: invisible)\n        '''\n        target_weight = np.ones((self.num_joints, 1), dtype=np.float32)\n        target_weight[:, 0] = joints_vis[:, 0]\n\n        assert self.target_type == 'gaussian', \\\n            'Only support gaussian map now!'\n\n        if self.target_type == 'gaussian':\n            target = np.zeros((self.num_joints,\n                               self.heatmap_size[1],\n                               self.heatmap_size[0]),\n                              dtype=np.float32)\n\n            tmp_size = self.sigma * 3\n\n            for joint_id in range(self.num_joints):\n                feat_stride = self.image_size / self.heatmap_size\n                mu_x = int(joints[joint_id][0] / feat_stride[0] + 0.5)\n                mu_y = int(joints[joint_id][1] / feat_stride[1] + 0.5)\n                # Check that any part of the gaussian is in-bounds\n                ul = [int(mu_x - tmp_size), int(mu_y - tmp_size)]\n                br = [int(mu_x + tmp_size + 1), int(mu_y + tmp_size + 1)]\n                if ul[0] >= self.heatmap_size[0] or ul[1] >= self.heatmap_size[1] \\\n                        or br[0] < 0 or br[1] < 0:\n                    # If not, just return the image as is\n                    target_weight[joint_id] = 0\n                    continue\n\n                # # Generate gaussian\n                size = 2 * tmp_size + 1\n                x = np.arange(0, size, 1, np.float32)\n                y = x[:, np.newaxis]\n                x0 = y0 = size // 2\n                # The gaussian is not normalized, we want the center value to equal 1\n                g = np.exp(- ((x - x0) ** 2 + (y - y0) ** 2) / (2 * self.sigma ** 2))\n\n                # Usable gaussian range\n                g_x = max(0, -ul[0]), min(br[0], self.heatmap_size[0]) - ul[0]\n                g_y = max(0, -ul[1]), min(br[1], self.heatmap_size[1]) - ul[1]\n                # Image range\n                img_x = max(0, ul[0]), min(br[0], self.heatmap_size[0])\n                img_y = max(0, ul[1]), min(br[1], self.heatmap_size[1])\n\n                v = target_weight[joint_id]\n                if v > 0.5:\n                    target[joint_id][img_y[0]:img_y[1], img_x[0]:img_x[1]] = \\\n                        g[g_y[0]:g_y[1], g_x[0]:g_x[1]]\n\n        if self.use_different_joints_weight:\n            target_weight = np.multiply(target_weight, self.joints_weight)\n\n        return target, target_weight\n"
  },
  {
    "path": "libs/dataset/h36m/pth_dataset.py",
    "content": "import logging\n\nimport torch\nimport torch.utils.data\nimport torch.nn.functional as F\nimport numpy as np\n\ndef normalize(vec):\n    # normalize a numpy vector\n    return (vec-vec.mean())/vec.std()\n\ndef unNormalizeData(normalized_data, data_mean, data_std, dimensions_to_ignore):\n    \"\"\"\n    Un-normalizes a matrix whose mean has been substracted and that has been divided by\n    standard deviation. Some dimensions might also be missing\n    \n    Args\n    normalized_data: nxd matrix to unnormalize\n    data_mean: np vector with the mean of the data\n    data_std: np vector with the standard deviation of the data\n    dimensions_to_ignore: list of dimensions that were removed from the original data\n    Returns\n    orig_data: the input normalized_data, but unnormalized\n    \"\"\"\n    T = normalized_data.shape[0] # Batch size\n    D = data_mean.shape[0] # Dimensionality\n    \n    orig_data = np.zeros((T, D), dtype=np.float32)\n    dimensions_to_use = np.array([dim for dim in range(D)\n                                    if dim not in dimensions_to_ignore])\n    \n    orig_data[:, dimensions_to_use] = normalized_data\n    \n    # Multiply times stdev and add the mean\n    stdMat = data_std.reshape((1, D))\n    stdMat = np.repeat(stdMat, T, axis=0)\n    meanMat = data_mean.reshape((1, D))\n    meanMat = np.repeat(meanMat, T, axis=0)\n    orig_data = np.multiply(orig_data, stdMat) + meanMat\n    return orig_data\n\nclass PoseDataset(torch.utils.data.Dataset):\n    def __init__(self, array_2d, array_3d, split, action_name=None, refine_3d=False):\n        \"\"\"\n        Args:\n            \n        \"\"\"\n        self.data_2d = array_2d\n        self.data_3d = array_3d\n        self.num_samples = len(self.data_2d)\n        self.split = split\n        self.action_name = action_name\n        self.refine_3d = refine_3d\n        self.stage_idx = 1\n        # initialize current estimate 3d pose\n        self.current_estimate = np.zeros(self.data_3d.shape, dtype=np.float32)\n        # initialize the regression target (starts with zero estimate)\n        self.regression_target = self.data_3d.copy()\n        assert len(self.data_2d) == len(self.data_3d)\n        \n    def __len__(self):\n        return self.num_samples\n\n    def __getitem__(self, idx):\n        if self.refine_3d and self.stage_idx > 1:\n            #return np.concatenate([self.data_2d[idx], self.current_estimate[idx]]), self.regression_target[idx]\n            # normalized version\n            return np.concatenate([self.data_2d[idx], normalize(self.current_estimate[idx])]), self.regression_target[idx]\n        else:\n            return self.data_2d[idx], self.regression_target[idx]\n    \n    def set_stage(self, stage_idx):\n        self.stage_idx = stage_idx\n        return\n    \n    def stage_update(self, model, stats, opt, verbose=False):\n        # update the dataset for cascaded regression\n        model.eval()\n        eval_loader = torch.utils.data.DataLoader(self, \n                                                  batch_size = opt.batch_size, \n                                                  shuffle = False, \n                                                  num_workers = opt.num_threads) \n        # vector to add at last\n        update_vector = []\n        total_loss = 0\n        all_distance = np.zeros((0))\n        for batch_idx, batch in enumerate(eval_loader):\n            data = batch[0]\n            target = batch[1]\n            if opt.cuda:\n                with torch.no_grad():\n                    # move to GPU\n                    data, target = data.cuda(), target.cuda()\n            # forward pass to get prediction\n            prediction = model(data)\n            # mean squared loss \n            loss = F.mse_loss(prediction, target, reduction='sum')\n            total_loss += loss.data.item()\n            # compute distance of body joints in un-normalized format\n            unnorm_target = unNormalizeData(target.data.cpu().numpy(), \n            stats['mean_3d'], stats['std_3d'], stats['dim_ignore_3d']) \n            # put the prediction into the update list\n            prediction = prediction.data.cpu().numpy()\n            update_vector.append(prediction)\n            unnorm_pred = unNormalizeData(prediction, \n            stats['mean_3d'], stats['std_3d'], stats['dim_ignore_3d'])\n            # pick the joints that are used\n            dim_use = stats['dim_use_3d']\n            unnorm_target_use = unnorm_target[:, dim_use]\n            unnorm_target_use = unnorm_target_use.reshape(-1,16,3)\n            unnorm_pred_use = unnorm_pred[:, dim_use]\n            unnorm_pred_use = unnorm_pred_use.reshape(-1,16,3)\n            distance = np.sum((unnorm_target_use - unnorm_pred_use)**2, axis=2)\n            distance = np.mean(np.sqrt(distance), axis=1)\n            all_distance = np.hstack([all_distance, distance])      \n        # update the current estimate and regression target\n        update_vector = np.concatenate(update_vector, axis=0)\n        self.current_estimate += update_vector\n        self.regression_target -= update_vector\n        # report statistics\n        avg_loss = total_loss/(self.num_samples*16*3)\n        avg_distance = all_distance.mean()\n        if verbose:\n            logging.info('Stage update finished.')\n            logging.info('{:s} set: average loss: {:.4f} '.format(self.split, avg_loss))\n            logging.info('{:s} set: average joint distance: {:.4f} '.format(self.split, avg_distance))\n        return avg_loss, avg_distance"
  },
  {
    "path": "libs/evolution/__init__.py",
    "content": "#!/usr/bin/env python3\n# -*- coding: utf-8 -*-\n\"\"\"\nEmpty file.\n\"\"\"\n\n\n"
  },
  {
    "path": "libs/evolution/genetic.py",
    "content": "\"\"\"\nUtility functions for genetic evolution.\n\"\"\"\nimport libs.dataset.h36m.cameras as cameras\nfrom libs.skeleton.anglelimits import \\\nto_local, to_global, get_skeleton, to_spherical, \\\nnt_parent_indices, nt_child_indices, \\\nis_valid_local, is_valid\n\nimport matplotlib.pyplot as plt\nimport os\nimport logging\n\nimport numpy as np\nimport torch\nfrom mpl_toolkits.mplot3d import Axes3D\nfrom tqdm import tqdm\nfrom scipy.spatial.transform import Rotation as R\n\nroot = \"../resources/constraints\"\n# Joints in H3.6M -- data has 32 joints, but only 17 that move\nH36M_NAMES = ['']*32\nH36M_NAMES[0]  = 'Hip'\nH36M_NAMES[1]  = 'RHip'\nH36M_NAMES[2]  = 'RKnee'\nH36M_NAMES[3]  = 'RFoot'\nH36M_NAMES[6]  = 'LHip'\nH36M_NAMES[7]  = 'LKnee'\nH36M_NAMES[8]  = 'LFoot'\nH36M_NAMES[12] = 'Spine'\nH36M_NAMES[13] = 'Thorax'\nH36M_NAMES[14] = 'Neck/Nose'\nH36M_NAMES[15] = 'Head'\nH36M_NAMES[17] = 'LShoulder'\nH36M_NAMES[18] = 'LElbow'\nH36M_NAMES[19] = 'LWrist'\nH36M_NAMES[25] = 'RShoulder'\nH36M_NAMES[26] = 'RElbow'\nH36M_NAMES[27] = 'RWrist'\ntotal_joints_num = len(H36M_NAMES)\n\n# this dictionary stores the parent indice for each joint\n# key:value -> child joint index:its parent joint index\nparent_idx = {1:0, 2:1, 3:2, 6:0, 7:6, 8:7, 12:0, 13:12, 14:13, 15:14, 17:13, \n              18:17, 19:18, 25:13, 26:25, 27:26\n              }\n\n# this dictionary stores the children indices for each parent joint\n# key:value -> parent index: joint indices for its children as a list  \nchildren_idx = {\n        0: [1, 6],\n        1: [2], 2: [3],\n        6: [7], 7: [8],\n        13: [14, 17, 25],\n        14: [15], 17: [18], 18:[19],\n        25: [26], 26:[27]\n        }\n\n# used roots for random selection\nroot_joints = [0, 1, 2, 6, 7, 13, 17, 18, 25, 26]\n# names of the bone vectors attached on the human torso\nbone_name = {\n 1: 'thorax to head top',\n 2: 'left shoulder to left elbow',\n 3: 'left elbow to left wrist',\n 4: 'right shoulder to right elbow',\n 5: 'right elbow to right wrist',\n 6: 'left hip to left knee',\n 7: 'left knee to left ankle',\n 8: 'right hip to right knee',\n 9: 'right knee to right ankle'\n}    \n# this dictionary stores the sub-tree rooted at each root joint\n# key:value->root joint index:list of bone vector indices \nbone_indices = {0: [5, 6, 7, 8],\n                1: [7, 8],\n                2: [8],\n                6: [5, 6],\n                7: [6],\n                13: [1, 2, 3, 4], # thorax\n                17: [1, 2],\n                18: [2],\n                25: [3, 4],\n                26: [4]\n                }\n\n# load template bone lengths that can be used during mutation\n# you can prepare your own bone length templates to represent \n# subjects with different size\nbl_templates = np.load(os.path.join(root, \"bones.npy\"), allow_pickle=True)\n\n# pre-compute the sub-tree joint indices for each joint\nsubtree_indices = {}\ndef get_subtree(joint_idx, children_idx):\n    if joint_idx not in children_idx:\n        return None\n    subtree = set()\n    for child_idx in children_idx[joint_idx]:\n        subtree.add(child_idx)\n        offsprings = get_subtree(child_idx, children_idx)\n        if offsprings is not None:\n            subtree = subtree.union(offsprings)\n    return subtree\nfor joint_idx in range(total_joints_num):\n    if H36M_NAMES[joint_idx] != '':\n        subtree_indices[joint_idx] = get_subtree(joint_idx, children_idx)\n\ndef swap_bones(bones_father, bones_mother, root_idx):\n    swap_indices = bone_indices[root_idx]\n    temp = bones_father.copy()\n    bones_father[swap_indices] = bones_mother[swap_indices].copy()\n    bones_mother[swap_indices] = temp[swap_indices].copy()\n    del temp\n    return bones_father, bones_mother, swap_indices\n\ndef get_bone_length(skeleton):\n    \"\"\"\n    Compute limb length for a given skeleton.\n    \"\"\"\n    bones = skeleton[nt_parent_indices, :] - skeleton[nt_child_indices, :]  \n    bone_lengths = to_spherical(bones)[:, 0]    \n    return bone_lengths\n\ndef get_h36m_bone_length(visualize=True):\n    #1: F 5:F 6:M 7:F 8:M 9:M 11:M\n    bl_dic = {1:[], 5:[], 6:[], 7:[], 8:[], 9:[], 11:[]}\n    train_dic = np.load('../data/human3.6M/h36m/numpy/threeDPose_train.npy').item()\n    test_dic = np.load('../data/human3.6M/h36m/numpy/threeDPose_test.npy').item()\n    def process_dic(data_dic, bl_dic, candicate=50):\n        for key in data_dic.keys():\n            subject = key[0]\n            indices = np.random.choice(len(data_dic[key]), candicate, replace=False)\n            selected = data_dic[key][indices]\n            for pose in selected:\n                bl_dic[subject].append(get_bone_length(pose.reshape(32,3)))\n        return\n    process_dic(train_dic, bl_dic)\n    process_dic(test_dic, bl_dic)\n    for key in bl_dic:\n        for array in bl_dic[key]:\n            array = array.reshape(9,1)\n        bl_dic[key] = np.vstack(bl_dic[key])    \n    if visualize:\n    # as can be observed, only bone length of idx 0 vary a lot. Others are almost fixed        \n#        for key in bl_dic.keys():    \n#            fig, axes = plt.subplots(3,3)\n#            plt.title('Subject: '+str(key))\n#            for row in range(3):\n#                for col in range(3):\n#                    axes[row][col].hist(bl_dic[key][:,3*row + col], bins=20)\n        fig, axes = plt.subplots(3,3)\n        all_lengths = np.vstack(list(bl_dic.values()))\n        for row in range(3):\n            for col in range(3):\n                axes[row][col].hist(all_lengths[:,3*row + col])        \n    return bl_dic\n\ndef get_random_rotation(sigma=60.):\n    angle = np.random.normal(scale=sigma)\n    axis_idx = np.random.choice(3, 1)\n    if axis_idx == 0:\n        r = R.from_euler('xyz', [angle, 0., 0.], degrees=True)\n    elif axis_idx == 1:\n        r = R.from_euler('xyz', [0., angle, 0.], degrees=True)\n    else:\n        r = R.from_euler('xyz', [0., 0., angle], degrees=True)    \n    return r\n\ndef rotate_bone_random(bone, sigma=10.):\n    r = get_random_rotation(sigma)\n    bone_rot = r.as_dcm() @ bone.reshape(3,1)\n    return bone_rot.reshape(3)\n\ndef rotate_pose_random(pose=None, sigma=60.):\n    # pose shape: [n_joints, 3]\n    if pose is None:\n        result = None\n    else:\n        r = get_random_rotation()\n        pose = pose.reshape(32, 3)\n        # rotate around hip\n        hip = pose[0].reshape(1, 3)\n        relative_pose = pose - hip\n        rotated = r.as_dcm() @ relative_pose.T\n        result = rotated.T + hip\n    return result\n\ndef re_order(skeleton):\n    # the ordering of coordinate used by the Prior was x,z and y\n    return skeleton[:, [0,2,1]]\n\ndef set_z(pose, target):\n    if pose is None:\n        return None\n    original_shape = pose.shape\n    pose = pose.reshape(32, 3)\n    min_val = pose[:, 2].min()\n    pose[:, 2] -= min_val - target\n    return pose.reshape(original_shape)\n    \ndef modify_pose(skeleton, local_bones, bone_length, ro=False):\n    # get a new pose by modify an existing pose with input local bone vectors\n    # and bone lengths\n    new_bones = to_global(skeleton, local_bones)['bg']\n    new_pose = get_skeleton(new_bones, skeleton, bone_length=bone_length)\n    if ro:\n        new_pose = re_order(new_pose)\n    return new_pose.reshape(-1)\n\ndef exploration(father, mother, opt, post_processing=True):\n    \"\"\"\n    Produce novel data by exploring the data space with evolutionary operators.\n    cross over operator in the local coordinate system\n    mutation: perturb the local joint angle\n    \"\"\"\n    # get local coordinate for each bone vector\n    father = re_order(father.reshape(total_joints_num, -1))\n    father_bone_length = get_bone_length(father)\n    mother = re_order(mother.reshape(total_joints_num, -1))\n    mother_bone_length = get_bone_length(mother)\n    bones_father = to_local(father)\n    bones_mother = to_local(mother)    \n    if opt.CV:\n        # crossover: exchange random sub-trees of two kinematic trees\n        root_idx = np.random.randint(0, len(root_joints))\n        root_selected = root_joints[root_idx]\n        bones_father, bones_mother, indices = swap_bones(bones_father, \n                                                         bones_mother, \n                                                         root_selected)\n    if opt.M:\n        # local mutation: apply random rotation to local limb\n        for bone_idx in indices:\n            if np.random.rand() <= opt.MRL:\n                bones_father[bone_idx] = rotate_bone_random(bones_father[bone_idx], sigma=opt.SDL)\n                bones_mother[bone_idx] = rotate_bone_random(bones_mother[bone_idx], sigma=opt.SDL)\n                \n    son_pose, daughter_pose = None, None\n    if opt.C:\n        # apply joint angle constraint as the fitness function\n        valid_vec_fa = is_valid_local(bones_father)\n        valid_vec_mo = is_valid_local(bones_mother)\n    if not opt.C or valid_vec_fa.sum() >= opt.Th:\n        son_pose = modify_pose(father, bones_father, mother_bone_length, ro=True)\n    if not opt.C or valid_vec_mo.sum() >= opt.Th:\n        daughter_pose = modify_pose(mother, bones_mother, father_bone_length, ro=True)\n    if opt.M:\n        # global mutation: rotate the whole 3D skeleton\n        if np.random.rand() <= opt.MRG:\n            son_pose = rotate_pose_random(son_pose, sigma=opt.SDG)\n        if np.random.rand() <= opt.MRG:\n            daughter_pose = rotate_pose_random(daughter_pose, sigma=opt.SDG)\n    if post_processing:\n        # move the poses to the ground plane\n        set_z(son_pose, np.random.normal(loc=20.0, scale=3.0))\n        set_z(daughter_pose, np.random.normal(loc=20.0, scale=3.0))\n    if opt.DE:\n        valid_vec_fa = is_valid_local(bones_father)\n        valid_vec_mo = is_valid_local(bones_mother)        \n        # re_order: order back to x, y, z\n        son_pose = modify_pose(father, bones_father, mother_bone_length, ro=True)\n        daughter_pose = modify_pose(mother, bones_mother, father_bone_length, ro=True)\n#        valid_vec_son = is_valid(son_pose)\n#        valid_vec_dau = is_valid(daughter_pose)             \n    if opt.DE and opt.V:\n        plt.figure()\n        ax1 = plt.subplot(1,4,1, projection='3d')\n        plt.title('father')\n        show3Dpose(re_order(father), ax1, add_labels=False, plot_dot=True)\n        ax2 = plt.subplot(1,4,2, projection='3d')\n        plt.title('mother')\n        show3Dpose(re_order(mother), ax2, add_labels=False, plot_dot=True)\n        ax3 = plt.subplot(1,4,3, projection='3d')\n        plt.title('son: ' + str(valid_vec_fa.sum()))\n        show3Dpose(son_pose, ax3, add_labels=False, plot_dot=True)\n        ax4 = plt.subplot(1,4,4, projection='3d')\n        plt.title('daughter: ' + str(valid_vec_mo.sum()))\n        show3Dpose(daughter_pose, ax4, add_labels=False, plot_dot=True)\n        plt.tight_layout() \n    return son_pose, daughter_pose\n\ndef show3Dpose(channels, \n               ax, \n               lcolor=\"#3498db\", \n               rcolor=\"#e74c3c\", \n               add_labels=True,\n               gt=False,\n               pred=False,\n               plot_dot=False\n               ): # blue, orange\n  \"\"\"\n  Visualize a 3d skeleton\n\n  Args\n    channels: 96x1 vector. The pose to plot.\n    ax: matplotlib 3d axis to draw on\n    lcolor: color for left part of the body\n    rcolor: color for right part of the body\n    add_labels: whether to add coordinate labels\n  Returns\n    Nothing. Draws on ax.\n  \"\"\"\n\n  if channels.shape[0] == 96:\n      vals = np.reshape( channels, (32, -1) )\n  else:\n      vals = channels\n  I   = np.array([1,2,3,1,7,8,1, 13,14,15,14,18,19,14,26,27])-1 # start points\n  J   = np.array([2,3,4,7,8,9,13,14,15,16,18,19,20,26,27,28])-1 # end points\n  LR  = np.array([1,1,1,0,0,0,0, 0, 0, 0, 0, 0, 0, 1, 1, 1], dtype=bool)\n  dim_use_3d = [3, 4, 5, 6, 7, 8, 9, 10, 11, 18, 19, 20, 21, 22, 23, 24, 25, \n                26, 36, 37, 38, 39, 40, 41, 42, 43, 44, 45, 46, 47, 51, 52, 53, \n                54, 55, 56, 57, 58, 59, 75, 76, 77, 78, 79, 80, 81, 82, 83]\n  # Make connection matrix\n  for i in np.arange( len(I) ):\n    x, y, z = [np.array( [vals[I[i], j], vals[J[i], j]] ) for j in range(3)]\n    if gt:\n        ax.plot(x,y, z,  lw=4, c='k')\n#        ax.plot(x,y, z,  lw=2, c='k')\n        \n    elif pred:\n        ax.plot(x,z, -y,  lw=4, c='r')\n#        ax.plot(x,y, z,  lw=2, c='r')\n\n    else:\n#        ax.plot(x,z, -y,  lw=2, c=lcolor if LR[i] else rcolor)\n        ax.plot(x,y, z,  lw=4, c=lcolor if LR[i] else rcolor)\n  if plot_dot:\n      joints = channels.reshape(96)\n      joints = joints[dim_use_3d].reshape(16,3)\n      ax.scatter(joints[:,0], joints[:,1], joints[:,2], c='k', marker='o')\n  RADIUS = 750 # space around the subject\n  xroot, yroot, zroot = vals[0,0], vals[0,1], vals[0,2]\n  ax.set_xlim3d([-RADIUS+xroot, RADIUS+xroot])\n  ax.set_zlim3d([-RADIUS+zroot, RADIUS+zroot])\n  ax.set_ylim3d([-RADIUS+yroot, RADIUS+yroot])\n\n  if add_labels:\n    ax.set_xlabel(\"x\")\n    ax.set_ylabel(\"y\")\n    ax.set_zlabel(\"z\")\n\n  ax.set_aspect('equal')\n#  ax.set_xticks([])\n#  ax.set_yticks([])\n#  ax.set_zticks([])\n\n#  ax.get_xaxis().set_ticklabels([])\n#  ax.get_yaxis().set_ticklabels([])\n#  ax.set_zticklabels([])\n  # Get rid of the panes (actually, make them white)\n#  white = (1.0, 1.0, 1.0, 0.0)\n#  ax.w_xaxis.set_pane_color(white)\n#  ax.w_yaxis.set_pane_color(white)\n  # Keep z pane\n\n  # Get rid of the lines in 3d\n#  ax.w_xaxis.line.set_color(white)\n#  ax.w_yaxis.line.set_color(white)\n#  ax.w_zaxis.line.set_color(white)\n  ax.view_init(10, -60)\n\ndef choose_best(population, fraction = 0.02, method='random'):\n    \"\"\"\n    Choose the best candidates to produce descendents.\n    \"\"\"\n    if method == 'random':\n        # this is a simple implementation by random sampling\n        num_total = len(population)\n        num_to_choose = int(fraction*num_total)\n        chosen_indices = np.random.choice(num_total, num_to_choose*2, replace=False)\n        father_indices = chosen_indices[:num_to_choose]\n        mother_indices = chosen_indices[num_to_choose:]\n    else:\n        raise NotImplementedError\n    return father_indices, mother_indices\n\ndef project_to_cameras(poses, cams):\n    \"\"\"\n    Project 3d poses using camera parameters\n    input: \n        3D poses: [n_pose, pose length]\n        cams: list of camera parameters\n    return: list of 2D projections for each camera\n    cams:\n    \"\"\"\n    p_2d = []\n    for cam in cams:\n        R, T, f, c, k, p, name = cam\n        pts2d, _, _, _, _ = cameras.project_point_radial(np.reshape(poses, [-1, 3]), R, T, f, c, k, p )\n        p_2d.append(np.reshape( pts2d, [-1, len(H36M_NAMES)*2]))\n    return p_2d\n\ndef transform_world_to_camera(poses, cams):\n    \"\"\"\n    Project 3d poses from world coordinate to camera coordinate system\n    return: list of 3D poses in camera coordinate systems\n    \"\"\"\n    p_3d_cam = []\n    for cam in cams:\n        R, T, f, c, k, p, name = cam\n        camera_coord = cameras.world_to_camera_frame( np.reshape(poses, [-1, 3]), R, T)\n        camera_coord = np.reshape( camera_coord, [-1, len(H36M_NAMES)*3] )\n        p_3d_cam.append(camera_coord)\n    return p_3d_cam\n\ndef normalize(data, mean=None, std=None):\n    if mean is not None and std is not None:\n        pass\n    elif mean is None and std is None:\n        mean = np.mean(data, axis=0).reshape(1, data.shape[1])\n        std = np.std(data, axis=0).reshape(1, data.shape[1])\n    else:\n        raise ValueError\n    return (data-mean)/std\n\ndef unnormalize(data, mean, std):\n    return (data*std) + mean\n\ndef postprocess_3d( poses):\n    return poses - np.tile( poses[:,:3], [1, len(H36M_NAMES)] ) \n\ndef calc_errors(pred_poses, gt_poses, protocol='mpjpe'):\n    # error after a regid alignment, corresponding to protocol #2 in the paper\n    # Compute Euclidean distance error per joint\n    sqerr = (pred_poses - gt_poses)**2 # Squared error between prediction and expected output\n    sqerr = sqerr.reshape(len(sqerr), -1, 3)\n    sqerr = np.sqrt(sqerr.sum(axis=2))\n    if protocol == 'mpjpe':\n        ret = sqerr.mean(axis=1)\n        ret = ret.reshape(len(ret), 1)\n    else:\n        raise NotImplementedError\n    return ret\n\ndef to_numpy(tensor):\n    return tensor.data.cpu().numpy()\n\ndef get_prediction(cascade, data):\n    data = torch.from_numpy(data.astype(np.float32))\n    if torch.cuda.is_available():\n        data = data.cuda()\n    # forward pass to get prediction for the first stage\n    prediction = cascade[0](data)\n    # prediction for later stages\n    for stage_idx in range(1, len(cascade)):\n        prediction += cascade[stage_idx](data)    \n    return prediction\n\ndef get_score(model_file, candidates):\n    \"\"\"\n    Obtain model inference errors for the candidates.\n    \"\"\"\n    cams = model_file['cams']\n    stats = model_file['stats']\n    model = model_file['model']\n    if torch.cuda.is_available():\n        model = model.cuda()\n    # project to 2D keypoints\n    p_2d = project_to_cameras(candidates, cams)\n    # convert to camera coordinate\n    p_3d_cam = transform_world_to_camera(candidates, cams)\n    # re-center relative to the hip\n    for idx in range(len(p_3d_cam)):\n        p_3d_cam[idx] = postprocess_3d(p_3d_cam[idx])\n    # normalize the inputs\n    dim_use_2d = stats['dim_use_2d']\n    dim_use_3d = stats['dim_use_3d']\n    mean_2d = stats['mean_2d'][dim_use_2d]\n    std_2d = stats['std_2d'][dim_use_2d]\n    mean_3d = stats['mean_3d'][dim_use_3d]\n    std_3d = stats['std_3d'][dim_use_3d]\n    for idx in range(len(p_2d)):\n        p_2d[idx] = p_2d[idx][:, dim_use_2d]\n        p_2d[idx] = normalize(p_2d[idx], mean_2d, std_2d)\n    # get output and calculate errors\n    output = []\n    for idx in range(len(p_2d)):\n        prediction = to_numpy(get_prediction(model, p_2d[idx]))\n        # unnormalize the prediction\n        prediction = unnormalize(prediction, mean_3d, std_3d)\n        output.append(prediction)\n    errors = []\n    for idx in range(len(output)):\n        gt_poses = p_3d_cam[idx][:, dim_use_3d]\n        errors.append(calc_errors(output[idx], gt_poses))\n    all_errors = np.concatenate(errors, axis = 1)\n    # mean error for all the cameras\n    mean_errors = all_errors.mean(axis = 1)\n    return mean_errors\n\ndef active_select(model_file, candidates, ratio):\n    \"\"\"\n    Actively select candidates that cause the model to fail.\n    \"\"\"\n    scores = get_score(model_file, candidates)\n    indices = np.argsort(scores) # error from low to high\n    indices = indices[-int(ratio*len(candidates)):]\n    mean_error = scores[indices].mean()\n    return candidates[indices], mean_error\n\ndef cast_to_float(dic, dtype=np.float32):\n    # cast to float 32 for space saving\n    for key in dic.keys():\n        dic[key] = dic[key].astype(dtype)\n    return dic\n\ndef xyz2spherical(xyz):\n    # convert cartesian coordinate to spherical coordinate\n    # return in r, phi, and theta (elevation angle from z axis down)\n    return_value = np.zeros(xyz.shape, dtype=xyz.dtype)\n    xy = xyz[:,:,0]**2 + xyz[:,:,1]**2\n    return_value[:,:,0] = np.sqrt(xy + xyz[:,:,2]**2) # r\n    return_value[:,:,1] = np.arctan2(np.sqrt(xy), xyz[:,:,2]) # phi\n    return_value[:,:,2] = np.arctan2(xyz[:,:,1], xyz[:,:,0]) #theta\n    return return_value \n\ndef spherical2xyz(rphitheta):\n    return_value = np.zeros(rphitheta.shape, dtype=rphitheta.dtype)\n    sinphi = np.sin(rphitheta[:,:,1])\n    cosphi = np.cos(rphitheta[:,:,1])\n    sintheta = np.sin(rphitheta[:,:,2])\n    costheta = np.cos(rphitheta[:,:,2])\n    return_value[:,:,0] = rphitheta[:,:,0]*sinphi*costheta # x\n    return_value[:,:,1] = rphitheta[:,:,0]*sinphi*sintheta # y\n    return_value[:,:,2] = rphitheta[:,:,0]*cosphi #z\n    return return_value\n\n# global variables\nparent_idx = [0, 6, 7, \\\n              0, 1, 2, \\\n              0, 12, 13, 14,\\\n              13, 17, 18,\\\n              13, 25, 26]\nchild_idx = [6, 7, 8, \\\n             1, 2, 3, \\\n             12, 13, 14, 15,\\\n             17, 18, 19,\\\n             25, 26, 27]\n\ndef position_to_angle(skeletons):\n    # transform 3d positions to joint angle representation \n    \n    # first compute the bone vectors\n    # a bone vector is the vector from on parent joint to one child joint \n    # hip->left hip->left knee->left foot,\n    # hip->right hip-> right knee-> right foot\n    # hip -> spine->thorax->nose->head\n    # thorax -> left shoulder->left elbow->left wrist\n    # thorax -> right shoulder-> right elbow->right wrist\n    num_sample = skeletons.shape[0]\n    skeletons = skeletons.reshape(num_sample, -1, 3)\n\n    parent_joints = skeletons[:, parent_idx, :]\n    child_joints = skeletons[:, child_idx, :]\n    bone_vectors = child_joints - parent_joints\n    # now compute the angles and bone lengths\n    rphitheta = xyz2spherical(bone_vectors)    \n    return rphitheta\n\ndef angle_to_position(rphitheta, skeletons):\n    # transform joint angle representation to 3d positions \n    # starting from the root, create joint one by one according to predefined\n    # hierarchical relation\n    num_sample = skeletons.shape[0]\n    skeletons = skeletons.reshape(num_sample, -1, 3)    \n    for bone_idx in range(len(parent_idx)):\n        offset = spherical2xyz(np.expand_dims(rphitheta[:, bone_idx, :], axis=1))\n        offset = offset[:,0,:]\n        skeletons[:, child_idx[bone_idx], :] = \\\n        skeletons[:, parent_idx[bone_idx], :] + offset\n    return skeletons\n\ndef mutate_bone_length(population, opt, gen_idx, method='simple'):\n    \"\"\"\n    Randomly mutate bone length in a population to increase variation in \n    subject size. \n    For example, H36M only contains adults yet you can modify bone\n    length to represent children. Since the posture and subject size are \n    independent, you can synthetize data for dancing kids for free if you already\n    have data for dancing adults. You only need little prior knowledge on human \n    bone length.\n    \"\"\"\n    # the camera parameters in H36M correspond to the five subjects\n    # Rename the synthetic population as these subjects so that the camera \n    # parameters can be used\n    psuedo_subject_names = [1, 5, 6, 7, 8]\n    dict_3d = {}\n    for i in range(len(population)):\n        if np.random.rand() > opt.MBLR:\n            angles = position_to_angle(population[i].reshape(1, -1))  \n            if method == 'simple':\n                # The simplest way is to change to bone length to some value\n                # according to prior knowledge about human bone size.\n                # In our experiment, we collect these values manually from our \n                # interactive visualization tool as well as cross validation. \n                idx = np.random.randint(0, len(bl_templates))\n                angles[0, :, 0] = bl_templates[idx]\n                population[i] = (angle_to_position(angles, population[i].reshape(1,-1))).reshape(-1)\n            elif method == 'addnoise':\n                # add Gaussian noise to current bone length to obtain new bone length\n                raise ValueError('Deprecated')\n            else:\n                raise NotImplementedError\n    poses_list = np.array_split(population, len(psuedo_subject_names))\n    for subject_idx in range(len(psuedo_subject_names)):\n        dict_3d[(psuedo_subject_names[subject_idx], 'n/a', 'n/a')] =\\\n        poses_list[subject_idx]       \n    save_path = get_save_path(opt, gen_idx)\n    np.save(save_path, cast_to_float(dict_3d))\n    logging.info('file saved at ' + save_path)\n    return\n\ndef one_iteration(population, opt, model_file=None):\n    \"\"\"\n    Run one iteration to produce the next generation.\n    \"\"\"\n    # select the best individuals\n    father_indices, mother_indices = choose_best(population, fraction=opt.F)\n    # produce next generation by evolutionary operators\n    offsprings = []\n    for idx in tqdm(range(len(father_indices))):\n        son, daughter = exploration(population[father_indices[idx]],\n                                    population[mother_indices[idx]],\n                                    opt)\n        if son is not None:\n            offsprings.append(son.reshape(1,-1))\n        if daughter is not None:\n            offsprings.append(daughter.reshape(1,-1))\n    offsprings = np.concatenate(offsprings, axis=0)\n    logging.info('{:d} out of {:d} poses survived.'.format(len(offsprings),\n                 len(father_indices)*2))\n    # select the synthetic data actively\n    if opt.A:\n        assert model_file is not None\n        num_before = len(offsprings)\n        offsprings, mean_error = active_select(model_file, offsprings, opt.AR)\n        logging.info('{:d} out of {:d} poses are selected actively with mean'\\\n                     'error {:.2f}'.format(len(offsprings), num_before, mean_error))        \n    if opt.Mer:\n        # merge the offsprings with the parents\n        population = np.vstack([population, offsprings])\n    else:\n        population = offsprings    \n    return population\n\ndef get_save_path(opt, gen_idx):\n    if opt.WS:\n        save_path = os.path.join(opt.SD, opt.SS, opt.SN)\n    else:\n        save_path = os.path.join(opt.SD, 'S15678', opt.SN)\n    if not os.path.exists(save_path):\n        os.makedirs(save_path)\n    save_path = os.path.join(save_path, 'generation_{:d}.npy'.format(gen_idx))\n    return save_path\n\ndef split_and_save(final_poses, parameters, gen_idx):\n    temp_subject_list = [1, 5, 6, 7, 8]\n    train_set_3d = {}\n    poses_list = np.array_split(final_poses, len(temp_subject_list))\n    for subject_idx in range(len(temp_subject_list)):\n        train_set_3d[(temp_subject_list[subject_idx], 'n/a', 'n/a')] =\\\n        poses_list[subject_idx]         \n    save_path = get_save_path(parameters, gen_idx)\n    np.save(save_path, cast_to_float(train_set_3d))     \n    print('file saved at {:s}!'.format(save_path))\n    return\n\ndef save_results(poses, opt, gen_idx):\n    # get save path\n    if opt.MBL:\n        mutate_bone_length(poses, opt, gen_idx)\n    else:\n        split_and_save(poses, opt, gen_idx)\n    return\n\ndef evolution(initial_population, opt, model_file=None):\n    \"\"\"\n    Dataset evolution.\n    \"\"\"\n    logging.basicConfig(level=logging.INFO,\n                        format=\"[%(asctime)s]: %(message)s\"\n                        )\n    population = initial_population\n    save_results(initial_population, opt, 0)\n    initial_num = len(initial_population)\n    for gen_idx in range(1, opt.G+1):\n        population = one_iteration(population, opt, model_file=model_file)\n    save_results(population, opt, gen_idx)\n    # if not enough\n    if opt.E and len(population) < initial_num * opt.T:\n        logging.info('Running extra generations to synthesize enough data...')\n        while len(population) < initial_num * opt.T:\n            gen_idx += 1\n            logging.info('Generation {:d}...'.format(gen_idx))\n            population = one_iteration(population, opt, model_file=model_file)\n            if opt.I:\n                save_results(population.copy(), opt, gen_idx)\n                logging.info('Generation {:d} saved.'.format(gen_idx))\n    save_results(population, opt, gen_idx)\n    logging.info('Final population saved.')\n    return population\n"
  },
  {
    "path": "libs/evolution/parameter.py",
    "content": "\"\"\"\nArguments and hyper-parameters used in dataset evolution. \n\"\"\"\nimport argparse\n\ndef parse_arg():\n    parser = argparse.ArgumentParser(description='evolve.py')\n    ##-----------------------------------------------------------------------##\n    ## Hyper-parameters    \n    # Number of generation to run\n    parser.add_argument('-G', type=int, default=1)\n    # Synthetize enough (E) data with a target ratio after G generations\n    parser.add_argument('-E', type=bool, default=True)\n    parser.add_argument('-T', type=float, default=2.5) # the target ratio\n    # Fraction\n    parser.add_argument('-F', type=float, default=0.05)\n    # Apply mutation on skeleton orientation\n    parser.add_argument('-M', type=bool, default=True)\n    # Apply mutation on bone vector length\n    parser.add_argument('-MBL', type=bool, default=True)    \n    # The mutation rate for bone vector length\n    parser.add_argument('-MBLR', type=float, default=0.5)   \n    # Mutation rate of changing local limb orientation\n    parser.add_argument('-MRL', type=float, default=0.3)\n    # Mutation rate of changing global skeleton orientation\n    parser.add_argument('-MRG', type=float, default=0.1)\n    # Standrd deviation of Guassian noise (in degrees) for local limb mutation\n    parser.add_argument('-SDL', type=float, default=10.0)\n    # Standrd deviation of Guassian noise for global orientation mutation\n    parser.add_argument('-SDG', type=float, default=30.0)\n    # Select the synthesized data in an active manner \n    parser.add_argument('-A', type=bool, default=False)\n    # The ratio for active selection\n    parser.add_argument('-AR', type=float, default=0.5)\n    # Merge the synthetic data with the initial population\n    parser.add_argument('-Mer', type=bool, default=True)\n    # Apply the crossover operator\n    parser.add_argument('-CV', type=bool, default=True)\n    # Apply constraint to rule out invalid poses\n    parser.add_argument('-C', type=bool, default=True)\n    # Threshold for valid bone vector\n    parser.add_argument('-Th', type=int, default=9)\n    # Visualize the synthetic skeleton during exploring the data space\n    parser.add_argument('-V', type=bool, default=True)\n    # Save the intermediate synthetic data after each generation\n    parser.add_argument('-I', type=bool, default=False)\n    # File name for saving\n    parser.add_argument('-SN', type=str, default='evolved_data')\n    # Sampling string used to down-sample the original data\n    # Examples: ['0.001S1', '0.01S1', '0.05S1', '0.1S1', '0.5S1', 'S1', 'S15', 'S156', 'S15678']\n    parser.add_argument('-SS', type=str, default='0.01S1')\n    # Down-sample the original data for weakly-supervised experiments\n    parser.add_argument('-WS', type=bool, default=False)\n    # Debug mode\n    parser.add_argument('-DE', type=bool, default=False)    \n    ##-----------------------------------------------------------------------##\n    ## paths    \n    # path to H36M data\n    parser.add_argument('-data_path', type=str, default=\"../data/human3.6M/threeDPose_train.npy\")\n    # Directory for saving the synthetic data\n    parser.add_argument('-SD', type=str, default='../data/human3.6M/evolved')\n    ##-----------------------------------------------------------------------##\n    ## Usages    \n    # Usage: generate data\n    parser.add_argument('-generate', type=bool, default=False)\n    # Usage: visualize data\n    parser.add_argument('-visualize', type=bool, default=False)\n    # Usage: split and save evolved dataset\n    parser.add_argument('-split', type=bool, default=False)\n    parser.add_argument('-split_ratio', type=float, default=0.9)\n    opt = parser.parse_args()\n    return opt"
  },
  {
    "path": "libs/hhr/__init__.py",
    "content": "#!/usr/bin/env python3\n# -*- coding: utf-8 -*-\n\"\"\"\nEmpty file.\n\"\"\"\n\n\n"
  },
  {
    "path": "libs/hhr/config/__init__.py",
    "content": "# ------------------------------------------------------------------------------\n# Copyright (c) Microsoft\n# Licensed under the MIT License.\n# Written by Bin Xiao (Bin.Xiao@microsoft.com)\n# ------------------------------------------------------------------------------\n\nfrom .default import _C as cfg\nfrom .default import update_config\nfrom .models import MODEL_EXTRAS\n"
  },
  {
    "path": "libs/hhr/config/default.py",
    "content": "\n# ------------------------------------------------------------------------------\n# Copyright (c) Microsoft\n# Licensed under the MIT License.\n# Written by Bin Xiao (Bin.Xiao@microsoft.com)\n# ------------------------------------------------------------------------------\n\nfrom __future__ import absolute_import\nfrom __future__ import division\nfrom __future__ import print_function\n\nimport os\n\nfrom yacs.config import CfgNode as CN\n\n\n_C = CN()\n\n_C.OUTPUT_DIR = ''\n_C.LOG_DIR = ''\n_C.DATA_DIR = ''\n_C.GPUS = (0,)\n_C.WORKERS = 4\n_C.PRINT_FREQ = 20\n_C.AUTO_RESUME = False\n_C.PIN_MEMORY = True\n_C.RANK = 0\n\n# Cudnn related params\n_C.CUDNN = CN()\n_C.CUDNN.BENCHMARK = True\n_C.CUDNN.DETERMINISTIC = False\n_C.CUDNN.ENABLED = True\n\n# common params for NETWORK\n_C.MODEL = CN()\n_C.MODEL.NAME = 'pose_hrnet'\n_C.MODEL.INIT_WEIGHTS = True\n_C.MODEL.PRETRAINED = ''\n_C.MODEL.NUM_JOINTS = 17\n_C.MODEL.TAG_PER_JOINT = True\n_C.MODEL.TARGET_TYPE = 'gaussian'\n_C.MODEL.IMAGE_SIZE = [256, 256]  # width * height, ex: 192 * 256\n_C.MODEL.HEATMAP_SIZE = [64, 64]  # width * height, ex: 24 * 32\n_C.MODEL.SIGMA = 2\n_C.MODEL.EXTRA = CN(new_allowed=True)\n\n_C.LOSS = CN()\n_C.LOSS.USE_OHKM = False\n_C.LOSS.TOPK = 8\n_C.LOSS.USE_TARGET_WEIGHT = True\n_C.LOSS.USE_DIFFERENT_JOINTS_WEIGHT = False\n\n# DATASET related params\n_C.DATASET = CN()\n_C.DATASET.ROOT = ''\n_C.DATASET.DATASET = 'mpii'\n_C.DATASET.TRAIN_SET = 'train'\n_C.DATASET.TEST_SET = 'valid'\n_C.DATASET.DATA_FORMAT = 'jpg'\n_C.DATASET.TRAIN_PATH = '/h36m/h36m_annot_train.npy'\n_C.DATASET.VALID_PATH = '/h36m/h36m_annot_test.npy'\n_C.DATASET.HYBRID_JOINTS_TYPE = ''\n_C.DATASET.SELECT_DATA = False\n\n# training data augmentation\n_C.DATASET.FLIP = True\n_C.DATASET.SCALE_FACTOR = 0.25\n_C.DATASET.ROT_FACTOR = 30\n_C.DATASET.PROB_HALF_BODY = 0.0\n_C.DATASET.NUM_JOINTS_HALF_BODY = 8\n_C.DATASET.COLOR_RGB = False\n\n# train\n_C.TRAIN = CN()\n\n_C.TRAIN.LR_FACTOR = 0.1\n_C.TRAIN.LR_STEP = [90, 110]\n_C.TRAIN.LR = 0.001\n\n_C.TRAIN.OPTIMIZER = 'adam'\n_C.TRAIN.MOMENTUM = 0.9\n_C.TRAIN.WD = 0.0001\n_C.TRAIN.NESTEROV = False\n_C.TRAIN.GAMMA1 = 0.99\n_C.TRAIN.GAMMA2 = 0.0\n\n_C.TRAIN.BEGIN_EPOCH = 0\n_C.TRAIN.END_EPOCH = 140\n\n_C.TRAIN.RESUME = False\n_C.TRAIN.CHECKPOINT = ''\n\n_C.TRAIN.BATCH_SIZE_PER_GPU = 32\n_C.TRAIN.SHUFFLE = True\n\n# testing\n_C.TEST = CN()\n\n# size of images for each device\n_C.TEST.BATCH_SIZE_PER_GPU = 32\n# Test Model Epoch\n_C.TEST.FLIP_TEST = False\n_C.TEST.POST_PROCESS = False\n_C.TEST.SHIFT_HEATMAP = False\n\n_C.TEST.USE_GT_BBOX = False\n\n# nms\n_C.TEST.IMAGE_THRE = 0.1\n_C.TEST.NMS_THRE = 0.6\n_C.TEST.SOFT_NMS = False\n_C.TEST.OKS_THRE = 0.5\n_C.TEST.IN_VIS_THRE = 0.0\n_C.TEST.COCO_BBOX_FILE = ''\n_C.TEST.BBOX_THRE = 1.0\n_C.TEST.MODEL_FILE = ''\n\n# debug\n_C.DEBUG = CN()\n_C.DEBUG.DEBUG = False\n_C.DEBUG.SAVE_BATCH_IMAGES_GT = False\n_C.DEBUG.SAVE_BATCH_IMAGES_PRED = False\n_C.DEBUG.SAVE_HEATMAPS_GT = False\n_C.DEBUG.SAVE_HEATMAPS_PRED = False\n\n\ndef update_config(cfg, args):\n    cfg.defrost()\n    if hasattr(args, 'cfg'):\n        cfg.merge_from_file(args.cfg)\n    if hasattr(args, 'opts'):\n        cfg.merge_from_list(args.opts)\n\n    if hasattr(args, 'modelDir') and args.modelDir:\n        cfg.OUTPUT_DIR = args.modelDir\n\n    if hasattr(args, 'logDir') and args.logDir:\n        cfg.LOG_DIR = args.logDir\n\n    if hasattr(args, 'dataDir') and args.dataDir:\n        cfg.DATA_DIR = args.dataDir\n\n    cfg.DATASET.ROOT = os.path.join(\n        cfg.DATA_DIR, cfg.DATASET.ROOT\n    )\n\n    cfg.MODEL.PRETRAINED = os.path.join(\n        cfg.DATA_DIR, cfg.MODEL.PRETRAINED\n    )\n\n    if cfg.TEST.MODEL_FILE:\n        cfg.TEST.MODEL_FILE = os.path.join(\n            cfg.DATA_DIR, cfg.TEST.MODEL_FILE\n        )\n\n    cfg.freeze()\n\n\nif __name__ == '__main__':\n    import sys\n    with open(sys.argv[1], 'w') as f:\n        print(_C, file=f)\n\n"
  },
  {
    "path": "libs/hhr/config/models.py",
    "content": "# ------------------------------------------------------------------------------\n# Copyright (c) Microsoft\n# Licensed under the MIT License.\n# Written by Bin Xiao (Bin.Xiao@microsoft.com)\n# ------------------------------------------------------------------------------\n\nfrom __future__ import absolute_import\nfrom __future__ import division\nfrom __future__ import print_function\n\nfrom yacs.config import CfgNode as CN\n\n\n# pose_resnet related params\nPOSE_RESNET = CN()\nPOSE_RESNET.NUM_LAYERS = 50\nPOSE_RESNET.DECONV_WITH_BIAS = False\nPOSE_RESNET.NUM_DECONV_LAYERS = 3\nPOSE_RESNET.NUM_DECONV_FILTERS = [256, 256, 256]\nPOSE_RESNET.NUM_DECONV_KERNELS = [4, 4, 4]\nPOSE_RESNET.FINAL_CONV_KERNEL = 1\nPOSE_RESNET.PRETRAINED_LAYERS = ['*']\n\n# pose_multi_resoluton_net related params\nPOSE_HIGH_RESOLUTION_NET = CN()\nPOSE_HIGH_RESOLUTION_NET.PRETRAINED_LAYERS = ['*']\nPOSE_HIGH_RESOLUTION_NET.STEM_INPLANES = 64\nPOSE_HIGH_RESOLUTION_NET.FINAL_CONV_KERNEL = 1\n\nPOSE_HIGH_RESOLUTION_NET.STAGE2 = CN()\nPOSE_HIGH_RESOLUTION_NET.STAGE2.NUM_MODULES = 1\nPOSE_HIGH_RESOLUTION_NET.STAGE2.NUM_BRANCHES = 2\nPOSE_HIGH_RESOLUTION_NET.STAGE2.NUM_BLOCKS = [4, 4]\nPOSE_HIGH_RESOLUTION_NET.STAGE2.NUM_CHANNELS = [32, 64]\nPOSE_HIGH_RESOLUTION_NET.STAGE2.BLOCK = 'BASIC'\nPOSE_HIGH_RESOLUTION_NET.STAGE2.FUSE_METHOD = 'SUM'\n\nPOSE_HIGH_RESOLUTION_NET.STAGE3 = CN()\nPOSE_HIGH_RESOLUTION_NET.STAGE3.NUM_MODULES = 1\nPOSE_HIGH_RESOLUTION_NET.STAGE3.NUM_BRANCHES = 3\nPOSE_HIGH_RESOLUTION_NET.STAGE3.NUM_BLOCKS = [4, 4, 4]\nPOSE_HIGH_RESOLUTION_NET.STAGE3.NUM_CHANNELS = [32, 64, 128]\nPOSE_HIGH_RESOLUTION_NET.STAGE3.BLOCK = 'BASIC'\nPOSE_HIGH_RESOLUTION_NET.STAGE3.FUSE_METHOD = 'SUM'\n\nPOSE_HIGH_RESOLUTION_NET.STAGE4 = CN()\nPOSE_HIGH_RESOLUTION_NET.STAGE4.NUM_MODULES = 1\nPOSE_HIGH_RESOLUTION_NET.STAGE4.NUM_BRANCHES = 4\nPOSE_HIGH_RESOLUTION_NET.STAGE4.NUM_BLOCKS = [4, 4, 4, 4]\nPOSE_HIGH_RESOLUTION_NET.STAGE4.NUM_CHANNELS = [32, 64, 128, 256]\nPOSE_HIGH_RESOLUTION_NET.STAGE4.BLOCK = 'BASIC'\nPOSE_HIGH_RESOLUTION_NET.STAGE4.FUSE_METHOD = 'SUM'\n\n\nMODEL_EXTRAS = {\n    'pose_resnet': POSE_RESNET,\n    'pose_high_resolution_net': POSE_HIGH_RESOLUTION_NET,\n}\n"
  },
  {
    "path": "libs/hhr/core/__init__.py",
    "content": "#!/usr/bin/env python3\n# -*- coding: utf-8 -*-\n\"\"\"\nEmpty file.\n\"\"\"\n\n\n"
  },
  {
    "path": "libs/hhr/core/evaluate.py",
    "content": "# ------------------------------------------------------------------------------\n# Copyright (c) Microsoft\n# Licensed under the MIT License.\n# Written by Bin Xiao (Bin.Xiao@microsoft.com)\n# Modified by Shichao Li (nicholas.li@connect.ust.hk)\n# ------------------------------------------------------------------------------\n\nfrom __future__ import absolute_import\nfrom __future__ import division\nfrom __future__ import print_function\n\nimport numpy as np\n\nfrom libs.hhr.core.inference import get_max_preds, get_max_preds_soft\nfrom libs.hhr.core.loss import get_max_preds_soft_pt\n\ndef calc_dists(preds, target, normalize):\n    preds = preds.astype(np.float32)\n    target = target.astype(np.float32)\n    dists = np.zeros((preds.shape[1], preds.shape[0]))\n    for n in range(preds.shape[0]):\n        for c in range(preds.shape[1]):\n            if target[n, c, 0] > 1 and target[n, c, 1] > 1:\n                normed_preds = preds[n, c, :] / normalize[n]\n                normed_targets = target[n, c, :] / normalize[n]\n                dists[c, n] = np.linalg.norm(normed_preds - normed_targets)\n            else:\n                dists[c, n] = -1\n    return dists\n\n\ndef dist_acc(dists, thr=0.5):\n    ''' Return percentage below threshold while ignoring values with a -1 '''\n    dist_cal = np.not_equal(dists, -1)\n    num_dist_cal = dist_cal.sum()\n    if num_dist_cal > 0:\n        return np.less(dists[dist_cal], thr).sum() * 1.0 / num_dist_cal\n    else:\n        return -1\n\n\ndef accuracy(output, target, hm_type='gaussian', thr=0.5):\n    '''\n    Calculate accuracy according to PCK,\n    but uses ground truth heatmap rather than x,y locations\n    First value to be returned is average accuracy across 'idxs',\n    followed by individual accuracies\n    '''\n    idx = list(range(output.shape[1]))\n    norm = 1.0\n    if hm_type == 'gaussian':\n        pred, _ = get_max_preds(output)\n        target, _ = get_max_preds(target)\n        h = output.shape[2]\n        w = output.shape[3]\n        norm = np.ones((pred.shape[0], 2)) * np.array([h, w]) / 10\n    dists = calc_dists(pred, target, norm)\n\n    acc = np.zeros((len(idx) + 1))\n    avg_acc = 0\n    cnt = 0\n\n    for i in range(len(idx)):\n        acc[i + 1] = dist_acc(dists[idx[i]])\n        if acc[i + 1] >= 0:\n            avg_acc = avg_acc + acc[i + 1]\n            cnt += 1\n\n    avg_acc = avg_acc / cnt if cnt != 0 else 0\n    if cnt != 0:\n        acc[0] = avg_acc\n    return acc, avg_acc, cnt, pred\n\nfrom libs.hhr.utils.transforms import get_affine_transform\nfrom libs.hhr.utils.transforms import affine_transform_modified\n\ndef get_distance(gt, pred):\n    # gt: [n_joints, 2]\n    # pred: [n_joints, 2]\n    sqerr = (gt - pred)**2\n    sqerr = sqerr.sum(axis = 1, keepdims=True)\n    dist = np.sqrt(sqerr)\n    return dist\n\ndef accuracy_pixel(output, meta_data, image_size = (288.0, 384.0), arg_max='hard'):\n    ''' \n    Report errors in terms of pixels in the original image plane.\n    '''\n    if arg_max == 'soft':\n        if isinstance(output, np.ndarray):\n            pred, max_vals = get_max_preds_soft(output)\n        else:\n            pred, max_vals = get_max_preds_soft_pt(output)\n    elif arg_max == 'hard':\n        if not isinstance(output, np.ndarray):\n            output = output.data.cpu().numpy()\n        pred, max_vals = get_max_preds(output)\n    else:\n        raise NotImplementedError\n        \n    # multiply by down-sample ratio\n    if not isinstance(pred, np.ndarray):\n        pred = pred.data.cpu().numpy()\n        max_vals = max_vals.data.cpu().numpy()\n    pred *= image_size[0]/output.shape[3]\n    # inverse transform and compare pixel didstance\n    centers, scales, rots = meta_data['center'], meta_data['scale'], meta_data['rotation']\n    centers = centers.data.cpu().numpy()\n    scales = scales.data.cpu().numpy()\n    rots = rots.data.cpu().numpy()\n    joints_original_batch = meta_data['j_original'].data.cpu().numpy()\n\n    distance_list = []\n    all_src_coordinates = []\n    for sample_idx in range(len(pred)):\n        trans_inv = get_affine_transform(centers[sample_idx], scales[sample_idx], rots[sample_idx], image_size, inv=1)\n        joints_original = joints_original_batch[sample_idx]        \n        pred_src_coordinates = affine_transform_modified(pred[sample_idx], trans_inv) \n        all_src_coordinates.append(pred_src_coordinates.reshape(1, len(pred_src_coordinates), 2))\n        distance_list.append(get_distance(joints_original, pred_src_coordinates))\n    all_distance = np.hstack(distance_list)\n    acc = all_distance\n    avg_acc = all_distance.mean()\n    cnt = len(distance_list) * len(all_distance)\n    return acc, avg_acc, cnt, np.concatenate(all_src_coordinates, axis=0), pred, max_vals\n"
  },
  {
    "path": "libs/hhr/core/function.py",
    "content": "# ------------------------------------------------------------------------------\n# Copyright (c) Microsoft\n# Licensed under the MIT License.\n# Written by Bin Xiao (Bin.Xiao@microsoft.com)\n# Modified by Shichao Li (nicholas.li@connect.ust.hk)\n# ------------------------------------------------------------------------------\n\nfrom __future__ import absolute_import\nfrom __future__ import division\nfrom __future__ import print_function\n \nimport time\nimport logging\nimport os\n\nimport numpy as np\nimport torch\nfrom torch import autograd\n\nfrom libs.hhr.core.evaluate import accuracy, accuracy_pixel\nfrom libs.hhr.core.inference import get_final_preds\nfrom libs.hhr.utils.transforms import flip_back, get_affine_transform, affine_transform_modified\nfrom libs.hhr.utils.vis import save_debug_images\n\nlogger = logging.getLogger(__name__)\n\ndef test_transformation(meta_data, image_size=(288, 384)):\n    joints_original = meta_data['j_original'].squeeze().data.cpu().numpy()\n    joints = meta_data['joints'].squeeze().data.cpu().numpy() # coordinates in 384*288 image box\n    center, scale = meta_data['center'], meta_data['scale']\n    center = center.data.cpu().numpy().reshape(2)\n    scale = scale.data.cpu().numpy().reshape(2)\n    trans = get_affine_transform(center=center, scale=scale, rot=0.0, output_size=image_size, inv=0)\n    trans_inv = get_affine_transform(center=center, scale=scale, rot=0.0, output_size=image_size, inv=1)\n    # calculate the distance in terms of pixels\n    transformed_coordinates = affine_transform_modified(joints_original, trans)\n    transformed_coordinates2 = affine_transform_modified(joints, trans_inv)\n    dif1 = joints - transformed_coordinates\n    dif2 = joints_original - transformed_coordinates2\n    # compute inverse matrix\n    inv_compute = np.zeros(trans.shape, trans.dtype)\n    inv_compute[:2, :2] = np.linalg.inv(trans[:2, :2])\n    inv_compute[:, 2] = - trans[:, 2]\n    transformed_coordinates3 = affine_transform_modified(joints, inv_compute)\n    dif3 = joints_original - transformed_coordinates3\n    print(dif1, dif2, dif3)\n    return\n\ndef train(config, train_loader, model, criterion, optimizer, epoch,\n          output_dir, tb_log_dir, writer_dict=None, total_iters=None):\n    batch_time = AverageMeter()\n    data_time = AverageMeter()\n    losses = AverageMeter()\n    acc = AverageMeter()\n\n    # switch to train mode\n    model.train()\n\n    end = time.time()\n    num_iters = 0\n    for i, (input, target, target_weight, meta) in enumerate(train_loader):\n        # measure data loading time\n        data_time.update(time.time() - end)\n        num_iters += 1\n        if total_iters is not None and num_iters > total_iters:\n            return\n        # compute output\n        outputs = model(input)\n        \n        target = target.cuda(non_blocking=True)\n        target_weight = target_weight.cuda(non_blocking=True)\n\n        if isinstance(outputs, list):\n            loss = criterion(outputs[0], target, target_weight)\n            for output in outputs[1:]:\n                loss += criterion(output, target, target_weight)\n        else:\n            output = outputs\n            loss = criterion(output, target, target_weight)\n\n        # compute gradient and do update step\n        optimizer.zero_grad()\n        loss.backward()\n        optimizer.step()\n        \n        # measure accuracy and record loss\n        losses.update(loss.item(), input.size(0))\n\n        _, avg_acc, cnt, src_coordinates, pred, max_vals = accuracy_pixel(outputs, meta)\n        \n        acc.update(avg_acc, cnt)\n\n        # measure elapsed time\n        batch_time.update(time.time() - end)\n        end = time.time()\n\n        if i % config.PRINT_FREQ == 0:\n            msg = 'Epoch: [{0}][{1}/{2}]\\t' \\\n                  'Time {batch_time.val:.3f}s ({batch_time.avg:.3f}s)\\t' \\\n                  'Speed {speed:.1f} samples/s\\t' \\\n                  'Data {data_time.val:.3f}s ({data_time.avg:.3f}s)\\t' \\\n                  'Loss {loss.val:.5f} ({loss.avg:.5f})\\t' \\\n                  'Accuracy {acc.val:.3f} ({acc.avg:.3f})'.format(\n                      epoch, i, len(train_loader), batch_time=batch_time,\n                      speed=input.size(0)/batch_time.val,\n                      data_time=data_time, loss=losses, acc=acc)\n            logger.info(msg)\n\n            prefix = '{}_{}'.format(os.path.join(output_dir, 'train'), i)\n            #TODO: save for L1 loss\n            save_debug_images(config, input, meta, target, pred, output,\n                              prefix)\n        del input, target\n    return\n\ndef validate(config, val_loader, val_dataset, model, criterion, output_dir,\n             tb_log_dir, writer_dict=None):\n    batch_time = AverageMeter()\n    losses = AverageMeter()\n    acc = AverageMeter()\n    ratio = config.MODEL.IMAGE_SIZE[0]/config.MODEL.HEATMAP_SIZE[0]\n    # switch to evaluate mode\n    model.eval()\n\n    num_samples = len(val_dataset)\n    all_preds = np.zeros(\n        (num_samples, config.MODEL.NUM_JOINTS, 3),\n        dtype=np.float32\n    )\n    all_boxes = np.zeros((num_samples, 6))\n    image_path = []\n    filenames = []\n    imgnums = []\n    idx = 0\n    with torch.no_grad():\n        end = time.time()\n        for i, (input, target, target_weight, meta) in enumerate(val_loader):\n            # compute output\n            outputs = model(input)\n            if isinstance(outputs, list):\n                output = outputs[-1]\n            else:\n                output = outputs\n\n            if config.TEST.FLIP_TEST:\n                input_flipped = np.flip(input.cpu().numpy(), 3).copy()\n                input_flipped = torch.from_numpy(input_flipped).cuda()\n                outputs_flipped = model(input_flipped)\n\n                if isinstance(outputs_flipped, list):\n                    output_flipped = outputs_flipped[-1]\n                else:\n                    output_flipped = outputs_flipped\n\n                output_flipped = flip_back(output_flipped.cpu().numpy(),\n                                           val_dataset.flip_pairs)\n                output_flipped = torch.from_numpy(output_flipped.copy()).cuda()\n\n\n                # feature is not aligned, shift flipped heatmap for higher accuracy\n                if config.TEST.SHIFT_HEATMAP:\n                    output_flipped[:, :, :, 1:] = \\\n                        output_flipped.clone()[:, :, :, 0:-1]\n\n                output = (output + output_flipped) * 0.5\n\n            target = target.cuda(non_blocking=True)\n            target_weight = target_weight.cuda(non_blocking=True)\n\n            loss = criterion(output, target, target_weight)\n\n            num_images = input.size(0)\n            # measure accuracy and record loss\n            losses.update(loss.item(), num_images)\n            _, avg_acc, cnt, pred = accuracy(output.cpu().numpy(),\n                                             target.cpu().numpy())\n\n            acc.update(avg_acc, cnt)\n\n            # measure elapsed time\n            batch_time.update(time.time() - end)\n            end = time.time()\n\n            c = meta['center'].numpy()\n            s = meta['scale'].numpy()\n            score = meta['score'].numpy()\n\n            preds, maxvals = get_final_preds(\n                config, output.clone().cpu().numpy(), c, s)\n\n            all_preds[idx:idx + num_images, :, 0:2] = preds[:, :, 0:2]\n            all_preds[idx:idx + num_images, :, 2:3] = maxvals\n            # double check this all_boxes parts\n            all_boxes[idx:idx + num_images, 0:2] = c[:, 0:2]\n            all_boxes[idx:idx + num_images, 2:4] = s[:, 0:2]\n            all_boxes[idx:idx + num_images, 4] = np.prod(s*200, 1)\n            all_boxes[idx:idx + num_images, 5] = score\n            image_path.extend(meta['image'])\n\n            idx += num_images\n\n\n            if i % config.PRINT_FREQ == 0:\n                msg = 'Test: [{0}/{1}]\\t' \\\n                      'Time {batch_time.val:.3f} ({batch_time.avg:.3f})\\t' \\\n                      'Loss {loss.val:.4f} ({loss.avg:.4f})\\t' \\\n                      'Accuracy {acc.val:.3f} ({acc.avg:.3f})'.format(\n                          i, len(val_loader), batch_time=batch_time,\n                          loss=losses, acc=acc)\n                logger.info(msg)\n\n                prefix = '{}_{}'.format(\n                    os.path.join(output_dir, 'val'), i\n                )\n                \n                save_debug_images(config, input, meta, target, pred*ratio, output,\n                                  prefix)\n            del input, target\n            \n        name_values, perf_indicator = val_dataset.evaluate(\n            config, all_preds, output_dir, all_boxes, image_path,\n            filenames, imgnums\n        )\n\n        model_name = config.MODEL.NAME\n        if isinstance(name_values, list):\n            for name_value in name_values:\n                _print_name_value(name_value, model_name)\n        else:\n            _print_name_value(name_values, model_name)\n\n    return perf_indicator\n\ndef validate_pixel(config, val_loader, val_dataset, model, criterion, output_dir,\n             tb_log_dir, writer_dict=None, total_batches = 410, save=False, split=None):\n    batch_time = AverageMeter()\n    losses = AverageMeter()\n    acc = AverageMeter()\n\n    # switch to evaluate mode\n    model.eval()\n    # save prediction results\n    if save:        \n        num_samples = len(val_dataset)\n        all_preds = np.zeros(\n            (num_samples, config.MODEL.NUM_JOINTS, 3),\n            dtype=np.float32\n        )\n        image_path = []\n        \n    num_iters = 0\n    idx = 0\n    with torch.no_grad():\n        end = time.time()\n        for i, (input, target, target_weight, meta) in enumerate(val_loader):\n            num_iters += 1\n            if num_iters > total_batches and not save:\n                break\n            # compute output\n            outputs = model(input)\n            if isinstance(outputs, list):\n                output = outputs[-1]\n            else:\n                output = outputs\n\n            if config.TEST.FLIP_TEST:\n                input_flipped = np.flip(input.cpu().numpy(), 3).copy()\n                input_flipped = torch.from_numpy(input_flipped).cuda()\n                outputs_flipped = model(input_flipped)\n\n                if isinstance(outputs_flipped, list):\n                    output_flipped = outputs_flipped[-1]\n                else:\n                    output_flipped = outputs_flipped\n\n                output_flipped = flip_back(output_flipped.cpu().numpy(),\n                                           val_dataset.flip_pairs)\n                output_flipped = torch.from_numpy(output_flipped.copy()).cuda()\n\n\n                # feature is not aligned, shift flipped heatmap for higher accuracy\n                if config.TEST.SHIFT_HEATMAP:\n                    output_flipped[:, :, :, 1:] = \\\n                        output_flipped.clone()[:, :, :, 0:-1]\n\n                output = (output + output_flipped) * 0.5\n\n            target = target.cuda(non_blocking=True)\n            target_weight = target_weight.cuda(non_blocking=True)\n\n            loss = criterion(output, target, target_weight)\n\n            num_images = input.size(0)\n            # measure accuracy and record loss\n            losses.update(loss.item(), num_images)\n            # pytorch version\n            _, avg_acc, cnt, src_coordinates, preds, max_vals = accuracy_pixel(output, meta)            \n            acc.update(avg_acc, cnt)\n\n            # measure elapsed time\n            batch_time.update(time.time() - end)\n            end = time.time()\n            if save:\n                all_preds[idx:idx + num_images, :, 0:2] = src_coordinates\n                all_preds[idx:idx + num_images, :, 2:3] = max_vals\n                # double check this all_boxes parts\n                image_path.extend(meta['image'])\n    \n                idx += num_images\n            \n            if i % config.PRINT_FREQ == 0:\n                msg = 'Test: [{0}/{1}]\\t' \\\n                      'Time {batch_time.val:.3f} ({batch_time.avg:.3f})\\t' \\\n                      'Loss {loss.val:.4f} ({loss.avg:.4f})\\t' \\\n                      'Accuracy {acc.val:.3f} ({acc.avg:.3f})'.format(\n                          i, len(val_loader), batch_time=batch_time,\n                          loss=losses, acc=acc)\n                logger.info(msg)\n\n                prefix = '{}_{}'.format(\n                    os.path.join(output_dir, 'val'), i\n                )\n                if not save:\n                    save_debug_images(config, input, meta, target, preds, output,\n                                      prefix)\n\n    if save:\n        save_path = './h36m_prediction_' + split + '.npy'\n        np.save(save_path, {'pred': all_preds, 'paths':image_path})\n    perf_indicator = acc.avg\n    return perf_indicator\n\n# markdown format output\ndef _print_name_value(name_value, full_arch_name):\n    names = name_value.keys()\n    values = name_value.values()\n    num_values = len(name_value)\n    logger.info(\n        '| Arch ' +\n        ' '.join(['| {}'.format(name) for name in names]) +\n        ' |'\n    )\n    logger.info('|---' * (num_values+1) + '|')\n\n    if len(full_arch_name) > 15:\n        full_arch_name = full_arch_name[:8] + '...'\n    logger.info(\n        '| ' + full_arch_name + ' ' +\n        ' '.join(['| {:.3f}'.format(value) for value in values]) +\n         ' |'\n    )\n\n\nclass AverageMeter(object):\n    \"\"\"Computes and stores the average and current value\"\"\"\n    def __init__(self):\n        self.reset()\n\n    def reset(self):\n        self.val = 0\n        self.avg = 0\n        self.sum = 0\n        self.count = 0\n\n    def update(self, val, n=1):\n        self.val = val\n        self.sum += val * n\n        self.count += n\n        self.avg = self.sum / self.count if self.count != 0 else 0\n"
  },
  {
    "path": "libs/hhr/core/inference.py",
    "content": "# ------------------------------------------------------------------------------\n# Copyright (c) Microsoft\n# Licensed under the MIT License.\n# Written by Bin Xiao (Bin.Xiao@microsoft.com)\n# Modified by Shichao Li (nicholas.li@connect.ust.hk)\n# ------------------------------------------------------------------------------\n\nfrom __future__ import absolute_import\nfrom __future__ import division\nfrom __future__ import print_function\n\nimport math\n\nimport numpy as np\nimport torch\n\nfrom libs.hhr.utils.transforms import transform_preds\n\ndef get_max_preds(batch_heatmaps):\n    '''\n    get predictions from score maps\n    heatmaps: numpy.ndarray([batch_size, num_joints, height, width])\n    '''\n    assert isinstance(batch_heatmaps, np.ndarray), \\\n        'batch_heatmaps should be numpy.ndarray'\n    assert batch_heatmaps.ndim == 4, 'batch_images should be 4-ndim'\n\n    batch_size = batch_heatmaps.shape[0]\n    num_joints = batch_heatmaps.shape[1]\n    width = batch_heatmaps.shape[3]\n    heatmaps_reshaped = batch_heatmaps.reshape((batch_size, num_joints, -1))\n    idx = np.argmax(heatmaps_reshaped, 2)\n    maxvals = np.amax(heatmaps_reshaped, 2)\n\n    maxvals = maxvals.reshape((batch_size, num_joints, 1))\n    idx = idx.reshape((batch_size, num_joints, 1))\n\n    preds = np.tile(idx, (1, 1, 2)).astype(np.float32)\n\n    preds[:, :, 0] = (preds[:, :, 0]) % width\n    preds[:, :, 1] = np.floor((preds[:, :, 1]) / width)\n\n    pred_mask = np.tile(np.greater(maxvals, 0.0), (1, 1, 2))\n    pred_mask = pred_mask.astype(np.float32)\n\n    preds *= pred_mask\n    return preds, maxvals\n\n\ndef get_final_preds(config, batch_heatmaps, center, scale):\n    coords, maxvals = get_max_preds(batch_heatmaps)\n\n    heatmap_height = batch_heatmaps.shape[2]\n    heatmap_width = batch_heatmaps.shape[3]\n\n    # post-processing\n    if config.TEST.POST_PROCESS:\n        for n in range(coords.shape[0]):\n            for p in range(coords.shape[1]):\n                hm = batch_heatmaps[n][p]\n                px = int(math.floor(coords[n][p][0] + 0.5))\n                py = int(math.floor(coords[n][p][1] + 0.5))\n                if 1 < px < heatmap_width-1 and 1 < py < heatmap_height-1:\n                    diff = np.array(\n                        [\n                            hm[py][px+1] - hm[py][px-1],\n                            hm[py+1][px]-hm[py-1][px]\n                        ]\n                    )\n                    coords[n][p] += np.sign(diff) * .25\n\n    preds = coords.copy()\n\n    # Transform back\n    for i in range(coords.shape[0]):\n        preds[i] = transform_preds(\n            coords[i], center[i], scale[i], [heatmap_width, heatmap_height]\n        )\n\n    return preds, maxvals\n\ndef get_max_preds_soft(batch_heatmaps):\n    assert isinstance(batch_heatmaps, np.ndarray), \\\n        'batch_heatmaps should be numpy.ndarray'\n    assert batch_heatmaps.ndim == 4, 'batch_images should be 4-ndim'\n    batch_size = batch_heatmaps.shape[0]\n    num_joints = batch_heatmaps.shape[1]\n    height = batch_heatmaps.shape[2]\n    width = batch_heatmaps.shape[3]\n    # get score/confidence for each joint\n    heatmaps_reshaped = batch_heatmaps.reshape((batch_size, num_joints, -1))\n    maxvals = np.amax(heatmaps_reshaped, 2)\n    maxvals = maxvals.reshape((batch_size, num_joints, 1))    \n    # normalize the heatmaps so that they sum to 1\n    #assert batch_heatmaps.min() >= 0.0\n    batch_heatmaps = np.clip(batch_heatmaps, a_min=0.0, a_max=None)\n    temp_sum = heatmaps_reshaped.sum(axis = 2, keepdims=True)\n    heatmaps_reshaped /= temp_sum\n    ## another normalization method: softmax\n    # spatial soft-max\n    #heatmaps_reshaped = softmax(heatmaps_reshaped, axis=2)\n    ##\n    batch_heatmaps = heatmaps_reshaped.reshape(batch_size, num_joints, height, width)\n    x = batch_heatmaps.sum(axis = 2)\n    y = batch_heatmaps.sum(axis = 3)\n    x_indices = np.arange(width).astype(np.float32).reshape(1,1,width)\n    y_indices = np.arange(height).astype(np.float32).reshape(1,1,height)\n    x *= x_indices\n    y *= y_indices\n    x = x.sum(axis = 2, keepdims=True)\n    y = y.sum(axis = 2, keepdims=True)\n    preds = np.concatenate([x, y], axis=2)\n    pred_mask = np.tile(np.greater(maxvals, 0.0), (1, 1, 2))\n    pred_mask = pred_mask.astype(np.float32)\n    preds *= pred_mask\n    return preds, maxvals"
  },
  {
    "path": "libs/hhr/core/loss.py",
    "content": "# ------------------------------------------------------------------------------\n# Copyright (c) Microsoft\n# Licensed under the MIT License.\n# Written by Bin Xiao (Bin.Xiao@microsoft.com)\n# Modified by Shichao Li (nicholas.li@connect.ust.hk)\n# ------------------------------------------------------------------------------\n\nfrom __future__ import absolute_import\nfrom __future__ import division\nfrom __future__ import print_function\n\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\n\nimport numpy as np\n\nclass JointsMSELoss(nn.Module):\n    def __init__(self, use_target_weight):\n        super(JointsMSELoss, self).__init__()\n        self.criterion = nn.MSELoss(reduction='mean')\n        self.use_target_weight = use_target_weight\n\n    def forward(self, output, target, target_weight):\n        batch_size = output.size(0)\n        num_joints = output.size(1)\n        heatmaps_pred = output.reshape((batch_size, num_joints, -1)).split(1, 1)\n        heatmaps_gt = target.reshape((batch_size, num_joints, -1)).split(1, 1)\n        loss = 0\n\n        for idx in range(num_joints):\n            heatmap_pred = heatmaps_pred[idx].squeeze()\n            heatmap_gt = heatmaps_gt[idx].squeeze()\n            if self.use_target_weight:\n                loss += 0.5 * self.criterion(\n                    heatmap_pred.mul(target_weight[:, idx]),\n                    heatmap_gt.mul(target_weight[:, idx])\n                )\n            else:\n                loss += 0.5 * self.criterion(heatmap_pred, heatmap_gt)\n\n        return loss / num_joints\n\n\nclass JointsOHKMMSELoss(nn.Module):\n    def __init__(self, use_target_weight, topk=8):\n        super(JointsOHKMMSELoss, self).__init__()\n        self.criterion = nn.MSELoss(reduction='none')\n        self.use_target_weight = use_target_weight\n        self.topk = topk\n\n    def ohkm(self, loss):\n        ohkm_loss = 0.\n        for i in range(loss.size()[0]):\n            sub_loss = loss[i]\n            topk_val, topk_idx = torch.topk(\n                sub_loss, k=self.topk, dim=0, sorted=False\n            )\n            tmp_loss = torch.gather(sub_loss, 0, topk_idx)\n            ohkm_loss += torch.sum(tmp_loss) / self.topk\n        ohkm_loss /= loss.size()[0]\n        return ohkm_loss\n\n    def forward(self, output, target, target_weight):\n        batch_size = output.size(0)\n        num_joints = output.size(1)\n        heatmaps_pred = output.reshape((batch_size, num_joints, -1)).split(1, 1)\n        heatmaps_gt = target.reshape((batch_size, num_joints, -1)).split(1, 1)\n\n        loss = []\n        for idx in range(num_joints):\n            heatmap_pred = heatmaps_pred[idx].squeeze()\n            heatmap_gt = heatmaps_gt[idx].squeeze()\n            if self.use_target_weight:\n                loss.append(0.5 * self.criterion(\n                    heatmap_pred.mul(target_weight[:, idx]),\n                    heatmap_gt.mul(target_weight[:, idx])\n                ))\n            else:\n                loss.append(\n                    0.5 * self.criterion(heatmap_pred, heatmap_gt)\n                )\n\n        loss = [l.mean(dim=1).unsqueeze(dim=1) for l in loss]\n        loss = torch.cat(loss, dim=1)\n\n        return self.ohkm(loss)\n\n# soft-argmax\ndef get_max_preds_soft_pt(batch_heatmaps):\n    # pytorch version of the above function using tensors\n    assert len(batch_heatmaps.shape) == 4, 'batch_images should be 4-ndim'\n    batch_size = batch_heatmaps.shape[0]\n    num_joints = batch_heatmaps.shape[1]\n    height = batch_heatmaps.shape[2]\n    width = batch_heatmaps.shape[3]\n    heatmaps_reshaped = batch_heatmaps.view((batch_size, num_joints, -1))\n    # get score/confidence for each joint    \n    maxvals = heatmaps_reshaped.max(dim=2)[0]\n    maxvals = maxvals.view((batch_size, num_joints, 1))       \n    heatmaps_reshaped = F.softmax(heatmaps_reshaped, dim=2)\n    batch_heatmaps = heatmaps_reshaped.view(batch_size, num_joints, height, width)\n    x = batch_heatmaps.sum(dim = 2)\n    y = batch_heatmaps.sum(dim = 3)\n    x_indices = torch.cuda.comm.broadcast(torch.arange(width).type(torch.cuda.FloatTensor), devices=[x.device.index])[0]\n    x_indices = x_indices.view(1,1,width)\n    y_indices = torch.cuda.comm.broadcast(torch.arange(height).type(torch.cuda.FloatTensor), devices=[x.device.index])[0]\n    y_indices = y_indices.view(1,1,height)    \n    x *= x_indices\n    y *= y_indices\n    x = x.sum(dim = 2, keepdim=True)\n    y = y.sum(dim = 2, keepdim=True)\n    preds = torch.cat([x, y], dim=2)\n    return preds, maxvals\n\nclass JointsCoordinateLoss(nn.Module):\n    def __init__(self, use_target_weight, loss_type='sl1', image_size=(384, 288)):\n        super(JointsCoordinateLoss, self).__init__()\n        self.use_target_weight = use_target_weight\n        self.loss_type = loss_type\n        self.image_size = image_size\n        return\n    \n    def forward(self, output, target, target_weight):\n        preds, _ = get_max_preds_soft_pt(output)\n        # normalize the coordinates to 0-1\n        preds[:, :, 0] /= self.image_size[1]\n        preds[:, :, 1] /= self.image_size[0]\n        target[:, :, 0] /= self.image_size[1]\n        target[:, :, 1] /= self.image_size[0]        \n        if self.loss_type == 'sl1':\n            loss = F.smooth_l1_loss(preds, target)\n        elif self.loss_type == 'wing':\n            raise NotImplementedError\n        else:\n            raise NotImplementedError\n        return loss\n\nclass WingLoss(nn.Module):\n    def __init__(self, use_target_weight, width=5, curvature=0.5, image_size=(384, 288)):\n        super(WingLoss, self).__init__()\n        self.width = width\n        self.curvature = curvature\n        self.C = self.width - self.width * np.log(1 + self.width / self.curvature)\n        self.image_size = image_size\n        \n    def forward(self, output, target, target_weight):\n        prediction, _ = get_max_preds_soft_pt(output)\n        # normalize the coordinates to 0-1\n        prediction[:, :, 0] /= self.image_size[1]\n        prediction[:, :, 1] /= self.image_size[0]\n        target[:, :, 0] /= self.image_size[1]\n        target[:, :, 1] /= self.image_size[0]  \n        diff = target - prediction\n        diff_abs = diff.abs()\n        loss = diff_abs.clone()\n\n        idx_smaller = diff_abs < self.width\n        idx_bigger = diff_abs >= self.width\n\n        loss[idx_smaller] = self.width * torch.log(1 + diff_abs[idx_smaller] / self.curvature)\n        loss[idx_bigger]  = loss[idx_bigger] - self.C\n        loss = loss.mean()\n        return loss\n"
  },
  {
    "path": "libs/hhr/utils/__init__.py",
    "content": ""
  },
  {
    "path": "libs/hhr/utils/transforms.py",
    "content": "# ------------------------------------------------------------------------------\n# Copyright (c) Microsoft\n# Licensed under the MIT License.\n# Written by Bin Xiao (Bin.Xiao@microsoft.com)\n# Modified by Shichao Li (nicholas.li@connect.ust.hk)\n# ------------------------------------------------------------------------------\n\nfrom __future__ import absolute_import\nfrom __future__ import division\nfrom __future__ import print_function\n\nimport numpy as np\nimport cv2\n\n\ndef flip_back(output_flipped, matched_parts):\n    '''\n    ouput_flipped: numpy.ndarray(batch_size, num_joints, height, width)\n    '''\n    assert output_flipped.ndim == 4,\\\n        'output_flipped should be [batch_size, num_joints, height, width]'\n\n    output_flipped = output_flipped[:, :, :, ::-1]\n\n    for pair in matched_parts:\n        tmp = output_flipped[:, pair[0], :, :].copy()\n        output_flipped[:, pair[0], :, :] = output_flipped[:, pair[1], :, :]\n        output_flipped[:, pair[1], :, :] = tmp\n\n    return output_flipped\n\n\ndef fliplr_joints(joints, joints_vis, width, matched_parts):\n    \"\"\"\n    flip coords\n    \"\"\"\n    # Flip horizontal\n    joints[:, 0] = width - joints[:, 0] - 1\n\n    # Change left-right parts\n    for pair in matched_parts:\n        joints[pair[0], :], joints[pair[1], :] = \\\n            joints[pair[1], :], joints[pair[0], :].copy()\n        joints_vis[pair[0], :], joints_vis[pair[1], :] = \\\n            joints_vis[pair[1], :], joints_vis[pair[0], :].copy()\n\n    return joints*joints_vis, joints_vis\n\n\ndef transform_preds(coords, center, scale, output_size):\n    target_coords = np.zeros(coords.shape)\n    trans = get_affine_transform(center, scale, 0, output_size, inv=1)\n    for p in range(coords.shape[0]):\n        target_coords[p, 0:2] = affine_transform(coords[p, 0:2], trans)\n    return target_coords\n\n\ndef get_affine_transform(\n        center, scale, rot, output_size,\n        shift=np.array([0, 0], dtype=np.float32), inv=0\n):\n    if not isinstance(scale, np.ndarray) and not isinstance(scale, list):\n        print(scale)\n        scale = np.array([scale, scale])\n\n    scale_tmp = scale * 200.0\n    src_w = scale_tmp[0]\n    dst_w = output_size[0]\n    dst_h = output_size[1]\n\n    rot_rad = np.pi * rot / 180\n    src_dir = get_dir([0, src_w * -0.5], rot_rad)\n    dst_dir = np.array([0, dst_w * -0.5], np.float32)\n\n    src = np.zeros((3, 2), dtype=np.float32)\n    dst = np.zeros((3, 2), dtype=np.float32)\n    src[0, :] = center + scale_tmp * shift\n    src[1, :] = center + src_dir + scale_tmp * shift\n    dst[0, :] = [dst_w * 0.5, dst_h * 0.5]\n    dst[1, :] = np.array([dst_w * 0.5, dst_h * 0.5]) + dst_dir\n\n    src[2:, :] = get_3rd_point(src[0, :], src[1, :])\n    dst[2:, :] = get_3rd_point(dst[0, :], dst[1, :])\n\n    if inv:\n        trans = cv2.getAffineTransform(np.float32(dst), np.float32(src))\n    else:\n        trans = cv2.getAffineTransform(np.float32(src), np.float32(dst))\n\n    return trans\n\n\ndef affine_transform(pt, t):\n    new_pt = np.array([pt[0], pt[1], 1.]).T\n    new_pt = np.dot(t, new_pt)\n    return new_pt[:2]\n\ndef affine_transform_modified(pts, t):\n    # pts of shape [n, 2]\n    new_pts = np.hstack([pts, np.ones((len(pts), 1))]).T\n    new_pts = t @ new_pts\n    return new_pts[:2, :].T\n\ndef get_3rd_point(a, b):\n    direct = a - b\n    return b + np.array([-direct[1], direct[0]], dtype=np.float32)\n\n\ndef get_dir(src_point, rot_rad):\n    sn, cs = np.sin(rot_rad), np.cos(rot_rad)\n\n    src_result = [0, 0]\n    src_result[0] = src_point[0] * cs - src_point[1] * sn\n    src_result[1] = src_point[0] * sn + src_point[1] * cs\n\n    return src_result\n\n\ndef crop(img, center, scale, output_size, rot=0):\n    trans = get_affine_transform(center, scale, rot, output_size)\n\n    dst_img = cv2.warpAffine(\n        img, trans, (int(output_size[0]), int(output_size[1])),\n        flags=cv2.INTER_LINEAR\n    )\n\n    return dst_img\n"
  },
  {
    "path": "libs/hhr/utils/utils.py",
    "content": "# ------------------------------------------------------------------------------\n# Copyright (c) Microsoft\n# Licensed under the MIT License.\n# Written by Bin Xiao (Bin.Xiao@microsoft.com)\n# Modified by Shichao Li (nicholas.li@connect.ust.hk)\n# ------------------------------------------------------------------------------\n\nfrom __future__ import absolute_import\nfrom __future__ import division\nfrom __future__ import print_function\n\nimport os\nimport logging\nimport time\nfrom collections import namedtuple\nfrom pathlib import Path\n\nimport torch\nimport torch.optim as optim\nimport torch.nn as nn\n\n\ndef create_logger(cfg, cfg_name, phase='train'):\n    root_output_dir = Path(cfg.OUTPUT_DIR)\n    # set up logger\n    if not root_output_dir.exists():\n        print('=> creating {}'.format(root_output_dir))\n        root_output_dir.mkdir()\n\n    dataset = cfg.DATASET.DATASET + '_' + cfg.DATASET.HYBRID_JOINTS_TYPE \\\n        if cfg.DATASET.HYBRID_JOINTS_TYPE else cfg.DATASET.DATASET\n    dataset = dataset.replace(':', '_')\n    model = cfg.MODEL.NAME\n    cfg_name = os.path.basename(cfg_name).split('.')[0]\n\n    final_output_dir = root_output_dir / dataset / model / cfg_name\n\n    print('=> creating {}'.format(final_output_dir))\n    final_output_dir.mkdir(parents=True, exist_ok=True)\n\n    time_str = time.strftime('%Y-%m-%d-%H-%M')\n    log_file = '{}_{}_{}.log'.format(cfg_name, time_str, phase)\n    final_log_file = final_output_dir / log_file\n    head = '%(asctime)-15s %(message)s'\n    logging.basicConfig(filename=str(final_log_file),\n                        format=head)\n    logger = logging.getLogger()\n    logger.setLevel(logging.INFO)\n    console = logging.StreamHandler()\n    logging.getLogger('').addHandler(console)\n\n    tensorboard_log_dir = Path(cfg.LOG_DIR) / dataset / model / \\\n        (cfg_name + '_' + time_str)\n\n    print('=> creating {}'.format(tensorboard_log_dir))\n    tensorboard_log_dir.mkdir(parents=True, exist_ok=True)\n\n    return logger, str(final_output_dir), str(tensorboard_log_dir)\n\n\ndef get_optimizer(cfg, model):\n    optimizer = None\n    if cfg.TRAIN.OPTIMIZER == 'sgd':\n        optimizer = optim.SGD(\n            model.parameters(),\n            lr=cfg.TRAIN.LR,\n            momentum=cfg.TRAIN.MOMENTUM,\n            weight_decay=cfg.TRAIN.WD,\n            nesterov=cfg.TRAIN.NESTEROV\n        )\n    elif cfg.TRAIN.OPTIMIZER == 'adam':\n        optimizer = optim.Adam(\n            model.parameters(),\n            lr=cfg.TRAIN.LR\n        )\n\n    return optimizer\n\n\ndef save_checkpoint(states, is_best, output_dir,\n                    filename='checkpoint.pth'):\n    torch.save(states, os.path.join(output_dir, filename))\n    if is_best and 'state_dict' in states:\n        torch.save(states['best_state_dict'],\n                   os.path.join(output_dir, 'model_best.pth'))\n\n\ndef get_model_summary(model, *input_tensors, item_length=26, verbose=False):\n    \"\"\"\n    :param model:\n    :param input_tensors:\n    :param item_length:\n    :return:\n    \"\"\"\n\n    summary = []\n\n    ModuleDetails = namedtuple(\n        \"Layer\", [\"name\", \"input_size\", \"output_size\", \"num_parameters\", \"multiply_adds\"])\n    hooks = []\n    layer_instances = {}\n\n    def add_hooks(module):\n\n        def hook(module, input, output):\n            class_name = str(module.__class__.__name__)\n\n            instance_index = 1\n            if class_name not in layer_instances:\n                layer_instances[class_name] = instance_index\n            else:\n                instance_index = layer_instances[class_name] + 1\n                layer_instances[class_name] = instance_index\n\n            layer_name = class_name + \"_\" + str(instance_index)\n\n            params = 0\n\n            if class_name.find(\"Conv\") != -1 or class_name.find(\"BatchNorm\") != -1 or \\\n               class_name.find(\"Linear\") != -1:\n                for param_ in module.parameters():\n                    params += param_.view(-1).size(0)\n\n            flops = \"Not Available\"\n            if class_name.find(\"Conv\") != -1 and hasattr(module, \"weight\"):\n                flops = (\n                    torch.prod(\n                        torch.LongTensor(list(module.weight.data.size()))) *\n                    torch.prod(\n                        torch.LongTensor(list(output.size())[2:]))).item()\n            elif isinstance(module, nn.Linear):\n                flops = (torch.prod(torch.LongTensor(list(output.size()))) \\\n                         * input[0].size(1)).item()\n\n            if isinstance(input[0], list):\n                input = input[0]\n            if isinstance(output, list):\n                output = output[0]\n\n            summary.append(\n                ModuleDetails(\n                    name=layer_name,\n                    input_size=list(input[0].size()),\n                    output_size=list(output.size()),\n                    num_parameters=params,\n                    multiply_adds=flops)\n            )\n\n        if not isinstance(module, nn.ModuleList) \\\n           and not isinstance(module, nn.Sequential) \\\n           and module != model:\n            hooks.append(module.register_forward_hook(hook))\n\n    model.eval()\n    model.apply(add_hooks)\n\n    space_len = item_length\n\n    model(*input_tensors)\n    for hook in hooks:\n        hook.remove()\n\n    details = ''\n    if verbose:\n        details = \"Model Summary\" + \\\n            os.linesep + \\\n            \"Name{}Input Size{}Output Size{}Parameters{}Multiply Adds (Flops){}\".format(\n                ' ' * (space_len - len(\"Name\")),\n                ' ' * (space_len - len(\"Input Size\")),\n                ' ' * (space_len - len(\"Output Size\")),\n                ' ' * (space_len - len(\"Parameters\")),\n                ' ' * (space_len - len(\"Multiply Adds (Flops)\"))) \\\n                + os.linesep + '-' * space_len * 5 + os.linesep\n\n    params_sum = 0\n    flops_sum = 0\n    for layer in summary:\n        params_sum += layer.num_parameters\n        if layer.multiply_adds != \"Not Available\":\n            flops_sum += layer.multiply_adds\n        if verbose:\n            details += \"{}{}{}{}{}{}{}{}{}{}\".format(\n                layer.name,\n                ' ' * (space_len - len(layer.name)),\n                layer.input_size,\n                ' ' * (space_len - len(str(layer.input_size))),\n                layer.output_size,\n                ' ' * (space_len - len(str(layer.output_size))),\n                layer.num_parameters,\n                ' ' * (space_len - len(str(layer.num_parameters))),\n                layer.multiply_adds,\n                ' ' * (space_len - len(str(layer.multiply_adds)))) \\\n                + os.linesep + '-' * space_len * 5 + os.linesep\n\n    details += os.linesep \\\n        + \"Total Parameters: {:,}\".format(params_sum) \\\n        + os.linesep + '-' * space_len * 5 + os.linesep\n    details += \"Total Multiply Adds (For Convolution and Linear Layers only): {:,} GFLOPs\".format(flops_sum/(1024**3)) \\\n        + os.linesep + '-' * space_len * 5 + os.linesep\n    details += \"Number of Layers\" + os.linesep\n    for layer in layer_instances:\n        details += \"{} : {} layers   \".format(layer, layer_instances[layer])\n\n    return details\n"
  },
  {
    "path": "libs/hhr/utils/vis.py",
    "content": "# ------------------------------------------------------------------------------\n# Copyright (c) Microsoft\n# Licensed under the MIT License.\n# Written by Bin Xiao (Bin.Xiao@microsoft.com)\n# Modified by Shichao Li (nicholas.li@connect.ust.hk)\n# ------------------------------------------------------------------------------\n\nfrom __future__ import absolute_import\nfrom __future__ import division\nfrom __future__ import print_function\n\nimport math\n\nimport numpy as np\nimport torchvision\nimport cv2\n\nfrom libs.hhr.core.inference import get_max_preds\n\n\ndef save_batch_image_with_joints(batch_image, batch_joints, batch_joints_vis,\n                                 file_name, nrow=8, padding=2):\n    '''\n    batch_image: [batch_size, channel, height, width]\n    batch_joints: [batch_size, num_joints, 3],\n    batch_joints_vis: [batch_size, num_joints, 1],\n    }\n    '''\n    grid = torchvision.utils.make_grid(batch_image, nrow, padding, True)\n    ndarr = grid.mul(255).clamp(0, 255).byte().permute(1, 2, 0).cpu().numpy()\n    ndarr = ndarr.copy()\n\n    nmaps = batch_image.size(0)\n    xmaps = min(nrow, nmaps)\n    ymaps = int(math.ceil(float(nmaps) / xmaps))\n    height = int(batch_image.size(2) + padding)\n    width = int(batch_image.size(3) + padding)\n    k = 0\n    for y in range(ymaps):\n        for x in range(xmaps):\n            if k >= nmaps:\n                break\n            joints = batch_joints[k]\n            joints_vis = batch_joints_vis[k]\n\n            for joint, joint_vis in zip(joints, joints_vis):\n                joint[0] = x * width + padding + joint[0]\n                joint[1] = y * height + padding + joint[1]\n                if joint_vis[0]:\n                    cv2.circle(ndarr, (int(joint[0]), int(joint[1])), 2, [255, 0, 0], 2)\n            k = k + 1\n    cv2.imwrite(file_name, ndarr)\n\n\ndef save_batch_heatmaps(batch_image, batch_heatmaps, file_name,\n                        normalize=True):\n    '''\n    batch_image: [batch_size, channel, height, width]\n    batch_heatmaps: ['batch_size, num_joints, height, width]\n    file_name: saved file name\n    '''\n    if normalize:\n        batch_image = batch_image.clone()\n        min = float(batch_image.min())\n        max = float(batch_image.max())\n\n        batch_image.add_(-min).div_(max - min + 1e-5)\n\n    batch_size = batch_heatmaps.size(0)\n    num_joints = batch_heatmaps.size(1)\n    heatmap_height = batch_heatmaps.size(2)\n    heatmap_width = batch_heatmaps.size(3)\n\n    grid_image = np.zeros((batch_size*heatmap_height,\n                           (num_joints+1)*heatmap_width,\n                           3),\n                          dtype=np.uint8)\n\n    preds, maxvals = get_max_preds(batch_heatmaps.detach().cpu().numpy())\n\n    for i in range(batch_size):\n        image = batch_image[i].mul(255)\\\n                              .clamp(0, 255)\\\n                              .byte()\\\n                              .permute(1, 2, 0)\\\n                              .cpu().numpy()\n        heatmaps = batch_heatmaps[i].mul(255)\\\n                                    .clamp(0, 255)\\\n                                    .byte()\\\n                                    .cpu().numpy()\n\n        resized_image = cv2.resize(image,\n                                   (int(heatmap_width), int(heatmap_height)))\n\n        height_begin = heatmap_height * i\n        height_end = heatmap_height * (i + 1)\n        for j in range(num_joints):\n            cv2.circle(resized_image,\n                       (int(preds[i][j][0]), int(preds[i][j][1])),\n                       1, [0, 0, 255], 1)\n            heatmap = heatmaps[j, :, :]\n            colored_heatmap = cv2.applyColorMap(heatmap, cv2.COLORMAP_JET)\n            masked_image = colored_heatmap*0.7 + resized_image*0.3\n            cv2.circle(masked_image,\n                       (int(preds[i][j][0]), int(preds[i][j][1])),\n                       1, [0, 0, 255], 1)\n\n            width_begin = heatmap_width * (j+1)\n            width_end = heatmap_width * (j+2)\n            grid_image[height_begin:height_end, width_begin:width_end, :] = \\\n                masked_image\n            # grid_image[height_begin:height_end, width_begin:width_end, :] = \\\n            #     colored_heatmap*0.7 + resized_image*0.3\n\n        grid_image[height_begin:height_end, 0:heatmap_width, :] = resized_image\n\n    cv2.imwrite(file_name, grid_image)\n\n\ndef save_debug_images(config, input, meta, target, joints_pred, output,\n                      prefix):\n    if not config.DEBUG.DEBUG:\n        return\n\n    if config.DEBUG.SAVE_BATCH_IMAGES_GT:\n        save_batch_image_with_joints(\n            input, meta['joints'], meta['joints_vis'],\n            '{}_gt.jpg'.format(prefix)\n        )\n    if config.DEBUG.SAVE_BATCH_IMAGES_PRED:\n        save_batch_image_with_joints(\n            input, joints_pred, meta['joints_vis'],\n            '{}_pred.jpg'.format(prefix)\n        )\n    if config.DEBUG.SAVE_HEATMAPS_GT:\n        save_batch_heatmaps(\n            input, target, '{}_hm_gt.jpg'.format(prefix)\n        )\n    if config.DEBUG.SAVE_HEATMAPS_PRED:\n        save_batch_heatmaps(\n            input, output, '{}_hm_pred.jpg'.format(prefix)\n        )\n"
  },
  {
    "path": "libs/model/__init__.py",
    "content": "#!/usr/bin/env python3\n# -*- coding: utf-8 -*-\n\"\"\"\nEmpty file.\n\"\"\"\n\n\n"
  },
  {
    "path": "libs/model/model.py",
    "content": "\"\"\"\nFully-connected residual network as a single deep learner.\n\"\"\"\n\nimport torch.nn as nn\nimport torch\n\nclass ResidualBlock(nn.Module):\n    \"\"\"\n    A residual block.\n    \"\"\"\n    def __init__(self, linear_size, p_dropout=0.5, kaiming=False, leaky=False):\n        super(ResidualBlock, self).__init__()\n        self.l_size = linear_size\n        if leaky:\n            self.relu = nn.LeakyReLU(inplace=True)\n        else:\n            self.relu = nn.ReLU(inplace=True)\n        self.dropout = nn.Dropout(p_dropout)\n\n        self.w1 = nn.Linear(self.l_size, self.l_size)\n        self.batch_norm1 = nn.BatchNorm1d(self.l_size)\n\n        self.w2 = nn.Linear(self.l_size, self.l_size)\n        self.batch_norm2 = nn.BatchNorm1d(self.l_size)\n        \n        if kaiming:\n            self.w1.weight.data = nn.init.kaiming_normal_(self.w1.weight.data)\n            self.w2.weight.data = nn.init.kaiming_normal_(self.w2.weight.data)\n            \n    def forward(self, x):\n        y = self.w1(x)\n        y = self.batch_norm1(y)\n        y = self.relu(y)\n        y = self.dropout(y)\n\n        y = self.w2(y)\n        y = self.batch_norm2(y)\n        y = self.relu(y)\n        y = self.dropout(y)\n\n        out = x + y\n\n        return out\n\nclass FCModel(nn.Module):\n    def __init__(self,\n                 stage_id=1,\n                 linear_size=1024,\n                 num_blocks=2,\n                 p_dropout=0.5,\n                 norm_twoD=False,\n                 kaiming=False,\n                 refine_3d=False, \n                 leaky=False,\n                 dm=False,\n                 input_size=32,\n                 output_size=64):\n        \"\"\"\n        Fully-connected network.\n        \"\"\"\n        super(FCModel, self).__init__()\n\n        self.linear_size = linear_size\n        self.p_dropout = p_dropout\n        self.num_blocks = num_blocks\n        self.stage_id = stage_id\n        self.refine_3d = refine_3d\n        self.leaky = leaky\n        self.dm = dm \n        self.input_size = input_size\n        if self.stage_id>1 and self.refine_3d:\n            self.input_size += 16 * 3        \n        # 3d joints\n        self.output_size = output_size\n\n        # process input to linear size\n        self.w1 = nn.Linear(self.input_size, self.linear_size)\n        self.batch_norm1 = nn.BatchNorm1d(self.linear_size)\n\n        self.res_blocks = []\n        for l in range(num_blocks):\n            self.res_blocks.append(ResidualBlock(self.linear_size, \n                                                 self.p_dropout,\n                                                 leaky=self.leaky))\n        self.res_blocks = nn.ModuleList(self.res_blocks)\n        \n        # output\n        self.w2 = nn.Linear(self.linear_size, self.output_size)\n        if self.leaky:\n            self.relu = nn.LeakyReLU(inplace=True)\n        else:\n            self.relu = nn.ReLU(inplace=True)\n        self.dropout = nn.Dropout(self.p_dropout)\n\n        if kaiming:\n            self.w1.weight.data = nn.init.kaiming_normal_(self.w1.weight.data)\n            self.w2.weight.data = nn.init.kaiming_normal_(self.w2.weight.data)\n            \n    def forward(self, x):\n        y = self.get_representation(x)\n        y = self.w2(y)\n        return y\n    \n    def get_representation(self, x):\n        # get the latent representation of an input vector\n        # first layer\n        y = self.w1(x)\n        y = self.batch_norm1(y)\n        y = self.relu(y)\n        y = self.dropout(y)\n\n        # residual blocks\n        for i in range(self.num_blocks):\n            y = self.res_blocks[i](y)        \n        \n        return y\n    \ndef get_model(stage_id, \n              refine_3d=False, \n              norm_twoD=False, \n              num_blocks=2,\n              input_size=32, \n              output_size=64, \n              linear_size=1024, \n              dropout=0.5,\n              leaky=False\n              ):\n    model = FCModel(stage_id=stage_id, \n                    refine_3d=refine_3d, \n                    norm_twoD=norm_twoD,\n                    num_blocks=num_blocks, \n                    input_size=input_size, \n                    output_size=output_size, \n                    linear_size=linear_size,\n                    p_dropout=dropout, \n                    leaky=leaky\n                    )\n    return model\n\ndef prepare_optim(model, opt):\n    \"\"\"\n    Prepare optimizer.\n    \"\"\"\n    params = [ p for p in model.parameters() if p.requires_grad]\n    if opt.optim_type == 'adam':\n        optimizer = torch.optim.Adam(params, \n                                     lr = opt.lr, \n                                     weight_decay = opt.weight_decay\n                                     )\n    elif opt.optim_type == 'sgd':\n        optimizer = torch.optim.SGD(params, \n                                    lr = opt.lr, \n                                    momentum = opt.momentum,\n                                    weight_decay = opt.weight_decay\n                                    )    \n    scheduler = torch.optim.lr_scheduler.MultiStepLR(optimizer, \n                                                     milestones = opt.milestones, \n                                                     gamma = opt.gamma)\n    return optimizer, scheduler\n\ndef get_cascade():\n    \"\"\"\n    Get an empty cascade.\n    \"\"\"\n    return nn.ModuleList([])"
  },
  {
    "path": "libs/model/pose_hrnet.py",
    "content": "# ------------------------------------------------------------------------------\n# Copyright (c) Microsoft\n# Licensed under the MIT License.\n# Written by Bin Xiao (Bin.Xiao@microsoft.com)\n# ------------------------------------------------------------------------------\n\nfrom __future__ import absolute_import\nfrom __future__ import division\nfrom __future__ import print_function\n\nimport os\nimport logging\n\nimport torch\nimport torch.nn as nn\n\n\nBN_MOMENTUM = 0.1\nlogger = logging.getLogger(__name__)\n\n\ndef conv3x3(in_planes, out_planes, stride=1):\n    \"\"\"3x3 convolution with padding\"\"\"\n    return nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride,\n                     padding=1, bias=False)\n\n\nclass BasicBlock(nn.Module):\n    expansion = 1\n\n    def __init__(self, inplanes, planes, stride=1, downsample=None):\n        super(BasicBlock, self).__init__()\n        self.conv1 = conv3x3(inplanes, planes, stride)\n        self.bn1 = nn.BatchNorm2d(planes, momentum=BN_MOMENTUM)\n        self.relu = nn.ReLU(inplace=True)\n        self.conv2 = conv3x3(planes, planes)\n        self.bn2 = nn.BatchNorm2d(planes, momentum=BN_MOMENTUM)\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\n\nclass Bottleneck(nn.Module):\n    expansion = 4\n\n    def __init__(self, inplanes, planes, stride=1, downsample=None):\n        super(Bottleneck, self).__init__()\n        self.conv1 = nn.Conv2d(inplanes, planes, kernel_size=1, bias=False)\n        self.bn1 = nn.BatchNorm2d(planes, momentum=BN_MOMENTUM)\n        self.conv2 = nn.Conv2d(planes, planes, kernel_size=3, stride=stride,\n                               padding=1, bias=False)\n        self.bn2 = nn.BatchNorm2d(planes, momentum=BN_MOMENTUM)\n        self.conv3 = nn.Conv2d(planes, planes * self.expansion, kernel_size=1,\n                               bias=False)\n        self.bn3 = nn.BatchNorm2d(planes * self.expansion,\n                                  momentum=BN_MOMENTUM)\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        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\n\nclass HighResolutionModule(nn.Module):\n    def __init__(self, num_branches, blocks, num_blocks, num_inchannels,\n                 num_channels, fuse_method, multi_scale_output=True):\n        super(HighResolutionModule, self).__init__()\n        self._check_branches(\n            num_branches, blocks, num_blocks, num_inchannels, num_channels)\n\n        self.num_inchannels = num_inchannels\n        self.fuse_method = fuse_method\n        self.num_branches = num_branches\n\n        self.multi_scale_output = multi_scale_output\n\n        self.branches = self._make_branches(\n            num_branches, blocks, num_blocks, num_channels)\n        self.fuse_layers = self._make_fuse_layers()\n        self.relu = nn.ReLU(True)\n\n    def _check_branches(self, num_branches, blocks, num_blocks,\n                        num_inchannels, num_channels):\n        if num_branches != len(num_blocks):\n            error_msg = 'NUM_BRANCHES({}) <> NUM_BLOCKS({})'.format(\n                num_branches, len(num_blocks))\n            logger.error(error_msg)\n            raise ValueError(error_msg)\n\n        if num_branches != len(num_channels):\n            error_msg = 'NUM_BRANCHES({}) <> NUM_CHANNELS({})'.format(\n                num_branches, len(num_channels))\n            logger.error(error_msg)\n            raise ValueError(error_msg)\n\n        if num_branches != len(num_inchannels):\n            error_msg = 'NUM_BRANCHES({}) <> NUM_INCHANNELS({})'.format(\n                num_branches, len(num_inchannels))\n            logger.error(error_msg)\n            raise ValueError(error_msg)\n\n    def _make_one_branch(self, branch_index, block, num_blocks, num_channels,\n                         stride=1):\n        downsample = None\n        if stride != 1 or \\\n           self.num_inchannels[branch_index] != num_channels[branch_index] * block.expansion:\n            downsample = nn.Sequential(\n                nn.Conv2d(\n                    self.num_inchannels[branch_index],\n                    num_channels[branch_index] * block.expansion,\n                    kernel_size=1, stride=stride, bias=False\n                ),\n                nn.BatchNorm2d(\n                    num_channels[branch_index] * block.expansion,\n                    momentum=BN_MOMENTUM\n                ),\n            )\n\n        layers = []\n        layers.append(\n            block(\n                self.num_inchannels[branch_index],\n                num_channels[branch_index],\n                stride,\n                downsample\n            )\n        )\n        self.num_inchannels[branch_index] = \\\n            num_channels[branch_index] * block.expansion\n        for i in range(1, num_blocks[branch_index]):\n            layers.append(\n                block(\n                    self.num_inchannels[branch_index],\n                    num_channels[branch_index]\n                )\n            )\n\n        return nn.Sequential(*layers)\n\n    def _make_branches(self, num_branches, block, num_blocks, num_channels):\n        branches = []\n\n        for i in range(num_branches):\n            branches.append(\n                self._make_one_branch(i, block, num_blocks, num_channels)\n            )\n\n        return nn.ModuleList(branches)\n\n    def _make_fuse_layers(self):\n        if self.num_branches == 1:\n            return None\n\n        num_branches = self.num_branches\n        num_inchannels = self.num_inchannels\n        fuse_layers = []\n        for i in range(num_branches if self.multi_scale_output else 1):\n            fuse_layer = []\n            for j in range(num_branches):\n                if j > i:\n                    fuse_layer.append(\n                        nn.Sequential(\n                            nn.Conv2d(\n                                num_inchannels[j],\n                                num_inchannels[i],\n                                1, 1, 0, bias=False\n                            ),\n                            nn.BatchNorm2d(num_inchannels[i]),\n                            nn.Upsample(scale_factor=2**(j-i), mode='nearest')\n                        )\n                    )\n                elif j == i:\n                    fuse_layer.append(None)\n                else:\n                    conv3x3s = []\n                    for k in range(i-j):\n                        if k == i - j - 1:\n                            num_outchannels_conv3x3 = num_inchannels[i]\n                            conv3x3s.append(\n                                nn.Sequential(\n                                    nn.Conv2d(\n                                        num_inchannels[j],\n                                        num_outchannels_conv3x3,\n                                        3, 2, 1, bias=False\n                                    ),\n                                    nn.BatchNorm2d(num_outchannels_conv3x3)\n                                )\n                            )\n                        else:\n                            num_outchannels_conv3x3 = num_inchannels[j]\n                            conv3x3s.append(\n                                nn.Sequential(\n                                    nn.Conv2d(\n                                        num_inchannels[j],\n                                        num_outchannels_conv3x3,\n                                        3, 2, 1, bias=False\n                                    ),\n                                    nn.BatchNorm2d(num_outchannels_conv3x3),\n                                    nn.ReLU(True)\n                                )\n                            )\n                    fuse_layer.append(nn.Sequential(*conv3x3s))\n            fuse_layers.append(nn.ModuleList(fuse_layer))\n\n        return nn.ModuleList(fuse_layers)\n\n    def get_num_inchannels(self):\n        return self.num_inchannels\n\n    def forward(self, x):\n        if self.num_branches == 1:\n            return [self.branches[0](x[0])]\n\n        for i in range(self.num_branches):\n            x[i] = self.branches[i](x[i])\n\n        x_fuse = []\n\n        for i in range(len(self.fuse_layers)):\n            y = x[0] if i == 0 else self.fuse_layers[i][0](x[0])\n            for j in range(1, self.num_branches):\n                if i == j:\n                    y = y + x[j]\n                else:\n                    y = y + self.fuse_layers[i][j](x[j])\n            x_fuse.append(self.relu(y))\n\n        return x_fuse\n\n\nblocks_dict = {\n    'BASIC': BasicBlock,\n    'BOTTLENECK': Bottleneck\n}\n\n\nclass PoseHighResolutionNet(nn.Module):\n\n    def __init__(self, cfg, **kwargs):\n        self.inplanes = 64\n        extra = cfg.MODEL.EXTRA\n        super(PoseHighResolutionNet, self).__init__()\n\n        # stem net\n        self.conv1 = nn.Conv2d(3, 64, kernel_size=3, stride=2, padding=1,\n                               bias=False)\n        self.bn1 = nn.BatchNorm2d(64, momentum=BN_MOMENTUM)\n        self.conv2 = nn.Conv2d(64, 64, kernel_size=3, stride=2, padding=1,\n                               bias=False)\n        self.bn2 = nn.BatchNorm2d(64, momentum=BN_MOMENTUM)\n        self.relu = nn.ReLU(inplace=True)\n        self.layer1 = self._make_layer(Bottleneck, 64, 4)\n\n        self.stage2_cfg = cfg['MODEL']['EXTRA']['STAGE2']\n        num_channels = self.stage2_cfg['NUM_CHANNELS']\n        block = blocks_dict[self.stage2_cfg['BLOCK']]\n        num_channels = [\n            num_channels[i] * block.expansion for i in range(len(num_channels))\n        ]\n        self.transition1 = self._make_transition_layer([256], num_channels)\n        self.stage2, pre_stage_channels = self._make_stage(\n            self.stage2_cfg, num_channels)\n\n        self.stage3_cfg = cfg['MODEL']['EXTRA']['STAGE3']\n        num_channels = self.stage3_cfg['NUM_CHANNELS']\n        block = blocks_dict[self.stage3_cfg['BLOCK']]\n        num_channels = [\n            num_channels[i] * block.expansion for i in range(len(num_channels))\n        ]\n        self.transition2 = self._make_transition_layer(\n            pre_stage_channels, num_channels)\n        self.stage3, pre_stage_channels = self._make_stage(\n            self.stage3_cfg, num_channels)\n\n        self.stage4_cfg = cfg['MODEL']['EXTRA']['STAGE4']\n        num_channels = self.stage4_cfg['NUM_CHANNELS']\n        block = blocks_dict[self.stage4_cfg['BLOCK']]\n        num_channels = [\n            num_channels[i] * block.expansion for i in range(len(num_channels))\n        ]\n        self.transition3 = self._make_transition_layer(\n            pre_stage_channels, num_channels)\n        self.stage4, pre_stage_channels = self._make_stage(\n            self.stage4_cfg, num_channels, multi_scale_output=False)\n\n        self.final_layer = nn.Conv2d(\n            in_channels=pre_stage_channels[0],\n            out_channels=cfg.MODEL.NUM_JOINTS,\n            kernel_size=extra.FINAL_CONV_KERNEL,\n            stride=1,\n            padding=1 if extra.FINAL_CONV_KERNEL == 3 else 0\n        )\n\n        self.pretrained_layers = cfg['MODEL']['EXTRA']['PRETRAINED_LAYERS']\n        \n        # add a pixel shuffle upsampling layer\n        self.upsample_layer = nn.Sequential(\n                nn.Conv2d(17, 17*16, kernel_size=1),\n                nn.BatchNorm2d(17*16),\n                nn.ReLU(inplace=True),                \n                nn.PixelShuffle(4)\n                )\n\n    def _make_transition_layer(\n            self, num_channels_pre_layer, num_channels_cur_layer):\n        num_branches_cur = len(num_channels_cur_layer)\n        num_branches_pre = len(num_channels_pre_layer)\n\n        transition_layers = []\n        for i in range(num_branches_cur):\n            if i < num_branches_pre:\n                if num_channels_cur_layer[i] != num_channels_pre_layer[i]:\n                    transition_layers.append(\n                        nn.Sequential(\n                            nn.Conv2d(\n                                num_channels_pre_layer[i],\n                                num_channels_cur_layer[i],\n                                3, 1, 1, bias=False\n                            ),\n                            nn.BatchNorm2d(num_channels_cur_layer[i]),\n                            nn.ReLU(inplace=True)\n                        )\n                    )\n                else:\n                    transition_layers.append(None)\n            else:\n                conv3x3s = []\n                for j in range(i+1-num_branches_pre):\n                    inchannels = num_channels_pre_layer[-1]\n                    outchannels = num_channels_cur_layer[i] \\\n                        if j == i-num_branches_pre else inchannels\n                    conv3x3s.append(\n                        nn.Sequential(\n                            nn.Conv2d(\n                                inchannels, outchannels, 3, 2, 1, bias=False\n                            ),\n                            nn.BatchNorm2d(outchannels),\n                            nn.ReLU(inplace=True)\n                        )\n                    )\n                transition_layers.append(nn.Sequential(*conv3x3s))\n\n        return nn.ModuleList(transition_layers)\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(\n                    self.inplanes, planes * block.expansion,\n                    kernel_size=1, stride=stride, bias=False\n                ),\n                nn.BatchNorm2d(planes * block.expansion, momentum=BN_MOMENTUM),\n            )\n\n        layers = []\n        layers.append(block(self.inplanes, planes, stride, downsample))\n        self.inplanes = planes * block.expansion\n        for i in range(1, blocks):\n            layers.append(block(self.inplanes, planes))\n\n        return nn.Sequential(*layers)\n\n    def _make_stage(self, layer_config, num_inchannels,\n                    multi_scale_output=True):\n        num_modules = layer_config['NUM_MODULES']\n        num_branches = layer_config['NUM_BRANCHES']\n        num_blocks = layer_config['NUM_BLOCKS']\n        num_channels = layer_config['NUM_CHANNELS']\n        block = blocks_dict[layer_config['BLOCK']]\n        fuse_method = layer_config['FUSE_METHOD']\n\n        modules = []\n        for i in range(num_modules):\n            # multi_scale_output is only used last module\n            if not multi_scale_output and i == num_modules - 1:\n                reset_multi_scale_output = False\n            else:\n                reset_multi_scale_output = True\n\n            modules.append(\n                HighResolutionModule(\n                    num_branches,\n                    block,\n                    num_blocks,\n                    num_inchannels,\n                    num_channels,\n                    fuse_method,\n                    reset_multi_scale_output\n                )\n            )\n            num_inchannels = modules[-1].get_num_inchannels()\n\n        return nn.Sequential(*modules), num_inchannels\n\n    def forward(self, x):\n        x = self.conv1(x)\n        x = self.bn1(x)\n        x = self.relu(x)\n        x = self.conv2(x)\n        x = self.bn2(x)\n        x = self.relu(x)\n        x = self.layer1(x)\n\n        x_list = []\n        for i in range(self.stage2_cfg['NUM_BRANCHES']):\n            if self.transition1[i] is not None:\n                x_list.append(self.transition1[i](x))\n            else:\n                x_list.append(x)\n        y_list = self.stage2(x_list)\n\n        x_list = []\n        for i in range(self.stage3_cfg['NUM_BRANCHES']):\n            if self.transition2[i] is not None:\n                x_list.append(self.transition2[i](y_list[-1]))\n            else:\n                x_list.append(y_list[i])\n        y_list = self.stage3(x_list)\n\n        x_list = []\n        for i in range(self.stage4_cfg['NUM_BRANCHES']):\n            if self.transition3[i] is not None:\n                x_list.append(self.transition3[i](y_list[-1]))\n            else:\n                x_list.append(y_list[i])\n        y_list = self.stage4(x_list)\n\n        x = self.final_layer(y_list[0])\n        \n        # Temp test: upsampling\n        x = self.upsample_layer(x)\n        return x\n\n    def init_weights(self, pretrained=''):\n        logger.info('=> init weights from normal distribution')\n        for m in self.modules():\n            if isinstance(m, nn.Conv2d):\n                # nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')\n                nn.init.normal_(m.weight, std=0.001)\n                for name, _ in m.named_parameters():\n                    if name in ['bias']:\n                        nn.init.constant_(m.bias, 0)\n            elif isinstance(m, nn.BatchNorm2d):\n                nn.init.constant_(m.weight, 1)\n                nn.init.constant_(m.bias, 0)\n            elif isinstance(m, nn.ConvTranspose2d):\n                nn.init.normal_(m.weight, std=0.001)\n                for name, _ in m.named_parameters():\n                    if name in ['bias']:\n                        nn.init.constant_(m.bias, 0)\n\n        if os.path.isfile(pretrained):\n            pretrained_state_dict = torch.load(pretrained)\n            logger.info('=> loading pretrained model {}'.format(pretrained))\n\n            need_init_state_dict = {}\n            for name, m in pretrained_state_dict.items():\n                if name.split('.')[0] in self.pretrained_layers \\\n                   or self.pretrained_layers[0] is '*':\n                    need_init_state_dict[name] = m\n            self.load_state_dict(need_init_state_dict, strict=False)\n        elif pretrained:\n            logger.error('=> please download pre-trained models first!')\n            raise ValueError('{} is not exist!'.format(pretrained))\n    \n    def load_my_state_dict(self, state_dict):\n        own_state = self.state_dict()\n        for name, param in state_dict.items():\n            if name not in own_state:\n                 continue\n            param = param.data\n            own_state[name].copy_(param)\n            \ndef get_pose_net(cfg, is_train, **kwargs):\n    model = PoseHighResolutionNet(cfg, **kwargs)\n\n    if is_train and cfg.MODEL.INIT_WEIGHTS:\n        model.init_weights(cfg.MODEL.PRETRAINED)\n\n    return model\n"
  },
  {
    "path": "libs/model/pose_resnet.py",
    "content": "# ------------------------------------------------------------------------------\n# Copyright (c) Microsoft\n# Licensed under the MIT License.\n# Written by Bin Xiao (Bin.Xiao@microsoft.com)\n# ------------------------------------------------------------------------------\n\nfrom __future__ import absolute_import\nfrom __future__ import division\nfrom __future__ import print_function\n\nimport os\nimport logging\n\nimport torch\nimport torch.nn as nn\n\n\nBN_MOMENTUM = 0.1\nlogger = logging.getLogger(__name__)\n\n\ndef conv3x3(in_planes, out_planes, stride=1):\n    \"\"\"3x3 convolution with padding\"\"\"\n    return nn.Conv2d(\n        in_planes, out_planes, kernel_size=3, stride=stride,\n        padding=1, bias=False\n    )\n\n\nclass BasicBlock(nn.Module):\n    expansion = 1\n\n    def __init__(self, inplanes, planes, stride=1, downsample=None):\n        super(BasicBlock, self).__init__()\n        self.conv1 = conv3x3(inplanes, planes, stride)\n        self.bn1 = nn.BatchNorm2d(planes, momentum=BN_MOMENTUM)\n        self.relu = nn.ReLU(inplace=True)\n        self.conv2 = conv3x3(planes, planes)\n        self.bn2 = nn.BatchNorm2d(planes, momentum=BN_MOMENTUM)\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\n\nclass Bottleneck(nn.Module):\n    expansion = 4\n\n    def __init__(self, inplanes, planes, stride=1, downsample=None):\n        super(Bottleneck, self).__init__()\n        self.conv1 = nn.Conv2d(inplanes, planes, kernel_size=1, bias=False)\n        self.bn1 = nn.BatchNorm2d(planes, momentum=BN_MOMENTUM)\n        self.conv2 = nn.Conv2d(planes, planes, kernel_size=3, stride=stride,\n                               padding=1, bias=False)\n        self.bn2 = nn.BatchNorm2d(planes, momentum=BN_MOMENTUM)\n        self.conv3 = nn.Conv2d(planes, planes * self.expansion, kernel_size=1,\n                               bias=False)\n        self.bn3 = nn.BatchNorm2d(planes * self.expansion,\n                                  momentum=BN_MOMENTUM)\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        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\n\nclass PoseResNet(nn.Module):\n\n    def __init__(self, block, layers, cfg, **kwargs):\n        self.inplanes = 64\n        extra = cfg.MODEL.EXTRA\n        self.deconv_with_bias = extra.DECONV_WITH_BIAS\n\n        super(PoseResNet, self).__init__()\n        self.conv1 = nn.Conv2d(3, 64, kernel_size=7, stride=2, padding=3,\n                               bias=False)\n        self.bn1 = nn.BatchNorm2d(64, momentum=BN_MOMENTUM)\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])\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        # used for deconv layers\n        self.deconv_layers = self._make_deconv_layer(\n            extra.NUM_DECONV_LAYERS,\n            extra.NUM_DECONV_FILTERS,\n            extra.NUM_DECONV_KERNELS,\n        )\n\n        self.final_layer = nn.Conv2d(\n            in_channels=extra.NUM_DECONV_FILTERS[-1],\n            out_channels=cfg.MODEL.NUM_JOINTS,\n            kernel_size=extra.FINAL_CONV_KERNEL,\n            stride=1,\n            padding=1 if extra.FINAL_CONV_KERNEL == 3 else 0\n        )\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, momentum=BN_MOMENTUM),\n            )\n\n        layers = []\n        layers.append(block(self.inplanes, planes, stride, downsample))\n        self.inplanes = planes * block.expansion\n        for i in range(1, blocks):\n            layers.append(block(self.inplanes, planes))\n\n        return nn.Sequential(*layers)\n\n    def _get_deconv_cfg(self, deconv_kernel, index):\n        if deconv_kernel == 4:\n            padding = 1\n            output_padding = 0\n        elif deconv_kernel == 3:\n            padding = 1\n            output_padding = 1\n        elif deconv_kernel == 2:\n            padding = 0\n            output_padding = 0\n\n        return deconv_kernel, padding, output_padding\n\n    def _make_deconv_layer(self, num_layers, num_filters, num_kernels):\n        assert num_layers == len(num_filters), \\\n            'ERROR: num_deconv_layers is different len(num_deconv_filters)'\n        assert num_layers == len(num_kernels), \\\n            'ERROR: num_deconv_layers is different len(num_deconv_filters)'\n\n        layers = []\n        for i in range(num_layers):\n            kernel, padding, output_padding = \\\n                self._get_deconv_cfg(num_kernels[i], i)\n\n            planes = num_filters[i]\n            layers.append(\n                nn.ConvTranspose2d(\n                    in_channels=self.inplanes,\n                    out_channels=planes,\n                    kernel_size=kernel,\n                    stride=2,\n                    padding=padding,\n                    output_padding=output_padding,\n                    bias=self.deconv_with_bias))\n            layers.append(nn.BatchNorm2d(planes, momentum=BN_MOMENTUM))\n            layers.append(nn.ReLU(inplace=True))\n            self.inplanes = planes\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.deconv_layers(x)\n        x = self.final_layer(x)\n\n        return x\n\n    def init_weights(self, pretrained=''):\n        if os.path.isfile(pretrained):\n            logger.info('=> init deconv weights from normal distribution')\n            for name, m in self.deconv_layers.named_modules():\n                if isinstance(m, nn.ConvTranspose2d):\n                    logger.info('=> init {}.weight as normal(0, 0.001)'.format(name))\n                    logger.info('=> init {}.bias as 0'.format(name))\n                    nn.init.normal_(m.weight, std=0.001)\n                    if self.deconv_with_bias:\n                        nn.init.constant_(m.bias, 0)\n                elif isinstance(m, nn.BatchNorm2d):\n                    logger.info('=> init {}.weight as 1'.format(name))\n                    logger.info('=> init {}.bias as 0'.format(name))\n                    nn.init.constant_(m.weight, 1)\n                    nn.init.constant_(m.bias, 0)\n            logger.info('=> init final conv weights from normal distribution')\n            for m in self.final_layer.modules():\n                if isinstance(m, nn.Conv2d):\n                    # nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')\n                    logger.info('=> init {}.weight as normal(0, 0.001)'.format(name))\n                    logger.info('=> init {}.bias as 0'.format(name))\n                    nn.init.normal_(m.weight, std=0.001)\n                    nn.init.constant_(m.bias, 0)\n\n            pretrained_state_dict = torch.load(pretrained)\n            logger.info('=> loading pretrained model {}'.format(pretrained))\n            self.load_state_dict(pretrained_state_dict, strict=False)\n        else:\n            logger.info('=> init weights from normal distribution')\n            for m in self.modules():\n                if isinstance(m, nn.Conv2d):\n                    # nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')\n                    nn.init.normal_(m.weight, std=0.001)\n                    # nn.init.constant_(m.bias, 0)\n                elif isinstance(m, nn.BatchNorm2d):\n                    nn.init.constant_(m.weight, 1)\n                    nn.init.constant_(m.bias, 0)\n                elif isinstance(m, nn.ConvTranspose2d):\n                    nn.init.normal_(m.weight, std=0.001)\n                    if self.deconv_with_bias:\n                        nn.init.constant_(m.bias, 0)\n\n\nresnet_spec = {\n    18: (BasicBlock, [2, 2, 2, 2]),\n    34: (BasicBlock, [3, 4, 6, 3]),\n    50: (Bottleneck, [3, 4, 6, 3]),\n    101: (Bottleneck, [3, 4, 23, 3]),\n    152: (Bottleneck, [3, 8, 36, 3])\n}\n\n\ndef get_pose_net(cfg, is_train, **kwargs):\n    num_layers = cfg.MODEL.EXTRA.NUM_LAYERS\n\n    block_class, layers = resnet_spec[num_layers]\n\n    model = PoseResNet(block_class, layers, cfg, **kwargs)\n\n    if is_train and cfg.MODEL.INIT_WEIGHTS:\n        model.init_weights(cfg.MODEL.PRETRAINED)\n\n    return model\n"
  },
  {
    "path": "libs/optimizer/__init__.py",
    "content": "#!/usr/bin/env python3\n# -*- coding: utf-8 -*-\n\"\"\"\nEmpty file.\n\"\"\"\n\n\n"
  },
  {
    "path": "libs/parser/__init__.py",
    "content": "#!/usr/bin/env python3\n# -*- coding: utf-8 -*-\n\"\"\"\nEmpty file.\n\"\"\"\n\n\n"
  },
  {
    "path": "libs/parser/parse.py",
    "content": "import argparse\n\ndef parse_arg():\n    parser = argparse.ArgumentParser(description='2Dto3Dnet.py')\n    ## paths\n    parser.add_argument('-save_root', type=str, default='../model/')\n    ##-----------------------------------------------------------------------##\n    ## model settings\n    parser.add_argument('-save_name', type=str, default=None)\n    # feed the current estimated 3D poses to the next stage\n    parser.add_argument('-refine_3d', type=bool, default=False)\n    parser.add_argument('-norm_twoD', type=bool, default=False)\n    parser.add_argument('-num_blocks', type=int, default=2)\n    # how many stages used for boosted regression\n    parser.add_argument('-num_stages', type=int, default=2)   \n    # the length of 3D pose representation used in the network\n    parser.add_argument('-linear_size', type=int, default=1024)  \n    # extra name for logging\n    parser.add_argument('-extra_str', type=str, default='')\n    # dropout\n    parser.add_argument('-dropout', type=float, default=0.5)\n    # leaky ReLu\n    parser.add_argument('-leaky', type=bool, default=False)\n    ##-----------------------------------------------------------------------##    \n    ## training settings\n    parser.add_argument('-batch_size', type=int, default=8192)\n    # random seed for reproduction of experiments\n    parser.add_argument('-seed', type=int, default=2019)\n    # number of threads to use when loading data\n    parser.add_argument('-num_threads', type=int, default=4)\n    # update leaf node distribution every certain number of network training\n    parser.add_argument('-gpuid', type=int, default=0)\n    parser.add_argument('-epochs', type=int, default=200)\n    # report_every: 10\n    parser.add_argument('-report_every', type=int, default=100)\n    # whether to perform evaluation on evaluation set during training\n    parser.add_argument('-eval', type=bool, default=False)\n    # whether to evaluate for each action during the training\n    parser.add_argument('-eval_action_wise', type=bool, default=True)   \n    # what protocol to use for evaluation\n    parser.add_argument('-protocols', type=list, default=['P1', 'P2'])\n    # whether to record and report loss history at the end of training  \n    parser.add_argument('-eval_every', type=int, default=350)\n    # path to the human3.6M dataset\n    parser.add_argument('-data_dir', type=str, default='../data/human3.6M/')\n    # actions to use for training\n    parser.add_argument('-actions', type=str, default='All')\n    # whether to do data augmentation for the training data\n    parser.add_argument('-augmentation', type=bool, default=False)    \n    # using virtual cameras\n    parser.add_argument('-virtual_cams', type=bool, default=False)    \n    # interpolate between 3D joints\n    parser.add_argument('-interpolate', type=bool, default=False)    \n    # what input to use, synthetic or detected\n    parser.add_argument('-twoD_source', type=str, default='synthetic')\n    # what dataset to use as the evaluation set\n    parser.add_argument('-test_source', type=str, default='h36m')\n    # whether to use pre-augmented training data\n    parser.add_argument('-pre_aug', type=bool, default=False)\n    # the path to the pre-augmented dataset\n    parser.add_argument('-pre_aug_dir', type=str, default='../data/augmented_evo_10.npy')\n    # the path of pre-trained check-point\n    parser.add_argument('-ckpt_dir', type=str)\n    ##-----------------------------------------------------------------------##    \n    ## dataset settings\n    # whether to only predict 14 joints\n    parser.add_argument('-pred14', type=bool, default=False)   \n    # whether to add 3D poses fitted by SMPL model\n    parser.add_argument('-SMPL', type=bool, default=False)      \n    # perform normalization for each pose instead of all the poses\n    parser.add_argument('-norm_single', type=bool, default=False) \n    # how much weight is given to the new poses of SMPL\n    parser.add_argument('-SMPL_weight', type=float, default=0.5)   \n    # whether to change the image size of the cameras\n    parser.add_argument('-change_size', type=bool, default=False)\n    # virtual image size if changed\n    parser.add_argument('-vir_img_size', type=int, default=256)\n    # use only a subset of training examples for weakly-supervised experiments\n    parser.add_argument('-ws', type=bool, default=False)\n    # the path to evolved training examples for weakly-supervised experiments\n    parser.add_argument('-evolved_path', type=str, default=None)\n    # the training sample used if no path is provided\n    parser.add_argument('-ws_name', type=str, default='S1')\n    # whether to visualize the dataset\n    parser.add_argument('-visualize', type=bool, default=False)\n    # whether to show the ambiguous pairs in the dataset\n    parser.add_argument('-show_ambi', type=bool, default=False)    \n    ##-----------------------------------------------------------------------##    \n    # Optimizer settings\n    parser.add_argument('-optim_type', type=str, default='adam')\n    parser.add_argument('-lr', type=float, default=0.001, help=\"sgd: 0.5, adam: 0.001\")\n    parser.add_argument('-weight_decay', type=float, default=0.0)\n    parser.add_argument('-momentum', type=float, default=0.9, help=\"sgd: 0.9\")\n    # reduce the learning rate after each milestone\n    #parser.add_argument('-milestones', type=list, default=[6, 12, 18])\n    parser.add_argument('-milestones', type=list, default=[50, 100, 150])\n    # how much to reduce the learning rate\n    parser.add_argument('-gamma', type=float, default=1)\n    ##-----------------------------------------------------------------------##  \n    ## usage configuration\n    # whether to train a model or deploy a trained model\n    parser.add_argument('-train', type=bool, default=False)\n    parser.add_argument('-evaluate', type=bool, default=False)\n    # evaluate a batch of models\n    parser.add_argument('-evaluate_batch', type=bool, default=False)\n    # whether to save the trained model\n    parser.add_argument('-save', type=bool, default=True)    \n    # evaluate for each action\n    parser.add_argument('-evaluate_action', type=bool, default=True)\n    parser.add_argument('-produce', type=bool, default=False)\n    opt = parser.parse_args()\n    return opt"
  },
  {
    "path": "libs/skeleton/__init__.py",
    "content": "#!/usr/bin/env python3\n# -*- coding: utf-8 -*-\n\"\"\"\nEmpty file.\n\"\"\"\n\n\n"
  },
  {
    "path": "libs/skeleton/anglelimits.py",
    "content": "\"\"\"\nUtility functions for the hierarchical human representation. \nA Python implementation for pose-conditioned joint angle limits is also included.\nReference: \"Pose-Conditioned Joint Angle Limits for 3D Human Pose Reconstruction\"\n\"\"\"\nimport logging\nimport os\nimport numpy as np\nimport scipy.io as sio\nimport matplotlib.pyplot as plt\nfrom mpl_toolkits.mplot3d import Axes3D\n#=============================================================================#\n# Load the joint angle constraints\n# These files are directly converted from .mat to .npy    \n# The MATLAB implementation of the CVPR 15 paper has detailed documentation.\nroot = \"../resources/constraints\"\nlogging.info(\"Loading files from \" + root)\nmodel_path = os.path.join(root, \"jointAngleModel_v2.npy\")\njoint_angle_limits = np.load(model_path, allow_pickle=True).item()\nangle_spread = joint_angle_limits['angleSprd']\n# separation plane for conditional joint angle\nsepPlane = joint_angle_limits['sepPlane']\nE2 = joint_angle_limits['E2']\nbounds = joint_angle_limits['bounds']\n# static pose and parameters used in coordinate transformation\nstatic_pose_path = os.path.join(root, \"staticPose.npy\")\nstatic_pose = np.load(static_pose_path, allow_pickle=True).item()\ndi = static_pose['di']\na = static_pose['a'].reshape(3)\n# load the pre-computed conditinal distribution \ncon_dis_path = os.path.join(root, \"conditional_dis.npy\")\ncon_dis = np.load(con_dis_path, allow_pickle=True).item()\n#=============================================================================#\n# joint names of the CVPR 15 paper\nPRIOR_NAMES = ['back-bone', \n               'R-shldr', \n               'R-Uarm', \n               'R-Larm', \n               'L-shldr', \n               'L-Uarm', \n               'L-Larm', \n               'head', \n               'R-hip', \n               'R-Uleg', \n               'R-Lleg', \n               'R-feet', \n               'L-hip',\n               'L-Uleg', \n               'L-Lleg', \n               'L-feet'\n               ]\n# Human 3.6M joint names are slightly different from the above\nH36M_NAMES = ['']*32\nH36M_NAMES[0]  = 'Hip'\nH36M_NAMES[1]  = 'RHip'\nH36M_NAMES[2]  = 'RKnee'\nH36M_NAMES[3]  = 'RFoot'\nH36M_NAMES[6]  = 'LHip'\nH36M_NAMES[7]  = 'LKnee'\nH36M_NAMES[8]  = 'LFoot'\nH36M_NAMES[12] = 'Spine'\nH36M_NAMES[13] = 'Thorax'\nH36M_NAMES[14] = 'Neck/Nose'\nH36M_NAMES[15] = 'Head'\nH36M_NAMES[17] = 'LShoulder'\nH36M_NAMES[18] = 'LElbow'\nH36M_NAMES[19] = 'LWrist'\nH36M_NAMES[25] = 'RShoulder'\nH36M_NAMES[26] = 'RElbow'\nH36M_NAMES[27] = 'RWrist'\n# correspondence of the joints \n# (key, value) -> (index in prior_names, index in H36M names)\ncorrespondence = {0:12, 1:13, 2:25, 3:26, 4:27, 5:17, 6:18, 7:19, 8:15, 9:1,\n                  10:2, 11:3, 13:6, 14:7, 15:8}\n# number of bone vectors attached to a torso\nnum_of_bones = 9        \n# descretization of spherical coordinates\n# bin edges for theta\ntheta_edges = np.arange(0.5, 122, 1)# theta values: 1 to 121 (integer)\n# bin edges for phi\nphi_edges = np.arange(0.5, 62, 1) # phi values: 1 to 61\n# color map used for visualization\ncmap = plt.cm.RdYlBu\n# indices used for computing bone vectors for non-torso bones\nnt_parent_indices = [13, 17, 18, 25, 26, 6, 7, 1, 2, 13]\nnt_child_indices = [15, 18, 19, 26, 27, 7, 8, 2, 3, 14]\n# map from bone index to the parent's di index\n# TODO\ndi_indices = {2:5, 4:2, 6:13, 8:9}\n# map from angle index to record index\nrecord_indices = {0:4, 1:2, 3:0, 5:8, 7:5, 2:3, 4:1, 6:9, 8:6}\n# name for the bone vectors\nbone_name = {\n 1: 'thorax to head top',\n 2: 'left shoulder to left elbow',\n 3: 'left elbow to left wrist',\n 4: 'right shoulder to right elbow',\n 5: 'right elbow to right wrist',\n 6: 'left hip to left knee',\n 7: 'left knee to left ankle',\n 8: 'right hip to right knee',\n 9: 'right knee to right ankle'\n}       \n#=============================================================================#\ndef is_valid_local(skeleton_local, return_ang = False):\n    \"\"\"\n    Check if the limbs represented in local coordinate system are valid or not.\n    \"\"\"\n    valid_vec = np.ones((num_of_bones), dtype=np.bool)\n    angles = to_spherical(skeleton_local)\n    angles[:,1:] *= 180/np.pi\n    # convert to valid range and discretize\n    # theta: -180~180 degrees discretized into 120 bins\n    # phi: -90~90 degrees discretized into 60 bins\n    angles[:, 1] = np.floor((angles[:, 1]+180)/3 + 1)\n    angles[:, 2] = np.floor((angles[:, 2]+90)/3 + 1)\n    # go through each bone and check the angle-limits\n    for angle_idx in range(len(angles)):\n        angle = angles[angle_idx]\n        record_idx = record_indices[angle_idx]        \n        theta, phi = int(angle[1]), int(angle[2])\n        if angle_idx in [0, 1, 3, 5, 7]:       \n            test_value = angle_spread[0, record_idx][theta-1, phi-1]\n            if test_value == 0:\n                valid_vec[angle_idx] == False\n        else:\n            angle_parent = angles[angle_idx - 1]\n            theta_p, phi_p = int(angle_parent[1]), int(angle_parent[2])\n            vector = normalize(sepPlane[0, record_idx][theta_p-1, phi_p-1])\n            for value in vector:\n                if np.isnan(value):\n                    valid_vec[angle_idx] = False\n                    continue\n            if np.dot(np.hstack([skeleton_local[angle_idx], 1]), vector) > 0:\n                valid_vec[angle_idx] = False\n            else:\n                e1 = vector[:-1]\n                e2 = E2[0, record_idx][theta_p-1, phi_p-1]\n                T = gram_schmidt_columns(np.hstack([e1.reshape(3,1),\n                                                    e2.reshape(3,1),\n                                                    np.cross(e1,e2).reshape(3,1)]))\n                bnd = bounds[0, record_idx][theta_p-1, phi_p-1]\n                u = (T[:, 1:]).T @ skeleton_local[angle_idx]\n                if u[0] < bnd[0] or u[0] > bnd[1] or u[1] < bnd[2] or u[1] > bnd[3]:\n                    valid_vec[angle_idx] = False    \n    if return_ang:\n        return valid_vec, angles\n    else:\n        return valid_vec\n\ndef is_valid(skeleton, return_ang = False, camera = False):\n    \"\"\"\n    args:\n        skeleton: input skeleton of shape [num_joints, 3] use the annotation \n        of Human 3.6M dataset\n    return:\n        valid_vec: boolean vector specifying the validity for each bone. \n        return 0 for invalid bones.\n        camera: relative orientation of camera and human\n    \"\"\"\n    \n    skeleton = skeleton.reshape(len(H36M_NAMES), -1)\n    # the ordering of coordinate used by the Prior was x,z and y\n    skeleton = skeleton[:, [0,2,1]]\n    # convert bone vectors into local coordinate\n    skeleton_local = to_local(skeleton)\n    ret = is_valid_local(skeleton_local, return_ang=return_ang)\n    if return_ang:\n        return ret[0], ret[1]\n    else:\n        return ret\n\ndef normalize(vector):\n    \"\"\"\n    Normalize a vector.\n    \"\"\"\n    return vector/np.linalg.norm(vector)\n\ndef to_spherical(xyz):\n    \"\"\"\n    Convert from Cartisian coordinate to spherical coordinate\n    theta: [-pi, pi]\n    phi: [-pi/2, pi/2]\n    note that xyz should be float number\n    \"\"\"\n    # return in r, phi, and theta (elevation angle from z axis down)\n    return_value = np.zeros(xyz.shape, dtype=xyz.dtype)\n    xy = xyz[:,0]**2 + xyz[:,1]**2\n    return_value[:,0] = np.sqrt(xy + xyz[:,2]**2) # r      \n    return_value[:,1] = np.arctan2(xyz[:,1], xyz[:,0]) # theta\n    return_value[:,2] = np.arctan2(xyz[:,2], np.sqrt(xy)) # phi\n    return return_value\n\ndef to_xyz(rthetaphi):\n    \"\"\"\n    Convert from spherical coordinate to Cartisian coordinate\n    theta: [0, 2*pi] or [-pi, pi]\n    phi: [-pi/2, pi/2]\n    \"\"\"\n    return_value = np.zeros(rthetaphi.shape, dtype=rthetaphi.dtype)\n    sintheta = np.sin(rthetaphi[:,1])\n    costheta = np.cos(rthetaphi[:,1])\n    sinphi = np.sin(rthetaphi[:,2])\n    cosphi = np.cos(rthetaphi[:,2])\n    return_value[:,0] = rthetaphi[:,0]*costheta*cosphi # x\n    return_value[:,1] = rthetaphi[:,0]*sintheta*cosphi # y\n    return_value[:,2] = rthetaphi[:,0]*sinphi #z\n    return return_value\n\ndef test_coordinate_conversion():\n    # theta: [-pi, pi] reference\n    xyz = np.random.rand(1, 3)*2 - 1\n    rthetaphi = to_spherical(xyz)\n    xyz2 = to_xyz(rthetaphi)\n    print('maximum error:', np.max(np.abs(xyz - xyz2)))\n    # theta: [0, 2*pi] reference\n    xyz = np.random.rand(1, 3)*2 - 1\n    rthetaphi = to_spherical(xyz)\n    indices = rthetaphi[:,1] < 0\n    rthetaphi[:,1][indices] += 2*np.pi\n    xyz2 = to_xyz(rthetaphi)\n    print('maximum error:', np.max(np.abs(xyz - xyz2)))    \n    return\n\ndef gram_schmidt_columns(X):\n    \"\"\"\n    Apply Gram-Schmidt orthogonalization to obtain basis vectors.\n    \"\"\"\n    B = np.zeros(X.shape)\n    B[:, 0] = (1/np.linalg.norm(X[:, 0]))*X[:, 0]\n    for i in range(1, 3):\n        v = X[:, i]\n        U = B[:, 0:i] # subspace basis which has already been orthonormalized\n        pc = U.T @ v # orthogonal projection coefficients of v onto U\n        p = U@pc\n        v = v - p\n        if np.linalg.norm(v) < 2e-16:\n            # vectors are not linearly independent!\n            raise ValueError\n        else:\n            v = normalize(v)\n            B[:, i] = v\n    return B\n\ndef direction_check(system, v1, v2, v3):\n    if system[:,0].dot(v1) <0:\n        system[:,0] *= -1\n    if system[:,1].dot(v2) <0:\n        system[:,1] *= -1\n    if system[:,2].dot(v3) <0:\n        system[:,2] *= -1        \n    return system\n\ndef get_normal(x1, a, x):\n    \"\"\"\n    Get normal vector.\n    \"\"\"\n    nth = 1e-4\n    # x and a are parallel\n    if np.linalg.norm(x - a) < nth or np.linalg.norm(x + a) < nth:   \n        n = np.cross(x, x1)\n        flag = True\n    else:\n        n = np.cross(a, x)\n        flag = False\n    return normalize(n), flag\n\ndef get_basis1(skeleton):\n    \"\"\"\n    Compute local coordinate system from 3D joint positions.\n    This system is used for upper-limbs.\n    \"\"\"\n    # compute the vector from the left shoulder to the right shoulder\n    left_shoulder = skeleton[17]\n    right_shoulder = skeleton[25]\n    v1 = normalize(right_shoulder - left_shoulder)    \n    # compute the backbone vector from the thorax to the spine \n    thorax = skeleton[13]\n    spine = skeleton[12]\n    v2 = normalize(spine - thorax)\n    # v3 is the cross product of v1 and v2 (front-facing vector for upper-body)\n    v3 = normalize(np.cross(v1, v2))    \n    return v1, v2, v3\n\ndef to_local(skeleton):\n    \"\"\"\n    Represent the bone vectors in the local coordinate systems.\n    \"\"\"\n    v1, v2, v3 = get_basis1(skeleton)\n    # compute the vector from the left hip to the right hip\n    left_hip = skeleton[6]\n    right_hip = skeleton[1]\n    v4 = normalize(right_hip - left_hip)\n    # v5 is the cross product of v4 and v2 (front-facing vector for lower-body)\n    v5 = normalize(np.cross(v4, v2))\n    # compute orthogonal coordinate systems using GramSchmidt\n    # for upper body, we use v1, v2 and v3\n    system1 = gram_schmidt_columns(np.hstack([v1.reshape(3,1), \n                                              v2.reshape(3,1), \n                                              v3.reshape(3,1)]))\n    # make sure the directions rougly align\n    #system1 = direction_check(system1, v1, v2, v3)\n    # for lower body, we use v4, v2 and v5\n    system2 = gram_schmidt_columns(np.hstack([v4.reshape(3,1), \n                                              v2.reshape(3,1), \n                                              v5.reshape(3,1)]))\n    #system2 = direction_check(system2, v4, v2, v5)\n\n    bones = skeleton[nt_parent_indices, :] - skeleton[nt_child_indices, :]    \n    # convert bone vector to local coordinate system\n    bones_local = np.zeros(bones.shape, dtype=bones.dtype)\n    for bone_idx in range(len(bones)):\n        # only compute bone vectors for non-torsos\n        # the order of the non-torso bone vector is: \n        # bone vector1: thorax to head top\n        # bone vector2: left shoulder to left elbow\n        # bone vector3: left elbow to left wrist\n        # bone vector4: right shoulder to right elbow\n        # bone vector5: right elbow to right wrist\n        # bone vector6: left hip to left knee\n        # bone vector7: left knee to left ankle\n        # bone vector8: right hip to right knee\n        # bone vector9: right knee to right ankle        \n        bone = normalize(bones[bone_idx])\n        if bone_idx in [0, 1, 3, 5, 7]:\n            # bones that are directly connected to the torso\n            if bone_idx in [0, 1, 3]:\n                # upper body\n                bones_local[bone_idx] = system1.T @ bone\n            else:\n                # lower body\n                bones_local[bone_idx] = system2.T @ bone\n        else:\n            if bone_idx in [2, 4]:\n                parent_R = system1\n            else:\n                parent_R = system2\n            # parent bone index is smaller than 1\n            vector_u = normalize(bones[bone_idx - 1])\n            di_index = di_indices[bone_idx]\n            vector_v, flag = get_normal(parent_R@di[:, di_index],\n                                        parent_R@a,\n                                        vector_u\n                                        )\n            vector_w = np.cross(vector_u, vector_v)\n            local_system = gram_schmidt_columns(np.hstack([vector_u.reshape(3,1),\n                                                           vector_v.reshape(3,1),\n                                                           vector_w.reshape(3,1)]\n                                                          )\n                                                )\n            bones_local[bone_idx] = local_system.T @ bone\n    return bones_local\n\ndef to_global(skeleton, bones_local, cache=False):\n    \"\"\"\n    Convert local coordinate back into global coordinate system.\n    cache: return intermeadiate results\n    \"\"\"\n    return_value = {}\n    v1, v2, v3 = get_basis1(skeleton)\n    # compute the vector from the left hip to the right hip\n    left_hip = skeleton[6]\n    right_hip = skeleton[1]\n    v4 = normalize(right_hip - left_hip)\n    # v5 is the cross product of v4 and v2 (front-facing vector for lower-body)\n    v5 = normalize(np.cross(v4, v2))\n    # compute orthogonal coordinate systems using GramSchmidt\n    # for upper body, we use v1, v2 and v3\n    system1 = gram_schmidt_columns(np.hstack([v1.reshape(3,1), \n                                              v2.reshape(3,1), \n                                              v3.reshape(3,1)]))\n    # make sure the directions rougly align\n    #system1 = direction_check(system1, v1, v2, v3)\n    # for lower body, we use v4, v2 and v5\n    system2 = gram_schmidt_columns(np.hstack([v4.reshape(3,1), \n                                              v2.reshape(3,1), \n                                              v5.reshape(3,1)]))\n    #system2 = direction_check(system2, v4, v2, v5)\n    if cache:\n        return_value['cache'] = [system1, system2]\n        return_value['bl'] = bones_local\n        \n    bones_global = np.zeros(bones_local.shape)   \n    # convert bone vector to local coordinate system\n    for bone_idx in [0,1,3,5,7,2,4,6,8]:\n        # the indices follow the order from torso to limbs\n        # only compute bone vectors for non-torsos      \n        bone = normalize(bones_local[bone_idx])\n        if bone_idx in [0, 1, 3, 5, 7]:\n            # bones that are directly connected to the torso\n            if bone_idx in [0, 1, 3]:\n                # upper body\n                # this is the inverse transformation compared to the to_local \n                # function\n                bones_global[bone_idx] = system1 @ bone\n            else:\n                # lower body\n                bones_global[bone_idx] = system2 @ bone\n        else:\n            if bone_idx in [2, 4]:\n                parent_R = system1\n            else:\n                parent_R = system2\n            # parent bone index is smaller than 1\n            vector_u = normalize(bones_global[bone_idx - 1])\n            di_index = di_indices[bone_idx]\n            vector_v, flag = get_normal(parent_R@di[:, di_index],\n                                  parent_R@a,\n                                  vector_u)\n            vector_w = np.cross(vector_u, vector_v)\n            local_system = gram_schmidt_columns(np.hstack([vector_u.reshape(3,1), \n                                              vector_v.reshape(3,1), \n                                              vector_w.reshape(3,1)]))\n            if cache:\n                return_value['cache'].append(local_system)\n            bones_global[bone_idx] = local_system @ bone\n    return_value['bg'] = bones_global\n    return return_value\n\ndef test_global_local_conversion():\n    \"\"\"\n    test for global and lobal coordinate conversion\n    \"\"\"\n    path='Your3DSkeleton.npy'\n    index = 0\n    pose = np.load(path, allow_pickle=True)[index] \n    pose = pose.reshape(32, -1)    \n    global_bones = pose[nt_parent_indices, :] - pose[nt_child_indices, :]    \n    for bone_idx in range(len(global_bones)):    \n        global_bones[bone_idx] = normalize(global_bones[bone_idx])    \n    local_c = to_local(pose)\n    global_c = to_global(pose, local_c)['bg']\n    maximum_error = np.max(np.abs(global_bones - global_c))\n    print('maximum error', maximum_error)\n    return maximum_error\n\ndef show3Dpose(channels, ax, lcolor=\"#3498db\", rcolor=\"#e74c3c\", add_labels=True,\n               gt=False,pred=False,inv_z=False): # blue, orange\n\n    vals = np.reshape( channels, (32, -1) )\n\n    I   = np.array([1,2,3,1,7,8,1, 13,14,15,14,18,19,14,26,27])-1 # start points\n    J   = np.array([2,3,4,7,8,9,13,14,15,16,18,19,20,26,27,28])-1 # end points\n    LR  = np.array([1,1,1,0,0,0,0, 0, 0, 0, 0, 0, 0, 1, 1, 1], dtype=bool)\n\n    # Make connection matrix\n    for i in np.arange( len(I) ):\n        x, y, z = [np.array( [vals[I[i], j], vals[J[i], j]] ) for j in range(3)]\n        if gt:\n            ax.plot(x,y, z,  lw=2, c='k')\n        \n        elif pred:\n            ax.plot(x,y, z,  lw=2, c='r')\n\n        else:\n            ax.plot(x,y, z,  lw=2, c=lcolor if LR[i] else rcolor)\n\n    RADIUS = 750 # space around the subject\n    xroot, yroot, zroot = vals[0,0], vals[0,1], vals[0,2]\n    ax.set_xlim3d([-RADIUS+xroot, RADIUS+xroot])\n    ax.set_zlim3d([-RADIUS+zroot, RADIUS+zroot])\n    ax.set_ylim3d([-RADIUS+yroot, RADIUS+yroot])\n\n#    if add_labels:\n#        ax.set_xlabel(\"x\")\n#        ax.set_ylabel(\"z\")\n#        ax.set_zlabel(\"y\")\n\n    if add_labels:\n        ax.set_xlabel(\"x\")\n        ax.set_ylabel(\"y\")\n        ax.set_zlabel(\"z\")\n        \n    ax.set_aspect('equal')\n\n  # Get rid of the panes (actually, make them white)\n    white = (1.0, 1.0, 1.0, 0.0)\n    ax.w_xaxis.set_pane_color(white)\n    ax.w_yaxis.set_pane_color(white)\n  # Keep z pane\n\n  # Get rid of the lines in 3d\n    ax.w_xaxis.line.set_color(white)\n    ax.w_yaxis.line.set_color(white)\n    ax.w_zaxis.line.set_color(white)\n    if inv_z:\n        ax.invert_zaxis()  \n        \ndef get_histogram2d(angles, validmap=None, smooth=True):\n    \"\"\"\n    Obtain a 2D histogram for discretized joint angles\n    \"\"\"\n    H, xedges, yedges = np.histogram2d(angles[:,0], \n                                       angles[:,1], \n                                       bins=(theta_edges, phi_edges)\n                                       )    \n    if validmap is not None:\n        # rule out outliers\n        mask = validmap != 0\n        H = H * mask\n    H = H.reshape(-1)\n    H = H/H.sum()\n    return H\n\ndef sample_from_histogram(histogram, \n                          x=np.arange(1,122,1), \n                          y=np.arange(1,62,1), \n                          total=1000, \n                          add_noise=False, \n                          bin_size=3\n                          ):\n    \"\"\"\n    Sample from a pre-computed histogram.\n    \"\"\"\n    assert histogram.shape[0] == len(x)\n    assert histogram.shape[1] == len(y)\n    # normalize the histogram\n    histogram = histogram/histogram.sum()\n    # multiply to get final counts\n    histogram = histogram*total\n    none_zeros = histogram!=0\n    histogram = histogram.astype(np.int)     \n    histogram[none_zeros] = histogram[none_zeros] + 1\n    data = []\n    for x_id in x:\n        for y_id in y:\n            counts = histogram[x_id-1, y_id-1]\n            if counts!=0:\n                temp = np.array([[x[x_id - 1], y[y_id - 1]]])\n                data.append(np.repeat(temp, counts, axis=0))\n    data = np.vstack(data)\n    if add_noise:\n        noise = np.random.rand(*(data.shape))*bin_size\n        data = (data - 1)*bin_size + noise\n    return data[:total,:]\n\ndef histogram_transform(histogram, gamma):\n    \"\"\"\n    Transform a distribution with power function.\n    \"\"\"\n    histogram = (histogram - histogram.min())/(histogram.max() - histogram.min())\n    histogram = np.power(histogram, gamma)\n    return histogram\n\n#=============================================================================#\n# Visualization utilities.\ndef smooth_histogram2d(data):\n    \"\"\"\n    Smooth a 2D histogram with kernel density estimation.\n    \"\"\"\n    from scipy.stats import kde    \n    # Evaluate a gaussian kde on a regular grid of nbins x nbins over data extents\n    k = kde.gaussian_kde(data.T)\n    xi, yi = np.mgrid[1:122, 1:62]\n    zi = k(np.vstack([xi.flatten(), yi.flatten()]))\n    # change the extent\n    xi = xi*3 - 180\n    yi = yi*3 - 180\n    fig, axes = plt.subplots(ncols=1, nrows=3)\n    # plot a density\n    axes[0].set_title('Calculate Gaussian KDE')\n    axes[0].pcolormesh(xi, yi, zi.reshape(xi.shape), cmap=plt.cm.BuGn_r)\n    axes[0].set_aspect('equal')\n    axes[0].invert_yaxis()\n    # add shading\n    axes[1].set_title('2D Density with shading')\n    axes[1].pcolormesh(xi, yi, zi.reshape(xi.shape), shading='gouraud', cmap=plt.cm.BuGn_r)\n    axes[1].set_aspect('equal') \n    axes[1].invert_yaxis()\n    # contour\n    axes[2].set_title('Contour')\n    axes[2].pcolormesh(xi, yi, zi.reshape(xi.shape), shading='gouraud', cmap=plt.cm.BuGn_r)\n    axes[2].contour(xi, yi, zi.reshape(xi.shape))    \n    axes[2].set_aspect('equal')\n    axes[2].invert_yaxis()\n    return\n\ndef decorate_axis(ax, title=None):\n    ax.set_xlabel('Theta: -180 to 180')\n    ax.set_label('Phi:-90 to 90')\n    if title is not None:\n        ax.set_title(title)\n    return\n\ndef adjust_figure(left = 0, \n                  right = 1, \n                  bottom = 0.01, \n                  top = 0.95,\n                  wspace = 0, \n                  hspace = 0.4\n                  ):  \n    plt.subplots_adjust(left, bottom, right, top, wspace, hspace)\n    return\n\ndef plot_relative_poses(skeletons, cameras=None):\n    \"\"\"\n    Visualize the distribution of front vector in camera coordinate system.\n    \"\"\"\n    # skeletons: 3D poses in world or camera coordinates\n    # cameras: camera parameters\n    vector_list = []\n    if cameras is None:\n        for pose_id in range(len(skeletons)):\n            _, _, front_vector = get_basis1(skeletons[pose_id].reshape(32, -1))\n            vector_list.append(front_vector)\n    else:\n        raise NotImplementedError\n    vector_list = np.vstack(vector_list)\n    spherical = to_spherical(vector_list)\n    # convert to valid range and discretize\n    spherical[:, 1:] *= 180/np.pi\n    spherical[:, 1] = np.floor((spherical[:, 1]+180)/3 + 1)\n    spherical[:, 2] = np.floor((spherical[:, 2]+90)/3 + 1)    \n    H, xedges, yedges = np.histogram2d(spherical[:,1], spherical[:,2], \n                                       bins=(theta_edges, phi_edges))    \n    plt.figure()\n    ax = plt.subplot(111)\n    plt.imshow(H.T, extent=[-180, 180, -90, 90], interpolation='bilinear')\n    decorate_axis(ax, 'Relative camera pose in H36M')\n    return vector_list\n\ndef plot_distribution(H_temp, \n                      angle_idx, \n                      dataset_name, \n                      gamma, \n                      save_path='../viz'\n                      ):\n    \"\"\"\n    Visualize distribution of limb orientation.\n    \"\"\"\n    # angles: [n_samples, 2] in theta and phi order\n    # plot the distribution of local joint angles and overlay valid regions\n    plt.ioff()\n    if not os.path.exists(save_path):\n        os.makedirs(save_path)\n    H = H_temp.copy()\n    # normalize\n    H = (H-H.min())/(H.max()-H.min())\n    # perform \"gamma correction\"\n    H = np.power(H, gamma)\n    # normalize again\n    H = (H-H.min())/(H.max()-H.min())\n    H_return = H.copy().reshape(121,61)\n    # map to color \n    H = cmap(H).reshape((121, 61, 4))\n    H = [np.expand_dims(H[:,:,i].T, axis=2) for i in range(4)]\n    H = np.concatenate(H, axis=2)\n    if angle_idx in [0,1,3,5,7]:\n        record_idx = record_indices[angle_idx]\n        mask_temp = angle_spread[0, record_idx]\n        mask = np.zeros((61, 121, 4))\n        mask[:,:,3] = mask_temp.T\n        mask[:,:,1] = mask_temp.T\n        f = plt.figure(figsize=(5, 6))\n        ax = plt.subplot(211)\n        ax.imshow(H, extent=[-180, 180, -90, 90], interpolation='bilinear', alpha=1)    \n        decorate_axis(ax, 'Distribution of ' + dataset_name + ' for bone: ' + bone_name[angle_idx+1]) \n        ax = plt.subplot(212)\n        ax.imshow(mask, extent=[-180, 180, -90, 90],alpha=0.5)\n        ax.imshow(H, extent=[-180, 180, -90, 90], interpolation='bilinear', alpha=0.5)\n        decorate_axis(ax, 'Overlayed')\n        plt.tight_layout()\n    else:\n        f = plt.figure()\n        ax = plt.subplot(111)\n        ax.imshow(H, extent=[-180, 180, -90, 90], interpolation='bilinear', alpha=1)    \n        decorate_axis(ax, 'Distribution of ' + dataset_name + ' for bone: ' + bone_name[angle_idx+1])  \n        plt.tight_layout()\n    adjust_figure(left = 0.135, \n                  right = 0.95, \n                  bottom = 0.05, \n                  top = 1,\n                  wspace = 0, \n                  hspace = 0\n                  )\n    save_name = dataset_name + bone_name[angle_idx+1] + '_gamma_' + str(gamma) + '.jpg'\n    f.savefig(os.path.join(save_path, save_name))\n    plt.close(f)\n    return save_name, H_return\n\n#=============================================================================#\n# sampling utilities: sample 3D human skeleton from a pre-computed distribution\ntemplate = np.load(os.path.join(root, 'template.npy'), allow_pickle=True).reshape(32,-1)  \ntemplate_bones = template[nt_parent_indices, :] - template[nt_child_indices, :]  \ntemplate_bone_lengths = to_spherical(template_bones)[:, 0]\nnt_parent_indices = [13, 17, 18, 25, 26, 6, 7, 1, 2]\nnt_child_indices = [15, 18, 19, 26, 27, 7, 8, 2, 3]    \ndef get_skeleton(bones, pose, bone_length=template_bone_lengths):\n    \"\"\"\n    Update the non-torso limb of a skeleton by specifying bone vectors.\n    \"\"\"\n    new_pose = pose.copy()\n    for bone_idx in [0,1,3,5,7,2,4,6,8]:\n        new_pose[nt_child_indices[bone_idx]] = new_pose[nt_parent_indices[bone_idx]] \\\n        - bones[bone_idx]*bone_length[bone_idx]            \n    return new_pose\n\ndef test_get_skeleton():\n    pose = template\n    nt_parent_indices = [13, 17, 18, 25, 26, 6, 7, 1, 2]\n    nt_child_indices = [15, 18, 19, 26, 27, 7, 8, 2, 3]    \n    global_bones = pose[nt_parent_indices, :] - pose[nt_child_indices, :]    \n    for bone_idx in range(len(global_bones)):    \n        global_bones[bone_idx] = normalize(global_bones[bone_idx])    \n    local_c = to_local(pose)\n    global_c = to_global(pose, local_c)['bg']\n    new_pose = get_skeleton(global_c, pose)\n    maximum_error = np.max(np.abs(new_pose - pose))\n    print('maximum error', maximum_error)    \n    return maximum_error\n\ndef grow_from_torso(poses, all_angles, cache=False):\n    \"\"\"\n    Update the non-torso limb of a skeleton by specifying limb orientations.\n    \"\"\"\n    new_poses = poses.copy()\n    return_value = {}\n    if cache:\n        return_value['cache'] = []\n        return_value['bl'] = []\n    for pose_id in range(len(poses)):\n        pose = poses[pose_id].reshape(32,-1)\n        angles = all_angles[:, pose_id, :]\n        # convert to spherical coordinate in radians\n        spherical = np.ones((len(angles), 3))\n        spherical[:, 1:] = angles/180*np.pi\n        spherical[:, 1] -= np.pi\n        spherical[:, 2] -= np.pi/2\n        # convert to local cartisian coordinate\n        local_xyz = to_xyz(spherical)\n        # convert to global coordinate\n        return_value_temp = to_global(pose, local_xyz, cache=cache)\n        if cache:\n            return_value['cache'].append(return_value_temp['cache'])\n            return_value['bl'].append(return_value_temp['bl'])\n        global_xyz = return_value_temp['bg']\n        new_poses[pose_id] = get_skeleton(global_xyz, new_poses[pose_id].reshape(32,-1)).reshape(-1)\n    return_value['np'] = new_poses\n    return return_value\n\ndef test_grow_from_torso():\n    poses = template.reshape(1, 96)\n    _, angles = is_valid(poses.reshape(32, -1), return_ang=True)\n    local_coordinate = to_local(poses.reshape(32, -1))\n    angles = angles[:,1:].reshape(9, 1, 2)\n    # to degrees\n    angles *= 3\n    return_dic = grow_from_torso(poses, angles, cache=True)\n    new_poses = return_dic['np']\n    bones_local = return_dic['bl']\n    print(np.max(np.abs(local_coordinate - bones_local)))\n    maximum_error = np.max(np.abs(new_poses - poses))\n    print('maximum error1', maximum_error)       \n    _, new_angles = is_valid(new_poses, return_ang=True)\n    maximum_error = np.max(np.abs(angles[:,0,:]/3 - new_angles[:,1:]))\n    print('maximum error2', maximum_error)    \n    return \n\ndef sample_lower_limbs(angles, bin_size=3):\n    \"\"\"\n    Sample limb orientations for lower limbs.\n    \"\"\"\n    # angles of shape [num_bones, sample_size, 2]\n    for sample_id in range(angles.shape[1]):\n        for angle_idx in [2, 4, 6, 8]:\n            record_idx = record_indices[angle_idx]\n            parent = angles[angle_idx - 1, sample_id, :]\n            theta = np.floor(parent[0]/3)\n            phi = np.floor(parent[1]/3) # convert to bins\n            candidate_length = len(con_dis[record_idx][(theta, phi)])\n            # change some boundary points\n            if candidate_length == 0:\n                keys = list(con_dis[record_idx].keys())\n                while candidate_length == 0:\n                    temp_idx = np.random.choice(len(keys), 1)\n                    theta, phi = keys[temp_idx[0]]\n                    candidate_length = len(con_dis[record_idx][(theta, phi)])\n            chosen_idx = np.random.choice(candidate_length, 1)\n            angles[angle_idx, sample_id, :] = con_dis[record_idx][(theta, phi)][chosen_idx]\n    # convert to degrees with some noise\n    angles[[2,4,6,8], :, :] = angles[[2,4,6,8], :, :]*bin_size \n    angles[[2,4,6,8], :, :] += np.random.rand(4, angles.shape[1], 2)*bin_size\n    return angles\n\ndef sample_upper_limbs(angles, sample_num):\n    \"\"\"\n    Sample limb orientation for upper limbs.\n    \"\"\"\n    # sample torso limbs from the valid maps uniformly\n    # angles of shape [num_bones, sample_size, 2]\n    for angle_idx in [0, 1, 3, 5, 7]:\n        valid_map = angle_spread[0, record_indices[angle_idx]]\n        valid_map = valid_map.astype(np.float16)\n        valid_map /= valid_map.sum()\n        bone_angles = sample_from_histogram(valid_map, total=sample_num, \n                                add_noise=True)\n        angles[angle_idx, :, :] = bone_angles\n    return angles"
  },
  {
    "path": "libs/trainer/__init__.py",
    "content": "#!/usr/bin/env python3\n# -*- coding: utf-8 -*-\n\"\"\"\nEmpty file.\n\"\"\"\n\n\n"
  },
  {
    "path": "libs/trainer/trainer.py",
    "content": "\"\"\"\nUtility functions for cascaded model training and evaluation.\n\"\"\"\nimport libs.dataset.h36m.data_utils as data_utils\nimport libs.model.model as model\nfrom libs.utils.utils import compute_similarity_transform\n\nimport torch.nn.functional as F\nimport torch\nimport numpy as np\nimport logging\n\ndef train_cascade(train_dataset, eval_dataset, stats, action_eval_list, opt):\n    \"\"\"\n    Train a cascade of deep neural networks that lift 2D key-points to 3D pose.\n    \"\"\"\n    # initialize an empty cascade\n    cascade = model.get_cascade()\n    stage_record = []\n    input_size = len(stats['dim_use_2d'])\n    output_size = len(stats['dim_use_3d'])\n    # train each deep learner in the cascade sequentially\n    for stage_id in range(opt.num_stages):\n        # initialize a single deep learner\n        stage_model = model.get_model(stage_id + 1,\n                                      refine_3d=opt.refine_3d,\n                                      norm_twoD=opt.norm_twoD, \n                                      num_blocks=opt.num_blocks,\n                                      input_size=input_size,\n                                      output_size=output_size,\n                                      linear_size=opt.linear_size,\n                                      dropout=opt.dropout,\n                                      leaky=opt.leaky)\n        # record the stage number\n        train_dataset.set_stage(stage_id+1)\n        eval_dataset.set_stage(stage_id+1)\n        for dataset in action_eval_list:\n            dataset.set_stage(stage_id+1)\n        # move the deep learner to GPU\n        if opt.cuda:\n            stage_model = stage_model.cuda()\n            \n        # prepare the optimizer and learning rate scheduler\n        optim, sche = model.prepare_optim(stage_model, opt)\n        \n        # train the model\n        record = train(train_dataset, \n                       eval_dataset, \n                       stage_model,\n                       optim, \n                       sche, \n                       stats, \n                       action_eval_list, \n                       opt)\n        stage_model = record['model']\n        # record the training history\n        stage_record.append((record['batch_idx'], record['loss']))\n        # update current estimates and regression target\n        train_dataset.stage_update(stage_model, stats, opt)\n        eval_dataset.stage_update(stage_model, stats, opt)\n        \n        # update evaluation datasets for each action\n        if opt.evaluate_action:\n            for dataset in action_eval_list:\n                dataset.stage_update(stage_model, stats, opt)\n        # put the trained model into the cascade\n        cascade.append(stage_model.cpu())     \n        \n        # release memory\n        del stage_model    \n    return {'cascade':cascade, 'record':stage_record}\n\ndef evaluate_cascade(cascade, \n                     eval_dataset, \n                     stats, \n                     opt, \n                     save=False, \n                     save_path=None,\n                     action_wise=False, \n                     action_eval_list=None, \n                     apply_dropout=False\n                     ):\n    \"\"\"\n    Evaluate a cascaded model given a dataset object.\n    \"\"\"\n    loss, distance = None, None\n    for stage_id in range(len(cascade)):\n        print(\"#\"+ \"=\"*60 + \"#\")\n        logging.info(\"Model performance after stage {:d}\".format(stage_id + 1))\n        stage_model = cascade[stage_id]\n        if opt.cuda:\n            stage_model = stage_model.cuda()    \n        if action_wise:\n            evaluate_action_wise(action_eval_list, stage_model, stats, opt)    \n            # update the current estimates and regression targets\n            for dataset in action_eval_list:\n                dataset.stage_update(stage_model, stats, opt)\n        else:\n            # evaluate up to this stage\n            loss, distance = evaluate(eval_dataset, \n                                      stage_model, \n                                      stats, \n                                      opt, \n                                      save=save, \n                                      save_path=save_path,\n                                      procrustes=False, \n                                      per_joint=True, \n                                      apply_dropout=apply_dropout\n                                      )\n\n            # update datasets\n            eval_dataset.stage_update(stage_model, stats, opt)        \n        # release memory\n        del stage_model   \n    return loss, distance\n\ndef logger_print(epoch, \n                 batch_idx, \n                 batch_size, \n                 total_sample, \n                 total_batches,\n                 loss\n                 ):         \n    \"\"\"\n    Log training history.\n    \"\"\"      \n    msg = 'Train Epoch: {} [{}/{} ({:.0f}%)]\\tLoss: {:.6f} '.format(\n            epoch, \n            batch_idx * batch_size, \n            total_sample,\n            100. * batch_idx / total_batches, \n            loss.data.item())\n    logging.info(msg)\n    return\n\ndef train(train_dataset, \n          eval_dataset, \n          model, \n          optim, \n          sche, \n          stats, \n          action_eval_list, \n          opt, \n          plot_loss=False):\n    \"\"\"\n    Train a single deep learner.\n    \"\"\"\n    x_data = []\n    y_data = []\n    eval_loss, eval_distance = evaluate(eval_dataset, model, stats, opt)\n    if plot_loss:\n        import matplotlib.pyplot as plt\n        # plot loss curve during training\n        ax = plt.subplot(111)\n        lines = ax.plot(x_data, y_data)\n        plt.xlabel('batch')\n        plt.ylabel('training loss')\n    for epoch in range(1, opt.epochs + 1):\n        model.train()\n        # update the learning rate according to the scheduler\n        sche.step()        \n        # data loader\n        train_loader = torch.utils.data.DataLoader(train_dataset, \n                                                   batch_size=opt.batch_size, \n                                                   shuffle=True, \n                                                   num_workers=opt.num_threads\n                                                   )    \n        num_batches = len(train_loader)\n        for batch_idx, batch in enumerate(train_loader):\n            data = batch[0]\n            target = batch[1]\n            if opt.cuda:\n                with torch.no_grad():\n                    # move to GPU\n                    data, target = data.cuda(), target.cuda()                    \n            # erase all computed gradient        \n            optim.zero_grad()            \n            # forward pass to get prediction\n            prediction = model(data)\n            # compute loss\n            loss = F.mse_loss(prediction, target)            \n            # smoothed l1 loss function\n            #loss = F.smooth_l1_loss(prediction, target)            \n            # compute gradient in the computational graph\n            loss.backward()            \n            # update parameters in the model \n            optim.step()            \n            # logging\n            if batch_idx % opt.report_every == 0:\n                logger_print(epoch,\n                             batch_idx,\n                             opt.batch_size,\n                             len(train_dataset),\n                             len(train_loader),\n                             loss)                \n                x_data.append(num_batches*(epoch-1) + batch_idx)\n                y_data.append(loss.data.item())\n                if plot_loss:\n                    lines[0].set_xdata(x_data)\n                    lines[0].set_ydata(y_data)\n                    ax.relim()\n                    # update ax.viewLim using the new dataLim\n                    ax.autoscale_view()\n                    plt.draw()\n                    plt.pause(0.05)\n            # optinal evaluation\n            if opt.eval and batch_idx!= 0 and batch_idx % opt.eval_every == 0:\n                eval_loss, eval_distance = evaluate(eval_dataset, model, stats, opt)\n\n                # update learning rate if needed\n                #sche.step(eval_loss)\n                # reset to training mode\n                model.train()\n        # evaluate after each epoch \n        if opt.eval_action_wise and epoch % 50 == 0:\n            evaluate_action_wise(action_eval_list, model, stats, opt)\n    logging.info('Training finished.')\n    return {'model':model, 'batch_idx':x_data, 'loss':y_data}  \n\ndef evaluate_action_wise(dataset_list, model, stats, opt):\n    \"\"\"\n    Evaluate for a list of dataset objects, where each contains inputs for one action.\n    \"\"\"\n    record_P1 = {}\n    record_P2 = {}\n    average_P1 = 0\n    average_P2 = 0\n    protocols = opt.protocols\n    for dataset in dataset_list:\n        action = dataset.action_name\n        if 'P1' in protocols:            \n            eval_loss, eval_distance = evaluate(dataset, model, stats, opt, \n                                                verbose=False, procrustes = False)\n            record_P1[action] = (eval_loss, eval_distance)\n            average_P1 += eval_distance\n        if 'P2' in protocols:\n            eval_loss, eval_distance = evaluate(dataset, model, stats, opt, \n                                                verbose=False, procrustes = True)\n            record_P2[action] = (eval_loss, eval_distance)\n            average_P2 += eval_distance            \n    average_P1 /= len(dataset_list)\n    average_P2 /= len(dataset_list)\n    # logging\n    for protocol in protocols:\n        logging.info(\"MPJPE under protocol {:s}\".format(protocol))\n        record = record_P1 if protocol == 'P1' else record_P2\n        average = average_P1 if protocol == 'P1' else average_P2\n        for key in record.keys():\n            logging.info(\"Action: {:s}, error: {:.2f}\".format(key, record[key][1]))\n        logging.info(\"Average error over actions: {:.2f}\".format(average))\n    return [record_P1, record_P2]\n\ndef align_skeleton(skeletons_pred, skeletons_gt, num_of_joints):\n    \"\"\"\n    Apply per-frame procrustes alignment before computing MPJPE.\n    \"\"\"\n    for j in range(len(skeletons_gt)):\n        gt  = np.reshape(skeletons_gt[j,:], [-1,3])\n        out = np.reshape(skeletons_pred[j,:],[-1,3])\n        _, Z, T, b, c = compute_similarity_transform(gt, \n                                                     out, \n                                                     compute_optimal_scale=True\n                                                     )\n        out = (b * out.dot(T)) + c\n        skeletons_pred[j,:] = np.reshape(out,[-1, (num_of_joints - 1) * 3])     \n    return skeletons_pred\n\ndef evaluate(eval_dataset, \n             model, \n             stats, \n             opt, \n             save = False, \n             save_path=None,\n             verbose = True, \n             procrustes = False, \n             per_joint = False, \n             apply_dropout=False\n             ):\n    \"\"\"\n    Evaluate a 2D-to-3D lifting model on a given PyTorch dataset.\n    Adapted from ICCV 2017 baseline\n    https://github.com/una-dinosauria/3d-pose-baseline    \n    \"\"\"\n    num_of_joints = 14 if opt.pred14 else 17 \n    all_dists = []\n    model.eval()\n    if apply_dropout:\n        def apply_dropout(m):\n            if type(m) == torch.nn.Dropout:\n                m.train()        \n        # enable the dropout layers to produce a loss similar to the training \n        # loss (only for debugging purpose)\n        model.apply(apply_dropout)\n    eval_loader = torch.utils.data.DataLoader(eval_dataset, \n                                              batch_size = opt.batch_size, \n                                              shuffle = False, \n                                              num_workers = opt.num_threads\n                                              )\n    total_loss = 0\n    for batch_idx, batch in enumerate(eval_loader):\n        data = batch[0]\n        target = batch[1]\n        if opt.cuda:\n            with torch.no_grad():\n                data, target = data.cuda(), target.cuda()\n        # forward pass to get prediction\n        prediction = model(data)\n        # mean squared loss \n        loss = F.mse_loss(prediction, target, reduction='sum')\n        total_loss += loss.data.item()\n        # unnormalize the data\n        skeleton_3d_gt = data_utils.unNormalizeData(target.data.cpu().numpy(),\n                                                    stats['mean_3d'], \n                                                    stats['std_3d'], \n                                                    stats['dim_ignore_3d']\n                                                    )    \n        skeleton_3d_pred = data_utils.unNormalizeData(prediction.data.cpu().numpy(),\n                                                      stats['mean_3d'], \n                                                      stats['std_3d'], \n                                                      stats['dim_ignore_3d']\n                                                      )\n        # pick the joints that are used\n        dim_use = stats['dim_use_3d']\n        skeleton_3d_gt_use = skeleton_3d_gt[:, dim_use]\n        skeleton_3d_pred_use = skeleton_3d_pred[:, dim_use]\n        # error after a regid alignment, corresponding to protocol #2 in the paper\n        if procrustes:\n            skeleton_3d_pred_use = align_skeleton(skeleton_3d_pred_use, \n                                                  skeleton_3d_gt_use, \n                                                  num_of_joints\n                                                  )\n        # Compute Euclidean distance error per joint\n        sqerr = (skeleton_3d_gt_use - skeleton_3d_pred_use)**2 # Squared error between prediction and expected output\n        dists = np.zeros((sqerr.shape[0], num_of_joints)) # Array with L2 error per joint in mm\n        dist_idx = 0\n        for k in np.arange(0, num_of_joints*3, 3):\n          # Sum across X,Y, and Z dimenstions to obtain L2 distance\n          dists[:,dist_idx] = np.sqrt(np.sum(sqerr[:, k:k+3], axis=1))\n          dist_idx = dist_idx + 1\n        all_dists.append(dists)  \n    all_dists = np.vstack(all_dists)\n    if per_joint:\n        # show average error for each joint\n        error_per_joint = all_dists.mean(axis = 0)    \n        logging.info('Average error for each joint: ')\n        print(error_per_joint)\n    avg_loss = total_loss/(len(eval_dataset)*16*3)\n    if save:\n        record = {'error':all_dists}\n        np.save(save_path, np.array(record))\n    avg_distance = all_dists.mean()\n    if verbose:\n        logging.info('Evaluation set: average loss: {:.4f} '.format(avg_loss))\n        logging.info('Evaluation set: average joint distance: {:.4f} '.format(avg_distance))\n    return avg_loss, avg_distance\n"
  },
  {
    "path": "libs/utils/__init__.py",
    "content": "#!/usr/bin/env python3\n# -*- coding: utf-8 -*-\n\"\"\"\nEmpty file.\n\"\"\"\n\n\n"
  },
  {
    "path": "libs/utils/utils.py",
    "content": "\"\"\"\nUtility functions.\n\"\"\"\nimport libs.dataset.h36m.data_utils as data_utils\nimport libs.dataset.h36m.cameras as cameras\nimport libs.dataset.h36m.pth_dataset as dataset\nimport libs.visualization.viz as viz\n\nimport logging\nimport time\nimport numpy as np\nimport torch\nimport matplotlib.pyplot as plt\nimport os\ninput_size = 32\noutput_size = 48\n\n# global configurations\ncamera_frame = True # camera coordinate\npredict_14 = False # predict 14 joints\n\ndef save_ckpt(opt, record, stats):\n    \"\"\"\n    Save training results.\n    \"\"\"\n    cascade = record['cascade']\n    if not opt.save:\n        return False\n    if opt.save_name is None:\n        save_name = time.asctime()\n        save_name += ('stages_' + str(opt.num_stages) + 'blocks_' \n                      + str(opt.num_blocks) + opt.extra_str\n                      )        \n    save_dir = os.path.join(opt.save_root, save_name)\n    if not os.path.exists(save_dir):\n        os.makedirs(save_dir)\n    torch.save(cascade, os.path.join(save_dir, 'model.th'))\n    np.save(os.path.join(save_dir, 'stats.npy'), stats)\n    print('Model saved at ' + save_dir)\n    return True\n\ndef load_ckpt(opt):\n    cascade = torch.load(os.path.join(opt.ckpt_dir, 'model.th'))\n    stats = np.load(os.path.join(opt.ckpt_dir, 'stats.npy'), allow_pickle=True).item()\n    if opt.cuda:\n        cascade.cuda()\n    return cascade, stats\n\ndef list_remove(list_a, list_b):\n    list_c = []\n    for item in list_a:\n        if item not in list_b:\n            list_c.append(item)\n    return list_c\n\n\ndef get_all_data(data_x, data_y, camera_frame):\n    \"\"\"\n    Obtain a list of all the batches, randomly permutted\n    Args\n      data_x: dictionary with 2d inputs\n      data_y: dictionary with 3d expected outputs\n      camera_frame: whether the 3d data is in camera coordinates\n      training: True if this is a training batch. False otherwise.\n    Returns\n      encoder_inputs: list of 2d batches\n      decoder_outputs: list of 3d batches\n    \"\"\"\n\n    # Figure out how many frames we have\n    n = 0\n    for key2d in data_x.keys():\n      n2d, _ = data_x[ key2d ].shape\n      n = n + n2d\n\n    encoder_inputs  = np.zeros((n, input_size), dtype=float)\n    decoder_outputs = np.zeros((n, output_size), dtype=float)\n\n    # Put all the data into big arrays\n    idx = 0\n    for key2d in data_x.keys():\n      (subj, b, fname) = key2d\n      # keys should be the same if 3d is in camera coordinates\n      key3d = key2d if (camera_frame) else (subj, b, '{0}.h5'.format(fname.split('.')[0]))\n      key3d = (subj, b, fname[:-3]) if fname.endswith('-sh') and camera_frame else key3d\n\n      n2d, _ = data_x[ key2d ].shape\n      encoder_inputs[idx:idx+n2d, :]  = data_x[ key2d ]\n      decoder_outputs[idx:idx+n2d, :] = data_y[ key3d ]\n      idx = idx + n2d\n\n    return encoder_inputs, decoder_outputs\n\ndef adjust_figure(left = 0, right = 1, bottom = 0.01, top = 0.95,\n                  wspace = 0, hspace = 0.4):  \n    plt.subplots_adjust(left, bottom, right, top, wspace, hspace)\n    return\n\ndef visualize(eval_dataset, model, stats, opt, save=False, save_dir=None):\n    # visualze model prediction batch by batch\n    batch_size = 9\n    # how many batches to save\n    if save:\n        num_batches = 10\n        current_batch = 1\n    model.eval()\n    eval_loader = torch.utils.data.DataLoader(eval_dataset, batch_size, \n                                              shuffle = True, num_workers = opt.num_threads)     \n    for batch_idx, batch in enumerate(eval_loader):\n        if save and current_batch > num_batches:\n            break\n        data = batch[0]\n        target = batch[1]\n        if opt.cuda:\n            with torch.no_grad():\n                # move to GPU\n                data, target = data.cuda(), target.cuda()\n        # forward pass to get prediction\n        prediction = model(data)\n        # un-normalize the data\n        skeleton_2d = data_utils.unNormalizeData(data.data.cpu().numpy(), \n        stats['mean_2d'], stats['std_2d'], stats['dim_ignore_2d'])\n        skeleton_3d_gt = data_utils.unNormalizeData(target.data.cpu().numpy(), \n        stats['mean_3d'], stats['std_3d'], stats['dim_ignore_3d'])    \n        skeleton_3d_pred = data_utils.unNormalizeData(prediction.data.cpu().numpy(), \n        stats['mean_3d'], stats['std_3d'], stats['dim_ignore_3d'])\n        # visualizing\n        if save:\n            plt.ioff()\n        f = plt.figure(figsize=(16, 8))\n        axes = []\n        for sample_idx in range(batch_size):\n            ax = plt.subplot(3, 9, 3*sample_idx + 1)\n            viz.show2Dpose(skeleton_2d[sample_idx], ax)\n            plt.gca().invert_yaxis()\n            ax = plt.subplot(3, 9, 3*sample_idx + 2, projection='3d')\n            viz.show3Dpose(skeleton_3d_gt[sample_idx], ax)\n            ax = plt.subplot(3, 9, 3*sample_idx + 3, projection='3d')\n            viz.show3Dpose(skeleton_3d_pred[sample_idx], ax, pred=True)  \n            viz.show3Dpose(skeleton_3d_gt[sample_idx], ax, gt=True)   \n            axes.append(ax)      \n        adjust_figure(left = 0.05, right = 0.95, bottom = 0.05, top = 0.95,\n                      wspace = 0.3, hspace = 0.3)              \n        if not save:\n            plt.pause(0.5)\n            # rotate the axes and update\n            for angle in range(0, 360, 5):\n                for ax in axes:\n                    ax.view_init(30, angle)\n                plt.draw()\n                plt.pause(.001)\n            input('Press enter to view next batch.')\n        else:\n            # save plot\n            f.savefig(save_dir +'/'+ str(current_batch) + '.png')\n        plt.close(f)     \n        del axes\n        if save:\n            current_batch += 1\n    return\n\ndef temp_visualize(eval_dataset, model, stats, opt):\n    # visualze model prediction batch by batch\n    model.eval()\n    data = np.load('./pics/pts1.npy').astype(np.float32)\n    data = data[:,2:]\n    # normalize the data\n    mean_vec = stats['mean_2d'][stats['dim_use_2d']]\n    std_vec = stats['std_2d'][stats['dim_use_2d']]\n    data = (data-mean_vec)/std_vec\n    data = torch.from_numpy(data.astype(np.float32))\n    data = data.cuda()\n    # forward pass to get prediction\n    prediction = model(data)\n    # un-normalize the data\n    skeleton_2d = data_utils.unNormalizeData(data.data.cpu().numpy(), \n    stats['mean_2d'], stats['std_2d'], stats['dim_ignore_2d'])   \n    skeleton_3d_pred = data_utils.unNormalizeData(prediction.data.cpu().numpy(), \n    stats['mean_3d'], stats['std_3d'], stats['dim_ignore_3d'])\n    # visualizing\n    plt.figure()\n    ax = plt.subplot(1, 2, 1)\n    viz.show2Dpose(skeleton_2d[0], ax)\n    plt.gca().invert_yaxis()\n    ax = plt.subplot(1, 2, 2, projection='3d')\n    viz.show3Dpose(skeleton_3d_pred[0], ax, pred=True)              \n    plt.show()\n            # rotate the axes and update\n#            for angle in range(0, 360, 5):\n#                for ax in axes:\n#                    ax.view_init(30, angle)\n#                plt.draw()\n#                plt.pause(.001)\n#            input('Press enter to view next batch.')\n    return\n\ndef visualize_cascade(eval_dataset, cascade, stats, opt, save=False, save_dir=None):\n    num_stages = len(cascade)\n    # visualze model prediction batch by batch\n    batch_size = 5\n    # how many batches to save\n    if save:\n        num_batches = 10\n        current_batch = 1\n    for stage_model in cascade:\n        stage_model.eval()\n    eval_loader = torch.utils.data.DataLoader(eval_dataset, batch_size, \n                                              shuffle = False, \n                                              num_workers = opt.num_threads)     \n    for batch_idx, batch in enumerate(eval_loader):\n        if save and current_batch > num_batches:\n            break\n        data = batch[0]\n        ## debug\n        # enc_in = np.array([[648., 266], [679, 311], [688, 320], [693, 161],\n        #            [620, 244], [526, 156], [642, 160], [590, 310],\n        #            [505, 350], [380, 375], [491, 285],\n        #            [543, 190], [572, 119], [515, 417], [518, 514],\n        #            [512, 638]],dtype=np.float32)\n        enc_in = data\n        enc_in = enc_in.reshape(1, 32)\n        # normalize\n        data_mean_2d = stats['mean_2d']\n        dim_to_use_2d = stats['dim_use_2d']\n        data_std_2d = stats['std_2d']\n        enc_in = (enc_in - data_mean_2d[dim_to_use_2d])/data_std_2d[dim_to_use_2d]\n        data = torch.from_numpy(enc_in.astype(np.float32))\n        ## End experiment 2019/10/16\n        target = batch[1]\n        # store predictions for each stage\n        prediction_stages = []\n        if opt.cuda:\n            with torch.no_grad():\n                # move to GPU\n                data, target = data.cuda(), target.cuda()\n        # forward pass to get prediction for the first stage\n        prediction = cascade[0](data)\n        prediction_stages.append(prediction)\n        # prediction for later stages\n        for stage_idx in range(1, num_stages):\n            prediction = cascade[stage_idx](data)\n            prediction_stages.append(prediction_stages[stage_idx-1] + prediction)\n        # un-normalize the data\n        skeleton_2d = data_utils.unNormalizeData(data.data.cpu().numpy(), \n        stats['mean_2d'], stats['std_2d'], stats['dim_ignore_2d'])\n        skeleton_3d_gt = data_utils.unNormalizeData(target.data.cpu().numpy(), \n        stats['mean_3d'], stats['std_3d'], stats['dim_ignore_3d'])\n        for stage_idx in range(num_stages):\n            prediction_stages[stage_idx] = data_utils.unNormalizeData(prediction_stages[stage_idx].data.cpu().numpy(), \n            stats['mean_3d'], stats['std_3d'], stats['dim_ignore_3d'])\n        ## save intermediate results\n        # import scipy.io as sio\n        # p3d = prediction_stages[0]\n        # sio.savemat('./teaser_pose3d.mat', {'pred_3d':p3d.reshape(32,3),\n        #             'pred_2d':np.array([[648., 266], [679, 311], [688, 320], [693, 161],\n        #            [620, 244], [526, 156], [642, 160], [590, 310],\n        #            [505, 350], [447, 348], [380, 375], [491, 285],\n        #            [543, 190], [572, 119], [515, 417], [518, 514],\n        #            [512, 638]])})        \n        ## End Experiment 2019/10/16\n        # visualizing\n        if save:\n            plt.ioff()\n        f = plt.figure(figsize=(16, 8))\n        axes = []\n        for sample_idx in range(batch_size):\n            for stage_idx in range(num_stages):\n                ax = plt.subplot(batch_size, num_stages+1, 1+(num_stages+1)*sample_idx)\n                viz.show2Dpose(skeleton_2d[sample_idx], ax)\n                plt.gca().invert_yaxis()\n                ax = plt.subplot(batch_size, num_stages+1, \n                                 2+stage_idx+(num_stages+1)*sample_idx, projection='3d')\n                viz.show3Dpose(prediction_stages[stage_idx][sample_idx], ax, pred=True)  \n                viz.show3Dpose(skeleton_3d_gt[sample_idx], ax, gt=True)   \n                axes.append(ax)      \n        adjust_figure(left = 0.05, right = 0.95, bottom = 0.05, top = 0.95,\n                      wspace = 0.3, hspace = 0.3)              \n        if not save:\n            plt.pause(0.5)\n            # rotate the axes and update\n#            for angle in range(0, 360, 5):\n#                for ax in axes:\n#                    ax.view_init(30, angle)\n#                plt.draw()\n#                plt.pause(.001)\n            input('Press enter to view next batch.')\n        else:\n            # save plot\n            f.savefig(save_dir +'/'+ str(current_batch) + '.png')\n        plt.close(f)     \n        del axes\n        if save:\n            current_batch += 1    \n    return\n\ndef compute_similarity_transform(X, Y, compute_optimal_scale=False):\n    \"\"\"\n    A port of MATLAB's `procrustes` function to Numpy.\n    Adapted from http://stackoverflow.com/a/18927641/1884420\n    Args\n      X: array NxM of targets, with N number of points and M point dimensionality\n      Y: array NxM of inputs\n      compute_optimal_scale: whether we compute optimal scale or force it to be 1\n    Returns:\n      d: squared error after transformation\n      Z: transformed Y\n      T: computed rotation\n      b: scaling\n      c: translation\n    \"\"\"\n    muX = X.mean(0)\n    muY = Y.mean(0)\n\n    X0 = X - muX\n    Y0 = Y - muY\n\n    ssX = (X0**2.).sum()\n    ssY = (Y0**2.).sum()\n\n    # centred Frobenius norm\n    normX = np.sqrt(ssX)\n    normY = np.sqrt(ssY)\n\n    # scale to equal (unit) norm\n    X0 = X0 / normX\n    Y0 = Y0 / normY\n\n    # optimum rotation matrix of Y\n    A = np.dot(X0.T, Y0)\n    U,s,Vt = np.linalg.svd(A,full_matrices=False)\n    V = Vt.T\n    T = np.dot(V, U.T)\n\n    # Make sure we have a rotation\n    detT = np.linalg.det(T)\n    V[:,-1] *= np.sign( detT )\n    s[-1]   *= np.sign( detT )\n    T = np.dot(V, U.T)\n\n    traceTA = s.sum()\n\n    if compute_optimal_scale:  # Compute optimum scaling of Y.\n        b = traceTA * normX / normY\n        d = 1 - traceTA**2\n        Z = normX*traceTA*np.dot(Y0, T) + muX\n    else:  # If no scaling allowed\n        b = 1\n        d = 1 + ssY/ssX - 2 * traceTA * normY / normX\n        Z = normY*np.dot(Y0, T) + muX\n\n    c = muX - b*np.dot(muY, T)\n    return d, Z, T, b, c\n"
  },
  {
    "path": "libs/visualization/__init__.py",
    "content": "#!/usr/bin/env python3\n# -*- coding: utf-8 -*-\n\"\"\"\nEmpty file.\n\"\"\"\n\n\n"
  },
  {
    "path": "libs/visualization/viz.py",
    "content": "\"\"\"\nFunctions to visualize human poses\n\"\"\"\n\nimport matplotlib.pyplot as plt\nimport libs.dataset.h36m.data_utils as data_utils\nimport numpy as np\n# import h5py\nimport os\nfrom mpl_toolkits.mplot3d import Axes3D\n\ndef show3Dpose(channels, ax, lcolor=\"#3498db\", rcolor=\"#e74c3c\", add_labels=True,\n               gt=False,pred=False): # blue, orange\n  \"\"\"\n  Visualize a 3d skeleton\n\n  Args\n    channels: 96x1 vector. The pose to plot.\n    ax: matplotlib 3d axis to draw on\n    lcolor: color for left part of the body\n    rcolor: color for right part of the body\n    add_labels: whether to add coordinate labels\n  Returns\n    Nothing. Draws on ax.\n  \"\"\"\n\n  assert channels.size == len(data_utils.H36M_NAMES)*3, \"channels should have 96 entries, it has %d instead\" % channels.size\n  vals = np.reshape( channels, (len(data_utils.H36M_NAMES), -1) )\n\n  I   = np.array([1,2,3,1,7,8,1, 13,14,15,14,18,19,14,26,27])-1 # start points\n  J   = np.array([2,3,4,7,8,9,13,14,15,16,18,19,20,26,27,28])-1 # end points\n  LR  = np.array([1,1,1,0,0,0,0, 0, 0, 0, 0, 0, 0, 1, 1, 1], dtype=bool)\n\n  # Make connection matrix\n  for i in np.arange( len(I) ):\n    x, y, z = [np.array( [vals[I[i], j], vals[J[i], j]] ) for j in range(3)]\n    if gt:\n        ax.plot(x,z, -y,  lw=2, c='k')\n#        ax.plot(x,y, z,  lw=2, c='k')\n        \n    elif pred:\n        ax.plot(x,z, -y,  lw=2, c='r')\n#        ax.plot(x,y, z,  lw=2, c='r')\n\n    else:\n#        ax.plot(x,z, -y,  lw=2, c=lcolor if LR[i] else rcolor)\n        ax.plot(x,y, z,  lw=2, c=lcolor if LR[i] else rcolor)\n\n  RADIUS = 750 # space around the subject\n  xroot, yroot, zroot = vals[0,0], vals[0,1], vals[0,2]\n  ax.set_xlim3d([-RADIUS+xroot, RADIUS+xroot])\n  ax.set_zlim3d([-RADIUS+zroot, RADIUS+zroot])\n  ax.set_ylim3d([-RADIUS+yroot, RADIUS+yroot])\n\n  if add_labels:\n    ax.set_xlabel(\"x\")\n    ax.set_ylabel(\"y\")\n    ax.set_zlabel(\"z\")\n\n  # Get rid of the ticks and tick labels\n#  ax.set_xticks([])\n#  ax.set_yticks([])\n#  ax.set_zticks([])\n#\n#  ax.get_xaxis().set_ticklabels([])\n#  ax.get_yaxis().set_ticklabels([])\n#  ax.set_zticklabels([])\n  ax.set_aspect('equal')\n\n  # Get rid of the panes (actually, make them white)\n  white = (1.0, 1.0, 1.0, 0.0)\n  ax.w_xaxis.set_pane_color(white)\n  ax.w_yaxis.set_pane_color(white)\n  # Keep z pane\n\n  # Get rid of the lines in 3d\n  ax.w_xaxis.line.set_color(white)\n  ax.w_yaxis.line.set_color(white)\n  ax.w_zaxis.line.set_color(white)\n\ndef show2Dpose(channels, ax, lcolor=\"#3498db\", rcolor=\"#e74c3c\", add_labels=False):\n  \"\"\"\n  Visualize a 2d skeleton\n\n  Args\n    channels: 64x1 vector. The pose to plot.\n    ax: matplotlib axis to draw on\n    lcolor: color for left part of the body\n    rcolor: color for right part of the body\n    add_labels: whether to add coordinate labels\n  Returns\n    Nothing. Draws on ax.\n  \"\"\"\n\n  assert channels.size == len(data_utils.H36M_NAMES)*2, \"channels should have 64 entries, it has %d instead\" % channels.size\n  vals = np.reshape( channels, (-1, 2) )\n  #plt.plot(vals[:,0], vals[:,1], 'ro')\n  I  = np.array([1,2,3,1,7,8,1, 13,14,14,18,19,14,26,27])-1 # start points\n  J  = np.array([2,3,4,7,8,9,13,14,16,18,19,20,26,27,28])-1 # end points\n  LR = np.array([1,1,1,0,0,0,0, 0, 0, 0, 0, 0, 1, 1, 1], dtype=bool)\n\n  # Make connection matrix\n  for i in np.arange( len(I) ):\n    x, y = [np.array( [vals[I[i], j], vals[J[i], j]] ) for j in range(2)]\n    ax.plot(x, y, lw=2, c=lcolor if LR[i] else rcolor)\n\n  # Get rid of the ticks\n#  ax.set_xticks([])\n#  ax.set_yticks([])\n#\n#  # Get rid of tick labels\n#  ax.get_xaxis().set_ticklabels([])\n#  ax.get_yaxis().set_ticklabels([])\n\n  RADIUS = 350 # space around the subject\n  xroot, yroot = vals[0,0], vals[0,1]\n  ax.set_xlim([-RADIUS+xroot, RADIUS+xroot])\n  ax.set_ylim([-RADIUS+yroot, RADIUS+yroot])\n  if add_labels:\n    ax.set_xlabel(\"x\")\n    ax.set_ylabel(\"z\")\n\n  ax.set_aspect('equal')\n"
  },
  {
    "path": "resources/.gitignore",
    "content": "# Ignore everything in this directory\n*\n# Except this file\n!.gitignore\n"
  },
  {
    "path": "spec-list.txt",
    "content": "# This file may be used to create an environment using:\n# $ conda create --name <env> --file <this file>\n# platform: linux-64\n@EXPLICIT\nhttps://repo.anaconda.com/pkgs/main/linux-64/conda-env-2.6.0-1.tar.bz2\nhttps://repo.anaconda.com/pkgs/free/linux-64/libgfortran-3.0.0-1.tar.bz2\nhttps://conda.anaconda.org/conda-forge/linux-64/ca-certificates-2019.3.9-hecc5488_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/cudatoolkit-9.0-h13b8566_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/free/linux-64/freeglut-2.8.1-0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/gmp-6.1.2-h6c8ec71_1.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/icu-58.2-h9c2bf20_1.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/intel-openmp-2019.3-199.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/libffi-3.2.1-hd88cf55_4.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/libgcc-ng-8.2.0-hdf63c60_1.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/libgfortran-ng-7.3.0-hdf63c60_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/libstdcxx-ng-8.2.0-hdf63c60_1.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/lzo-2.10-h49e0be7_2.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/snappy-1.1.7-hbae5bb6_3.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/yaml-0.1.7-had09818_2.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/bzip2-1.0.6-h14c3975_5.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/cudnn-7.3.1-cuda9.0_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/expat-2.2.6-he6710b0_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/fribidi-1.0.5-h7b6447c_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/gettext-0.19.8.1-hd7bead4_3.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/giflib-5.1.4-h14c3975_1.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/graphite2-1.3.13-h23475e2_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/jbig-2.1-hdba287a_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/jpeg-9b-h024ee3a_2.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/libgcc-7.2.0-h69d50b8_2.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/libglu-9.0.0-hf484d3e_1.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/libiconv-1.15-h63c8f33_5.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/liblief-0.9.0-h7725739_2.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/libopenblas-0.3.3-h5a2b251_3.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/libopus-1.3-h7b6447c_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/libsodium-1.0.16-h1bed415_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/libtool-2.4.6-h7b6447c_5.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/libuuid-1.0.3-h1bed415_2.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/libvpx-1.7.0-h439df22_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/libxcb-1.13-h1bed415_1.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/lz4-c-1.8.1.2-h14c3975_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/mkl-2018.0.3-1.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/mpfr-3.1.5-h11a74b3_2.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/nccl-1.3.5-cuda9.0_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/ncurses-6.1-he6710b0_1.tar.bz2\nhttps://conda.anaconda.org/conda-forge/linux-64/nettle-3.3-0.tar.bz2\nhttps://conda.anaconda.org/conda-forge/linux-64/openblas-0.3.5-h9ac9557_1001.tar.bz2\nhttps://conda.anaconda.org/conda-forge/linux-64/openssl-1.1.1b-h14c3975_1.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/patchelf-0.9-he6710b0_3.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/pcre-8.43-he6710b0_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/pixman-0.38.0-h7b6447c_0.tar.bz2\nhttps://conda.anaconda.org/conda-forge/linux-64/x264-1!152.20180806-h14c3975_0.tar.bz2\nhttps://conda.anaconda.org/conda-forge/linux-64/xorg-kbproto-1.0.7-h14c3975_1002.tar.bz2\nhttps://conda.anaconda.org/conda-forge/linux-64/xorg-libice-1.0.9-h516909a_1004.tar.bz2\nhttps://conda.anaconda.org/conda-forge/linux-64/xorg-renderproto-0.11.1-h14c3975_1002.tar.bz2\nhttps://conda.anaconda.org/conda-forge/linux-64/xorg-xextproto-7.3.0-h14c3975_1002.tar.bz2\nhttps://conda.anaconda.org/conda-forge/linux-64/xorg-xproto-7.0.31-h14c3975_1007.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/xz-5.2.4-h14c3975_4.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/zlib-1.2.11-h7b6447c_3.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/blosc-1.15.0-hd408876_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/glib-2.56.2-hd408876_0.tar.bz2\nhttps://conda.anaconda.org/conda-forge/linux-64/gnutls-3.5.19-h2a4e5f8_1.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/hdf5-1.10.2-hba1933b_1.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/jasper-2.0.14-h07fcdf6_0.tar.bz2\nhttps://conda.anaconda.org/conda-forge/linux-64/libblas-3.8.0-4_openblas.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/libedit-3.1.20181209-hc058e9b_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/libpng-1.6.36-hbc83047_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/libprotobuf-3.5.2-h6f1eeef_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/libssh2-1.8.0-h1ba5d50_4.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/libxml2-2.9.9-he19cac6_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/mpc-1.0.3-hec55b23_5.tar.bz2\nhttps://conda.anaconda.org/conda-forge/linux-64/openh264-1.7.0-0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/pandoc-2.2.3.2-0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/readline-7.0-h7b6447c_5.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/tk-8.6.8-hbc83047_0.tar.bz2\nhttps://conda.anaconda.org/conda-forge/linux-64/xorg-libsm-1.2.2-h470a237_5.tar.bz2\nhttps://conda.anaconda.org/conda-forge/linux-64/xorg-libx11-1.6.7-h14c3975_1000.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/zeromq-4.2.5-hf484d3e_1.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/zstd-1.3.7-h0b5b093_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/dbus-1.13.6-h746ee38_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/freetype-2.9.1-h8a8886c_1.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/gstreamer-1.14.0-hb453b48_1.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/krb5-1.16.1-h173b8e3_7.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/libarchive-3.3.3-h5d8350f_5.tar.bz2\nhttps://conda.anaconda.org/conda-forge/linux-64/libcblas-3.8.0-4_openblas.tar.bz2\nhttps://conda.anaconda.org/conda-forge/linux-64/liblapack-3.8.0-4_openblas.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/libtiff-4.0.10-h2733197_2.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/libxslt-1.1.33-h7d1a2b0_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/sqlite-3.27.2-h7b6447c_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/unixodbc-2.3.7-h14c3975_0.tar.bz2\nhttps://conda.anaconda.org/conda-forge/linux-64/xorg-libxext-1.3.4-h516909a_0.tar.bz2\nhttps://conda.anaconda.org/conda-forge/linux-64/xorg-libxrender-0.9.10-h516909a_1002.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/ffmpeg-4.0-hcdf2ecd_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/fontconfig-2.13.0-h9420a91_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/gst-plugins-base-1.14.0-hbbd80ab_1.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/libcurl-7.64.0-h20c2e04_2.tar.bz2\nhttps://conda.anaconda.org/conda-forge/linux-64/liblapacke-3.8.0-4_openblas.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/libwebp-1.0.0-h222930b_1.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/python-3.6.8-h0371630_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/alabaster-0.7.12-py36_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/appdirs-1.4.3-py36h28b3542_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/asn1crypto-0.24.0-py36_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/noarch/atomicwrites-1.3.0-py_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/noarch/attrs-19.1.0-py_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/backcall-0.1.0-py36_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/backports-1.0-py36_1.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/bitarray-0.8.3-py36h14c3975_0.tar.bz2\nhttps://conda.anaconda.org/conda-forge/linux-64/blas-2.4-openblas.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/boto-2.49.0-py36_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/cairo-1.14.12-h8948797_3.tar.bz2\nhttps://conda.anaconda.org/conda-forge/linux-64/certifi-2019.3.9-py36_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/chardet-3.0.4-py36_1.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/click-7.0-py36_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/noarch/cloudpickle-0.7.0-py_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/colorama-0.4.1-py36_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/constantly-15.1.0-py36h28b3542_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/contextlib2-0.5.5-py36_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/noarch/cryptography-vectors-2.6.1-py_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/curl-7.64.0-hbc83047_2.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/dask-core-1.1.4-py36_1.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/noarch/decorator-4.4.0-py_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/defusedxml-0.5.0-py36_1.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/docutils-0.14-py36_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/entrypoints-0.3-py36_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/et_xmlfile-1.0.1-py36_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/fastcache-1.0.2-py36h14c3975_2.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/filelock-3.0.10-py36_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/future-0.17.1-py36_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/glob2-0.6-py36_1.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/gmpy2-2.0.8-py36hc8893dd_2.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/greenlet-0.4.15-py36h7b6447c_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/heapdict-1.0.0-py36_2.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/idna-2.8-py36_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/imagesize-1.1.0-py36_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/incremental-17.5.0-py36_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/ipython_genutils-0.2.0-py36_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/itsdangerous-1.1.0-py36_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/jdcal-1.4-py36_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/jeepney-0.4-py36_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/kiwisolver-1.0.1-py36hf484d3e_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/lazy-object-proxy-1.3.1-py36h14c3975_2.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/llvmlite-0.28.0-py36hd408876_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/locket-0.2.0-py36_1.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/lxml-4.3.2-py36hefd8a0e_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/markupsafe-1.1.1-py36h7b6447c_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/mccabe-0.6.1-py36_1.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/mistune-0.8.4-py36h7b6447c_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/mkl-service-1.1.2-py36h651fb7a_4.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/more-itertools-6.0.0-py36_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/mpmath-1.1.0-py36_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/msgpack-python-0.6.1-py36hfd86e86_1.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/ninja-1.9.0-py36hfd86e86_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/olefile-0.46-py36_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/pandocfilters-1.4.2-py36_1.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/parso-0.3.4-py36_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/pep8-1.7.1-py36_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/pickleshare-0.7.5-py36_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/pkginfo-1.5.0.1-py36_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/pluggy-0.9.0-py36_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/ply-3.11-py36_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/prometheus_client-0.6.0-py36_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/psutil-5.6.1-py36h7b6447c_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/ptyprocess-0.6.0-py36_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/py-1.8.0-py36_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/py-lief-0.9.0-py36h7725739_2.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/noarch/pyasn1-0.4.5-py_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/pycodestyle-2.5.0-py36_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/pycosat-0.6.3-py36h14c3975_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/pycparser-2.19-py36_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/pycrypto-2.6.1-py36h14c3975_9.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/pycurl-7.43.0.2-py36h1ba5d50_0.tar.bz2\nhttps://conda.anaconda.org/conda-forge/noarch/pydicom-1.2.2-py_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/pyflakes-2.1.1-py36_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/pyodbc-4.0.26-py36he6710b0_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/pyparsing-2.3.1-py36_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/pysocks-1.6.8-py36_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/python-libarchive-c-2.8-py36_6.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/pytz-2018.9-py36_0.tar.bz2\nhttps://conda.anaconda.org/conda-forge/linux-64/pyyaml-5.1-py36h14c3975_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/pyzmq-17.1.2-py36h14c3975_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/qt-5.9.7-h5867ecd_1.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/noarch/qtpy-1.7.0-py_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/rope-0.12.0-py36_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/ruamel_yaml-0.15.46-py36h14c3975_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/send2trash-1.5.0-py36_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/simplegeneric-0.8.1-py36_2.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/sip-4.19.8-py36hf484d3e_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/six-1.12.0-py36_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/snowballstemmer-1.2.1-py36_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/sortedcontainers-2.1.0-py36_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/soupsieve-1.8-py36_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/sphinxcontrib-1.0-py36_1.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/sqlalchemy-1.3.1-py36h7b6447c_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/tblib-1.3.2-py36_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/testpath-0.4.2-py36_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/toolz-0.9.0-py36_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/tornado-6.0.2-py36h7b6447c_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/noarch/tqdm-4.31.1-py_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/typed-ast-1.3.1-py36h7b6447c_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/typing-3.6.4-py36_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/unicodecsv-0.14.1-py36_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/wcwidth-0.1.7-py36_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/webencodings-0.5.1-py36_1.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/werkzeug-0.14.1-py36_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/wrapt-1.11.1-py36h7b6447c_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/wurlitzer-1.0.2-py36_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/xlrd-1.2.0-py36_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/xlsxwriter-1.1.5-py36_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/xlwt-1.3.0-py36_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/zipp-0.3.3-py36_1.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/zope-1.0-py36_1.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/automat-0.7.0-py36_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/babel-2.6.0-py36_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/backports.os-0.1.1-py36_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/backports.shutil_get_terminal_size-1.0.0-py36_2.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/beautifulsoup4-4.7.1-py36_1.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/cffi-1.12.2-py36h2e261b9_1.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/cycler-0.10.0-py36_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/cytoolz-0.9.0.1-py36h14c3975_1.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/harfbuzz-1.8.8-hffaf4a1_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/html5lib-1.0.1-py36_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/hyperlink-18.0.0-py36_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/importlib_metadata-0.8-py36_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/jedi-0.13.3-py36_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/multipledispatch-0.6.0-py36_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/nltk-3.4-py36_1.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/nomkl-3.0-0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/numpy-base-1.16.2-py36h2f8d375_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/noarch/openpyxl-2.6.1-py_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/packaging-19.0-py36_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/noarch/partd-0.3.10-py_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/pathlib2-2.3.3-py36_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/pexpect-4.6.0-py36_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/pillow-5.4.1-py36h34e0f95_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/protobuf-3.5.2-py36hf484d3e_1.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/pyasn1-modules-0.2.4-py36_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/pyhamcrest-1.9.0-py36_2.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/pyqt-5.9.2-py36h05f1152_2.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/pyrsistent-0.14.11-py36h7b6447c_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/python-dateutil-2.8.0-py36_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/noarch/qtawesome-0.5.7-py_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/setuptools-40.8.0-py36_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/singledispatch-3.4.0.3-py36_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/snakeviz-1.0.0-py36_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/sortedcollections-1.1.2-py36_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/sphinxcontrib-websupport-1.1.0-py36_1.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/sympy-1.3-py36_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/terminado-0.8.1-py36_1.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/traitlets-4.3.2-py36_0.tar.bz2\nhttps://conda.anaconda.org/conda-forge/noarch/yacs-0.1.6-py_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/zict-0.1.4-py36_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/zope.interface-4.6.0-py36h7b6447c_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/astroid-2.2.5-py36_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/bleach-3.1.0-py36_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/clyent-1.2.2-py36_1.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/cryptography-2.6.1-py36h1ba5d50_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/cython-0.29.6-py36he6710b0_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/distributed-1.26.0-py36_1.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/get_terminal_size-1.0.0-haa9412d_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/gevent-1.4.0-py36h7b6447c_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/isort-4.3.16-py36_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/jinja2-2.10-py36_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/jsonschema-3.0.1-py36_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/jupyter_core-4.4.0-py36_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/libopencv-3.4.2-hb342d67_1.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/networkx-2.2-py36_1.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/nose-1.3.7-py36_2.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/numpy-1.16.2-py36h99e49ec_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/openblas-devel-0.3.3-3.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/pango-1.42.4-h049681c_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/path.py-11.5.0-py36_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/pygments-2.3.1-py36_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/pytest-4.3.1-py36_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/wheel-0.33.1-py36_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/conda-verify-3.1.1-py36_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/flask-1.0.2-py36_1.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/graphviz-2.40.1-h21bd128_2.tar.bz2\nhttps://conda.anaconda.org/anaconda/linux-64/h5py-2.8.0-py36h989c5e5_3.tar.bz2\nhttps://conda.anaconda.org/conda-forge/linux-64/imageio-2.5.0-py36_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/jupyter_client-5.2.4-py36_0.tar.bz2\nhttps://conda.anaconda.org/anaconda/linux-64/matplotlib-3.0.3-py36h5429711_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/nbformat-4.4.0-py36_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/pip-19.0.3-py36_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/prompt_toolkit-2.0.9-py36_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/pylint-2.3.1-py36_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/pyopenssl-19.0.0-py36_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/pytest-openfiles-0.3.2-py36_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/pytest-remotedata-0.3.1-py36_0.tar.bz2\nhttps://conda.anaconda.org/pytorch/linux-64/pytorch-1.0.1-py3.6_cuda9.0.176_cudnn7.4.2_2.tar.bz2\nhttps://conda.anaconda.org/anaconda/linux-64/pywavelets-1.0.2-py36hdd07704_0.tar.bz2\nhttps://conda.anaconda.org/anaconda/linux-64/scipy-1.2.1-py36he2b7bc3_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/secretstorage-3.1.1-py36_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/service_identity-18.1.0-py36h28b3542_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/flask-cors-3.0.7-py36_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/ipython-7.4.0-py36h39e3cac_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/keyring-18.0.0-py36_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/noarch/nbconvert-5.4.1-py_2.tar.bz2\nhttps://conda.anaconda.org/anaconda/linux-64/scikit-image-0.14.2-py36he6710b0_0.tar.bz2\nhttps://conda.anaconda.org/anaconda/linux-64/scikit-learn-0.20.3-py36h22eb022_0.tar.bz2\nhttps://conda.anaconda.org/pytorch/noarch/torchvision-0.2.2-py_3.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/twisted-18.9.0-py36h7b6447c_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/urllib3-1.24.1-py36_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/ipykernel-5.1.0-py36h39e3cac_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/requests-2.21.0-py36_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/anaconda-client-1.7.2-py36_0.tar.bz2\nhttps://conda.anaconda.org/conda-forge/linux-64/conda-4.6.14-py36_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/jupyter_console-6.0.0-py36_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/notebook-5.7.8-py36_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/qtconsole-4.4.3-py36_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/sphinx-1.8.5-py36_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/spyder-kernels-0.4.2-py36_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/anaconda-navigator-1.9.7-py36_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/anaconda-project-0.8.2-py36_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/conda-build-3.17.8-py36_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/numpydoc-0.8.0-py36_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/widgetsnbextension-3.4.2-py36_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/ipywidgets-7.4.2-py36_0.tar.bz2\nhttps://repo.anaconda.com/pkgs/main/linux-64/spyder-3.3.3-py36_0.tar.bz2\n"
  },
  {
    "path": "tools/2Dto3Dnet.py",
    "content": "\"\"\"\n3D pose estimation based on 2D key-point coordinates as inputs.\nAuthor: Shichao Li\nEmail: nicholas.li@connect.ust.hk\n\"\"\"\n\nimport logging\nimport os\nimport sys\nsys.path.append(\"../\")\n\nimport torch\nimport numpy as np\n\nimport libs.parser.parse as parse\nimport libs.utils.utils as utils\nimport libs.dataset.h36m.data_utils as data_utils\nimport libs.trainer.trainer as trainer\n\ndef main():\n    # logging configuration\n    logging.basicConfig(level=logging.INFO,\n                        format=\"[%(asctime)s]: %(message)s\"\n                        )        \n    \n    # parse command line input\n    opt = parse.parse_arg()\n    \n    # Set GPU\n    opt.cuda = opt.gpuid >= 0\n    if opt.cuda:\n        torch.cuda.set_device(opt.gpuid)\n    else:\n        logging.info(\"GPU is disabled.\")\n\n    # dataset preparation\n    train_dataset, eval_dataset, stats, action_eval_list = \\\n    data_utils.prepare_dataset(opt)\n    \n    if opt.train:\n        # train a cascaded 2D-to-3D pose estimation model\n        record = trainer.train_cascade(train_dataset, \n                                       eval_dataset, \n                                       stats, \n                                       action_eval_list, \n                                       opt\n                                       )\n        utils.save_ckpt(opt, record, stats)\n        \n    if opt.visualize:\n        # visualize the inference results of a pre-trained model\n        cascade = torch.load(opt.ckpt_path)        \n        if opt.cuda:\n            cascade.cuda()  \n        utils.visualize_cascade(eval_dataset, cascade, stats, opt)\n    if opt.evaluate:\n        # evalaute a pre-trained cascade\n        cascade, stats = utils.load_ckpt(opt)\n        trainer.evaluate_cascade(cascade, \n                                 eval_dataset, \n                                 stats, \n                                 opt, \n                                 action_wise=opt.eval_action_wise,\n                                 action_eval_list=action_eval_list\n                                 ) \n        \nif __name__ == \"__main__\":\n    main()"
  },
  {
    "path": "tools/annotate_2D.py",
    "content": "\"\"\"\nAnnotate 2D key-points interactively.\nPress Q to exit the tool.\nPress C to remove the annotation.\nPress N to go to the next image.\nPress Z to save the annotation.\nMouse click to annotate 2D key-points.\n\"\"\"\n\nimport imageio\nimport matplotlib.pyplot as plt \nimport numpy as np \nimport argparse\nimport os\nimport time\n\nfrom glob import glob\n\n''' ANNOTATION CONTROLS '''\nnames = ['ra',      # 0 right ankle\n         'rk',      # 1 right knee2\n         'rh',      # 2 right hip\n         'lh',      # 3 left hip\n         'lk',      # 4 left knee\n         'la',      # 5 left ankle\n         'rw',      # 6 right wrist\n         're',      # 7 right elbow\n         'rs',      # 8 right shoulder\n         'ls',      # 9 left shoulder\n         'le',      # 10 left elbow\n         'lw',      # 11 left wrist\n         'ne',      # 12 neck\n         'ht',      # 13 head top\n         'sp',      # 14 spine\n         'th',      # 15 thorax\n         'ns']      # 16 nose\n\nMAX_CLICK_LENGTH = 0.3 # in seconds; anything longer is a pan/zoom motion\n\njoints = np.array([]).reshape(0, 2)\nfig = None\nax = None\ncid = None\nplots = None\nimg_path_list = None\nimg_idx = None\nannotation = None\nannotation_path = None\n# whether the image has veen updated or not\nupdated = False\nn = len(names)\ntime_onclick = None\n\ndef initialization(opt):\n    global img_path_list, img_idx, annotation,annotation_path, updated\n    img_path_list = sorted(glob(os.path.join(opt.dataset_dir, '*.jpg')))\n    img_path_list += sorted(glob(os.path.join(opt.dataset_dir, '*.png')))\n    assert len(img_path_list) != 0, \"Can not find image files.\"\n    img_idx = -1\n    annotation_path = os.path.join(opt.dataset_dir, 'annotation.npy')\n    if os.path.exists(annotation_path):\n        annotation = np.load(annotation_path).item()\n    else:\n        annotation = {}    \n    return\n\ndef plot_image(img_path):\n    global plots, updated, joints\n    ax.clear()\n    img = imageio.imread(img_path)\n    joints = np.array([]).reshape(0, 2)\n    ax.imshow(img)\n    plots = ax.plot([], [], 'ro') \n    fig.canvas.draw()\n    updated = True  \n    return\n\ndef onclick(event):\n    global time_onclick\n    time_onclick = time.time()\n    \ndef onrelease(event):\n    global joints, fig, plots, updated, time_onclick\n    \n    if event.button == 1 and ((time.time() - time_onclick) < MAX_CLICK_LENGTH):\n        if len(joints) < n and updated:\n            ind = len(joints)\n            joint = np.array([event.xdata, event.ydata])\n            joints = np.vstack((joints, joint))\n            plots[0].remove()\n            plots = plt.plot(joints[:, 0], joints[:, 1], 'ro')\n            fig.canvas.draw()\n            print(names[ind] + \": \" + str(joint))\n            if len(joints) == n:\n                # record the annotation\n                img_name = img_path_list[img_idx].split(os.sep)[-1]\n                annotation[img_name] = {'p2d':joints}\n                joints = np.array([]).reshape(0, 2)\n                updated = False\n                print('Please go on to the next image.')\n\ndef save_results():\n    np.save(annotation_path, annotation)\n    print('A Python dictionary has been saved at ' + annotation_path)\n    \ndef onkey(event):\n    global joints, fig, cid, plots, img_idx\n\n    if event.key == 'c': \n        # remove the annotation on the image\n        if len(joints) > 0:\n            joints = np.array([]).reshape(0, 2)\n            plots[0].remove()\n            plots = plt.plot([], [], 'ro')\n            fig.canvas.draw()\n        \n    if event.key == 'n':\n        # go to next image\n        if img_idx <= len(img_path_list) - 1:\n            # look for the next unannotated image\n            img_idx += 1\n            while img_path_list[img_idx].split(os.sep)[-1] in annotation:\n                img_idx += 1     \n            img_idx = len(img_path_list) - 1 if img_idx == len(img_path_list) else img_idx\n            plot_image(img_path_list[img_idx])\n        else:\n            print('Already the last image.')\n            save_results()\n    \n    if event.key == 'z':\n        # save the annotation\n        save_results()\n        \ndef main(opt):\n    global joints, fig, ax, cid, plots, img_idx, img_path_list, annotation\n    # show one unannotated image\n    for idx in range(len(img_path_list)):\n        img_name = img_path_list[idx].split(os.sep)[-1]\n        if img_name not in annotation:\n            # start with this unannotated image\n            img_idx = idx\n            break\n    if img_idx == -1:\n        print('No unannotated image found.')\n        return\n    fig = plt.gcf() \n    ax = plt.gca() \n    plot_image(img_path_list[img_idx])\n    fig.canvas.mpl_connect('button_press_event', onclick)\n    fig.canvas.mpl_connect('button_release_event', onrelease)\n    cid = fig.canvas.mpl_connect('key_press_event', onkey)\n    plt.show() \n    return\n\nif __name__ == '__main__':\n    parser = argparse.ArgumentParser(description='2D Annotation')\n    parser.add_argument('-d', '--dataset_dir', default=None, type=str)\n    opt = parser.parse_args()\n    initialization(opt)\n    main(opt)"
  },
  {
    "path": "tools/annotate_3D.py",
    "content": "\"\"\"\nInteractive annotation tool for 3D human pose estimation.\nGiven an image and a coarse 3D skeleton estimation, the user can interactively\nmodify the 3D parameters and save them as the ground truth.\n\"\"\"\n\nimport matplotlib.pyplot as plt\nimport numpy as np\nimport argparse\nimport imageio\nimport sys\nimport os\n\nfrom mpl_toolkits.mplot3d import Axes3D\nfrom scipy.spatial.transform import Rotation as R\nfrom cv2 import projectPoints\n\nsys.path.append(\"../\")\n\nfrom libs.skeleton.anglelimits import get_basis1, normalize, gram_schmidt_columns\nfrom libs.skeleton.anglelimits import nt_parent_indices, nt_child_indices, di_indices\nfrom libs.skeleton.anglelimits import get_normal, di, a, to_spherical, to_xyz, bone_name\n\n''' GLOBAL VARIABLES '''\nangle_idx = 0 # Bone angle to adjust\ndirection = 0 # Direction to rotate, (0 - x, 1 - y, 2 - z) for upper arm only\nstep = 3 # 3 degrees for step size\nstep_radian = step * np.pi / 180\nlocal_system_map = {1:0, 3:0, 5:1, 7:1, 2:2, 4:3, 6:4, 8:5}\nline_index_map = {1:11, 3:14, 5:4, 7:1, 2:12, 4:15, 6:5, 8:2}\nparent_indices = np.array([1,2,3,1,7,8,1, 13,14,15,14,18,19,14,26,27])-1\nchild_indices = np.array([2,3,4,7,8,9,13,14,15,16,18,19,20,26,27,28])-1\ndirection_name = ['x', 'y', 'z']\n\n# translation vector of the camera\nt = None\n# focal length of the camera\nf = None\n# intrinsic matrix for camera projection\nintrinsic_mat = None\n\n# Objects for ploting\nfig = None\nplot_ax = None\nimg_ax = None\nskeleton = None\nlines = None \npoints = None\nRADIUS = 1 # Space around the subject\n\n# hierarchical representation\nlocal_systems = None\nneed_to_update_lc = False\nbones_global = None\nbones_local = None\nangles = None\n\n# file path\nannotation_path = None\nannotation = None\nimg_name = None\n\n# some joint correspondence\nindex_list = [13, 14, 129, 145]\nH36M_IDS = [0, 2, 5, 8, 1, 4, 7, 3, 12, 15, 24, 16, 18, 20, 17, 19, 21]\nUSE_DIMS = [0, 1, 2, 3, 6, 7, 8, 12, 13, 14, 15, 17, 18, 19, 25, 26, 27]\n\n# keyboard inputs\nbone_idx_keys = ['1', '2', '3', '4', '5', '6', '7', '8', '9']\nglobal_rot_key ='0'\ninc_step_key = 'd'\ndec_step_key = 'f'\nang_inc_key = 'up'\nang_dec_key = 'down'\nang_cw_key = 'right'\nang_ccw_key = 'left'\nsave_key = 'm'\n\ndef press(event):\n    \"\"\"\n    Call-back function when user press any key.\n    \"\"\"\n    global angle_idx, direction, need_to_update_lc\n    global bones_global, bones_local, skeleton, angles, local_systems\n\n    if event.key == 'p':\n        plot_ax.plot([np.random.rand()], [np.random.rand()], [np.random.rand()], 'ro')\n        fig.canvas.draw()\n\n    if event.key in bone_idx_keys:  angle_idx = int(event.key) - 1 \n    if event.key == global_rot_key: angle_idx = None\n    if event.key == inc_step_key:   direction = (direction + 1) % 3\n    if event.key == dec_step_key:   direction = (direction - 1) % 3\n\n    if event.key == ang_inc_key or event.key == ang_dec_key:\n        update_skeleton(angle_idx, event.key)\n\n    if event.key == ang_cw_key or event.key == ang_ccw_key:\n        if angle_idx in [2, 4, 6, 8]:\n            update_skeleton(angle_idx, event.key)\n\n    if event.key == save_key:\n        save_skeleton()\n\n    if angle_idx is not None:\n        notes = 'current limb: ' + bone_name[angle_idx + 1]\n        # update local coordinate systems if needed\n        if need_to_update_lc:\n            # compute the local coordinate system\n            bones_global, bones_local, local_systems = to_local(skeleton)\n            # convert the local coordinates into spherical coordinates\n            angles = to_spherical(bones_local)\n            angles[:,1:] *= 180/np.pi\n            # need to update local coordinate system once after global rotation\n            need_to_update_lc = False            \n    else:\n        notes = 'global rotation: '\n\n    if angle_idx in [None, 1, 3, 5, 7]:\n        notes += ' direction: ' + direction_name[direction]\n    if event.key not in ['up', 'down', 'right', 'left']:\n        print(notes)\n    plot_ax.set_xlabel(notes)\n    fig.canvas.draw_idle()\n        \ndef show3Dpose(channels, \n               ax, \n               lcolor=\"#3498db\", \n               rcolor=\"#e74c3c\", \n               add_labels=True,\n               gt=False,\n               pred=False,\n               inv_z=False\n               ): \n\n    vals = np.reshape( channels, (32, -1) )\n\n    I   = parent_indices\n    J   = child_indices\n    LR  = np.array([1,1,1,0,0,0,0, 0, 0, 0, 0, 0, 0, 1, 1, 1], dtype=bool)\n    \n    lines = []\n\n    # Make connection matrix\n    for i in np.arange( len(I) ):\n        x, y, z = [np.array( [vals[I[i], j], vals[J[i], j]] ) for j in range(3)]\n        line = ax.plot(x,y, z,  lw=2, c=lcolor if LR[i] else rcolor)\n        lines.append(line)\n\n    xroot, yroot, zroot = vals[0,0], vals[0,1], vals[0,2]\n\n    ax.set_xlim3d([-RADIUS+xroot, RADIUS+xroot])\n    ax.set_zlim3d([-RADIUS+zroot, RADIUS+zroot])\n    ax.set_ylim3d([-RADIUS+yroot, RADIUS+yroot])\n\n    if add_labels:\n        ax.set_xlabel(\"x\")\n        ax.set_ylabel(\"y\")\n        ax.set_zlabel(\"z\")\n        \n    ax.set_aspect('auto')\n\n    # Get rid of the panes (actually, make them white)\n    white = (1.0, 1.0, 1.0, 0.0)\n    ax.w_xaxis.set_pane_color(white)\n    ax.w_yaxis.set_pane_color(white)\n\n    # Get rid of the lines in 3d\n    ax.w_xaxis.line.set_color(white)\n    ax.w_yaxis.line.set_color(white)\n    ax.w_zaxis.line.set_color(white)\n\n    if inv_z:\n        ax.invert_zaxis() \n\n    return lines\n\ndef to_local(skeleton):\n    \"\"\"\n    Convert bone vector in skeleton format to local coordinate system\n    \"\"\"\n    global local_systems\n    \n    v1, v2, v3 = get_basis1(skeleton)\n\n    # Compute vector of left hip to right hip\n    left_hip = skeleton[6]\n    right_hip = skeleton[1]\n    v4 = normalize(right_hip - left_hip)\n\n    # v5 is the cross product of v4 and v2 (front-facing vector for lower-body)\n    v5 = normalize(np.cross(v4, v2))\n\n    # Compute orthogonal coordinate systems using GramSchmidt\n    # Make sure the directions roughly align\n    # For upper body, we use v1, v2 and v3\n    # For lower body, we use v4, v2 and v5\n    system1 = gram_schmidt_columns(np.hstack([v1.reshape(3,1), \n                                              v2.reshape(3,1), \n                                              v3.reshape(3,1)]))\n\n    system2 = gram_schmidt_columns(np.hstack([v4.reshape(3,1), \n                                              v2.reshape(3,1), \n                                              v5.reshape(3,1)]))\n\n    local_systems = [system1, system2]\n    bones = skeleton[nt_parent_indices, :] - skeleton[nt_child_indices, :]\n\n    # convert bone vector to local coordinate system\n    bones_local = np.zeros(bones.shape, dtype=bones.dtype)\n\n    for bone_idx in range(len(bones)):\n        # only compute bone vectors for non-torsos\n        # the order of the non-torso bone vector is: \n        # bone vector1: thorax to head top\n        # bone vector2: left shoulder to left elbow\n        # bone vector3: left elbow to left wrist\n        # bone vector4: right shoulder to right elbow\n        # bone vector5: right elbow to right wrist\n        # bone vector6: left hip to left knee\n        # bone vector7: left knee to left ankle\n        # bone vector8: right hip to right knee\n        # bone vector9: right knee to right ankle\n\n        bone = bones[bone_idx]\n        if bone_idx in [0, 1, 3, 5, 7]:\n            # Bones directly connected to torso\n            # Upper body - 0, 1, 3\n            # Lower body - 5, 7\n            if bone_idx in [0, 1, 3]:   bones_local[bone_idx] = system1.T @ bone\n            else:                       bones_local[bone_idx] = system2.T @ bone\n        else:\n            if bone_idx in [2, 4]:  parent_R = system1\n            else:                   parent_R = system2\n\n            # parent bone index is smaller than 1\n            vector_u = normalize(bones[bone_idx - 1])\n            di_index = di_indices[bone_idx]\n\n            vector_v, flag = get_normal(parent_R@di[:, di_index],\n                                  parent_R@a,\n                                  vector_u)\n            vector_w = np.cross(vector_u, vector_v)\n\n            local_system = gram_schmidt_columns(np.hstack([vector_u.reshape(3,1), \n                                              vector_v.reshape(3,1), \n                                              vector_w.reshape(3,1)]))\n\n            local_systems.append(local_system)\n            bones_local[bone_idx] = local_system.T @ bone\n\n    return bones, bones_local, local_systems\n\ndef update_line(line_idx, parent_idx, child_idx):\n    \"\"\"\n    Update 3D line segments.\n    \"\"\"\n    global lines\n\n    # update 3D lines\n    parent, child = skeleton[parent_idx], skeleton[child_idx]\n\n    x = np.array([parent[0], child[0]])\n    y = np.array([parent[1], child[1]])\n    z = np.array([parent[2], child[2]])\n\n    lines[line_idx][0].set_data(x, y)\n    lines[line_idx][0].set_3d_properties(z)\n\n    fig.canvas.draw_idle()\n\ndef update_global(angle_idx):\n    \"\"\"\n    Update bone vectors.\n    \"\"\"\n    global bones_global, bones_local, local_systems, skeleton\n    bones_global[angle_idx] = local_systems[local_system_map[angle_idx]] @ bones_local[angle_idx]\n    skeleton[nt_child_indices[angle_idx]] = skeleton[nt_parent_indices[angle_idx]] \\\n                                            - bones_global[angle_idx]  \n\n    line_idx = line_index_map[angle_idx]\n    parent_idx = nt_parent_indices[angle_idx]\n    child_idx = nt_child_indices[angle_idx]\n\n    update_line(line_idx, parent_idx, child_idx)\n\ndef rotate_global(rot):\n    \"\"\"\n    Change the global orientation of the 3D skeleton.\n    \"\"\"\n    global skeleton\n    hip = skeleton[0].reshape(1,3)\n    temp_skeleton = skeleton - hip\n    skeleton = (rot.as_matrix() @ temp_skeleton.T).T + hip\n\n    for line_idx in range(len(parent_indices)):\n        update_line(line_idx, \n                    parent_indices[line_idx], \n                    child_indices[line_idx]\n                    )\n\ndef update_skeleton(angle_idx, key_name):\n    \"\"\"\n    Update the 3D skeleton with a specified keyboard input.\n    \"\"\"\n    global need_to_update_lc, local_systems\n\n    # Rotate the lower-limb\n    if angle_idx in [2, 4, 6, 8]:\n        if key_name == 'up':        angles[angle_idx, 1] += step\n        elif key_name == 'down':    angles[angle_idx, 1] -= step\n        elif key_name == 'left':    angles[angle_idx, 2] += step\n        elif key_name == 'right':   angles[angle_idx, 2] -= step \n        \n        temp = angles[angle_idx].copy()\n        temp[1:] *= np.pi / 180\n        bones_local[angle_idx] = to_xyz(temp.reshape(1,3))\n        \n        update_global(angle_idx)\n\n    # Rotate the upper-limb with respect to the torso coordinate system\n    if angle_idx in [1, 3, 5, 7]:\n\n        # Local rotation vector\n        rot_vec = np.array([0., 0., 0.])\n        rot_vec[direction] = 1. if key_name == 'up' else -1.\n        rot = R.from_rotvec(rot_vec*step_radian)\n        bones_local[angle_idx] = rot.apply(bones_local[angle_idx]) \n\n        # Global rotation vector\n        rot_vec2 = local_systems[local_system_map[angle_idx]][:, direction].copy()\n        rot_vec2 *= 1. if key_name == 'up' else -1.\n        rot2 = R.from_rotvec(rot_vec2*step_radian)\n\n        # Local rotation vector for child/lower limb\n        temp = local_systems[local_system_map[angle_idx + 1]]\n        local_systems[local_system_map[angle_idx + 1]] = rot2.as_matrix() @ temp\n\n        # update parent and child bone\n        update_global(angle_idx)\n        update_global(angle_idx + 1)\n\n    # Global rotation\n    if angle_idx is None and key_name in ['up', 'down']:\n        need_to_update_lc = True\n        rot_vec = np.array([0., 0., 0.])\n        rot_vec[direction] = 1. if key_name == 'up' else -1.\n        rot = R.from_rotvec(rot_vec*step_radian)\n        rotate_global(rot)\n\n    # Update the 2D Projection\n    update_projection()\n\ndef update_projection():\n    \"\"\"\n    Update the 2D projection of the 3D key-points.\n    \"\"\"\n    global points\n    points.pop(0).remove()\n    proj2d = projectPoints(skeleton, \n                           np.zeros((3)), \n                           t, \n                           intrinsic_mat, \n                           np.zeros((5))\n                           )\n    proj2d = proj2d[0].reshape(-1,2)\n    points = img_ax.plot(proj2d[:,0], proj2d[:,1], 'ro')   \n    fig.canvas.draw_idle()\n\ndef save_skeleton():\n    \"\"\"\n    Save the annotation file.\n    \"\"\"\n    global annotation\n    annotation[img_name]['p3d'] = skeleton\n    np.save(annotation_path, annotation)\n    print('Annotated 3D parameters saved at ' + annotation_path)\n\ndef visualize(pose, skeleton, img):\n    \"\"\"\n    Initialize the 3D and 2D plots.\n    \"\"\"\n    global lines, points, fig, plot_ax, img_ax, intrinsic_mat\n    fig = plt.figure() \n    fig.canvas.mpl_disconnect(fig.canvas.manager.key_press_handler_id)\n    \n    # 3D pose plot\n    plot_ax = plt.subplot(121, projection='3d')\n\n    lines = show3Dpose(pose, plot_ax)\n    fig.canvas.mpl_connect('key_press_event', press)\n    plot_ax.set_title('1-9: limb selection, 0: global rotation, arrow keys: rotation')\n    # Image plot\n    img_ax = plt.subplot(122)\n    img_ax.imshow(img)\n    intrinsic_mat = np.array([[f[0], 0.00e+00, float(img.shape[1])/2],\n                              [0.00e+00, f[1], float(img.shape[0])/2],\n                              [0.00e+00, 0.00e+00, 1.00e+00]])\n    proj2d = projectPoints(skeleton, \n                           np.zeros((3)), \n                           t, \n                           intrinsic_mat, \n                           np.zeros((5))\n                           )\n    proj2d = proj2d[0].reshape(-1,2)\n    points = img_ax.plot(proj2d[:,0], proj2d[:,1], 'ro')\n    # Show the plot\n    plt.show()\n\ndef create_python3_file(opt):\n    \"\"\"\n    The fitted parameters are stored using python 2. \n    Create a Python 3 file from it when this script is executed for the first time.\n    \"\"\"\n    fitted_path = os.path.join(opt.dataset_dir, \"fitted.npy\")\n    annotation_path = os.path.join(opt.dataset_dir, \"annot_3d.npy\")\n    if not os.path.exists(annotation_path):\n    # The fitting parameters are obtianed in Python 2 environment,\n    # thus the encoding argument is used here        \n        fitted = np.load(fitted_path, encoding='latin1', allow_pickle=True).item()   \n        np.save(annotation_path, fitted)\n    return\n\ndef main(opt):\n    global t, f, angles, bones_global, bones_local, need_to_update_lc, \\\n    local_system, skeleton, annotation_path, img_name, annotation\n    create_python3_file(opt)\n    annotation_path = os.path.join(opt.dataset_dir, \"annot_3d.npy\")\n    annotation = np.load(annotation_path, allow_pickle=True).item()    \n    for img_name in annotation.keys():\n        # e.g., img_name = '260.jpg'\n        img_name = '260.jpg'\n        # select one unannotated image and start the interactive annotation\n        if 'p3d' in annotation[img_name]:\n            continue\n        fitting_params = annotation[img_name]['fitting_params']\n        img_path = os.path.join(opt.dataset_dir, img_name)\n        \n        img = imageio.imread(img_path)\n\n        # Convert smpl format to Human 3.6M Format\n        skeleton_smpl = fitting_params['v'].reshape(-1, 3)\n        skeleton = np.zeros((32, 3))\n        skeleton[USE_DIMS] = skeleton_smpl[H36M_IDS]\n\n        pose = skeleton.reshape(-1)\n\n        t = fitting_params['cam_t']\n        f = fitting_params['f']\n\n        # Convert skeleton to local coordinate system\n        bones_global, bones_local, local_system = to_local(skeleton)\n\n        # Convert the local coordinates to spherical coordinates\n        angles = to_spherical(bones_local)\n        angles[:, 1:] *= 180 / np.pi\n\n        # Set update local coordinate flag only after global rotation\n        need_to_update_lc = False\n\n        # Visualize\n        visualize(pose, skeleton, img)\n        \n        # annotate only one image at a time\n        break\n\nif __name__ == '__main__':\n    parser = argparse.ArgumentParser(description='3D Interactive Tool')\n    parser.add_argument('-d', '--dataset_dir', default=None, type=str)\n    opt = parser.parse_args()\n    main(opt)"
  },
  {
    "path": "tools/evolve.py",
    "content": "\"\"\"\nEvolution of 3D human skeleton.\nauthor: Nicholas Li\ncontact: nicholas.li@connect.ust.hk\n\"\"\"\nimport sys\nsys.path.append(\"../\")\n\nfrom libs.evolution.genetic import evolution\nfrom libs.evolution.parameter import parse_arg\n\nimport os\nimport logging\nfrom scipy.spatial.transform import Rotation as R\nimport numpy as np\n\ndef cast_to_float(dic, dtype=np.float32):\n    # cast to float 32 for space saving\n    for key in dic.keys():\n        dic[key] = dic[key].astype(dtype)\n    return dic\n\ndef random_rotation(pose, sigma=360):\n    # apply random rotation to equivalently augment viewpoints\n    pose = pose.reshape(32, 3)\n    hip = pose[0].copy().reshape(1, 3)\n    x = np.random.normal(scale=sigma)\n    y = np.random.normal(scale=sigma)\n    z = np.random.normal(scale=sigma)\n    r = R.from_euler('xyz', [x, y, z], degrees=True)\n    rotated = r.as_dcm() @ (pose-hip).T\n    return (rotated.T + hip).reshape(-1)\n\ndef initialize_population(data_dic, opt):\n    \"\"\"\n    Initialize a population for later evolution.\n    \"\"\"\n    # down-sample the raw data if used for weakly-supervised experiments\n    if opt.WS and opt.SS.startswith(\"0.\") and opt.SS.endswith(\"S1\"):\n        # a fraction of S1 data for H36M\n        ratio = float(opt.SS.split('S')[0])\n        # randomly sample a portion of 3D data\n        sampled_dic = {}\n        # sample each video\n        for key in data_dic.keys():\n            if key[0] != 1:\n                continue\n            total = len(data_dic[key])\n            sampled_num = int(ratio*total)\n            chosen_indices = np.random.choice(total, sampled_num, replace=False)\n            sampled_dic[key] = data_dic[key][chosen_indices].copy()\n        initial_population = np.concatenate(list(sampled_dic.values()), axis=0)   \n    elif opt.WS and opt.SS.startswith(\"S\"):\n        # a collection of data from a few subjects\n        # delete unused subjects\n        sub_list = [int(opt.SS[i]) for i in range(1, len(opt.SS))]\n        keys_to_delete = []\n        for key in data_dic.keys():\n            if key[0] not in sub_list:\n                keys_to_delete.append(key)\n        for key in keys_to_delete:    \n            del data_dic[key]        \n        initial_population = np.concatenate(list(data_dic.values()), axis=0)               \n    else:\n        # do not perform down-sampling\n        initial_population = np.concatenate(list(data_dic.values()), axis=0)    \n    return initial_population\n\ndef initialize_model_file(opt):\n    if opt.A:\n        import torch\n        model = torch.load(os.path.join(opt.ckpt_dir, \"model.th\"))\n        stats = np.load(os.path.join(opt.ckpt_dir, \"stats.npy\")).item()\n        cameras = np.load(\"../data/human3.6M/cameras.npy\").item()\n        model_file = {\"model\":model, \"stats\":stats, \"cams\":list(cameras.values())}\n    else:\n        model_file = None  \n    return model_file\n\n\ndef split_and_save(evolved_population):\n    \"\"\"\n    Split and save the evolved dataset into training and validation set.\n    \"\"\"\n    training_indices = np.random.choice(len(evolved_population), int(0.95*len(evolved_population)), replace=False)\n    testing_indices = np.delete(np.arange(len(evolved_population)), training_indices)\n    training_poses = evolved_population[training_indices]\n    testing_poses = evolved_population[testing_indices]\n\n    temp_subject_list = [1, 5, 6, 7, 8]\n    train_set_3d = {}\n    poses_list = np.array_split(training_poses, len(temp_subject_list))\n    for subject_idx in range(len(temp_subject_list)):\n        train_set_3d[(temp_subject_list[subject_idx], 'n/a', 'n/a')] =\\\n        poses_list[subject_idx] \n    # testing\n    testing_poses = evolved_population[testing_indices]\n    temp_subject_list = [9,11]\n    test_set_3d = {}\n    poses_list = np.array_split(testing_poses, len(temp_subject_list))\n    for subject_idx in range(len(temp_subject_list)):\n        test_set_3d[(temp_subject_list[subject_idx], 'n/a', 'n/a')] =\\\n        poses_list[subject_idx] \n    np.save('../data/human3.6M/h36m/numpy/threeDPose_train_split.npy', train_set_3d)\n    np.save('../data/human3.6M/h36m/numpy/threeDPose_test.npy', test_set_3d)    \n    return\n\ndef visualize(initial_population, evolved_population):\n    \"\"\"\n    Visualize the augmented dataset\n    \"\"\"\n    import matplotlib.pyplot as plt\n    from genetic import show3Dpose\n    def get_zmin(pose):\n        return pose.reshape(32,3)[:,2].min()\n    # initial population\n    chosen_indices = np.random.choice(len(initial_population), 9, replace=False)\n    plt.figure()\n    for idx in range(9):\n        ax = plt.subplot(3, 3, idx+1, projection='3d')\n        pose = initial_population[chosen_indices[idx]]\n        show3Dpose(pose, ax) \n        plt.title(\"{:d}:{:.2f}\".format(chosen_indices[idx], get_zmin(pose)))\n    plt.tight_layout() \n    # after evolution\n    chosen_indices = np.random.choice(len(evolved_population) - len(initial_population), 9, replace=False)\n    plt.figure()\n    for idx in range(9):\n        ax = plt.subplot(3, 3, idx+1, projection='3d')\n        pose = evolved_population[chosen_indices[idx] + len(initial_population)]\n        show3Dpose(pose, ax) \n        plt.title(\"{:d}:{:.2f}\".format(chosen_indices[idx] + len(initial_population), get_zmin(pose)))\n    plt.tight_layout()     \n    return\n\ndef main():\n    logging.basicConfig(level=logging.INFO,\n                        format=\"[%(asctime)s]: %(message)s\"\n                        ) \n    # parse command line input\n    opt = parse_arg()    \n    if opt.generate:\n        # get the training set of human 3.6M \n        data_dic = np.load(opt.data_path, allow_pickle=True).item()\n        initial_population = initialize_population(data_dic, opt)\n        # load a pre-trained model for active searching (optional)\n        model_file = initialize_model_file(opt)      \n        evolved_population = evolution(initial_population,\n                                       opt,\n                                       model_file=model_file\n                                       )        \n        if opt.split:\n            split_and_save(evolved_population)\n    \n    if opt.visualize:\n        visualize(initial_population, evolved_population)\n    \nif __name__ == \"__main__\":\n    main()"
  },
  {
    "path": "tools/imgto2Dnet.py",
    "content": "# ------------------------------------------------------------------------------\n# Copyright (c) Microsoft\n# Licensed under the MIT License.\n# Written by Bin Xiao (Bin.Xiao@microsoft.com)\n# Modified by Shichao Li (nicholas.li@connect.ust.hk)\n# ------------------------------------------------------------------------------\n\n\"\"\"\nTraining and inference of a high-resolution heatmap regression model.\n\"\"\"\n\nfrom __future__ import absolute_import\nfrom __future__ import division\nfrom __future__ import print_function\n\nimport argparse\nimport os\nimport pprint\nimport shutil\n\nimport torch\nimport torch.nn.parallel\nimport torch.backends.cudnn as cudnn\nimport torch.optim\nimport torch.utils.data\nimport torch.utils.data.distributed\nimport torchvision.transforms as transforms\n\nimport sys\nsys.path.append(\"../\")\n\nfrom libs.hhr.config import cfg\nfrom libs.hhr.config import update_config\nfrom libs.hhr.core.loss import JointsMSELoss, JointsCoordinateLoss, WingLoss\nfrom libs.hhr.core.function import train\nfrom libs.hhr.core.function import validate_pixel\nfrom libs.hhr.utils.utils import get_optimizer\nfrom libs.hhr.utils.utils import save_checkpoint\nfrom libs.hhr.utils.utils import create_logger\nfrom libs.hhr.utils.utils import get_model_summary\nfrom libs.model.pose_hrnet import get_pose_net\n\nimport libs.dataset.h36m\nfrom libs.dataset.h36m.h36m_pose import H36MDataset\n# run with your configuration file as follows:\n# --cfg \"./models/experiments/h36m/hrnet/w48_384x288_adam_lr1e-3.yaml\" \n\ndef parse_args():\n    parser = argparse.ArgumentParser(description='Train keypoints network')\n    # general\n    parser.add_argument('--cfg',\n                        help='experiment configure file name',\n                        required=True,\n                        type=str)\n\n    parser.add_argument('opts',\n                        help=\"Modify config options using the command-line\",\n                        default=None,\n                        nargs=argparse.REMAINDER)\n\n    # philly\n    parser.add_argument('--modelDir',\n                        help='model directory',\n                        type=str,\n                        default='')\n    parser.add_argument('--logDir',\n                        help='log directory',\n                        type=str,\n                        default='')\n    parser.add_argument('--dataDir',\n                        help='data directory',\n                        type=str,\n                        default='')\n    parser.add_argument('--prevModelDir',\n                        help='prev Model directory',\n                        type=str,\n                        default='')\n    args = parser.parse_args()\n\n    return args\n\n\ndef main():\n    args = parse_args()\n    update_config(cfg, args)\n\n    logger, final_output_dir, tb_log_dir = create_logger(cfg, args.cfg, 'train')\n\n    logger.info(pprint.pformat(args))\n    logger.info(cfg)\n\n    # cudnn related setting\n    cudnn.benchmark = cfg.CUDNN.BENCHMARK\n    torch.backends.cudnn.deterministic = cfg.CUDNN.DETERMINISTIC\n    torch.backends.cudnn.enabled = cfg.CUDNN.ENABLED\n\n    model = get_pose_net(cfg, is_train=True)\n\n    dump_input = torch.rand(\n        (1, 3, cfg.MODEL.IMAGE_SIZE[1], cfg.MODEL.IMAGE_SIZE[0])\n    )\n\n    logger.info(get_model_summary(model, dump_input))\n\n    model = torch.nn.DataParallel(model, device_ids=cfg.GPUS).cuda()\n\n    # define loss function (criterion) and optimizer\n    criterion = JointsMSELoss(\n        use_target_weight=cfg.LOSS.USE_TARGET_WEIGHT\n    ).cuda()\n    \n    # coordinate loss with soft arg-max\n#    criterion = JointsCoordinateLoss(\n#        use_target_weight=cfg.LOSS.USE_TARGET_WEIGHT\n#    ).cuda()    \n    \n    # Wing Loss\n#    criterion = WingLoss(\n#        use_target_weight=cfg.LOSS.USE_TARGET_WEIGHT\n#    ).cuda()  \n    \n    # Data loading code\n    normalize = transforms.Normalize(\n        mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]\n    )\n    \n    train_dataset = H36MDataset(\n        cfg, True, cfg.DATASET.TRAIN_PATH,\n        transforms.Compose([\n            transforms.ToTensor(),\n            normalize,\n        ])\n    )\n    valid_dataset = H36MDataset(\n        cfg, False, cfg.DATASET.VALID_PATH,\n        transforms.Compose([\n            transforms.ToTensor(),\n            normalize,\n        ])\n    )\n\n    train_loader = torch.utils.data.DataLoader(\n        train_dataset,\n        batch_size=cfg.TRAIN.BATCH_SIZE_PER_GPU*len(cfg.GPUS),\n        shuffle=True,\n        num_workers=cfg.WORKERS,\n        pin_memory=cfg.PIN_MEMORY,\n    )\n    valid_loader = torch.utils.data.DataLoader(\n        valid_dataset,\n        batch_size=cfg.TEST.BATCH_SIZE_PER_GPU*len(cfg.GPUS),\n        shuffle=False,\n        num_workers=cfg.WORKERS,\n        pin_memory=cfg.PIN_MEMORY\n    )\n\n    best_perf = 0.0\n    best_model = False\n    last_epoch = -1\n    optimizer = get_optimizer(cfg, model)\n    begin_epoch = cfg.TRAIN.BEGIN_EPOCH\n    checkpoint_file = os.path.join(\n        final_output_dir, 'checkpoint.pth'\n    )\n    \n    if cfg.AUTO_RESUME and os.path.exists(checkpoint_file):\n        logger.info(\"=> loading checkpoint '{}'\".format(checkpoint_file))\n        checkpoint = torch.load(checkpoint_file)\n        begin_epoch = checkpoint['epoch']\n        best_perf = checkpoint['perf']\n        last_epoch = checkpoint['epoch']\n        model.load_state_dict(checkpoint['state_dict'])\n        optimizer.load_state_dict(checkpoint['optimizer'])\n        logger.info(\"=> loaded checkpoint '{}' (epoch {})\".format(\n            checkpoint_file, checkpoint['epoch']))\n\n    lr_scheduler = torch.optim.lr_scheduler.MultiStepLR(\n        optimizer, cfg.TRAIN.LR_STEP, cfg.TRAIN.LR_FACTOR,\n        last_epoch=last_epoch\n    )\n    \n    # inference\n    # perf_indicator = validate_pixel(\n    #     cfg, valid_loader, valid_dataset, model, criterion,\n    #     final_output_dir, tb_log_dir, save=True, split='test')    \n    # return     \n    \n    # training\n    # train with hard arg-max with MSE loss first and then fine-tune with\n    # soft-argmax coordinate loss works well in practice\n    #iterations = [3000, 3000, 6000, 3000, 3000]\n    # fine-tune with L1 loss\n    #iterations = [6000, 3000, 3000]\n    iterations = [6000, 6000, 6000, 3000, 3000]\n\n    for epoch in range(begin_epoch, cfg.TRAIN.END_EPOCH):\n        lr_scheduler.step()\n        \n        # set total iterations\n        if epoch - begin_epoch < len(iterations):\n            total_iters = iterations[epoch - begin_epoch]\n            logger.info(\"Total iterations to train: {}\".format(total_iters))\n        else:\n            total_iters = None\n            \n#       perform validation during training \n#        perf_indicator = validate_pixel(\n#            cfg, valid_loader, valid_dataset, model, criterion,\n#            final_output_dir, tb_log_dir)\n        \n        # train for one epoch\n        train(cfg, train_loader, model, criterion, optimizer, epoch,\n              final_output_dir, tb_log_dir, total_iters=total_iters)\n\n        #perf_indicator = 0.0\n        # evaluate on validation set\n        perf_indicator = validate_pixel(\n            cfg, valid_loader, valid_dataset, model, criterion,\n            final_output_dir, tb_log_dir)\n\n        if perf_indicator >= best_perf:\n            best_perf = perf_indicator\n            best_model = True\n        else:\n            best_model = False\n\n        logger.info('=> saving checkpoint to {}'.format(final_output_dir))\n        save_checkpoint({\n            'epoch': epoch + 1,\n            'model': cfg.MODEL.NAME,\n            'state_dict': model.state_dict(),\n            'best_state_dict': model.module.state_dict(),\n            'perf': perf_indicator,\n            'optimizer': optimizer.state_dict(),\n        }, best_model, final_output_dir)\n\n    final_model_state_file = os.path.join(\n        final_output_dir, 'final_state.pth'\n    )\n    logger.info('=> saving final model state to {}'.format(\n        final_model_state_file)\n    )\n    torch.save(model.module.state_dict(), final_model_state_file)\n\nif __name__ == '__main__':\n    main()\n"
  }
]