[
{
"path": ".gitignore",
"content": "*.pyc\n.torch\n_ext\n*.o\n_ext/\n*.png\n*.jpg\n*.tar\nlog/*\n"
},
{
"path": "LICENSE",
"content": "\n Apache License\n Version 2.0, January 2004\n http://www.apache.org/licenses/\n\n TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION\n\n 1. Definitions.\n\n \"License\" shall mean the terms and conditions for use, reproduction,\n and distribution as defined by Sections 1 through 9 of this document.\n\n \"Licensor\" shall mean the copyright owner or entity authorized by\n the copyright owner that is granting the License.\n\n \"Legal Entity\" shall mean the union of the acting entity and all\n other entities that control, are controlled by, or are under common\n control with that entity. For the purposes of this definition,\n \"control\" means (i) the power, direct or indirect, to cause the\n direction or management of such entity, whether by contract or\n otherwise, or (ii) ownership of fifty percent (50%) or more of the\n outstanding shares, or (iii) beneficial ownership of such entity.\n\n \"You\" (or \"Your\") shall mean an individual or Legal Entity\n exercising permissions granted by this License.\n\n \"Source\" form shall mean the preferred form for making modifications,\n including but not limited to software source code, documentation\n source, and configuration files.\n\n \"Object\" form shall mean any form resulting from mechanical\n transformation or translation of a Source form, including but\n not limited to compiled object code, generated documentation,\n and conversions to other media types.\n\n \"Work\" shall mean the work of authorship, whether in Source or\n Object form, made available under the License, as indicated by a\n copyright notice that is included in or attached to the work\n (an example is provided in the Appendix below).\n\n \"Derivative Works\" shall mean any work, whether in Source or Object\n form, that is based on (or derived from) the Work and for which the\n editorial revisions, annotations, elaborations, or other modifications\n represent, as a whole, an original work of authorship. For the purposes\n of this License, Derivative Works shall not include works that remain\n separable from, or merely link (or bind by name) to the interfaces of,\n the Work and Derivative Works thereof.\n\n \"Contribution\" shall mean any work of authorship, including\n the original version of the Work and any modifications or additions\n to that Work or Derivative Works thereof, that is intentionally\n submitted to Licensor for inclusion in the Work by the copyright owner\n or by an individual or Legal Entity authorized to submit on behalf of\n the copyright owner. For the purposes of this definition, \"submitted\"\n means any form of electronic, verbal, or written communication sent\n to the Licensor or its representatives, including but not limited to\n communication on electronic mailing lists, source code control systems,\n and issue tracking systems that are managed by, or on behalf of, the\n Licensor for the purpose of discussing and improving the Work, but\n excluding communication that is conspicuously marked or otherwise\n designated in writing by the copyright owner as \"Not a Contribution.\"\n\n \"Contributor\" shall mean Licensor and any individual or Legal Entity\n on behalf of whom a Contribution has been received by Licensor and\n subsequently incorporated within the Work.\n\n 2. Grant of Copyright License. Subject to the terms and conditions of\n this License, each Contributor hereby grants to You a perpetual,\n worldwide, non-exclusive, no-charge, royalty-free, irrevocable\n copyright license to reproduce, prepare Derivative Works of,\n publicly display, publicly perform, sublicense, and distribute the\n Work and such Derivative Works in Source or Object form.\n\n 3. Grant of Patent License. Subject to the terms and conditions of\n this License, each Contributor hereby grants to You a perpetual,\n worldwide, non-exclusive, no-charge, royalty-free, irrevocable\n (except as stated in this section) patent license to make, have made,\n use, offer to sell, sell, import, and otherwise transfer the Work,\n where such license applies only to those patent claims licensable\n by such Contributor that are necessarily infringed by their\n Contribution(s) alone or by combination of their Contribution(s)\n with the Work to which such Contribution(s) was submitted. If You\n institute patent litigation against any entity (including a\n cross-claim or counterclaim in a lawsuit) alleging that the Work\n or a Contribution incorporated within the Work constitutes direct\n or contributory patent infringement, then any patent licenses\n granted to You under this License for that Work shall terminate\n as of the date such litigation is filed.\n\n 4. Redistribution. You may reproduce and distribute copies of the\n Work or Derivative Works thereof in any medium, with or without\n modifications, and in Source or Object form, provided that You\n meet the following conditions:\n\n (a) You must give any other recipients of the Work or\n Derivative Works a copy of this License; and\n\n (b) You must cause any modified files to carry prominent notices\n stating that You changed the files; and\n\n (c) You must retain, in the Source form of any Derivative Works\n that You distribute, all copyright, patent, trademark, and\n attribution notices from the Source form of the Work,\n excluding those notices that do not pertain to any part of\n the Derivative Works; and\n\n (d) If the Work includes a \"NOTICE\" text file as part of its\n distribution, then any Derivative Works that You distribute must\n include a readable copy of the attribution notices contained\n within such NOTICE file, excluding those notices that do not\n pertain to any part of the Derivative Works, in at least one\n of the following places: within a NOTICE text file distributed\n as part of the Derivative Works; within the Source form or\n documentation, if provided along with the Derivative Works; or,\n within a display generated by the Derivative Works, if and\n wherever such third-party notices normally appear. The contents\n of the NOTICE file are for informational purposes only and\n do not modify the License. You may add Your own attribution\n notices within Derivative Works that You distribute, alongside\n or as an addendum to the NOTICE text from the Work, provided\n that such additional attribution notices cannot be construed\n as modifying the License.\n\n You may add Your own copyright statement to Your modifications and\n may provide additional or different license terms and conditions\n for use, reproduction, or distribution of Your modifications, or\n for any such Derivative Works as a whole, provided Your use,\n reproduction, and distribution of the Work otherwise complies with\n the conditions stated in this License.\n\n 5. Submission of Contributions. Unless You explicitly state otherwise,\n any Contribution intentionally submitted for inclusion in the Work\n by You to the Licensor shall be under the terms and conditions of\n this License, without any additional terms or conditions.\n Notwithstanding the above, nothing herein shall supersede or modify\n the terms of any separate license agreement you may have executed\n with Licensor regarding such Contributions.\n\n 6. Trademarks. This License does not grant permission to use the trade\n names, trademarks, service marks, or product names of the Licensor,\n except as required for reasonable and customary use in describing the\n origin of the Work and reproducing the content of the NOTICE file.\n\n 7. Disclaimer of Warranty. Unless required by applicable law or\n agreed to in writing, Licensor provides the Work (and each\n Contributor provides its Contributions) on an \"AS IS\" BASIS,\n WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or\n implied, including, without limitation, any warranties or conditions\n of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A\n PARTICULAR PURPOSE. You are solely responsible for determining the\n appropriateness of using or redistributing the Work and assume any\n risks associated with Your exercise of permissions under this License.\n\n 8. Limitation of Liability. In no event and under no legal theory,\n whether in tort (including negligence), contract, or otherwise,\n unless required by applicable law (such as deliberate and grossly\n negligent acts) or agreed to in writing, shall any Contributor be\n liable to You for damages, including any direct, indirect, special,\n incidental, or consequential damages of any character arising as a\n result of this License or out of the use or inability to use the\n Work (including but not limited to damages for loss of goodwill,\n work stoppage, computer failure or malfunction, or any and all\n other commercial damages or losses), even if such Contributor\n has been advised of the possibility of such damages.\n\n 9. Accepting Warranty or Additional Liability. While redistributing\n the Work or Derivative Works thereof, You may choose to offer,\n and charge a fee for, acceptance of support, warranty, indemnity,\n or other liability obligations and/or rights consistent with this\n License. However, in accepting such obligations, You may act only\n on Your own behalf and on Your sole responsibility, not on behalf\n of any other Contributor, and only if You agree to indemnify,\n defend, and hold each Contributor harmless for any liability\n incurred by, or claims asserted against, such Contributor by reason\n of your accepting any such warranty or additional liability.\n\n END OF TERMS AND CONDITIONS\n\n APPENDIX: How to apply the Apache License to your work.\n\n To apply the Apache License to your work, attach the following\n boilerplate notice, with the fields enclosed by brackets \"[]\"\n replaced with your own identifying information. (Don't include\n the brackets!) The text should be enclosed in the appropriate\n comment syntax for the file format. We also recommend that a\n file or class name and description of purpose be included on the\n same \"printed page\" as the copyright notice for easier\n identification within third-party archives.\n\n Copyright [yyyy] [name of copyright owner]\n\n Licensed under the Apache License, Version 2.0 (the \"License\");\n you may not use this file except in compliance with the License.\n You may obtain a copy of the License at\n\n http://www.apache.org/licenses/LICENSE-2.0\n\n Unless required by applicable law or agreed to in writing, software\n distributed under the License is distributed on an \"AS IS\" BASIS,\n WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n See the License for the specific language governing permissions and\n limitations under the License.\n"
},
{
"path": "README.md",
"content": "# Deep Online Fused Video Stabilization\n\n[[Paper]](https://openaccess.thecvf.com/content/WACV2022/papers/Shi_Deep_Online_Fused_Video_Stabilization_WACV_2022_paper.pdf)[[Supplementary]](https://zhmeishi.github.io/dvs/paper/dvs_supp.pdf) [[Project Page]](https://zhmeishi.github.io/dvs/) [[Dataset]](https://storage.googleapis.com/dataset_release/all.zip) [[Our Result]](https://storage.googleapis.com/dataset_release/inference_result_release.zip) [[More Results]](https://zhmeishi.github.io/dvs/supp/results.html) \n\nThis repository contains the Pytorch implementation of our method in the paper \"Deep Online Fused Video Stabilization\".\n\n## Environment Setting\nPython version >= 3.6\nPytorch with CUDA >= 1.0.0 (guide is [here](https://pytorch.org/get-started/locally/))\nInstall other used packages:\n```\ncd dvs\npip install -r requirements.txt --ignore-installed\n```\n\n## Data Preparation\nDownload sample video [here](https://drive.google.com/file/d/1PpF3-6BbQKy9fldjIfwa5AlbtQflx3sG/view?usp=sharing).\nUncompress the *video* folder under the *dvs* folder.\n```\npython load_frame_sensor_data.py \n```\nDemo of curve visualization:\nThe **gyro/OIS curve visualization** can be found at *dvs/video/s_114_outdoor_running_trail_daytime/ControlCam_20200930_104820_real.jpg*.\n\n\n## FlowNet2 Preparation\nNote, we provide optical flow result of one test video in our Data Preparation. If you would like to generate them for all test videos, please follow [FlowNet2 official website](https://github.com/NVIDIA/flownet2-pytorch) and guide below. Otherwise, you can skip this section. \n\nNote, FlowNet2 installation is tricky. Please use Python=3.6 and Pytorch=1.0.0. More details are [here](https://github.com/NVIDIA/flownet2-pytorch/issues/156) or contact us for any questions.\n\nDownload FlowNet2 model *FlowNet2_checkpoint.pth.tar* [here](https://drive.google.com/file/d/1hF8vS6YeHkx3j2pfCeQqqZGwA_PJq_Da/view). Move it under folder *dvs/flownet2*.\n```\npython warp/read_write.py # video2frames\ncd flownet2\nbash install.sh # install package\nbash run.sh # generate optical flow file for dataset\n``` \n\n## Running Inference \n```\npython inference.py\npython metrics.py\n``` \nThe loss and metric information will be printed in the terminal. The metric numbers can be slightly different due to difference on opencv/pytorch versions. \n\nThe result is under *dvs/test/stabilzation*. \nIn *s_114_outdoor_running_trail_daytime.jpg*, the blue curve is the output of our models, and the green curve is the input. \n*s_114_outdoor_running_trail_daytime_stab.mp4* is uncropped stabilized video. \n*s_114_outdoor_running_trail_daytime_stab_crop.mp4* is cropped stabilized video. Note, the cropped video is generated after running the metrics code. \n\n## Training\nDownload dataset for training and test [here](https://storage.googleapis.com/dataset_release/all.zip). \nUncompress *all.zip* and move *dataset_release* folder under the *dvs* folder.\n\nFollow FlowNet2 Preparation Section.\n```\npython warp/read_write.py --dir_path ./dataset_release # video2frames\ncd flownet2\nbash run_release.sh # generate optical flow file for dataset\n``` \n\nRun training code.\n```\npython train.py\n``` \nThe model is saved in *checkpoint/stabilzation_train*.\n\n## Citation \nIf you use this code or dataset for your research, please cite our paper.\n```\n@inproceedings{shi2022deep,\n title={Deep Online Fused Video Stabilization},\n author={Shi, Zhenmei and Shi, Fuhao and Lai, Wei-Sheng and Liang, Chia-Kai and Liang, Yingyu},\n booktitle={Proceedings of the IEEE/CVF Winter Conference on Applications of Computer Vision},\n pages={1250--1258},\n year={2022}\n}\n```\n"
},
{
"path": "docs/code-of-conduct.md",
"content": "# Google Open Source Community Guidelines\n\nAt Google, we recognize and celebrate the creativity and collaboration of open\nsource contributors and the diversity of skills, experiences, cultures, and\nopinions they bring to the projects and communities they participate in.\n\nEvery one of Google's open source projects and communities are inclusive\nenvironments, based on treating all individuals respectfully, regardless of\ngender identity and expression, sexual orientation, disabilities,\nneurodiversity, physical appearance, body size, ethnicity, nationality, race,\nage, religion, or similar personal characteristic.\n\nWe value diverse opinions, but we value respectful behavior more.\n\nRespectful behavior includes:\n\n* Being considerate, kind, constructive, and helpful.\n* Not engaging in demeaning, discriminatory, harassing, hateful, sexualized, or\n physically threatening behavior, speech, and imagery.\n* Not engaging in unwanted physical contact.\n\nSome Google open source projects [may adopt][] an explicit project code of\nconduct, which may have additional detailed expectations for participants. Most\nof those projects will use our [modified Contributor Covenant][].\n\n[may adopt]: https://opensource.google/docs/releasing/preparing/#conduct\n[modified Contributor Covenant]: https://opensource.google/docs/releasing/template/CODE_OF_CONDUCT/\n\n## Resolve peacefully\n\nWe do not believe that all conflict is necessarily bad; healthy debate and\ndisagreement often yields positive results. However, it is never okay to be\ndisrespectful.\n\nIf you see someone behaving disrespectfully, you are encouraged to address the\nbehavior directly with those involved. Many issues can be resolved quickly and\neasily, and this gives people more control over the outcome of their dispute.\nIf you are unable to resolve the matter for any reason, or if the behavior is\nthreatening or harassing, report it. We are dedicated to providing an\nenvironment where participants feel welcome and safe.\n\n## Reporting problems\n\nSome Google open source projects may adopt a project-specific code of conduct.\nIn those cases, a Google employee will be identified as the Project Steward,\nwho will receive and handle reports of code of conduct violations. In the event\nthat a project hasn’t identified a Project Steward, you can report problems by\nemailing opensource@google.com.\n\nWe will investigate every complaint, but you may not receive a direct response.\nWe will use our discretion in determining when and how to follow up on reported\nincidents, which may range from not taking action to permanent expulsion from\nthe project and project-sponsored spaces. We will notify the accused of the\nreport and provide them an opportunity to discuss it before any action is\ntaken. The identity of the reporter will be omitted from the details of the\nreport supplied to the accused. In potentially harmful situations, such as\nongoing harassment or threats to anyone's safety, we may take action without\nnotice.\n\n*This document was adapted from the [IndieWeb Code of Conduct][] and can also\nbe found at .*\n\n[IndieWeb Code of Conduct]: https://indieweb.org/code-of-conduct\n"
},
{
"path": "docs/contributing.md",
"content": "# How to Contribute\n\nWe'd love to accept your patches and contributions to this project. There are\njust a few small guidelines you need to follow.\n\n## Contributor License Agreement\n\nContributions to this project must be accompanied by a Contributor License\nAgreement. You (or your employer) retain the copyright to your contribution;\nthis simply gives us permission to use and redistribute your contributions as\npart of the project. Head over to to see\nyour current agreements on file or to sign a new one.\n\nYou generally only need to submit a CLA once, so if you've already submitted one\n(even if it was for a different project), you probably don't need to do it\nagain.\n\n## Code reviews\n\nAll submissions, including submissions by project members, require review. We\nuse GitHub pull requests for this purpose. Consult\n[GitHub Help](https://help.github.com/articles/about-pull-requests/) for more\ninformation on using pull requests.\n\n## Community Guidelines\n\nThis project follows [Google's Open Source Community\nGuidelines](https://opensource.google/conduct/).\n"
},
{
"path": "dvs/conf/stabilzation.yaml",
"content": "data:\n exp: 'stabilzation'\n checkpoints_dir: './checkpoint'\n log: './log'\n data_dir: './video' \n use_cuda: true\n batch_size: 16\n resize_ratio: 0.25\n number_real: 10\n number_virtual: 2\n time_train: 2000 # ms\n sample_freq: 40 # ms\n channel_size: 1\n num_workers: 16 # num_workers for data_loader\nmodel:\n load_model: null\n cnn:\n activate_function: relu # sigmoid, relu, tanh, quadratic\n batch_norm: true\n gap: false\n layers:\n rnn:\n layers: \n - - 512 \n - true \n - - 512 \n - true \n fc:\n activate_function: relu\n batch_norm: false # (batch_norm and drop_out) is False\n layers: \n - - 256 \n - true \n - - 4 # last layer should be equal to nr_class\n - true\n drop_out: 0\ntrain:\n optimizer: \"adam\" # adam or sgd\n momentum: 0.9 # for sgd\n decay_epoch: null\n epoch: 400\n snapshot: 2\n init_lr: 0.0001\n lr_decay: 0.5\n lr_step: 200 # if > 0 decay_epoch should be null\n seed: 1\n weight_decay: 0.0001\n clip_norm: False\n init: \"xavier_uniform\" # xavier_uniform or xavier_normal\nloss:\n follow: 10\n angle: 1\n smooth: 10 #10\n c2_smooth: 200 #20\n undefine: 2.0\n opt: 0.1\n stay: 0"
},
{
"path": "dvs/conf/stabilzation_train.yaml",
"content": "data:\n exp: 'stabilzation_train'\n checkpoints_dir: './checkpoint'\n log: './log'\n data_dir: './dataset_release' \n use_cuda: true\n batch_size: 16\n resize_ratio: 0.25\n number_real: 10\n number_virtual: 2\n time_train: 2000 # ms\n sample_freq: 40 # ms\n channel_size: 1\n num_workers: 16 # num_workers for data_loader\nmodel:\n load_model: null\n cnn:\n activate_function: relu # sigmoid, relu, tanh, quadratic\n batch_norm: true\n gap: false\n layers:\n rnn:\n layers: \n - - 512 \n - true \n - - 512 \n - true \n fc:\n activate_function: relu\n batch_norm: false # (batch_norm and drop_out) is False\n layers: \n - - 256 \n - true \n - - 4 # last layer should be equal to nr_class\n - true\n drop_out: 0\ntrain:\n optimizer: \"adam\" # adam or sgd\n momentum: 0.9 # for sgd\n decay_epoch: null\n epoch: 400\n snapshot: 2\n init_lr: 0.0001\n lr_decay: 0.5\n lr_step: 200 # if > 0 decay_epoch should be null\n seed: 1\n weight_decay: 0.0001\n clip_norm: False\n init: \"xavier_uniform\" # xavier_uniform or xavier_normal\nloss:\n follow: 10\n angle: 1\n smooth: 10 #10\n c2_smooth: 200 #20\n undefine: 2.0\n opt: 0.1\n stay: 0"
},
{
"path": "dvs/dataset.py",
"content": "from torch.utils.data import Dataset\nimport os\nimport collections\nfrom gyro import (\n LoadGyroData, \n LoadOISData, \n LoadFrameData, \n GetGyroAtTimeStamp, \n get_static, \n GetMetadata, \n GetProjections, \n train_GetGyroAtTimeStamp,\n QuaternionProduct,\n QuaternionReciprocal,\n FindOISAtTimeStamp,\n norm_quat\n )\nimport random\nimport numpy as np\nimport torchvision.transforms as transforms\nimport torch\nfrom flownet2 import flow_utils\nfrom scipy import ndimage, misc\nfrom numpy import linalg as LA\n\ndef get_data_loader(cf, no_flo = False):\n size = cf[\"data\"][\"batch_size\"]\n num_workers = cf[\"data\"][\"num_workers\"]\n train_data, test_data = get_dataset(cf, no_flo)\n trainloader = torch.utils.data.DataLoader(train_data, batch_size=size,shuffle=True, pin_memory=True, num_workers=num_workers)\n testloader = torch.utils.data.DataLoader(test_data, batch_size=size,shuffle=False, pin_memory=True, num_workers=num_workers)\n return trainloader,testloader\n\ndef get_dataset(cf, no_flo = False):\n resize_ratio = cf[\"data\"][\"resize_ratio\"]\n train_transform, test_transform = _data_transforms()\n train_path = os.path.join(cf[\"data\"][\"data_dir\"], \"training\")\n test_path = os.path.join(cf[\"data\"][\"data_dir\"], \"test\")\n if not os.path.exists(train_path):\n train_path = cf[\"data\"][\"data_dir\"]\n if not os.path.exists(test_path):\n test_path = cf[\"data\"][\"data_dir\"]\n\n train_data = Dataset_Gyro(\n train_path, sample_freq = cf[\"data\"][\"sample_freq\"]*1000000, number_real = cf[\"data\"][\"number_real\"], \n time_train = cf[\"data\"][\"time_train\"]*1000000, transform = train_transform, resize_ratio = resize_ratio, no_flo = no_flo)\n test_data = Dataset_Gyro(\n test_path, sample_freq = cf[\"data\"][\"sample_freq\"]*1000000, number_real = cf[\"data\"][\"number_real\"], \n time_train = cf[\"data\"][\"time_train\"]*1000000, transform = test_transform, resize_ratio = resize_ratio, no_flo = no_flo)\n return train_data, test_data\n\ndef get_inference_data_loader(cf, data_path, no_flo = False):\n test_data = get_inference_dataset(cf, data_path, no_flo)\n testloader = torch.utils.data.DataLoader(test_data, batch_size=1,shuffle=False, pin_memory=True, num_workers=1)\n return testloader\n\ndef get_inference_dataset(cf, data_path, no_flo = False):\n resize_ratio = cf[\"data\"][\"resize_ratio\"]\n _, test_transform = _data_transforms()\n test_data = Dataset_Gyro(\n data_path, sample_freq = cf[\"data\"][\"sample_freq\"]*1000000, number_real = cf[\"data\"][\"number_real\"], \n time_train = cf[\"data\"][\"time_train\"]*1000000, transform = test_transform, resize_ratio = resize_ratio,\n inference_only = True, no_flo = no_flo)\n return test_data\n\ndef _data_transforms():\n\n test_transform = transforms.Compose(\n [transforms.ToTensor(),\n ])\n train_transform = transforms.Compose(\n [transforms.ToTensor(),\n ])\n\n return train_transform, test_transform\n\nclass DVS_data():\n def __init__(self):\n self.gyro = None\n self.ois = None\n self.frame = None\n self.length = 0\n self.flo_path = None\n self.flo_shape = None\n self.flo_back_path = None\n\nclass Dataset_Gyro(Dataset):\n def __init__(self, path, sample_freq = 33*1000000, number_real = 10, time_train = 2000*1000000, \\\n transform = None, inference_only = False, no_flo = False, resize_ratio = 1): \n r\"\"\"\n Arguments:\n sample_freq: real quaternions [t-sample_freq*number_real, t+sample_freq*number_real] ns\n number_real: real gyro num in half time_interval\n time_train: time for a batch ns\n \"\"\"\n self.sample_freq = sample_freq\n self.number_real = number_real\n self.no_flo = no_flo\n self.resize_ratio = resize_ratio\n self.static_options = get_static()\n self.inference_only = inference_only\n\n self.ois_ratio = np.array([self.static_options[\"crop_window_width\"] / self.static_options[\"width\"], \\\n self.static_options[\"crop_window_height\"] / self.static_options[\"height\"]]) * 0.01\n self.unit_size = 4\n\n if inference_only:\n self.length = 1\n self.data = [self.process_one_video(path)]\n self.number_train = self.data[0].length \n return\n\n self.time_train = time_train\n self.number_train = time_train//self.sample_freq\n \n self.data_name = sorted(os.listdir(path))\n self.length = len(self.data_name)\n self.data = []\n for i in range(self.length):\n self.data.append(self.process_one_video(os.path.join(path,self.data_name[i])))\n \n def process_one_video(self, path):\n dvs_data = DVS_data()\n files = sorted(os.listdir(path))\n print(path)\n for f in files:\n file_path = os.path.join(path,f)\n if \"gimbal\" in file_path.lower():\n continue\n if \"frame\" in f and \"txt\" in f:\n dvs_data.frame = LoadFrameData(file_path)\n print(\"frame:\", dvs_data.frame.shape, end=\" \")\n elif \"gyro\" in f:\n dvs_data.gyro = LoadGyroData(file_path)\n dvs_data.gyro = preprocess_gyro(dvs_data.gyro) \n print(\"gyro:\", dvs_data.gyro.shape, end=\" \")\n elif \"ois\" in f and \"txt\" in f:\n dvs_data.ois = LoadOISData(file_path)\n print(\"ois:\", dvs_data.ois.shape, end=\" \")\n elif f == \"flo\":\n dvs_data.flo_path, dvs_data.flo_shape = LoadFlow(file_path)\n print(\"flo_path:\", len(dvs_data.flo_path), end=\" \")\n print(\"flo_shape:\", dvs_data.flo_shape, end=\" \")\n elif f == \"flo_back\":\n dvs_data.flo_back_path, _ = LoadFlow(file_path)\n \n print()\n if dvs_data.flo_path is not None:\n dvs_data.length = min(dvs_data.frame.shape[0] - 1, len(dvs_data.flo_path))\n else:\n dvs_data.length = dvs_data.frame.shape[0] - 1 \n return dvs_data\n\n def generate_quaternions(self, dvs_data):\n first_id = random.randint(0, dvs_data.length - self.number_train) + 1 # skip the first frame\n\n sample_data = np.zeros((self.number_train, 2 * self.number_real + 1, self.unit_size), dtype=np.float32)\n sample_ois = np.zeros((self.number_train, 2), dtype=np.float32)\n\n sample_time = np.zeros((self.number_train+1), dtype=np.float32)\n sample_time[0] = get_timestamp(dvs_data.frame, first_id - 1)\n\n real_postion = np.zeros((self.number_train, 4), dtype=np.float32)\n\n time_start = sample_time[0]\n\n for i in range(self.number_train):\n sample_time[i+1] = get_timestamp(dvs_data.frame, first_id + i)\n real_postion[i] = GetGyroAtTimeStamp(dvs_data.gyro, sample_time[i+1] - self.sample_freq)\n sample_ois[i] = self.get_ois_at_timestamp(dvs_data.ois, sample_time[i+1])\n for j in range(-self.number_real, self.number_real+1):\n index = j + self.number_real\n time_stamp = sample_time[i+1] + self.sample_freq * j \n sample_data[i, index] = self.get_data_at_timestamp(dvs_data.gyro, dvs_data.ois, time_stamp, real_postion[i])\n \n sample_data = np.reshape(sample_data, (self.number_train, (2*self.number_real+1) * self.unit_size))\n return sample_data, sample_time, first_id, real_postion, sample_ois\n\n def load_flo(self, idx, first_id):\n shape = self.data[idx].flo_shape\n h, w = shape[0], shape[1]\n flo = np.zeros((self.number_train, h, w, 2))\n flo_back = np.zeros((self.number_train, h, w, 2))\n\n for i in range(self.number_train):\n frame_id = i + first_id\n f = flow_utils.readFlow(self.data[idx].flo_path[frame_id-1]).astype(np.float32) \n flo[i] = f\n\n f_b = flow_utils.readFlow(self.data[idx].flo_back_path[frame_id-1]).astype(np.float32) \n flo_back[i] = f_b\n\n return flo, flo_back\n\n def load_real_projections(self, idx, first_id):\n real_projections = np.zeros((self.number_train + 1, self.static_options[\"num_grid_rows\"], 3, 3))\n for i in range(self.number_train + 1):\n frame_id = i + first_id\n metadata = GetMetadata(self.data[idx].frame, frame_id - 1)\n real_projections[i] = np.array(GetProjections(self.static_options, metadata, self.data[idx].gyro, np.zeros(self.data[idx].ois.shape), no_shutter = True))\n return real_projections\n\n def __getitem__(self, idx):\n inputs, times, first_id, real_postion, ois = self.generate_quaternions(self.data[idx]) \n real_projections = self.load_real_projections(idx, first_id)\n if self.no_flo:\n flo, flo_back = 0, 0\n else:\n flo, flo_back = self.load_flo(idx, first_id)\n return inputs, times, flo, flo_back, real_projections, real_postion, ois, idx\n\n def __len__(self):\n return self.length\n\n def get_virtual_data(self, virtual_queue, real_queue_idx, pre_times, cur_times, time_start, batch_size, number_virtual, quat_t_1):\n # virtual_queue: [batch_size, num, 5 (timestamp, quats)]\n # eular angle, \n # deta R angular velocity [Q't-1, Q't-2] \n # output virtual angular velocity, x, x*dtime => detaQt\n virtual_data = np.zeros((batch_size, number_virtual, 4), dtype=np.float32)\n vt_1 = np.zeros((batch_size, 4), dtype=np.float32)\n quat_t_1 = quat_t_1.numpy()\n for i in range(batch_size):\n sample_time = cur_times[i]\n for j in range(number_virtual):\n time_stamp = sample_time - self.sample_freq * (number_virtual - j) \n virtual_data[i, j] = get_virtual_at_timestamp(virtual_queue[i], self.data[real_queue_idx[i]].gyro, time_stamp, time_start[i], quat_t_1[i])\n vt_1[i] = get_virtual_at_timestamp(virtual_queue[i], self.data[real_queue_idx[i]].gyro, pre_times[i], time_start[i], None)\n virtual_data = np.reshape(virtual_data, (batch_size, number_virtual * 4))\n return torch.tensor(virtual_data, dtype=torch.float), torch.tensor(vt_1, dtype=torch.float)\n\n def update_virtual_queue(self, batch_size, virtual_queue, out, times):\n virtual_data = np.zeros((batch_size, 5))\n virtual_data[:,0] = times\n virtual_data[:, 1:] = out\n virtual_data = np.expand_dims(virtual_data, axis = 1)\n\n if None in virtual_queue:\n virtual_queue = virtual_data\n else:\n virtual_queue = np.concatenate((virtual_queue, virtual_data), axis = 1)\n return virtual_queue\n\n def random_init_virtual_queue(self, batch_size, real_postion, times):\n virtual_queue = np.zeros((batch_size, 3, 5))\n virtual_queue[:, 2, 0] = times - 0.1 * self.sample_freq\n virtual_queue[:, 1, 0] = times - 1.1 * self.sample_freq\n virtual_queue[:, 0, 0] = times - 2.1 * self.sample_freq\n for i in range(batch_size):\n quat = np.random.uniform(low=-0.06, high= 0.06, size=4) # transfer to angle # 0.05\n quat[3] = 1\n quat = quat / LA.norm(quat)\n quat = norm_quat(QuaternionProduct(real_postion[i], quat))\n virtual_queue[i, 2, 1:] = quat\n virtual_queue[i, 1, 1:] = quat\n virtual_queue[i, 0, 1:] = quat\n return virtual_queue\n\n def get_data_at_timestamp(self, gyro_data, ois_data, time_stamp, quat_t_1):\n quat_t = GetGyroAtTimeStamp(gyro_data, time_stamp)\n quat_dif = QuaternionProduct(quat_t, QuaternionReciprocal(quat_t_1)) \n return quat_dif\n\n def get_ois_at_timestamp(self, ois_data, time_stamp):\n ois_t = FindOISAtTimeStamp(ois_data, time_stamp)\n ois_t = np.array(ois_t) / self.ois_ratio\n return ois_t\n\ndef get_timestamp(frame_data, idx):\n sample_time = frame_data[idx, 0]\n metadata = GetMetadata(frame_data, idx)\n timestmap_ns = metadata[\"timestamp_ns\"] + metadata[\"rs_time_ns\"] * 0.5\n return timestmap_ns\n\ndef preprocess_gyro(gyro, extend = 200):\n fake_gyro = np.zeros((extend, 5))\n time_start = gyro[0,0]\n for i in range(extend):\n fake_gyro[-i-1, 0] = time_start - (gyro[i+1, 0] - time_start)\n fake_gyro[-i-1, 4] = gyro[i+1, 4]\n fake_gyro[-i-1, 1:4] = -gyro[i+1, 1:4]\n\n new_gyro = np.concatenate((fake_gyro, gyro), axis = 0)\n return new_gyro\n\ndef LoadFlow(path):\n file_names = sorted(os.listdir(path))\n file_path =[]\n for n in file_names:\n file_path.append(os.path.join(path, n))\n return file_path, flow_utils.readFlow(file_path[0]).shape\n\ndef get_virtual_at_timestamp(virtual_queue, real_queue, time_stamp, time_start, quat_t_1 = None, sample_freq = None):\n if virtual_queue is None:\n quat_t = GetGyroAtTimeStamp(real_queue, time_stamp)\n else:\n quat_t = train_GetGyroAtTimeStamp(virtual_queue, time_stamp)\n if quat_t is None:\n quat_t = GetGyroAtTimeStamp(real_queue, time_stamp)\n \n if quat_t_1 is None:\n return quat_t\n else:\n quat_dif = QuaternionProduct(quat_t, QuaternionReciprocal(quat_t_1)) \n return quat_dif\n"
},
{
"path": "dvs/flownet2/LICENSE",
"content": "Copyright 2017 NVIDIA CORPORATION\n\nLicensed under the Apache License, Version 2.0 (the \"License\");\nyou may not use this file except in compliance with the License.\nYou may obtain a copy of the License at\n\n http://www.apache.org/licenses/LICENSE-2.0\n\nUnless required by applicable law or agreed to in writing, software\ndistributed under the License is distributed on an \"AS IS\" BASIS,\nWITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\nSee the License for the specific language governing permissions and\nlimitations under the License."
},
{
"path": "dvs/flownet2/README.md",
"content": "# flownet2-pytorch \n\nPytorch implementation of [FlowNet 2.0: Evolution of Optical Flow Estimation with Deep Networks](https://arxiv.org/abs/1612.01925). \n\nMultiple GPU training is supported, and the code provides examples for training or inference on [MPI-Sintel](http://sintel.is.tue.mpg.de/) clean and final datasets. The same commands can be used for training or inference with other datasets. See below for more detail.\n\nInference using fp16 (half-precision) is also supported.\n\nFor more help, type
\n \n python main.py --help\n\n## Network architectures\nBelow are the different flownet neural network architectures that are provided.
\nA batchnorm version for each network is also available.\n\n - **FlowNet2S**\n - **FlowNet2C**\n - **FlowNet2CS**\n - **FlowNet2CSS**\n - **FlowNet2SD**\n - **FlowNet2**\n\n## Custom layers\n\n`FlowNet2` or `FlowNet2C*` achitectures rely on custom layers `Resample2d` or `Correlation`.
\nA pytorch implementation of these layers with cuda kernels are available at [./networks](./networks).
\nNote : Currently, half precision kernels are not available for these layers.\n\n## Data Loaders\n\nDataloaders for FlyingChairs, FlyingThings, ChairsSDHom and ImagesFromFolder are available in [datasets.py](./datasets.py).
\n\n## Loss Functions\n\nL1 and L2 losses with multi-scale support are available in [losses.py](./losses.py).
\n\n## Installation \n\n # get flownet2-pytorch source\n git clone https://github.com/NVIDIA/flownet2-pytorch.git\n cd flownet2-pytorch\n\n # install custom layers\n bash install.sh\n \n### Python requirements \nCurrently, the code supports python 3\n* numpy \n* PyTorch ( == 0.4.1, for <= 0.4.0 see branch [python36-PyTorch0.4](https://github.com/NVIDIA/flownet2-pytorch/tree/python36-PyTorch0.4))\n* scipy \n* scikit-image\n* tensorboardX\n* colorama, tqdm, setproctitle \n\n## Converted Caffe Pre-trained Models\nWe've included caffe pre-trained models. Should you use these pre-trained weights, please adhere to the [license agreements](https://drive.google.com/file/d/1TVv0BnNFh3rpHZvD-easMb9jYrPE2Eqd/view?usp=sharing). \n\n* [FlowNet2](https://drive.google.com/file/d/1hF8vS6YeHkx3j2pfCeQqqZGwA_PJq_Da/view?usp=sharing)[620MB]\n* [FlowNet2-C](https://drive.google.com/file/d/1BFT6b7KgKJC8rA59RmOVAXRM_S7aSfKE/view?usp=sharing)[149MB]\n* [FlowNet2-CS](https://drive.google.com/file/d/1iBJ1_o7PloaINpa8m7u_7TsLCX0Dt_jS/view?usp=sharing)[297MB]\n* [FlowNet2-CSS](https://drive.google.com/file/d/157zuzVf4YMN6ABAQgZc8rRmR5cgWzSu8/view?usp=sharing)[445MB]\n* [FlowNet2-CSS-ft-sd](https://drive.google.com/file/d/1R5xafCIzJCXc8ia4TGfC65irmTNiMg6u/view?usp=sharing)[445MB]\n* [FlowNet2-S](https://drive.google.com/file/d/1V61dZjFomwlynwlYklJHC-TLfdFom3Lg/view?usp=sharing)[148MB]\n* [FlowNet2-SD](https://drive.google.com/file/d/1QW03eyYG_vD-dT-Mx4wopYvtPu_msTKn/view?usp=sharing)[173MB]\n \n## Inference\n # Example on MPISintel Clean \n python main.py --inference --model FlowNet2 --save_flow --inference_dataset MpiSintelClean \\\n --inference_dataset_root /path/to/mpi-sintel/clean/dataset \\\n --resume /path/to/checkpoints \n \n## Training and validation\n\n # Example on MPISintel Final and Clean, with L1Loss on FlowNet2 model\n python main.py --batch_size 8 --model FlowNet2 --loss=L1Loss --optimizer=Adam --optimizer_lr=1e-4 \\\n --training_dataset MpiSintelFinal --training_dataset_root /path/to/mpi-sintel/final/dataset \\\n --validation_dataset MpiSintelClean --validation_dataset_root /path/to/mpi-sintel/clean/dataset\n\n # Example on MPISintel Final and Clean, with MultiScale loss on FlowNet2C model \n python main.py --batch_size 8 --model FlowNet2C --optimizer=Adam --optimizer_lr=1e-4 --loss=MultiScale --loss_norm=L1 \\\n --loss_numScales=5 --loss_startScale=4 --optimizer_lr=1e-4 --crop_size 384 512 \\\n --training_dataset FlyingChairs --training_dataset_root /path/to/flying-chairs/dataset \\\n --validation_dataset MpiSintelClean --validation_dataset_root /path/to/mpi-sintel/clean/dataset\n \n## Results on MPI-Sintel\n[](https://www.youtube.com/watch?v=HtBmabY8aeU \"Predicted flows on MPI-Sintel\")\n\n## Reference \nIf you find this implementation useful in your work, please acknowledge it appropriately and cite the paper:\n````\n@InProceedings{IMKDB17,\n author = \"E. Ilg and N. Mayer and T. Saikia and M. Keuper and A. Dosovitskiy and T. Brox\",\n title = \"FlowNet 2.0: Evolution of Optical Flow Estimation with Deep Networks\",\n booktitle = \"IEEE Conference on Computer Vision and Pattern Recognition (CVPR)\",\n month = \"Jul\",\n year = \"2017\",\n url = \"http://lmb.informatik.uni-freiburg.de//Publications/2017/IMKDB17\"\n}\n````\n```\n@misc{flownet2-pytorch,\n author = {Fitsum Reda and Robert Pottorff and Jon Barker and Bryan Catanzaro},\n title = {flownet2-pytorch: Pytorch implementation of FlowNet 2.0: Evolution of Optical Flow Estimation with Deep Networks},\n year = {2017},\n publisher = {GitHub},\n journal = {GitHub repository},\n howpublished = {\\url{https://github.com/NVIDIA/flownet2-pytorch}}\n}\n```\n## Related Optical Flow Work from Nvidia \nCode (in Caffe and Pytorch): [PWC-Net](https://github.com/NVlabs/PWC-Net)
\nPaper : [PWC-Net: CNNs for Optical Flow Using Pyramid, Warping, and Cost Volume](https://arxiv.org/abs/1709.02371). \n\n## Acknowledgments\nParts of this code were derived, as noted in the code, from [ClementPinard/FlowNetPytorch](https://github.com/ClementPinard/FlowNetPytorch).\n"
},
{
"path": "dvs/flownet2/__init__.py",
"content": "from .utils import flow_utils, tools"
},
{
"path": "dvs/flownet2/convert.py",
"content": "#!/usr/bin/env python2.7\n\nimport caffe\nfrom caffe.proto import caffe_pb2\nimport sys, os\n\nimport torch\nimport torch.nn as nn\n\nimport argparse, tempfile\nimport numpy as np\n\nparser = argparse.ArgumentParser()\nparser.add_argument('caffe_model', help='input model in hdf5 or caffemodel format')\nparser.add_argument('prototxt_template',help='prototxt template')\nparser.add_argument('flownet2_pytorch', help='path to flownet2-pytorch')\n\nargs = parser.parse_args()\n\nargs.rgb_max = 255\nargs.fp16 = False\nargs.grads = {}\n\n# load models\nsys.path.append(args.flownet2_pytorch)\n\nimport models\nfrom utils.param_utils import *\n\nwidth = 256\nheight = 256\nkeys = {'TARGET_WIDTH': width, \n 'TARGET_HEIGHT': height,\n 'ADAPTED_WIDTH':width,\n 'ADAPTED_HEIGHT':height,\n 'SCALE_WIDTH':1.,\n 'SCALE_HEIGHT':1.,}\n\ntemplate = '\\n'.join(np.loadtxt(args.prototxt_template, dtype=str, delimiter='\\n'))\nfor k in keys:\n template = template.replace('$%s$'%(k),str(keys[k]))\n\nprototxt = tempfile.NamedTemporaryFile(mode='w', delete=True)\nprototxt.write(template)\nprototxt.flush()\n\nnet = caffe.Net(prototxt.name, args.caffe_model, caffe.TEST)\n\nweights = {}\nbiases = {}\n\nfor k, v in list(net.params.items()):\n weights[k] = np.array(v[0].data).reshape(v[0].data.shape)\n biases[k] = np.array(v[1].data).reshape(v[1].data.shape)\n print((k, weights[k].shape, biases[k].shape))\n\nif 'FlowNet2/' in args.caffe_model:\n model = models.FlowNet2(args)\n\n parse_flownetc(model.flownetc.modules(), weights, biases)\n parse_flownets(model.flownets_1.modules(), weights, biases, param_prefix='net2_')\n parse_flownets(model.flownets_2.modules(), weights, biases, param_prefix='net3_')\n parse_flownetsd(model.flownets_d.modules(), weights, biases, param_prefix='netsd_')\n parse_flownetfusion(model.flownetfusion.modules(), weights, biases, param_prefix='fuse_')\n\n state = {'epoch': 0,\n 'state_dict': model.state_dict(),\n 'best_EPE': 1e10}\n torch.save(state, os.path.join(args.flownet2_pytorch, 'FlowNet2_checkpoint.pth.tar'))\n\nelif 'FlowNet2-C/' in args.caffe_model:\n model = models.FlowNet2C(args)\n\n parse_flownetc(model.modules(), weights, biases)\n state = {'epoch': 0,\n 'state_dict': model.state_dict(),\n 'best_EPE': 1e10}\n torch.save(state, os.path.join(args.flownet2_pytorch, 'FlowNet2-C_checkpoint.pth.tar'))\n\nelif 'FlowNet2-CS/' in args.caffe_model:\n model = models.FlowNet2CS(args)\n\n parse_flownetc(model.flownetc.modules(), weights, biases)\n parse_flownets(model.flownets_1.modules(), weights, biases, param_prefix='net2_')\n\n state = {'epoch': 0,\n 'state_dict': model.state_dict(),\n 'best_EPE': 1e10}\n torch.save(state, os.path.join(args.flownet2_pytorch, 'FlowNet2-CS_checkpoint.pth.tar'))\n\nelif 'FlowNet2-CSS/' in args.caffe_model:\n model = models.FlowNet2CSS(args)\n\n parse_flownetc(model.flownetc.modules(), weights, biases)\n parse_flownets(model.flownets_1.modules(), weights, biases, param_prefix='net2_')\n parse_flownets(model.flownets_2.modules(), weights, biases, param_prefix='net3_')\n\n state = {'epoch': 0,\n 'state_dict': model.state_dict(),\n 'best_EPE': 1e10}\n torch.save(state, os.path.join(args.flownet2_pytorch, 'FlowNet2-CSS_checkpoint.pth.tar'))\n\nelif 'FlowNet2-CSS-ft-sd/' in args.caffe_model:\n model = models.FlowNet2CSS(args)\n\n parse_flownetc(model.flownetc.modules(), weights, biases)\n parse_flownets(model.flownets_1.modules(), weights, biases, param_prefix='net2_')\n parse_flownets(model.flownets_2.modules(), weights, biases, param_prefix='net3_')\n\n state = {'epoch': 0,\n 'state_dict': model.state_dict(),\n 'best_EPE': 1e10}\n torch.save(state, os.path.join(args.flownet2_pytorch, 'FlowNet2-CSS-ft-sd_checkpoint.pth.tar'))\n\nelif 'FlowNet2-S/' in args.caffe_model:\n model = models.FlowNet2S(args)\n\n parse_flownetsonly(model.modules(), weights, biases, param_prefix='')\n state = {'epoch': 0,\n 'state_dict': model.state_dict(),\n 'best_EPE': 1e10}\n torch.save(state, os.path.join(args.flownet2_pytorch, 'FlowNet2-S_checkpoint.pth.tar'))\n\nelif 'FlowNet2-SD/' in args.caffe_model:\n model = models.FlowNet2SD(args)\n\n parse_flownetsd(model.modules(), weights, biases, param_prefix='')\n\n state = {'epoch': 0,\n 'state_dict': model.state_dict(),\n 'best_EPE': 1e10}\n torch.save(state, os.path.join(args.flownet2_pytorch, 'FlowNet2-SD_checkpoint.pth.tar'))\n\nelse:\n print(('model type cound not be determined from input caffe model %s'%(args.caffe_model)))\n quit()\nprint((\"done converting \", args.caffe_model))"
},
{
"path": "dvs/flownet2/datasets.py",
"content": "import torch\nimport torch.utils.data as data\n\nimport os, math, random\nfrom os.path import *\nimport numpy as np\n\nfrom glob import glob\nimport utils.frame_utils as frame_utils\n\nfrom imageio import imread\n\nclass StaticRandomCrop(object):\n def __init__(self, image_size, crop_size):\n self.th, self.tw = crop_size\n h, w = image_size\n self.h1 = random.randint(0, h - self.th)\n self.w1 = random.randint(0, w - self.tw)\n\n def __call__(self, img):\n return img[self.h1:(self.h1+self.th), self.w1:(self.w1+self.tw),:]\n\nclass StaticCenterCrop(object):\n def __init__(self, image_size, crop_size):\n self.th, self.tw = crop_size\n self.h, self.w = image_size\n def __call__(self, img):\n return img[(self.h-self.th)//2:(self.h+self.th)//2, (self.w-self.tw)//2:(self.w+self.tw)//2,:]\n\nclass Padding(object):\n def __init__(self, image_size, pad_size):\n self.th, self.tw = pad_size\n self.h, self.w = image_size\n def __call__(self, img):\n out = np.zeros((self.th, self.tw, 3))\n out[:self.h, :self.w,:] = img\n return out\n\nclass MpiSintel(data.Dataset):\n def __init__(self, args, is_cropped = False, root = '', dstype = 'clean', replicates = 1):\n self.args = args\n self.is_cropped = is_cropped\n self.crop_size = args.crop_size\n self.render_size = args.inference_size\n self.replicates = replicates\n\n flow_root = join(root, 'flow')\n image_root = join(root, dstype)\n\n file_list = sorted(glob(join(flow_root, '*/*.flo')))\n\n self.flow_list = []\n self.image_list = []\n\n for file in file_list:\n if 'test' in file:\n # print file\n continue\n\n fbase = file[len(flow_root)+1:]\n fprefix = fbase[:-8]\n fnum = int(fbase[-8:-4])\n\n img1 = join(image_root, fprefix + \"%04d\"%(fnum+0) + '.png')\n img2 = join(image_root, fprefix + \"%04d\"%(fnum+1) + '.png')\n\n if not isfile(img1) or not isfile(img2) or not isfile(file):\n continue\n\n self.image_list += [[img1, img2]]\n self.flow_list += [file]\n\n self.size = len(self.image_list)\n\n self.frame_size = frame_utils.read_gen(self.image_list[0][0]).shape\n\n if (self.render_size[0] < 0) or (self.render_size[1] < 0) or (self.frame_size[0]%64) or (self.frame_size[1]%64):\n self.render_size[0] = ( (self.frame_size[0])//64 ) * 64\n self.render_size[1] = ( (self.frame_size[1])//64 ) * 64\n\n args.inference_size = self.render_size\n\n assert (len(self.image_list) == len(self.flow_list))\n\n def __getitem__(self, index):\n\n index = index % self.size\n\n img1 = frame_utils.read_gen(self.image_list[index][0])\n img2 = frame_utils.read_gen(self.image_list[index][1])\n\n flow = frame_utils.read_gen(self.flow_list[index])\n\n images = [img1, img2]\n image_size = img1.shape[:2]\n\n if self.is_cropped:\n cropper = StaticRandomCrop(image_size, self.crop_size)\n else:\n cropper = StaticCenterCrop(image_size, self.render_size)\n images = list(map(cropper, images))\n flow = cropper(flow)\n\n images = np.array(images).transpose(3,0,1,2)\n flow = flow.transpose(2,0,1)\n\n images = torch.from_numpy(images.astype(np.float32))\n flow = torch.from_numpy(flow.astype(np.float32))\n\n return [images], [flow]\n\n def __len__(self):\n return self.size * self.replicates\n\nclass MpiSintelClean(MpiSintel):\n def __init__(self, args, is_cropped = False, root = '', replicates = 1):\n super(MpiSintelClean, self).__init__(args, is_cropped = is_cropped, root = root, dstype = 'clean', replicates = replicates)\n\nclass MpiSintelFinal(MpiSintel):\n def __init__(self, args, is_cropped = False, root = '', replicates = 1):\n super(MpiSintelFinal, self).__init__(args, is_cropped = is_cropped, root = root, dstype = 'final', replicates = replicates)\n\nclass FlyingChairs(data.Dataset):\n def __init__(self, args, is_cropped, root = '/path/to/FlyingChairs_release/data', replicates = 1):\n self.args = args\n self.is_cropped = is_cropped\n self.crop_size = args.crop_size\n self.render_size = args.inference_size\n self.replicates = replicates\n\n images = sorted( glob( join(root, '*.ppm') ) )\n\n self.flow_list = sorted( glob( join(root, '*.flo') ) )\n\n assert (len(images)//2 == len(self.flow_list))\n\n self.image_list = []\n for i in range(len(self.flow_list)):\n im1 = images[2*i]\n im2 = images[2*i + 1]\n self.image_list += [ [ im1, im2 ] ]\n\n assert len(self.image_list) == len(self.flow_list)\n\n self.size = len(self.image_list)\n\n self.frame_size = frame_utils.read_gen(self.image_list[0][0]).shape\n\n if (self.render_size[0] < 0) or (self.render_size[1] < 0) or (self.frame_size[0]%64) or (self.frame_size[1]%64):\n self.render_size[0] = ( (self.frame_size[0])//64 ) * 64\n self.render_size[1] = ( (self.frame_size[1])//64 ) * 64\n\n args.inference_size = self.render_size\n\n def __getitem__(self, index):\n index = index % self.size\n\n img1 = frame_utils.read_gen(self.image_list[index][0])\n img2 = frame_utils.read_gen(self.image_list[index][1])\n\n flow = frame_utils.read_gen(self.flow_list[index])\n\n images = [img1, img2]\n image_size = img1.shape[:2]\n if self.is_cropped:\n cropper = StaticRandomCrop(image_size, self.crop_size)\n else:\n cropper = StaticCenterCrop(image_size, self.render_size)\n images = list(map(cropper, images))\n flow = cropper(flow)\n\n\n images = np.array(images).transpose(3,0,1,2)\n flow = flow.transpose(2,0,1)\n\n images = torch.from_numpy(images.astype(np.float32))\n flow = torch.from_numpy(flow.astype(np.float32))\n\n return [images], [flow]\n\n def __len__(self):\n return self.size * self.replicates\n\nclass FlyingThings(data.Dataset):\n def __init__(self, args, is_cropped, root = '/path/to/flyingthings3d', dstype = 'frames_cleanpass', replicates = 1):\n self.args = args\n self.is_cropped = is_cropped\n self.crop_size = args.crop_size\n self.render_size = args.inference_size\n self.replicates = replicates\n\n image_dirs = sorted(glob(join(root, dstype, 'TRAIN/*/*')))\n image_dirs = sorted([join(f, 'left') for f in image_dirs] + [join(f, 'right') for f in image_dirs])\n\n flow_dirs = sorted(glob(join(root, 'optical_flow_flo_format/TRAIN/*/*')))\n flow_dirs = sorted([join(f, 'into_future/left') for f in flow_dirs] + [join(f, 'into_future/right') for f in flow_dirs])\n\n assert (len(image_dirs) == len(flow_dirs))\n\n self.image_list = []\n self.flow_list = []\n\n for idir, fdir in zip(image_dirs, flow_dirs):\n images = sorted( glob(join(idir, '*.png')) )\n flows = sorted( glob(join(fdir, '*.flo')) )\n for i in range(len(flows)):\n self.image_list += [ [ images[i], images[i+1] ] ]\n self.flow_list += [flows[i]]\n\n assert len(self.image_list) == len(self.flow_list)\n\n self.size = len(self.image_list)\n\n self.frame_size = frame_utils.read_gen(self.image_list[0][0]).shape\n\n if (self.render_size[0] < 0) or (self.render_size[1] < 0) or (self.frame_size[0]%64) or (self.frame_size[1]%64):\n self.render_size[0] = ( (self.frame_size[0])//64 ) * 64\n self.render_size[1] = ( (self.frame_size[1])//64 ) * 64\n\n args.inference_size = self.render_size\n\n def __getitem__(self, index):\n index = index % self.size\n\n img1 = frame_utils.read_gen(self.image_list[index][0])\n img2 = frame_utils.read_gen(self.image_list[index][1])\n\n flow = frame_utils.read_gen(self.flow_list[index])\n\n images = [img1, img2]\n image_size = img1.shape[:2]\n if self.is_cropped:\n cropper = StaticRandomCrop(image_size, self.crop_size)\n else:\n cropper = StaticCenterCrop(image_size, self.render_size)\n images = list(map(cropper, images))\n flow = cropper(flow)\n\n\n images = np.array(images).transpose(3,0,1,2)\n flow = flow.transpose(2,0,1)\n\n images = torch.from_numpy(images.astype(np.float32))\n flow = torch.from_numpy(flow.astype(np.float32))\n\n return [images], [flow]\n\n def __len__(self):\n return self.size * self.replicates\n\nclass FlyingThingsClean(FlyingThings):\n def __init__(self, args, is_cropped = False, root = '', replicates = 1):\n super(FlyingThingsClean, self).__init__(args, is_cropped = is_cropped, root = root, dstype = 'frames_cleanpass', replicates = replicates)\n\nclass FlyingThingsFinal(FlyingThings):\n def __init__(self, args, is_cropped = False, root = '', replicates = 1):\n super(FlyingThingsFinal, self).__init__(args, is_cropped = is_cropped, root = root, dstype = 'frames_finalpass', replicates = replicates)\n\nclass ChairsSDHom(data.Dataset):\n def __init__(self, args, is_cropped, root = '/path/to/chairssdhom/data', dstype = 'train', replicates = 1):\n self.args = args\n self.is_cropped = is_cropped\n self.crop_size = args.crop_size\n self.render_size = args.inference_size\n self.replicates = replicates\n\n image1 = sorted( glob( join(root, dstype, 't0/*.png') ) )\n image2 = sorted( glob( join(root, dstype, 't1/*.png') ) )\n self.flow_list = sorted( glob( join(root, dstype, 'flow/*.flo') ) )\n\n assert (len(image1) == len(self.flow_list))\n\n self.image_list = []\n for i in range(len(self.flow_list)):\n im1 = image1[i]\n im2 = image2[i]\n self.image_list += [ [ im1, im2 ] ]\n\n assert len(self.image_list) == len(self.flow_list)\n\n self.size = len(self.image_list)\n\n self.frame_size = frame_utils.read_gen(self.image_list[0][0]).shape\n\n if (self.render_size[0] < 0) or (self.render_size[1] < 0) or (self.frame_size[0]%64) or (self.frame_size[1]%64):\n self.render_size[0] = ( (self.frame_size[0])//64 ) * 64\n self.render_size[1] = ( (self.frame_size[1])//64 ) * 64\n\n args.inference_size = self.render_size\n\n def __getitem__(self, index):\n index = index % self.size\n\n img1 = frame_utils.read_gen(self.image_list[index][0])\n img2 = frame_utils.read_gen(self.image_list[index][1])\n\n flow = frame_utils.read_gen(self.flow_list[index])\n flow = flow[::-1,:,:]\n\n images = [img1, img2]\n image_size = img1.shape[:2]\n if self.is_cropped:\n cropper = StaticRandomCrop(image_size, self.crop_size)\n else:\n cropper = StaticCenterCrop(image_size, self.render_size)\n images = list(map(cropper, images))\n flow = cropper(flow)\n\n\n images = np.array(images).transpose(3,0,1,2)\n flow = flow.transpose(2,0,1)\n\n images = torch.from_numpy(images.astype(np.float32))\n flow = torch.from_numpy(flow.astype(np.float32))\n\n return [images], [flow]\n\n def __len__(self):\n return self.size * self.replicates\n\nclass ChairsSDHomTrain(ChairsSDHom):\n def __init__(self, args, is_cropped = False, root = '', replicates = 1):\n super(ChairsSDHomTrain, self).__init__(args, is_cropped = is_cropped, root = root, dstype = 'train', replicates = replicates)\n\nclass ChairsSDHomTest(ChairsSDHom):\n def __init__(self, args, is_cropped = False, root = '', replicates = 1):\n super(ChairsSDHomTest, self).__init__(args, is_cropped = is_cropped, root = root, dstype = 'test', replicates = replicates)\n\nclass ImagesFromFolder(data.Dataset):\n def __init__(self, args, is_cropped, root = '/path/to/frames/only/folder', iext = 'png', replicates = 1):\n self.args = args\n self.is_cropped = is_cropped\n self.crop_size = args.crop_size\n self.render_size = args.inference_size\n self.replicates = replicates\n\n images = sorted( glob( join(root, '*.' + iext) ) )\n self.image_list = []\n for i in range(len(images)-1):\n im1 = images[i]\n im2 = images[i+1]\n self.image_list += [ [ im1, im2 ] ]\n\n self.size = len(self.image_list)\n\n self.frame_size = frame_utils.read_gen(self.image_list[0][0]).shape\n\n if (self.render_size[0] < 0) or (self.render_size[1] < 0) or (self.frame_size[0]%64) or (self.frame_size[1]%64):\n self.render_size[0] = ( (self.frame_size[0])//64 ) * 64\n self.render_size[1] = ( (self.frame_size[1])//64 ) * 64\n\n args.inference_size = self.render_size\n\n def __getitem__(self, index):\n index = index % self.size\n\n img1 = frame_utils.read_gen(self.image_list[index][0])\n img2 = frame_utils.read_gen(self.image_list[index][1])\n\n images = [img1, img2]\n image_size = img1.shape[:2]\n if self.is_cropped:\n cropper = StaticRandomCrop(image_size, self.crop_size)\n else:\n cropper = StaticCenterCrop(image_size, self.render_size)\n images = list(map(cropper, images))\n \n images = np.array(images).transpose(3,0,1,2)\n images = torch.from_numpy(images.astype(np.float32))\n\n return [images], [torch.zeros(images.size()[0:1] + (2,) + images.size()[-2:])]\n\n def __len__(self):\n return self.size * self.replicates\n\n\nclass Google(data.Dataset):\n def __init__(self, args, is_cropped = False, root = '', dstype = 'frames', replicates = 1):\n self.args = args\n self.is_cropped = is_cropped\n self.crop_size = args.crop_size\n self.render_size = args.inference_size\n self.replicates = replicates\n\n image_root = join(root, dstype)\n\n file_list = sorted(glob(join(image_root, '*.png')))\n\n self.image_list = []\n\n for i in range(len(file_list)-1):\n\n img1 = join(file_list[i])\n img2 = join(file_list[i+1])\n\n if not isfile(img1) or not isfile(img2):\n continue\n\n self.image_list += [[img1, img2]]\n\n self.size = len(self.image_list)\n\n self.frame_size = frame_utils.read_gen(self.image_list[0][0]).shape\n\n if (self.render_size[0] < 0) or (self.render_size[1] < 0) or (self.frame_size[0]%64) or (self.frame_size[1]%64):\n self.render_size[0] = ( math.ceil(self.frame_size[0]/64) ) * 64\n self.render_size[1] = ( math.ceil(self.frame_size[1]/64) ) * 64\n\n\n args.inference_size = self.render_size\n\n def __getitem__(self, index):\n\n index = index % self.size\n\n img1 = frame_utils.read_gen(self.image_list[index][0])\n img2 = frame_utils.read_gen(self.image_list[index][1])\n\n images = [img1, img2]\n image_size = img1.shape[:2]\n\n if self.is_cropped:\n cropper = StaticRandomCrop(image_size, self.crop_size)\n else:\n cropper = Padding(image_size, self.render_size)\n images = list(map(cropper, images))\n\n images = np.array(images).transpose(3,0,1,2)\n\n images = torch.from_numpy(images.astype(np.float32))\n\n return [images]\n\n def __len__(self):\n return self.size * self.replicates\n\n'''\nimport argparse\nimport sys, os\nimport importlib\nfrom scipy.misc import imsave\nimport numpy as np\n\nimport datasets\nreload(datasets)\n\nparser = argparse.ArgumentParser()\nargs = parser.parse_args()\nargs.inference_size = [1080, 1920]\nargs.crop_size = [384, 512]\nargs.effective_batch_size = 1\n\nindex = 500\nv_dataset = datasets.MpiSintelClean(args, True, root='../MPI-Sintel/flow/training')\na, b = v_dataset[index]\nim1 = a[0].numpy()[:,0,:,:].transpose(1,2,0)\nim2 = a[0].numpy()[:,1,:,:].transpose(1,2,0)\nimsave('./img1.png', im1)\nimsave('./img2.png', im2)\nflow_utils.writeFlow('./flow.flo', b[0].numpy().transpose(1,2,0))\n\n'''\n"
},
{
"path": "dvs/flownet2/install.sh",
"content": "#!/bin/bash\ncd ./networks/correlation_package\nrm -rf *_cuda.egg-info build dist __pycache__\npython3 setup.py install --user\n\ncd ../resample2d_package\nrm -rf *_cuda.egg-info build dist __pycache__\npython3 setup.py install --user\n\ncd ../channelnorm_package\nrm -rf *_cuda.egg-info build dist __pycache__\npython3 setup.py install --user\n\ncd ..\n"
},
{
"path": "dvs/flownet2/losses.py",
"content": "'''\r\nPortions of this code copyright 2017, Clement Pinard\r\n'''\r\n\r\n# freda (todo) : adversarial loss \r\n\r\nimport torch\r\nimport torch.nn as nn\r\nimport math\r\n\r\ndef EPE(input_flow, target_flow):\r\n return torch.norm(target_flow-input_flow,p=2,dim=1).mean()\r\n\r\nclass L1(nn.Module):\r\n def __init__(self):\r\n super(L1, self).__init__()\r\n def forward(self, output, target):\r\n lossvalue = torch.abs(output - target).mean()\r\n return lossvalue\r\n\r\nclass L2(nn.Module):\r\n def __init__(self):\r\n super(L2, self).__init__()\r\n def forward(self, output, target):\r\n lossvalue = torch.norm(output-target,p=2,dim=1).mean()\r\n return lossvalue\r\n\r\nclass L1Loss(nn.Module):\r\n def __init__(self, args):\r\n super(L1Loss, self).__init__()\r\n self.args = args\r\n self.loss = L1()\r\n self.loss_labels = ['L1', 'EPE']\r\n\r\n def forward(self, output, target):\r\n lossvalue = self.loss(output, target)\r\n epevalue = EPE(output, target)\r\n return [lossvalue, epevalue]\r\n\r\nclass L2Loss(nn.Module):\r\n def __init__(self, args):\r\n super(L2Loss, self).__init__()\r\n self.args = args\r\n self.loss = L2()\r\n self.loss_labels = ['L2', 'EPE']\r\n\r\n def forward(self, output, target):\r\n lossvalue = self.loss(output, target)\r\n epevalue = EPE(output, target)\r\n return [lossvalue, epevalue]\r\n\r\nclass MultiScale(nn.Module):\r\n def __init__(self, args, startScale = 4, numScales = 5, l_weight= 0.32, norm= 'L1'):\r\n super(MultiScale,self).__init__()\r\n\r\n self.startScale = startScale\r\n self.numScales = numScales\r\n self.loss_weights = torch.FloatTensor([(l_weight / 2 ** scale) for scale in range(self.numScales)])\r\n self.args = args\r\n self.l_type = norm\r\n self.div_flow = 0.05\r\n assert(len(self.loss_weights) == self.numScales)\r\n\r\n if self.l_type == 'L1':\r\n self.loss = L1()\r\n else:\r\n self.loss = L2()\r\n\r\n self.multiScales = [nn.AvgPool2d(self.startScale * (2**scale), self.startScale * (2**scale)) for scale in range(self.numScales)]\r\n self.loss_labels = ['MultiScale-'+self.l_type, 'EPE'],\r\n\r\n def forward(self, output, target):\r\n lossvalue = 0\r\n epevalue = 0\r\n\r\n if type(output) is tuple:\r\n target = self.div_flow * target\r\n for i, output_ in enumerate(output):\r\n target_ = self.multiScales[i](target)\r\n epevalue += self.loss_weights[i]*EPE(output_, target_)\r\n lossvalue += self.loss_weights[i]*self.loss(output_, target_)\r\n return [lossvalue, epevalue]\r\n else:\r\n epevalue += EPE(output, target)\r\n lossvalue += self.loss(output, target)\r\n return [lossvalue, epevalue]\r\n\r\n"
},
{
"path": "dvs/flownet2/main.py",
"content": "#!/usr/bin/env python\nimport os\nos.environ[\"CUDA_VISIBLE_DEVICES\"] = \"0\"\nimport torch\nimport torch.nn as nn\nfrom torch.utils.data import DataLoader\nfrom torch.autograd import Variable\nfrom tensorboardX import SummaryWriter\n\nimport argparse, os, sys, subprocess\nimport colorama\nimport numpy as np\nfrom tqdm import tqdm\nfrom glob import glob\nfrom os.path import *\n\nimport models, datasets\nfrom utils import flow_utils, tools\nimport time\n\n # Reusable function for inference\ndef inference(args, epoch, data_path, data_loader, model, offset=0):\n\n model.eval()\n \n if args.save_flow or args.render_validation:\n flow_folder = \"{}/flo\".format(data_path)\n flow_back_folder = \"{}/flo_back\".format(data_path)\n if not os.path.exists(flow_folder):\n os.makedirs(flow_folder)\n if not os.path.exists(flow_back_folder):\n os.makedirs(flow_back_folder)\n \n # visualization folder\n if args.inference_visualize:\n flow_vis_folder = \"{}/flo_vis\".format(data_path)\n if not os.path.exists(flow_vis_folder):\n os.makedirs(flow_vis_folder)\n flow_back_vis_folder = \"{}/flo_back_vis\".format(data_path)\n if not os.path.exists(flow_back_vis_folder):\n os.makedirs(flow_back_vis_folder)\n \n args.inference_n_batches = np.inf if args.inference_n_batches < 0 else args.inference_n_batches\n\n progress = tqdm(data_loader, ncols=100, total=np.minimum(len(data_loader), args.inference_n_batches), desc='Inferencing ', \n leave=True, position=offset)\n\n for batch_idx, (data) in enumerate(progress):\n data = data[0]\n data_back = torch.cat((data[:,:,1:,:,:], data[:,:,:1,:,:]), dim = 2)\n if args.cuda:\n data_forward = data.cuda(non_blocking=True)\n data_back = data_back.cuda(non_blocking=True)\n data_forward = Variable(data_forward)\n data_back = Variable(data_back)\n\n flo_path = join(flow_folder, '%06d.flo'%(batch_idx))\n flo_back_path = join(flow_back_folder, '%06d.flo'%(batch_idx))\n frame_size = data_loader.dataset.frame_size\n if not os.path.exists(flo_path):\n with torch.no_grad():\n output = model(data_forward)[:,:,:frame_size[0], :frame_size[1]]\n if args.save_flow or args.render_validation:\n _pflow = output[0].data.cpu().numpy().transpose(1, 2, 0)\n flow_utils.writeFlow( flo_path, _pflow)\n if args.inference_visualize:\n flow_utils.visulize_flow_file(\n join(flow_folder, '%06d.flo' % (batch_idx)),flow_vis_folder)\n\n if not os.path.exists(flo_back_path):\n with torch.no_grad():\n output = model(data_back)[:,:,:frame_size[0], :frame_size[1]]\n if args.save_flow or args.render_validation:\n _pflow = output[0].data.cpu().numpy().transpose(1, 2, 0)\n flow_utils.writeFlow( flo_back_path, _pflow)\n if args.inference_visualize:\n flow_utils.visulize_flow_file(\n join(flow_back_folder, '%06d.flo' % (batch_idx)), flow_back_vis_folder)\n \n progress.update(1)\n\n if batch_idx == (args.inference_n_batches - 1):\n break\n progress.close()\n return\n\nif __name__ == '__main__':\n parser = argparse.ArgumentParser()\n parser.add_argument('--fp16', action='store_true', help='Run model in pseudo-fp16 mode (fp16 storage fp32 math).')\n parser.add_argument('--fp16_scale', type=float, default=1024., help='Loss scaling, positive power of 2 values can improve fp16 convergence.')\n\n parser.add_argument('--start_epoch', type=int, default=1)\n parser.add_argument('--batch_size', '-b', type=int, default=8, help=\"Batch size\")\n parser.add_argument('--crop_size', type=int, nargs='+', default = [256, 256], help=\"Spatial dimension to crop training samples for training\")\n parser.add_argument(\"--rgb_max\", type=float, default = 255.)\n\n parser.add_argument('--number_workers', '-nw', '--num_workers', type=int, default=8)\n parser.add_argument('--number_gpus', '-ng', type=int, default=-1, help='number of GPUs to use')\n parser.add_argument('--no_cuda', action='store_true')\n\n parser.add_argument('--save', '-s', default='./Google', type=str, help='directory for saving')\n\n parser.add_argument('--inference', action='store_true')\n parser.add_argument('--inference_visualize', action='store_true',\n help=\"visualize the optical flow during inference\")\n parser.add_argument('--inference_size', type=int, nargs='+', default = [-1,-1], help='spatial size divisible by 64. default (-1,-1) - largest possible valid size would be used')\n parser.add_argument('--inference_batch_size', type=int, default=1)\n parser.add_argument('--inference_n_batches', type=int, default=-1)\n parser.add_argument('--save_flow', action='store_true', help='save predicted flows to file')\n\n parser.add_argument('--resume', default='', type=str, metavar='PATH', help='path to latest checkpoint (default: none)')\n parser.add_argument('--log_frequency', '--summ_iter', type=int, default=1, help=\"Log every n batches\")\n\n tools.add_arguments_for_module(parser, models, argument_for_class='model', default='FlowNet2')\n \n tools.add_arguments_for_module(parser, datasets, argument_for_class='inference_dataset', default='Google', \n skip_params=['is_cropped'],\n parameter_defaults={'root': './Google/train',\n 'replicates': 1})\n\n main_dir = os.path.dirname(os.path.realpath(__file__))\n os.chdir(main_dir)\n\n # Parse the official arguments\n with tools.TimerBlock(\"Parsing Arguments\") as block:\n args = parser.parse_args()\n if args.number_gpus < 0 : args.number_gpus = torch.cuda.device_count()\n\n # Get argument defaults (hastag #thisisahack)\n parser.add_argument('--IGNORE', action='store_true')\n defaults = vars(parser.parse_args(['--IGNORE']))\n\n # Print all arguments, color the non-defaults\n for argument, value in sorted(vars(args).items()):\n reset = colorama.Style.RESET_ALL\n color = reset if value == defaults[argument] else colorama.Fore.MAGENTA\n block.log('{}{}: {}{}'.format(color, argument, value, reset))\n\n args.model_class = tools.module_to_dict(models)[args.model]\n\n args.inference_dataset_class = tools.module_to_dict(datasets)[args.inference_dataset]\n\n args.cuda = not args.no_cuda and torch.cuda.is_available()\n # args.current_hash = subprocess.check_output([\"git\", \"rev-parse\", \"HEAD\"]).rstrip()\n args.log_file = join(args.save, 'args.txt')\n\n # dict to collect activation gradients (for training debug purpose)\n args.grads = {}\n\n args.total_epochs = 1\n args.inference_dir = \"{}/inference\".format(args.save)\n\n print('Source Code')\n # print((' Current Git Hash: {}\\n'.format(args.current_hash)))\n\n # Dynamically load the dataset class with parameters passed in via \"--argument_[param]=[value]\" arguments\n with tools.TimerBlock(\"Initializing Datasets\") as block:\n args.effective_batch_size = args.batch_size * args.number_gpus\n args.effective_inference_batch_size = args.inference_batch_size * args.number_gpus\n args.effective_number_workers = args.number_workers * args.number_gpus\n gpuargs = {'num_workers': args.effective_number_workers, \n 'pin_memory': True, \n 'drop_last' : True} if args.cuda else {}\n inf_gpuargs = gpuargs.copy()\n inf_gpuargs['num_workers'] = args.number_workers\n\n block.log('Inference Dataset: {}'.format(args.inference_dataset))\n\n dataset_root = args.inference_dataset_root \n data_name = sorted(os.listdir(dataset_root))\n\n block.log(data_name)\n inference_loaders = {}\n for i in range(len(data_name)):\n dataset_path = os.path.join(dataset_root, data_name[i])\n args.inference_dataset_root = dataset_path\n inference_dataset = args.inference_dataset_class(args, False, **tools.kwargs_from_args(args, 'inference_dataset'))\n inference_loaders[dataset_path] = DataLoader(inference_dataset, batch_size=args.effective_inference_batch_size, shuffle=False, **inf_gpuargs)\n block.log('Inference Input: {}'.format(' '.join([str([d for d in x.size()]) for x in inference_dataset[0][0]])))\n\n # Dynamically load model and loss class with parameters passed in via \"--model_[param]=[value]\" or \"--loss_[param]=[value]\" arguments\n with tools.TimerBlock(\"Building {} model\".format(args.model)) as block:\n class Model(nn.Module):\n def __init__(self, args):\n super(Model, self).__init__()\n kwargs = tools.kwargs_from_args(args, 'model')\n self.model = args.model_class(args, **kwargs)\n \n def forward(self, data):\n output = self.model(data)\n return output\n\n model = Model(args)\n\n block.log('Effective Batch Size: {}'.format(args.effective_batch_size))\n block.log('Number of parameters: {}'.format(sum([p.data.nelement() if p.requires_grad else 0 for p in model.parameters()])))\n\n if args.cuda and args.number_gpus > 0:\n block.log('Initializing CUDA')\n model = model.cuda()\n block.log('Parallelizing')\n model = nn.parallel.DataParallel(model, device_ids=list(range(args.number_gpus)))\n\n # Load weights if needed, otherwise randomly initialize\n if args.resume and os.path.isfile(args.resume):\n block.log(\"Loading checkpoint '{}'\".format(args.resume))\n checkpoint = torch.load(args.resume)\n model.module.model.load_state_dict(checkpoint['state_dict'])\n block.log(\"Loaded checkpoint '{}' (at epoch {})\".format(args.resume, checkpoint['epoch']))\n\n elif args.resume and args.inference:\n block.log(\"No checkpoint found at '{}'\".format(args.resume))\n quit()\n\n else:\n block.log(\"Random initialization\")\n\n block.log(\"Initializing save directory: {}\".format(args.save))\n if not os.path.exists(args.save):\n os.makedirs(args.save)\n\n # Log all arguments to file\n for argument, value in sorted(vars(args).items()):\n block.log2file(args.log_file, '{}: {}'.format(argument, value))\n\n for data_path in inference_loaders:\n # Primary epoch loop\n progress = tqdm(list(range(args.start_epoch, args.total_epochs + 1)), miniters=1, ncols=100, desc='Overall Progress', leave=True, position=0)\n offset = 1\n\n for epoch in progress:\n stats = inference(args=args, epoch=epoch - 1, data_path = data_path, data_loader=inference_loaders[data_path], model=model, offset=offset)\n offset += 1\n print(\"\\n\")"
},
{
"path": "dvs/flownet2/models.py",
"content": "import torch\r\nimport torch.nn as nn\r\nfrom torch.nn import init\r\n\r\nimport math\r\nimport numpy as np\r\n\r\ntry:\r\n from networks.resample2d_package.resample2d import Resample2d\r\n from networks.channelnorm_package.channelnorm import ChannelNorm\r\n\r\n from networks import FlowNetC\r\n from networks import FlowNetS\r\n from networks import FlowNetSD\r\n from networks import FlowNetFusion\r\n\r\n from networks.submodules import *\r\nexcept:\r\n from .networks.resample2d_package.resample2d import Resample2d\r\n from .networks.channelnorm_package.channelnorm import ChannelNorm\r\n\r\n from .networks import FlowNetC\r\n from .networks import FlowNetS\r\n from .networks import FlowNetSD\r\n from .networks import FlowNetFusion\r\n\r\n from .networks.submodules import *\r\n'Parameter count = 162,518,834'\r\n\r\nclass FlowNet2(nn.Module):\r\n\r\n def __init__(self, args, batchNorm=False, div_flow = 20.):\r\n super(FlowNet2,self).__init__()\r\n self.batchNorm = batchNorm\r\n self.div_flow = div_flow\r\n self.rgb_max = args.rgb_max\r\n self.args = args\r\n\r\n self.channelnorm = ChannelNorm()\r\n\r\n # First Block (FlowNetC)\r\n self.flownetc = FlowNetC.FlowNetC(args, batchNorm=self.batchNorm)\r\n self.upsample1 = nn.Upsample(scale_factor=4, mode='bilinear')\r\n\r\n if args.fp16:\r\n self.resample1 = nn.Sequential(\r\n tofp32(), \r\n Resample2d(),\r\n tofp16()) \r\n else:\r\n self.resample1 = Resample2d()\r\n\r\n # Block (FlowNetS1)\r\n self.flownets_1 = FlowNetS.FlowNetS(args, batchNorm=self.batchNorm)\r\n self.upsample2 = nn.Upsample(scale_factor=4, mode='bilinear')\r\n if args.fp16:\r\n self.resample2 = nn.Sequential(\r\n tofp32(), \r\n Resample2d(),\r\n tofp16()) \r\n else:\r\n self.resample2 = Resample2d()\r\n\r\n\r\n # Block (FlowNetS2)\r\n self.flownets_2 = FlowNetS.FlowNetS(args, batchNorm=self.batchNorm)\r\n\r\n # Block (FlowNetSD)\r\n self.flownets_d = FlowNetSD.FlowNetSD(args, batchNorm=self.batchNorm) \r\n self.upsample3 = nn.Upsample(scale_factor=4, mode='nearest') \r\n self.upsample4 = nn.Upsample(scale_factor=4, mode='nearest') \r\n\r\n if args.fp16:\r\n self.resample3 = nn.Sequential(\r\n tofp32(), \r\n Resample2d(),\r\n tofp16()) \r\n else:\r\n self.resample3 = Resample2d()\r\n\r\n if args.fp16:\r\n self.resample4 = nn.Sequential(\r\n tofp32(), \r\n Resample2d(),\r\n tofp16()) \r\n else:\r\n self.resample4 = Resample2d()\r\n\r\n # Block (FLowNetFusion)\r\n self.flownetfusion = FlowNetFusion.FlowNetFusion(args, batchNorm=self.batchNorm)\r\n\r\n for m in self.modules():\r\n if isinstance(m, nn.Conv2d):\r\n if m.bias is not None:\r\n init.uniform_(m.bias)\r\n init.xavier_uniform_(m.weight)\r\n\r\n if isinstance(m, nn.ConvTranspose2d):\r\n if m.bias is not None:\r\n init.uniform_(m.bias)\r\n init.xavier_uniform_(m.weight)\r\n # init_deconv_bilinear(m.weight)\r\n\r\n def init_deconv_bilinear(self, weight):\r\n f_shape = weight.size()\r\n heigh, width = f_shape[-2], f_shape[-1]\r\n f = np.ceil(width/2.0)\r\n c = (2 * f - 1 - f % 2) / (2.0 * f)\r\n bilinear = np.zeros([heigh, width])\r\n for x in range(width):\r\n for y in range(heigh):\r\n value = (1 - abs(x / f - c)) * (1 - abs(y / f - c))\r\n bilinear[x, y] = value\r\n min_dim = min(f_shape[0], f_shape[1])\r\n weight.data.fill_(0.)\r\n for i in range(min_dim):\r\n weight.data[i,i,:,:] = torch.from_numpy(bilinear)\r\n return \r\n\r\n def forward(self, inputs):\r\n rgb_mean = inputs.contiguous().view(inputs.size()[:2]+(-1,)).mean(dim=-1).view(inputs.size()[:2] + (1,1,1,))\r\n \r\n x = (inputs - rgb_mean) / self.rgb_max\r\n x1 = x[:,:,0,:,:]\r\n x2 = x[:,:,1,:,:]\r\n x = torch.cat((x1,x2), dim = 1)\r\n\r\n # flownetc\r\n flownetc_flow2 = self.flownetc(x)[0]\r\n flownetc_flow = self.upsample1(flownetc_flow2*self.div_flow)\r\n \r\n # warp img1 to img0; magnitude of diff between img0 and and warped_img1, \r\n resampled_img1 = self.resample1(x[:,3:,:,:], flownetc_flow)\r\n diff_img0 = x[:,:3,:,:] - resampled_img1 \r\n norm_diff_img0 = self.channelnorm(diff_img0)\r\n\r\n # concat img0, img1, img1->img0, flow, diff-mag ; \r\n concat1 = torch.cat((x, resampled_img1, flownetc_flow/self.div_flow, norm_diff_img0), dim=1)\r\n \r\n # flownets1\r\n flownets1_flow2 = self.flownets_1(concat1)[0]\r\n flownets1_flow = self.upsample2(flownets1_flow2*self.div_flow) \r\n\r\n # warp img1 to img0 using flownets1; magnitude of diff between img0 and and warped_img1\r\n resampled_img1 = self.resample2(x[:,3:,:,:], flownets1_flow)\r\n diff_img0 = x[:,:3,:,:] - resampled_img1\r\n norm_diff_img0 = self.channelnorm(diff_img0)\r\n\r\n # concat img0, img1, img1->img0, flow, diff-mag\r\n concat2 = torch.cat((x, resampled_img1, flownets1_flow/self.div_flow, norm_diff_img0), dim=1)\r\n\r\n # flownets2\r\n flownets2_flow2 = self.flownets_2(concat2)[0]\r\n flownets2_flow = self.upsample4(flownets2_flow2 * self.div_flow)\r\n norm_flownets2_flow = self.channelnorm(flownets2_flow)\r\n\r\n diff_flownets2_flow = self.resample4(x[:,3:,:,:], flownets2_flow)\r\n # if not diff_flownets2_flow.volatile:\r\n # diff_flownets2_flow.register_hook(save_grad(self.args.grads, 'diff_flownets2_flow'))\r\n\r\n diff_flownets2_img1 = self.channelnorm((x[:,:3,:,:]-diff_flownets2_flow))\r\n # if not diff_flownets2_img1.volatile:\r\n # diff_flownets2_img1.register_hook(save_grad(self.args.grads, 'diff_flownets2_img1'))\r\n\r\n # flownetsd\r\n flownetsd_flow2 = self.flownets_d(x)[0]\r\n flownetsd_flow = self.upsample3(flownetsd_flow2 / self.div_flow)\r\n norm_flownetsd_flow = self.channelnorm(flownetsd_flow)\r\n \r\n diff_flownetsd_flow = self.resample3(x[:,3:,:,:], flownetsd_flow)\r\n # if not diff_flownetsd_flow.volatile:\r\n # diff_flownetsd_flow.register_hook(save_grad(self.args.grads, 'diff_flownetsd_flow'))\r\n\r\n diff_flownetsd_img1 = self.channelnorm((x[:,:3,:,:]-diff_flownetsd_flow))\r\n # if not diff_flownetsd_img1.volatile:\r\n # diff_flownetsd_img1.register_hook(save_grad(self.args.grads, 'diff_flownetsd_img1'))\r\n\r\n # concat img1 flownetsd, flownets2, norm_flownetsd, norm_flownets2, diff_flownetsd_img1, diff_flownets2_img1\r\n concat3 = torch.cat((x[:,:3,:,:], flownetsd_flow, flownets2_flow, norm_flownetsd_flow, norm_flownets2_flow, diff_flownetsd_img1, diff_flownets2_img1), dim=1)\r\n flownetfusion_flow = self.flownetfusion(concat3)\r\n\r\n # if not flownetfusion_flow.volatile:\r\n # flownetfusion_flow.register_hook(save_grad(self.args.grads, 'flownetfusion_flow'))\r\n\r\n return flownetfusion_flow\r\n\r\nclass FlowNet2C(FlowNetC.FlowNetC):\r\n def __init__(self, args, batchNorm=False, div_flow=20):\r\n super(FlowNet2C,self).__init__(args, batchNorm=batchNorm, div_flow=20)\r\n self.rgb_max = args.rgb_max\r\n\r\n def forward(self, inputs):\r\n rgb_mean = inputs.contiguous().view(inputs.size()[:2]+(-1,)).mean(dim=-1).view(inputs.size()[:2] + (1,1,1,))\r\n \r\n x = (inputs - rgb_mean) / self.rgb_max\r\n x1 = x[:,:,0,:,:]\r\n x2 = x[:,:,1,:,:]\r\n\r\n # FlownetC top input stream\r\n out_conv1a = self.conv1(x1)\r\n out_conv2a = self.conv2(out_conv1a)\r\n out_conv3a = self.conv3(out_conv2a)\r\n\r\n # FlownetC bottom input stream\r\n out_conv1b = self.conv1(x2)\r\n \r\n out_conv2b = self.conv2(out_conv1b)\r\n out_conv3b = self.conv3(out_conv2b)\r\n\r\n # Merge streams\r\n out_corr = self.corr(out_conv3a, out_conv3b) # False\r\n out_corr = self.corr_activation(out_corr)\r\n\r\n # Redirect top input stream and concatenate\r\n out_conv_redir = self.conv_redir(out_conv3a)\r\n\r\n in_conv3_1 = torch.cat((out_conv_redir, out_corr), 1)\r\n\r\n # Merged conv layers\r\n out_conv3_1 = self.conv3_1(in_conv3_1)\r\n\r\n out_conv4 = self.conv4_1(self.conv4(out_conv3_1))\r\n\r\n out_conv5 = self.conv5_1(self.conv5(out_conv4))\r\n out_conv6 = self.conv6_1(self.conv6(out_conv5))\r\n\r\n flow6 = self.predict_flow6(out_conv6)\r\n flow6_up = self.upsampled_flow6_to_5(flow6)\r\n out_deconv5 = self.deconv5(out_conv6)\r\n\r\n concat5 = torch.cat((out_conv5,out_deconv5,flow6_up),1)\r\n\r\n flow5 = self.predict_flow5(concat5)\r\n flow5_up = self.upsampled_flow5_to_4(flow5)\r\n out_deconv4 = self.deconv4(concat5)\r\n concat4 = torch.cat((out_conv4,out_deconv4,flow5_up),1)\r\n\r\n flow4 = self.predict_flow4(concat4)\r\n flow4_up = self.upsampled_flow4_to_3(flow4)\r\n out_deconv3 = self.deconv3(concat4)\r\n concat3 = torch.cat((out_conv3_1,out_deconv3,flow4_up),1)\r\n\r\n flow3 = self.predict_flow3(concat3)\r\n flow3_up = self.upsampled_flow3_to_2(flow3)\r\n out_deconv2 = self.deconv2(concat3)\r\n concat2 = torch.cat((out_conv2a,out_deconv2,flow3_up),1)\r\n\r\n flow2 = self.predict_flow2(concat2)\r\n\r\n if self.training:\r\n return flow2,flow3,flow4,flow5,flow6\r\n else:\r\n return self.upsample1(flow2*self.div_flow)\r\n\r\nclass FlowNet2S(FlowNetS.FlowNetS):\r\n def __init__(self, args, batchNorm=False, div_flow=20):\r\n super(FlowNet2S,self).__init__(args, input_channels = 6, batchNorm=batchNorm)\r\n self.rgb_max = args.rgb_max\r\n self.div_flow = div_flow\r\n \r\n def forward(self, inputs):\r\n rgb_mean = inputs.contiguous().view(inputs.size()[:2]+(-1,)).mean(dim=-1).view(inputs.size()[:2] + (1,1,1,))\r\n x = (inputs - rgb_mean) / self.rgb_max\r\n x = torch.cat( (x[:,:,0,:,:], x[:,:,1,:,:]), dim = 1)\r\n\r\n out_conv1 = self.conv1(x)\r\n\r\n out_conv2 = self.conv2(out_conv1)\r\n out_conv3 = self.conv3_1(self.conv3(out_conv2))\r\n out_conv4 = self.conv4_1(self.conv4(out_conv3))\r\n out_conv5 = self.conv5_1(self.conv5(out_conv4))\r\n out_conv6 = self.conv6_1(self.conv6(out_conv5))\r\n\r\n flow6 = self.predict_flow6(out_conv6)\r\n flow6_up = self.upsampled_flow6_to_5(flow6)\r\n out_deconv5 = self.deconv5(out_conv6)\r\n \r\n concat5 = torch.cat((out_conv5,out_deconv5,flow6_up),1)\r\n flow5 = self.predict_flow5(concat5)\r\n flow5_up = self.upsampled_flow5_to_4(flow5)\r\n out_deconv4 = self.deconv4(concat5)\r\n \r\n concat4 = torch.cat((out_conv4,out_deconv4,flow5_up),1)\r\n flow4 = self.predict_flow4(concat4)\r\n flow4_up = self.upsampled_flow4_to_3(flow4)\r\n out_deconv3 = self.deconv3(concat4)\r\n \r\n concat3 = torch.cat((out_conv3,out_deconv3,flow4_up),1)\r\n flow3 = self.predict_flow3(concat3)\r\n flow3_up = self.upsampled_flow3_to_2(flow3)\r\n out_deconv2 = self.deconv2(concat3)\r\n\r\n concat2 = torch.cat((out_conv2,out_deconv2,flow3_up),1)\r\n flow2 = self.predict_flow2(concat2)\r\n\r\n if self.training:\r\n return flow2,flow3,flow4,flow5,flow6\r\n else:\r\n return self.upsample1(flow2*self.div_flow)\r\n\r\nclass FlowNet2SD(FlowNetSD.FlowNetSD):\r\n def __init__(self, args, batchNorm=False, div_flow=20):\r\n super(FlowNet2SD,self).__init__(args, batchNorm=batchNorm)\r\n self.rgb_max = args.rgb_max\r\n self.div_flow = div_flow\r\n\r\n def forward(self, inputs):\r\n rgb_mean = inputs.contiguous().view(inputs.size()[:2]+(-1,)).mean(dim=-1).view(inputs.size()[:2] + (1,1,1,))\r\n x = (inputs - rgb_mean) / self.rgb_max\r\n x = torch.cat( (x[:,:,0,:,:], x[:,:,1,:,:]), dim = 1)\r\n\r\n out_conv0 = self.conv0(x)\r\n out_conv1 = self.conv1_1(self.conv1(out_conv0))\r\n out_conv2 = self.conv2_1(self.conv2(out_conv1))\r\n\r\n out_conv3 = self.conv3_1(self.conv3(out_conv2))\r\n out_conv4 = self.conv4_1(self.conv4(out_conv3))\r\n out_conv5 = self.conv5_1(self.conv5(out_conv4))\r\n out_conv6 = self.conv6_1(self.conv6(out_conv5))\r\n\r\n flow6 = self.predict_flow6(out_conv6)\r\n flow6_up = self.upsampled_flow6_to_5(flow6)\r\n out_deconv5 = self.deconv5(out_conv6)\r\n \r\n concat5 = torch.cat((out_conv5,out_deconv5,flow6_up),1)\r\n out_interconv5 = self.inter_conv5(concat5)\r\n flow5 = self.predict_flow5(out_interconv5)\r\n\r\n flow5_up = self.upsampled_flow5_to_4(flow5)\r\n out_deconv4 = self.deconv4(concat5)\r\n \r\n concat4 = torch.cat((out_conv4,out_deconv4,flow5_up),1)\r\n out_interconv4 = self.inter_conv4(concat4)\r\n flow4 = self.predict_flow4(out_interconv4)\r\n flow4_up = self.upsampled_flow4_to_3(flow4)\r\n out_deconv3 = self.deconv3(concat4)\r\n \r\n concat3 = torch.cat((out_conv3,out_deconv3,flow4_up),1)\r\n out_interconv3 = self.inter_conv3(concat3)\r\n flow3 = self.predict_flow3(out_interconv3)\r\n flow3_up = self.upsampled_flow3_to_2(flow3)\r\n out_deconv2 = self.deconv2(concat3)\r\n\r\n concat2 = torch.cat((out_conv2,out_deconv2,flow3_up),1)\r\n out_interconv2 = self.inter_conv2(concat2)\r\n flow2 = self.predict_flow2(out_interconv2)\r\n\r\n if self.training:\r\n return flow2,flow3,flow4,flow5,flow6\r\n else:\r\n return self.upsample1(flow2*self.div_flow)\r\n\r\nclass FlowNet2CS(nn.Module):\r\n\r\n def __init__(self, args, batchNorm=False, div_flow = 20.):\r\n super(FlowNet2CS,self).__init__()\r\n self.batchNorm = batchNorm\r\n self.div_flow = div_flow\r\n self.rgb_max = args.rgb_max\r\n self.args = args\r\n\r\n self.channelnorm = ChannelNorm()\r\n\r\n # First Block (FlowNetC)\r\n self.flownetc = FlowNetC.FlowNetC(args, batchNorm=self.batchNorm)\r\n self.upsample1 = nn.Upsample(scale_factor=4, mode='bilinear')\r\n\r\n if args.fp16:\r\n self.resample1 = nn.Sequential(\r\n tofp32(), \r\n Resample2d(),\r\n tofp16()) \r\n else:\r\n self.resample1 = Resample2d()\r\n\r\n # Block (FlowNetS1)\r\n self.flownets_1 = FlowNetS.FlowNetS(args, batchNorm=self.batchNorm)\r\n self.upsample2 = nn.Upsample(scale_factor=4, mode='bilinear')\r\n\r\n for m in self.modules():\r\n if isinstance(m, nn.Conv2d):\r\n if m.bias is not None:\r\n init.uniform(m.bias)\r\n init.xavier_uniform(m.weight)\r\n\r\n if isinstance(m, nn.ConvTranspose2d):\r\n if m.bias is not None:\r\n init.uniform(m.bias)\r\n init.xavier_uniform(m.weight)\r\n # init_deconv_bilinear(m.weight)\r\n\r\n def forward(self, inputs):\r\n rgb_mean = inputs.contiguous().view(inputs.size()[:2]+(-1,)).mean(dim=-1).view(inputs.size()[:2] + (1,1,1,))\r\n \r\n x = (inputs - rgb_mean) / self.rgb_max\r\n x1 = x[:,:,0,:,:]\r\n x2 = x[:,:,1,:,:]\r\n x = torch.cat((x1,x2), dim = 1)\r\n\r\n # flownetc\r\n flownetc_flow2 = self.flownetc(x)[0]\r\n flownetc_flow = self.upsample1(flownetc_flow2*self.div_flow)\r\n \r\n # warp img1 to img0; magnitude of diff between img0 and and warped_img1, \r\n resampled_img1 = self.resample1(x[:,3:,:,:], flownetc_flow)\r\n diff_img0 = x[:,:3,:,:] - resampled_img1 \r\n norm_diff_img0 = self.channelnorm(diff_img0)\r\n\r\n # concat img0, img1, img1->img0, flow, diff-mag ; \r\n concat1 = torch.cat((x, resampled_img1, flownetc_flow/self.div_flow, norm_diff_img0), dim=1)\r\n \r\n # flownets1\r\n flownets1_flow2 = self.flownets_1(concat1)[0]\r\n flownets1_flow = self.upsample2(flownets1_flow2*self.div_flow) \r\n\r\n return flownets1_flow\r\n\r\nclass FlowNet2CSS(nn.Module):\r\n\r\n def __init__(self, args, batchNorm=False, div_flow = 20.):\r\n super(FlowNet2CSS,self).__init__()\r\n self.batchNorm = batchNorm\r\n self.div_flow = div_flow\r\n self.rgb_max = args.rgb_max\r\n self.args = args\r\n\r\n self.channelnorm = ChannelNorm()\r\n\r\n # First Block (FlowNetC)\r\n self.flownetc = FlowNetC.FlowNetC(args, batchNorm=self.batchNorm)\r\n self.upsample1 = nn.Upsample(scale_factor=4, mode='bilinear')\r\n\r\n if args.fp16:\r\n self.resample1 = nn.Sequential(\r\n tofp32(), \r\n Resample2d(),\r\n tofp16()) \r\n else:\r\n self.resample1 = Resample2d()\r\n\r\n # Block (FlowNetS1)\r\n self.flownets_1 = FlowNetS.FlowNetS(args, batchNorm=self.batchNorm)\r\n self.upsample2 = nn.Upsample(scale_factor=4, mode='bilinear')\r\n if args.fp16:\r\n self.resample2 = nn.Sequential(\r\n tofp32(), \r\n Resample2d(),\r\n tofp16()) \r\n else:\r\n self.resample2 = Resample2d()\r\n\r\n\r\n # Block (FlowNetS2)\r\n self.flownets_2 = FlowNetS.FlowNetS(args, batchNorm=self.batchNorm)\r\n self.upsample3 = nn.Upsample(scale_factor=4, mode='nearest') \r\n\r\n for m in self.modules():\r\n if isinstance(m, nn.Conv2d):\r\n if m.bias is not None:\r\n init.uniform(m.bias)\r\n init.xavier_uniform(m.weight)\r\n\r\n if isinstance(m, nn.ConvTranspose2d):\r\n if m.bias is not None:\r\n init.uniform(m.bias)\r\n init.xavier_uniform(m.weight)\r\n # init_deconv_bilinear(m.weight)\r\n\r\n def forward(self, inputs):\r\n rgb_mean = inputs.contiguous().view(inputs.size()[:2]+(-1,)).mean(dim=-1).view(inputs.size()[:2] + (1,1,1,))\r\n \r\n x = (inputs - rgb_mean) / self.rgb_max\r\n x1 = x[:,:,0,:,:]\r\n x2 = x[:,:,1,:,:]\r\n x = torch.cat((x1,x2), dim = 1)\r\n\r\n # flownetc\r\n flownetc_flow2 = self.flownetc(x)[0]\r\n flownetc_flow = self.upsample1(flownetc_flow2*self.div_flow)\r\n \r\n # warp img1 to img0; magnitude of diff between img0 and and warped_img1, \r\n resampled_img1 = self.resample1(x[:,3:,:,:], flownetc_flow)\r\n diff_img0 = x[:,:3,:,:] - resampled_img1 \r\n norm_diff_img0 = self.channelnorm(diff_img0)\r\n\r\n # concat img0, img1, img1->img0, flow, diff-mag ; \r\n concat1 = torch.cat((x, resampled_img1, flownetc_flow/self.div_flow, norm_diff_img0), dim=1)\r\n \r\n # flownets1\r\n flownets1_flow2 = self.flownets_1(concat1)[0]\r\n flownets1_flow = self.upsample2(flownets1_flow2*self.div_flow) \r\n\r\n # warp img1 to img0 using flownets1; magnitude of diff between img0 and and warped_img1\r\n resampled_img1 = self.resample2(x[:,3:,:,:], flownets1_flow)\r\n diff_img0 = x[:,:3,:,:] - resampled_img1\r\n norm_diff_img0 = self.channelnorm(diff_img0)\r\n\r\n # concat img0, img1, img1->img0, flow, diff-mag\r\n concat2 = torch.cat((x, resampled_img1, flownets1_flow/self.div_flow, norm_diff_img0), dim=1)\r\n\r\n # flownets2\r\n flownets2_flow2 = self.flownets_2(concat2)[0]\r\n flownets2_flow = self.upsample3(flownets2_flow2 * self.div_flow)\r\n\r\n return flownets2_flow\r\n\r\n"
},
{
"path": "dvs/flownet2/networks/FlowNetC.py",
"content": "import torch\nimport torch.nn as nn\nfrom torch.nn import init\n\nimport math\nimport numpy as np\n\nfrom .correlation_package.correlation import Correlation\n\nfrom .submodules import *\n'Parameter count , 39,175,298 '\n\nclass FlowNetC(nn.Module):\n def __init__(self,args, batchNorm=True, div_flow = 20):\n super(FlowNetC,self).__init__()\n\n self.batchNorm = batchNorm\n self.div_flow = div_flow\n\n self.conv1 = conv(self.batchNorm, 3, 64, kernel_size=7, stride=2)\n self.conv2 = conv(self.batchNorm, 64, 128, kernel_size=5, stride=2)\n self.conv3 = conv(self.batchNorm, 128, 256, kernel_size=5, stride=2)\n self.conv_redir = conv(self.batchNorm, 256, 32, kernel_size=1, stride=1)\n\n if args.fp16:\n self.corr = nn.Sequential(\n tofp32(),\n Correlation(pad_size=20, kernel_size=1, max_displacement=20, stride1=1, stride2=2, corr_multiply=1),\n tofp16())\n else:\n self.corr = Correlation(pad_size=20, kernel_size=1, max_displacement=20, stride1=1, stride2=2, corr_multiply=1)\n\n self.corr_activation = nn.LeakyReLU(0.1,inplace=True)\n self.conv3_1 = conv(self.batchNorm, 473, 256)\n self.conv4 = conv(self.batchNorm, 256, 512, stride=2)\n self.conv4_1 = conv(self.batchNorm, 512, 512)\n self.conv5 = conv(self.batchNorm, 512, 512, stride=2)\n self.conv5_1 = conv(self.batchNorm, 512, 512)\n self.conv6 = conv(self.batchNorm, 512, 1024, stride=2)\n self.conv6_1 = conv(self.batchNorm,1024, 1024)\n\n self.deconv5 = deconv(1024,512)\n self.deconv4 = deconv(1026,256)\n self.deconv3 = deconv(770,128)\n self.deconv2 = deconv(386,64)\n\n self.predict_flow6 = predict_flow(1024)\n self.predict_flow5 = predict_flow(1026)\n self.predict_flow4 = predict_flow(770)\n self.predict_flow3 = predict_flow(386)\n self.predict_flow2 = predict_flow(194)\n\n self.upsampled_flow6_to_5 = nn.ConvTranspose2d(2, 2, 4, 2, 1, bias=True)\n self.upsampled_flow5_to_4 = nn.ConvTranspose2d(2, 2, 4, 2, 1, bias=True)\n self.upsampled_flow4_to_3 = nn.ConvTranspose2d(2, 2, 4, 2, 1, bias=True)\n self.upsampled_flow3_to_2 = nn.ConvTranspose2d(2, 2, 4, 2, 1, bias=True)\n\n for m in self.modules():\n if isinstance(m, nn.Conv2d):\n if m.bias is not None:\n init.uniform_(m.bias)\n init.xavier_uniform_(m.weight)\n\n if isinstance(m, nn.ConvTranspose2d):\n if m.bias is not None:\n init.uniform_(m.bias)\n init.xavier_uniform_(m.weight)\n # init_deconv_bilinear(m.weight)\n self.upsample1 = nn.Upsample(scale_factor=4, mode='bilinear')\n\n def forward(self, x):\n x1 = x[:,0:3,:,:]\n x2 = x[:,3::,:,:]\n\n out_conv1a = self.conv1(x1)\n out_conv2a = self.conv2(out_conv1a)\n out_conv3a = self.conv3(out_conv2a)\n\n # FlownetC bottom input stream\n out_conv1b = self.conv1(x2)\n \n out_conv2b = self.conv2(out_conv1b)\n out_conv3b = self.conv3(out_conv2b)\n\n # Merge streams\n out_corr = self.corr(out_conv3a, out_conv3b) # False\n out_corr = self.corr_activation(out_corr)\n\n # Redirect top input stream and concatenate\n out_conv_redir = self.conv_redir(out_conv3a)\n\n in_conv3_1 = torch.cat((out_conv_redir, out_corr), 1)\n\n # Merged conv layers\n out_conv3_1 = self.conv3_1(in_conv3_1)\n\n out_conv4 = self.conv4_1(self.conv4(out_conv3_1))\n\n out_conv5 = self.conv5_1(self.conv5(out_conv4))\n out_conv6 = self.conv6_1(self.conv6(out_conv5))\n\n flow6 = self.predict_flow6(out_conv6)\n flow6_up = self.upsampled_flow6_to_5(flow6)\n out_deconv5 = self.deconv5(out_conv6)\n\n concat5 = torch.cat((out_conv5,out_deconv5,flow6_up),1)\n\n flow5 = self.predict_flow5(concat5)\n flow5_up = self.upsampled_flow5_to_4(flow5)\n out_deconv4 = self.deconv4(concat5)\n concat4 = torch.cat((out_conv4,out_deconv4,flow5_up),1)\n\n flow4 = self.predict_flow4(concat4)\n flow4_up = self.upsampled_flow4_to_3(flow4)\n out_deconv3 = self.deconv3(concat4)\n concat3 = torch.cat((out_conv3_1,out_deconv3,flow4_up),1)\n\n flow3 = self.predict_flow3(concat3)\n flow3_up = self.upsampled_flow3_to_2(flow3)\n out_deconv2 = self.deconv2(concat3)\n concat2 = torch.cat((out_conv2a,out_deconv2,flow3_up),1)\n\n flow2 = self.predict_flow2(concat2)\n\n if self.training:\n return flow2,flow3,flow4,flow5,flow6\n else:\n return flow2,\n"
},
{
"path": "dvs/flownet2/networks/FlowNetFusion.py",
"content": "import torch\nimport torch.nn as nn\nfrom torch.nn import init\n\nimport math\nimport numpy as np\n\nfrom .submodules import *\n'Parameter count = 581,226'\n\nclass FlowNetFusion(nn.Module):\n def __init__(self,args, batchNorm=True):\n super(FlowNetFusion,self).__init__()\n\n self.batchNorm = batchNorm\n self.conv0 = conv(self.batchNorm, 11, 64)\n self.conv1 = conv(self.batchNorm, 64, 64, stride=2)\n self.conv1_1 = conv(self.batchNorm, 64, 128)\n self.conv2 = conv(self.batchNorm, 128, 128, stride=2)\n self.conv2_1 = conv(self.batchNorm, 128, 128)\n\n self.deconv1 = deconv(128,32)\n self.deconv0 = deconv(162,16)\n\n self.inter_conv1 = i_conv(self.batchNorm, 162, 32)\n self.inter_conv0 = i_conv(self.batchNorm, 82, 16)\n\n self.predict_flow2 = predict_flow(128)\n self.predict_flow1 = predict_flow(32)\n self.predict_flow0 = predict_flow(16)\n\n self.upsampled_flow2_to_1 = nn.ConvTranspose2d(2, 2, 4, 2, 1)\n self.upsampled_flow1_to_0 = nn.ConvTranspose2d(2, 2, 4, 2, 1)\n\n for m in self.modules():\n if isinstance(m, nn.Conv2d):\n if m.bias is not None:\n init.uniform_(m.bias)\n init.xavier_uniform_(m.weight)\n\n if isinstance(m, nn.ConvTranspose2d):\n if m.bias is not None:\n init.uniform_(m.bias)\n init.xavier_uniform_(m.weight)\n # init_deconv_bilinear(m.weight)\n\n def forward(self, x):\n out_conv0 = self.conv0(x)\n out_conv1 = self.conv1_1(self.conv1(out_conv0))\n out_conv2 = self.conv2_1(self.conv2(out_conv1))\n\n flow2 = self.predict_flow2(out_conv2)\n flow2_up = self.upsampled_flow2_to_1(flow2)\n out_deconv1 = self.deconv1(out_conv2)\n \n concat1 = torch.cat((out_conv1,out_deconv1,flow2_up),1)\n out_interconv1 = self.inter_conv1(concat1)\n flow1 = self.predict_flow1(out_interconv1)\n flow1_up = self.upsampled_flow1_to_0(flow1)\n out_deconv0 = self.deconv0(concat1)\n \n concat0 = torch.cat((out_conv0,out_deconv0,flow1_up),1)\n out_interconv0 = self.inter_conv0(concat0)\n flow0 = self.predict_flow0(out_interconv0)\n\n return flow0\n\n"
},
{
"path": "dvs/flownet2/networks/FlowNetS.py",
"content": "'''\nPortions of this code copyright 2017, Clement Pinard\n'''\n\nimport torch\nimport torch.nn as nn\nfrom torch.nn import init\n\nimport math\nimport numpy as np\n\nfrom .submodules import *\n'Parameter count : 38,676,504 '\n\nclass FlowNetS(nn.Module):\n def __init__(self, args, input_channels = 12, batchNorm=True):\n super(FlowNetS,self).__init__()\n\n self.batchNorm = batchNorm\n self.conv1 = conv(self.batchNorm, input_channels, 64, kernel_size=7, stride=2)\n self.conv2 = conv(self.batchNorm, 64, 128, kernel_size=5, stride=2)\n self.conv3 = conv(self.batchNorm, 128, 256, kernel_size=5, stride=2)\n self.conv3_1 = conv(self.batchNorm, 256, 256)\n self.conv4 = conv(self.batchNorm, 256, 512, stride=2)\n self.conv4_1 = conv(self.batchNorm, 512, 512)\n self.conv5 = conv(self.batchNorm, 512, 512, stride=2)\n self.conv5_1 = conv(self.batchNorm, 512, 512)\n self.conv6 = conv(self.batchNorm, 512, 1024, stride=2)\n self.conv6_1 = conv(self.batchNorm,1024, 1024)\n\n self.deconv5 = deconv(1024,512)\n self.deconv4 = deconv(1026,256)\n self.deconv3 = deconv(770,128)\n self.deconv2 = deconv(386,64)\n\n self.predict_flow6 = predict_flow(1024)\n self.predict_flow5 = predict_flow(1026)\n self.predict_flow4 = predict_flow(770)\n self.predict_flow3 = predict_flow(386)\n self.predict_flow2 = predict_flow(194)\n\n self.upsampled_flow6_to_5 = nn.ConvTranspose2d(2, 2, 4, 2, 1, bias=False)\n self.upsampled_flow5_to_4 = nn.ConvTranspose2d(2, 2, 4, 2, 1, bias=False)\n self.upsampled_flow4_to_3 = nn.ConvTranspose2d(2, 2, 4, 2, 1, bias=False)\n self.upsampled_flow3_to_2 = nn.ConvTranspose2d(2, 2, 4, 2, 1, bias=False)\n\n for m in self.modules():\n if isinstance(m, nn.Conv2d):\n if m.bias is not None:\n init.uniform_(m.bias)\n init.xavier_uniform_(m.weight)\n\n if isinstance(m, nn.ConvTranspose2d):\n if m.bias is not None:\n init.uniform_(m.bias)\n init.xavier_uniform_(m.weight)\n # init_deconv_bilinear(m.weight)\n self.upsample1 = nn.Upsample(scale_factor=4, mode='bilinear')\n\n def forward(self, x):\n out_conv1 = self.conv1(x)\n\n out_conv2 = self.conv2(out_conv1)\n out_conv3 = self.conv3_1(self.conv3(out_conv2))\n out_conv4 = self.conv4_1(self.conv4(out_conv3))\n out_conv5 = self.conv5_1(self.conv5(out_conv4))\n out_conv6 = self.conv6_1(self.conv6(out_conv5))\n\n flow6 = self.predict_flow6(out_conv6)\n flow6_up = self.upsampled_flow6_to_5(flow6)\n out_deconv5 = self.deconv5(out_conv6)\n \n concat5 = torch.cat((out_conv5,out_deconv5,flow6_up),1)\n flow5 = self.predict_flow5(concat5)\n flow5_up = self.upsampled_flow5_to_4(flow5)\n out_deconv4 = self.deconv4(concat5)\n \n concat4 = torch.cat((out_conv4,out_deconv4,flow5_up),1)\n flow4 = self.predict_flow4(concat4)\n flow4_up = self.upsampled_flow4_to_3(flow4)\n out_deconv3 = self.deconv3(concat4)\n \n concat3 = torch.cat((out_conv3,out_deconv3,flow4_up),1)\n flow3 = self.predict_flow3(concat3)\n flow3_up = self.upsampled_flow3_to_2(flow3)\n out_deconv2 = self.deconv2(concat3)\n\n concat2 = torch.cat((out_conv2,out_deconv2,flow3_up),1)\n flow2 = self.predict_flow2(concat2)\n\n if self.training:\n return flow2,flow3,flow4,flow5,flow6\n else:\n return flow2,\n\n"
},
{
"path": "dvs/flownet2/networks/FlowNetSD.py",
"content": "import torch\nimport torch.nn as nn\nfrom torch.nn import init\n\nimport math\nimport numpy as np\n\nfrom .submodules import *\n'Parameter count = 45,371,666'\n\nclass FlowNetSD(nn.Module):\n def __init__(self, args, batchNorm=True):\n super(FlowNetSD,self).__init__()\n\n self.batchNorm = batchNorm\n self.conv0 = conv(self.batchNorm, 6, 64)\n self.conv1 = conv(self.batchNorm, 64, 64, stride=2)\n self.conv1_1 = conv(self.batchNorm, 64, 128)\n self.conv2 = conv(self.batchNorm, 128, 128, stride=2)\n self.conv2_1 = conv(self.batchNorm, 128, 128)\n self.conv3 = conv(self.batchNorm, 128, 256, stride=2)\n self.conv3_1 = conv(self.batchNorm, 256, 256)\n self.conv4 = conv(self.batchNorm, 256, 512, stride=2)\n self.conv4_1 = conv(self.batchNorm, 512, 512)\n self.conv5 = conv(self.batchNorm, 512, 512, stride=2)\n self.conv5_1 = conv(self.batchNorm, 512, 512)\n self.conv6 = conv(self.batchNorm, 512, 1024, stride=2)\n self.conv6_1 = conv(self.batchNorm,1024, 1024)\n\n self.deconv5 = deconv(1024,512)\n self.deconv4 = deconv(1026,256)\n self.deconv3 = deconv(770,128)\n self.deconv2 = deconv(386,64)\n\n self.inter_conv5 = i_conv(self.batchNorm, 1026, 512)\n self.inter_conv4 = i_conv(self.batchNorm, 770, 256)\n self.inter_conv3 = i_conv(self.batchNorm, 386, 128)\n self.inter_conv2 = i_conv(self.batchNorm, 194, 64)\n\n self.predict_flow6 = predict_flow(1024)\n self.predict_flow5 = predict_flow(512)\n self.predict_flow4 = predict_flow(256)\n self.predict_flow3 = predict_flow(128)\n self.predict_flow2 = predict_flow(64)\n\n self.upsampled_flow6_to_5 = nn.ConvTranspose2d(2, 2, 4, 2, 1)\n self.upsampled_flow5_to_4 = nn.ConvTranspose2d(2, 2, 4, 2, 1)\n self.upsampled_flow4_to_3 = nn.ConvTranspose2d(2, 2, 4, 2, 1)\n self.upsampled_flow3_to_2 = nn.ConvTranspose2d(2, 2, 4, 2, 1)\n\n for m in self.modules():\n if isinstance(m, nn.Conv2d):\n if m.bias is not None:\n init.uniform_(m.bias)\n init.xavier_uniform_(m.weight)\n\n if isinstance(m, nn.ConvTranspose2d):\n if m.bias is not None:\n init.uniform_(m.bias)\n init.xavier_uniform_(m.weight)\n # init_deconv_bilinear(m.weight)\n self.upsample1 = nn.Upsample(scale_factor=4, mode='bilinear')\n\n\n\n def forward(self, x):\n out_conv0 = self.conv0(x)\n out_conv1 = self.conv1_1(self.conv1(out_conv0))\n out_conv2 = self.conv2_1(self.conv2(out_conv1))\n\n out_conv3 = self.conv3_1(self.conv3(out_conv2))\n out_conv4 = self.conv4_1(self.conv4(out_conv3))\n out_conv5 = self.conv5_1(self.conv5(out_conv4))\n out_conv6 = self.conv6_1(self.conv6(out_conv5))\n\n flow6 = self.predict_flow6(out_conv6)\n flow6_up = self.upsampled_flow6_to_5(flow6)\n out_deconv5 = self.deconv5(out_conv6)\n \n concat5 = torch.cat((out_conv5,out_deconv5,flow6_up),1)\n out_interconv5 = self.inter_conv5(concat5)\n flow5 = self.predict_flow5(out_interconv5)\n\n flow5_up = self.upsampled_flow5_to_4(flow5)\n out_deconv4 = self.deconv4(concat5)\n \n concat4 = torch.cat((out_conv4,out_deconv4,flow5_up),1)\n out_interconv4 = self.inter_conv4(concat4)\n flow4 = self.predict_flow4(out_interconv4)\n flow4_up = self.upsampled_flow4_to_3(flow4)\n out_deconv3 = self.deconv3(concat4)\n \n concat3 = torch.cat((out_conv3,out_deconv3,flow4_up),1)\n out_interconv3 = self.inter_conv3(concat3)\n flow3 = self.predict_flow3(out_interconv3)\n flow3_up = self.upsampled_flow3_to_2(flow3)\n out_deconv2 = self.deconv2(concat3)\n\n concat2 = torch.cat((out_conv2,out_deconv2,flow3_up),1)\n out_interconv2 = self.inter_conv2(concat2)\n flow2 = self.predict_flow2(out_interconv2)\n\n if self.training:\n return flow2,flow3,flow4,flow5,flow6\n else:\n return flow2,\n"
},
{
"path": "dvs/flownet2/networks/__init__.py",
"content": ""
},
{
"path": "dvs/flownet2/networks/channelnorm_package/__init__.py",
"content": ""
},
{
"path": "dvs/flownet2/networks/channelnorm_package/channelnorm.py",
"content": "from torch.autograd import Function, Variable\nfrom torch.nn.modules.module import Module\nimport channelnorm_cuda\n\nclass ChannelNormFunction(Function):\n\n @staticmethod\n def forward(ctx, input1, norm_deg=2):\n assert input1.is_contiguous()\n b, _, h, w = input1.size()\n output = input1.new(b, 1, h, w).zero_()\n\n channelnorm_cuda.forward(input1, output, norm_deg)\n ctx.save_for_backward(input1, output)\n ctx.norm_deg = norm_deg\n\n return output\n\n @staticmethod\n def backward(ctx, grad_output):\n input1, output = ctx.saved_tensors\n\n grad_input1 = Variable(input1.new(input1.size()).zero_())\n\n channelnorm_cuda.backward(input1, output, grad_output.data,\n grad_input1.data, ctx.norm_deg)\n\n return grad_input1, None\n\n\nclass ChannelNorm(Module):\n\n def __init__(self, norm_deg=2):\n super(ChannelNorm, self).__init__()\n self.norm_deg = norm_deg\n\n def forward(self, input1):\n return ChannelNormFunction.apply(input1, self.norm_deg)\n\n"
},
{
"path": "dvs/flownet2/networks/channelnorm_package/channelnorm_cuda.cc",
"content": "#include \n#include \n\n#include \"channelnorm_kernel.cuh\"\n\nint channelnorm_cuda_forward(\n at::Tensor& input1, \n at::Tensor& output,\n int norm_deg) {\n\n channelnorm_kernel_forward(input1, output, norm_deg);\n return 1;\n}\n\n\nint channelnorm_cuda_backward(\n at::Tensor& input1, \n at::Tensor& output,\n at::Tensor& gradOutput,\n at::Tensor& gradInput1,\n int norm_deg) {\n\n channelnorm_kernel_backward(input1, output, gradOutput, gradInput1, norm_deg);\n return 1;\n}\n\nPYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {\n m.def(\"forward\", &channelnorm_cuda_forward, \"Channel norm forward (CUDA)\");\n m.def(\"backward\", &channelnorm_cuda_backward, \"Channel norm backward (CUDA)\");\n}\n\n"
},
{
"path": "dvs/flownet2/networks/channelnorm_package/channelnorm_kernel.cu",
"content": "#include \n#include \n#include \n\n#include \"channelnorm_kernel.cuh\"\n\n#define CUDA_NUM_THREADS 512 \n\n#define DIM0(TENSOR) ((TENSOR).x)\n#define DIM1(TENSOR) ((TENSOR).y)\n#define DIM2(TENSOR) ((TENSOR).z)\n#define DIM3(TENSOR) ((TENSOR).w)\n\n#define DIM3_INDEX(TENSOR, xx, yy, zz, ww) ((TENSOR)[((xx) * (TENSOR##_stride.x)) + ((yy) * (TENSOR##_stride.y)) + ((zz) * (TENSOR##_stride.z)) + ((ww) * (TENSOR##_stride.w))])\n\nusing at::Half;\n\ntemplate \n__global__ void kernel_channelnorm_update_output(\n const int n, \n const scalar_t* __restrict__ input1,\n const long4 input1_size,\n const long4 input1_stride,\n scalar_t* __restrict__ output, \n const long4 output_size,\n const long4 output_stride,\n int norm_deg) {\n\n int index = blockIdx.x * blockDim.x + threadIdx.x;\n\n if (index >= n) {\n return;\n }\n\n int dim_b = DIM0(output_size);\n int dim_c = DIM1(output_size);\n int dim_h = DIM2(output_size);\n int dim_w = DIM3(output_size);\n int dim_chw = dim_c * dim_h * dim_w;\n\n int b = ( index / dim_chw ) % dim_b;\n int y = ( index / dim_w ) % dim_h;\n int x = ( index ) % dim_w;\n\n int i1dim_c = DIM1(input1_size);\n int i1dim_h = DIM2(input1_size);\n int i1dim_w = DIM3(input1_size);\n int i1dim_chw = i1dim_c * i1dim_h * i1dim_w;\n int i1dim_hw = i1dim_h * i1dim_w;\n\n float result = 0.0;\n\n for (int c = 0; c < i1dim_c; ++c) {\n int i1Index = b * i1dim_chw + c * i1dim_hw + y * i1dim_w + x;\n scalar_t val = input1[i1Index];\n result += static_cast(val * val);\n }\n result = sqrt(result);\n output[index] = static_cast(result);\n}\n\n\ntemplate \n__global__ void kernel_channelnorm_backward_input1(\n const int n,\n const scalar_t* __restrict__ input1, const long4 input1_size, const long4 input1_stride,\n const scalar_t* __restrict__ output, const long4 output_size, const long4 output_stride, \n const scalar_t* __restrict__ gradOutput, const long4 gradOutput_size, const long4 gradOutput_stride,\n scalar_t* __restrict__ gradInput, const long4 gradInput_size, const long4 gradInput_stride, \n int norm_deg) {\n\n int index = blockIdx.x * blockDim.x + threadIdx.x;\n\n if (index >= n) {\n return;\n }\n\n float val = 0.0;\n\n int dim_b = DIM0(gradInput_size);\n int dim_c = DIM1(gradInput_size);\n int dim_h = DIM2(gradInput_size);\n int dim_w = DIM3(gradInput_size);\n int dim_chw = dim_c * dim_h * dim_w;\n int dim_hw = dim_h * dim_w;\n\n int b = ( index / dim_chw ) % dim_b;\n int y = ( index / dim_w ) % dim_h;\n int x = ( index ) % dim_w;\n\n\n int outIndex = b * dim_hw + y * dim_w + x;\n val = static_cast(gradOutput[outIndex]) * static_cast(input1[index]) / (static_cast(output[outIndex])+1e-9);\n gradInput[index] = static_cast(val);\n\n}\n\nvoid channelnorm_kernel_forward(\n at::Tensor& input1, \n at::Tensor& output, \n int norm_deg) {\n\n const long4 input1_size = make_long4(input1.size(0), input1.size(1), input1.size(2), input1.size(3));\n const long4 input1_stride = make_long4(input1.stride(0), input1.stride(1), input1.stride(2), input1.stride(3));\n\n const long4 output_size = make_long4(output.size(0), output.size(1), output.size(2), output.size(3));\n const long4 output_stride = make_long4(output.stride(0), output.stride(1), output.stride(2), output.stride(3));\n\n int n = output.numel();\n\n AT_DISPATCH_FLOATING_TYPES_AND_HALF(input1.type(), \"channelnorm_forward\", ([&] {\n\n kernel_channelnorm_update_output<<< (n + CUDA_NUM_THREADS - 1)/CUDA_NUM_THREADS, CUDA_NUM_THREADS, 0, at::cuda::getCurrentCUDAStream() >>>(\n//at::globalContext().getCurrentCUDAStream() >>>(\n n,\n input1.data(), \n input1_size,\n input1_stride, \n output.data(),\n output_size,\n output_stride, \n norm_deg);\n\n }));\n\n // TODO: ATen-equivalent check\n\n // THCudaCheck(cudaGetLastError());\n}\n\nvoid channelnorm_kernel_backward(\n at::Tensor& input1, \n at::Tensor& output,\n at::Tensor& gradOutput, \n at::Tensor& gradInput1, \n int norm_deg) {\n\n const long4 input1_size = make_long4(input1.size(0), input1.size(1), input1.size(2), input1.size(3));\n const long4 input1_stride = make_long4(input1.stride(0), input1.stride(1), input1.stride(2), input1.stride(3));\n\n const long4 output_size = make_long4(output.size(0), output.size(1), output.size(2), output.size(3));\n const long4 output_stride = make_long4(output.stride(0), output.stride(1), output.stride(2), output.stride(3));\n\n const long4 gradOutput_size = make_long4(gradOutput.size(0), gradOutput.size(1), gradOutput.size(2), gradOutput.size(3));\n const long4 gradOutput_stride = make_long4(gradOutput.stride(0), gradOutput.stride(1), gradOutput.stride(2), gradOutput.stride(3));\n\n const long4 gradInput1_size = make_long4(gradInput1.size(0), gradInput1.size(1), gradInput1.size(2), gradInput1.size(3));\n const long4 gradInput1_stride = make_long4(gradInput1.stride(0), gradInput1.stride(1), gradInput1.stride(2), gradInput1.stride(3));\n\n int n = gradInput1.numel();\n\n AT_DISPATCH_FLOATING_TYPES_AND_HALF(input1.type(), \"channelnorm_backward_input1\", ([&] {\n\n kernel_channelnorm_backward_input1<<< (n + CUDA_NUM_THREADS - 1)/CUDA_NUM_THREADS, CUDA_NUM_THREADS, 0, at::cuda::getCurrentCUDAStream() >>>(\n//at::globalContext().getCurrentCUDAStream() >>>(\n n, \n input1.data(),\n input1_size,\n input1_stride,\n output.data(),\n output_size,\n output_stride,\n gradOutput.data(),\n gradOutput_size,\n gradOutput_stride, \n gradInput1.data(),\n gradInput1_size,\n gradInput1_stride,\n norm_deg\n );\n\n }));\n\n // TODO: Add ATen-equivalent check\n\n// THCudaCheck(cudaGetLastError());\n}\n"
},
{
"path": "dvs/flownet2/networks/channelnorm_package/channelnorm_kernel.cuh",
"content": "#pragma once\n\n#include \n\nvoid channelnorm_kernel_forward(\n at::Tensor& input1,\n at::Tensor& output, \n int norm_deg);\n\n\nvoid channelnorm_kernel_backward(\n at::Tensor& input1,\n at::Tensor& output,\n at::Tensor& gradOutput,\n at::Tensor& gradInput1,\n int norm_deg);\n"
},
{
"path": "dvs/flownet2/networks/channelnorm_package/setup.py",
"content": "#!/usr/bin/env python3\nimport os\nimport torch\n\nfrom setuptools import setup\nfrom torch.utils.cpp_extension import BuildExtension, CUDAExtension\n\ncxx_args = ['-std=c++11']\n\nnvcc_args = [\n '-gencode', 'arch=compute_52,code=sm_52',\n '-gencode', 'arch=compute_60,code=sm_60',\n '-gencode', 'arch=compute_61,code=sm_61',\n '-gencode', 'arch=compute_70,code=sm_70',\n '-gencode', 'arch=compute_70,code=compute_70'\n]\n\nsetup(\n name='channelnorm_cuda',\n ext_modules=[\n CUDAExtension('channelnorm_cuda', [\n 'channelnorm_cuda.cc',\n 'channelnorm_kernel.cu'\n ], extra_compile_args={'cxx': cxx_args, 'nvcc': nvcc_args})\n ],\n cmdclass={\n 'build_ext': BuildExtension\n })\n"
},
{
"path": "dvs/flownet2/networks/correlation_package/__init__.py",
"content": ""
},
{
"path": "dvs/flownet2/networks/correlation_package/correlation.py",
"content": "import torch\nfrom torch.nn.modules.module import Module\nfrom torch.autograd import Function\nimport correlation_cuda\n\nclass CorrelationFunction(Function):\n\n @staticmethod\n def forward(ctx, input1, input2, pad_size=3, kernel_size=3, max_displacement=20, stride1=1, stride2=2, corr_multiply=1):\n ctx.save_for_backward(input1, input2)\n\n ctx.pad_size = pad_size\n ctx.kernel_size = kernel_size\n ctx.max_displacement = max_displacement\n ctx.stride1 = stride1\n ctx.stride2 = stride2\n ctx.corr_multiply = corr_multiply\n\n with torch.cuda.device_of(input1):\n rbot1 = input1.new()\n rbot2 = input2.new()\n output = input1.new()\n\n correlation_cuda.forward(input1, input2, rbot1, rbot2, output,\n ctx.pad_size, ctx.kernel_size, ctx.max_displacement, ctx.stride1, ctx.stride2, ctx.corr_multiply)\n\n return output\n\n @staticmethod\n def backward(ctx, grad_output):\n input1, input2 = ctx.saved_tensors\n\n with torch.cuda.device_of(input1):\n rbot1 = input1.new()\n rbot2 = input2.new()\n\n grad_input1 = input1.new()\n grad_input2 = input2.new()\n\n correlation_cuda.backward(input1, input2, rbot1, rbot2, grad_output, grad_input1, grad_input2,\n ctx.pad_size, ctx.kernel_size, ctx.max_displacement, ctx.stride1, ctx.stride2, ctx.corr_multiply)\n\n return grad_input1, grad_input2, None, None, None, None, None, None\n\n\nclass Correlation(Module):\n def __init__(self, pad_size=0, kernel_size=0, max_displacement=0, stride1=1, stride2=2, corr_multiply=1):\n super(Correlation, self).__init__()\n self.pad_size = pad_size\n self.kernel_size = kernel_size\n self.max_displacement = max_displacement\n self.stride1 = stride1\n self.stride2 = stride2\n self.corr_multiply = corr_multiply\n\n def forward(self, input1, input2):\n\n result = CorrelationFunction.apply(input1, input2, self.pad_size, self.kernel_size, self.max_displacement, self.stride1, self.stride2, self.corr_multiply)\n\n return result\n\n"
},
{
"path": "dvs/flownet2/networks/correlation_package/correlation_cuda.cc",
"content": "#include \n#include \n#include \n#include \n#include \n#include \n\n#include \"correlation_cuda_kernel.cuh\"\n\nint correlation_forward_cuda(at::Tensor& input1, at::Tensor& input2, at::Tensor& rInput1, at::Tensor& rInput2, at::Tensor& output,\n int pad_size,\n int kernel_size,\n int max_displacement,\n int stride1,\n int stride2,\n int corr_type_multiply)\n{\n\n int batchSize = input1.size(0);\n\n int nInputChannels = input1.size(1);\n int inputHeight = input1.size(2);\n int inputWidth = input1.size(3);\n\n int kernel_radius = (kernel_size - 1) / 2;\n int border_radius = kernel_radius + max_displacement;\n\n int paddedInputHeight = inputHeight + 2 * pad_size;\n int paddedInputWidth = inputWidth + 2 * pad_size;\n\n int nOutputChannels = ((max_displacement/stride2)*2 + 1) * ((max_displacement/stride2)*2 + 1);\n\n int outputHeight = ceil(static_cast(paddedInputHeight - 2 * border_radius) / static_cast(stride1));\n int outputwidth = ceil(static_cast(paddedInputWidth - 2 * border_radius) / static_cast(stride1));\n\n rInput1.resize_({batchSize, paddedInputHeight, paddedInputWidth, nInputChannels});\n rInput2.resize_({batchSize, paddedInputHeight, paddedInputWidth, nInputChannels});\n output.resize_({batchSize, nOutputChannels, outputHeight, outputwidth});\n\n rInput1.fill_(0);\n rInput2.fill_(0);\n output.fill_(0);\n\n int success = correlation_forward_cuda_kernel(\n output,\n output.size(0), \n output.size(1),\n output.size(2),\n output.size(3),\n output.stride(0),\n output.stride(1),\n output.stride(2),\n output.stride(3),\n input1,\n input1.size(1),\n input1.size(2),\n input1.size(3),\n input1.stride(0),\n input1.stride(1),\n input1.stride(2),\n input1.stride(3),\n input2,\n input2.size(1),\n input2.stride(0),\n input2.stride(1),\n input2.stride(2),\n input2.stride(3),\n rInput1,\n rInput2,\n pad_size, \n kernel_size,\n max_displacement,\n stride1,\n stride2,\n corr_type_multiply,\n\tat::cuda::getCurrentCUDAStream()\n //at::globalContext().getCurrentCUDAStream()\n );\n\n //check for errors\n if (!success) {\n AT_ERROR(\"CUDA call failed\");\n }\n\n return 1;\n\n}\n\nint correlation_backward_cuda(at::Tensor& input1, at::Tensor& input2, at::Tensor& rInput1, at::Tensor& rInput2, at::Tensor& gradOutput, \n at::Tensor& gradInput1, at::Tensor& gradInput2,\n int pad_size,\n int kernel_size,\n int max_displacement,\n int stride1,\n int stride2,\n int corr_type_multiply)\n{\n\n int batchSize = input1.size(0);\n int nInputChannels = input1.size(1);\n int paddedInputHeight = input1.size(2)+ 2 * pad_size;\n int paddedInputWidth = input1.size(3)+ 2 * pad_size;\n\n int height = input1.size(2);\n int width = input1.size(3);\n\n rInput1.resize_({batchSize, paddedInputHeight, paddedInputWidth, nInputChannels});\n rInput2.resize_({batchSize, paddedInputHeight, paddedInputWidth, nInputChannels});\n gradInput1.resize_({batchSize, nInputChannels, height, width});\n gradInput2.resize_({batchSize, nInputChannels, height, width});\n\n rInput1.fill_(0);\n rInput2.fill_(0);\n gradInput1.fill_(0);\n gradInput2.fill_(0);\n\n int success = correlation_backward_cuda_kernel(gradOutput,\n gradOutput.size(0),\n gradOutput.size(1),\n gradOutput.size(2),\n gradOutput.size(3),\n gradOutput.stride(0),\n gradOutput.stride(1),\n gradOutput.stride(2),\n gradOutput.stride(3),\n input1,\n input1.size(1),\n input1.size(2),\n input1.size(3),\n input1.stride(0),\n input1.stride(1),\n input1.stride(2),\n input1.stride(3),\n input2, \n input2.stride(0),\n input2.stride(1),\n input2.stride(2),\n input2.stride(3),\n gradInput1,\n gradInput1.stride(0),\n gradInput1.stride(1),\n gradInput1.stride(2),\n gradInput1.stride(3),\n gradInput2,\n gradInput2.size(1),\n gradInput2.stride(0),\n gradInput2.stride(1),\n gradInput2.stride(2),\n gradInput2.stride(3),\n rInput1,\n rInput2,\n pad_size,\n kernel_size,\n max_displacement,\n stride1, \n stride2,\n corr_type_multiply,\n\t\t\t\t\t\t\t\t\t\t\t\tat::cuda::getCurrentCUDAStream()\n //at::globalContext().getCurrentCUDAStream()\n );\n\n if (!success) {\n AT_ERROR(\"CUDA call failed\");\n }\n\n return 1;\n}\n\nPYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {\n m.def(\"forward\", &correlation_forward_cuda, \"Correlation forward (CUDA)\");\n m.def(\"backward\", &correlation_backward_cuda, \"Correlation backward (CUDA)\");\n}\n\n"
},
{
"path": "dvs/flownet2/networks/correlation_package/correlation_cuda_kernel.cu",
"content": "#include \n\n#include \"correlation_cuda_kernel.cuh\"\n\n#define CUDA_NUM_THREADS 1024\n#define THREADS_PER_BLOCK 32\n#define FULL_MASK 0xffffffff\n\n#include \n#include \n#include \n#include \n\nusing at::Half;\n\ntemplate\n__forceinline__ __device__ scalar_t warpReduceSum(scalar_t val) {\n for (int offset = 16; offset > 0; offset /= 2)\n val += __shfl_down_sync(FULL_MASK, val, offset);\n return val;\n}\n\ntemplate\n__forceinline__ __device__ scalar_t blockReduceSum(scalar_t val) {\n\n static __shared__ scalar_t shared[32];\n int lane = threadIdx.x % warpSize;\n int wid = threadIdx.x / warpSize;\n\n val = warpReduceSum(val);\n\n if (lane == 0)\n shared[wid] = val;\n\n __syncthreads();\n\n val = (threadIdx.x < blockDim.x / warpSize) ? shared[lane] : 0;\n\n if (wid == 0)\n val = warpReduceSum(val);\n\n return val;\n}\n\n\ntemplate \n__global__ void channels_first(const scalar_t* __restrict__ input, scalar_t* rinput, int channels, int height, int width, int pad_size)\n{\n\n // n (batch size), c (num of channels), y (height), x (width)\n int n = blockIdx.x;\n int y = blockIdx.y;\n int x = blockIdx.z;\n\n int ch_off = threadIdx.x;\n scalar_t value;\n\n int dimcyx = channels * height * width;\n int dimyx = height * width;\n\n int p_dimx = (width + 2 * pad_size);\n int p_dimy = (height + 2 * pad_size);\n int p_dimyxc = channels * p_dimy * p_dimx;\n int p_dimxc = p_dimx * channels;\n\n for (int c = ch_off; c < channels; c += THREADS_PER_BLOCK) {\n value = input[n * dimcyx + c * dimyx + y * width + x];\n rinput[n * p_dimyxc + (y + pad_size) * p_dimxc + (x + pad_size) * channels + c] = value;\n }\n}\n\n\ntemplate\n__global__ void correlation_forward(scalar_t* __restrict__ output, const int nOutputChannels,\n const int outputHeight, const int outputWidth, const scalar_t* __restrict__ rInput1,\n const int nInputChannels, const int inputHeight, const int inputWidth,\n const scalar_t* __restrict__ rInput2, const int pad_size, const int kernel_size,\n const int max_displacement, const int stride1, const int stride2) {\n\n int32_t pInputWidth = inputWidth + 2 * pad_size;\n int32_t pInputHeight = inputHeight + 2 * pad_size;\n\n int32_t kernel_rad = (kernel_size - 1) / 2;\n\n int32_t displacement_rad = max_displacement / stride2;\n\n int32_t displacement_size = 2 * displacement_rad + 1;\n\n int32_t n = blockIdx.x;\n int32_t y1 = blockIdx.y * stride1 + max_displacement;\n int32_t x1 = blockIdx.z * stride1 + max_displacement;\n int32_t c = threadIdx.x;\n\n int32_t pdimyxc = pInputHeight * pInputWidth * nInputChannels;\n\n int32_t pdimxc = pInputWidth * nInputChannels;\n\n int32_t pdimc = nInputChannels;\n\n int32_t tdimcyx = nOutputChannels * outputHeight * outputWidth;\n int32_t tdimyx = outputHeight * outputWidth;\n int32_t tdimx = outputWidth;\n\n int32_t nelems = kernel_size * kernel_size * pdimc;\n\n // element-wise product along channel axis\n for (int tj = -displacement_rad; tj <= displacement_rad; ++tj) {\n for (int ti = -displacement_rad; ti <= displacement_rad; ++ti) {\n int x2 = x1 + ti * stride2;\n int y2 = y1 + tj * stride2;\n\n float acc0 = 0.0f;\n\n for (int j = -kernel_rad; j <= kernel_rad; ++j) {\n for (int i = -kernel_rad; i <= kernel_rad; ++i) {\n // THREADS_PER_BLOCK\n #pragma unroll\n for (int ch = c; ch < pdimc; ch += blockDim.x) {\n\n int indx1 = n * pdimyxc + (y1 + j) * pdimxc\n + (x1 + i) * pdimc + ch;\n int indx2 = n * pdimyxc + (y2 + j) * pdimxc\n + (x2 + i) * pdimc + ch;\n acc0 += static_cast(rInput1[indx1] * rInput2[indx2]);\n }\n }\n }\n\n if (blockDim.x == warpSize) {\n __syncwarp();\n acc0 = warpReduceSum(acc0);\n } else {\n __syncthreads();\n acc0 = blockReduceSum(acc0);\n }\n\n if (threadIdx.x == 0) {\n\n int tc = (tj + displacement_rad) * displacement_size\n + (ti + displacement_rad);\n const int tindx = n * tdimcyx + tc * tdimyx + blockIdx.y * tdimx\n + blockIdx.z;\n output[tindx] = static_cast(acc0 / nelems);\n }\n }\n }\n}\n\n\ntemplate \n__global__ void correlation_backward_input1(int item, scalar_t* gradInput1, int nInputChannels, int inputHeight, int inputWidth, \n const scalar_t* __restrict__ gradOutput, int nOutputChannels, int outputHeight, int outputWidth, \n const scalar_t* __restrict__ rInput2, \n int pad_size,\n int kernel_size,\n int max_displacement,\n int stride1,\n int stride2)\n {\n // n (batch size), c (num of channels), y (height), x (width)\n\n int n = item; \n int y = blockIdx.x * stride1 + pad_size;\n int x = blockIdx.y * stride1 + pad_size;\n int c = blockIdx.z;\n int tch_off = threadIdx.x;\n\n int kernel_rad = (kernel_size - 1) / 2;\n int displacement_rad = max_displacement / stride2;\n int displacement_size = 2 * displacement_rad + 1;\n\n int xmin = (x - kernel_rad - max_displacement) / stride1;\n int ymin = (y - kernel_rad - max_displacement) / stride1;\n\n int xmax = (x + kernel_rad - max_displacement) / stride1;\n int ymax = (y + kernel_rad - max_displacement) / stride1;\n\n if (xmax < 0 || ymax < 0 || xmin >= outputWidth || ymin >= outputHeight) {\n // assumes gradInput1 is pre-allocated and zero filled\n return;\n }\n\n if (xmin > xmax || ymin > ymax) {\n // assumes gradInput1 is pre-allocated and zero filled\n return;\n }\n\n xmin = max(0,xmin);\n xmax = min(outputWidth-1,xmax);\n\n ymin = max(0,ymin);\n ymax = min(outputHeight-1,ymax);\n\n int pInputWidth = inputWidth + 2 * pad_size;\n int pInputHeight = inputHeight + 2 * pad_size;\n\n int pdimyxc = pInputHeight * pInputWidth * nInputChannels;\n int pdimxc = pInputWidth * nInputChannels;\n int pdimc = nInputChannels;\n\n int tdimcyx = nOutputChannels * outputHeight * outputWidth;\n int tdimyx = outputHeight * outputWidth;\n int tdimx = outputWidth;\n\n int odimcyx = nInputChannels * inputHeight* inputWidth;\n int odimyx = inputHeight * inputWidth;\n int odimx = inputWidth;\n\n scalar_t nelems = kernel_size * kernel_size * nInputChannels;\n\n __shared__ scalar_t prod_sum[THREADS_PER_BLOCK];\n prod_sum[tch_off] = 0;\n\n for (int tc = tch_off; tc < nOutputChannels; tc += THREADS_PER_BLOCK) {\n\n int i2 = (tc % displacement_size - displacement_rad) * stride2;\n int j2 = (tc / displacement_size - displacement_rad) * stride2;\n\n int indx2 = n * pdimyxc + (y + j2)* pdimxc + (x + i2) * pdimc + c;\n \n scalar_t val2 = rInput2[indx2];\n\n for (int j = ymin; j <= ymax; ++j) {\n for (int i = xmin; i <= xmax; ++i) {\n int tindx = n * tdimcyx + tc * tdimyx + j * tdimx + i;\n prod_sum[tch_off] += gradOutput[tindx] * val2;\n }\n }\n }\n __syncthreads();\n\n if(tch_off == 0) {\n scalar_t reduce_sum = 0;\n for(int idx = 0; idx < THREADS_PER_BLOCK; idx++) {\n reduce_sum += prod_sum[idx];\n }\n const int indx1 = n * odimcyx + c * odimyx + (y - pad_size) * odimx + (x - pad_size);\n gradInput1[indx1] = reduce_sum / nelems;\n }\n\n}\n\ntemplate \n__global__ void correlation_backward_input2(int item, scalar_t* gradInput2, int nInputChannels, int inputHeight, int inputWidth,\n const scalar_t* __restrict__ gradOutput, int nOutputChannels, int outputHeight, int outputWidth,\n const scalar_t* __restrict__ rInput1,\n int pad_size,\n int kernel_size,\n int max_displacement,\n int stride1,\n int stride2)\n{\n // n (batch size), c (num of channels), y (height), x (width)\n\n int n = item;\n int y = blockIdx.x * stride1 + pad_size;\n int x = blockIdx.y * stride1 + pad_size;\n int c = blockIdx.z;\n\n int tch_off = threadIdx.x;\n\n int kernel_rad = (kernel_size - 1) / 2;\n int displacement_rad = max_displacement / stride2;\n int displacement_size = 2 * displacement_rad + 1;\n\n int pInputWidth = inputWidth + 2 * pad_size;\n int pInputHeight = inputHeight + 2 * pad_size;\n\n int pdimyxc = pInputHeight * pInputWidth * nInputChannels;\n int pdimxc = pInputWidth * nInputChannels;\n int pdimc = nInputChannels;\n\n int tdimcyx = nOutputChannels * outputHeight * outputWidth;\n int tdimyx = outputHeight * outputWidth;\n int tdimx = outputWidth;\n\n int odimcyx = nInputChannels * inputHeight* inputWidth;\n int odimyx = inputHeight * inputWidth;\n int odimx = inputWidth;\n\n scalar_t nelems = kernel_size * kernel_size * nInputChannels;\n\n __shared__ scalar_t prod_sum[THREADS_PER_BLOCK];\n prod_sum[tch_off] = 0;\n\n for (int tc = tch_off; tc < nOutputChannels; tc += THREADS_PER_BLOCK) {\n int i2 = (tc % displacement_size - displacement_rad) * stride2;\n int j2 = (tc / displacement_size - displacement_rad) * stride2;\n\n int xmin = (x - kernel_rad - max_displacement - i2) / stride1;\n int ymin = (y - kernel_rad - max_displacement - j2) / stride1;\n\n int xmax = (x + kernel_rad - max_displacement - i2) / stride1;\n int ymax = (y + kernel_rad - max_displacement - j2) / stride1;\n\n if (xmax < 0 || ymax < 0 || xmin >= outputWidth || ymin >= outputHeight) {\n // assumes gradInput2 is pre-allocated and zero filled\n continue;\n }\n\n if (xmin > xmax || ymin > ymax) {\n // assumes gradInput2 is pre-allocated and zero filled\n continue;\n }\n\n xmin = max(0,xmin);\n xmax = min(outputWidth-1,xmax);\n\n ymin = max(0,ymin);\n ymax = min(outputHeight-1,ymax);\n \n int indx1 = n * pdimyxc + (y - j2)* pdimxc + (x - i2) * pdimc + c;\n scalar_t val1 = rInput1[indx1];\n\n for (int j = ymin; j <= ymax; ++j) {\n for (int i = xmin; i <= xmax; ++i) {\n int tindx = n * tdimcyx + tc * tdimyx + j * tdimx + i;\n prod_sum[tch_off] += gradOutput[tindx] * val1;\n }\n }\n }\n\n __syncthreads();\n\n if(tch_off == 0) {\n scalar_t reduce_sum = 0;\n for(int idx = 0; idx < THREADS_PER_BLOCK; idx++) {\n reduce_sum += prod_sum[idx];\n }\n const int indx2 = n * odimcyx + c * odimyx + (y - pad_size) * odimx + (x - pad_size);\n gradInput2[indx2] = reduce_sum / nelems;\n }\n\n}\n\nint correlation_forward_cuda_kernel(at::Tensor& output,\n int ob,\n int oc,\n int oh,\n int ow,\n int osb,\n int osc,\n int osh,\n int osw,\n\n at::Tensor& input1,\n int ic,\n int ih,\n int iw,\n int isb,\n int isc,\n int ish,\n int isw,\n\n at::Tensor& input2,\n int gc,\n int gsb,\n int gsc,\n int gsh,\n int gsw,\n\n at::Tensor& rInput1,\n at::Tensor& rInput2,\n int pad_size,\n int kernel_size,\n int max_displacement,\n int stride1,\n int stride2,\n int corr_type_multiply,\n cudaStream_t stream) \n{\n\n int batchSize = ob;\n\n int nInputChannels = ic;\n int inputWidth = iw;\n int inputHeight = ih;\n\n int nOutputChannels = oc;\n int outputWidth = ow;\n int outputHeight = oh;\n\n dim3 blocks_grid(batchSize, inputHeight, inputWidth);\n dim3 threads_block(THREADS_PER_BLOCK);\n\n AT_DISPATCH_FLOATING_TYPES_AND_HALF(input1.type(), \"channels_first_fwd_1\", ([&] {\n\n channels_first<<>>(\n input1.data(), rInput1.data(), nInputChannels, inputHeight, inputWidth, pad_size);\n\n }));\n\n AT_DISPATCH_FLOATING_TYPES_AND_HALF(input2.type(), \"channels_first_fwd_2\", ([&] {\n\n channels_first<<>> (\n input2.data(), rInput2.data(), nInputChannels, inputHeight, inputWidth, pad_size);\n\n }));\n\n dim3 threadsPerBlock(THREADS_PER_BLOCK);\n dim3 totalBlocksCorr(batchSize, outputHeight, outputWidth);\n\n AT_DISPATCH_FLOATING_TYPES_AND_HALF(input1.type(), \"correlation_forward\", ([&] {\n\n correlation_forward<<>> \n (output.data(), nOutputChannels, outputHeight, outputWidth,\n rInput1.data(), nInputChannels, inputHeight, inputWidth,\n rInput2.data(),\n pad_size,\n kernel_size,\n max_displacement,\n stride1,\n stride2);\n\n }));\n\n cudaError_t err = cudaGetLastError();\n\n\n // check for errors\n if (err != cudaSuccess) {\n printf(\"error in correlation_forward_cuda_kernel: %s\\n\", cudaGetErrorString(err));\n return 0;\n }\n\n return 1;\n}\n\n\nint correlation_backward_cuda_kernel(\n at::Tensor& gradOutput,\n int gob,\n int goc,\n int goh,\n int gow,\n int gosb,\n int gosc,\n int gosh,\n int gosw,\n\n at::Tensor& input1,\n int ic,\n int ih,\n int iw,\n int isb,\n int isc,\n int ish,\n int isw,\n\n at::Tensor& input2,\n int gsb,\n int gsc,\n int gsh,\n int gsw,\n\n at::Tensor& gradInput1,\n int gisb,\n int gisc,\n int gish,\n int gisw,\n\n at::Tensor& gradInput2,\n int ggc,\n int ggsb,\n int ggsc,\n int ggsh,\n int ggsw,\n\n at::Tensor& rInput1,\n at::Tensor& rInput2,\n int pad_size,\n int kernel_size,\n int max_displacement,\n int stride1,\n int stride2,\n int corr_type_multiply,\n cudaStream_t stream)\n{\n\n int batchSize = gob;\n int num = batchSize;\n\n int nInputChannels = ic;\n int inputWidth = iw;\n int inputHeight = ih;\n\n int nOutputChannels = goc;\n int outputWidth = gow;\n int outputHeight = goh;\n\n dim3 blocks_grid(batchSize, inputHeight, inputWidth);\n dim3 threads_block(THREADS_PER_BLOCK);\n\n\n AT_DISPATCH_FLOATING_TYPES_AND_HALF(input1.type(), \"lltm_forward_cuda\", ([&] {\n\n channels_first<<>>(\n input1.data(),\n rInput1.data(),\n nInputChannels,\n inputHeight,\n inputWidth,\n pad_size\n );\n }));\n\n AT_DISPATCH_FLOATING_TYPES_AND_HALF(input2.type(), \"lltm_forward_cuda\", ([&] {\n\n channels_first<<>>(\n input2.data(),\n rInput2.data(),\n nInputChannels,\n inputHeight,\n inputWidth,\n pad_size\n );\n }));\n\n dim3 threadsPerBlock(THREADS_PER_BLOCK);\n dim3 totalBlocksCorr(inputHeight, inputWidth, nInputChannels);\n\n for (int n = 0; n < num; ++n) {\n\n AT_DISPATCH_FLOATING_TYPES_AND_HALF(input2.type(), \"lltm_forward_cuda\", ([&] {\n\n\n correlation_backward_input1<<>> (\n n, gradInput1.data(), nInputChannels, inputHeight, inputWidth,\n gradOutput.data(), nOutputChannels, outputHeight, outputWidth,\n rInput2.data(),\n pad_size,\n kernel_size,\n max_displacement,\n stride1,\n stride2);\n }));\n }\n\n for(int n = 0; n < batchSize; n++) {\n\n AT_DISPATCH_FLOATING_TYPES_AND_HALF(rInput1.type(), \"lltm_forward_cuda\", ([&] {\n\n correlation_backward_input2<<>>(\n n, gradInput2.data(), nInputChannels, inputHeight, inputWidth,\n gradOutput.data(), nOutputChannels, outputHeight, outputWidth,\n rInput1.data(),\n pad_size,\n kernel_size,\n max_displacement,\n stride1,\n stride2);\n\n }));\n }\n\n // check for errors\n cudaError_t err = cudaGetLastError();\n if (err != cudaSuccess) {\n printf(\"error in correlation_backward_cuda_kernel: %s\\n\", cudaGetErrorString(err));\n return 0;\n }\n\n return 1;\n}\n"
},
{
"path": "dvs/flownet2/networks/correlation_package/correlation_cuda_kernel.cuh",
"content": "#pragma once\n\n#include \n#include \n#include \n\nint correlation_forward_cuda_kernel(at::Tensor& output,\n int ob,\n int oc,\n int oh,\n int ow,\n int osb,\n int osc,\n int osh,\n int osw,\n\n at::Tensor& input1,\n int ic,\n int ih,\n int iw,\n int isb,\n int isc,\n int ish,\n int isw,\n\n at::Tensor& input2,\n int gc,\n int gsb,\n int gsc,\n int gsh,\n int gsw,\n\n at::Tensor& rInput1,\n at::Tensor& rInput2,\n int pad_size,\n int kernel_size,\n int max_displacement,\n int stride1,\n int stride2,\n int corr_type_multiply,\n cudaStream_t stream);\n\n\nint correlation_backward_cuda_kernel( \n at::Tensor& gradOutput,\n int gob,\n int goc,\n int goh,\n int gow,\n int gosb,\n int gosc,\n int gosh,\n int gosw,\n\n at::Tensor& input1,\n int ic,\n int ih,\n int iw,\n int isb,\n int isc,\n int ish,\n int isw,\n\n at::Tensor& input2,\n int gsb,\n int gsc,\n int gsh,\n int gsw,\n\n at::Tensor& gradInput1, \n int gisb,\n int gisc,\n int gish,\n int gisw,\n\n at::Tensor& gradInput2,\n int ggc,\n int ggsb,\n int ggsc,\n int ggsh,\n int ggsw,\n\n at::Tensor& rInput1,\n at::Tensor& rInput2,\n int pad_size,\n int kernel_size,\n int max_displacement,\n int stride1,\n int stride2,\n int corr_type_multiply,\n cudaStream_t stream);\n"
},
{
"path": "dvs/flownet2/networks/correlation_package/setup.py",
"content": "#!/usr/bin/env python3\nimport os\nimport torch\n\nfrom setuptools import setup, find_packages\nfrom torch.utils.cpp_extension import BuildExtension, CUDAExtension\n\ncxx_args = ['-std=c++11']\n\nnvcc_args = [\n '-gencode', 'arch=compute_50,code=sm_50',\n '-gencode', 'arch=compute_52,code=sm_52',\n '-gencode', 'arch=compute_60,code=sm_60',\n '-gencode', 'arch=compute_61,code=sm_61',\n '-gencode', 'arch=compute_70,code=sm_70',\n '-gencode', 'arch=compute_70,code=compute_70'\n]\n\nsetup(\n name='correlation_cuda',\n ext_modules=[\n CUDAExtension('correlation_cuda', [\n 'correlation_cuda.cc',\n 'correlation_cuda_kernel.cu'\n ], extra_compile_args={'cxx': cxx_args, 'nvcc': nvcc_args})\n ],\n cmdclass={\n 'build_ext': BuildExtension\n })\n"
},
{
"path": "dvs/flownet2/networks/resample2d_package/__init__.py",
"content": ""
},
{
"path": "dvs/flownet2/networks/resample2d_package/resample2d.py",
"content": "from torch.nn.modules.module import Module\nfrom torch.autograd import Function, Variable\nimport resample2d_cuda\n\nclass Resample2dFunction(Function):\n\n @staticmethod\n def forward(ctx, input1, input2, kernel_size=1, bilinear= True):\n assert input1.is_contiguous()\n assert input2.is_contiguous()\n\n ctx.save_for_backward(input1, input2)\n ctx.kernel_size = kernel_size\n ctx.bilinear = bilinear\n\n _, d, _, _ = input1.size()\n b, _, h, w = input2.size()\n output = input1.new(b, d, h, w).zero_()\n\n resample2d_cuda.forward(input1, input2, output, kernel_size, bilinear)\n\n return output\n\n @staticmethod\n def backward(ctx, grad_output):\n grad_output = grad_output.contiguous()\n assert grad_output.is_contiguous()\n\n input1, input2 = ctx.saved_tensors\n\n grad_input1 = Variable(input1.new(input1.size()).zero_())\n grad_input2 = Variable(input1.new(input2.size()).zero_())\n\n resample2d_cuda.backward(input1, input2, grad_output.data,\n grad_input1.data, grad_input2.data,\n ctx.kernel_size, ctx.bilinear)\n\n return grad_input1, grad_input2, None, None\n\nclass Resample2d(Module):\n\n def __init__(self, kernel_size=1, bilinear = True):\n super(Resample2d, self).__init__()\n self.kernel_size = kernel_size\n self.bilinear = bilinear\n\n def forward(self, input1, input2):\n input1_c = input1.contiguous()\n return Resample2dFunction.apply(input1_c, input2, self.kernel_size, self.bilinear)\n"
},
{
"path": "dvs/flownet2/networks/resample2d_package/resample2d_cuda.cc",
"content": "#include \n#include \n\n#include \"resample2d_kernel.cuh\"\n\nint resample2d_cuda_forward(\n at::Tensor& input1,\n at::Tensor& input2, \n at::Tensor& output,\n int kernel_size, bool bilinear) {\n resample2d_kernel_forward(input1, input2, output, kernel_size, bilinear);\n return 1;\n}\n\nint resample2d_cuda_backward(\n at::Tensor& input1, \n at::Tensor& input2,\n at::Tensor& gradOutput,\n at::Tensor& gradInput1, \n at::Tensor& gradInput2, \n int kernel_size, bool bilinear) {\n resample2d_kernel_backward(input1, input2, gradOutput, gradInput1, gradInput2, kernel_size, bilinear);\n return 1;\n}\n\n\n\nPYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {\n m.def(\"forward\", &resample2d_cuda_forward, \"Resample2D forward (CUDA)\");\n m.def(\"backward\", &resample2d_cuda_backward, \"Resample2D backward (CUDA)\");\n}\n\n"
},
{
"path": "dvs/flownet2/networks/resample2d_package/resample2d_kernel.cu",
"content": "#include \n#include \n#include \n\n#define CUDA_NUM_THREADS 512 \n#define THREADS_PER_BLOCK 64 \n\n#define DIM0(TENSOR) ((TENSOR).x)\n#define DIM1(TENSOR) ((TENSOR).y)\n#define DIM2(TENSOR) ((TENSOR).z)\n#define DIM3(TENSOR) ((TENSOR).w)\n\n#define DIM3_INDEX(TENSOR, xx, yy, zz, ww) ((TENSOR)[((xx) * (TENSOR##_stride.x)) + ((yy) * (TENSOR##_stride.y)) + ((zz) * (TENSOR##_stride.z)) + ((ww) * (TENSOR##_stride.w))])\n\ntemplate \n__global__ void kernel_resample2d_update_output(const int n, \n const scalar_t* __restrict__ input1, const long4 input1_size, const long4 input1_stride,\n const scalar_t* __restrict__ input2, const long4 input2_size, const long4 input2_stride, \n scalar_t* __restrict__ output, const long4 output_size, const long4 output_stride, int kernel_size, bool bilinear) {\n int index = blockIdx.x * blockDim.x + threadIdx.x;\n\n if (index >= n) {\n return;\n }\n\n scalar_t val = 0.0f;\n\n int dim_b = DIM0(output_size);\n int dim_c = DIM1(output_size);\n int dim_h = DIM2(output_size);\n int dim_w = DIM3(output_size);\n int dim_chw = dim_c * dim_h * dim_w;\n int dim_hw = dim_h * dim_w;\n\n int b = ( index / dim_chw ) % dim_b;\n int c = ( index / dim_hw ) % dim_c;\n int y = ( index / dim_w ) % dim_h;\n int x = ( index ) % dim_w;\n\n scalar_t dx = DIM3_INDEX(input2, b, 0, y, x);\n scalar_t dy = DIM3_INDEX(input2, b, 1, y, x);\n\n scalar_t xf = static_cast(x) + dx;\n scalar_t yf = static_cast(y) + dy;\n scalar_t alpha = xf - floor(xf); // alpha\n scalar_t beta = yf - floor(yf); // beta\n\n if (bilinear) {\n int xL = max(min( int (floor(xf)), dim_w-1), 0);\n int xR = max(min( int (floor(xf)+1), dim_w -1), 0);\n int yT = max(min( int (floor(yf)), dim_h-1), 0);\n int yB = max(min( int (floor(yf)+1), dim_h-1), 0);\n\n for (int fy = 0; fy < kernel_size; fy += 1) {\n for (int fx = 0; fx < kernel_size; fx += 1) {\n val += static_cast((1. - alpha)*(1. - beta) * DIM3_INDEX(input1, b, c, yT + fy, xL + fx));\n val += static_cast((alpha)*(1. - beta) * DIM3_INDEX(input1, b, c, yT + fy, xR + fx));\n val += static_cast((1. - alpha)*(beta) * DIM3_INDEX(input1, b, c, yB + fy, xL + fx));\n val += static_cast((alpha)*(beta) * DIM3_INDEX(input1, b, c, yB + fy, xR + fx));\n }\n }\n\n output[index] = val;\n }\n else {\n int xN = max(min( int (floor(xf + 0.5)), dim_w - 1), 0);\n int yN = max(min( int (floor(yf + 0.5)), dim_h - 1), 0);\n\n output[index] = static_cast ( DIM3_INDEX(input1, b, c, yN, xN) );\n }\n\n}\n\n\ntemplate \n__global__ void kernel_resample2d_backward_input1(\n const int n, const scalar_t* __restrict__ input1, const long4 input1_size, const long4 input1_stride,\n const scalar_t* __restrict__ input2, const long4 input2_size, const long4 input2_stride,\n const scalar_t* __restrict__ gradOutput, const long4 gradOutput_size, const long4 gradOutput_stride,\n scalar_t* __restrict__ gradInput, const long4 gradInput_size, const long4 gradInput_stride, int kernel_size, bool bilinear) {\n\n int index = blockIdx.x * blockDim.x + threadIdx.x;\n\n if (index >= n) {\n return;\n }\n\n int dim_b = DIM0(gradOutput_size);\n int dim_c = DIM1(gradOutput_size);\n int dim_h = DIM2(gradOutput_size);\n int dim_w = DIM3(gradOutput_size);\n int dim_chw = dim_c * dim_h * dim_w;\n int dim_hw = dim_h * dim_w;\n\n int b = ( index / dim_chw ) % dim_b;\n int c = ( index / dim_hw ) % dim_c;\n int y = ( index / dim_w ) % dim_h;\n int x = ( index ) % dim_w;\n\n scalar_t dx = DIM3_INDEX(input2, b, 0, y, x);\n scalar_t dy = DIM3_INDEX(input2, b, 1, y, x);\n\n scalar_t xf = static_cast(x) + dx;\n scalar_t yf = static_cast(y) + dy;\n scalar_t alpha = xf - int(xf); // alpha\n scalar_t beta = yf - int(yf); // beta\n\n int idim_h = DIM2(input1_size);\n int idim_w = DIM3(input1_size);\n\n int xL = max(min( int (floor(xf)), idim_w-1), 0);\n int xR = max(min( int (floor(xf)+1), idim_w -1), 0);\n int yT = max(min( int (floor(yf)), idim_h-1), 0);\n int yB = max(min( int (floor(yf)+1), idim_h-1), 0);\n\n for (int fy = 0; fy < kernel_size; fy += 1) {\n for (int fx = 0; fx < kernel_size; fx += 1) {\n atomicAdd(&DIM3_INDEX(gradInput, b, c, (yT + fy), (xL + fx)), (1-alpha)*(1-beta) * DIM3_INDEX(gradOutput, b, c, y, x));\n atomicAdd(&DIM3_INDEX(gradInput, b, c, (yT + fy), (xR + fx)), (alpha)*(1-beta) * DIM3_INDEX(gradOutput, b, c, y, x));\n atomicAdd(&DIM3_INDEX(gradInput, b, c, (yB + fy), (xL + fx)), (1-alpha)*(beta) * DIM3_INDEX(gradOutput, b, c, y, x));\n atomicAdd(&DIM3_INDEX(gradInput, b, c, (yB + fy), (xR + fx)), (alpha)*(beta) * DIM3_INDEX(gradOutput, b, c, y, x));\n }\n }\n\n}\n\ntemplate \n__global__ void kernel_resample2d_backward_input2(\n const int n, const scalar_t* __restrict__ input1, const long4 input1_size, const long4 input1_stride,\n const scalar_t* __restrict__ input2, const long4 input2_size, const long4 input2_stride,\n const scalar_t* __restrict__ gradOutput, const long4 gradOutput_size, const long4 gradOutput_stride,\n scalar_t* __restrict__ gradInput, const long4 gradInput_size, const long4 gradInput_stride, int kernel_size, bool bilinear) {\n\n int index = blockIdx.x * blockDim.x + threadIdx.x;\n\n if (index >= n) {\n return;\n }\n\n scalar_t output = 0.0;\n int kernel_rad = (kernel_size - 1)/2;\n\n int dim_b = DIM0(gradInput_size);\n int dim_c = DIM1(gradInput_size);\n int dim_h = DIM2(gradInput_size);\n int dim_w = DIM3(gradInput_size);\n int dim_chw = dim_c * dim_h * dim_w;\n int dim_hw = dim_h * dim_w;\n\n int b = ( index / dim_chw ) % dim_b;\n int c = ( index / dim_hw ) % dim_c;\n int y = ( index / dim_w ) % dim_h;\n int x = ( index ) % dim_w;\n\n int odim_c = DIM1(gradOutput_size);\n\n scalar_t dx = DIM3_INDEX(input2, b, 0, y, x);\n scalar_t dy = DIM3_INDEX(input2, b, 1, y, x);\n\n scalar_t xf = static_cast(x) + dx;\n scalar_t yf = static_cast(y) + dy;\n\n int xL = max(min( int (floor(xf)), dim_w-1), 0);\n int xR = max(min( int (floor(xf)+1), dim_w -1), 0);\n int yT = max(min( int (floor(yf)), dim_h-1), 0);\n int yB = max(min( int (floor(yf)+1), dim_h-1), 0);\n \n if (c % 2) {\n float gamma = 1 - (xf - floor(xf)); // alpha\n for (int i = 0; i <= 2*kernel_rad; ++i) {\n for (int j = 0; j <= 2*kernel_rad; ++j) {\n for (int ch = 0; ch < odim_c; ++ch) {\n output += (gamma) * DIM3_INDEX(gradOutput, b, ch, y, x) * DIM3_INDEX(input1, b, ch, (yB + j), (xL + i));\n output -= (gamma) * DIM3_INDEX(gradOutput, b, ch, y, x) * DIM3_INDEX(input1, b, ch, (yT + j), (xL + i));\n output += (1-gamma) * DIM3_INDEX(gradOutput, b, ch, y, x) * DIM3_INDEX(input1, b, ch, (yB + j), (xR + i));\n output -= (1-gamma) * DIM3_INDEX(gradOutput, b, ch, y, x) * DIM3_INDEX(input1, b, ch, (yT + j), (xR + i));\n }\n }\n }\n }\n else {\n float gamma = 1 - (yf - floor(yf)); // alpha\n for (int i = 0; i <= 2*kernel_rad; ++i) {\n for (int j = 0; j <= 2*kernel_rad; ++j) {\n for (int ch = 0; ch < odim_c; ++ch) {\n output += (gamma) * DIM3_INDEX(gradOutput, b, ch, y, x) * DIM3_INDEX(input1, b, ch, (yT + j), (xR + i));\n output -= (gamma) * DIM3_INDEX(gradOutput, b, ch, y, x) * DIM3_INDEX(input1, b, ch, (yT + j), (xL + i));\n output += (1-gamma) * DIM3_INDEX(gradOutput, b, ch, y, x) * DIM3_INDEX(input1, b, ch, (yB + j), (xR + i));\n output -= (1-gamma) * DIM3_INDEX(gradOutput, b, ch, y, x) * DIM3_INDEX(input1, b, ch, (yB + j), (xL + i));\n }\n }\n }\n\n }\n\n gradInput[index] = output;\n\n}\n\nvoid resample2d_kernel_forward(\n at::Tensor& input1, \n at::Tensor& input2,\n at::Tensor& output, \n int kernel_size,\n bool bilinear) {\n\n int n = output.numel();\n\n const long4 input1_size = make_long4(input1.size(0), input1.size(1), input1.size(2), input1.size(3));\n const long4 input1_stride = make_long4(input1.stride(0), input1.stride(1), input1.stride(2), input1.stride(3));\n\n const long4 input2_size = make_long4(input2.size(0), input2.size(1), input2.size(2), input2.size(3));\n const long4 input2_stride = make_long4(input2.stride(0), input2.stride(1), input2.stride(2), input2.stride(3));\n\n const long4 output_size = make_long4(output.size(0), output.size(1), output.size(2), output.size(3));\n const long4 output_stride = make_long4(output.stride(0), output.stride(1), output.stride(2), output.stride(3));\n\n // TODO: when atomicAdd gets resolved, change to AT_DISPATCH_FLOATING_TYPES_AND_HALF\n// AT_DISPATCH_FLOATING_TYPES(input1.type(), \"resample_forward_kernel\", ([&] {\n\n kernel_resample2d_update_output<<< (n + CUDA_NUM_THREADS - 1)/CUDA_NUM_THREADS, CUDA_NUM_THREADS, 0, at::cuda::getCurrentCUDAStream() >>>(\n//at::globalContext().getCurrentCUDAStream() >>>(\n n,\n input1.data(),\n input1_size,\n input1_stride, \n input2.data(),\n input2_size,\n input2_stride,\n output.data(),\n output_size,\n output_stride,\n kernel_size,\n bilinear);\n\n// }));\n\n // TODO: ATen-equivalent check\n\n // THCudaCheck(cudaGetLastError());\n\n}\n\nvoid resample2d_kernel_backward(\n at::Tensor& input1,\n at::Tensor& input2,\n at::Tensor& gradOutput,\n at::Tensor& gradInput1,\n at::Tensor& gradInput2,\n int kernel_size,\n bool bilinear) {\n\n int n = gradOutput.numel();\n\n const long4 input1_size = make_long4(input1.size(0), input1.size(1), input1.size(2), input1.size(3));\n const long4 input1_stride = make_long4(input1.stride(0), input1.stride(1), input1.stride(2), input1.stride(3));\n\n const long4 input2_size = make_long4(input2.size(0), input2.size(1), input2.size(2), input2.size(3));\n const long4 input2_stride = make_long4(input2.stride(0), input2.stride(1), input2.stride(2), input2.stride(3));\n\n const long4 gradOutput_size = make_long4(gradOutput.size(0), gradOutput.size(1), gradOutput.size(2), gradOutput.size(3));\n const long4 gradOutput_stride = make_long4(gradOutput.stride(0), gradOutput.stride(1), gradOutput.stride(2), gradOutput.stride(3));\n\n const long4 gradInput1_size = make_long4(gradInput1.size(0), gradInput1.size(1), gradInput1.size(2), gradInput1.size(3));\n const long4 gradInput1_stride = make_long4(gradInput1.stride(0), gradInput1.stride(1), gradInput1.stride(2), gradInput1.stride(3));\n\n// AT_DISPATCH_FLOATING_TYPES(input1.type(), \"resample_backward_input1\", ([&] {\n\n kernel_resample2d_backward_input1<<< (n + CUDA_NUM_THREADS - 1)/CUDA_NUM_THREADS, CUDA_NUM_THREADS, 0, at::cuda::getCurrentCUDAStream() >>>(\n//at::globalContext().getCurrentCUDAStream() >>>(\n n, \n input1.data(), \n input1_size,\n input1_stride,\n input2.data(),\n input2_size, \n input2_stride,\n gradOutput.data(),\n gradOutput_size,\n gradOutput_stride,\n gradInput1.data(),\n gradInput1_size,\n gradInput1_stride, \n kernel_size,\n bilinear\n );\n\n// }));\n\n const long4 gradInput2_size = make_long4(gradInput2.size(0), gradInput2.size(1), gradInput2.size(2), gradInput2.size(3));\n const long4 gradInput2_stride = make_long4(gradInput2.stride(0), gradInput2.stride(1), gradInput2.stride(2), gradInput2.stride(3));\n\n n = gradInput2.numel();\n\n// AT_DISPATCH_FLOATING_TYPES(gradInput2.type(), \"resample_backward_input2\", ([&] {\n\n\n kernel_resample2d_backward_input2<<< (n + CUDA_NUM_THREADS - 1)/CUDA_NUM_THREADS, CUDA_NUM_THREADS, 0, at::cuda::getCurrentCUDAStream() >>>(\n//at::globalContext().getCurrentCUDAStream() >>>(\n n, \n input1.data(), \n input1_size, \n input1_stride,\n input2.data(), \n input2_size,\n input2_stride,\n gradOutput.data(),\n gradOutput_size,\n gradOutput_stride,\n gradInput2.data(),\n gradInput2_size,\n gradInput2_stride,\n kernel_size,\n bilinear\n );\n\n// }));\n\n // TODO: Use the ATen equivalent to get last error\n\n // THCudaCheck(cudaGetLastError());\n\n}\n"
},
{
"path": "dvs/flownet2/networks/resample2d_package/resample2d_kernel.cuh",
"content": "#pragma once\n\n#include \n\nvoid resample2d_kernel_forward(\n at::Tensor& input1,\n at::Tensor& input2,\n at::Tensor& output,\n int kernel_size,\n bool bilinear);\n\nvoid resample2d_kernel_backward(\n at::Tensor& input1,\n at::Tensor& input2,\n at::Tensor& gradOutput,\n at::Tensor& gradInput1, \n at::Tensor& gradInput2, \n int kernel_size,\n bool bilinear);"
},
{
"path": "dvs/flownet2/networks/resample2d_package/setup.py",
"content": "#!/usr/bin/env python3\nimport os\nimport torch\n\nfrom setuptools import setup\nfrom torch.utils.cpp_extension import BuildExtension, CUDAExtension\n\ncxx_args = ['-std=c++11']\n\nnvcc_args = [\n '-gencode', 'arch=compute_50,code=sm_50',\n '-gencode', 'arch=compute_52,code=sm_52',\n '-gencode', 'arch=compute_60,code=sm_60',\n '-gencode', 'arch=compute_61,code=sm_61',\n '-gencode', 'arch=compute_70,code=sm_70',\n '-gencode', 'arch=compute_70,code=compute_70'\n]\n\nsetup(\n name='resample2d_cuda',\n ext_modules=[\n CUDAExtension('resample2d_cuda', [\n 'resample2d_cuda.cc',\n 'resample2d_kernel.cu'\n ], extra_compile_args={'cxx': cxx_args, 'nvcc': nvcc_args})\n ],\n cmdclass={\n 'build_ext': BuildExtension\n })\n"
},
{
"path": "dvs/flownet2/networks/submodules.py",
"content": "# freda (todo) : \r\n\r\nimport torch.nn as nn\r\nimport torch\r\nimport numpy as np \r\n\r\ndef conv(batchNorm, in_planes, out_planes, kernel_size=3, stride=1):\r\n if batchNorm:\r\n return nn.Sequential(\r\n nn.Conv2d(in_planes, out_planes, kernel_size=kernel_size, stride=stride, padding=(kernel_size-1)//2, bias=False),\r\n nn.BatchNorm2d(out_planes),\r\n nn.LeakyReLU(0.1,inplace=True)\r\n )\r\n else:\r\n return nn.Sequential(\r\n nn.Conv2d(in_planes, out_planes, kernel_size=kernel_size, stride=stride, padding=(kernel_size-1)//2, bias=True),\r\n nn.LeakyReLU(0.1,inplace=True)\r\n )\r\n\r\ndef i_conv(batchNorm, in_planes, out_planes, kernel_size=3, stride=1, bias = True):\r\n if batchNorm:\r\n return nn.Sequential(\r\n nn.Conv2d(in_planes, out_planes, kernel_size=kernel_size, stride=stride, padding=(kernel_size-1)//2, bias=bias),\r\n nn.BatchNorm2d(out_planes),\r\n )\r\n else:\r\n return nn.Sequential(\r\n nn.Conv2d(in_planes, out_planes, kernel_size=kernel_size, stride=stride, padding=(kernel_size-1)//2, bias=bias),\r\n )\r\n\r\ndef predict_flow(in_planes):\r\n return nn.Conv2d(in_planes,2,kernel_size=3,stride=1,padding=1,bias=True)\r\n\r\ndef deconv(in_planes, out_planes):\r\n return nn.Sequential(\r\n nn.ConvTranspose2d(in_planes, out_planes, kernel_size=4, stride=2, padding=1, bias=True),\r\n nn.LeakyReLU(0.1,inplace=True)\r\n )\r\n\r\nclass tofp16(nn.Module):\r\n def __init__(self):\r\n super(tofp16, self).__init__()\r\n\r\n def forward(self, input):\r\n return input.half()\r\n\r\n\r\nclass tofp32(nn.Module):\r\n def __init__(self):\r\n super(tofp32, self).__init__()\r\n\r\n def forward(self, input):\r\n return input.float()\r\n\r\n\r\ndef init_deconv_bilinear(weight):\r\n f_shape = weight.size()\r\n heigh, width = f_shape[-2], f_shape[-1]\r\n f = np.ceil(width/2.0)\r\n c = (2 * f - 1 - f % 2) / (2.0 * f)\r\n bilinear = np.zeros([heigh, width])\r\n for x in range(width):\r\n for y in range(heigh):\r\n value = (1 - abs(x / f - c)) * (1 - abs(y / f - c))\r\n bilinear[x, y] = value\r\n weight.data.fill_(0.)\r\n for i in range(f_shape[0]):\r\n for j in range(f_shape[1]):\r\n weight.data[i,j,:,:] = torch.from_numpy(bilinear)\r\n\r\n\r\ndef save_grad(grads, name):\r\n def hook(grad):\r\n grads[name] = grad\r\n return hook\r\n\r\n'''\r\ndef save_grad(grads, name):\r\n def hook(grad):\r\n grads[name] = grad\r\n return hook\r\nimport torch\r\nfrom channelnorm_package.modules.channelnorm import ChannelNorm \r\nmodel = ChannelNorm().cuda()\r\ngrads = {}\r\na = 100*torch.autograd.Variable(torch.randn((1,3,5,5)).cuda(), requires_grad=True)\r\na.register_hook(save_grad(grads, 'a'))\r\nb = model(a)\r\ny = torch.mean(b)\r\ny.backward()\r\n\r\n'''\r\n"
},
{
"path": "dvs/flownet2/run.sh",
"content": "#!/bin/bash\npython main.py --inference --model FlowNet2 --save_flow --inference_dataset Google \\\n\t--inference_dataset_root ./../video \\\n\t--resume ./FlowNet2_checkpoint.pth.tar \\\n\t--inference_visualize\n"
},
{
"path": "dvs/flownet2/run_release.sh",
"content": "#!/bin/bash\npython main.py --inference --model FlowNet2 --save_flow --inference_dataset Google \\\n\t--inference_dataset_root ./../dataset_release/test \\\n\t--resume ./FlowNet2_checkpoint.pth.tar \\\n\t--inference_visualize\n\npython main.py --inference --model FlowNet2 --save_flow --inference_dataset Google \\\n\t--inference_dataset_root ./../dataset_release/training \\\n\t--resume ./FlowNet2_checkpoint.pth.tar \\\n\t--inference_visualize"
},
{
"path": "dvs/flownet2/utils/__init__.py",
"content": ""
},
{
"path": "dvs/flownet2/utils/flow_utils.py",
"content": "import numpy as np\nimport matplotlib.pyplot as plt\nimport os.path\n\nTAG_CHAR = np.array([202021.25], np.float32)\n\ndef readFlow(fn):\n \"\"\" Read .flo file in Middlebury format\"\"\"\n # Code adapted from:\n # http://stackoverflow.com/questions/28013200/reading-middlebury-flow-files-with-python-bytes-array-numpy\n\n # WARNING: this will work on little-endian architectures (eg Intel x86) only!\n # print 'fn = %s'%(fn)\n with open(fn, 'rb') as f:\n magic = np.fromfile(f, np.float32, count=1)\n if 202021.25 != magic:\n print('Magic number incorrect. Invalid .flo file')\n return None\n else:\n w = np.fromfile(f, np.int32, count=1)\n h = np.fromfile(f, np.int32, count=1)\n # print 'Reading %d x %d flo file\\n' % (w, h)\n data = np.fromfile(f, np.float32, count=2*int(w)*int(h))\n # Reshape data into 3D array (columns, rows, bands)\n # The reshape here is for visualization, the original code is (w,h,2)\n return np.resize(data, (int(h), int(w), 2))\n\ndef writeFlow(filename,uv,v=None):\n \"\"\" Write optical flow to file.\n \n If v is None, uv is assumed to contain both u and v channels,\n stacked in depth.\n Original code by Deqing Sun, adapted from Daniel Scharstein.\n \"\"\"\n nBands = 2\n\n if v is None:\n assert(uv.ndim == 3)\n assert(uv.shape[2] == 2)\n u = uv[:,:,0]\n v = uv[:,:,1]\n else:\n u = uv\n\n assert(u.shape == v.shape)\n height,width = u.shape\n f = open(filename,'wb')\n # write the header\n f.write(TAG_CHAR)\n np.array(width).astype(np.int32).tofile(f)\n np.array(height).astype(np.int32).tofile(f)\n # arrange into matrix form\n tmp = np.zeros((height, width*nBands))\n tmp[:,np.arange(width)*2] = u\n tmp[:,np.arange(width)*2 + 1] = v\n tmp.astype(np.float32).tofile(f)\n f.close()\n\n\n# ref: https://github.com/sampepose/flownet2-tf/\n# blob/18f87081db44939414fc4a48834f9e0da3e69f4c/src/flowlib.py#L240\ndef visulize_flow_file(flow_filename, save_dir=None):\n\tflow_data = readFlow(flow_filename)\n\timg = flow2img(flow_data)\n\t# plt.imshow(img)\n\t# plt.show()\n\tif save_dir:\n\t\tidx = flow_filename.rfind(\"/\") + 1\n\t\tplt.imsave(os.path.join(save_dir, \"%s-vis.png\" % flow_filename[idx:-4]), img)\n\n\ndef flow2img(flow_data):\n\t\"\"\"\n\tconvert optical flow into color image\n\t:param flow_data:\n\t:return: color image\n\t\"\"\"\n\t# print(flow_data.shape)\n\t# print(type(flow_data))\n\tu = flow_data[:, :, 0]\n\tv = flow_data[:, :, 1]\n\n\tUNKNOW_FLOW_THRESHOLD = 1e7\n\tpr1 = abs(u) > UNKNOW_FLOW_THRESHOLD\n\tpr2 = abs(v) > UNKNOW_FLOW_THRESHOLD\n\tidx_unknown = (pr1 | pr2)\n\tu[idx_unknown] = v[idx_unknown] = 0\n\n\t# get max value in each direction\n\tmaxu = -999.\n\tmaxv = -999.\n\tminu = 999.\n\tminv = 999.\n\tmaxu = max(maxu, np.max(u))\n\tmaxv = max(maxv, np.max(v))\n\tminu = min(minu, np.min(u))\n\tminv = min(minv, np.min(v))\n\n\trad = np.sqrt(u ** 2 + v ** 2)\n\tmaxrad = max(-1, np.max(rad))\n\tu = u / maxrad + np.finfo(float).eps\n\tv = v / maxrad + np.finfo(float).eps\n\n\timg = compute_color(u, v)\n\n\tidx = np.repeat(idx_unknown[:, :, np.newaxis], 3, axis=2)\n\timg[idx] = 0\n\n\treturn np.uint8(img)\n\n\ndef compute_color(u, v):\n\t\"\"\"\n\tcompute optical flow color map\n\t:param u: horizontal optical flow\n\t:param v: vertical optical flow\n\t:return:\n\t\"\"\"\n\n\theight, width = u.shape\n\timg = np.zeros((height, width, 3))\n\n\tNAN_idx = np.isnan(u) | np.isnan(v)\n\tu[NAN_idx] = v[NAN_idx] = 0\n\n\tcolorwheel = make_color_wheel()\n\tncols = np.size(colorwheel, 0)\n\n\trad = np.sqrt(u ** 2 + v ** 2)\n\n\ta = np.arctan2(-v, -u) / np.pi\n\n\tfk = (a + 1) / 2 * (ncols - 1) + 1\n\n\tk0 = np.floor(fk).astype(int)\n\n\tk1 = k0 + 1\n\tk1[k1 == ncols + 1] = 1\n\tf = fk - k0\n\n\tfor i in range(0, np.size(colorwheel, 1)):\n\t\ttmp = colorwheel[:, i]\n\t\tcol0 = tmp[k0 - 1] / 255\n\t\tcol1 = tmp[k1 - 1] / 255\n\t\tcol = (1 - f) * col0 + f * col1\n\n\t\tidx = rad <= 1\n\t\tcol[idx] = 1 - rad[idx] * (1 - col[idx])\n\t\tnotidx = np.logical_not(idx)\n\n\t\tcol[notidx] *= 0.75\n\t\timg[:, :, i] = np.uint8(np.floor(255 * col * (1 - NAN_idx)))\n\n\treturn img\n\n\ndef make_color_wheel():\n\t\"\"\"\n\tGenerate color wheel according Middlebury color code\n\t:return: Color wheel\n\t\"\"\"\n\tRY = 15\n\tYG = 6\n\tGC = 4\n\tCB = 11\n\tBM = 13\n\tMR = 6\n\n\tncols = RY + YG + GC + CB + BM + MR\n\n\tcolorwheel = np.zeros([ncols, 3])\n\n\tcol = 0\n\n\t# RY\n\tcolorwheel[0:RY, 0] = 255\n\tcolorwheel[0:RY, 1] = np.transpose(np.floor(255 * np.arange(0, RY) / RY))\n\tcol += RY\n\n\t# YG\n\tcolorwheel[col:col + YG, 0] = 255 - np.transpose(np.floor(255 * np.arange(0, YG) / YG))\n\tcolorwheel[col:col + YG, 1] = 255\n\tcol += YG\n\n\t# GC\n\tcolorwheel[col:col + GC, 1] = 255\n\tcolorwheel[col:col + GC, 2] = np.transpose(np.floor(255 * np.arange(0, GC) / GC))\n\tcol += GC\n\n\t# CB\n\tcolorwheel[col:col + CB, 1] = 255 - np.transpose(np.floor(255 * np.arange(0, CB) / CB))\n\tcolorwheel[col:col + CB, 2] = 255\n\tcol += CB\n\n\t# BM\n\tcolorwheel[col:col + BM, 2] = 255\n\tcolorwheel[col:col + BM, 0] = np.transpose(np.floor(255 * np.arange(0, BM) / BM))\n\tcol += + BM\n\n\t# MR\n\tcolorwheel[col:col + MR, 2] = 255 - np.transpose(np.floor(255 * np.arange(0, MR) / MR))\n\tcolorwheel[col:col + MR, 0] = 255\n\n\treturn colorwheel\n"
},
{
"path": "dvs/flownet2/utils/frame_utils.py",
"content": "import numpy as np\nfrom os.path import *\nfrom imageio import imread\nfrom . import flow_utils \n\ndef read_gen(file_name):\n ext = splitext(file_name)[-1]\n if ext == '.png' or ext == '.jpeg' or ext == '.ppm' or ext == '.jpg':\n im = imread(file_name)\n if im.shape[2] > 3:\n return im[:,:,:3]\n else:\n return im\n elif ext == '.bin' or ext == '.raw':\n return np.load(file_name)\n elif ext == '.flo':\n return flow_utils.readFlow(file_name).astype(np.float32)\n return []\n"
},
{
"path": "dvs/flownet2/utils/param_utils.py",
"content": "import torch\nimport torch.nn as nn\nimport numpy as np\n\ndef parse_flownetc(modules, weights, biases):\n keys = [\n 'conv1',\n 'conv2',\n 'conv3',\n 'conv_redir',\n 'conv3_1',\n 'conv4',\n 'conv4_1',\n 'conv5',\n 'conv5_1',\n 'conv6',\n 'conv6_1',\n \n 'deconv5',\n 'deconv4',\n 'deconv3',\n 'deconv2',\n \n 'Convolution1',\n 'Convolution2',\n 'Convolution3',\n 'Convolution4',\n 'Convolution5',\n\n 'upsample_flow6to5',\n 'upsample_flow5to4',\n 'upsample_flow4to3',\n 'upsample_flow3to2',\n \n ]\n i = 0\n for m in modules:\n if isinstance(m, nn.Conv2d) or isinstance(m, nn.ConvTranspose2d):\n weight = weights[keys[i]].copy()\n bias = biases[keys[i]].copy()\n if keys[i] == 'conv1':\n m.weight.data[:,:,:,:] = torch.from_numpy(np.flip(weight, axis=1).copy())\n m.bias.data[:] = torch.from_numpy(bias)\n else:\n m.weight.data[:,:,:,:] = torch.from_numpy(weight)\n m.bias.data[:] = torch.from_numpy(bias) \n\n i = i + 1\n return\n\ndef parse_flownets(modules, weights, biases, param_prefix='net2_'):\n keys = [\n 'conv1',\n 'conv2',\n 'conv3',\n 'conv3_1',\n 'conv4',\n 'conv4_1',\n 'conv5',\n 'conv5_1',\n 'conv6',\n 'conv6_1',\n \n 'deconv5',\n 'deconv4',\n 'deconv3',\n 'deconv2',\n \n 'predict_conv6',\n 'predict_conv5',\n 'predict_conv4',\n 'predict_conv3',\n 'predict_conv2',\n\n 'upsample_flow6to5',\n 'upsample_flow5to4',\n 'upsample_flow4to3',\n 'upsample_flow3to2',\n ]\n for i, k in enumerate(keys):\n if 'upsample' in k:\n keys[i] = param_prefix + param_prefix + k\n else:\n keys[i] = param_prefix + k\n i = 0\n for m in modules:\n if isinstance(m, nn.Conv2d) or isinstance(m, nn.ConvTranspose2d):\n weight = weights[keys[i]].copy()\n bias = biases[keys[i]].copy()\n if keys[i] == param_prefix+'conv1':\n m.weight.data[:,0:3,:,:] = torch.from_numpy(np.flip(weight[:,0:3,:,:], axis=1).copy())\n m.weight.data[:,3:6,:,:] = torch.from_numpy(np.flip(weight[:,3:6,:,:], axis=1).copy())\n m.weight.data[:,6:9,:,:] = torch.from_numpy(np.flip(weight[:,6:9,:,:], axis=1).copy())\n m.weight.data[:,9::,:,:] = torch.from_numpy(weight[:,9:,:,:].copy())\n if m.bias is not None:\n m.bias.data[:] = torch.from_numpy(bias)\n else:\n m.weight.data[:,:,:,:] = torch.from_numpy(weight)\n if m.bias is not None:\n m.bias.data[:] = torch.from_numpy(bias)\n i = i + 1\n return\n\ndef parse_flownetsonly(modules, weights, biases, param_prefix=''):\n keys = [\n 'conv1',\n 'conv2',\n 'conv3',\n 'conv3_1',\n 'conv4',\n 'conv4_1',\n 'conv5',\n 'conv5_1',\n 'conv6',\n 'conv6_1',\n \n 'deconv5',\n 'deconv4',\n 'deconv3',\n 'deconv2',\n \n 'Convolution1',\n 'Convolution2',\n 'Convolution3',\n 'Convolution4',\n 'Convolution5',\n\n 'upsample_flow6to5',\n 'upsample_flow5to4',\n 'upsample_flow4to3',\n 'upsample_flow3to2',\n ]\n for i, k in enumerate(keys):\n if 'upsample' in k:\n keys[i] = param_prefix + param_prefix + k\n else:\n keys[i] = param_prefix + k\n i = 0\n for m in modules:\n if isinstance(m, nn.Conv2d) or isinstance(m, nn.ConvTranspose2d):\n weight = weights[keys[i]].copy()\n bias = biases[keys[i]].copy()\n if keys[i] == param_prefix+'conv1':\n # print (\"%s :\"%(keys[i]), m.weight.size(), m.bias.size(), tf_w[keys[i]].shape[::-1])\n m.weight.data[:,0:3,:,:] = torch.from_numpy(np.flip(weight[:,0:3,:,:], axis=1).copy())\n m.weight.data[:,3:6,:,:] = torch.from_numpy(np.flip(weight[:,3:6,:,:], axis=1).copy())\n if m.bias is not None:\n m.bias.data[:] = torch.from_numpy(bias)\n else:\n m.weight.data[:,:,:,:] = torch.from_numpy(weight)\n if m.bias is not None:\n m.bias.data[:] = torch.from_numpy(bias)\n i = i + 1\n return\n\ndef parse_flownetsd(modules, weights, biases, param_prefix='netsd_'):\n keys = [\n 'conv0',\n 'conv1',\n 'conv1_1',\n 'conv2',\n 'conv2_1',\n 'conv3',\n 'conv3_1',\n 'conv4',\n 'conv4_1',\n 'conv5',\n 'conv5_1',\n 'conv6',\n 'conv6_1',\n \n 'deconv5',\n 'deconv4',\n 'deconv3',\n 'deconv2',\n\n 'interconv5',\n 'interconv4',\n 'interconv3',\n 'interconv2',\n \n 'Convolution1',\n 'Convolution2',\n 'Convolution3',\n 'Convolution4',\n 'Convolution5',\n\n 'upsample_flow6to5',\n 'upsample_flow5to4',\n 'upsample_flow4to3',\n 'upsample_flow3to2',\n ]\n for i, k in enumerate(keys):\n keys[i] = param_prefix + k\n\n i = 0\n for m in modules:\n if isinstance(m, nn.Conv2d) or isinstance(m, nn.ConvTranspose2d):\n weight = weights[keys[i]].copy()\n bias = biases[keys[i]].copy()\n if keys[i] == param_prefix+'conv0':\n m.weight.data[:,0:3,:,:] = torch.from_numpy(np.flip(weight[:,0:3,:,:], axis=1).copy())\n m.weight.data[:,3:6,:,:] = torch.from_numpy(np.flip(weight[:,3:6,:,:], axis=1).copy())\n if m.bias is not None:\n m.bias.data[:] = torch.from_numpy(bias)\n else:\n m.weight.data[:,:,:,:] = torch.from_numpy(weight)\n if m.bias is not None:\n m.bias.data[:] = torch.from_numpy(bias)\n i = i + 1\n\n return\n\ndef parse_flownetfusion(modules, weights, biases, param_prefix='fuse_'):\n keys = [\n 'conv0',\n 'conv1',\n 'conv1_1',\n 'conv2',\n 'conv2_1',\n\n 'deconv1',\n 'deconv0',\n\n 'interconv1',\n 'interconv0',\n \n '_Convolution5',\n '_Convolution6',\n '_Convolution7',\n\n 'upsample_flow2to1',\n 'upsample_flow1to0',\n ]\n for i, k in enumerate(keys):\n keys[i] = param_prefix + k\n\n i = 0\n for m in modules:\n if isinstance(m, nn.Conv2d) or isinstance(m, nn.ConvTranspose2d):\n weight = weights[keys[i]].copy()\n bias = biases[keys[i]].copy()\n if keys[i] == param_prefix+'conv0':\n m.weight.data[:,0:3,:,:] = torch.from_numpy(np.flip(weight[:,0:3,:,:], axis=1).copy())\n m.weight.data[:,3::,:,:] = torch.from_numpy(weight[:,3:,:,:].copy())\n if m.bias is not None:\n m.bias.data[:] = torch.from_numpy(bias)\n else:\n m.weight.data[:,:,:,:] = torch.from_numpy(weight)\n if m.bias is not None:\n m.bias.data[:] = torch.from_numpy(bias)\n i = i + 1\n\n return\n"
},
{
"path": "dvs/flownet2/utils/tools.py",
"content": "# freda (todo) : \n\nimport os, time, sys, math\nimport subprocess, shutil\nfrom os.path import *\nimport numpy as np\nfrom inspect import isclass\nfrom pytz import timezone\nfrom datetime import datetime\nimport inspect\nimport torch\n\ndef datestr():\n pacific = timezone('US/Pacific')\n now = datetime.now(pacific)\n return '{}{:02}{:02}_{:02}{:02}'.format(now.year, now.month, now.day, now.hour, now.minute)\n\ndef module_to_dict(module, exclude=[]):\n return dict([(x, getattr(module, x)) for x in dir(module)\n if isclass(getattr(module, x))\n and x not in exclude\n and getattr(module, x) not in exclude])\n\nclass TimerBlock: \n def __init__(self, title):\n print((\"{}\".format(title)))\n\n def __enter__(self):\n self.start = time.clock()\n return self\n\n def __exit__(self, exc_type, exc_value, traceback):\n self.end = time.clock()\n self.interval = self.end - self.start\n\n if exc_type is not None:\n self.log(\"Operation failed\\n\")\n else:\n self.log(\"Operation finished\\n\")\n\n\n def log(self, string):\n duration = time.clock() - self.start\n units = 's'\n if duration > 60:\n duration = duration / 60.\n units = 'm'\n print((\" [{:.3f}{}] {}\".format(duration, units, string)))\n \n def log2file(self, fid, string):\n fid = open(fid, 'a')\n fid.write(\"%s\\n\"%(string))\n fid.close()\n\ndef add_arguments_for_module(parser, module, argument_for_class, default, skip_params=[], parameter_defaults={}):\n argument_group = parser.add_argument_group(argument_for_class.capitalize())\n\n module_dict = module_to_dict(module)\n argument_group.add_argument('--' + argument_for_class, type=str, default=default, choices=list(module_dict.keys()))\n \n args, unknown_args = parser.parse_known_args()\n class_obj = module_dict[vars(args)[argument_for_class]]\n\n argspec = inspect.getargspec(class_obj.__init__)\n\n defaults = argspec.defaults[::-1] if argspec.defaults else None\n\n args = argspec.args[::-1]\n for i, arg in enumerate(args):\n cmd_arg = '{}_{}'.format(argument_for_class, arg)\n if arg not in skip_params + ['self', 'args']:\n if arg in list(parameter_defaults.keys()):\n argument_group.add_argument('--{}'.format(cmd_arg), type=type(parameter_defaults[arg]), default=parameter_defaults[arg])\n elif (defaults is not None and i < len(defaults)):\n argument_group.add_argument('--{}'.format(cmd_arg), type=type(defaults[i]), default=defaults[i])\n else:\n print((\"[Warning]: non-default argument '{}' detected on class '{}'. This argument cannot be modified via the command line\"\n .format(arg, module.__class__.__name__)))\n # We don't have a good way of dealing with inferring the type of the argument\n # TODO: try creating a custom action and using ast's infer type?\n # else:\n # argument_group.add_argument('--{}'.format(cmd_arg), required=True)\n\ndef kwargs_from_args(args, argument_for_class):\n argument_for_class = argument_for_class + '_'\n return {key[len(argument_for_class):]: value for key, value in list(vars(args).items()) if argument_for_class in key and key != argument_for_class + 'class'}\n\ndef format_dictionary_of_losses(labels, values):\n try:\n string = ', '.join([('{}: {:' + ('.3f' if value >= 0.001 else '.1e') +'}').format(name, value) for name, value in zip(labels, values)])\n except (TypeError, ValueError) as e:\n print((list(zip(labels, values))))\n string = '[Log Error] ' + str(e)\n\n return string\n\n\nclass IteratorTimer():\n def __init__(self, iterable):\n self.iterable = iterable\n self.iterator = self.iterable.__iter__()\n\n def __iter__(self):\n return self\n\n def __len__(self):\n return len(self.iterable)\n\n def __next__(self):\n start = time.time()\n n = next(self.iterator)\n self.last_duration = (time.time() - start)\n return n\n\n next = __next__\n\ndef gpumemusage():\n gpu_mem = subprocess.check_output(\"nvidia-smi | grep MiB | cut -f 3 -d '|'\", shell=True).replace(' ', '').replace('\\n', '').replace('i', '')\n all_stat = [float(a) for a in gpu_mem.replace('/','').split('MB')[:-1]]\n\n gpu_mem = ''\n for i in range(len(all_stat)/2):\n curr, tot = all_stat[2*i], all_stat[2*i+1]\n util = \"%1.2f\"%(100*curr/tot)+'%'\n cmem = str(int(math.ceil(curr/1024.)))+'GB'\n gmem = str(int(math.ceil(tot/1024.)))+'GB'\n gpu_mem += util + '--' + join(cmem, gmem) + ' '\n return gpu_mem\n\n\ndef update_hyperparameter_schedule(args, epoch, global_iteration, optimizer):\n if args.schedule_lr_frequency > 0:\n for param_group in optimizer.param_groups:\n if (global_iteration + 1) % args.schedule_lr_frequency == 0:\n param_group['lr'] /= float(args.schedule_lr_fraction)\n param_group['lr'] = float(np.maximum(param_group['lr'], 0.000001))\n\ndef save_checkpoint(state, is_best, path, prefix, filename='checkpoint.pth.tar'):\n prefix_save = os.path.join(path, prefix)\n name = prefix_save + '_' + filename\n torch.save(state, name)\n if is_best:\n shutil.copyfile(name, prefix_save + '_model_best.pth.tar')\n\n"
},
{
"path": "dvs/gyro/__init__.py",
"content": "from .gyro_function import (\n GetGyroAtTimeStamp,\n QuaternionProduct,\n QuaternionReciprocal,\n ConvertQuaternionToAxisAngle,\n FindOISAtTimeStamp,\n GetMetadata,\n GetProjections,\n GetVirtualProjection,\n GetForwardGrid,\n CenterZoom,\n GetWarpingFlow,\n torch_norm_quat,\n torch_QuaternionProduct, \n torch_QuaternionReciprocal,\n torch_GetVirtualProjection,\n get_static,\n torch_GetForwardGrid,\n torch_GetWarpingFlow,\n train_GetGyroAtTimeStamp,\n train_ConvertQuaternionToAxisAngle,\n ConvertAxisAngleToQuaternion,\n torch_ConvertAxisAngleToQuaternion,\n torch_ConvertQuaternionToAxisAngle,\n ConvertAxisAngleToQuaternion_no_angle,\n ConvertQuaternionToAxisAngle_no_angle,\n torch_GetHomographyTransformFromProjections,\n torch_ApplyTransform,\n norm_quat,\n SlerpWithDefault\n )\nfrom .gyro_io import (\n LoadGyroData, \n LoadOISData, \n LoadFrameData, \n LoadStabResult,\n get_grid, \n get_rotations, \n visual_rotation\n )"
},
{
"path": "dvs/gyro/gyro_function.py",
"content": "import numpy as np\nfrom numpy import linalg as LA\nimport matplotlib.pyplot as plt\nimport torch\nfrom torch.autograd import Variable\n\ndef get_static(height = 1080, width = 1920, ratio = 0.1):\n static_options = {}\n static_options[\"active_array_width\"] = 4032\n static_options[\"active_array_height\"] = 3024\n static_options[\"crop_window_width\"] = 4032\n static_options[\"crop_window_height\"] = 2272\n static_options[\"num_grid_rows\"] = 12\n static_options[\"num_grid_cols\"] = 12\n static_options[\"dim_homography\"] = 9\n static_options[\"width\"] = width # frame width.\n static_options[\"height\"] = height # frame height\n # static_options[\"fov\"] = 1.27 # sensor_width/sensor_focal_length\n static_options[\"cropping_ratio\"] = 0.0 #ratio # normalized cropping ratio at each side. \n return static_options\n\n# Quaternion: [x, y, z, w]\n\ndef norm_quat(quat):\n norm_quat = LA.norm(quat) \n if norm_quat > 1e-6:\n quat = quat / norm_quat \n # [0 norm_quat norm_quat - 1e-6]\n else:\n # print('bad len for Reciprocal')\n quat = np.array([0,0,0,1])\n return quat\n\ndef torch_norm_quat(quat, USE_CUDA = True):\n # Method 1:\n batch_size = quat.size()[0]\n quat_out = Variable(torch.zeros((batch_size, 4), requires_grad=True))\n if USE_CUDA == True:\n quat_out = quat_out.cuda()\n for i in range(batch_size):\n norm_quat = torch.norm(quat[i]) \n if norm_quat > 1e-6: \n quat_out[i] = quat[i] / norm_quat \n # [0 norm_quat norm_quat - 1e-6]\n else:\n quat_out[i,:3] = quat[i,:3] * 0\n quat_out[i,3] = quat[i,3] / quat[i,3]\n\n # Method 2:\n # quat = quat / (torch.unsqueeze(torch.norm(quat, dim = 1), 1) + 1e-6) # check norm\n return quat_out\n\ndef ConvertAxisAngleToQuaternion(axis, angle):\n if LA.norm(axis) > 1e-6 and angle > 1e-6: \n axis = axis/LA.norm(axis) \n half_angle = angle*0.5 \n sin_half_angle = np.sin(half_angle)\n quat = np.array([sin_half_angle* axis[0], sin_half_angle* axis[1], sin_half_angle* axis[2], np.cos(half_angle)])\n\n return norm_quat(quat)\n\ndef ConvertAxisAngleToQuaternion_no_angle(axis):\n angle = LA.norm(axis) \n if LA.norm(axis) > 1e-6: \n axis = axis/LA.norm(axis) \n half_angle = angle*0.5 \n sin_half_angle = np.sin(half_angle)\n quat = np.array([sin_half_angle* axis[0], sin_half_angle* axis[1], sin_half_angle* axis[2], np.cos(half_angle)])\n\n return norm_quat(quat)\n\ndef torch_ConvertAxisAngleToQuaternion(axis, USE_CUDA = True):\n batch_size = axis.size()[0]\n\n angle = torch.norm(axis[:,:3], dim = 1)\n\n half_angle = angle * 0.5 \n sin_half_angle = torch.sin(half_angle)\n quats = Variable(torch.zeros((batch_size, 4), requires_grad=True))\n norm_axis = axis[:,:3] * 1\n if USE_CUDA:\n quats = quats.cuda()\n for i in range(batch_size):\n if angle[i] > 1e-6:\n norm_axis[i] = axis[i,:3]/angle[i]\n quats[:, :3] = sin_half_angle * norm_axis\n quats[:, 3] = torch.cos(half_angle)\n return torch_norm_quat(quats)\n\ndef ConvertQuaternionToAxisAngle(quat):\n quat = quat/LA.norm(quat) \n axis_norm = LA.norm(quat[0:3])\n axis = np.array([0.0, 0.0, 0.0])\n if axis_norm < 1e-6:\n angle = 0 \n else:\n axis_norm_reciprocal = 1/axis_norm \n axis[0] = quat[0] * axis_norm_reciprocal \n axis[1] = quat[1] * axis_norm_reciprocal \n axis[2] = quat[2] * axis_norm_reciprocal \n angle = 2 * np.arccos(quat[3])\n return [axis, angle]\n\ndef ConvertQuaternionToAxisAngle_no_angle(quat):\n quat = quat/LA.norm(quat) \n axis_norm = LA.norm(quat[0:3])\n axis = np.array([0.0, 0.0, 0.0]) \n if axis_norm > 1e-6:\n axis_norm_reciprocal = 1 / axis_norm * 2 * np.arccos(quat[3])\n axis[0] = quat[0] * axis_norm_reciprocal \n axis[1] = quat[1] * axis_norm_reciprocal \n axis[2] = quat[2] * axis_norm_reciprocal \n return axis\n\ndef torch_ConvertQuaternionToAxisAngle(quat, USE_CUDA = True):\n batch_size = quat.size()[0]\n axis_angle = Variable(torch.zeros((batch_size, 4), requires_grad=True))\n if USE_CUDA:\n axis_angle = axis_angle.cuda()\n for i in range(batch_size): \n axis_norm = torch.norm(quat[i, 0:3])\n if axis_norm > 1e-6:\n axis_norm_reciprocal = 1/axis_norm * 2 * torch.acos(quat[i,3])\n axis_angle[i,0] = quat[i,0] * axis_norm_reciprocal \n axis_angle[i,1] = quat[i,1] * axis_norm_reciprocal \n axis_angle[i,2] = quat[i,2] * axis_norm_reciprocal \n return axis_angle\n\ndef train_ConvertQuaternionToAxisAngle(quat):\n out = np.zeros(4)\n out[:3] = ConvertQuaternionToAxisAngle_no_angle(quat)\n return out\n\ndef AngularVelocityToQuat(angular_v, dt):\n length = LA.norm(angular_v) \n if length < 1e-6:\n angular_v = np.array([1, 0, 0]) \n print('bad length')\n else:\n angular_v = angular_v/length \n quat = ConvertAxisAngleToQuaternion(angular_v, length*dt) \n return quat\n\ndef QuaternionProduct(q1, q2):\n x1 = q1[0] \n y1 = q1[1] \n z1 = q1[2] \n w1 = q1[3] \n\n x2 = q2[0] \n y2 = q2[1] \n z2 = q2[2] \n w2 = q2[3] \n\n quat = np.zeros(4)\n quat[3] = w1 * w2 - x1 * x2 - y1 * y2 - z1 * z2 \n quat[0] = w1 * x2 + x1 * w2 + y1 * z2 - z1 * y2 \n quat[1] = w1 * y2 - x1 * z2 + y1 * w2 + z1 * x2 \n quat[2] = w1 * z2 + x1 * y2 - y1 * x2 + z1 * w2 \n\n return norm_quat(quat)\n\ndef torch_QuaternionProduct(q1, q2, USE_CUDA = True):\n x1 = q1[:,0] \n y1 = q1[:,1] \n z1 = q1[:,2] \n w1 = q1[:,3] \n\n x2 = q2[:,0] \n y2 = q2[:,1] \n z2 = q2[:,2] \n w2 = q2[:,3] \n\n batch_size = q1.size()[0]\n quat = Variable(torch.zeros((batch_size, 4), requires_grad=True))\n if USE_CUDA == True:\n quat = quat.cuda()\n \n quat[:,3] = w1 * w2 - x1 * x2 - y1 * y2 - z1 * z2 \n quat[:,0] = w1 * x2 + x1 * w2 + y1 * z2 - z1 * y2 \n quat[:,1] = w1 * y2 - x1 * z2 + y1 * w2 + z1 * x2 \n quat[:,2] = w1 * z2 + x1 * y2 - y1 * x2 + z1 * w2 \n\n quat = torch_norm_quat(quat)\n\n return quat\n\ndef ProcessGyroRotation(gyro_data):\n num_inputs = np.shape(gyro_data)[0]\n quats = np.zeros((num_inputs, 4)) \n quats[0,:] = np.array([0, 0, 0, 1])\n for i in range(1, num_inputs):\n dt = (gyro_data[i, 0] - gyro_data[i-1, 0])*1e-9 \n quat = AngularVelocityToQuat(gyro_data[i, 1:4], dt) \n quats[i,:] = QuaternionProduct(quat, quats[i-1,:]) # R_t = delta R_t * R_t-1\n quats[i,:] = quats[i,:] / LA.norm(quats[i,:]) \n return quats \n\ndef QuaternionReciprocal(q):\n quat = np.array([-q[0], -q[1], -q[2], q[3]]) \n return norm_quat(quat)\n\ndef torch_QuaternionReciprocal(q, USE_CUDA = True):\n quat = torch.cat((-q[:,0:1], -q[:,1:2], -q[:,2:3], q[:,3:]), dim = 1) \n batch_size = quat.size()[0]\n\n quat = torch_norm_quat(quat)\n return quat\n\ndef ProcessGyroData(gyro_data):\n quats = ProcessGyroRotation(gyro_data) \n size = np.shape(gyro_data)[0]\n axis_dif = np.zeros((size,3)) \n for i in range(1, size):\n quat_dif = QuaternionProduct(quats[i,:], QuaternionReciprocal(quats[i-1,:])) \n [axis_dif_cur, angles_cur] = ConvertQuaternionToAxisAngle(quat_dif) \n axis_dif[i,:] = axis_dif_cur*angles_cur \n return [axis_dif, quats]\n\n\ndef SlerpWithDefault(q1, q2, t, q_default):\n t = max(min(t, 1.0), 0.0) \n kEpsilon = 1e-6 \n kSlerpLinearThresh = 0.9995 \n \n q1 = q1/LA.norm(q1) \n q2 = q2/LA.norm(q2) \n\n if t < kEpsilon:\n q3 = q1 \n return q3\n elif t > 1-kEpsilon:\n q3 = q2 \n return q3\n\n dot_prodcut = np.sum(q1*q2) \n\n if abs(dot_prodcut) >= 1:\n q3= q_default \n return q3\n elif abs(dot_prodcut) > kSlerpLinearThresh:\n q3 = q1*(1-t) + q2*t \n q3 = q3/LA.norm(q3) \n return q3\n\n sign = 1 \n if dot_prodcut < 0:\n sign = -1 \n dot_prodcut = -dot_prodcut \n\n theta = np.arccos(dot_prodcut) \n sin_theta = np.sin(theta) \n inv_sin_theta = 1.0 / sin_theta \n coeff1 = np.sin((1.0 - t) * theta) * inv_sin_theta \n coeff2 = sign * np.sin(t * theta) * inv_sin_theta \n q3 = q1 * coeff1 + q2 * coeff2 \n return q3\n\n\ndef GetGyroAtTimeStamp(gyro_data, timestamp):\n z = np.array([0,0,0,1]) \n if len(gyro_data) >= 2 and (not(timestamp < gyro_data[0,0] or timestamp > gyro_data[-1, 0])):\n ind = np.where(gyro_data[:,0] >= timestamp)\n ind = np.squeeze( ind, axis = 0)\n if gyro_data[ind[0], 0] == timestamp:\n z = gyro_data[ind[0],1:]\n else:\n start_index = ind[0] -1 \n end_index = ind[0] \n ratio = (timestamp - gyro_data[start_index,0])/(gyro_data[end_index,0]-gyro_data[start_index,0])\n z = SlerpWithDefault(gyro_data[start_index,1:], gyro_data[end_index, 1:], ratio, gyro_data[start_index,1:]) \n z = z / (LA.norm(z) + 1e-6)\n return z\n\ndef train_GetGyroAtTimeStamp(gyro_data, timestamp, check = False):\n if len(gyro_data) >= 2 and (not(timestamp < gyro_data[0,0] or timestamp > gyro_data[-1, 0])):\n ind = np.where(gyro_data[:,0] >= timestamp)\n ind = np.squeeze( ind, axis = 0)\n if gyro_data[ind[0], 0] == timestamp:\n z = gyro_data[ind[0],1:]\n else:\n start_index = ind[0] -1 \n end_index = ind[0] \n ratio = (timestamp - gyro_data[start_index,0])/(gyro_data[end_index,0]-gyro_data[start_index,0])\n z = SlerpWithDefault(gyro_data[start_index,1:], gyro_data[end_index, 1:], ratio, gyro_data[start_index,1:]) \n return z / (LA.norm(z) + 1e-6)\n if check:\n print(\"bad value\")\n return None\n\ndef FindOISAtTimeStamp(ois_log, time):\n ois_time = ois_log[:,2] \n if time <= ois_time[0]:\n ois_data = ois_log[0, 0:2] \n elif time > ois_time[-1]:\n ois_data = ois_log[-1, 0:2]\n else:\n ind = np.where(ois_time >= time)\n ind = np.squeeze( ind, axis = 0)\n first_ind = ind[0]\n if ois_time[first_ind] == ind[0]:\n ois_data = ois_log[first_ind, 0:2]\n else:\n cur_time = ois_time[first_ind] \n last_timestamp = ois_time[first_ind - 1]\n ratio = (time - last_timestamp) / (cur_time - last_timestamp) \n ois_data = ois_log[first_ind - 1,0:2] * (1-ratio) + ois_log[first_ind,0:2]*ratio \n\n return ois_data\n\ndef GetMetadata(frame_data, frame_index, result_poses = {} ):\n # global static_options\n # We can just use 1.27 as fov and virtual fov for videos in the data set.\n metadata = {}\n metadata[\"frame_id\"] = frame_index\n metadata[\"timestamp_ns\"] = frame_data[frame_index, 0]\n metadata[\"timestamp_ois_ns\"] = frame_data[frame_index, 4]\n metadata[\"rs_time_ns\"] = frame_data[frame_index, 3]\n if \"real fov\" in result_poses:\n metadata[\"fov\"] = result_poses['real fov'][frame_index,:] \n else:\n metadata[\"fov\"] = 1.27\n if \"virtual fov\" in result_poses:\n metadata[\"virtual_fov\"] = result_poses['virtual fov'][frame_index,:] \n else:\n metadata[\"virtual_fov\"] = 1.27\n\n return metadata\n\ndef GetProjections(static_options, metadata, quats_data, ois_data, no_shutter = False):\n num_rows = static_options[\"num_grid_rows\"]\n real_projections = []\n for i in range(num_rows):\n if no_shutter:\n timestmap_ns = metadata[\"timestamp_ns\"] + metadata[\"rs_time_ns\"] * 0.5\n timestamp_ois_ns = metadata[\"timestamp_ois_ns\"] + metadata[\"rs_time_ns\"] * 0.5\n else:\n timestmap_ns = metadata[\"timestamp_ns\"] + metadata[\"rs_time_ns\"] * i / (num_rows-1)\n timestamp_ois_ns = metadata[\"timestamp_ois_ns\"] + metadata[\"rs_time_ns\"] * i / (num_rows-1)\n real_projections.append(GetRealProjection(\n static_options, quats_data, ois_data, metadata[\"fov\"], timestmap_ns, timestamp_ois_ns))\n return real_projections\n\ndef GetRealProjection(static_options, quats_data, ois_data, fov, timestamp_ns, timestamp_ois_ns):\n quat = GetGyroAtTimeStamp(quats_data, timestamp_ns)\n ois_offset = FindOISAtTimeStamp(ois_data, timestamp_ois_ns) \n # ois is w.r.t. active array size, thus we need to convert it to normalzied space.\n\n ois_offset = np.array(ois_offset) / np.array([static_options[\"crop_window_width\"], static_options[\"crop_window_height\"]])\n \n projection = GetProjectionHomography(quat, fov, ois_offset, static_options[\"width\"], static_options[\"height\"])\n return projection\n\ndef GetProjectionHomography(rot, fov, offset, width, height):\n # rot: rotation in quaternion\n # fov: sensor_width / focal_length.\n # offset: additional ois offset at normalized domain.\n # width/height: frame size.\n focal_length = width / fov\n rotation = ConvertQuaternionToRotationMatrix(rot)\n intrinsics = GetIntrinsics(focal_length, offset, width, height)\n projection_homography = np.matmul(intrinsics, rotation)\n return projection_homography\n\ndef torch_GetProjectionHomography(rot, fov, width, height, USE_CUDA = True):\n # rot: rotation in quaternion\n # fov: sensor_width / focal_length.\n # offset: additional ois offset at normalized domain.\n # width/height: frame size.\n focal_length = width / fov\n rotation = torch_ConvertQuaternionToRotationMatrix(rot)\n batch_size = rotation.size()[0]\n offset = np.array([0,0])\n intrinsics = GetIntrinsics(focal_length, offset, width, height)\n intrinsics = torch.Tensor(np.repeat(np.expand_dims(intrinsics, axis = 0), batch_size, axis = 0))\n if USE_CUDA == True:\n intrinsics = intrinsics.cuda()\n projection_homography = torch.matmul(intrinsics, rotation)\n return projection_homography\n\ndef ConvertQuaternionToRotationMatrix(quat):\n x = quat[0]\n y = quat[1]\n z = quat[2]\n w = quat[3]\n rotation = np.zeros(9)\n rotation[0] = 1 - 2 * y * y - 2 * z * z\n rotation[1] = 2 * x * y - 2 * z * w\n rotation[2] = 2 * x * z + 2 * y * w\n rotation[3] = 2 * x * y + 2 * z * w\n rotation[4] = 1 - 2 * x * x - 2 * z * z\n rotation[5] = 2 * y * z - 2 * x * w\n rotation[6] = 2 * x * z - 2 * y * w\n rotation[7] = 2 * y * z + 2 * x * w\n rotation[8] = 1 - 2 * x * x - 2 * y * y\n rotation = np.reshape(rotation, (3, 3)) # Note reshape is different with matlab\n return rotation\n\ndef torch_ConvertQuaternionToRotationMatrix(quat, USE_CUDA = True):\n x = quat[:,0]\n y = quat[:,1]\n z = quat[:,2]\n w = quat[:,3]\n\n batch_size = quat.size()[0]\n rotation = Variable(torch.zeros((batch_size, 9), requires_grad=True))\n if USE_CUDA == True:\n rotation = rotation.cuda()\n\n rotation[:,0] = 1 - 2 * y * y - 2 * z * z\n rotation[:,1] = 2 * x * y - 2 * z * w\n rotation[:,2] = 2 * x * z + 2 * y * w\n rotation[:,3] = 2 * x * y + 2 * z * w\n rotation[:,4] = 1 - 2 * x * x - 2 * z * z\n rotation[:,5] = 2 * y * z - 2 * x * w\n rotation[:,6] = 2 * x * z - 2 * y * w\n rotation[:,7] = 2 * y * z + 2 * x * w\n rotation[:,8] = 1 - 2 * x * x - 2 * y * y\n rotation = rotation.view(batch_size, 3, 3) # Note reshape is different with matlab\n return rotation\n\ndef ConvertRotationMatrixToQuaternion(m):\n tr = m[0,0] + m[1,1] + m[2,2]\n if tr > 0 :\n S = 2 * (tr+1.0)**0.5\n qw = 0.25 * S\n qx = (m[2,1] - m[1,2]) / S\n qy = (m[0,2] - m[2,0]) / S\n qz = (m[1,0] - m[0,1]) / S\n elif m[0,0] > m[1,1] and m[0,0] > m[2,2]:\n S = 2* (1.0 + m[0,0] - m[1,1] - m[2,2]) ** 0.5\n qw = (m[2,1] - m[1,2]) / S\n qx = 0.25 * S\n qy = (m[0,1] + m[1,0]) / S\n qz = (m[0,2] + m[2,0]) / S\n elif m[1,1] > m[2,2]: \n S = 2* (1.0 - m[0,0] + m[1,1] - m[2,2]) ** 0.5\n qw = (m[0,2] - m[2,0]) / S\n qx = (m[0,1] + m[1,0]) / S\n qy = 0.25 * S\n qz = (m[1,2] + m[2,1]) / S\n else: \n S = 2* (1.0 - m[0,0] - m[1,1] + m[2,2]) ** 0.5\n qw = (m[1,0] - m[0,1]) / S\n qx = (m[0,2] + m[2,0]) / S\n qy = (m[1,2] + m[2,1]) / S\n qz = 0.25 * S\n return np.array([qx,qy,qz,qw])\n\ndef GetIntrinsics(focal_length, offset, width, height):\n intrinsics = [\n [float(focal_length), 0.0, 0.5*(width-1)+offset[0]*width], \n [0.0, float(focal_length), 0.5*(height-1)+offset[1]*height], \n [0.0, 0.0, 1.0]\n ]\n return np.array(intrinsics)\n\n\ndef GetVirtualProjection(static_options, result_pose, metadata, frame_index):\n # debug only, for getting results and references for comparisons.\n quat = result_pose['virtual pose'][frame_index,:]\n if 'vitual lens offset' in result_pose:\n virutal_lens_offset = result_pose['vitual lens offset'][frame_index,:]\n else:\n virutal_lens_offset = np.array([0,0])\n virtual_projection = GetProjectionHomography(\n quat, metadata[\"virtual_fov\"], virutal_lens_offset, static_options[\"width\"], static_options[\"height\"])\n return virtual_projection\n\ndef torch_GetVirtualProjection(static_options, quat, virtual_fov = 1.27):\n virtual_projection = torch_GetProjectionHomography(\n quat, virtual_fov, static_options[\"width\"], static_options[\"height\"])\n return virtual_projection\n\n\ndef GetForwardGrid(static_options, real_projections, virtual_projection):\n # real_projections: a set of 3x3 projections.\n # virtual_projection: a single 3x3 projection.\n\n grid = np.zeros((4, static_options[\"num_grid_cols\"], static_options[\"num_grid_rows\"]))\n width = static_options[\"width\"]\n height = static_options[\"height\"]\n\n row_step = 1/ (static_options[\"num_grid_rows\"] - 1)\n col_step = 1/ (static_options[\"num_grid_cols\"] - 1)\n\n for i in range(static_options[\"num_grid_rows\"]):\n transform = GetHomographyTransformFromProjections(real_projections[i], virtual_projection)\n v = i * row_step\n for j in range(static_options[\"num_grid_cols\"]):\n u = j * col_step\n point = np.array([u * width, v * height, 1]).T\n warped_point = ApplyTransform(transform, point)\n warped_point = warped_point / np.array([width, height, 1]) # normalize\n grid[:, j, i] = np.array([warped_point[0], warped_point[1], u, v])\n return grid\n\ndef torch_GetForwardGrid(static_options, real_projections, virtual_projection, USE_CUDA = True):\n # real_projections: a set of 3x3 projections.\n # virtual_projection: a single 3x3 projection.\n batch_size = real_projections.size()[0]\n\n grid = torch.zeros((batch_size, 4, static_options[\"num_grid_cols\"], static_options[\"num_grid_rows\"]))\n if USE_CUDA:\n grid = grid.cuda()\n width = static_options[\"width\"]\n height = static_options[\"height\"]\n\n row_step = 1/ (static_options[\"num_grid_rows\"] - 1)\n col_step = 1/ (static_options[\"num_grid_cols\"] - 1)\n\n for i in range(static_options[\"num_grid_rows\"]):\n transform = torch_GetHomographyTransformFromProjections(real_projections[:, i], virtual_projection)\n v = i * row_step\n for j in range(static_options[\"num_grid_cols\"]):\n u = j * col_step\n point = torch.Tensor([u * width, v * height, 1])\n norm = torch.Tensor([width, height, 1])\n if USE_CUDA == True:\n point = point.cuda()\n norm = norm.cuda()\n warped_point = torch_ApplyTransform(transform, point)\n warped_point = warped_point / norm # normalize\n grid[:, 0, j, i] = warped_point[:,0]\n grid[:, 1, j, i] = warped_point[:,1]\n grid[:, 2, j, i] = u\n grid[:, 3, j, i] = v\n return grid\n\ndef GetWarpingFlow(real_projections_src, real_projections_dst, num_rows, num_cols, frame_width, frame_height):\n # num_rows: rows of the flow.\n # num_cols: cols of the flow.\n grid = np.zeros((4, num_cols, num_rows))\n\n row_step = 1/ (num_rows - 1)\n col_step = 1/ (num_cols - 1)\n\n for i in range(num_rows):\n transform = GetHomographyTransformFromProjections(real_projections_src[i], real_projections_dst[i])\n v = i * row_step\n for j in range(num_cols):\n u = j * col_step\n point = np.array([u * frame_width, v * frame_height, 1]).T\n warped_point = ApplyTransform(transform, point)\n warped_point = warped_point / np.array([frame_width, frame_height, 1]) # normalize\n grid[:, j, i] = np.array([warped_point[0], warped_point[1], u, v])\n return grid\n\ndef torch_GetWarpingFlow(static_options, real_projections_src, real_projections_dst, USE_CUDA = True):\n # real_projections: a set of 3x3 projections.\n # virtual_projection: a single 3x3 projection.\n batch_size = real_projections_src.size()[0]\n\n grid = torch.zeros((batch_size, 4, static_options[\"num_grid_cols\"], static_options[\"num_grid_rows\"]))\n if USE_CUDA:\n grid = grid.cuda()\n width = static_options[\"width\"]\n height = static_options[\"height\"]\n\n row_step = 1/ (static_options[\"num_grid_rows\"] - 1)\n col_step = 1/ (static_options[\"num_grid_cols\"] - 1)\n\n for i in range(static_options[\"num_grid_rows\"]):\n transform = torch_GetHomographyTransformFromProjections(real_projections_src[:, i], real_projections_dst[:, i])\n v = i * row_step\n for j in range(static_options[\"num_grid_cols\"]):\n u = j * col_step\n point = torch.Tensor([u * width, v * height, 1])\n norm = torch.Tensor([width, height, 1])\n if USE_CUDA == True:\n point = point.cuda()\n norm = norm.cuda()\n warped_point = torch_ApplyTransform(transform, point)\n warped_point = warped_point / norm # normalize\n grid[:, 0, j, i] = warped_point[:,0]\n grid[:, 1, j, i] = warped_point[:,1]\n grid[:, 2, j, i] = u\n grid[:, 3, j, i] = v\n return grid\n\ndef GetHomographyTransformFromProjections(proj_src, proj_dst):\n return np.matmul(proj_dst, LA.inv(proj_src))\n\ndef torch_GetHomographyTransformFromProjections(proj_src, proj_dst):\n return torch.matmul(proj_dst, torch.inverse(proj_src))\n\ndef ApplyTransform(transform, point):\n # Warps a 2D point ([x y 1]) using a homography transform.\n # Returns the warped 2D point ([warped_x, warped_y, 1]).\n z = np.matmul(transform, point)\n z = z / z[2]\n return z\n\ndef torch_ApplyTransform(transform, point):\n # Warps a 2D point ([x y 1]) using a homography transform.\n # Returns the warped 2D point ([warped_x, warped_y, 1]).\n z = torch.matmul(transform, point)\n z = z / z[:,2:]\n return z\n\ndef CenterZoom(grid, ratio):\n grid[:, 0:2, :, :] = (grid[:, 0:2, :, :] - 0.5) * ratio + 0.5\n return grid\n\n"
},
{
"path": "dvs/gyro/gyro_io.py",
"content": "import numpy as np\nfrom numpy import linalg as LA\nimport matplotlib.pyplot as plt\nimport scipy.io as sio\nfrom .gyro_function import (\n ProcessGyroData, QuaternionProduct, QuaternionReciprocal, \n ConvertQuaternionToAxisAngle, FindOISAtTimeStamp, GetMetadata,\n GetProjections, GetVirtualProjection, GetForwardGrid,\n CenterZoom, GetGyroAtTimeStamp, get_static, ConvertAxisAngleToQuaternion,\n ConvertAxisAngleToQuaternion_no_angle, ConvertQuaternionToAxisAngle_no_angle\n )\n\ndef load_gyro_mesh(input_name):\n data = LoadStabResult(input_name)\n w, h = data[\"vertex_grid_size\"][0]\n data[\"warping grid\"] = np.reshape(data[\"warping grid\"],(-1,int(w),int(h),4))\n return data\n\ndef get_grid(static_options, frame_data, quats_data, ois_data, virtual_data, no_shutter = False):\n grid = []\n result_poses = {}\n result_poses['virtual pose'] = virtual_data\n for i in range(len(virtual_data)):\n metadata = GetMetadata(frame_data, i)\n real_projections = GetProjections(static_options, metadata, quats_data, ois_data, no_shutter = no_shutter)\n virtual_projection = GetVirtualProjection(static_options, result_poses, metadata, i) \n grid.append(GetForwardGrid(static_options, real_projections, virtual_projection))\n grid = np.array(grid)\n zoom_ratio = 1 / (1 - 2 * static_options[\"cropping_ratio\"])\n curr_grid = CenterZoom(grid, zoom_ratio)\n curr_grid = np.transpose(curr_grid,(0,3,2,1))\n return curr_grid\n\ndef get_rotations(frame_data, quats_data, ois_data, num_frames):\n quats = np.zeros((num_frames, 4)) \n for i in range(num_frames):\n quats[i,:] = GetGyroAtTimeStamp(quats_data, frame_data[i,0])\n\n rotations = np.zeros((num_frames,3))\n lens_offsets = np.zeros((num_frames, 2)) \n for i in range(num_frames):\n if i != 0:\n quat_dif = QuaternionProduct(quats[i,:], QuaternionReciprocal(quats[i-1,:])) \n axis_dif_cur = ConvertQuaternionToAxisAngle_no_angle(quat_dif)\n rotations[i,:] = axis_dif_cur\n lens_offsets[i,:] = FindOISAtTimeStamp(ois_data, frame_data[i, 4]) \n\n return rotations, lens_offsets\n\ndef visual_rotation(rotations_real, lens_offsets_real, rotations_virtual, lens_offsets_virtual, rotations_virtual2, lens_offsets_virtual2, path):\n # figure('units','normalized','outerposition',[0 0 1 1])\n plt.clf()\n plt.figure(figsize=(8,16))\n \n plt.subplot(5,1,1)\n plt.plot(rotations_real[:,0], \"g\")\n if rotations_virtual is not None:\n plt.plot(rotations_virtual[:,0], \"b\")\n if rotations_virtual2 is not None:\n plt.plot(rotations_virtual2[:,0], \"r\")\n plt.ylim(-0.02, 0.02)\n plt.xlabel('frame id')\n plt.ylabel('gyro x')\n\n plt.subplot(5,1,2)\n plt.plot(rotations_real[:,1], \"g\")\n if rotations_virtual is not None:\n plt.plot(rotations_virtual[:,1], \"b\")\n if rotations_virtual2 is not None:\n plt.plot(rotations_virtual2[:,1], \"r\")\n plt.ylim(-0.02, 0.02)\n plt.xlabel('frame id')\n plt.ylabel('gyro y')\n\n plt.subplot(5,1,3)\n plt.plot(rotations_real[:,2], \"g\")\n if rotations_virtual is not None:\n plt.plot(rotations_virtual[:,2], \"b\")\n if rotations_virtual2 is not None:\n plt.plot(rotations_virtual2[:,2], \"r\")\n plt.ylim(-0.02, 0.02)\n plt.xlabel('frame id')\n plt.ylabel('gyro z')\n \n plt.subplot(5,1,4)\n plt.plot(lens_offsets_real[:,0], \"g\")\n if lens_offsets_virtual is not None:\n plt.plot(lens_offsets_virtual[:,0], \"b\")\n if rotations_virtual2 is not None:\n plt.plot(lens_offsets_virtual2[:,0], \"r\")\n plt.xlabel('frame id')\n plt.ylabel('ois x')\n\n plt.subplot(5,1,5)\n plt.plot(lens_offsets_real[:,1], \"g\")\n if lens_offsets_virtual is not None:\n plt.plot(lens_offsets_virtual[:,1], \"b\")\n if rotations_virtual2 is not None:\n plt.plot(lens_offsets_virtual2[:,1], \"r\")\n plt.xlabel('frame id')\n plt.ylabel('ois y')\n \n plt.savefig(path[:-4]+\".jpg\")\n return\n\ndef LoadOISData(ois_name):\n ois_log = np.loadtxt(ois_name)\n ois_log = ois_log[:, -3:]\n return ois_log\n\ndef LoadFrameData(frame_log_name):\n frame_data = np.loadtxt(frame_log_name)\n frame_data[:, [0,4]] = frame_data[:, [0,4]] - np.expand_dims(frame_data[:,1]/2, axis = 1)\n return frame_data\n\n\ndef LoadGyroData(gyro_log_name):\n raw_gyro_data = np.loadtxt(gyro_log_name) \n raw_gyro_data[:,0] = raw_gyro_data[:,0] * 1000 \n raw_gyro_data = raw_gyro_data[:,[0, 2, 1, 3]]\n\n [_, quats_data] = ProcessGyroData(raw_gyro_data) \n quats_data = np.concatenate((raw_gyro_data[:, 0, None], quats_data), axis = 1)\n return quats_data\n\ndef LoadStabResult(input_name):\n fid = open(input_name)\n data = {}\n while True:\n name, val = ReadLine(fid)\n if name == None:\n break\n if name in data:\n data[name] = np.concatenate((data[name], val), axis=0)\n else:\n data[name] = val\n fid.close()\n print(\"Mesh length: \", len(list(data.values())[0]))\n return data\n\n\ndef ReadLine(fid):\n name = ''\n val = 0\n tline = fid.readline()\n if len(tline) == 0:\n return None, None\n if tline[-1] == \"\\n\":\n tline = tline[:-1]\n ind = tline.find(':')\n name = tline[:ind]\n tmp_val= str2num(tline[ind+1:])\n if len(tmp_val) > 0:\n val = tmp_val\n else:\n tline = fid.readline()\n if tline[-1] == \"\\n\":\n tline = tline[:-1]\n val = str2num(tline)\n return name, np.expand_dims(np.array(val), axis=0)\n\ndef str2num(string):\n nums = string.split(\" \")\n nums = [float(_) for _ in nums if _ != \"\"]\n return nums\n \n "
},
{
"path": "dvs/inference.py",
"content": "import os\nimport sys\nimport torch\nimport torchvision\nimport torch.nn as nn\nfrom torch.autograd import Variable\n\nimport time\nimport yaml\nimport argparse\nimport numpy as np\nfrom printer import Printer\nfrom dataset import get_data_loader, get_inference_data_loader\nfrom model import Model\nimport datetime\nimport copy\nfrom util import make_dir, get_optimizer, norm_flow\nfrom gyro import (\n get_grid, \n get_rotations, \n visual_rotation,\n torch_QuaternionProduct,\n torch_norm_quat\n )\nfrom warp import warp_video\n\nos.environ[\"CUDA_VISIBLE_DEVICES\"] = \"0\"\n\ndef run(model, loader, cf, USE_CUDA=True):\n no_flo = False\n number_virtual, number_real = cf['data'][\"number_virtual\"], cf['data'][\"number_real\"]\n model.net.eval()\n model.unet.eval()\n activation = nn.Softshrink(0.0006) # 0.0036\n for i, data in enumerate(loader, 0):\n # get the inputs; data is a list of [inputs, labels]\n real_inputs, times, flo, flo_back, real_projections, real_postion, ois, real_queue_idx = data\n print(\"Fininsh Load data\")\n\n real_inputs = real_inputs.type(torch.float) #[b,60,84=21*4]\n real_projections = real_projections.type(torch.float) \n flo = flo.type(torch.float) \n flo_back = flo_back.type(torch.float) \n ois = ois.type(torch.float)\n\n batch_size, step, dim = real_inputs.size()\n times = times.numpy()\n real_queue_idx = real_queue_idx.numpy()\n virtual_queue = [None] * batch_size\n\n run_loss = 0\n model.net.init_hidden(batch_size)\n count = 0\n for j in range(step):\n if (j+1) % 100 == 0:\n print(\"Step: \"+str(j+1)+\"/\"+str(step))\n virtual_inputs, vt_1 = loader.dataset.get_virtual_data(\n virtual_queue, real_queue_idx, times[:, j], times[:, j+1], times[:, 0], batch_size, number_virtual, real_postion[:,j]) \n real_inputs_step = real_inputs[:,j,:]\n inputs = torch.cat((real_inputs_step,virtual_inputs), dim = 1) \n\n # inputs = Variable(real_inputs_step)\n if USE_CUDA:\n real_inputs_step = real_inputs_step.cuda()\n virtual_inputs = virtual_inputs.cuda()\n inputs = inputs.cuda()\n if no_flo is False:\n flo_step = flo[:,j].cuda()\n flo_back_step = flo_back[:,j].cuda()\n else:\n flo_step = None\n flo_back_step = None\n vt_1 = vt_1.cuda()\n real_projections_t = real_projections[:,j+1].cuda()\n real_projections_t_1 = real_projections[:,j].cuda()\n real_postion_anchor = real_postion[:,j].cuda()\n ois_step = ois[:,j].cuda()\n\n if no_flo is False:\n b, h, w, _ = flo_step.size()\n flo_step = norm_flow(flo_step, h, w)\n flo_back_step = norm_flow(flo_back_step, h, w)\n\n with torch.no_grad():\n if no_flo is False:\n flo_out = model.unet(flo_step, flo_back_step)\n else:\n flo_out = None\n if j < 1:\n for i in range(2):\n out = model.net(inputs, flo_out, ois_step)\n else:\n out = model.net(inputs, flo_out, ois_step)\n\n real_position = real_inputs_step[:,40:44]\n virtual_position = virtual_inputs[:, -4:]\n\n out[:, :3] = activation(out[:, :3])\n out = torch_norm_quat(out)\n\n pos = torch_QuaternionProduct(virtual_position, real_postion_anchor)\n loss_step = model.loss(out, vt_1, virtual_inputs, real_inputs_step, \\\n flo_step, flo_back_step, real_projections_t, real_projections_t_1, real_postion_anchor, \\\n follow = True, optical = True, undefine = True)\n run_loss += loss_step\n\n out = torch_QuaternionProduct(out, pos)\n\n if USE_CUDA:\n out = out.cpu().detach().numpy() \n\n virtual_queue = loader.dataset.update_virtual_queue(batch_size, virtual_queue, out, times[:,j+1])\n \n run_loss /= step\n print( \"\\nLoss: follow, angle, smooth, c2_smooth, undefine, optical\")\n print(run_loss.cpu().numpy()[:-1], \"\\n\")\n return np.squeeze(virtual_queue, axis=0)\n\n\ndef inference(cf, data_path, USE_CUDA):\n checkpoints_dir = cf['data']['checkpoints_dir']\n checkpoints_dir = make_dir(checkpoints_dir, cf)\n files = os.listdir(data_path)\n for f in files:\n if f[-3:] == \"mp4\" and \"no_ois\" not in f and \"no_shutter\" not in f and \"gimbal\" not in f.lower() and \"grid\" not in f.lower() and \"flo\" not in f.lower():\n video_name = f[:-4]\n\n # Define the model\n model = Model(cf) \n load_model = cf[\"model\"][\"load_model\"]\n\n print(\"------Load Pretrined Model--------\")\n if load_model is not None:\n checkpoint = torch.load(load_model)\n print(load_model)\n else:\n load_last = os.path.join(checkpoints_dir, cf['data']['exp']+'_last.checkpoint')\n checkpoint = torch.load(load_last)\n print(load_last)\n model.net.load_state_dict(checkpoint['state_dict'])\n model.unet.load_state_dict(checkpoint['unet'])\n \n if USE_CUDA:\n model.net.cuda()\n model.unet.cuda()\n\n print(\"-----------Load Dataset----------\")\n test_loader = get_inference_data_loader(cf, data_path, no_flo = False)\n data = test_loader.dataset.data[0]\n\n start_time = time.time()\n virtual_queue= run(model, test_loader, cf, USE_CUDA=USE_CUDA)\n\n virtual_data = np.zeros((1,5))\n virtual_data[:,1:] = virtual_queue[0, 1:]\n virtual_data[:,0] = data.frame[0,0]\n virtual_queue = np.concatenate((virtual_data, virtual_queue), axis = 0)\n\n print(virtual_queue.shape)\n time_used = (time.time() - start_time) / 60\n\n print(\"Time_used: %.4f minutes\" % (time_used))\n\n \n virtual_path = os.path.join(\"./test\", cf['data']['exp'], data_path.split(\"/\")[-1]+'.txt')\n np.savetxt(virtual_path, virtual_queue, delimiter=' ')\n\n print(\"------Start Warping Video--------\")\n grid = get_grid(test_loader.dataset.static_options, \\\n data.frame[:data.length], data.gyro, data.ois, virtual_queue[:data.length,1:], no_shutter = False)\n return data, virtual_queue, video_name, grid\n\ndef visual_result(cf, data, video_name, virtual_queue, virtual_queue2 = None, compare_exp = None):\n print(\"------Start Visual Result--------\")\n rotations_virtual, lens_offsets_virtual = get_rotations(data.frame[:data.length], virtual_queue, np.zeros(data.ois.shape), data.length)\n rotations_real, lens_offsets_real = get_rotations(data.frame[:data.length], data.gyro, data.ois, data.length)\n if virtual_queue2 is not None:\n rotations_virtual2, lens_offsets_virtual2 = get_rotations(data.frame[:data.length], virtual_queue2, np.zeros(data.ois.shape), data.length)\n path = os.path.join(\"./test\", cf['data']['exp'], video_name+'_'+compare_exp+'.jpg')\n else:\n rotations_virtual2, lens_offsets_virtual2 = None, None\n path = os.path.join(\"./test\", cf['data']['exp'], video_name+'.jpg')\n \n visual_rotation(rotations_real, lens_offsets_real, rotations_virtual, lens_offsets_virtual, rotations_virtual2, lens_offsets_virtual2, path)\n\n\ndef main(args = None):\n config_file = args.config\n dir_path = args.dir_path\n cf = yaml.load(open(config_file, 'r'))\n\n USE_CUDA = cf['data'][\"use_cuda\"]\n\n log_file = open(os.path.join(cf[\"data\"][\"log\"], cf['data']['exp']+'_test.log'), 'w+')\n printer = Printer(sys.stdout, log_file).open()\n\n data_name = sorted(os.listdir(dir_path))\n for i in range(len(data_name)):\n print(\"Running Inference: \" + str(i+1) + \"/\" + str(len(data_name)))\n save_path = os.path.join(\"./test\", cf['data']['exp'], data_name[i]+'_stab.mp4')\n\n data_path = os.path.join(dir_path, data_name[i])\n data, virtual_queue, video_name, grid= inference(cf, data_path, USE_CUDA)\n\n virtual_queue2 = None\n visual_result(cf, data, data_name[i], virtual_queue, virtual_queue2 = virtual_queue2, compare_exp = None)\n\n video_path = os.path.join(data_path, video_name+\".mp4\")\n warp_video(grid, video_path, save_path, frame_number = False)\n return \n\nif __name__ == '__main__':\n parser = argparse.ArgumentParser(\"Training model\")\n parser.add_argument(\"--config\", default=\"./conf/stabilzation.yaml\", help=\"Config file.\")\n parser.add_argument(\"--dir_path\", default=\"./video\")\n args = parser.parse_args()\n main(args = args)"
},
{
"path": "dvs/load_frame_sensor_data.py",
"content": "import os\nos.environ[\"CUDA_VISIBLE_DEVICES\"] = \"0\"\nimport sys\nimport torch\nimport torchvision\nimport torch.nn as nn\nfrom torch.autograd import Variable\n\nimport time\nimport yaml\nimport argparse\nimport numpy as np\nfrom printer import Printer\nfrom dataset import get_data_loader, get_inference_data_loader\nfrom model import Model\nimport datetime\nimport copy\nfrom util import make_dir, get_optimizer, norm_flow\nfrom gyro import (\n get_grid, \n get_rotations, \n visual_rotation, \n GetGyroAtTimeStamp, \n torch_ConvertQuaternionToAxisAngle, \n torch_ConvertAxisAngleToQuaternion,\n torch_QuaternionProduct,\n get_static\n )\nfrom warp import warp_video\n\ndef run(loader, cf, USE_CUDA=True):\n number_virtual, number_real = cf['data'][\"number_virtual\"], cf['data'][\"number_real\"]\n for i, data in enumerate(loader, 0):\n # get the inputs; data is a list of [inputs, labels]\n real_inputs, times, flo, flo_back, real_projections, real_postion, ois, real_queue_idx = data\n print(\"Fininsh Load data\")\n\n real_inputs = real_inputs.type(torch.float) #[b,60,84=21*4]\n real_projections = real_projections.type(torch.float) \n \n batch_size, step, dim = real_inputs.size()\n times = times.numpy()\n real_queue_idx = real_queue_idx.numpy()\n virtual_queue = [None] * batch_size\n\n for j in range(step):\n virtual_inputs, vt_1 = loader.dataset.get_virtual_data(\n virtual_queue, real_queue_idx, times[:, j], times[:, j+1], times[:, 0], batch_size, number_virtual, real_postion[:,j]) \n real_inputs_step = real_inputs[:,j,:]\n if USE_CUDA:\n real_inputs_step = real_inputs_step.cuda()\n virtual_inputs = virtual_inputs.cuda()\n real_postion_anchor = real_postion[:,j].cuda()\n\n out = real_inputs_step[:,40:44]\n \n virtual_position = virtual_inputs[:, -4:]\n pos = torch_QuaternionProduct(virtual_position, real_postion_anchor)\n\n out = torch_QuaternionProduct(out, pos)\n\n if USE_CUDA:\n out = out.cpu().detach().numpy() \n\n virtual_queue = loader.dataset.update_virtual_queue(batch_size, virtual_queue, out, times[:,j+1])\n return np.squeeze(virtual_queue, axis=0)\n\ndef inference(cf, data_path, USE_CUDA):\n print(\"-----------Load Dataset----------\")\n test_loader = get_inference_data_loader(cf, data_path)\n data = test_loader.dataset.data[0]\n test_loader.dataset.no_flo = True\n test_loader.dataset.static_options = get_static(ratio = 0)\n\n start_time = time.time()\n virtual_queue = run(test_loader, cf, USE_CUDA=USE_CUDA)\n\n virtual_data = np.zeros((1,5))\n virtual_data[:,1:] = virtual_queue[0, 1:]\n virtual_data[:,0] = data.frame[0,0]\n virtual_queue = np.concatenate((virtual_data, virtual_queue), axis = 0)\n \n files = os.listdir(data_path)\n for f in files:\n if f[-3:] == \"mp4\" and \"no_ois\" not in f and \"gimbal\" not in f.lower():\n video_name = f[:-4]\n print(video_name)\n virtual_path = os.path.join(\"./test\", cf['data']['exp'], video_name+'.txt')\n\n print(\"------Start Visual Result--------\")\n rotations_real, lens_offsets_real = get_rotations(data.frame[:data.length], data.gyro, data.ois, data.length)\n fig_path = os.path.join(data_path, video_name+\"_real.jpg\")\n visual_rotation(rotations_real, lens_offsets_real, None, None, None, None, fig_path)\n \n return\n\ndef main(args = None):\n config_file = args.config\n dir_path = args.dir_path\n cf = yaml.load(open(config_file, 'r'))\n \n USE_CUDA = cf['data'][\"use_cuda\"]\n\n checkpoints_dir = cf['data']['checkpoints_dir']\n checkpoints_dir = make_dir(checkpoints_dir, cf)\n\n data_name = sorted(os.listdir(dir_path))\n for i in range(len(data_name)):\n print(\"Running: \" + str(i+1) + \"/\" + str(len(data_name)))\n inference(cf, os.path.join(dir_path, data_name[i]), USE_CUDA)\n return \n\nif __name__ == '__main__':\n parser = argparse.ArgumentParser(\"Training model\")\n parser.add_argument(\"--config\", default=\"./conf/stabilzation.yaml\", help=\"Config file.\")\n parser.add_argument(\"--dir_path\", default=\"./video\")\n args = parser.parse_args()\n main(args = args)"
},
{
"path": "dvs/loss.py",
"content": "import torch\nimport numpy as np\nfrom torch.autograd import Variable\nimport operator\nimport torch.nn.functional as F\nimport matplotlib.pyplot as plt\nfrom gyro import (\n torch_QuaternionProduct, \n torch_QuaternionReciprocal, \n get_static, \n torch_GetVirtualProjection,\n torch_GetForwardGrid,\n torch_GetWarpingFlow,\n torch_ConvertAxisAngleToQuaternion,\n torch_ConvertQuaternionToAxisAngle,\n torch_norm_quat,\n torch_GetHomographyTransformFromProjections,\n torch_ApplyTransform\n)\n \nclass C2_Smooth_loss(torch.nn.Module):\n def __init__(self):\n super(C2_Smooth_loss, self).__init__()\n self.MSE = torch.nn.MSELoss()\n\n def forward(self, Qt, Qt_1, Qt_2):\n detaQt_1 = torch_QuaternionProduct(Qt_1, torch_QuaternionReciprocal(Qt_2))\n return self.MSE(Qt, detaQt_1)\n\nclass C1_Smooth_loss(torch.nn.Module):\n def __init__(self):\n super(C1_Smooth_loss, self).__init__()\n self.MSE = torch.nn.MSELoss()\n\n def forward(self, v_r_axis, v_axis_t_1 = None, real_postion = None):\n quat_zero = torch.zeros(v_r_axis.shape).cuda()\n quat_zero[:,3] = 1\n return self.MSE(v_r_axis, quat_zero)\n\nclass Follow_loss(torch.nn.Module):\n def __init__(self):\n super(Follow_loss, self).__init__()\n self.MSE = torch.nn.MSELoss()\n\n def forward(self, virtual_quat, real_quat, real_postion = None):\n if real_postion is not None:\n real_quat = torch_QuaternionProduct(real_quat, real_postion)\n return self.MSE(virtual_quat, real_quat)\n\nclass Stay_loss(torch.nn.Module):\n def __init__(self):\n super(Stay_loss, self).__init__()\n self.zero = torch.tensor([0.0,0.0,0.0,1.0]).cuda()\n\n def forward(self, virtual_quat):\n return torch.mean(torch.abs(virtual_quat - self.zero))\n\n\nclass Angle_loss(torch.nn.Module):\n def __init__(self):\n super(Angle_loss, self).__init__()\n self.MSE = torch.nn.MSELoss()\n\n def forward(self, Q1, Q2, threshold = 0.5236, logistic_beta1 = 100):\n batch_size = Q1.shape[0]\n Q3 = torch_norm_quat(torch_QuaternionProduct(Q2, torch_QuaternionReciprocal(Q1)))\n theta = torch.zeros(batch_size).cuda()\n index = (Q3[:,3] < 1).nonzero()\n theta[index] = torch.acos(Q3[index,3]) * 2\n loss = torch.mean(theta * (1 / (1 + torch.exp(-logistic_beta1 * (theta - threshold)))))\n return loss, theta\n\nclass Optical_loss(torch.nn.Module):\n def __init__(self):\n super(Optical_loss, self).__init__()\n self.static_options = get_static()\n self.mesh = get_mesh()\n\n def forward(self, Vt, Vt_1, flo, flo_back, real_projection_t, real_projection_t_1):\n virtual_projection_t = torch_GetVirtualProjection(self.static_options, Vt) \n virtual_projection_t_1 = torch_GetVirtualProjection(self.static_options, Vt_1) \n\n b, h, w = flo.size()[:3]\n\n grid_t = torch_GetForwardGrid(self.static_options, real_projection_t, virtual_projection_t)[:,:2,:,:].permute(0,1,3,2)\n grid_t = torch.nn.functional.upsample_bilinear(grid_t, size = (h, w)) # [B,C(xy),H,W]\n\n grid_t_1 = torch_GetForwardGrid(self.static_options, real_projection_t_1, virtual_projection_t_1)[:,:2,:,:].permute(0,1,3,2) \n grid_t_1 = torch.nn.functional.upsample_bilinear(grid_t_1, size = (h, w)) # [B,C(xy),H,W]\n \n mesh = self.mesh.repeat(b, 1, 1, 1)\n flo = flo + mesh \n flo_back = flo_back + mesh # [B,H,W,C]\n\n valid = (flo[:,:,:,0] > 0) * (flo[:,:,:,1] > 0) * (flo[:,:,:,0] < 1) * (flo[:,:,:,1] < 1)\n valid_f = torch.unsqueeze(valid, dim = 3).type(torch.cuda.FloatTensor)\n valid = torch.unsqueeze(valid, dim = 1).type(torch.cuda.FloatTensor)\n\n valid_back = (flo_back[:,:,:,0] > 0) * (flo_back[:,:,:,1] > 0) * (flo_back[:,:,:,0] < 1) * (flo_back[:,:,:,1] < 1)\n valid_back_f = torch.unsqueeze(valid_back, dim = 3).type(torch.cuda.FloatTensor) \n valid_back = torch.unsqueeze(valid_back, dim = 1).type(torch.cuda.FloatTensor) # [B,C,H,W]\n\n flo = (flo * 2 - 1) * valid_f\n flo_back = (flo_back * 2 - 1) * valid_back_f\n\n forward_t = torch.nn.functional.grid_sample(grid_t, flo, padding_mode=\"reflection\") # default bilinear\n backward_t_1 = torch.nn.functional.grid_sample(grid_t_1, flo_back, padding_mode=\"reflection\") # default bilinear\n\n forward_diff = ((forward_t - grid_t_1) * valid) ** 2 \n backward_diff = ((backward_t_1 - grid_t) * valid_back) ** 2\n\n forward_loss = torch.sum(forward_diff, dim = (1,2,3)) / torch.sum(valid, dim = (1,2,3))\n backward_loss = torch.sum(backward_diff, dim = (1,2,3)) / torch.sum(valid_back, dim = (1,2,3))\n\n loss = forward_loss + backward_loss\n loss = torch.min(loss, loss - loss + 1) #[0]\n loss = torch.sum(loss) / b\n\n return loss \n\n\ndef get_mesh(height = 270, width = 480, USE_CUDA = True):\n xs = np.linspace(0, 1, width, endpoint = False) + 0.5 / height\n ys = np.linspace(0, 1, height, endpoint = False) + 0.5 / width\n xmesh, ymesh = np.meshgrid(xs, ys)\n # Reshape the sampling positions to a H x W x 2 tensor\n mesh = torch.Tensor(np.expand_dims(np.moveaxis(np.array(list(zip(xmesh, ymesh))), 1, 2),axis=0))\n if USE_CUDA:\n mesh = mesh.cuda()\n return mesh\n\nclass Undefine_loss(torch.nn.Module):\n def __init__(self, ratio = 0.08, inner_ratio = 0.04, USE_CUDA = True):\n super(Undefine_loss, self).__init__()\n self.static_options = get_static() \n self.inner_ratio = inner_ratio\n width = self.static_options[\"width\"]\n height = self.static_options[\"height\"]\n x0, x1, y0, y1 = \\\n int(width*ratio), int(width*(1-ratio)), int(height*ratio), int(height*(1-ratio))\n self.norm = torch.Tensor([width, height, 1])\n self.p00 = torch.Tensor([x0, y0, 1])\n self.p01 = torch.Tensor([x0, y1, 1])\n self.p10 = torch.Tensor([x1, y0, 1])\n self.p11 = torch.Tensor([x1, y1, 1])\n if USE_CUDA == True:\n self.p00 = self.p00.cuda()\n self.p01 = self.p01.cuda()\n self.p10 = self.p10.cuda()\n self.p11 = self.p11.cuda()\n self.norm = self.norm.cuda()\n\n def forward(self, Vt, Rt, ratio = 0.04):\n batch_size = Vt.size()[0]\n\n row_mid = self.static_options[\"num_grid_rows\"] // 2\n virtual_projection_t = torch_GetVirtualProjection(self.static_options, Vt) \n\n real_projection_t = torch_GetVirtualProjection(self.static_options, Rt) \n\n # virtual projection and real projection\n transform = torch_GetHomographyTransformFromProjections(real_projection_t, virtual_projection_t)\n \n p00 = (torch_ApplyTransform(transform, self.p00) / self.norm)[:,:2]\n p01 = (torch_ApplyTransform(transform, self.p01) / self.norm)[:,:2]\n p10 = (torch_ApplyTransform(transform, self.p10) / self.norm)[:,:2]\n p11 = (torch_ApplyTransform(transform, self.p11) / self.norm)[:,:2]\n\n loss = torch.stack((self.get_loss(p00), self.get_loss(p01), self.get_loss(p10), self.get_loss(p11)),dim = 1)\n loss,_ = torch.max(loss, dim = 1)\n\n loss = torch.min(loss, loss - loss + 1) #[0]\n loss = torch.sum(loss) / batch_size\n\n return loss\n \n def get_loss(self, p):\n d = (p - self.inner_ratio) * (p < self.inner_ratio).type(torch.cuda.FloatTensor) + \\\n (1 - self.inner_ratio - p) * (p > (1 - self.inner_ratio)).type(torch.cuda.FloatTensor)\n return torch.sum(d**2, dim = 1) \n"
},
{
"path": "dvs/metrics.py",
"content": "import os\nimport sys\nimport numpy as np\nimport cv2\nimport math\nimport pdb\nimport matplotlib.pyplot as plt\nfrom printer import Printer\nfrom warp import video2frame_one_seq\nimport datetime\nimport torch\nimport copy\nimport csv\nimport copyreg\nimport shutil\nimport matplotlib.pyplot as plt\nfrom util import crop_video\n\ndef _pickle_keypoints(point):\n return cv2.KeyPoint, (*point.pt, point.size, point.angle,\n point.response, point.octave, point.class_id)\n\ncopyreg.pickle(cv2.KeyPoint().__class__, _pickle_keypoints)\n\nos.environ[\"CUDA_VISIBLE_DEVICES\"] = \"0\"\n\nh_size = 480\nw_size = 640\n\ndef crop_metric(M):\n points = np.array([[0,0,1],[0,h_size,1], [w_size,0,1], [w_size,h_size,1]]).T\n result = np.matmul(M,points).T\n result = result[:,:2]/result[:,2:]\n w_out = 1 - max(result[0,0], result[1,0], w_size - result[2,0], w_size - result[3,0], 0)/w_size\n h_out = 1 - max(result[0,1], result[2,1], h_size - result[1,1], h_size - result[3,1], 0)/h_size\n return w_out, h_out\n\n# https://stackoverflow.com/questions/34389125/how-to-get-the-scale-factor-of-getperspectivetransform-in-opencv\ndef get_scale(M):\n h1 = M[0, 0]\n h2 = M[0, 1]\n h3 = M[0, 2]\n h4 = M[1, 0]\n h5 = M[1, 1]\n h6 = M[1, 2]\n h7 = M[2, 0]\n h8 = M[2, 1]\n QR = np.array([[h1-(h7*h3), h2-(h8*h3)], [h4-(h7*h6), h5-(h8*h6)]])\n Q, R = np.linalg.qr(QR)\n return abs(R[0,0]), abs(R[1,1])\n\n# https://stackoverflow.com/questions/21019338/how-to-change-the-homography-with-the-scale-of-the-image\ndef get_rescale_matrix(M, sx, sy):\n S = np.eye(3, dtype = float)\n S[0,0] = sx\n S[1,1] = sy\n\n S1 = np.eye(3, dtype = float)\n S1[0,0] = 1/sx\n S1[1,1] = 1/sy\n return np.matmul(M, S1)\n\n# Part of code reference from https://github.com/jinsc37/DIFRINT/blob/master/metrics.py\ndef metrics(in_src, out_src, package, crop_scale = False, re_compute = False):\n load_dic = None\n if re_compute and os.path.exists(package):\n print(\"Start load\")\n load_dic = torch.load(package)\n print(\"Finish load\")\n dic = {\n 'M': None,\n 'CR_seq': [],\n 'DV_seq': [],\n 'SS_t': None,\n 'SS_r': None,\n 'w_crop':[],\n 'h_crop':[],\n 'distortion': [],\n 'count': 0,\n 'in_sift': {},\n 'out_sift': {},\n 'fft_t': {},\n 'fft_r': {}\n }\n\n if load_dic is not None:\n dic[\"in_sift\"] = load_dic[\"in_sift\"]\n dic[\"out_sift\"] = load_dic[\"out_sift\"]\n\n frameList_in = sorted(os.listdir(in_src))\n frameList = sorted(os.listdir(out_src))\n frameList = frameList[:min(len(frameList_in),len(frameList))]\n\n # Create brute-force matcher object\n bf = cv2.BFMatcher()\n\n # Apply the homography transformation if we have enough good matches \n MIN_MATCH_COUNT = 10 #10\n\n ratio = 0.7 #0.7\n thresh = 5.0 #5.0\n\n Pt = np.asarray([[1.0,0.0,0.0],[0.0,1.0,0.0],[0.0,0.0,1.0]])\n P_seq = []\n count = 1\n for index, f in enumerate(frameList, 0):\n if f.endswith('.png'):\n # Load the images in gray scale\n img1 = cv2.imread(os.path.join(in_src, f), 0) \n img1 = cv2.resize(img1, (w_size,h_size), interpolation = cv2.INTER_LINEAR)\n\n img1o = cv2.imread(os.path.join(out_src, f), 0)\n img1o = cv2.resize(img1o, (w_size,h_size), interpolation = cv2.INTER_LINEAR)\n sift = cv2.SIFT_create() \n \n if f in dic[\"in_sift\"]:\n keyPoints1, descriptors1 = dic[\"in_sift\"][f]\n else:\n # Detect the SIFT key points and compute the descriptors for the two images\n keyPoints1, descriptors1 = sift.detectAndCompute(img1, None)\n dic[\"in_sift\"][f] = (keyPoints1, descriptors1)\n\n if f in dic[\"out_sift\"]:\n keyPoints1o, descriptors1o = dic[\"out_sift\"][f]\n else:\n keyPoints1o, descriptors1o = sift.detectAndCompute(img1o, None)\n dic[\"out_sift\"][f] = (keyPoints1o, descriptors1o)\n\n # Match the descriptors\n matches = bf.knnMatch(descriptors1, descriptors1o, k=2)\n\n # Select the good matches using the ratio test\n goodMatches = []\n\n for m, n in matches:\n if m.distance < ratio * n.distance:\n goodMatches.append(m)\n\n if len(goodMatches) > MIN_MATCH_COUNT:\n # Get the good key points positions\n sourcePoints = np.float32([ keyPoints1[m.queryIdx].pt for m in goodMatches ]).reshape(-1, 1, 2)\n destinationPoints = np.float32([ keyPoints1o[m.trainIdx].pt for m in goodMatches ]).reshape(-1, 1, 2)\n \n M, mask = cv2.findHomography(sourcePoints, destinationPoints, method=cv2.RANSAC, ransacReprojThreshold=thresh)\n im_dst = cv2.warpPerspective(img1, M, (w_size,h_size)) \n\n cm = []\n for i in range(6):\n for j in range(6):\n hs = int(h_size * (0.2 + 0.1 * i))\n he = int(h_size * (0.3 + 0.1 * i))\n ws = int(w_size * (0.2 + 0.1 * j))\n we = int(w_size * (0.3 + 0.1 * j))\n cm.append(np.corrcoef(img1o[hs:he, ws:we].flat, im_dst[hs:he, ws:we].flat))\n dic[\"distortion\"].append(cm)\n\n if crop_scale:\n sx, sy = get_scale(M)\n M_scale = get_rescale_matrix(M, sx, sy)\n w_crop, h_crop = crop_metric(M_scale)\n else:\n w_crop, h_crop = crop_metric(M)\n dic[\"w_crop\"].append(w_crop)\n dic[\"h_crop\"].append(h_crop)\n\n # Obtain Scale, Translation, Rotation, Distortion value\n sx = M[0, 0]\n sy = M[1, 1]\n scaleRecovered = math.sqrt(np.abs(sx*sy))\n\n w, _ = np.linalg.eig(M[0:2,0:2])\n w = np.sort(w)[::-1]\n DV = w[1]/w[0]\n #pdb.set_trace()\n\n dic[\"CR_seq\"].append(1.0/scaleRecovered)\n dic[\"DV_seq\"].append(DV) \n\n # For Stability score calculation\n if count < len(frameList):\n f_path = f[:-9] + '%05d.png' % (int(f[-9:-4])+1)\n if f_path in dic[\"out_sift\"]:\n keyPoints2o, descriptors2o = dic[\"out_sift\"][f_path]\n else:\n img2o = cv2.imread(os.path.join(out_src, f_path), 0)\n img2o = cv2.resize(img2o, (w_size,h_size), interpolation = cv2.INTER_LINEAR)\n keyPoints2o, descriptors2o = sift.detectAndCompute(img2o, None)\n dic[\"out_sift\"][f_path] = (keyPoints2o, descriptors2o)\n \n matches = bf.knnMatch(descriptors1o, descriptors2o, k=2)\n goodMatches = []\n\n for m, n in matches:\n if m.distance < ratio * n.distance:\n goodMatches.append(m)\n\n if len(goodMatches) > MIN_MATCH_COUNT:\n # Get the good key points positions\n sourcePoints = np.float32([ keyPoints1o[m.queryIdx].pt for m in goodMatches ]).reshape(-1, 1, 2)\n destinationPoints = np.float32([ keyPoints2o[m.trainIdx].pt for m in goodMatches ]).reshape(-1, 1, 2)\n \n # Obtain the homography matrix\n M, mask = cv2.findHomography(sourcePoints, destinationPoints, method=cv2.RANSAC, ransacReprojThreshold=thresh)\n\n P_seq.append(np.matmul(Pt, M))\n Pt = np.matmul(Pt, M)\n if count % 10 ==0:\n sys.stdout.write('\\rFrame: ' + str(count) + '/' + str(len(frameList)))\n sys.stdout.flush()\n dic[\"count\"] = count\n count += 1\n\n # Make 1D temporal signals\n P_seq_t = np.asarray([1])\n P_seq_r = np.asarray([1])\n\n #pdb.set_trace()\n for Mp in P_seq:\n sx = Mp[0, 0]\n sy = Mp[1, 1]\n c = Mp[0, 2]\n f = Mp[1, 2]\n\n transRecovered = math.sqrt(c*c + f*f)\n thetaRecovered = math.atan2(sx, sy) * 180 / math.pi\n\n P_seq_t = np.concatenate((P_seq_t, [transRecovered]), axis=0)\n P_seq_r = np.concatenate((P_seq_r, [thetaRecovered]), axis=0)\n\n P_seq_t = np.delete(P_seq_t, 0)\n P_seq_r = np.delete(P_seq_r, 0)\n\n # FFT\n fft_t = np.fft.fft(P_seq_t)\n fft_r = np.fft.fft(P_seq_r)\n fft_t = abs(fft_t)**2\n fft_r = abs(fft_r)**2\n\n fft_t = np.delete(fft_t, 0)\n fft_r = np.delete(fft_r, 0)\n fft_t = fft_t[:int(len(fft_t)/2)]\n fft_r = fft_r[:int(len(fft_r)/2)]\n\n dic[\"fft_t\"] = fft_t\n dic[\"fft_r\"] = fft_r\n\n SS_t = np.sum(fft_t[:5])/np.sum(fft_t) \n SS_r = np.sum(fft_r[:5])/np.sum(fft_r)\n\n dic[\"CR_seq\"] = np.array(dic[\"CR_seq\"])\n dic[\"DV_seq\"] = np.array(dic[\"DV_seq\"])\n dic[\"w_crop\"] = np.array(dic[\"w_crop\"])\n dic[\"h_crop\"] = np.array(dic[\"h_crop\"])\n dic[\"distortion\"] = np.array(dic[\"distortion\"])\n dic[\"SS_t\"] = SS_t\n dic[\"SS_r\"] = SS_r\n \n if not (re_compute and os.path.exists(package)):\n torch.save(dic, package)\n\n DV_seq = np.absolute(dic[\"DV_seq\"])\n DV_seq = DV_seq[np.where((DV_seq >= 0.5) & (DV_seq <= 1))]\n Distortion = str.format('{0:.4f}', np.nanmin(DV_seq))\n Distortion_avg = str.format('{0:.4f}', np.nanmean(DV_seq))\n\n Trans = str.format('{0:.4f}', dic[\"SS_t\"])\n Rot = str.format('{0:.4f}', dic[\"SS_r\"])\n\n w_crop = crop_rm_outlier(dic[\"w_crop\"])\n h_crop = crop_rm_outlier(dic[\"h_crop\"])\n\n FOV = str.format( '{0:.4f}', min(np.nanmin(w_crop), np.nanmin(h_crop)) )\n FOV_avg = str.format( '{0:.4f}', (np.nanmean(w_crop)+np.nanmean(h_crop)) / 2 )\n\n Correlation_avg = str.format( '{0:.4f}', np.nanmean(dic[\"distortion\"][10:]) )\n Correlation_min = str.format( '{0:.4f}', np.nanmin(dic[\"distortion\"][10:]) )\n\n # Print results\n print('\\n***Distortion value (Avg, Min):')\n print(Distortion_avg +' | '+ Distortion)\n print('***Stability Score (Avg, Trans, Rot):')\n print(str.format('{0:.4f}', (dic[\"SS_t\"]+dic[\"SS_r\"])/2) +' | '+ Trans +' | '+ Rot )\n print(\"=================\")\n print('***FOV ratio (Avg, Min):')\n print( FOV_avg +' | '+ FOV )\n print('***Correlation value (Avg, Min):')\n print( Correlation_avg +' | '+ Correlation_min , \"\\n\") \n\n dic['in_sift'] = 0\n dic['out_sift'] = 0\n torch.save(dic, package[:-3]+\"_light.pt\") \n return float(FOV)\n\ndef crop_rm_outlier(crop):\n crop = np.array(crop)\n crop = crop[crop >= 0.5]\n return sorted(crop)[5:]\n\nif __name__ == '__main__':\n metric_path = os.path.join(\"./test/stabilzation/metric\")\n if not os.path.exists(metric_path):\n os.makedirs(metric_path)\n\n in_video = \"./video/s_114_outdoor_running_trail_daytime/ControlCam_20200930_104820.mp4\"\n in_folder = os.path.join(metric_path, \"in_frame\")\n if not os.path.exists(in_folder):\n os.makedirs(in_folder)\n print(\"Convert video to frames\")\n video2frame_one_seq(in_video, in_folder)\n \n out_video = \"./test/stabilzation/s_114_outdoor_running_trail_daytime_stab.mp4\"\n out_folder = os.path.join(metric_path, \"out_frame\")\n if not os.path.exists(out_folder):\n os.makedirs(out_folder)\n print(\"Convert video to frames\")\n video2frame_one_seq(out_video, out_folder)\n \n package = os.path.join(metric_path, \"stabilzation.pt\")\n FOV = metrics(in_folder, out_folder, package)\n\n crop_path = out_video[:-4] + \"_crop.mp4\"\n crop_video(out_video, crop_path, FOV)\n"
},
{
"path": "dvs/model.py",
"content": "import math\nimport torch\nfrom collections import OrderedDict\n\nimport torch.nn as nn\nimport numpy as np\nimport util\nimport yaml\nimport os\nfrom loss import C2_Smooth_loss, C1_Smooth_loss, Optical_loss, Undefine_loss, Angle_loss, Follow_loss, Stay_loss\nfrom gyro import torch_norm_quat, torch_QuaternionProduct\nimport torch.nn.functional as F\n\nActivates = {\"sigmoid\": nn.Sigmoid, \"relu\": nn.ReLU, \"tanh\": nn.Tanh}\n\nclass LayerLSTM(nn.Module):\n def __init__(self, input_size, hidden_size, bias):\n super(LayerLSTM, self).__init__()\n self.LSTM = nn.LSTMCell(input_size, hidden_size, bias)\n self.hidden_size = hidden_size\n \n def init_hidden(self, batch_size):\n self.hx = torch.zeros((batch_size, self.hidden_size)).cuda()\n self.cx = torch.zeros((batch_size, self.hidden_size)).cuda()\n\n def forward(self, x):\n self.hx, self.cx = self.LSTM(x, (self.hx, self.cx))\n return self.hx\n \n\nclass LayerCNN(nn.Module):\n def __init__(self, in_channel, out_channel, kernel_size, stride, padding, pooling_size=None, \n activation_function=nn.ReLU, batch_norm=True):\n super(LayerCNN, self).__init__()\n self.conv = nn.Conv2d(in_channel, out_channel, kernel_size=kernel_size, stride=stride, padding=padding)\n self.batch_norm = nn.BatchNorm2d(out_channel) if batch_norm else None\n self.activation = activation_function(inplace=True)\n if pooling_size is not None:\n self.pooling = nn.MaxPool2d(pooling_size)\n else:\n self.pooling = None\n \n def forward(self, x):\n x = self.conv(x) #x->[batch,channel,height,width]\n if self.batch_norm is not None:\n x = self.batch_norm(x)\n x = self.activation(x)\n if self.pooling is not None:\n x = self.pooling(x)\n return x\n\nclass LayerFC(nn.Module):\n def __init__(self, in_features, out_features, bias, drop_out=0, activation_function=nn.ReLU, batch_norm = False):\n super(LayerFC, self).__init__()\n self.fc = nn.Linear(in_features, out_features, bias=bias)\n # self.activation = activation_function(inplace=True) if activation_function is not None else None\n self.activation = activation_function() if activation_function is not None else None\n self.dropout = nn.Dropout(p=drop_out,inplace=False) if drop_out else None\n self.batch_norm = nn.BatchNorm1d(out_features) if batch_norm else None\n \n def forward(self, x):\n if self.dropout is not None:\n x = self.dropout(x)\n x = self.fc(x)\n if self.batch_norm is not None:\n x = self.batch_norm(x)\n if self.activation is not None:\n x = self.activation(x)\n return x\n\nclass Net(nn.Module):\n def __init__(self, cf):\n super(Net, self).__init__()\n self.cnn_param = cf[\"model\"][\"cnn\"]\n self.rnn_param = cf[\"model\"][\"rnn\"]\n self.fc_param = cf[\"model\"][\"fc\"]\n self.unit_size = 4\n self.no_flo = False\n\n if self.no_flo is False:\n self._rnn_input_size = (2*cf[\"data\"][\"number_real\"]+1+cf[\"data\"][\"number_virtual\"]) * 4 + 64\n else:\n self._rnn_input_size = (2*cf[\"data\"][\"number_real\"]+1+cf[\"data\"][\"number_virtual\"]) * self.unit_size\n\n #CNN Layers\n cnns = []\n cnn_activation = Activates[self.cnn_param[\"activate_function\"]]\n cnn_batch_norm = self.cnn_param[\"batch_norm\"]\n cnn_layer_param = self.cnn_param[\"layers\"]\n if cnn_layer_param is not None:\n cnn_layers = len(cnn_layer_param)\n for layer in range(cnn_layers):\n in_channel = eval(cnn_layer_param[layer][0])[0]\n out_channel = eval(cnn_layer_param[layer][0])[1]\n kernel_size = eval(cnn_layer_param[layer][1])\n stride = eval(cnn_layer_param[layer][2])\n padding = eval(cnn_layer_param[layer][3])\n pooling_size = eval(cnn_layer_param[layer][4])\n\n cnn = None\n cnn = LayerCNN(in_channel, out_channel, kernel_size, stride, padding, pooling_size, \n activation_function=cnn_activation, batch_norm=cnn_batch_norm)\n cnns.append(('%d' % layer, cnn))\n \n self._rnn_input_size = int(math.floor((self._rnn_input_size+2*padding[1]-kernel_size[1])/stride[1])+1)\n if pooling_size is not None:\n self._rnn_input_size = int(math.floor((self._rnn_input_size-pooling_size[1])/pooling_size[1])+1)\n self.convs = nn.Sequential(OrderedDict(cnns))\n\n else:\n self.convs = None\n out_channel = cf[\"data\"][\"channel_size\"]\n \n self.gap = nn.AvgPool2d(self._rnn_input_size) if self.cnn_param[\"gap\"] else None\n self._rnn_input_size = out_channel if self.cnn_param[\"gap\"] else out_channel*(self._rnn_input_size)\n\n #RNN Layers\n rnns = []\n rnn_layer_param = self.rnn_param[\"layers\"]\n rnn_layers = len(rnn_layer_param)\n \n for layer in range(rnn_layers):\n if layer:\n rnn = LayerLSTM(rnn_layer_param[layer-1][0], rnn_layer_param[layer][0], rnn_layer_param[layer][1])\n else:\n rnn = LayerLSTM(self._rnn_input_size, rnn_layer_param[layer][0], rnn_layer_param[layer][1])\n rnns.append(('%d'%layer, rnn))\n self.rnns = nn.Sequential(OrderedDict(rnns))\n\n self._fc_input_size = rnn_layer_param[rnn_layers-1][0] #* 2 # ois\n \n #FC Layers\n fcs = []\n fc_activation = Activates[self.fc_param[\"activate_function\"]]\n fc_batch_norm = self.fc_param[\"batch_norm\"]\n fc_layer_param = self.fc_param[\"layers\"]\n fc_drop_out = self.fc_param[\"drop_out\"]\n fc_layers = len(fc_layer_param)\n \n if fc_layers == 1:\n fc = LayerFC(self._fc_input_size,fc_layer_param[0][0],fc_layer_param[0][1],\n fc_drop_out, None, fc_batch_norm)\n fcs.append(('%d'%(fc_layers-1), fc))\n else:\n for layer in range(fc_layers-1):\n if layer:\n fc = LayerFC(fc_layer_param[layer-1][0],fc_layer_param[layer][0],fc_layer_param[layer][1],\n fc_drop_out, fc_activation, fc_batch_norm)\n else:\n fc = LayerFC(self._fc_input_size,fc_layer_param[layer][0],fc_layer_param[layer][1],\n fc_drop_out,fc_activation, fc_batch_norm)\n fcs.append(('%d'%layer, fc))\n fc = LayerFC(fc_layer_param[fc_layers-2][0],fc_layer_param[fc_layers-1][0],fc_layer_param[fc_layers-1][1],\n fc_drop_out,None, fc_batch_norm) # Modified\n fcs.append(('%d'%(fc_layers-1), fc))\n\n self.class_num = fc_layer_param[fc_layers-1][0]\n self.fcs = nn.Sequential(OrderedDict(fcs))\n\n def init_hidden(self, batch_size):\n for i in range(len(self.rnns)):\n self.rnns[i].init_hidden(batch_size)\n\n def forward(self, x, flo, ois):\n b,c = x.size() #x->[batch,channel,height,width]\n if self.convs is not None:\n x = self.convs(x)\n if self.gap is not None:\n x = self.gap(x)\n x = x.view(b,-1)\n if self.no_flo is False:\n x = torch.cat((x, flo), dim = 1) \n x = self.rnns(x)\n x = self.fcs(x) # [b, 4]\n x = torch_norm_quat(x)\n return x\n\nclass Model():\n def __init__(self, cf):\n super().__init__()\n self.net = Net(cf)\n self.unet = UNet()\n self.init_weights(cf)\n \n self.loss_smooth = C1_Smooth_loss()\n self.loss_follow = Follow_loss()\n self.loss_c2_smooth = C2_Smooth_loss()\n self.loss_optical = Optical_loss()\n self.loss_undefine = Undefine_loss(ratio = 0.08)\n self.loss_angle = Angle_loss()\n self.loss_stay = Stay_loss()\n\n self.loss_smooth_w = cf[\"loss\"][\"smooth\"]\n self.loss_angle_w = cf[\"loss\"][\"angle\"]\n self.loss_follow_w = cf[\"loss\"][\"follow\"]\n self.loss_c2_smooth_w = cf[\"loss\"][\"c2_smooth\"]\n self.loss_undefine_w = cf[\"loss\"][\"undefine\"]\n self.loss_opt_w = cf[\"loss\"][\"opt\"]\n self.loss_stay_w = cf[\"loss\"][\"stay\"]\n\n self.gaussian_weight = np.array([0.072254, 0.071257, 0.068349, 0.063764, 0.057856, 0.051058, 0.043824, 0.036585, 0.029705, 0.023457, 0.01801])\n\n def loss(\n self, out, vt_1, virtual_inputs, real_inputs, flo, flo_back, \n real_projections_t, real_projections_t_1, real_postion_anchor, \n follow = True, undefine = True, optical = True, stay = False\n ):\n unit_size = self.net.unit_size\n mid = real_inputs.size()[1]//(2*unit_size) \n\n Rt = real_inputs[:,unit_size*(mid):unit_size*(mid)+4] \n v_pos = torch_QuaternionProduct(out, virtual_inputs[:, -4:])\n r_pos = torch_QuaternionProduct(v_pos, real_postion_anchor)\n\n loss = torch.zeros(7).cuda()\n if self.loss_follow_w > 0 and follow:\n for i in range(-2,3):\n loss[0] += self.loss_follow_w * self.loss_follow(v_pos, real_inputs[:,unit_size*(i+mid):unit_size*(i+mid)+4], None)\n if self.loss_angle_w > 0 and follow:\n threshold = 6 / 180 * 3.1415926\n loss_angle, theta = self.loss_angle(v_pos, Rt, threshold = threshold)\n loss[1] = self.loss_angle_w * loss_angle\n if self.loss_smooth_w > 0:\n loss_smooth = self.loss_smooth(out)\n loss[2] = self.loss_smooth_w * loss_smooth\n if self.loss_c2_smooth_w > 0: \n loss[3] = self.loss_c2_smooth_w * self.loss_c2_smooth(out, virtual_inputs[:, -4:], virtual_inputs[:, -8:-4])\n if self.loss_undefine_w > 0 and undefine:\n Vt_undefine = v_pos.clone() \n for i in range(0, 10, 2):\n Rt_undefine = real_inputs[:,unit_size*(mid+i):unit_size*(mid+i)+4]\n loss_undefine_w = self.loss_undefine_w * self.gaussian_weight[i]\n loss[4] += loss_undefine_w * self.loss_undefine(Vt_undefine, Rt_undefine)\n Vt_undefine = torch_QuaternionProduct(out, Vt_undefine)\n Vt_undefine = torch_QuaternionProduct(out, Vt_undefine)\n if self.loss_opt_w > 0 and optical:\n loss[5] = self.loss_opt_w * self.loss_optical(r_pos, vt_1, flo, flo_back, real_projections_t, real_projections_t_1) \n if self.loss_stay_w > 0 and stay:\n loss[6] = self.loss_stay_w * self.loss_stay(out) \n return loss\n\n\n def init_weights(self, cf):\n for m in self.net.modules():\n if isinstance(m, nn.Conv1d) or isinstance(m, nn.Conv2d) or isinstance(m, nn.Conv3d) or isinstance(m, nn.Linear):\n if cf[\"train\"][\"init\"] == \"xavier_uniform\":\n nn.init.xavier_uniform_(m.weight.data)\n elif cf[\"train\"][\"init\"] == \"xavier_normal\":\n nn.init.xavier_normal_(m.weight.data)\n\n for m in self.unet.modules():\n if isinstance(m, nn.Conv1d) or isinstance(m, nn.Conv2d) or isinstance(m, nn.Conv3d) or isinstance(m, nn.Linear):\n if cf[\"train\"][\"init\"] == \"xavier_uniform\":\n nn.init.xavier_uniform_(m.weight.data)\n elif cf[\"train\"][\"init\"] == \"xavier_normal\":\n nn.init.xavier_normal_(m.weight.data)\n\n def save_checkpoint(self, epoch = 0, optimizer=None):\n package = {\n 'cnn': self.net.cnn_param,\n 'fc': self.net.fc_param,\n 'state_dict': self.net.state_dict(),\n }\n if optimizer is not None:\n package['optim_dict'] = optimizer.state_dict()\n if self.unet is not None:\n package['unet'] = self.unet.state_dict()\n package[\"epoch\"] = epoch\n return package\n\n\nclass UNet(nn.Module):\n def __init__(self, n_channels = 4, n_classes = 16, bilinear=True):\n super(UNet, self).__init__()\n self.n_channels = n_channels\n self.n_classes = n_classes\n self.bilinear = bilinear\n\n self.inc = DoubleConv(n_channels, 8)\n self.down1 = Down(8, 16)\n self.down2 = Down(16, 32)\n self.down3 = Down(32, 64)\n # factor = 2 if bilinear else 1\n self.down4 = Down(64, 128)\n self._fc_input_size = 128 * 1 * 1\n self.fc = LayerFC(self._fc_input_size, 64, bias = True)\n\n def forward(self, x, x_back = None):\n if x_back is not None:\n x = torch.cat((x,x_back), dim =3)\n x = x.permute(0,3,1,2)\n b,c,h,w = x.size()\n\n x1 = self.inc(x)\n x2 = self.down1(x1)\n x3 = self.down2(x2)\n x4 = self.down3(x3)\n x5 = self.down4(x4)\n x = torch.reshape(x5, (b, -1))\n x = self.fc(x)\n return x\n\n\nclass DoubleConv(nn.Module):\n \"\"\"(convolution => [BN] => ReLU) * 2\"\"\"\n\n def __init__(self, in_channels, out_channels, mid_channels=None):\n super().__init__()\n if not mid_channels:\n mid_channels = out_channels\n self.double_conv = nn.Sequential(\n nn.Conv2d(in_channels, mid_channels, kernel_size=3, padding=1),\n nn.ReLU(inplace=True),\n )\n\n def forward(self, x):\n return self.double_conv(x)\n\n\nclass Down(nn.Module):\n \"\"\"Downscaling with maxpool then double conv\"\"\"\n\n def __init__(self, in_channels, out_channels):\n super().__init__()\n self.maxpool_conv = nn.Sequential(\n nn.MaxPool2d(4),\n DoubleConv(in_channels, out_channels)\n )\n\n def forward(self, x):\n return self.maxpool_conv(x)\n\n\nclass Up(nn.Module):\n \"\"\"Upscaling then double conv\"\"\"\n\n def __init__(self, in_channels, out_channels, bilinear=True):\n super().__init__()\n\n # if bilinear, use the normal convolutions to reduce the number of channels\n if bilinear:\n self.up = nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True)\n self.conv = DoubleConv(in_channels, out_channels, in_channels // 2)\n else:\n self.up = nn.ConvTranspose2d(in_channels , in_channels // 2, kernel_size=2, stride=2)\n self.conv = DoubleConv(in_channels, out_channels)\n\n\n def forward(self, x1, x2):\n x1 = self.up(x1)\n # input is CHW\n diffY = x2.size()[2] - x1.size()[2]\n diffX = x2.size()[3] - x1.size()[3]\n\n x1 = F.pad(x1, [diffX // 2, diffX - diffX // 2,\n diffY // 2, diffY - diffY // 2])\n # if you have padding issues, see\n # https://github.com/HaiyongJiang/U-Net-Pytorch-Unstructured-Buggy/commit/0e854509c2cea854e247a9c615f175f76fbb2e3a\n # https://github.com/xiaopeng-liao/Pytorch-UNet/commit/8ebac70e633bac59fc22bb5195e513d5832fb3bd\n x = torch.cat([x2, x1], dim=1)\n return self.conv(x)\n\n\nclass OutConv(nn.Module):\n def __init__(self, in_channels, out_channels):\n super(OutConv, self).__init__()\n self.conv = nn.Conv2d(in_channels, out_channels, kernel_size=1)\n\n def forward(self, x):\n return self.conv(x)"
},
{
"path": "dvs/printer.py",
"content": "import sys\n\nclass Printer(object):\n def __init__(self, *files):\n self.files = files\n \n #Redirect Printer\n def open(self):\n if not hasattr(sys, '_stdout'):\n sys._stdout = sys.stdout\n sys.stdout = self\n return self\n\n #Restore the Default Printer\n def close(self):\n stdout = sys._stdout\n for f in self.files:\n if f != stdout:\n f.close()\n sys.stdout = stdout\n\n #Overloading write() Function\n def write(self, obj):\n for f in self.files:\n f.write(obj)\n f.flush()\n\n def flush(self):\n pass\n\nif __name__ == '__main__':\n print(\"Start testing\")\n t = Printer(sys.stdout, open('./test.txt', 'w+')).open()\n print(\"In files\")\n t.close()\n print(\"Not in files\")"
},
{
"path": "dvs/requirements.txt",
"content": "colorama==0.4.4\nffmpeg==1.4\nimageio==2.9.0\nmatplotlib==3.3.4\nopencv-contrib-python==4.5.1.48\nopencv-python==4.5.1.48\npytz==2021.1\nPyYAML==5.4.1\nscipy==1.5.4\ntensorboardX==2.1\ntqdm==4.59.0"
},
{
"path": "dvs/train.py",
"content": "import os\nimport sys\nimport torch\nimport torchvision\nimport torch.nn as nn\nfrom torch.autograd import Variable\n\nimport time\nimport yaml\nimport argparse\nimport numpy as np\nfrom printer import Printer\nfrom dataset import get_data_loader\nfrom model import Model\nimport datetime\nimport copy\nfrom util import make_dir, get_optimizer, AverageMeter, save_train_info, norm_flow\nfrom gyro import torch_QuaternionProduct, torch_QuaternionReciprocal, torch_norm_quat\n\nos.environ[\"CUDA_VISIBLE_DEVICES\"] = \"0\"\n\ndef run_epoch(model, loader, cf, epoch, lr, optimizer=None, is_training=True, USE_CUDA=True, clip_norm=0):\n no_flo = False\n number_virtual, number_real = cf['data'][\"number_virtual\"], cf['data'][\"number_real\"]\n avg_loss = AverageMeter()\n if is_training:\n model.net.train()\n model.unet.train()\n else:\n model.net.eval()\n model.unet.eval()\n for i, data in enumerate(loader, 0):\n # get the inputs; data is a list of [inputs, labels]\n real_inputs, times, flo, flo_back, real_projections, real_postion, ois, real_queue_idx = data\n print(\"Fininsh Load data\")\n\n real_inputs = real_inputs.type(torch.float) #[b,60,84=21*4]\n real_projections = real_projections.type(torch.float) \n flo = flo.type(torch.float) \n flo_back = flo_back.type(torch.float) \n ois = ois.type(torch.float) \n\n batch_size, step, dim = real_inputs.size()\n times = times.numpy()\n real_queue_idx = real_queue_idx.numpy()\n virtual_queue = loader.dataset.random_init_virtual_queue(batch_size, real_postion[:,0,:].numpy(), times[:,1]) # TODO\n # virtual_queue = [None] * batch_size\n loss = 0\n model.net.init_hidden(batch_size)\n for j in range(step):\n virtual_inputs, vt_1 = loader.dataset.get_virtual_data(\n virtual_queue, real_queue_idx, times[:, j], times[:, j+1], times[:, 0], batch_size, number_virtual, real_postion[:,j]) \n \n real_inputs_step = real_inputs[:,j,:]\n inputs = torch.cat((real_inputs_step,virtual_inputs), dim = 1) \n\n # inputs = Variable(real_inputs_step)\n if USE_CUDA:\n real_inputs_step = real_inputs_step.cuda()\n virtual_inputs = virtual_inputs.cuda()\n inputs = inputs.cuda()\n if no_flo is False:\n flo_step = flo[:,j].cuda()\n flo_back_step = flo_back[:,j].cuda()\n else:\n flo_step = None\n flo_back_step = None\n vt_1 = vt_1.cuda()\n real_projections_t = real_projections[:,j+1].cuda()\n real_projections_t_1 = real_projections[:,j].cuda()\n real_postion_anchor = real_postion[:,j].cuda()\n ois_step = ois[:,j].cuda()\n\n if no_flo is False:\n b, h, w, _ = flo_step.size()\n flo_step = norm_flow(flo_step, h, w)\n flo_back_step = norm_flow(flo_back_step, h, w)\n\n if is_training:\n if no_flo is False:\n flo_out = model.unet(flo_step, flo_back_step)\n else:\n flo_out = None\n\n if j < 1:\n for i in range(2):\n out = model.net(inputs, flo_out, ois_step)\n else:\n out = model.net(inputs, flo_out, ois_step)\n else:\n with torch.no_grad():\n if no_flo is False:\n flo_out = model.unet(flo_step, flo_back_step)\n else:\n flo_out = None\n\n if j < 1:\n for i in range(2):\n out = model.net(inputs, flo_out, ois_step)\n else:\n out = model.net(inputs, flo_out, ois_step)\n\n if epoch <= 30:\n follow = True\n else:\n follow = False\n\n if epoch > 30:\n undefine = True\n else:\n undefine = False\n\n if epoch > 40:\n optical = True\n else:\n optical = False\n \n loss_step = model.loss(out, vt_1, virtual_inputs, real_inputs_step, \\\n flo_step, flo_back_step, real_projections_t, real_projections_t_1, real_postion_anchor, \\\n follow = follow, undefine = undefine, optical = optical, stay = optical)\n\n loss = loss_step\n \n virtual_position = virtual_inputs[:, -4:]\n pos = torch_QuaternionProduct(virtual_position, real_postion_anchor)\n out = torch_QuaternionProduct(out, pos)\n\n if USE_CUDA:\n out = out.cpu().detach().numpy() \n\n virtual_queue = loader.dataset.update_virtual_queue(batch_size, virtual_queue, out, times[:,j+1])\n\n if (j+1) % 10 == 0:\n print(\"Step: \"+str(j+1)+\"/\"+str(step))\n print(loss)\n loss = torch.sum(loss)\n if is_training:\n optimizer.zero_grad()\n loss.backward(retain_graph=True)\n if clip_norm:\n nn.utils.clip_grad_norm_(model.net.parameters(), max_norm=clip_norm)\n nn.utils.clip_grad_norm_(model.unet.parameters(), max_norm=clip_norm)\n optimizer.step()\n\n avg_loss.update(loss.item(), batch_size) \n \n return avg_loss.avg\n\n\ndef train(args = None):\n torch.autograd.set_detect_anomaly(True)\n config_file = args.config\n cf = yaml.load(open(config_file, 'r'))\n \n USE_CUDA = cf['data'][\"use_cuda\"]\n seed = cf['train'][\"seed\"]\n \n torch.manual_seed(seed)\n if USE_CUDA:\n torch.cuda.manual_seed(seed)\n\n checkpoints_dir = cf['data']['checkpoints_dir']\n epochs = cf[\"train\"][\"epoch\"]\n snapshot = cf[\"train\"][\"snapshot\"]\n decay_epoch = cf['train']['decay_epoch']\n init_lr = cf[\"train\"][\"init_lr\"]\n lr_decay = cf[\"train\"][\"lr_decay\"]\n lr_step = cf[\"train\"][\"lr_step\"]\n clip_norm = cf[\"train\"][\"clip_norm\"]\n load_model = cf[\"model\"][\"load_model\"]\n\n checkpoints_dir = make_dir(checkpoints_dir, cf)\n\n if load_model is None:\n log_file = open(os.path.join(cf[\"data\"][\"log\"], cf['data']['exp']+'.log'), 'w+')\n else:\n log_file = open(os.path.join(cf[\"data\"][\"log\"], cf['data']['exp']+'.log'), 'a')\n printer = Printer(sys.stdout, log_file).open()\n \n print('----Print Arguments Setting------') \n for key in cf:\n print('{}:'.format(key))\n for para in cf[key]:\n print('{:50}:{}'.format(para,cf[key][para]))\n print('\\n')\n\n # Define the model\n model = Model(cf) \n optimizer = get_optimizer(cf[\"train\"][\"optimizer\"], model, init_lr, cf)\n\n for idx, m in enumerate(model.net.children()):\n print('{}:{}'.format(idx,m))\n for idx, m in enumerate(model.unet.children()):\n print('{}:{}'.format(idx,m))\n\n if load_model is not None:\n print(\"------Load Pretrined Model--------\")\n checkpoint = torch.load(load_model)\n model.net.load_state_dict(checkpoint['state_dict'])\n model.unet.load_state_dict(checkpoint['unet'])\n print(\"------Resume Training Process-----\")\n optimizer.load_state_dict(checkpoint['optim_dict'])\n epoch_load = checkpoint['epoch']\n print(\"Epoch load: \", epoch_load)\n else:\n epoch_load = 0\n \n if USE_CUDA:\n model.net.cuda()\n model.unet.cuda()\n if load_model is not None:\n for state in optimizer.state.values():\n for k, v in state.items():\n if isinstance(v, torch.Tensor):\n state[k] = v.cuda()\n for param in optimizer.param_groups:\n init_lr = param['lr']\n\n print(\"-----------Load Dataset----------\")\n train_loader, test_loader = get_data_loader(cf, no_flo = False)\n\n print(\"----------Start Training----------\")\n currentDT = datetime.datetime.now()\n print(currentDT.strftime(\" %Y-%m-%d %H:%M:%S\"))\n \n start_time = time.time()\n\n if lr_step:\n decay_epoch = list(range(1+lr_step, epochs+1, lr_step))\n \n lr = init_lr\n\n for count in range(epoch_load+1, epochs+1):\n if decay_epoch != None and count in decay_epoch:\n lr *= lr_decay\n for param in optimizer.param_groups:\n param['lr'] *= lr_decay\n \n print(\"Epoch: %d, learning_rate: %.5f\" % (count, lr))\n\n train_loss = run_epoch(model, train_loader, cf, count, lr, optimizer=optimizer, clip_norm=clip_norm, is_training=True, USE_CUDA=USE_CUDA)\n\n test_loss = run_epoch(model, test_loader, cf, count, lr, is_training=False, USE_CUDA=USE_CUDA)\n\n time_used = (time.time() - start_time) / 60\n print(\"Epoch %d done | TrLoss: %.4f | TestLoss: %.4f | Time_used: %.4f minutes\" % (\n count, train_loss, test_loss, time_used))\n\n if count % snapshot == 0:\n save_train_info(\"epoch\", checkpoints_dir, cf, model, count, optimizer)\n save_train_info(\"last\", checkpoints_dir, cf, model, count, optimizer)\n print(\"Model stored at epoch %d\"%count)\n\n currentDT = datetime.datetime.now()\n print(currentDT.strftime(\" %Y-%m-%d %H:%M:%S\"))\n print(\"------------End Training----------\")\n return \n\nif __name__ == '__main__':\n parser = argparse.ArgumentParser(\"Training model\")\n parser.add_argument(\"--config\", default=\"./conf/stabilzation_train.yaml\", help=\"Config file.\")\n args = parser.parse_args()\n train(args = args)"
},
{
"path": "dvs/util.py",
"content": "import os\nimport torch\nimport cv2\nfrom itertools import chain\nfrom warp import load_video, save_video\nimport numpy as np\nimport matplotlib.pyplot as plt\nfrom gyro import get_rotations\nimport shutil\n\ndef save_train_info(name, checkpoints_dir, cf, model, count, optimizer = None):\n path = None\n if name == \"last\":\n path = os.path.join(checkpoints_dir, cf['data']['exp']+'_last.checkpoint')\n elif name == \"best\":\n path = os.path.join(checkpoints_dir, cf['data']['exp']+'_best.checkpoint')\n else:\n path = os.path.join(checkpoints_dir, cf['data']['exp']+'_epoch%d.checkpoint'%count)\n torch.save(model.save_checkpoint(epoch = count, optimizer=optimizer), path)\n\ndef make_dir(checkpoints_dir ,cf):\n inference_path = \"./test\"\n if not os.path.exists(checkpoints_dir):\n os.makedirs(checkpoints_dir)\n if not os.path.exists(cf[\"data\"][\"log\"]):\n os.makedirs(cf[\"data\"][\"log\"])\n if not os.path.exists(inference_path):\n os.makedirs(inference_path)\n \n inference_path = os.path.join(inference_path, cf['data']['exp'])\n if not os.path.exists(inference_path):\n os.makedirs(inference_path)\n checkpoints_dir = os.path.join(checkpoints_dir, cf['data']['exp'])\n if not os.path.exists(checkpoints_dir):\n os.makedirs(checkpoints_dir)\n return checkpoints_dir\n\ndef get_optimizer(optimizer, model, init_lr, cf):\n if optimizer == \"adam\":\n optimizer = torch.optim.Adam(chain(model.net.parameters(), model.unet.parameters()), lr=init_lr, weight_decay=cf[\"train\"][\"weight_decay\"])\n elif optimizer == \"sgd\":\n optimizer = torch.optim.SGD(chain(model.net.parameters(), model.unet.parameters()), lr=init_lr, momentum=cf[\"train\"][\"momentum\"])\n return optimizer\n\ndef crop_video(in_path, out_path, crop_ratio):\n frame_array, fps, size = load_video(in_path)\n hs = int((1-crop_ratio)*1080) + 1\n he = int(crop_ratio*1080) - 1\n ws = int((1-crop_ratio)*1920) + 1\n we = int(crop_ratio*1920) - 1\n for i in range(len(frame_array)):\n frame_array[i] = cv2.resize(frame_array[i][hs:he,ws:we,:], size, interpolation = cv2.INTER_LINEAR)\n save_video(out_path, frame_array, fps, size= size)\n\ndef norm_flow(flow, h, w):\n if flow.shape[2] == 2:\n flow[:,:,0] /= h\n flow[:,:,1] /= w\n else:\n flow[:,:,:,0] /= h\n flow[:,:,:,1] /= w\n return flow\n\nclass AverageMeter(object):\n def __init__(self):\n self.reset()\n\n def reset(self):\n self.avg = 0\n self.sum = 0\n self.cnt = 0\n\n def update(self, val, n=1):\n self.sum += val * n\n self.cnt += n\n if self.cnt > 0:\n self.avg = self.sum / self.cnt"
},
{
"path": "dvs/warp/__init__.py",
"content": "from .warping import (\n warp_video\n )\nfrom .read_write import (\n save_video,\n load_video,\n video2frame_one_seq\n )"
},
{
"path": "dvs/warp/rasterizer.py",
"content": "import numpy as np\nimport matplotlib.pyplot as plt\nfrom numpy import array\nimport torch\nimport cv2\nimport time\n\ndevice = torch.device(\"cuda\")\n\ndef Rasterization(image, grid, get_mesh_only = False):\n # grid xy WH\n shape = image.size()\n height = shape[1]\n width = shape[2]\n wapper_upper_triangle, wapper_lower_triangle = grid_to_triangle(grid[:,:,:2])\n origin_upper_triangle, origin_lower_triangle = grid_to_triangle(grid[:,:,2:])\n\n\n [xmax, xmin, ymax, ymin], xlength, ylength = grid_size(wapper_upper_triangle, wapper_lower_triangle, height, width)\n\n xratio = xlength / width\n yratio = ylength / height\n\n wapper_triangle = torch.stack((wapper_upper_triangle,wapper_lower_triangle),dim = 1).to(device) # grid * upper/lower * point * xy\n origin_triangle = torch.stack((origin_upper_triangle,origin_lower_triangle),dim = 1).to(device) # grid * upper/lower * point * xy\n\n tran_triangle = torch.zeros(wapper_triangle.size()).to(device)\n\n tran_triangle[:,:,:,0] = (wapper_triangle[:,:,:,0] - xmin.view(-1,1,1).to(device)/width) / xratio\n tran_triangle[:,:,:,1] = (wapper_triangle[:,:,:,1] - ymin.view(-1,1,1).to(device)/height) / yratio\n\n mask = triangle2mask(tran_triangle, ylength, xlength) # consuming\n\n mask = torch.unsqueeze(mask, 4)\n origin_triangle = torch.unsqueeze(origin_triangle, 1)\n\n grid_sample = origin_triangle * mask # consuming\n grid_sample = torch.sum(torch.sum(grid_sample, dim = 3), dim = 2).view(-1,ylength,xlength,2) # consuming\n\n gxmin = min(0, int(torch.min(xmin)))\n gxmax = int(torch.max(xmin) + xlength)\n gymin = min(0, int(torch.min(ymin)))\n gymax = int(torch.max(ymin) + ylength)\n grid_merge = torch.zeros((max(gymax-gymin, height, height - gymin),max(gxmax - gxmin, width, width - gxmin),2)).to(device)\n for i in range(grid_sample.size()[0]):\n x_s = int(xmin[i] - gxmin)\n x_e = int(xmin[i] + xlength - gxmin)\n y_s = int(ymin[i] - gymin)\n y_e = int(ymin[i] + ylength -gymin)\n grid_merge[ y_s:y_e, x_s:x_e, :] += grid_sample[i, :, :, :]\n\n # grid_merge = grid_merge[min(-gxmin,0):min(-gxmin,0)+height, min(-gymin,0):min(-gymin,0)+width, :] \n grid_merge = grid_merge[-gymin:-gymin+height, -gxmin:-gxmin+width, :] \n # if get_mesh_only:\n # grid_merge = grid_merge.cpu().numpy()\n # mesh_grid = generate_mesh_grid(height, width)\n # out = grid_merge - mesh_grid\n # return np.concatenate((out[:,:,1:],out[:,:,:1]),2)\n \n shift = torch.tensor([0.5/height,0.5/width])[None, None, :].to(device)\n grid_merge = (grid_merge + 1*shift) * 2 - 1\n\n image[:3,:2,:2] = 0\n\n image = torch.unsqueeze(image, 0).to(device)\n grid_merge = torch.unsqueeze(grid_merge, 0)\n\n image = torch.nn.functional.grid_sample(image, grid_merge) # default bilinear\n\n image = torch.squeeze(image, 0)\n return image.cpu()\n\ndef grid_to_triangle(grid):\n grid_shape = grid.size()\n num = (grid_shape[0] - 1) * (grid_shape[1] - 1)\n\n upper_triangle = grid[:-1, :-1, :, None]\n upper_triangle = torch.cat(( upper_triangle, grid[1:, :-1, :, None]), dim = 3)\n upper_triangle = torch.cat(( upper_triangle, grid[:-1, 1:, :, None]), dim = 3)\n upper_triangle = upper_triangle.view(num, 2, 3)\n upper_triangle = torch.transpose(upper_triangle, 1, 2) # grid * point * xy\n \n lower_triangle = grid[:-1, 1:, :, None]\n lower_triangle = torch.cat(( lower_triangle, grid[1:, :-1, :, None]), dim = 3)\n lower_triangle = torch.cat(( lower_triangle, grid[1:, 1:, :, None]), dim = 3)\n lower_triangle = lower_triangle.view(num, 2, 3)\n lower_triangle = torch.transpose(lower_triangle, 1, 2)\n \n return upper_triangle, lower_triangle # grid * point * xy\n\ndef grid_size(upper_triangle, lower_triangle, height, width):\n wapper_grid = torch.cat((upper_triangle, lower_triangle),dim =1)\n xmax = torch.floor(torch.max(wapper_grid[:,:,0]*width, 1)[0]) + 1\n ymax = torch.floor(torch.max(wapper_grid[:,:,1]*height, 1)[0]) + 1\n xmin = torch.floor(torch.min(wapper_grid[:,:,0]*width, 1)[0])\n ymin = torch.floor(torch.min(wapper_grid[:,:,1]*height, 1)[0])\n\n xlength = int(torch.max(xmax - xmin))\n ylength = int(torch.max(ymax - ymin))\n\n return [xmax, xmin, ymax, ymin], xlength, ylength\n\ndef generate_mesh_grid(height, width):\n # Create a grid of sampling positions\n xs = np.linspace(0, 1, width, endpoint=False)\n ys = np.linspace(0, 1, height, endpoint=False)\n xmesh, ymesh = np.meshgrid(xs, ys)\n # Reshape the sampling positions to a H x W x 2 tensor\n return np.moveaxis(array(list(zip(xmesh, ymesh))), 1, 2)\n\ndef triangle2mask(d, height, width): # d: [N x T x 3 x 2]\n N = d.size()[0] # batch size\n T = d.size()[1] # triangle number\n P = height * width # The number of pixels in the output image.\n\n area = edgefunc(d[:, :, 1, :], d[:, :, 2, :], d[:, :, None, 0, :])\n\n gridcpu = generate_mesh_grid(height, width)\n \n gridcpu = np.reshape(gridcpu, (height*width, 2))\n\n grid = torch.Tensor(gridcpu)\n grid = grid.unsqueeze(0).repeat((N, T, 1, 1)) # [N x T x P x 2]\n\n grid = grid.to(device)\n\n # Evaluate the edge functions at every position.\n # We should get a [N x P] vector out of each.\n w0 = edgefunc(d[:, :, 1, :], d[:, :, 2, :], grid) / area\n w1 = edgefunc(d[:, :, 2, :], d[:, :, 0, :], grid) / area\n w2 = edgefunc(d[:, :, 0, :], d[:, :, 1, :], grid) / area\n\n # Only pixels inside the triangles will have color\n # [N x P]\n\n mask = (w0 > 0) & (w1 > 0) & (w2 > 0)\n mask = torch.unsqueeze(mask, 3).type(torch.cuda.FloatTensor)\n\n w = torch.stack((w0,w1,w2),dim = 3) * mask\n\n return torch.transpose(w, 1, 2) # [N x P x T x 3]\n \n\ndef edgefunc(v0, v1, p):\n \"\"\"\n let P = H * W\n v0 and v1 have vertex positions for all T triangles.\n Their shapes are [N x T X 2]\n p is a list of sampling points as a [N x T X P x 2] tensor.\n Each of the T triangles has an [P x 2] matrix of sampling points.\n returns a [N x T x P] matrix\n \"\"\"\n P = p.size()[2]\n \n # Take all the x and y coordinates of all the positions as a\n # [N x S] tensor\n py = p[:, :, :, 1]\n px = p[:, :, :, 0]\n\n # We need to manually broadcast the vector to cover all sample points\n x10 = v0[:, :, 0] - v1[:, :, 0] # [N x T]\n y01 = v1[:, :, 1] - v0[:, :, 1] # [N x T]\n\n x10 = x10.unsqueeze(2).repeat((1, 1, P)) # [N x T x P]\n y01 = y01.unsqueeze(2).repeat((1, 1, P)) # [N x T x P]\n\n cross = v0[:,:,1]*v1[:,:,0] - v0[:,:,0]*v1[:,:,1] # [N x T]\n cross = cross.unsqueeze(2).repeat((1, 1, P)) # [N x T x P]\n\n return y01*px + x10*py + cross\n\nif __name__ == '__main__':\n print(generate_mesh_grid(2,3))"
},
{
"path": "dvs/warp/read_write.py",
"content": "import numpy as np\nimport cv2\nimport os\nfrom PIL import Image, ImageDraw, ImageFont\nimport matplotlib.pyplot as plt\nimport ffmpeg \nimport json\nimport torch \nimport argparse\n \ndef load_video(path, save_dir = None, resize = None, length = -1): # N x H x W x C\n vidcap = cv2.VideoCapture(path)\n fps = vidcap.get(cv2.CAP_PROP_FPS)\n success,image = vidcap.read()\n print(image.shape)\n height, width, layers = image.shape\n if resize is None:\n size = (width,height)\n elif type(resize) is int:\n size = (width//resize,height//resize)\n else:\n size = resize\n count = 0\n frames = []\n while success: \n if resize is not None:\n image = cv2.resize(image, size, interpolation = cv2.INTER_LINEAR)\n if save_dir != None:\n path = os.path.join(save_dir, \"frame_\" + str(count).zfill(4) + \".png\")\n cv2.imwrite(path, image) \n frames.append(image)\n success,image = vidcap.read()\n count += 1\n if length > 0 and count >= length:\n break\n print(\"Video length: \", len(frames))\n return frames, fps, size\n\ndef video2frame(path, resize = None):\n data_name = sorted(os.listdir(path))\n for i in range(len(data_name)):\n print(str(i+1)+\" / \" + str(len(data_name)))\n data_folder = os.path.join(path, data_name[i])\n print(data_folder)\n files = os.listdir(data_folder)\n for f in files:\n if f[-4:] == \".mp4\":\n video_name = f\n video_path = os.path.join(data_folder, video_name)\n frame_folder = os.path.join(data_folder, \"frames\")\n if not os.path.exists(frame_folder):\n os.makedirs(frame_folder)\n load_video(video_path, save_dir = frame_folder, resize=resize)\n\ndef video2frame_one_seq(path, save_dir = None, resize = None): # N x H x W x C\n vidcap = cv2.VideoCapture(path)\n fps = vidcap.get(cv2.CAP_PROP_FPS)\n success,image = vidcap.read()\n print(path)\n print(image.shape)\t\n height, width, layers = image.shape\n if resize is None:\n size = (width,height)\n elif type(resize) is int:\n size = (width//resize,height//resize)\n else:\n size = resize\n count = 0\n while success: \n if resize is not None:\n image = cv2.resize(image, size, interpolation = cv2.INTER_LINEAR)\n if save_dir != None:\n path = os.path.join(save_dir, \"frame_\" + str(count).zfill(5) + \".png\")\n cv2.imwrite(path, image) \n success,image = vidcap.read()\n count += 1\n return fps, size\n\ndef save_video(path,frame_array, fps, size, losses = None, frame_number = False, writer = None):\n if writer is None:\n if path[-3:] == \"mp4\":\n out = cv2.VideoWriter(path,cv2.VideoWriter_fourcc(*'mp4v'), fps, size)\n else:\n out = cv2.VideoWriter(path,cv2.VideoWriter_fourcc('M','J','P','G'), fps, size)\n else:\n out = writer\n for i in range(len(frame_array)):\n # writing to a image array\n if frame_number:\n frame_array[i] = draw_number(np.asarray(frame_array[i]), i)\n if losses is not None:\n frame_array[i] = draw_number(np.asarray(frame_array[i]), losses[i], x = 900, message = \"Loss: \")\n out.write(frame_array[i])\n if writer is None:\n out.release()\n\ndef draw_number(frame, num, x = 10, y = 10, message = \"Frame: \"):\n image=Image.fromarray(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))\n draw = ImageDraw.Draw(image)\n font = ImageFont.truetype(\"./data/arial.ttf\", 45)\n \n message = message + str(num)\n color = 'rgb(0, 0, 0)' # black color\n \n draw.text((x, y), message, fill=color, font=font)\n return cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR)\n\nif __name__ == \"__main__\":\n parser = argparse.ArgumentParser(\"FlowNet2 Preparation\")\n parser.add_argument(\"--dir_path\", default=\"./video\")\n args = parser.parse_args()\n dir_path = args.dir_path\n if dir_path == \"./video\":\n video2frame(dir_path, resize = 4)\n else:\n video2frame(os.path.join(dir_path, \"test\"), resize = 4)\n video2frame(os.path.join(dir_path, \"training\"), resize = 4)"
},
{
"path": "dvs/warp/warping.py",
"content": "import numpy as np\nfrom .read_write import load_video, save_video\nimport torch\nimport cv2\nfrom .rasterizer import Rasterization\nimport time\nimport os\n\ndef warp_video(mesh_path, video_path, save_path, losses = None, frame_number = False, fps_fix = None):\n if type(mesh_path) == str:\n print(\"Error\")\n else:\n grid_data = mesh_path\n\n frame_array, fps, size = load_video(video_path, length = grid_data.shape[0])\n if fps_fix is not None:\n fps = fps_fix\n length = min(grid_data.shape[0], len(frame_array))\n seq_length = 100\n seq = length//seq_length\n writer = cv2.VideoWriter(save_path,cv2.VideoWriter_fourcc(*'mp4v'), fps, size)\n for i in range(seq+1):\n if seq_length*i==length:\n break\n print(\"Frame: \"+str(i*seq_length)+\"/\"+str(length))\n frame_array_save = warpping_rast(grid_data[seq_length*i:min(seq_length*(i+1),length)], frame_array[seq_length*i:min(seq_length*(i+1),length)], losses = losses)\n save_video(save_path,frame_array_save, fps, size, losses = losses, frame_number = frame_number, writer = writer)\n writer.release()\n\ndef warpping_rast(grid_data, frame_array, losses = None):\n output = []\n for i in range(0, min(len(frame_array), grid_data.shape[0])):\n frame = warpping_one_frame_rast(frame_array[i], grid_data[i])\n output.append(frame)\n return output\n\ndef warpping_one_frame_rast(image, grid):\n img = torch.Tensor(image).permute(2,0,1)/255\n grid = torch.Tensor(grid)\n output_image = Rasterization(img, grid)\n return np.clip(output_image.permute(1,2,0).numpy() * 255, 0, 255).astype(\"uint8\")\n"
}
]