├── .gitignore ├── .vscode ├── launch.json └── settings.json ├── LICENSE ├── README.md ├── __init__.py ├── apps ├── Cal_norm.py ├── generate-dataset-canny.py ├── kinect2grasp.py ├── pointclouds.py ├── read_file_sdf.py ├── read_grasps_from_file.py ├── show_pcd.py ├── voxelgrid.py └── ycb_cloud_generate.py ├── data ├── 1v_500_2class_ourway2sample.model ├── 1v_750_2class_ourway2sample.model ├── 1v_750_2class_ourway2sample0.model ├── grasp_pipeline.svg ├── image-20210311111731964.png ├── image-20210311162402868.png ├── 在线检测时的夹爪各项参数定义.png └── 在线检测时的夹爪数学模型各点以及夹爪坐标系定义.png ├── main_1v.py ├── main_1v_gpd.py ├── main_1v_mc.py ├── main_fullv.py ├── main_fullv_gpd.py ├── main_fullv_mc.py ├── main_test.py ├── model ├── __init__.py ├── dataset.py ├── gpd.py └── pointnet.py ├── requirements.txt └── requirements2.txt /.gitignore: -------------------------------------------------------------------------------- 1 | .idea/ 2 | assets/learned_models/ 3 | assets/log/ 4 | dataset/ycb_grasp/ 5 | dataset/ycb_meshes_google/ 6 | dataset/ycb_rgbd/ 7 | PointNetGPD/OLD/ 8 | PointNetGPD/chann_12/ 9 | PointNetGPDchann_3/ 10 | PointNetGPD/fullpc/ 11 | PointNetGPD/1view/ 12 | exception.txt 13 | *.npy 14 | # Created by https://www.gitignore.io/api/vim,python,pycharm,jupyternotebook 15 | 16 | ### JupyterNotebook ### 17 | .ipynb_checkpoints 18 | */.ipynb_checkpoints/* 19 | 20 | # Remove previous ipynb_checkpoints 21 | # git rm -r .ipynb_checkpoints/ 22 | # 23 | ### PyCharm ### 24 | # Covers JetBrains IDEs: IntelliJ, RubyMine, PhpStorm, AppCode, PyCharm, CLion, Android Studio and Webstorm 25 | # Reference: https://intellij-support.jetbrains.com/hc/en-us/articles/206544839 26 | 27 | # User-specific stuff: 28 | .idea/**/workspace.xml 29 | .idea/**/tasks.xml 30 | .idea/dictionaries 31 | 32 | # Sensitive or high-churn files: 33 | .idea/**/dataSources/ 34 | .idea/**/dataSources.ids 35 | .idea/**/dataSources.xml 36 | .idea/**/dataSources.local.xml 37 | .idea/**/sqlDataSources.xml 38 | .idea/**/dynamic.xml 39 | .idea/**/uiDesigner.xml 40 | 41 | # Gradle: 42 | .idea/**/gradle.xml 43 | .idea/**/libraries 44 | 45 | # CMake 46 | cmake-build-debug/ 47 | 48 | # Mongo Explorer plugin: 49 | .idea/**/mongoSettings.xml 50 | 51 | ## File-based project format: 52 | *.iws 53 | 54 | ## Plugin-specific files: 55 | 56 | # IntelliJ 57 | /out/ 58 | 59 | # mpeltonen/sbt-idea plugin 60 | .idea_modules/ 61 | 62 | # JIRA plugin 63 | atlassian-ide-plugin.xml 64 | 65 | # Cursive Clojure plugin 66 | .idea/replstate.xml 67 | 68 | # Ruby plugin and RubyMine 69 | /.rakeTasks 70 | 71 | # Crashlytics plugin (for Android Studio and IntelliJ) 72 | com_crashlytics_export_strings.xml 73 | crashlytics.properties 74 | crashlytics-build.properties 75 | fabric.properties 76 | 77 | ### PyCharm Patch ### 78 | # Comment Reason: https://github.com/joeblau/gitignore.io/issues/186#issuecomment-215987721 79 | 80 | # *.iml 81 | # modules.xml 82 | # .idea/misc.xml 83 | # *.ipr 84 | 85 | # Sonarlint plugin 86 | .idea/sonarlint 87 | 88 | ### Python ### 89 | # Byte-compiled / optimized / DLL files 90 | __pycache__/ 91 | *.py[cod] 92 | *$py.class 93 | 94 | # C extensions 95 | *.so 96 | 97 | # Distribution / packaging 98 | .Python 99 | build/ 100 | develop-eggs/ 101 | dist/ 102 | downloads/ 103 | eggs/ 104 | .eggs/ 105 | lib/ 106 | lib64/ 107 | parts/ 108 | sdist/ 109 | var/ 110 | wheels/ 111 | *.egg-info/ 112 | .installed.cfg 113 | *.egg 114 | 115 | # PyInstaller 116 | # Usually these files are written by a python script from a template 117 | # before PyInstaller builds the exe, so as to inject date/other infos into it. 118 | *.manifest 119 | *.spec 120 | 121 | # Installer logs 122 | pip-log.txt 123 | pip-delete-this-directory.txt 124 | 125 | # Unit test / coverage reports 126 | htmlcov/ 127 | .tox/ 128 | .coverage 129 | .coverage.* 130 | .cache 131 | .pytest_cache/ 132 | nosetests.xml 133 | coverage.xml 134 | *.cover 135 | .hypothesis/ 136 | 137 | # Translations 138 | *.mo 139 | *.pot 140 | 141 | # Flask stuff: 142 | instance/ 143 | .webassets-cache 144 | 145 | # Scrapy stuff: 146 | .scrapy 147 | 148 | # Sphinx documentation 149 | docs/_build/ 150 | 151 | # PyBuilder 152 | target/ 153 | 154 | # Jupyter Notebook 155 | 156 | # pyenv 157 | .python-version 158 | 159 | # celery beat schedule file 160 | celerybeat-schedule.* 161 | 162 | # SageMath parsed files 163 | *.sage.py 164 | 165 | # Environments 166 | .env 167 | .venv 168 | env/ 169 | venv/ 170 | ENV/ 171 | env.bak/ 172 | venv.bak/ 173 | 174 | # Spyder project settings 175 | .spyderproject 176 | .spyproject 177 | 178 | # Rope project settings 179 | .ropeproject 180 | 181 | # mkdocs documentation 182 | /site 183 | 184 | # mypy 185 | .mypy_cache/ 186 | 187 | ### Vim ### 188 | # swap 189 | .sw[a-p] 190 | .*.sw[a-p] 191 | # session 192 | Session.vim 193 | # temporary 194 | .netrwhist 195 | *~ 196 | # auto-generated tag files 197 | tags 198 | 199 | 200 | # End of https://www.gitignore.io/api/vim,python,pycharm,jupyternotebook 201 | 202 | # ignore pickle file 203 | *.pickle 204 | 205 | # for latex files 206 | *.aux 207 | *.log 208 | *.nav 209 | *.out 210 | *.snm 211 | *.toc 212 | *.synctex.gz 213 | -------------------------------------------------------------------------------- /.vscode/launch.json: -------------------------------------------------------------------------------- 1 | { 2 | // 使用 IntelliSense 了解相关属性。 3 | // 悬停以查看现有属性的描述。 4 | // 欲了解更多信息，请访问: https://go.microsoft.com/fwlink/?linkid=830387 5 | "version": "0.2.0", 6 | "configurations": [ 7 | { 8 | "name": "离线训练", 9 | "type": "python", 10 | "request": "launch", 11 | "program": "${workspaceFolder}/main_1v.py", 12 | "console": "integratedTerminal", 13 | "args": [ 14 | "--epoch", 15 | "200", 16 | "--mode", 17 | "train", 18 | "--batch-size", 19 | "32" 20 | ] 21 | }, 22 | { 23 | "name": "Python: 当前文件", 24 | "type": "python", 25 | "request": "launch", 26 | "program": "${file}", 27 | "console": "integratedTerminal" 28 | }, 29 | { 30 | "name": "在线GPD", 31 | "type": "python", 32 | "request": "launch", 33 | "program": "${workspaceFolder}/apps/kinect2grasp.py", 34 | "console": "integratedTerminal", 35 | "args": [ 36 | "--cuda", 37 | "--gpu", 38 | "0", 39 | "--load-model", 40 | "./data/1v_500_2class_ourway2sample.model", 41 | "--using_mp", 42 | "--model_type", 43 | "750" 44 | ] 45 | } 46 | ] 47 | } -------------------------------------------------------------------------------- /.vscode/settings.json: -------------------------------------------------------------------------------- 1 | { 2 | "python.pythonPath": "/home/wgk/anaconda3/bin/python" 3 | } -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | MIT License 2 | 3 | Copyright (c) 2018 Hongzhuo Liang 4 | 5 | Permission is hereby granted, free of charge, to any person obtaining a copy 6 | of this software and associated documentation files (the "Software"), to deal 7 | in the Software without restriction, including without limitation the rights 8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 9 | copies of the Software, and to permit persons to whom the Software is 10 | furnished to do so, subject to the following conditions: 11 | 12 | The above copyright notice and this permission notice shall be included in all 13 | copies or substantial portions of the Software. 14 | 15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 21 | SOFTWARE. 22 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | --- 2 | typora-root-url: ./ 3 | --- 4 | 5 | # PointNetGPD: Detecting Grasp Configurations from Point Set 6 | 7 | 感谢原作者开源：https://github.com/lianghongzhuo/PointNetGPD 8 | 9 | PointNetGPD (ICRA 2019, [arXiv](https://arxiv.org/abs/1809.06267)) is an end-to-end grasp evaluation model to address the challenging problem of localizing robot grasp configurations directly from the point cloud. 10 | 11 | **一个题外话：** 关于抓取采样的问题，很多人希望在混杂场景中也能够进行6-Dof候选抓取的采样，以便于用于一些regression-base算法；目前这种采样方式考虑到标注成本，一般是在虚拟场景中采样，然后用机理法标注；我个人也开源了一个[“半成品的”](https://github.com/Hymwgk/simulate_dataset_generation)仓库，有时间有精力的同学可以看看，也许有些代码片段可以帮助到你; 由于是自己手写的，很多是从小白开始写的，希望也许会有一些借鉴的价值吧。 12 | 13 | 14 | 15 | 简单介绍一下PointNet的代码流程： 16 | - 离线阶段 17 | 18 | 1. **候选抓取采样与打分：** 针对[YCB objects Dataset](http://ycb-benchmarks.s3-website-us-east-1.amazonaws.com/) 中的CAD模型，利用传统采样方法（默认Antipods采样）对进行候选抓取姿态（夹爪6D姿态）采样，并通过Force Closure结合 GWS 方法对所检测出的姿态进行打分，生成候选抓取数据集存起来备用； 19 | 2. **点云原始数据生成：** [YCB objects Dataset](http://ycb-benchmarks.s3-website-us-east-1.amazonaws.com/) 提供了一些物体的CAD模型，以及某些视角下这些物体的激光深度图；但是我们在训练网络时候，使用的是点云数据，因此，这里提前要把某些视角下的深度图转换为点云数据，以备后续使用； 20 | 3. **PointNet模型训练：** 利用生成的某候选夹爪姿态（步骤1.i），结合CAD模型的多视角点云（步骤1.ii），提取出来该候选抓取姿态下夹爪闭合区域内的点云；将夹爪闭合区域内的点云以及该抓取姿态的分数送入PointNet中进行训练； 21 | 22 | - 在线阶段 23 | 1. **抓取采样：** [从点云中直接采样候选抓取姿态](https://www.researchgate.net/publication/318107041_Grasp_Pose_Detection_in_Point_Clouds) ，并剔除掉与桌面碰撞、与场景点云碰撞的非法抓取； 24 | 2. **打分：** 提取剩余的抓取姿态夹爪内部的点，进一步剔除掉不合理的数据之后，将点集送入训练好的PointNet网络中打分； 25 | 3. **排序：** 将候选的抓取按照分数从高到低排序，输出分数最高的抓取。 26 | 27 | 28 | 29 | ## Video 30 | - 作者的实验视频 31 | [![Video for PointNetGPD](https://img.youtube.com/vi/RBFFCLiWhRw/0.jpg )](https://www.youtube.com/watch?v=RBFFCLiWhRw) 32 | - 本人的复现实验视频(未加速) 33 | 34 | [![PointNetGPD on Franka Panda](https://img.youtube.com/vi/OfvJ-HpKjI4/0.jpg)](https://www.youtube.com/watch?v=OfvJ-HpKjI4) 35 | 在实验中发现，gpd效果还是很不错的；但是夹爪经常撞到目标物体上，这是受到了手眼标定的精度以及panda夹爪构型的影响（panda夹爪的深度比较浅，最大张开距离也比较小） 36 | ## 关于数据集 37 | 38 | 代码中数据集的生成部分比较混乱，这里解释一下具体数据集的生成逻辑： 39 | 40 | YCB数据集主要分为两个部分，一部分是数据集物体的CAD模型，一部分是各个视角下物体的rgbd深度图，因此在前期的准备过程中，主要分为两个独立的部分： 41 | 42 | - 对YCB数据集的CAD模型进行候选抓取姿态采样，获取候选grasp数据集； 43 | - 对YCB数据集中各个视角的深度图像进行计算，生成该视角下的点云数据集； 44 | 45 | 之后两部分的计算结果将会被放在同一个文件夹中，Dataloader将会对该文件夹进行处理，并计算出样本数据送入PointNet中进行训练。 46 | 47 | image-20210311111731964 48 | 49 | 以下是稍微具体一些的解释图： 50 | 51 | image-20210311162402868 52 | 53 | 为了防止硬盘不够大，本代码将最终Dataloader用到的数据集，单独放在了以下目录，可以挂载到一个单独的硬盘里，详细见后文的设置 54 | 55 | ```bash 56 | mkdir -p $HOME/dataset/PointNetGPD 57 | ``` 58 | 59 | 60 | 61 | ## 本代码的修改 62 | 63 | 原代码中 64 | 65 | 离线阶段：使用Antipod算法对YCB 的CAD模型进行候选grasp pose采样，并结合夹爪数学模型和模型视角点云来生成夹爪内部点云，进而输入PointNet中训练； 66 | 67 | 在线阶段：使用基于点云的gpg算法，对场景点云进行候选grasp pose采样，并结合夹爪数学模型和场景点云来生成夹爪内部点云，进而输入PointNet中进行预测。 68 | 69 | 然而，实际代码中两个阶段使用的夹爪数学模型不一致，个人感觉统一模型会好一些，于是本代码中统一dataloader中相同的数学模型； 70 | 71 | ## Before Install 72 | 73 | 在使用前，clone的代码文件夹需要放在如下的code文件夹中: 74 | ```bash 75 | mkdir -p $HOME/code/ 76 | cd $HOME/code/ 77 | ``` 78 | 79 | 80 | ## Install all the requirements (Using a virtual environment is recommended) 81 | 82 | 1. Make sure in your Python environment do not have same package named ```meshpy``` or ```dexnet```. 83 | 84 | 2. 确保已经安装了ROS以及相机预处理相关的[程序包](https://github.com/Hymwgk/point_cloud_process) 85 | 86 | 3. Clone this repository: 87 | ```bash 88 | cd $HOME/code 89 | git clone https://github.com/hymwgk/PointNetGPD.git 90 | ``` 91 | 92 | 4. Install our requirements in `requirements.txt` (在python2以及python3环境中都需要安装) 93 | ```bash 94 | cd $HOME/code/PointNetGPD 95 | pip install -r requirements.txt #python3环境中 96 | pip install -r requirements2.txt #python2环境中 97 | ``` 98 | 99 | 5. Install our modified meshpy (Modify from [Berkeley Automation Lab: meshpy](https://github.com/BerkeleyAutomation/meshpy)) 100 | ```bash 101 | cd $HOME/code 102 | git clone https://github.com/Hymwgk/meshpy.git 103 | cd meshpy 104 | #分别在python2和python3环境下执行一遍 105 | python setup.py develop 106 | ``` 107 | 108 | 6. Install our modified dex-net (Modify from [Berkeley Automation Lab: dex-net](https://github.com/BerkeleyAutomation/dex-net)) 109 | ```bash 110 | cd $HOME/code 111 | git clone https://github.com/Hymwgk/dex-net.git 112 | cd dex-net 113 | # 分别在python2和python3环境下执行一遍 114 | python setup.py develop 115 | ``` 116 | 117 | 7. 设置夹爪数学模型，之后离线以及在线节点的候选夹爪姿态计算都会依据此模型来生成；你可以直接在以下文件中根据自己的实际夹爪尺寸来修改对应参数。 118 | 119 | ```bash 120 | vim $HOME/code/dex-net/data/grippers/robotiq_85/params.json 121 | ``` 122 | 以下是`params.json`中，原作者离线阶段用到的夹爪数学模型各个参数（本代码中已经将其废弃） 123 | 124 | ```bash 125 | "min_width": 夹爪的最小闭合角度 126 | "force_limit": 抓取力度限制 127 | "max_width": 夹爪最大张开距离 128 | "finger_radius": 用于软体手，指定软体手的弯曲角度（弧度制），一般用不到，补上去就行了 129 | "max_depth": 夹爪的最大深度，竖向的距离 130 | ``` 131 | 以下本代码中离线训练和在线检测统一使用的夹爪参数定义: 132 | ```bash 133 | "finger_width": 夹持器的两个夹爪的“厚度” 134 | "real_finger_width": 也是两个夹爪的厚度，和上面写一样就行（仅仅用于显示，影响不大，不用于姿态检测） 135 | "hand_height": 夹爪的另一侧厚度，一会儿看图 136 | "hand_height_two_finger_side": 没有用到，代码中没有找到调用，所以不用管 137 | "hand_outer_diameter": 夹爪最大的可以张开的距离，从最外侧量（包括两个爪子的厚度） 138 | "hand_depth": 夹爪的竖向深度 139 | "real_hand_depth": 和hand_depth保持一致，代码中两者是相同的 140 | "init_bite": 这个是用于在线抓取检测时，定义的一个后撤距离，主要是避免由于点云误差之类的，导致夹爪和物体碰撞，以米为单位，一般设置1cm就行了 141 | ``` 142 | 143 | `params.json`参数的具体定义示意图，修改后的本代码，离线的夹爪参数仅作为候选抓取姿态的采样，而不涉及到夹爪内部点云的提取。 144 | 145 | 在线检测时的夹爪各项参数定义 146 | 在线检测时的夹爪数学模型各点以及夹爪坐标系定义 147 | 148 | 149 | 150 | ## 对YCB数据集的CAD模型进行候选抓取姿态采样 151 | 152 | 1. 下载 YCB object set from [YCB Dataset](http://ycb-benchmarks.s3-website-us-east-1.amazonaws.com/)，该数据集提供了物体的CAD模型和一定角度下的深度图； 153 | 2. 原代码中，将YCB的数据集放在了如下位置: 154 | ```bash 155 | cd $HOME/dataset/PointNetGPD/ycb_meshes_google/objects 156 | ``` 157 | 158 | 如果你的Home盘分区不够大，希望换一个位置，可以: 159 | >```bash 160 | >cd $HOME/code/PointNetGPD/apps 161 | >vim generate-dataset-canny.py 162 | >``` 163 | >修改YCB dataset 路径 164 | >```python 165 | >219 #存放CAD模型的文件夹 166 | >220 file_dir = home_dir + "/dataset/ycb_meshes_google/objects" #获取模型的路径 167 | >``` 168 | >也可以修改计算结果文件的存放位置 169 | >```python 170 | >61 #将gpg得到的候选抓取文件存放起来 171 | >62 good_grasp_file_name = "./generated_grasps/{}_{}_{}".format(filename_prefix, str(object_name), str(len(good_grasp))) 172 | >``` 173 | 174 | 每个物体的文件夹结构都应该如下所示: 175 | ```bash 176 | ├002_master_chef_can 177 | |└── google_512k 178 | | ├── kinbody.xml (no use) 179 | | ├── nontextured.obj 180 | | ├── nontextured.ply 181 | | ├── nontextured.sdf (generated by SDFGen) 182 | | ├── nontextured.stl 183 | | ├── textured.dae (no use) 184 | | ├── textured.mtl (no use) 185 | | ├── textured.obj (no use) 186 | | ├── textured.sdf (no use) 187 | | └── texture_map.png (no use) 188 | ├003_cracker_box 189 | └004_sugar_box 190 | ... 191 | ``` 192 | 3. Install SDFGen from [GitHub](https://github.com/jeffmahler/SDFGen.git): 193 | ```bash 194 | cd code 195 | git clone https://github.com/jeffmahler/SDFGen.git 196 | cd SDFGen 197 | sudo sh install.sh 198 | ``` 199 | 200 | 201 | 4. 安装python pcl library [python-pcl](https://github.com/strawlab/python-pcl)，python pcl在离线训练(python3)和在线pgd(python2)时均有使用，以下要求Ubuntu18.04，PCL1.8.1（源安装在系统路径下）: 202 | ```bash 203 | git clone https://github.com/lianghongzhuo/python-pcl.git 204 | pip install --upgrade pip 205 | pip install cython==0.25.2 #python2 206 | pip install cython #python3 207 | pip install numpy 208 | cd python-pcl 209 | python setup.py build_ext -i #python2和3环境中都要执行 210 | python setup.py develop #python2和3环境中都要执行 211 | ``` 212 | 5. 为默认路径`$HOME/dataset/PointNetGPD/ycb_meshes_google/objects/`下的.ply文件生成.sdf 文件（放在同一文件夹下）: 213 | ```bash 214 | cd $HOME/code/PointNetGPD/apps 215 | python read_file_sdf.py #anaconda3环境下python3 216 | ``` 217 | 218 | 219 | 6. 为默认路径`$HOME/dataset/PointNetGPD/ycb_meshes_google/objects/`下的CAD模型使用Antipod进行候选抓取姿态采样,以及利用ForceClosure&GWS对生成抓取姿态进行打分，这部分的执行时间极长，主要花费时间在抓取采样之上： 220 | ```bash 221 | cd $HOME/code/PointNetGPD/apps 222 | python generate-dataset-canny.py [prefix] #anaconda3环境下python3 223 | ``` 224 | 225 | 计算结束后将会把结果以`.npy`文件形式保存在默认的`$HOME/dataset/PointNetGPD/ycb_grasp/`路径下；这里的`[prefix]`可以根据自己的夹爪类型，添加一个标签，也可以选择不加，那么就会自动被替换成为`default` 226 | 227 | 7. 作者还给出了一个根据roboticq85夹爪模型采样好的候选grasp pose [结果文件](https://tams.informatik.uni-hamburg.de/research/datasets/PointNetGPD_grasps_dataset.zip) 228 | 229 | 230 | ## 对YCB数据集中各个视角的深度图像生成点云 231 | 232 | 1. 将下载的YCB数据集文件夹`ycb_rgbd`拷贝至如下路径 233 | 234 | ```bash 235 | cp .../ycb_rgbd $HOME/dataset/PointNetGPD/ 236 | ``` 237 | 238 | 239 | 2. 将默认路径`$HOME/dataset/PointNetGPD/ycb_rgbd/*/`下的深度图转换为点云数据，并放在`$HOME/dataset/PointNetGPD/ycb_rgbd/*/clouds`文件夹中。 240 | 241 | ```bash 242 | cd $HOME/code/PointNetGPD/apps/ 243 | python ycb_cloud_generate.py #anaconda3 python3 244 | ``` 245 | 246 | 247 | ## 准备Dataloader需要的数据文件夹 248 | 249 | 此该文件夹将会提供给PointNet的Dataloader，该Dataloader将会在训练时结合候选grasp pose&点云 提取“夹爪内部的点云”（详细解释见作者论文） 250 | 251 | 1. 进入Dataloader需要的文件夹: 252 | 253 | ```bash 254 | cd $HOME/dataset/PointNetGPD/ 255 | ``` 256 | 确保该文件夹下有如下文件 257 | ``` 258 | ├── google2cloud.csv (Transform from google_ycb model to ycb_rgbd model) 259 | ├── google2cloud.pkl (Transform from google_ycb model to ycb_rgbd model) 260 | ├── ycb_grasp (里面就是离线Antipod采样到的候选grasp pose) 261 | ├── ycb_meshes_google (YCB dataset) 262 | └── ycb_rgbd (上面已经生成了各模型各视角点云) 263 | ``` 264 | 265 | 其中，`ycb_grasp`文件夹需要手动创建为如下结构，对之前生成的候选抓取样本分成两部分一部分训练一部分验证（作者好像没说咋分），每个文件夹中都是之前`generate-dataset-canny.py`采样到的grasp pose（`.npy`） 266 | 267 | ```bash 268 | ├── ycb_grasp 269 | │   ├── test # 270 | │   └── train #(训练) 271 | ``` 272 | 273 | 274 | ## 训练模型 275 | 276 | 1. Run the experiments: 277 | 278 | ```bash 279 | cd PointNetGPD 280 | ``` 281 | 282 | 283 | Launch a tensorboard for monitoring 284 | ```bash 285 | tensorboard --log-dir ./assets/log --port 8080 286 | ``` 287 | 288 | and run an experiment for 200 epoch 289 | ```bash 290 | python main_1v.py --epoch 200 --mode train --batch-size x (x>1) #anaconda3 python3 291 | ``` 292 | 293 | File name and corresponding experiment: 294 | ```bash 295 | main_1v.py --- 1-viewed point cloud, 2 class 296 | main_1v_mc.py --- 1-viewed point cloud, 3 class 297 | main_1v_gpd.py --- 1-viewed point cloud, GPD 298 | main_fullv.py --- Full point cloud, 2 class 299 | main_fullv_mc.py --- Full point cloud, 3 class 300 | main_fullv_gpd.py --- Full point cloud, GPD 301 | ``` 302 | For GPD experiments, you may change the input channel number by modifying `input_chann` in the experiment scripts(only 3 and 12 channels are available) 303 | 304 | ## 使用模型，执行抓取 305 | 306 | 需要注意的是： 307 | 308 | - 为了能够脱离机械臂实物，仅仅进行GPD的实验，同时还能够相对容易地剔除掉桌面点云；代码中选择将场景点云变换到桌面标签（ar_marker_6）坐标系中，该部分的变换处理[参看此处](https://github.com/Hymwgk/point_cloud_process) 309 | - 如果使用机械臂实物（以panda为例）所有的指令运行的窗口都需要运行`source panda_client.sh`指令确保本机ROS_MASTER指向远程工控机，如果仅仅进行gpd不实际进行抓取则不需要这样做。[参看此处](https://github.com/Hymwgk/panda_moveit_config) 310 | 311 | 1. （不实际抓取不需要）手眼标定并发布“手眼”变换关系，关于Panda手眼标定和发布步骤[参看此处](https://github.com/Hymwgk/panda_hand_eye_calibrate) 312 | 313 | ```bash 314 | roslaunch panda_hand_eye_calibrate publish_panda_eob.launch 315 | ``` 316 | 317 | 2. 安装预处理包 318 | 319 | ```bash 320 | cd ~/catkin_ws/src 321 | git clone https://github.com/Hymwgk/point_cloud_process.git 322 | catkin build 323 | ``` 324 | 325 | 3. 启动点云采集与预处理 326 | ```bash 327 | roslaunch kinect2_bridge kinect2_bridge.launch publish_tf:=true #启动相机 328 | roslaunch point_cloud_process get_table_top_points.launch #启动标签追踪以及点云预处理 329 | ``` 330 | 331 | 4. 运行感知节点 332 | 333 | 这部分就是实际使用PointNetGPD的部分，读取预处理后桌面上的目标区域点云，基于点云进行gpg，之后将夹爪内部的点云送入pointNet中打分，并以ROS消息的形式输出good grasp 334 | 335 | ```bash 336 | cd $HOME/code/PointNetGPD/apps 337 | python kinect2grasp.py #anaconda2 python2 需要在conda2环境中也安装pytorch 338 | ``` 339 | 以下是用到的参数 340 | ```bash 341 | arguments: 342 | -h, --help show this help message and exit 343 | --cuda 使用CUDA进行计算 344 | --gpu GPU 指定使用的GPU编号 345 | --load-model LOAD_MODEL 设置使用了哪个训练好的网络 (这个参数其实没有效果，被后面的model_type MODEL_TYPE覆盖了) 346 | --show_final_grasp 设置是否显示最终抓取（修改了多线程不显示的问题） 347 | --tray_grasp not finished grasp type（还没搞好） 348 | --using_mp 是否使用多线程去进行抓取采样 349 | --model_type MODEL_TYPE 从三种模型中选择使用哪个模型 350 | ``` 351 | 举个栗子： 352 | ```bash 353 | python kinect2grasp.py --cuda --gpu 0 --load-model ../data/1v_500_2class_ourway2sample.model --using_mp --model_type 750 354 | ``` 355 | 356 | 357 | 358 | ------ 359 | 360 | ​ **以下是使用机械臂进行真正的抓取执行需要的步骤，可以不做** 361 | 362 | 5. 安装ROS抓取消息包 363 | 364 | ```bash 365 | cd ~/catkin_ws/src 366 | git clone https://github.com/TAMS-Group/gpd_grasp_msgs.git 367 | cd .. 368 | catkin build 369 | ``` 370 | 371 | 6. 获取机械臂当前状态，并准备接收生成的抓取（以franka panda为例）需提前安装[panda_go_grasp包](https://github.com/Hymwgk/panda_go_grasp)，向ROS参数服务器发布一个参数，指明机械臂的当前是在移动状态还是已经返回home状态，机械臂在移动时，将暂时禁止gpd。 372 | 373 | ```bash 374 | roslaunch panda_go_grasp state_checker.launch #anaconda2 python2 375 | ``` 376 | 377 | 7. 执行抓取 378 | 379 | ```bash 380 | roslaunch panda_go_grasp go_grasp.launch #anaconda2 python2 381 | ``` 382 | 383 | 384 | 385 | 386 | 387 | ## 辅助工具：查看之前采样的候选grasp pose（可跳过） 388 | - Visualization grasps 389 | ```bash 390 | cd $HOME/code/PointNetGPD/apps 391 | python read_grasps_from_file.py #anaconda3 python3 392 | ``` 393 | 394 | Note: 395 | 396 | - This file will visualize the grasps in `$HOME/code/PointNetGPD/PointNetGPD/data/ycb_grasp/` folder 397 | 398 | - Visualization object normals 399 | ```bash 400 | cd $HOME/code/PointNetGPD/apps 401 | python Cal_norm.py #anaconda3 python3 402 | ``` 403 | 404 | This code will check the norm calculated by meshpy and pcl library. 405 | 406 | 407 | ## Citation 408 | If you found PointNetGPD useful in your research, please consider citing: 409 | 410 | ```plain 411 | @inproceedings{liang2019pointnetgpd, 412 | title={{PointNetGPD}: Detecting Grasp Configurations from Point Sets}, 413 | author={Liang, Hongzhuo and Ma, Xiaojian and Li, Shuang and G{\"o}rner, Michael and Tang, Song and Fang, Bin and Sun, Fuchun and Zhang, Jianwei}, 414 | booktitle={IEEE International Conference on Robotics and Automation (ICRA)}, 415 | year={2019} 416 | } 417 | ``` 418 | 419 | ## Acknowledgement 420 | - [gpg](https://github.com/atenpas/gpg) 421 | - [gpd](https://github.com/atenpas/gpd) 422 | - [dex-net](https://github.com/BerkeleyAutomation/dex-net) 423 | - [meshpy](https://github.com/BerkeleyAutomation/meshpy) 424 | - [SDFGen](https://github.com/christopherbatty/SDFGen) 425 | - [pyntcloud](https://github.com/daavoo/pyntcloud) 426 | - [metu-ros-pkg](https://github.com/kadiru/metu-ros-pkg) 427 | - [mayavi](https://github.com/enthought/mayavi) 428 | 429 | -------------------------------------------------------------------------------- /__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Hymwgk/PointNetGPD/e33b37f8da0b6326976bb339faa976f724980d82/__init__.py -------------------------------------------------------------------------------- /apps/Cal_norm.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python 2 | # -*- coding: utf-8 -*- 3 | # Author : Hongzhuo Liang 4 | # E-mail : liang@informatik.uni-hamburg.de 5 | # Description: 6 | # Date : 09/06/2018 7:47 PM 7 | # File Name : Cal_norm.py 8 | import os 9 | from meshpy.obj_file import ObjFile 10 | from meshpy.sdf_file import SdfFile 11 | from dexnet.grasping import GraspableObject3D 12 | import matplotlib.pyplot as plt 13 | import numpy as np 14 | import pcl 15 | import multiprocessing 16 | import time 17 | from mayavi import mlab 18 | 19 | 20 | def show_obj(surface_points_, color='b'): 21 | if color == 'b': 22 | color_f = (0, 0, 1) 23 | elif color == 'r': 24 | color_f = (1, 0, 0) 25 | elif color == 'g': 26 | color_f = (0, 1, 0) 27 | else: 28 | color_f = (1, 1, 1) 29 | points = surface_points_ 30 | mlab.points3d(points[:, 0], points[:, 1], points[:, 2], color=color_f, scale_factor=.0007) 31 | 32 | 33 | def get_file_name(file_dir_): 34 | file_list = [] 35 | for root, dirs, files in os.walk(file_dir_): 36 | # print(root) # current path 37 | if root.count('/') == file_dir_.count('/')+1: 38 | file_list.append(root) 39 | # print(dirs) # all the directories in current path 40 | # print(files) # all the files in current path 41 | file_list.sort() 42 | return file_list 43 | 44 | 45 | def show_grasp_norm(grasp_bottom_center, grasp_axis): 46 | un1 = grasp_bottom_center - 0.25 * grasp_axis * 0.25 47 | un2 = grasp_bottom_center # - 0.25 * grasp_axis * 0.05 48 | mlab.plot3d([un1[0], un2[0]], [un1[1], un2[1]], [un1[2], un2[2]], color=(0, 1, 0), tube_radius=0.0005) 49 | 50 | 51 | def show_pcl_norm(grasp_bottom_center, normal_, color='r', clear=False): 52 | if clear: 53 | plt.figure() 54 | plt.clf() 55 | plt.gcf() 56 | plt.ion() 57 | 58 | ax = plt.gca(projection='3d') 59 | un1 = grasp_bottom_center + 0.5 * normal_ * 0.05 60 | ax.scatter(un1[0], un1[1], un1[2], marker='x', c=color) 61 | un2 = grasp_bottom_center 62 | ax.scatter(un2[0], un2[1], un2[2], marker='^', c='g') 63 | ax.plot([un1[0], un2[0]], [un1[1], un2[1]], [un1[2], un2[2]], 'b-', linewidth=1) # bi normal 64 | 65 | 66 | def do_job(job_i): 67 | ii = np.random.choice(all_p.shape[0]) 68 | show_grasp_norm(all_p[ii], surface_normal[ii]) 69 | print("done job", job_i, ii) 70 | 71 | 72 | if __name__ == '__main__': 73 | home_dir = os.environ['HOME'] 74 | file_dir = home_dir + "/dataset/ycb_meshes_google/objects" 75 | mlab.figure(bgcolor=(1.0, 1.0, 1.0), fgcolor=(0.7, 0.7, 0.7)) 76 | file_list_all = get_file_name(file_dir) 77 | object_numbers = file_list_all.__len__() 78 | i = 0 # index of objects to define which object to show 79 | if os.path.exists(str(file_list_all[i]) + "/google_512k/nontextured.obj"): 80 | of = ObjFile(str(file_list_all[i]) + "/google_512k/nontextured.obj") 81 | sf = SdfFile(str(file_list_all[i]) + "/google_512k/nontextured.sdf") 82 | 83 | else: 84 | print("can't find any obj or sdf file!") 85 | raise NameError("can't find any obj or sdf file!") 86 | mesh = of.read() 87 | sdf = sf.read() 88 | graspable = GraspableObject3D(sdf, mesh) 89 | print("Log: opened object") 90 | begin_time = time.time() 91 | surface_points, _ = graspable.sdf.surface_points(grid_basis=False) 92 | all_p = surface_points 93 | method = "voxel" 94 | if method == "random": 95 | surface_points = surface_points[np.random.choice(surface_points.shape[0], 1000, replace=False)] 96 | surface_normal = [] 97 | elif method == "voxel": 98 | surface_points = surface_points.astype(np.float32) 99 | p = pcl.PointCloud(surface_points) 100 | voxel = p.make_voxel_grid_filter() 101 | voxel.set_leaf_size(*([graspable.sdf.resolution * 5] * 3)) 102 | surface_points = voxel.filter().to_array() 103 | 104 | # cal normal with pcl 105 | use_voxel = False # use voxelized point to get norm is not accurate 106 | if use_voxel: 107 | norm = voxel.filter().make_NormalEstimation() 108 | else: 109 | norm = p.make_NormalEstimation() 110 | 111 | norm.set_KSearch(10) 112 | normals = norm.compute() 113 | surface_normal = normals.to_array() 114 | surface_normal = surface_normal[:, 0:3] 115 | surface_gg = norm.compute().to_array()[:, 3] 116 | # from IPython import embed;embed() 117 | use_pcl = True 118 | if use_pcl: 119 | # for ii in range(surface_normal[0]): 120 | show_grasp_norm(all_p[0], surface_normal[0]) 121 | # FIXME: multi processing is not working, Why? 122 | # mayavi do not support 123 | # cores = multiprocessing.cpu_count() 124 | # pool = multiprocessing.Pool(processes=cores) 125 | # pool.map(multiprocessing_jobs.do_job, range(2000)) 126 | sample_points = 500 127 | for _ in range(sample_points): 128 | do_job(_) 129 | mlab.pipeline.surface(mlab.pipeline.open(str(file_list_all[i]) + "/google_512k/nontextured.ply") 130 | , opacity=1) 131 | mlab.show() 132 | print(time.time() - begin_time) 133 | # show_obj(all_p) 134 | else: 135 | raise ValueError("No such method", method) 136 | 137 | use_meshpy = False 138 | if use_meshpy: 139 | normal = [] 140 | # show norm 141 | surface_points = surface_points[:100] 142 | for ind in range(len(surface_points)): 143 | # use meshpy cal norm: 144 | p_grid = graspable.sdf.transform_pt_obj_to_grid(surface_points[ind]) 145 | normal_tmp = graspable.sdf.surface_normal(p_grid) 146 | # use py pcl cal norm, Wrong. 147 | # normal_tmp = surface_normal[ind] 148 | if normal_tmp is not None: 149 | normal.append(normal_tmp) 150 | show_grasp_norm(surface_points[ind], normal_tmp) 151 | else: 152 | print(len(normal)) 153 | mlab.pipeline.surface(mlab.pipeline.open(str(file_list_all[i]) + "/google_512k/nontextured.ply")) 154 | mlab.show() 155 | -------------------------------------------------------------------------------- /apps/generate-dataset-canny.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python 2 | # -*- coding: utf-8 -*- 3 | # python3 4 | # Author : Hongzhuo Liang 5 | # E-mail : liang@informatik.uni-hamburg.de 6 | # Description: 7 | # Date : 20/05/2018 2:45 PM 8 | # File Name : generate-dataset-canny.py 9 | 10 | import numpy as np 11 | import sys 12 | import pickle 13 | from dexnet.grasping.quality import PointGraspMetrics3D 14 | from dexnet.grasping import GaussianGraspSampler, AntipodalGraspSampler, UniformGraspSampler, GpgGraspSampler 15 | from dexnet.grasping import RobotGripper, GraspableObject3D, GraspQualityConfigFactory, PointGraspSampler 16 | import dexnet 17 | from autolab_core import YamlConfig 18 | from meshpy.obj_file import ObjFile 19 | from meshpy.sdf_file import SdfFile 20 | import os 21 | 22 | import multiprocessing 23 | import matplotlib.pyplot as plt 24 | plt.switch_backend('agg') # for the convenient of run on remote computer 25 | 26 | #sys.path() 27 | 28 | #输入文件夹地址，返回一个列表，其中保存的是文件夹中的文件名称 29 | def get_file_name(file_dir_): 30 | file_list = [] 31 | for root, dirs, files in os.walk(file_dir_): 32 | #将下一层子文件夹的地址保存到 file_list 中 33 | if root.count('/') == file_dir_.count('/') + 1: 34 | file_list.append(root) 35 | #排序 36 | file_list.sort() 37 | return file_list 38 | """ 39 | \brief: 创建多线程，多线程调用worker函数，处理指定的模型， 40 | """ 41 | def do_job(i): #处理函数 处理第i个模型 42 | #根据id号，截取对应的目标模型名称 43 | object_name = file_list_all[i].split('/')[-1] 44 | 45 | """ 46 | good_grasp = multiprocessing.Manager().list() #列表 47 | 48 | # grasp_amount per friction: 20*40 共开50个子线程做这件事 49 | p_set = [multiprocessing.Process(target=worker, args=(i, 100, 20, good_grasp)) for _ in 50 | range(50)] 51 | 52 | #开始多线程，并等待结束 53 | [p.start() for p in p_set] 54 | [p.join() for p in p_set] 55 | good_grasp = list(good_grasp) 56 | """ 57 | #创建空列表 58 | good_grasp=[] 59 | #执行worker，将采样的抓取结果放到good_grasp中 60 | worker(i, 100, 20, good_grasp) 61 | 62 | #对CAD模型进行Antipod采样并筛选得到的候选抓取 63 | good_grasp_file_name = os.environ['HOME']+"/dataset/PointNetGPD/ycb_grasp/{}_{}_{}".format(filename_prefix, str(object_name), str(len(good_grasp))) 64 | 65 | #创建一个pickle文件，将good_grasp保存起来 66 | with open(good_grasp_file_name + '.pickle', 'wb') as f: 67 | pickle.dump(good_grasp, f) 68 | 69 | tmp = [] 70 | for grasp in good_grasp: 71 | grasp_config = grasp[0].configuration 72 | score_friction = grasp[1] 73 | score_canny = grasp[2] 74 | tmp.append(np.concatenate([grasp_config, [score_friction, score_canny]])) 75 | np.save(good_grasp_file_name + '.npy', np.array(tmp)) 76 | print("finished job ", object_name) 77 | 78 | def do_jobs(i): 79 | print("tesk id", i) 80 | 81 | def worker(i, sample_nums, grasp_amount, good_grasp): #主要是抓取采样器以及打分 100 20 82 | """ 83 | brief: 对制定的模型，利用随机采样算法，进行抓取姿态的检测和打分 84 | param [in] i 处理第i个mesh模型 85 | param [in] sample_nums 每个对象模型返回的目标抓取数量 86 | param [in] grasp_amount 87 | """ 88 | 89 | #截取目标对象名称 90 | object_name = file_list_all[i][len(home_dir) + 35:] 91 | print('a worker of task {} start'.format(object_name)) 92 | 93 | #读取初始配置文件，读取并在内存中复制了一份 94 | yaml_config = YamlConfig(home_dir + "/code/dex-net/test/config.yaml") 95 | #设置夹名称 96 | gripper_name = 'panda' 97 | #根据设置的夹爪名称加载夹爪配置 98 | gripper = RobotGripper.load(gripper_name, home_dir + "/code/dex-net/data/grippers") 99 | #设置抓取采样的方法 100 | grasp_sample_method = "antipodal" 101 | if grasp_sample_method == "uniform": 102 | ags = UniformGraspSampler(gripper, yaml_config) 103 | elif grasp_sample_method == "gaussian": 104 | ags = GaussianGraspSampler(gripper, yaml_config) 105 | elif grasp_sample_method == "antipodal": 106 | #使用对映点抓取，输入夹爪与配置文件 107 | ags = AntipodalGraspSampler(gripper, yaml_config) 108 | elif grasp_sample_method == "gpg": 109 | ags = GpgGraspSampler(gripper, yaml_config) 110 | elif grasp_sample_method == "point": 111 | ags = PointGraspSampler(gripper, yaml_config) 112 | else: 113 | raise NameError("Can't support this sampler") 114 | print("Log: do job", i) 115 | #设置obj模型文件与sdf文件路径 116 | if os.path.exists(str(file_list_all[i]) + "/google_512k/nontextured.obj") and os.path.exists(str(file_list_all[i]) + "/google_512k/nontextured.sdf"): 117 | of = ObjFile(str(file_list_all[i]) + "/google_512k/nontextured.obj") 118 | sf = SdfFile(str(file_list_all[i]) + "/google_512k/nontextured.sdf") 119 | else: 120 | print("can't find any obj or sdf file!") 121 | raise NameError("can't find any obj or sdf file!") 122 | 123 | #根据路径读取模型与sdf文件 124 | mesh = of.read() 125 | sdf = sf.read() 126 | #构建抓取模型类 127 | obj = GraspableObject3D(sdf, mesh) 128 | print("Log: opened object", i + 1, object_name) 129 | 130 | ######################################### 131 | #设置 132 | force_closure_quality_config = {} #设置力闭合 字典 133 | canny_quality_config = {} 134 | #生成一个起点是2.0终点是0.75 步长为-0.4 （递减）的等距数列fc_list_sub1 (2.0, 0.75, -0.4) 135 | fc_list_sub1 = np.arange(2.0, 0.75, -0.3) 136 | #生成一个起点是0.5终点是0.36 步长为-0.05的等距数列fc_list_sub2 (0.5, 0.36, -0.05) 137 | fc_list_sub2 = np.arange(0.5, 0.36, -0.1) 138 | 139 | #将上面两个向量接起来，变成一个长条向量，使用不同的步长，目的是为了在更小摩擦力的时候，有更多的分辨率 140 | fc_list = np.concatenate([fc_list_sub1, fc_list_sub2]) 141 | print("判断摩擦系数") 142 | print(fc_list) 143 | for value_fc in fc_list: 144 | #对value_fc保留2位小数，四舍五入 145 | value_fc = round(value_fc, 2) 146 | #更改内存中配置中的摩擦系数，而没有修改硬盘中的yaml文件 147 | yaml_config['metrics']['force_closure']['friction_coef'] = value_fc 148 | yaml_config['metrics']['robust_ferrari_canny']['friction_coef'] = value_fc 149 | #把每个摩擦力值当成键， 150 | force_closure_quality_config[value_fc] = GraspQualityConfigFactory.create_config( 151 | yaml_config['metrics']['force_closure']) 152 | canny_quality_config[value_fc] = GraspQualityConfigFactory.create_config( 153 | yaml_config['metrics']['robust_ferrari_canny']) 154 | 155 | #####################准备开始采样############################ 156 | #填充一个与摩擦数量相同的数组，每个对应的元素都是0 157 | good_count_perfect = np.zeros(len(fc_list)) 158 | count = 0 159 | #设置每个摩擦值需要计算的最少抓取数量 （根据指定输入值20） 160 | minimum_grasp_per_fc = grasp_amount 161 | #如果每个摩擦系数下，有效的抓取(满足力闭合或者其他判断标准)小于要求值，就一直循环查找，直到所有摩擦系数条件下至少都存在20个有效抓取 162 | while np.sum(good_count_perfect < minimum_grasp_per_fc) != 0: 163 | #开始使用antipodes sample获得对映随机抓取，此时并不判断是否满足力闭合，只是先采集满足夹爪条件的抓取 164 | #如果一轮多次随机采样之后，发现无法获得指定数量的随机抓取，就会重复迭代计算3次，之后放弃，并把已经找到的抓取返回来 165 | grasps = ags.generate_grasps(obj, target_num_grasps=sample_nums, grasp_gen_mult=10,max_iter=10, 166 | vis=False, random_approach_angle=True) 167 | count += len(grasps) 168 | #循环对每个采样抓取进行判断 169 | for j in grasps: 170 | tmp, is_force_closure = False, False 171 | #循环对某个采样抓取应用不同的抓取摩擦系数，判断是否是力闭合 172 | for ind_, value_fc in enumerate(fc_list): 173 | value_fc = round(value_fc, 2) 174 | tmp = is_force_closure 175 | #判断在当前给定的摩擦系数下，抓取是否是力闭合的 176 | is_force_closure = PointGraspMetrics3D.grasp_quality(j, obj, 177 | force_closure_quality_config[value_fc], vis=False) 178 | #假设当前,1号摩擦力为1.6 抓取不是力闭合的，但是上一个0号摩擦系数2.0 条件下抓取是力闭合的 179 | if tmp and not is_force_closure: 180 | #当0号2.0摩擦系数条件下采样的good抓取数量还不足指定的最低数量20 181 | if good_count_perfect[ind_ - 1] < minimum_grasp_per_fc: 182 | #以0号摩擦系数作为边界 183 | canny_quality = PointGraspMetrics3D.grasp_quality(j, obj, 184 | canny_quality_config[ 185 | round(fc_list[ind_ - 1], 2)], 186 | vis=False) 187 | good_grasp.append((j, round(fc_list[ind_ - 1], 2), canny_quality)) 188 | #在0号系数的good抓取下计数加1 189 | good_count_perfect[ind_ - 1] += 1 190 | #当前抓取j的边界摩擦系数找到了，退出摩擦循环，判断下一个抓取 191 | break 192 | #如果当前1号摩擦系数1.6条件下，该抓取j本身就是力闭合的，且摩擦系数是列表中的最后一个（所有的摩擦系数都判断完了） 193 | elif is_force_closure and value_fc == fc_list[-1]: 194 | if good_count_perfect[ind_] < minimum_grasp_per_fc: 195 | #以当前摩擦系数作为边界 196 | canny_quality = PointGraspMetrics3D.grasp_quality(j, obj, 197 | canny_quality_config[value_fc], vis=False) 198 | good_grasp.append((j, value_fc, canny_quality)) 199 | good_count_perfect[ind_] += 1 200 | #当前抓取j关于当前摩擦系数1.6判断完毕，而且满足所有的摩擦系数，就换到下一个摩擦系数 201 | break 202 | print('Object:{} GoodGrasp:{}'.format(object_name, good_count_perfect)) #判断 203 | 204 | object_name_len = len(object_name) 205 | object_name_ = str(object_name) + " " * (25 - object_name_len) 206 | if count == 0: 207 | good_grasp_rate = 0 208 | else: 209 | good_grasp_rate = len(good_grasp) / count 210 | print('Gripper:{} Object:{} Rate:{:.4f} {}/{}'. 211 | format(gripper_name, object_name_, good_grasp_rate, len(good_grasp), count)) 212 | 213 | 214 | if __name__ == '__main__': 215 | if len(sys.argv) > 1: 216 | filename_prefix = sys.argv[1] 217 | else: 218 | filename_prefix = "default" 219 | home_dir = os.environ['HOME'] 220 | #存放CAD模型的文件夹 221 | file_dir = home_dir + "/dataset/PointNetGPD/ycb_meshes_google/objects" #获取模型的路径 222 | file_list_all = get_file_name(file_dir) #返回一个列表，包含物体 223 | object_numbers = file_list_all.__len__() #获取文件夹中物体数量 224 | 225 | job_list = np.arange(object_numbers) #返回一个长度为object_numbers的元组 0 1 2 3 ... 226 | job_list = list(job_list) #转换为列表 227 | #设置同时对几个模型进行采样 228 | pool_size = 40 229 | assert (pool_size <= len(job_list)) 230 | # Initialize pool 231 | pool = [] #创建列表 232 | 233 | for _ in range(pool_size): #想多线程处理多个模型，但是实际上本代码每次只处理一个 234 | job_i = job_list.pop(0) #删除掉第0号值，并把job_i赋予0号值 235 | pool.append(multiprocessing.Process(target=do_job, args=(job_i,))) #在pool末尾添加元素 236 | [p.start() for p in pool] #启动多线程 237 | # refill 238 | 239 | while len(job_list) > 0: #如果有些没处理完 240 | for ind, p in enumerate(pool): 241 | if not p.is_alive(): 242 | pool.pop(ind) 243 | job_i = job_list.pop(0) 244 | p = multiprocessing.Process(target=do_job, args=(job_i,)) 245 | p.start() 246 | pool.append(p) 247 | break 248 | print('All job done.') 249 | -------------------------------------------------------------------------------- /apps/kinect2grasp.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python 2 | # -*- coding: utf-8 -*- 3 | # Author : Hongzhuo Liang 4 | # E-mail : liang@informatik.uni-hamburg.de 5 | # Description: 6 | # Date : 05/08/2018 6:04 PM 7 | # File Name : kinect2grasp.py 8 | 9 | import torch 10 | 11 | import rospy 12 | from sensor_msgs.msg import PointCloud2 13 | from visualization_msgs.msg import MarkerArray 14 | from visualization_msgs.msg import Marker 15 | import tf 16 | import moveit_commander 17 | import numpy as np 18 | #自定义pointcloud包 19 | import pointclouds 20 | #from pcl import PointCloud 21 | #自定义 22 | import voxelgrid 23 | 24 | 25 | import pcl 26 | #容易报错无法导入ruamel.yaml，需要使用命令 conda install ruamel.yaml 来安装，不能使用pip安装 27 | from autolab_core import YamlConfig 28 | from dexnet.grasping import RobotGripper 29 | from dexnet.grasping import GpgGraspSamplerPcl 30 | import os 31 | from pyquaternion import Quaternion 32 | import sys 33 | from os import path 34 | import time 35 | from scipy.stats import mode 36 | import multiprocessing as mp 37 | try: 38 | from gpd_grasp_msgs.msg import GraspConfig 39 | from gpd_grasp_msgs.msg import GraspConfigList 40 | except ImportError: 41 | print("Please install grasp msgs from https://github.com/TAMS-Group/gpd_grasp_msgs in your ROS workspace") 42 | exit() 43 | 44 | try: 45 | from mayavi import mlab 46 | except ImportError: 47 | print("Can not import mayavi") 48 | mlab = None 49 | #sys.path.append(path.dirname(path.dirname(path.dirname(path.abspath("__file__"))))) 50 | sys.path.append(os.environ['HOME'] + "/code/PointNetGPD") 51 | 52 | #把读取网络模型，外部命令参数的指令写在了main_test中 53 | from main_test import test_network, model, args 54 | 55 | # global config:全局的配置文件 56 | yaml_config = YamlConfig(os.environ['HOME'] + "/code/dex-net/test/config.yaml") 57 | gripper_name = 'panda' 58 | #加载夹爪 59 | gripper = RobotGripper.load(gripper_name, os.environ['HOME'] + "/code/dex-net/data/grippers") 60 | ags = GpgGraspSamplerPcl(gripper, yaml_config) 61 | 62 | value_fc = 0.4 # no use, set a random number 63 | num_grasps_single_worker = 20 64 | #如果使用多线程，每个线程采样num_grasps_p_worker个抓取 65 | num_grasps_p_worker=6 66 | #如果使用多线程，将一共使用num_workers个线程 67 | num_workers = 20 68 | max_num_samples = 150 69 | n_voxel = 500 70 | 71 | #输入pointnet的最小点数 72 | minimal_points_send_to_point_net = 20 73 | marker_life_time = 8 74 | 75 | show_bad_grasp = False 76 | save_grasp_related_file = False 77 | 78 | using_mp = args.using_mp 79 | show_final_grasp = args.show_final_grasp 80 | 81 | 82 | tray_grasp = args.tray_grasp 83 | single_obj_testing = False # if True, it will wait for input before get pointcloud 84 | 85 | #指定输入点云的点数 number of points put into neural network 86 | if args.model_type == "500": # minimal points send for training 87 | input_points_num = 500 88 | elif args.model_type == "750": 89 | input_points_num = 750 90 | elif args.model_type == "3class": 91 | input_points_num = 500 92 | else: 93 | input_points_num = 0 94 | 95 | #去除支撑桌面点 96 | def remove_table_points(points_voxel_, vis=False): 97 | xy_unique = np.unique(points_voxel_[:, 0:2], axis=0) 98 | new_points_voxel_ = points_voxel_ 99 | pre_del = np.zeros([1]) 100 | for i in range(len(xy_unique)): 101 | tmp = [] 102 | for j in range(len(points_voxel_)): 103 | if np.array_equal(points_voxel_[j, 0:2], xy_unique[i]): 104 | tmp.append(j) 105 | print(len(tmp)) 106 | if len(tmp) < 3: 107 | tmp = np.array(tmp) 108 | pre_del = np.hstack([pre_del, tmp]) 109 | if len(pre_del) != 1: 110 | pre_del = pre_del[1:] 111 | new_points_voxel_ = np.delete(points_voxel_, pre_del, 0) 112 | print("Success delete [[ {} ]] points from the table!".format(len(points_voxel_) - len(new_points_voxel_))) 113 | 114 | if vis: 115 | p = points_voxel_ 116 | mlab.points3d(p[:, 0], p[:, 1], p[:, 2], scale_factor=0.002, color=(1, 0, 0)) 117 | p = new_points_voxel_ 118 | mlab.points3d(p[:, 0], p[:, 1], p[:, 2], scale_factor=0.002, color=(0, 0, 1)) 119 | mlab.points3d(0, 0, 0, scale_factor=0.01, color=(0, 1, 0)) # plot 0 point 120 | mlab.show() 121 | return new_points_voxel_ 122 | 123 | 124 | def remove_white_pixel(msg, points_, vis=False): 125 | points_with_c_ = pointclouds.pointcloud2_to_array(msg) 126 | points_with_c_ = pointclouds.split_rgb_field(points_with_c_) 127 | r = np.asarray(points_with_c_['r'], dtype=np.uint32) 128 | g = np.asarray(points_with_c_['g'], dtype=np.uint32) 129 | b = np.asarray(points_with_c_['b'], dtype=np.uint32) 130 | rgb_colors = np.vstack([r, g, b]).T 131 | # rgb = rgb_colors.astype(np.float) / 255 132 | ind_good_points_ = np.sum(rgb_colors[:] < 210, axis=-1) == 3 133 | ind_good_points_ = np.where(ind_good_points_ == 1)[0] 134 | new_points_ = points_[ind_good_points_] 135 | if vis: 136 | p = points_ 137 | mlab.points3d(p[:, 0], p[:, 1], p[:, 2], scale_factor=0.002, color=(1, 0, 0)) 138 | p = new_points_ 139 | mlab.points3d(p[:, 0], p[:, 1], p[:, 2], scale_factor=0.002, color=(0, 0, 1)) 140 | mlab.points3d(0, 0, 0, scale_factor=0.01, color=(0, 1, 0)) # plot 0 point 141 | mlab.show() 142 | return new_points_ 143 | 144 | 145 | def get_voxel_fun(points_, n): 146 | get_voxel = voxelgrid.VoxelGrid(points_, n_x=n, n_y=n, n_z=n) 147 | get_voxel.compute() 148 | points_voxel_ = get_voxel.voxel_centers[get_voxel.voxel_n] 149 | points_voxel_ = np.unique(points_voxel_, axis=0) 150 | return points_voxel_ 151 | 152 | 153 | def cal_grasp(msg, cam_pos_): 154 | """根据在线采集的点云计算候选的抓取姿态 155 | """ 156 | #把pointcloud2类型的消息点云，转换为ndarray points_ 157 | points_ = pointclouds.pointcloud2_to_xyz_array(msg) 158 | #复制一份points_ ndarray对象，并将所有的点坐标转换为float32类型 159 | points_ = points_.astype(np.float32) 160 | 161 | remove_white = False 162 | if remove_white: 163 | points_ = remove_white_pixel(msg, points_, vis=True) 164 | # begin voxel points 165 | n = n_voxel # parameter related to voxel method 166 | # gpg improvements, highlights: flexible n parameter for voxelizing. 167 | #这一句话执行的时候，不能打开虚拟机，否则容易卡住 168 | points_voxel_ = get_voxel_fun(points_, n) 169 | 170 | #当点云点数小于2000时 171 | if len(points_) < 2000: # should be a parameter 172 | while len(points_voxel_) < len(points_)-15: 173 | points_voxel_ = get_voxel_fun(points_, n) 174 | n = n + 100 175 | rospy.loginfo("the voxel has {} points, we want get {} points".format(len(points_voxel_), len(points_))) 176 | 177 | rospy.loginfo("the voxel has {} points, we want get {} points".format(len(points_voxel_), len(points_))) 178 | # 179 | points_ = points_voxel_ 180 | remove_points = False 181 | #是否剔除支撑平面 182 | if remove_points: 183 | points_ = remove_table_points(points_, vis=True) 184 | #重新构造经过“降采样”的点云 185 | point_cloud = pcl.PointCloud(points_) 186 | 187 | print(len(points_)) 188 | #构造法向量估计对象 189 | norm = point_cloud.make_NormalEstimation() 190 | tree=point_cloud.make_kdtree() 191 | norm.set_SearchMethod(tree) 192 | #以周边30个点作为法向量计算点 193 | norm.set_KSearch(10) # critical parameter when calculating the norms 194 | normals = norm.compute() 195 | 196 | 197 | #将点云法向量转换为ndarry类型 198 | surface_normal = normals.to_array() 199 | 200 | surface_normal = surface_normal[:, 0:3] 201 | 202 | #每个点到 相机位置（无姿态）的向量 但是，感觉是相机到点的向量 203 | vector_p2cam = cam_pos_ - points_ 204 | #print(vector_p2cam) 205 | #print(cam_pos_) 206 | 207 | """ 208 | np.linalg.norm(vector_p2cam, axis=1) 默认求2范数，axis=1 代表按行向量处理，求多个行向量的2范数（求模） 209 | np.linalg.norm(vector_p2cam, axis=1).reshape(-1, 1) 将其调整为m行 1列 210 | 211 | 整句话的含义是，将vector_p2cam归一化，单位化 212 | """ 213 | vector_p2cam = vector_p2cam / np.linalg.norm(vector_p2cam, axis=1).reshape(-1, 1) 214 | 215 | 216 | 217 | #将表面法相与表面法相（都是单位向量）点乘，以备后面计算向量夹角 218 | tmp = np.dot(vector_p2cam, surface_normal.T).diagonal() 219 | #print(vector_p2cam) 220 | #print(surface_normal.T) 221 | #print(tmp) 222 | 223 | """ 224 | np.clip(tmp, -1.0, 1.0) 截取函数，将tmp中的值，都限制在-1.0到1.0之间，大于1的变成1，小于-1的记为-1 225 | np.arccos() 求解反余弦，求夹角 226 | """ 227 | angel = np.arccos(np.clip(tmp, -1.0, 1.0)) 228 | #print(angel) 229 | 230 | #找到与视角向量夹角大于90度的角（认为法向量计算错误） 231 | wrong_dir_norm = np.where(angel > np.pi * 0.5)[0] 232 | #print(np.where(angel > np.pi * 0.5)) 233 | #print(wrong_dir_norm) 234 | #print(len(wrong_dir_norm)) 235 | 236 | #创建一个len(angel)行，3列的ndarry对象 237 | tmp = np.ones([len(angel), 3]) 238 | #将法向量错误的行的元素都改写为-1 239 | tmp[wrong_dir_norm, :] = -1 240 | #与表面法相元素对元素相乘，作用是将"错误的"法向量的方向 扭转过来 241 | surface_normal = surface_normal * tmp 242 | #选取桌子以上2cm处的点作为检测点 243 | select_point_above_table = 0.020 244 | #modify of gpg: make it as a parameter. avoid select points near the table. 245 | #查看每个点的z方向，如果它们的点z轴方向的值大于select_point_above_table，就把他们抽出来 246 | points_for_sample = points_[np.where(points_[:, 2] > select_point_above_table)[0] ] 247 | print(len(points_for_sample)) 248 | if len(points_for_sample) == 0: 249 | rospy.loginfo("Can not seltect point, maybe the point cloud is too low?") 250 | return [], points_, surface_normal 251 | yaml_config['metrics']['robust_ferrari_canny']['friction_coef'] = value_fc 252 | 253 | grasps_together_ = [] 254 | 255 | if rospy.get_param("/robot_at_home") == "false": 256 | robot_at_home = False 257 | else: 258 | robot_at_home = True 259 | 260 | if not robot_at_home: 261 | rospy.loginfo("Robot is moving, waiting the robot go home.") 262 | elif not using_mp: 263 | rospy.loginfo("Begin cal grasps using single thread, slow!") 264 | grasps_together_ = ags.sample_grasps(point_cloud, points_for_sample, surface_normal, num_grasps_single_worker, 265 | max_num_samples=max_num_samples, show_final_grasp=show_final_grasp) 266 | else: 267 | # begin parallel grasp: 268 | rospy.loginfo("Begin cal grasps using parallel!") 269 | 270 | def grasp_task(num_grasps_, ags_, queue_): 271 | ret = ags_.sample_grasps(point_cloud, points_for_sample, surface_normal, num_grasps_, 272 | max_num_samples=max_num_samples, show_final_grasp=False) 273 | queue_.put(ret) 274 | 275 | queue = mp.Queue() 276 | 277 | #num_grasps_p_worker = int(num_grasps/num_workers) 278 | workers = [mp.Process(target=grasp_task, args=(num_grasps_p_worker, ags, queue)) for _ in range(num_workers)] 279 | [i.start() for i in workers] 280 | 281 | 282 | grasps_together_ = [] 283 | for i in range(num_workers): 284 | grasps_together_ = grasps_together_ + queue.get() 285 | rospy.loginfo("Finish mp processing!") 286 | 287 | if show_final_grasp and using_mp: 288 | ags.show_all_grasps(points_, grasps_together_) 289 | ags.show_points(points_, scale_factor=0.002) 290 | mlab.show() 291 | 292 | 293 | rospy.loginfo("Grasp sampler finish, generated {} grasps.".format(len(grasps_together_))) 294 | #返回抓取 场景的点 以及点云的表面法向量 295 | return grasps_together_, points_, surface_normal 296 | 297 | 298 | #检查碰撞 299 | def check_collision_square(grasp_bottom_center, approach_normal, binormal, 300 | minor_pc, points_, p, way="p_open"): 301 | #抓取坐标系 轴单位化 302 | approach_normal = approach_normal.reshape(1, 3) 303 | approach_normal = approach_normal / np.linalg.norm(approach_normal) 304 | binormal = binormal.reshape(1, 3) 305 | binormal = binormal / np.linalg.norm(binormal) 306 | minor_pc = minor_pc.reshape(1, 3) 307 | minor_pc = minor_pc / np.linalg.norm(minor_pc) 308 | #构建旋转矩阵 309 | matrix_ = np.hstack([approach_normal.T, binormal.T, minor_pc.T]) 310 | #正交矩阵的逆，就是它的转置 311 | grasp_matrix = matrix_.T 312 | #center=grasp_bottom_center+approach_normal*ags.gripper.hand_depth 313 | points_ = points_ - grasp_bottom_center.reshape(1, 3) 314 | #points_ = points_ - center.reshape(1, 3) 315 | 316 | tmp = np.dot(grasp_matrix, points_.T) 317 | points_g = tmp.T 318 | 319 | #选择是否使用与数据集采样相同的方式来采集点云 320 | use_dataset_py = False 321 | if not use_dataset_py: 322 | if way == "p_open": 323 | s1, s2, s4, s8 = p[1], p[2], p[4], p[8] 324 | elif way == "p_left": 325 | s1, s2, s4, s8 = p[9], p[1], p[10], p[12] 326 | elif way == "p_right": 327 | s1, s2, s4, s8 = p[2], p[13], p[3], p[7] 328 | elif way == "p_bottom": 329 | s1, s2, s4, s8 = p[11], p[15], p[12], p[20] 330 | else: 331 | raise ValueError('No way!') 332 | a1 = s1[1] < points_g[:, 1] 333 | a2 = s2[1] > points_g[:, 1] 334 | a3 = s1[2] > points_g[:, 2] 335 | a4 = s4[2] < points_g[:, 2] 336 | a5 = s4[0] > points_g[:, 0] 337 | a6 = s8[0] < points_g[:, 0] 338 | 339 | a = np.vstack([a1, a2, a3, a4, a5, a6]) 340 | points_in_area = np.where(np.sum(a, axis=0) == len(a))[0] 341 | if len(points_in_area) == 0: 342 | has_p = False 343 | else: 344 | has_p = True 345 | # for the way of pointGPD/dataset.py: 346 | else: 347 | width = ags.gripper.hand_outer_diameter - 2 * ags.gripper.finger_width 348 | x_limit = ags.gripper.hand_depth 349 | z_limit = width / 4 350 | y_limit = width / 2 351 | x1 = points_g[:, 0] > 0 352 | x2 = points_g[:, 0] < x_limit 353 | y1 = points_g[:, 1] > -y_limit 354 | y2 = points_g[:, 1] < y_limit 355 | z1 = points_g[:, 2] > -z_limit 356 | z2 = points_g[:, 2] < z_limit 357 | a = np.vstack([x1, x2, y1, y2, z1, z2]) 358 | points_in_area = np.where(np.sum(a, axis=0) == len(a))[0] 359 | if len(points_in_area) == 0: 360 | has_p = False 361 | else: 362 | has_p = True 363 | 364 | vis = False 365 | if vis: 366 | p = points_g 367 | mlab.points3d(p[:, 0], p[:, 1], p[:, 2], scale_factor=0.002, color=(0, 0, 1)) 368 | p = points_g[points_in_area] 369 | mlab.points3d(p[:, 0], p[:, 1], p[:, 2], scale_factor=0.002, color=(1, 0, 0)) 370 | p = ags.get_hand_points(np.array([0, 0, 0]), np.array([1, 0, 0]), np.array([0, 1, 0])) 371 | mlab.points3d(p[:, 0], p[:, 1], p[:, 2], scale_factor=0.005, color=(0, 1, 0)) 372 | mlab.show() 373 | 374 | return has_p, points_in_area, points_g 375 | 376 | 377 | def collect_pc(grasp_, pc): 378 | """ 379 | grasp_bottom_center, normal, major_pc, minor_pc 380 | grasp_是一个list形式的数据 381 | """ 382 | #获取抓取的数量 383 | grasp_num = len(grasp_) 384 | #将抓取转换为ndarry的形式，其实本身也就是 385 | grasp_ = np.array(grasp_) 386 | #print(grasp_) 387 | 388 | #变成n个，5行3列的 矩阵, 每个矩阵块，都是一个候选抓取坐标系 389 | grasp_ = grasp_.reshape(-1, 5, 3) # prevent to have grasp that only have number 1 390 | #print(grasp_) 391 | 392 | #使用每个块的第五行数据，作为抓取中心点，因为，第5个是修正后的夹爪姿态中心 393 | grasp_bottom_center = grasp_[:,4, :] 394 | approach_normal = grasp_[:, 1,:] 395 | binormal = grasp_[:, 2,:] 396 | minor_pc = grasp_[:, 3,:] 397 | 398 | in_ind_ = [] 399 | in_ind_points_ = [] 400 | #获取夹爪角点 401 | p = ags.get_hand_points(np.array([0, 0, 0]), np.array([1, 0, 0]), np.array([0, 1, 0])) 402 | 403 | for i_ in range(grasp_num): 404 | #通过实验发现，这个函数并不是只返回了夹爪open区域内部的点，而是"扩大"了一些， 405 | # 估计是为了返回足够多的点 406 | #但是，问题是这需要与训练时候的样本点相对应啊，如果训练的时候，就没有"扩大"，那网络的打分会高么？ 407 | #所以需要去看训练部分的采样部分两点 408 | # 1.生成的样本点云是否是以固定夹爪为参考系的？ 409 | # 2.训练时，采样点的提取部分，是否进行了"扩大"? 410 | has_p, in_ind_tmp, points_g = check_collision_square(grasp_bottom_center[i_], approach_normal[i_], 411 | binormal[i_], minor_pc[i_], pc, p) 412 | #把索引添加进list中保存 413 | in_ind_.append(in_ind_tmp) 414 | #从下面这句话可以看出 415 | # 这些返回的夹爪内部点，是已经旋转过后的点，而不是原始的点，为什么？ 416 | #是因为，网络输入的点要求是以固定夹爪为参考系的么？ 417 | #猜测是的，为什么？因为，这样子的话，就可以在传入网络点云的同时， 418 | # 相当于也传入了夹爪的姿态 419 | 420 | 421 | 422 | #这一点，他想把grasp_g中的点坐标也都保存下来，in_ind_points_是一个list 423 | in_ind_points_.append(points_g[in_ind_tmp]) 424 | 425 | 426 | #显示出截取的夹爪内部区域的点云（红色）&夹爪&整体点云 427 | if 0: 428 | #p = ags.get_hand_points(np.array([0, 0, 0]), np.array([1, 0, 0]), np.array([0, 1, 0])) 429 | ags.show_grasp_3d(p) 430 | ags.show_points(points_g) 431 | ags.show_points(in_ind_points_[i_],color='r',scale_factor=0.005) 432 | #table_points = np.array([[-1, 1, 0], [1, 1, 0], [1, -1, 0], [-1, -1, 0]]) * 0.5 433 | #triangles = [(1, 2, 3), (0, 1, 3)] 434 | #mlab.triangular_mesh(table_points[:, 0], table_points[:, 1], table_points[:, 2], 435 | #triangles, color=(0.8, 0.8, 0.8), opacity=0.5) 436 | mlab.show() 437 | 438 | return in_ind_, in_ind_points_ 439 | 440 | 441 | def show_marker(marker_array_, pos_, ori_, scale_, color_, lifetime_): 442 | """显示标注物体 443 | """ 444 | marker_ = Marker() 445 | marker_.header.frame_id = "/ar_marker_6" 446 | # marker_.header.stamp = rospy.Time.now() 447 | marker_.type = marker_.CUBE 448 | marker_.action = marker_.ADD 449 | 450 | marker_.pose.position.x = pos_[0] 451 | marker_.pose.position.y = pos_[1] 452 | marker_.pose.position.z = pos_[2] 453 | marker_.pose.orientation.x = ori_[1] 454 | marker_.pose.orientation.y = ori_[2] 455 | marker_.pose.orientation.z = ori_[3] 456 | marker_.pose.orientation.w = ori_[0] 457 | 458 | marker_.lifetime = rospy.Duration.from_sec(lifetime_) 459 | marker_.scale.x = scale_[0] 460 | marker_.scale.y = scale_[1] 461 | marker_.scale.z = scale_[2] 462 | marker_.color.a = 0.5 463 | red_, green_, blue_ = color_ 464 | marker_.color.r = red_ 465 | marker_.color.g = green_ 466 | marker_.color.b = blue_ 467 | marker_array_.markers.append(marker_) 468 | 469 | 470 | def show_grasp_marker(marker_array_, real_grasp_, gripper_, color_, lifetime_): 471 | """ 472 | show grasp using marker使用marker来显示抓取 473 | :param marker_array_: marker array 474 | :param real_grasp_: [0] position, [1] approach [2] binormal [3] minor pc 475 | :param gripper_: gripper parameter of a grasp 476 | :param color_: color of the gripper 显示夹爪的颜色 477 | :param lifetime_: time for showing the maker marker的显示时间长短 478 | :return: return add makers to the maker array 479 | 480 | """ 481 | hh = gripper_.hand_height 482 | fw = gripper_.real_finger_width 483 | hod = gripper_.hand_outer_diameter 484 | hd = gripper_.real_hand_depth 485 | open_w = hod - fw * 2 486 | 487 | approach = real_grasp_[1] 488 | binormal = real_grasp_[2] 489 | minor_pc = real_grasp_[3] 490 | grasp_bottom_center = real_grasp_[4] - approach * (gripper_.real_hand_depth - gripper_.hand_depth) 491 | 492 | rotation = np.vstack([approach, binormal, minor_pc]).T 493 | qua = Quaternion(matrix=rotation) 494 | 495 | marker_bottom_pos = grasp_bottom_center - approach * hh * 0.5 496 | marker_left_pos = grasp_bottom_center - binormal * (open_w * 0.5 + fw * 0.5) + hd * 0.5 * approach 497 | marker_right_pos = grasp_bottom_center + binormal * (open_w * 0.5 + fw * 0.5) + hd * 0.5 * approach 498 | show_marker(marker_array_, marker_bottom_pos, qua, np.array([hh, hod, hh]), color_, lifetime_) 499 | show_marker(marker_array_, marker_left_pos, qua, np.array([hd, fw, hh]), color_, lifetime_) 500 | show_marker(marker_array_, marker_right_pos, qua, np.array([hd, fw, hh]), color_, lifetime_) 501 | 502 | 503 | def check_hand_points_fun(real_grasp_): 504 | """该函数是计算处于夹爪内部有多少个点（只求数量） 505 | """ 506 | ind_points_num = [] 507 | for i in range(len(real_grasp_)): 508 | #修正后的抓取中心点坐标（位于第5个） 509 | grasp_bottom_center = real_grasp_[i][4] 510 | approach_normal = real_grasp_[i][1] 511 | binormal = real_grasp_[i][2] 512 | minor_pc = real_grasp_[i][3] 513 | #固定了手，而去变换点云，原因是，这样更容易计算碰撞，可以查看ags.check_collision_square函数内部 514 | local_hand_points = ags.get_hand_points(np.array([0, 0, 0]), np.array([1, 0, 0]), np.array([0, 1, 0])) 515 | #检查open区域内的点数量，返回点的索引list 516 | has_points_tmp, ind_points_tmp = ags.check_collision_square(grasp_bottom_center, approach_normal, 517 | binormal, minor_pc, points, 518 | local_hand_points, "p_open") 519 | #打印list中的点数 520 | ind_points_num.append(len(ind_points_tmp)) 521 | print(ind_points_num) 522 | #这是啥意思？ 523 | file_name = "./generated_grasps/real_points/" + str(np.random.randint(300)) + str(len(real_grasp_)) + ".npy" 524 | np.save(file_name, np.array(ind_points_num)) 525 | 526 | 527 | def get_grasp_msg(real_good_grasp_, score_value_): 528 | """ 529 | 创建一个ROS官方的抓取配置消息 530 | 里面涵盖了抓取坐标系的位置与姿态定义 531 | """ 532 | grasp_bottom_center_modify = real_good_grasp_[4] 533 | approach = real_good_grasp_[1] 534 | binormal = real_good_grasp_[2] 535 | minor_pc = real_good_grasp_[3] 536 | 537 | grasp_config_ = GraspConfig() 538 | # 539 | top_p_ = grasp_bottom_center_modify + approach * ags.gripper.hand_depth 540 | grasp_config_.bottom.x = grasp_bottom_center_modify[0] 541 | grasp_config_.bottom.y = grasp_bottom_center_modify[1] 542 | grasp_config_.bottom.z = grasp_bottom_center_modify[2] 543 | grasp_config_.top.x = top_p_[0] 544 | grasp_config_.top.y = top_p_[1] 545 | grasp_config_.top.z = top_p_[2] 546 | #抓取的三个坐标轴向量 547 | grasp_config_.approach.x = approach[0] 548 | grasp_config_.approach.y = approach[1] 549 | grasp_config_.approach.z = approach[2] 550 | grasp_config_.binormal.x = binormal[0] 551 | grasp_config_.binormal.y = binormal[1] 552 | grasp_config_.binormal.z = binormal[2] 553 | grasp_config_.axis.x = minor_pc[0] 554 | grasp_config_.axis.y = minor_pc[1] 555 | grasp_config_.axis.z = minor_pc[2] 556 | #该抓取的分数 557 | grasp_config_.score.data = score_value_ 558 | 559 | return grasp_config_ 560 | 561 | 562 | def remove_grasp_outside_tray(grasps_, points_): 563 | x_min = points_[:, 0].min() 564 | x_max = points_[:, 0].max() 565 | y_min = points_[:, 1].min() 566 | y_max = points_[:, 1].max() 567 | valid_grasp_ind_ = [] 568 | for i in range(len(grasps_)): 569 | grasp_bottom_center = grasps_[i][4] 570 | approach_normal = grasps_[i][1] 571 | major_pc = grasps_[i][2] 572 | hand_points_ = ags.get_hand_points(grasp_bottom_center, approach_normal, major_pc) 573 | finger_points_ = hand_points_[[1, 2, 3, 4, 9, 10, 13, 14], :] 574 | # aa = points_[:, :2] - finger_points_[0][:2] # todo： work of remove outside grasp not finished. 575 | 576 | # from IPython import embed;embed() 577 | a = finger_points_[:, 0] < x_min 578 | b = finger_points_[:, 0] > x_max 579 | c = finger_points_[:, 1] < y_min 580 | d = finger_points_[:, 1] > y_max 581 | if np.sum(a) + np.sum(b) + np.sum(c) + np.sum(d) == 0: 582 | valid_grasp_ind_.append(i) 583 | grasps_inside_ = [grasps_[i] for i in valid_grasp_ind_] 584 | rospy.loginfo("gpg got {} grasps, after remove grasp outside tray, {} grasps left".format(len(grasps_), 585 | len(grasps_inside_))) 586 | return grasps_inside_ 587 | 588 | 589 | if __name__ == '__main__': 590 | """ 591 | definition of gotten grasps: 592 | 593 | grasp_bottom_center = grasp_[0] 594 | approach_normal = grasp_[1] 595 | binormal = grasp_[2] 596 | """ 597 | #初始化节点，把节点名称写为grasp_tf_broadcaster 抓取发布器，anonymous参数在 598 | # 为True的时候会在原本节点名字的后面加一串随机数，来保证可以同时开启多个同样的 599 | # 节点，如果为false的话就只能开一个 600 | rospy.init_node('grasp_tf_broadcaster', anonymous=True) 601 | #创建发布器，用于显示抓取 602 | pub1 = rospy.Publisher('gripper_vis', MarkerArray, queue_size=1) 603 | #发布检测到的抓取，用于抓取 604 | pub2 = rospy.Publisher('/detect_grasps/clustered_grasps', GraspConfigList, queue_size=1) 605 | 606 | # 607 | rate = rospy.Rate(10) 608 | #在ros参数服务器上，设置一个参数 609 | rospy.set_param("/robot_at_home", "true") # only use when in simulation test. 610 | rospy.loginfo("getting transform from kinect2 to table top") 611 | 612 | #cam_pos列表 613 | cam_pos = [] 614 | #创建TF监听器 615 | listener = tf.TransformListener() 616 | #是否得到变换的标志位 617 | get_transform = False 618 | #等待tf中查找到'/table_top'与'/kinect2_ir_optical_frame' 619 | # 两个坐标系之间的变换关系 620 | while not get_transform: 621 | try: 622 | #查看kinect2相机与桌子之间的关系，确保相机已经能够看到标签ar_marker_6 623 | #cam_pos代表的是相机的trans，不是rot 624 | cam_pos, _ = listener.lookupTransform('/ar_marker_6', '/kinect2_rgb_optical_frame', rospy.Time(0)) 625 | print(cam_pos) 626 | print(_) 627 | get_transform = True 628 | rospy.loginfo("got transform complete") 629 | except (tf.LookupException, tf.ConnectivityException, tf.ExtrapolationException): 630 | continue 631 | 632 | 633 | while not rospy.is_shutdown(): 634 | #检测机械臂是否处于home状态 635 | if rospy.get_param("/robot_at_home") == "false": 636 | robot_at_home = False 637 | else: 638 | robot_at_home = True 639 | #如果机械臂状态不在home（移动中） 640 | if not robot_at_home: 641 | rospy.loginfo("Robot is moving, waiting the robot go home.") 642 | #就跳过后面一直循环等待移动完毕 643 | continue 644 | else: 645 | #在home状态，可采集点云 646 | rospy.loginfo("Robot is at home, safely catching point cloud data.") 647 | if single_obj_testing: 648 | input("Pleas put object on table and press any number to continue!") 649 | 650 | rospy.loginfo("rospy is waiting for message: /table_top_points") 651 | 652 | """点云数据的名称是/table_top_points 653 | 对象是kinect_data 类形是 PointCloud2类型 654 | """ 655 | #kinect_data = rospy.wait_for_message("/table_top_points", PointCloud2) 656 | kinect_data_ = rospy.wait_for_message("/table_top_points", PointCloud2) 657 | kinect_data = rospy.wait_for_message("/table_top_points_subsampled", PointCloud2) 658 | 659 | real_good_grasp = [] 660 | real_bad_grasp = [] 661 | real_score_value = [] 662 | 663 | 664 | #设置每个候选抓取的夹爪内部的点 665 | # 需要送入pointnet中训练几次 666 | repeat = 3 # speed up this try 10 time is too time consuming 667 | ###########################开始抓取姿态检测############################### 668 | # begin of grasp detection 开始抓取检测 669 | # if there is no point cloud on table, waiting for point cloud. 670 | if kinect_data.data == '': 671 | rospy.loginfo("There is no points on the table, waiting...") 672 | continue 673 | 674 | 675 | #获取当前文件所在目录（文件夹）的绝对路径 676 | path=os.path.dirname(os.path.abspath(__file__)) 677 | #更改（确保）当前所在目录是工作目录 678 | os.chdir(path) 679 | 680 | """根据Kinect读取到的场景点云，使用gpd检测候选的抓取？ 681 | 输入： 682 | kinect_data读取的点云数据 683 | cam_pos Kinect与桌子标签之间的距离 684 | grasp_sampled 生成的所有候选抓取（不与桌子以及场景物体碰撞的抓取） 685 | points 函数中处理后，用于计算抓取的点云，ndarry形式 686 | normals_cal points对应的所有法相量 ndarry形式 687 | """ 688 | grasp_sampled, points, normals_cal = cal_grasp(kinect_data, cam_pos) 689 | #检查生成的抓取数量是否为0 690 | if len(grasp_sampled)==0: 691 | rospy.loginfo("Nice try!") 692 | continue 693 | 694 | 695 | #托盘，如果有托盘？ 696 | if tray_grasp: 697 | #grasp_sampled 去除外部托盘导致的抓取； 698 | grasp_sampled = remove_grasp_outside_tray(grasp_sampled, points) 699 | #估计一个抓取中的点数 700 | check_grasp_points_num = False # evaluate the number of points in a grasp 701 | 702 | """ 703 | 等效于 704 | if check_grasp_points_num: 705 | check_hand_points_fun(grasp_sampled) 706 | else: 707 | 0 708 | """ 709 | check_hand_points_fun(grasp_sampled) if check_grasp_points_num else 0 710 | 711 | #计算，每个抓取，夹爪内部的点云 712 | in_ind, in_ind_points= collect_pc(grasp_sampled, points) 713 | 714 | 715 | #保存在线抓取抓取相关的文件 716 | if save_grasp_related_file: 717 | np.save("./generated_grasps/points.npy", points) 718 | np.save("./generated_grasps/in_ind.npy", in_ind) 719 | np.save("./generated_grasps/grasp_sampled.npy", grasp_sampled) 720 | np.save("./generated_grasps/cal_norm.npy", normals_cal) 721 | #打分 722 | score = [] # should be 0 or 1 723 | score_value = [] # should be float [0, 1] 724 | 725 | ind_good_grasp = [] 726 | ind_bad_grasp = [] 727 | ############################把检测到的抓取送到pointnet中打分#################################### 728 | rospy.loginfo("Begin send grasp into pointnet, cal grasp score") 729 | 730 | #in_ind_points 是一个 list 731 | for ii in range(len(in_ind_points)): 732 | """ 733 | 首先需要保证机械臂处于home状态 734 | """ 735 | if rospy.get_param("/robot_at_home") == "false": 736 | robot_at_home = False 737 | else: 738 | robot_at_home = True 739 | if not robot_at_home: 740 | rospy.loginfo("robot is not at home, stop calculating the grasp score") 741 | break 742 | """ 743 | 判断，夹爪内部的点数量，是否满足最小点数（20），如果小于20的话就别扩充啥的了 744 | 太少了，可能就只是噪点抓取而已 745 | """ 746 | if in_ind_points[ii].shape[0] < minimal_points_send_to_point_net: 747 | rospy.loginfo("Mark as bad grasp! Only {} points, should be at least {} points.".format( 748 | in_ind_points[ii].shape[0], minimal_points_send_to_point_net)) 749 | #第一个分数记为0 750 | score.append(0) 751 | score_value.append(0.0) 752 | #如果需要记下来，哪个是bad抓取的话，就把哪个抓取给记下来 753 | if show_bad_grasp: 754 | ind_bad_grasp.append(ii) 755 | 756 | #内部点数满足要求 757 | else: 758 | predict = [] 759 | grasp_score = [] 760 | """for _ in range() 其中_代表临时变量，说明只关心完成循环而已，不关心循环到第几个 761 | 给定一个抓取，找到夹爪内部的点，然后重复repeat次输入pointnet中，进行多次打分 762 | 因为下面的采集点具有部分随机性，所以需要重复输入几次，但是稍微比较费时 763 | """ 764 | for _ in range(repeat): 765 | #如果在线采集的夹爪内部点云点数量是大于 指定的点数量 766 | if len(in_ind_points[ii]) >= input_points_num: 767 | #随机抽选其中的一些点，保证和要求的点数量是一致的 768 | points_modify = in_ind_points[ii][np.random.choice(len(in_ind_points[ii]), 769 | input_points_num, replace=False)] 770 | #如果在线采集的夹爪内部点云点数量是小于 指定的点数量 771 | else: 772 | #就是补上去一些没有用的点，但是怎么补的，没看清楚？？？？？？？？？？？？？？？？？？ 773 | points_modify = in_ind_points[ii][np.random.choice(len(in_ind_points[ii]), 774 | input_points_num, replace=True)] 775 | if 0: 776 | p = ags.get_hand_points(np.array([0, 0, 0]), np.array([1, 0, 0]), np.array([0, 1, 0])) 777 | ags.show_grasp_3d(p) 778 | print("原始点数{},修改后点数{}".format(len(in_ind_points[ii]),len(points_modify))) 779 | ags.show_points(points_modify,color='b',scale_factor=0.005) 780 | ags.show_points(in_ind_points[ii],color='r',scale_factor=0.003) 781 | #table_points = np.array([[-1, 1, 0], [1, 1, 0], [1, -1, 0], [-1, -1, 0]]) * 0.5 782 | #triangles = [(1, 2, 3), (0, 1, 3)] 783 | #mlab.triangular_mesh(table_points[:, 0], table_points[:, 1], table_points[:, 2], 784 | #triangles, color=(0.8, 0.8, 0.8), opacity=0.5) 785 | mlab.show() 786 | 787 | 788 | 789 | """在这里输入网络打分！ 790 | 返回的是？，grasp打分 791 | """ 792 | if_good_grasp, grasp_score_tmp = test_network(model.eval(), points_modify) 793 | 794 | 795 | predict.append(if_good_grasp.item()) 796 | #保存分数，添加到list中（这个list中都是同一个抓取的分数，重复了几次） 797 | grasp_score.append(grasp_score_tmp) 798 | 799 | predict_vote = mode(predict)[0][0] # vote from all the "repeat" results. 800 | #把list转换为np.array 801 | grasp_score = np.array(grasp_score) 802 | #如果是3class的？？？？？？？？？？？？？？？ 803 | if args.model_type == "3class": # the best in 3 class classification is the last column, third column 804 | which_one_is_best = 2 # should set as 2 805 | #两分类是第二个 806 | else: # for two class classification best is the second column (also the last column) 807 | which_one_is_best = 1 # should set as 1 808 | 809 | """np.mean()函数功能：求取均值numpy.mean(a, axis, dtype, out，keepdims ) 810 | 经常操作的参数为axis，以m * n矩阵举例： 811 | axis 不设置值，对 m*n 个数求均值，返回一个实数 812 | axis = 0：压缩行，对各列求均值，返回 1* n 矩阵 813 | axis =1 ：压缩列，对各行求均值，返回 m *1 矩阵 814 | """ 815 | score_vote = np.mean(grasp_score[np.where(predict == predict_vote)][:, 0, which_one_is_best]) 816 | score.append(predict_vote) 817 | score_value.append(score_vote) 818 | 819 | if score[ii] == which_one_is_best: 820 | ind_good_grasp.append(ii) 821 | else: 822 | if show_bad_grasp: 823 | #存放“坏抓取”的抓取编号 824 | ind_bad_grasp.append(ii) 825 | #记录好的抓取的数量，以及坏的抓取数量 826 | print("Got {} good grasps, and {} bad grasps".format(len(ind_good_grasp), 827 | len(in_ind_points) - len(ind_good_grasp))) 828 | 829 | 830 | if len(ind_good_grasp) != 0: 831 | real_good_grasp = [grasp_sampled[i] for i in ind_good_grasp] 832 | real_score_value = [score_value[i] for i in ind_good_grasp] 833 | if show_bad_grasp: 834 | real_bad_grasp = [grasp_sampled[i] for i in ind_bad_grasp] 835 | # end of grasp detection抓取检测部分结束 836 | 837 | 838 | #real_bad_grasp = [grasp_sampled[i] for i in ind_bad_grasp] 839 | 840 | 841 | # get sorted ind by the score values对检测到的抓取通过打分进行排序 842 | sorted_value_ind = list(index for index, item in sorted(enumerate(real_score_value), 843 | key=lambda item: item[1], 844 | reverse=True)) 845 | # sort grasps using the ind 846 | sorted_real_good_grasp = [real_good_grasp[i] for i in sorted_value_ind] 847 | real_good_grasp = sorted_real_good_grasp 848 | # get the sorted score value, from high to low 849 | real_score_value = sorted(real_score_value, reverse=True) 850 | 851 | 852 | 853 | marker_array = MarkerArray() 854 | marker_array_single = MarkerArray() 855 | 856 | 857 | grasp_msg_list = GraspConfigList() 858 | 859 | #按照顺序将所有的好的抓取添加显示，应该是第0个是最优的抓取 860 | for i in range(len(real_good_grasp)): 861 | grasp_msg = get_grasp_msg(real_good_grasp[i], real_score_value[i]) 862 | grasp_msg_list.grasps.append(grasp_msg) 863 | #把所有的抓取都显示成为绿色的，都添加进入那个显示序列中 864 | for i in range(len(real_good_grasp)): 865 | show_grasp_marker(marker_array, real_good_grasp[i], gripper, (0, 1, 0), marker_life_time) 866 | 867 | if show_bad_grasp: 868 | for i in range(len(real_bad_grasp)): 869 | show_grasp_marker(marker_array, real_bad_grasp[i], gripper, (1, 0, 0), marker_life_time) 870 | 871 | id_ = 0 872 | for m in marker_array.markers: 873 | m.id = id_ 874 | id_ += 1 875 | 876 | grasp_msg_list.header.stamp = rospy.Time.now() 877 | grasp_msg_list.header.frame_id = "/ar_marker_6" 878 | 879 | # from IPython import embed;embed() 880 | if len(real_good_grasp) != 0: 881 | #if 0: 882 | i = 0 883 | #第0个抓取得分是最高的，以红色显示出来 884 | single_grasp_list_pub = GraspConfigList() 885 | single_grasp_list_pub.header.stamp = rospy.Time.now() 886 | single_grasp_list_pub.header.frame_id = "/ar_marker_6" 887 | grasp_msg = get_grasp_msg(real_good_grasp[i], real_score_value[i]) 888 | single_grasp_list_pub.grasps.append(grasp_msg) 889 | show_grasp_marker(marker_array_single, real_good_grasp[i], gripper, (1, 0, 0), marker_life_time+5) 890 | 891 | for m in marker_array_single.markers: 892 | m.id = id_ 893 | id_ += 1 894 | pub1.publish(marker_array) 895 | rospy.sleep(4) 896 | pub2.publish(single_grasp_list_pub) 897 | pub1.publish(marker_array_single) 898 | print(grasp_msg.approach) 899 | print(grasp_msg.binormal) 900 | print(grasp_msg.axis) 901 | 902 | # pub2.publish(grasp_msg_list) 903 | rospy.loginfo(" Publishing grasp pose to rviz using marker array and good grasp pose") 904 | 905 | rate.sleep() 906 | -------------------------------------------------------------------------------- /apps/pointclouds.py: -------------------------------------------------------------------------------- 1 | # Software License Agreement (BSD License) 2 | # 3 | # Copyright (c) 2008, Willow Garage, Inc. 4 | # All rights reserved. 5 | # 6 | # Redistribution and use in source and binary forms, with or without 7 | # modification, are permitted provided that the following conditions 8 | # are met: 9 | # 10 | # * Redistributions of source code must retain the above copyright 11 | # notice, this list of conditions and the following disclaimer. 12 | # * Redistributions in binary form must reproduce the above 13 | # copyright notice, this list of conditions and the following 14 | # disclaimer in the documentation and/or other materials provided 15 | # with the distribution. 16 | # * Neither the name of Willow Garage, Inc. nor the names of its 17 | # contributors may be used to endorse or promote products derived 18 | # from this software without specific prior written permission. 19 | # 20 | # THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS 21 | # "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT 22 | # LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS 23 | # FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE 24 | # COPYRIGHT OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, 25 | # INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, 26 | # BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; 27 | # LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER 28 | # CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT 29 | # LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN 30 | # ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE 31 | # POSSIBILITY OF SUCH DAMAGE. 32 | # 33 | # Author: Jon Binney 34 | # Original author: Jon Binney 35 | # Modify by: Hongzhuo Liang 36 | # Time: Aug 6 2018 37 | # Download from https://github.com/kfu/metu-ros-pkg.git 38 | """ 39 | Functions for working with PointCloud2. 40 | Tested on kinect2 and works well. 41 | Python2 is needed. Do not support python3 42 | """ 43 | 44 | __docformat__ = "restructuredtext en" 45 | 46 | 47 | import numpy as np 48 | from sensor_msgs.msg import PointCloud2, PointField 49 | 50 | # prefix to the names of dummy fields we add to get byte alignment correct. this needs to not 51 | # clash with any actual field names 52 | DUMMY_FIELD_PREFIX = '__' 53 | 54 | # mappings between PointField types and numpy types 55 | type_mappings = [(PointField.INT8, np.dtype('int8')), (PointField.UINT8, np.dtype('uint8')), 56 | (PointField.INT16, np.dtype('int16')), 57 | (PointField.UINT16, np.dtype('uint16')), (PointField.INT32, np.dtype('int32')), 58 | (PointField.UINT32, np.dtype('uint32')), 59 | (PointField.FLOAT32, np.dtype('float32')), (PointField.FLOAT64, np.dtype('float64'))] 60 | pftype_to_nptype = dict(type_mappings) 61 | nptype_to_pftype = dict((nptype, pftype) for pftype, nptype in type_mappings) 62 | 63 | # sizes (in bytes) of PointField types 64 | pftype_sizes = {PointField.INT8: 1, PointField.UINT8: 1, PointField.INT16: 2, PointField.UINT16: 2, 65 | PointField.INT32: 4, PointField.UINT32: 4, PointField.FLOAT32: 4, PointField.FLOAT64: 8} 66 | 67 | 68 | def pointcloud2_to_dtype(cloud_msg): 69 | """Convert a list of PointFields to a numpy record datatype. 70 | """ 71 | offset = 0 72 | np_dtype_list = [] 73 | for f in cloud_msg.fields: 74 | while offset < f.offset: 75 | # might be extra padding between fields 76 | np_dtype_list.append(('%s%d' % (DUMMY_FIELD_PREFIX, offset), np.uint8)) 77 | offset += 1 78 | np_dtype_list.append((f.name, pftype_to_nptype[f.datatype])) 79 | offset += pftype_sizes[f.datatype] 80 | 81 | # might be extra padding between points 82 | while offset < cloud_msg.point_step: 83 | np_dtype_list.append(('%s%d' % (DUMMY_FIELD_PREFIX, offset), np.uint8)) 84 | offset += 1 85 | 86 | return np_dtype_list 87 | 88 | 89 | def arr_to_fields(cloud_arr): 90 | """Convert a numpy record datatype into a list of PointFields. 91 | """ 92 | fields = [] 93 | for field_name in cloud_arr.dtype.names: 94 | np_field_type, field_offset = cloud_arr.dtype.fields[field_name] 95 | pf = PointField() 96 | pf.name = field_name 97 | pf.datatype = nptype_to_pftype[np_field_type] 98 | pf.offset = field_offset 99 | pf.count = 1 # is this ever more than one? 100 | fields.append(pf) 101 | return fields 102 | 103 | 104 | def pointcloud2_to_array(cloud_msg, split_rgb=False): 105 | """ 106 | Converts a rospy PointCloud2 message to a numpy recordarray 107 | 108 | Reshapes the returned array to have shape (height, width), even if the height is 1. 109 | 110 | The reason for using np.fromstring rather than struct.unpack is speed... especially 111 | for large point clouds, this will be faster. 112 | """ 113 | # construct a numpy record type equivalent to the point type of this cloud 114 | dtype_list = pointcloud2_to_dtype(cloud_msg) 115 | 116 | # parse the cloud into an array 117 | cloud_arr = np.fromstring(cloud_msg.data, dtype_list) 118 | 119 | # remove the dummy fields that were added 120 | cloud_arr = cloud_arr[ 121 | [fname for fname, _type in dtype_list if not (fname[:len(DUMMY_FIELD_PREFIX)] == DUMMY_FIELD_PREFIX)]] 122 | 123 | if split_rgb: 124 | cloud_arr = split_rgb_field(cloud_arr) 125 | 126 | return np.reshape(cloud_arr, (cloud_msg.height, cloud_msg.width)) 127 | 128 | 129 | def array_to_pointcloud2(cloud_arr, stamp=None, frame_id=None, merge_rgb=False): 130 | """Converts a numpy record array to a sensor_msgs.msg.PointCloud2. 131 | """ 132 | if merge_rgb: 133 | cloud_arr = merge_rgb_fields(cloud_arr) 134 | 135 | # make it 2d (even if height will be 1) 136 | cloud_arr = np.atleast_2d(cloud_arr) 137 | 138 | cloud_msg = PointCloud2() 139 | 140 | if stamp is not None: 141 | cloud_msg.header.stamp = stamp 142 | if frame_id is not None: 143 | cloud_msg.header.frame_id = frame_id 144 | cloud_msg.height = cloud_arr.shape[0] 145 | cloud_msg.width = cloud_arr.shape[1] 146 | cloud_msg.fields = arr_to_fields(cloud_arr) 147 | cloud_msg.is_bigendian = False # assumption 148 | cloud_msg.point_step = cloud_arr.dtype.itemsize 149 | cloud_msg.row_step = cloud_msg.point_step*cloud_arr.shape[1] 150 | cloud_msg.is_dense = all([np.isfinite(cloud_arr[fname]).all() for fname in cloud_arr.dtype.names]) 151 | cloud_msg.data = cloud_arr.tostring() 152 | return cloud_msg 153 | 154 | 155 | def merge_rgb_fields(cloud_arr): 156 | """Takes an array with named np.uint8 fields 'r', 'g', and 'b', and returns an array in 157 | which they have been merged into a single np.float32 'rgb' field. The first byte of this 158 | field is the 'r' uint8, the second is the 'g', uint8, and the third is the 'b' uint8. 159 | 160 | This is the way that pcl likes to handle RGB colors for some reason. 161 | """ 162 | r = np.asarray(cloud_arr['r'], dtype=np.uint32) 163 | g = np.asarray(cloud_arr['g'], dtype=np.uint32) 164 | b = np.asarray(cloud_arr['b'], dtype=np.uint32) 165 | rgb_arr = np.array((r << 16) | (g << 8) | (b << 0), dtype=np.uint32) 166 | 167 | # not sure if there is a better way to do this. i'm changing the type of the array 168 | # from uint32 to float32, but i don't want any conversion to take place -jdb 169 | rgb_arr.dtype = np.float32 170 | 171 | # create a new array, without r, g, and b, but with rgb float32 field 172 | new_dtype = [] 173 | for field_name in cloud_arr.dtype.names: 174 | field_type, field_offset = cloud_arr.dtype.fields[field_name] 175 | if field_name not in ('r', 'g', 'b'): 176 | new_dtype.append((field_name, field_type)) 177 | new_dtype.append(('rgb', np.float32)) 178 | new_cloud_arr = np.zeros(cloud_arr.shape, new_dtype) 179 | 180 | # fill in the new array 181 | for field_name in new_cloud_arr.dtype.names: 182 | if field_name == 'rgb': 183 | new_cloud_arr[field_name] = rgb_arr 184 | else: 185 | new_cloud_arr[field_name] = cloud_arr[field_name] 186 | 187 | return new_cloud_arr 188 | 189 | 190 | def split_rgb_field(cloud_arr): 191 | """Takes an array with a named 'rgb' float32 field, and returns an array in which 192 | this has been split into 3 uint 8 fields: 'r', 'g', and 'b'. 193 | 194 | (pcl stores rgb in packed 32 bit floats) 195 | """ 196 | rgb_arr = cloud_arr['rgb'].copy() 197 | rgb_arr.dtype = np.uint32 198 | r = np.asarray((rgb_arr >> 16) & 255, dtype=np.uint8) 199 | g = np.asarray((rgb_arr >> 8) & 255, dtype=np.uint8) 200 | b = np.asarray(rgb_arr & 255, dtype=np.uint8) 201 | 202 | # create a new array, without rgb, but with r, g, and b fields 203 | new_dtype = [] 204 | for field_name in cloud_arr.dtype.names: 205 | field_type, field_offset = cloud_arr.dtype.fields[field_name] 206 | if not field_name == 'rgb': 207 | new_dtype.append((field_name, field_type)) 208 | new_dtype.append(('r', np.uint8)) 209 | new_dtype.append(('g', np.uint8)) 210 | new_dtype.append(('b', np.uint8)) 211 | new_cloud_arr = np.zeros(cloud_arr.shape, new_dtype) 212 | 213 | # fill in the new array 214 | for field_name in new_cloud_arr.dtype.names: 215 | if field_name == 'r': 216 | new_cloud_arr[field_name] = r 217 | elif field_name == 'g': 218 | new_cloud_arr[field_name] = g 219 | elif field_name == 'b': 220 | new_cloud_arr[field_name] = b 221 | else: 222 | new_cloud_arr[field_name] = cloud_arr[field_name] 223 | return new_cloud_arr 224 | 225 | 226 | def get_xyz_points(cloud_array, remove_nans=True, dtype=np.float): 227 | """Pulls out x, y, and z columns from the cloud recordarray, and returns 228 | a 3xN matrix. 229 | """ 230 | # remove crap points 231 | if remove_nans: 232 | mask = np.isfinite(cloud_array['x']) & np.isfinite(cloud_array['y']) & np.isfinite(cloud_array['z']) 233 | cloud_array = cloud_array[mask] 234 | 235 | # pull out x, y, and z values 236 | points = np.zeros(list(cloud_array.shape) + [3], dtype=dtype) 237 | points[..., 0] = cloud_array['x'] 238 | points[..., 1] = cloud_array['y'] 239 | points[..., 2] = cloud_array['z'] 240 | 241 | return points 242 | 243 | 244 | def pointcloud2_to_xyz_array(cloud_msg, remove_nans=True): 245 | return get_xyz_points(pointcloud2_to_array(cloud_msg), remove_nans=remove_nans) 246 | -------------------------------------------------------------------------------- /apps/read_file_sdf.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python 2 | # -*- coding: utf-8 -*- 3 | # Author : Hongzhuo Liang 4 | # E-mail : liang@informatik.uni-hamburg.de 5 | # Description: 6 | # Date : 22/05/2018 12:37 AM 7 | # File Name : read_file_sdf.py 8 | import os 9 | import multiprocessing 10 | import subprocess 11 | 12 | """ 13 | This file convert obj file to sdf file automatically and multiprocessingly. 14 | All the cores of a computer can do the job parallel. 15 | """ 16 | 17 | 18 | def get_file_name(file_dir_): 19 | file_list = [] 20 | for root, dirs, files in os.walk(file_dir_): 21 | # print(root) # current path 22 | if root.count('/') == file_dir_.count('/')+1: 23 | file_list.append(root) 24 | # print(dirs) # all the directories in current path 25 | # print(files) # all the files in current path 26 | file_list.sort() 27 | return file_list 28 | 29 | 30 | def generate_sdf(path_to_sdfgen, obj_filename, dim, padding): 31 | """ Converts mesh to an sdf object """ 32 | 33 | # create the SDF using binary tools 34 | sdfgen_cmd = '%s \"%s\" %d %d' % (path_to_sdfgen, obj_filename, dim, padding) 35 | os.system(sdfgen_cmd) 36 | # print('SDF Command: %s' % sdfgen_cmd) 37 | return 38 | 39 | 40 | def do_job_convert_obj_to_sdf(x): 41 | # file_list_all = get_file_name(file_dir) 42 | if os.path.exists(str(file_list_all[x])+"/google_512k/nontextured.sdf"): 43 | print(str(file_list_all[x])+"/google_512k/nontextured.sdf","pass") 44 | else: 45 | generate_sdf(path_sdfgen, str(file_list_all[x])+"/google_512k/nontextured.obj", 100, 5) # for google scanner 46 | print("Done job number", x) 47 | 48 | 49 | def generate_obj_from_ply(file_name_): 50 | base = file_name_.split(".")[0] 51 | p = subprocess.Popen(["pcl_ply2obj", base + ".ply", base + ".obj"]) 52 | p.wait() 53 | 54 | 55 | if __name__ == '__main__': 56 | home_dir = os.environ['HOME'] 57 | #设置obj ply文件的路径 58 | file_dir = home_dir + "/dataset/PointNetGPD/ycb_meshes_google/objects" # for google ycb 59 | # file_dir = home_dir + "/dataset/ycb_meshes" # for low quality ycb 60 | path_sdfgen = home_dir + "/code/SDFGen/bin/SDFGen" 61 | file_list_all = get_file_name(file_dir) 62 | object_numbers = file_list_all.__len__() 63 | 64 | # generate obj from ply file 65 | for i in file_list_all: 66 | if os.path.exists(i+"/google_512k/nontextured.ply"): 67 | print(i+"/google_512k/nontextured.ply","pass") 68 | else: 69 | generate_obj_from_ply(i+"/google_512k/nontextured.ply") 70 | print("finish", i) 71 | # The operation for the multi core 72 | cores = multiprocessing.cpu_count() 73 | pool = multiprocessing.Pool(processes=cores) 74 | pool.map(do_job_convert_obj_to_sdf, range(object_numbers)) 75 | -------------------------------------------------------------------------------- /apps/read_grasps_from_file.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python 2 | # -*- coding: utf-8 -*- 3 | # Author : Hongzhuo Liang 4 | # E-mail : liang@informatik.uni-hamburg.de 5 | # Description: 6 | # Date : 30/05/2018 9:57 AM 7 | # File Name : read_grasps_from_file.py 8 | # 使用python3 9 | import os 10 | import sys 11 | import re 12 | from meshpy.obj_file import ObjFile 13 | from meshpy.sdf_file import SdfFile 14 | from dexnet.grasping import GraspableObject3D 15 | import numpy as np 16 | from dexnet.visualization.visualizer3d import DexNetVisualizer3D as Vis 17 | from dexnet.grasping import RobotGripper 18 | from autolab_core import YamlConfig 19 | from mayavi import mlab 20 | from dexnet.grasping import GpgGraspSampler # temporary way for show 3D gripper using mayavi 21 | import glob 22 | 23 | # global configurations: 24 | home_dir = os.environ["HOME"] + "/code/PointNetGPD" 25 | yaml_config = YamlConfig(home_dir + "/dex-net/test/config.yaml") 26 | gripper_name = "robotiq_85" 27 | gripper = RobotGripper.load(gripper_name, home_dir + "/dex-net/data/grippers") 28 | ags = GpgGraspSampler(gripper, yaml_config) 29 | 30 | save_fig = False # save fig as png file 31 | show_fig = True # show the mayavi figure 32 | generate_new_file = False # whether generate new file for collision free grasps 33 | check_pcd_grasp_points = False 34 | 35 | 36 | def open_npy_and_obj(name_to_open_): 37 | """从路径指定的npy文件中读取并返回 38 | npy_m_: 抓取序列（只是针对某单个物体的抓取序列） 39 | obj_: 待抓取物体的模型mesh（单个模型） 40 | ply_name_: 待抓取物体的路径 41 | object_name_: 待抓取物体的名称 42 | """ 43 | npy_m_ = np.load(name_to_open_) 44 | file_dir = home_dir + "/PointNetGPD/data/ycb_meshes_google/" 45 | object_name_ = name_to_open_.split("/")[-1][:-4] 46 | #object_name_=object_name_.split("_")[-1][:-3] 47 | print(object_name_) 48 | ply_name_ = file_dir + object_name_ + "/google_512k/nontextured.ply" 49 | if not check_pcd_grasp_points: 50 | of = ObjFile(file_dir + object_name_ + "/google_512k/nontextured.obj") 51 | sf = SdfFile(file_dir + object_name_ + "/google_512k/nontextured.sdf") 52 | mesh = of.read() 53 | sdf = sf.read() 54 | obj_ = GraspableObject3D(sdf, mesh) 55 | else: 56 | cloud_path = home_dir + "/dataset/ycb_rgbd/" + object_name_ + "/clouds/" 57 | pcd_files = glob.glob(cloud_path + "*.pcd") 58 | obj_ = pcd_files 59 | obj_.sort() 60 | return npy_m_, obj_, ply_name_, object_name_ 61 | 62 | 63 | def display_object(obj_): 64 | """display object only using mayavi""" 65 | Vis.figure(bgcolor=(1, 1, 1), size=(1000, 1000)) 66 | Vis.mesh(obj_.mesh.trimesh, color=(0.5, 0.5, 0.5), style="surface") 67 | Vis.show() 68 | 69 | 70 | def display_gripper_on_object(obj_, grasp_): 71 | """display both object and gripper using mayavi""" 72 | # transfer wrong was fixed by the previews comment of meshpy modification. 73 | # gripper_name = "robotiq_85" 74 | # gripper = RobotGripper.load(gripper_name, home_dir + "/dex-net/data/grippers") 75 | # stable_pose = self.dataset.stable_pose(object.key, "pose_1") 76 | # T_obj_world = RigidTransform(from_frame="obj", to_frame="world") 77 | t_obj_gripper = grasp_.gripper_pose(gripper) 78 | 79 | stable_pose = t_obj_gripper 80 | grasp_ = grasp_.perpendicular_table(stable_pose) 81 | 82 | Vis.figure(bgcolor=(1, 1, 1), size=(1000, 1000)) 83 | Vis.gripper_on_object(gripper, grasp_, obj_, 84 | gripper_color=(0.25, 0.25, 0.25), 85 | # stable_pose=stable_pose, # .T_obj_world, 86 | plot_table=False) 87 | Vis.show() 88 | 89 | 90 | def display_grasps(grasp, graspable, color): 91 | """ 92 | 显示给定的抓取，做一下碰撞检测，和指定颜色 93 | grasp：单个抓取 94 | graspable：目标物体（仅做碰撞） 95 | color：夹爪的颜色 96 | """ 97 | center_point = grasp[0:3] #夹爪中心(指尖中心) 98 | major_pc = grasp[3:6] # binormal 99 | width = grasp[6] 100 | angle = grasp[7] 101 | level_score, refine_score = grasp[-2:] 102 | # cal approach 103 | cos_t = np.cos(angle) 104 | sin_t = np.sin(angle) 105 | R1 = np.c_[[cos_t, 0, sin_t],[0, 1, 0],[-sin_t, 0, cos_t]] 106 | 107 | axis_y = major_pc 108 | axis_x = np.array([axis_y[1], -axis_y[0], 0]) 109 | if np.linalg.norm(axis_x) == 0: 110 | axis_x = np.array([1, 0, 0]) 111 | 112 | axis_x = axis_x / np.linalg.norm(axis_x) 113 | axis_y = axis_y / np.linalg.norm(axis_y) 114 | 115 | axis_z = np.cross(axis_x, axis_y) 116 | 117 | R2 = np.c_[axis_x, np.c_[axis_y, axis_z]] 118 | approach_normal = R2.dot(R1)[:, 0] 119 | approach_normal = approach_normal / np.linalg.norm(approach_normal) 120 | minor_pc = np.cross(major_pc, approach_normal) 121 | 122 | #计算夹爪bottom_center 123 | grasp_bottom_center = -ags.gripper.hand_depth * approach_normal + center_point 124 | #以世界坐标系为参考系时 125 | hand_points = ags.get_hand_points(grasp_bottom_center, approach_normal, major_pc) 126 | #固定夹爪作为参考系时 127 | local_hand_points = ags.get_hand_points(np.array([0, 0, 0]), np.array([1, 0, 0]), np.array([0, 1, 0])) 128 | #检测一下夹爪的碰撞（这里使用原始函数，不检测与桌面的碰撞） 129 | if_collide = ags.check_collide(grasp_bottom_center, approach_normal, 130 | major_pc, minor_pc, graspable, local_hand_points) 131 | 132 | #只显示夹爪，不显示点云 133 | if not if_collide and (show_fig or save_fig): 134 | ags.show_grasp_3d(hand_points, color=color) 135 | return True 136 | elif not if_collide: 137 | return True 138 | else: 139 | return False 140 | 141 | 142 | def show_selected_grasps_with_color(m, ply_name_, title, obj_): 143 | """显示选择出的抓取，将抓取序列中good的抓取和bad的抓取区分开来，并分别显示在两个窗口中 144 | m： 针对被抓物体的带有打分的grasp序列 145 | ply_name_ ： 被抓物体的路径 146 | title：被抓物体的名称 147 | obj_：被抓物体的mesh文件 148 | """ 149 | m_good = m[m[:, -2] <= 0.4] 150 | m_good = m_good[np.random.choice(len(m_good), size=25, replace=True)] 151 | 152 | m_bad = m[m[:, -2] >= 1.8] 153 | m_bad = m_bad[np.random.choice(len(m_bad), size=25, replace=True)] 154 | 155 | collision_grasp_num = 0 156 | if save_fig or show_fig: 157 | # fig 1: good grasps 158 | mlab.figure(bgcolor=(1, 1, 1), fgcolor=(0.7, 0.7, 0.7), size=(1000, 1000)) 159 | mlab.pipeline.surface(mlab.pipeline.open(ply_name_)) 160 | for a in m_good: 161 | # display_gripper_on_object(obj, a) # real gripper 162 | collision_free = display_grasps(a, obj_, color="d") # simulated gripper 163 | if not collision_free: 164 | collision_grasp_num += 1 165 | 166 | if save_fig: 167 | mlab.savefig("good_"+title+".png") 168 | mlab.close() 169 | elif show_fig: 170 | mlab.title(title, size=0.5) 171 | 172 | # fig 2: bad grasps 173 | mlab.figure(bgcolor=(1, 1, 1), fgcolor=(0.7, 0.7, 0.7), size=(1000, 1000)) 174 | mlab.pipeline.surface(mlab.pipeline.open(ply_name_)) 175 | 176 | for a in m_bad: 177 | # display_gripper_on_object(obj, a) # real gripper 178 | collision_free = display_grasps(a, obj_, color=(1, 0, 0)) 179 | if not collision_free: 180 | collision_grasp_num += 1 181 | 182 | if save_fig: 183 | mlab.savefig("bad_"+title+".png") 184 | mlab.close() 185 | elif show_fig: 186 | mlab.title(title, size=0.5) 187 | mlab.show() 188 | elif generate_new_file: 189 | # only to calculate collision: 190 | collision_grasp_num = 0 191 | ind_good_grasp_ = [] 192 | for i_ in range(len(m)): 193 | collision_free = display_grasps(m[i_][0], obj_, color=(1, 0, 0)) 194 | if not collision_free: 195 | collision_grasp_num += 1 196 | else: 197 | ind_good_grasp_.append(i_) 198 | collision_grasp_num = str(collision_grasp_num) 199 | collision_grasp_num = (4-len(collision_grasp_num))*" " + collision_grasp_num 200 | print("collision_grasp_num =", collision_grasp_num, "| object name:", title) 201 | return ind_good_grasp_ 202 | 203 | 204 | def get_grasp_points_num(m, obj_): 205 | """获取夹爪内部点云的数量 206 | """ 207 | has_points_ = [] 208 | ind_points_ = [] 209 | for i_ in range(len(m)): 210 | grasps = m[i_][0] 211 | approach_normal = grasps.rotated_full_axis[:, 0] 212 | approach_normal = approach_normal / np.linalg.norm(approach_normal) 213 | major_pc = grasps.configuration[3:6] 214 | major_pc = major_pc / np.linalg.norm(major_pc) 215 | minor_pc = np.cross(approach_normal, major_pc) 216 | center_point = grasps.center 217 | 218 | #计算bottom_center 219 | grasp_bottom_center = -ags.gripper.hand_depth * approach_normal + center_point 220 | 221 | # hand_points = ags.get_hand_points(grasp_bottom_center, approach_normal, major_pc) 222 | local_hand_points = ags.get_hand_points(np.array([0, 0, 0]), np.array([1, 0, 0]), np.array([0, 1, 0])) 223 | 224 | has_points_tmp, ind_points_tmp = ags.check_collision_square(grasp_bottom_center, approach_normal, 225 | major_pc, minor_pc, obj_, local_hand_points, 226 | "p_open") 227 | ind_points_tmp = len(ind_points_tmp) # here we only want to know the number of in grasp points. 228 | has_points_.append(has_points_tmp) 229 | ind_points_.append(ind_points_tmp) 230 | return has_points_, ind_points_ 231 | 232 | 233 | if __name__ == "__main__": 234 | #获取npy的名称 235 | npy_names = glob.glob(home_dir + "/PointNetGPD/data/ycb_grasp/train/*.npy") 236 | npy_names.sort() 237 | for i in range(len(npy_names)): 238 | #把夹爪姿态和打分，先从npy中读取出来 239 | grasps_with_score, obj, ply_name, obj_name = open_npy_and_obj(npy_names[i]) 240 | print("load file {}".format(npy_names[i])) 241 | #显示抓取 242 | ind_good_grasp = show_selected_grasps_with_color(grasps_with_score, ply_name, obj_name, obj) 243 | 244 | -------------------------------------------------------------------------------- /apps/show_pcd.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python 2 | # -*- coding: utf-8 -*- 3 | # Author : Hongzhuo Liang 4 | # E-mail : liang@informatik.uni-hamburg.de 5 | # Description: 6 | # Date : 21/08/2018 9:08 PM 7 | # File Name : show_pcd.py 8 | 9 | import numpy as np 10 | from mayavi import mlab 11 | import glob 12 | import os 13 | import pcl 14 | 15 | 16 | #需要提前安装pyside 17 | def show_obj(surface_points_, title, color='b'): 18 | mlab.figure(bgcolor=(0, 0, 0), fgcolor=(0.7, 0.7, 0.7), size=(1000, 1000)) 19 | if color == 'b': 20 | color_f = (0, 0, 1) 21 | elif color == 'r': 22 | color_f = (1, 0, 0) 23 | elif color == 'g': 24 | color_f = (0, 1, 0) 25 | else: 26 | color_f = (1, 1, 1) 27 | points_ = surface_points_ 28 | mlab.points3d(points_[:, 0], points_[:, 1], points_[:, 2], color=color_f, scale_factor=.0007) 29 | mlab.title(title, size=0.5) 30 | 31 | 32 | path = os.environ['HOME'] + "/win10/data/ycb_rgbd/003_cracker_box/clouds" 33 | pcd_files = glob.glob(os.path.join(path, '*.pcd')) 34 | pcd_files.sort() 35 | for i in range(450, len(pcd_files)): 36 | pcd = pcd_files[i] 37 | file_name = pcd.split('/')[-1] 38 | print(file_name) 39 | points_c = pcl.load(pcd) 40 | points = points_c.to_array() 41 | show_obj(points, file_name) 42 | mlab.show() 43 | -------------------------------------------------------------------------------- /apps/voxelgrid.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python 2 | # -*- coding: utf-8 -*- 3 | # Author : Hongzhuo Liang 4 | # E-mail : liang@informatik.uni-hamburg.de 5 | # Description: 6 | # Date : 05/08/2018 6:04 PM 7 | # File Name : kinect2grasp_python2.py 8 | # Note: this file is inspired by PyntCloud 9 | # Reference web: https://github.com/daavoo/pyntcloud 10 | 11 | import numpy as np 12 | from scipy.spatial import cKDTree 13 | from numba import jit 14 | 15 | is_numba_avaliable = True 16 | 17 | 18 | @jit 19 | def groupby_count(xyz, indices, out): 20 | for i in range(xyz.shape[0]): 21 | out[indices[i]] += 1 22 | return out 23 | 24 | 25 | @jit 26 | def groupby_sum(xyz, indices, N, out): 27 | for i in range(xyz.shape[0]): 28 | out[indices[i]] += xyz[i][N] 29 | return out 30 | 31 | 32 | @jit 33 | def groupby_max(xyz, indices, N, out): 34 | for i in range(xyz.shape[0]): 35 | if xyz[i][N] > out[indices[i]]: 36 | out[indices[i]] = xyz[i][N] 37 | return out 38 | 39 | 40 | def cartesian(arrays, out=None): 41 | """Generate a cartesian product of input arrays. 42 | 43 | Parameters 44 | ---------- 45 | arrays : list of array-like 46 | 1-D arrays to form the cartesian product of. 47 | out : ndarray 48 | Array to place the cartesian product in. 49 | 50 | Returns 51 | ------- 52 | out : ndarray 53 | 2-D array of shape (M, len(arrays)) containing cartesian products 54 | formed of input arrays. 55 | 56 | Examples 57 | -------- 58 | >>> cartesian(([1, 2, 3], [4, 5], [6, 7])) 59 | array([[1, 4, 6], 60 | [1, 4, 7], 61 | [1, 5, 6], 62 | [1, 5, 7], 63 | [2, 4, 6], 64 | [2, 4, 7], 65 | [2, 5, 6], 66 | [2, 5, 7], 67 | [3, 4, 6], 68 | [3, 4, 7], 69 | [3, 5, 6], 70 | [3, 5, 7]]) 71 | 72 | """ 73 | arrays = [np.asarray(x) for x in arrays] 74 | shape = (len(x) for x in arrays) 75 | dtype = arrays[0].dtype 76 | 77 | ix = np.indices(shape) 78 | ix = ix.reshape(len(arrays), -1).T 79 | 80 | if out is None: 81 | out = np.empty_like(ix, dtype=dtype) 82 | 83 | for n, arr in enumerate(arrays): 84 | out[:, n] = arrays[n][ix[:, n]] 85 | 86 | return out 87 | 88 | 89 | class VoxelGrid: 90 | 91 | def __init__(self, points, n_x=1, n_y=1, n_z=1, size_x=None, size_y=None, size_z=None, regular_bounding_box=True): 92 | """Grid of voxels with support for different build methods. 93 | 94 | Parameters 95 | ---------- 96 | points: (N, 3) numpy.array 97 | n_x, n_y, n_z : int, optional 98 | Default: 1 99 | The number of segments in which each axis will be divided. 100 | Ignored if corresponding size_x, size_y or size_z is not None. 101 | size_x, size_y, size_z : float, optional 102 | Default: None 103 | The desired voxel size along each axis. 104 | If not None, the corresponding n_x, n_y or n_z will be ignored. 105 | regular_bounding_box : bool, optional 106 | Default: True 107 | If True, the bounding box of the point cloud will be adjusted 108 | in order to have all the dimensions of equal length. 109 | """ 110 | self._points = points 111 | self.x_y_z = [n_x, n_y, n_z] 112 | self.sizes = [size_x, size_y, size_z] 113 | self.regular_bounding_box = regular_bounding_box 114 | 115 | def compute(self): 116 | xyzmin = self._points.min(0) 117 | xyzmax = self._points.max(0) 118 | 119 | if self.regular_bounding_box: 120 | #: adjust to obtain a minimum bounding box with all sides of equal length 121 | margin = max(xyzmax - xyzmin) - (xyzmax - xyzmin) 122 | xyzmin = xyzmin - margin / 2 123 | xyzmax = xyzmax + margin / 2 124 | 125 | for n, size in enumerate(self.sizes): 126 | if size is None: 127 | continue 128 | margin = (((self._points.ptp(0)[n] // size) + 1) * size) - self._points.ptp(0)[n] 129 | xyzmin[n] -= margin / 2 130 | xyzmax[n] += margin / 2 131 | self.x_y_z[n] = ((xyzmax[n] - xyzmin[n]) / size).astype(int) 132 | 133 | self.xyzmin = xyzmin 134 | self.xyzmax = xyzmax 135 | 136 | segments = [] 137 | shape = [] 138 | for i in range(3): 139 | # note the +1 in num 140 | s, step = np.linspace(xyzmin[i], xyzmax[i], num=(self.x_y_z[i] + 1), retstep=True) 141 | segments.append(s) 142 | shape.append(step) 143 | 144 | self.segments = segments 145 | self.shape = shape 146 | 147 | self.n_voxels = self.x_y_z[0] * self.x_y_z[1] * self.x_y_z[2] 148 | 149 | self.id = "V({},{},{})".format(self.x_y_z, self.sizes, self.regular_bounding_box) 150 | 151 | # find where each point lies in corresponding segmented axis 152 | # -1 so index are 0-based; clip for edge cases 153 | self.voxel_x = np.clip(np.searchsorted(self.segments[0], self._points[:, 0]) - 1, 0, self.x_y_z[0]) 154 | self.voxel_y = np.clip(np.searchsorted(self.segments[1], self._points[:, 1]) - 1, 0, self.x_y_z[1]) 155 | self.voxel_z = np.clip(np.searchsorted(self.segments[2], self._points[:, 2]) - 1, 0, self.x_y_z[2]) 156 | self.voxel_n = np.ravel_multi_index([self.voxel_x, self.voxel_y, self.voxel_z], self.x_y_z) 157 | 158 | # compute center of each voxel 159 | midsegments = [(self.segments[i][1:] + self.segments[i][:-1]) / 2 for i in range(3)] 160 | self.voxel_centers = cartesian(midsegments).astype(np.float32) 161 | 162 | def query(self, points): 163 | """ABC API. Query structure. 164 | 165 | TODO Make query_voxelgrid an independent function, and add a light 166 | save mode where only segments and x_y_z are saved. 167 | """ 168 | voxel_x = np.clip(np.searchsorted( 169 | self.segments[0], points[:, 0]) - 1, 0, self.x_y_z[0]) 170 | voxel_y = np.clip(np.searchsorted( 171 | self.segments[1], points[:, 1]) - 1, 0, self.x_y_z[1]) 172 | voxel_z = np.clip(np.searchsorted( 173 | self.segments[2], points[:, 2]) - 1, 0, self.x_y_z[2]) 174 | voxel_n = np.ravel_multi_index([voxel_x, voxel_y, voxel_z], self.x_y_z) 175 | 176 | return voxel_n 177 | 178 | def get_feature_vector(self, mode="binary"): 179 | """Return a vector of size self.n_voxels. See mode options below. 180 | 181 | Parameters 182 | ---------- 183 | mode: str in available modes. See Notes 184 | Default "binary" 185 | 186 | Returns 187 | ------- 188 | feature_vector: [n_x, n_y, n_z] ndarray 189 | See Notes. 190 | 191 | Notes 192 | ----- 193 | Available modes are: 194 | 195 | binary 196 | 0 for empty voxels, 1 for occupied. 197 | 198 | density 199 | number of points inside voxel / total number of points. 200 | 201 | TDF 202 | Truncated Distance Function. Value between 0 and 1 indicating the distance 203 | between the voxel's center and the closest point. 1 on the surface, 204 | 0 on voxels further than 2 * voxel side. 205 | 206 | x_max, y_max, z_max 207 | Maximum coordinate value of points inside each voxel. 208 | 209 | x_mean, y_mean, z_mean 210 | Mean coordinate value of points inside each voxel. 211 | """ 212 | vector = np.zeros(self.n_voxels) 213 | 214 | if mode == "binary": 215 | vector[np.unique(self.voxel_n)] = 1 216 | 217 | elif mode == "density": 218 | count = np.bincount(self.voxel_n) 219 | vector[:len(count)] = count 220 | vector /= len(self.voxel_n) 221 | 222 | elif mode == "TDF": 223 | # truncation = np.linalg.norm(self.shape) 224 | kdt = cKDTree(self._points) 225 | vector, i = kdt.query(self.voxel_centers, n_jobs=-1) 226 | 227 | elif mode.endswith("_max"): 228 | if not is_numba_avaliable: 229 | raise ImportError("numba is required to compute {}".format(mode)) 230 | axis = {"x_max": 0, "y_max": 1, "z_max": 2} 231 | vector = groupby_max(self._points, self.voxel_n, axis[mode], vector) 232 | 233 | elif mode.endswith("_mean"): 234 | if not is_numba_avaliable: 235 | raise ImportError("numba is required to compute {}".format(mode)) 236 | axis = {"x_mean": 0, "y_mean": 1, "z_mean": 2} 237 | voxel_sum = groupby_sum(self._points, self.voxel_n, axis[mode], np.zeros(self.n_voxels)) 238 | voxel_count = groupby_count(self._points, self.voxel_n, np.zeros(self.n_voxels)) 239 | vector = np.nan_to_num(voxel_sum / voxel_count) 240 | 241 | else: 242 | raise NotImplementedError("{} is not a supported feature vector mode".format(mode)) 243 | 244 | return vector.reshape(self.x_y_z) 245 | 246 | def get_voxel_neighbors(self, voxel): 247 | """Get valid, non-empty 26 neighbors of voxel. 248 | 249 | Parameters 250 | ---------- 251 | voxel: int in self.set_voxel_n 252 | 253 | Returns 254 | ------- 255 | neighbors: list of int 256 | Indices of the valid, non-empty 26 neighborhood around voxel. 257 | """ 258 | 259 | x, y, z = np.unravel_index(voxel, self.x_y_z) 260 | 261 | valid_x = [] 262 | valid_y = [] 263 | valid_z = [] 264 | if x - 1 >= 0: 265 | valid_x.append(x - 1) 266 | if y - 1 >= 0: 267 | valid_y.append(y - 1) 268 | if z - 1 >= 0: 269 | valid_z.append(z - 1) 270 | 271 | valid_x.append(x) 272 | valid_y.append(y) 273 | valid_z.append(z) 274 | 275 | if x + 1 < self.x_y_z[0]: 276 | valid_x.append(x + 1) 277 | if y + 1 < self.x_y_z[1]: 278 | valid_y.append(y + 1) 279 | if z + 1 < self.x_y_z[2]: 280 | valid_z.append(z + 1) 281 | 282 | valid_neighbor_indices = cartesian((valid_x, valid_y, valid_z)) 283 | 284 | ravel_indices = np.ravel_multi_index((valid_neighbor_indices[:, 0], 285 | valid_neighbor_indices[:, 1], 286 | valid_neighbor_indices[:, 2]), self.x_y_z) 287 | 288 | return [x for x in ravel_indices if x in np.unique(self.voxel_n)] 289 | 290 | -------------------------------------------------------------------------------- /apps/ycb_cloud_generate.py: -------------------------------------------------------------------------------- 1 | import os 2 | import numpy as np 3 | import h5py as h5 4 | from scipy.misc import imread 5 | from struct import pack, unpack 6 | import math 7 | import multiprocessing as mp 8 | import glob 9 | 10 | 11 | #将rgbd图像转换为点云图，并转换到目标物体坐标系系统 12 | # extract pointcloud from rgb-d, then convert it into obj coordinate system(tsdf/poisson reconstructed object) 13 | 14 | def im2col(im, psize): 15 | n_channels = 1 if len(im.shape) == 2 else im.shape[0] 16 | (n_channels, rows, cols) = (1,) * (3 - len(im.shape)) + im.shape 17 | 18 | im_pad = np.zeros((n_channels, 19 | int(math.ceil(1.0 * rows / psize) * psize), 20 | int(math.ceil(1.0 * cols / psize) * psize))) 21 | im_pad[:, 0:rows, 0:cols] = im 22 | 23 | final = np.zeros((im_pad.shape[1], im_pad.shape[2], n_channels, 24 | psize, psize)) 25 | for c in range(n_channels): 26 | for x in range(psize): 27 | for y in range(psize): 28 | im_shift = np.vstack( 29 | (im_pad[c, x:], im_pad[c, :x])) 30 | im_shift = np.column_stack( 31 | (im_shift[:, y:], im_shift[:, :y])) 32 | final[x::psize, y::psize, c] = np.swapaxes( 33 | im_shift.reshape(int(im_pad.shape[1] / psize), psize, 34 | int(im_pad.shape[2] / psize), psize), 1, 2) 35 | 36 | return np.squeeze(final[0:rows - psize + 1, 0:cols - psize + 1]) 37 | 38 | #滤波器，率除掉断续的图像？ 39 | def filterDiscontinuities(depthMap): 40 | filt_size = 7 41 | thresh = 1000 42 | 43 | # Ensure that filter sizes are okay 44 | assert filt_size % 2 == 1, "Can only use odd filter sizes." 45 | 46 | # Compute discontinuities 47 | offset = int((filt_size - 1) / 2) 48 | patches = 1.0 * im2col(depthMap, filt_size) 49 | mids = patches[:, :, offset, offset] 50 | mins = np.min(patches, axis=(2, 3)) 51 | maxes = np.max(patches, axis=(2, 3)) 52 | 53 | discont = np.maximum(np.abs(mins - mids), 54 | np.abs(maxes - mids)) 55 | mark = discont > thresh 56 | 57 | # Account for offsets 58 | final_mark = np.zeros((480, 640), dtype=np.uint16) 59 | final_mark[offset:offset + mark.shape[0], 60 | offset:offset + mark.shape[1]] = mark 61 | 62 | return depthMap * (1 - final_mark) 63 | 64 | #把深度图和rgb图像相互配准 65 | def registerDepthMap(unregisteredDepthMap, 66 | rgbImage, 67 | depthK, 68 | rgbK, 69 | H_RGBFromDepth): 70 | 71 | 72 | unregisteredHeight = unregisteredDepthMap.shape[0] 73 | unregisteredWidth = unregisteredDepthMap.shape[1] 74 | 75 | registeredHeight = rgbImage.shape[0] 76 | registeredWidth = rgbImage.shape[1] 77 | 78 | registeredDepthMap = np.zeros((registeredHeight, registeredWidth)) 79 | 80 | xyzDepth = np.empty((4,1)) 81 | xyzRGB = np.empty((4,1)) 82 | 83 | # Ensure that the last value is 1 (homogeneous coordinates) 84 | xyzDepth[3] = 1 85 | 86 | invDepthFx = 1.0 / depthK[0,0] 87 | invDepthFy = 1.0 / depthK[1,1] 88 | depthCx = depthK[0,2] 89 | depthCy = depthK[1,2] 90 | 91 | rgbFx = rgbK[0,0] 92 | rgbFy = rgbK[1,1] 93 | rgbCx = rgbK[0,2] 94 | rgbCy = rgbK[1,2] 95 | 96 | undistorted = np.empty(2) 97 | for v in range(unregisteredHeight): 98 | for u in range(unregisteredWidth): 99 | 100 | depth = unregisteredDepthMap[v,u] 101 | if depth == 0: 102 | continue 103 | 104 | xyzDepth[0] = ((u - depthCx) * depth) * invDepthFx 105 | xyzDepth[1] = ((v - depthCy) * depth) * invDepthFy 106 | xyzDepth[2] = depth 107 | 108 | xyzRGB[0] = (H_RGBFromDepth[0,0] * xyzDepth[0] + 109 | H_RGBFromDepth[0,1] * xyzDepth[1] + 110 | H_RGBFromDepth[0,2] * xyzDepth[2] + 111 | H_RGBFromDepth[0,3]) 112 | xyzRGB[1] = (H_RGBFromDepth[1,0] * xyzDepth[0] + 113 | H_RGBFromDepth[1,1] * xyzDepth[1] + 114 | H_RGBFromDepth[1,2] * xyzDepth[2] + 115 | H_RGBFromDepth[1,3]) 116 | xyzRGB[2] = (H_RGBFromDepth[2,0] * xyzDepth[0] + 117 | H_RGBFromDepth[2,1] * xyzDepth[1] + 118 | H_RGBFromDepth[2,2] * xyzDepth[2] + 119 | H_RGBFromDepth[2,3]) 120 | 121 | invRGB_Z = 1.0 / xyzRGB[2] 122 | undistorted[0] = (rgbFx * xyzRGB[0]) * invRGB_Z + rgbCx 123 | undistorted[1] = (rgbFy * xyzRGB[1]) * invRGB_Z + rgbCy 124 | 125 | uRGB = int(undistorted[0] + 0.5) 126 | vRGB = int(undistorted[1] + 0.5) 127 | 128 | if (uRGB < 0 or uRGB >= registeredWidth) or (vRGB < 0 or vRGB >= registeredHeight): 129 | continue 130 | 131 | registeredDepth = xyzRGB[2] 132 | if registeredDepth > registeredDepthMap[vRGB,uRGB]: 133 | registeredDepthMap[vRGB,uRGB] = registeredDepth 134 | 135 | return registeredDepthMap 136 | 137 | 138 | #把rgb配准后的深度图转换为点云 139 | def registeredDepthMapToPointCloud(depthMap, rgbImage, rgbK, refFromRGB, objFromref, organized=False): 140 | rgbCx = rgbK[0,2] 141 | rgbCy = rgbK[1,2] 142 | invRGBFx = 1.0/rgbK[0,0] 143 | invRGBFy = 1.0/rgbK[1,1] 144 | 145 | height = depthMap.shape[0] 146 | width = depthMap.shape[1] 147 | 148 | if organized: 149 | cloud = np.empty((height, width, 6), dtype=np.float) 150 | else: 151 | cloud = np.empty((1, height*width, 6), dtype=np.float) 152 | 153 | goodPointsCount = 0 154 | for v in range(height): 155 | for u in range(width): 156 | 157 | depth = depthMap[v,u] 158 | 159 | if organized: 160 | row = v 161 | col = u 162 | else: 163 | row = 0 164 | col = goodPointsCount 165 | 166 | if depth <= 0: 167 | if organized: 168 | if depth <= 0: 169 | cloud[row,col,0] = float('nan') 170 | cloud[row,col,1] = float('nan') 171 | cloud[row,col,2] = float('nan') 172 | cloud[row,col,3] = 0 173 | cloud[row,col,4] = 0 174 | cloud[row,col,5] = 0 175 | continue 176 | 177 | x = (u - rgbCx) * depth * invRGBFx 178 | y = (v - rgbCy) * depth * invRGBFy 179 | z = depth 180 | 181 | #refFromRGB 182 | x1 = (refFromRGB[0,0] * x + 183 | refFromRGB[0,1] * y + 184 | refFromRGB[0,2] * z + 185 | refFromRGB[0,3]) 186 | y1 = (refFromRGB[1,0] * x + 187 | refFromRGB[1,1] * y + 188 | refFromRGB[1,2] * z + 189 | refFromRGB[1,3]) 190 | z1 = (refFromRGB[2,0] * x + 191 | refFromRGB[2,1] * y + 192 | refFromRGB[2,2] * z + 193 | refFromRGB[2,3]) 194 | 195 | x, y, z = x1, y1, z1 196 | 197 | # obj from ref 198 | cloud[row,col,0] = (objFromref[0,0] * x + 199 | objFromref[0,1] * y + 200 | objFromref[0,2] * z + 201 | objFromref[0,3]) 202 | cloud[row,col,1] = (objFromref[1,0] * x + 203 | objFromref[1,1] * y + 204 | objFromref[1,2] * z + 205 | objFromref[1,3]) 206 | cloud[row,col,2] = (objFromref[2,0] * x + 207 | objFromref[2,1] * y + 208 | objFromref[2,2] * z + 209 | objFromref[2,3]) 210 | 211 | cloud[row,col,3] = rgbImage[v,u,0] 212 | cloud[row,col,4] = rgbImage[v,u,1] 213 | cloud[row,col,5] = rgbImage[v,u,2] 214 | if not organized: 215 | goodPointsCount += 1 216 | 217 | if not organized: 218 | cloud = cloud[:,:goodPointsCount,:] 219 | 220 | return cloud 221 | 222 | #把点云写到ply文件中 223 | def writePLY(filename, cloud, faces=[]): 224 | if len(cloud.shape) != 3: 225 | print("Expected pointCloud to have 3 dimensions. Got %d instead" % len(cloud.shape)) 226 | return 227 | 228 | color = True if cloud.shape[2] == 6 else False 229 | num_points = cloud.shape[0]*cloud.shape[1] 230 | 231 | header_lines = [ 232 | 'ply', 233 | 'format ascii 1.0', 234 | 'element vertex %d' % num_points, 235 | 'property float x', 236 | 'property float y', 237 | 'property float z', 238 | ] 239 | if color: 240 | header_lines.extend([ 241 | 'property uchar diffuse_red', 242 | 'property uchar diffuse_green', 243 | 'property uchar diffuse_blue', 244 | ]) 245 | if faces != None: 246 | header_lines.extend([ 247 | 'element face %d' % len(faces), 248 | 'property list uchar int vertex_indices' 249 | ]) 250 | 251 | header_lines.extend([ 252 | 'end_header', 253 | ]) 254 | 255 | f = open(filename, 'w+') 256 | f.write('\n'.join(header_lines)) 257 | f.write('\n') 258 | 259 | lines = [] 260 | for i in range(cloud.shape[0]): 261 | for j in range(cloud.shape[1]): 262 | if color: 263 | lines.append('%s %s %s %d %d %d' % tuple(cloud[i, j, :].tolist())) 264 | else: 265 | lines.append('%s %s %s' % tuple(cloud[i, j, :].tolist())) 266 | 267 | for face in faces: 268 | lines.append(('%d' + ' %d'*len(face)) % tuple([len(face)] + list(face))) 269 | 270 | f.write('\n'.join(lines) + '\n') 271 | f.close() 272 | 273 | 274 | #把点云写到pcd文件中 275 | def writePCD(filename, pointCloud, ascii=True): 276 | if len(pointCloud.shape) != 3: 277 | print("Expected pointCloud to have 3 dimensions. Got %d instead" % len(pointCloud.shape)) 278 | return 279 | with open(filename, 'w') as f: 280 | height = pointCloud.shape[0] 281 | width = pointCloud.shape[1] 282 | f.write("# .PCD v.7 - Point Cloud Data file format\n") 283 | f.write("VERSION .7\n") 284 | if pointCloud.shape[2] == 3: 285 | f.write("FIELDS x y z\n") 286 | f.write("SIZE 4 4 4\n") 287 | f.write("TYPE F F F\n") 288 | f.write("COUNT 1 1 1\n") 289 | else: 290 | f.write("FIELDS x y z rgb\n") 291 | f.write("SIZE 4 4 4 4\n") 292 | f.write("TYPE F F F F\n") 293 | f.write("COUNT 1 1 1 1\n") 294 | f.write("WIDTH %d\n" % width) 295 | f.write("HEIGHT %d\n" % height) 296 | f.write("VIEWPOINT 0 0 0 1 0 0 0\n") 297 | f.write("POINTS %d\n" % (height * width)) 298 | if ascii: 299 | f.write("DATA ascii\n") 300 | for row in range(height): 301 | for col in range(width): 302 | if pointCloud.shape[2] == 3: 303 | f.write("%f %f %f\n" % tuple(pointCloud[row, col, :])) 304 | else: 305 | f.write("%f %f %f" % tuple(pointCloud[row, col, :3])) 306 | r = int(pointCloud[row, col, 3]) 307 | g = int(pointCloud[row, col, 4]) 308 | b = int(pointCloud[row, col, 5]) 309 | rgb_int = (r << 16) | (g << 8) | b 310 | packed = pack('i', rgb_int) 311 | rgb = unpack('f', packed)[0] 312 | f.write(" %.12e\n" % rgb) 313 | else: 314 | f.write("DATA binary\n") 315 | if pointCloud.shape[2] == 6: 316 | # These are written as bgr because rgb is interpreted as a single 317 | # little-endian float. 318 | dt = np.dtype([('x', np.float32), 319 | ('y', np.float32), 320 | ('z', np.float32), 321 | ('b', np.uint8), 322 | ('g', np.uint8), 323 | ('r', np.uint8), 324 | ('I', np.uint8)]) 325 | pointCloud_tmp = np.zeros((height*width, 1), dtype=dt) 326 | for i, k in enumerate(['x', 'y', 'z', 'r', 'g', 'b']): 327 | pointCloud_tmp[k] = pointCloud[:, :, i].reshape((height*width, 1)) 328 | pointCloud_tmp.tofile(f) 329 | else: 330 | dt = np.dtype([('x', np.float32), 331 | ('y', np.float32), 332 | ('z', np.float32), 333 | ('I', np.uint8)]) 334 | pointCloud_tmp = np.zeros((height*width, 1), dtype=dt) 335 | for i, k in enumerate(['x', 'y', 'z']): 336 | pointCloud_tmp[k] = pointCloud[:, :, i].reshape((height*width, 1)) 337 | pointCloud_tmp.tofile(f) 338 | 339 | def getRGBFromDepthTransform(calibration, camera, referenceCamera): 340 | irKey = "H_{0}_ir_from_{1}".format(camera, referenceCamera) 341 | rgbKey = "H_{0}_from_{1}".format(camera, referenceCamera) 342 | 343 | rgbFromRef = calibration[rgbKey][:] 344 | irFromRef = calibration[irKey][:] 345 | 346 | return np.dot(rgbFromRef, np.linalg.inv(irFromRef)), np.linalg.inv(rgbFromRef) 347 | 348 | 349 | def generate(path): 350 | path = path.split('/') 351 | # Parameters 352 | ycb_data_folder = os.path.join('/',*path[:-2]) # Folder that contains the ycb data. 353 | target_object = path[-2] # Full name of the target object. 354 | viewpoint_camera = path[-1].split('_')[0] # Camera which the viewpoint will be generated. 355 | viewpoint_angle = path[-1].split('_')[1].split('.')[0] # Relative angle of the object w.r.t the camera (angle of the turntable). 356 | 357 | referenceCamera = "NP5" # can only be NP5 358 | 359 | ply_fname = os.path.join(ycb_data_folder, target_object, "clouds", "pc_"+viewpoint_camera+"_"+referenceCamera+"_"+viewpoint_angle+".ply") 360 | pcd_fname = os.path.join(ycb_data_folder, target_object, "clouds", "pc_"+viewpoint_camera+"_"+referenceCamera+"_"+viewpoint_angle+".pcd") 361 | npy_fname = os.path.join(ycb_data_folder, target_object, "clouds", "pc_"+viewpoint_camera+"_"+referenceCamera+"_"+viewpoint_angle+".npy") 362 | 363 | if os.path.exists(ply_fname) and os.path.exists(pcd_fname): 364 | print(ycb_data_folder, target_object, viewpoint_camera, viewpoint_angle, 'pass') 365 | return 366 | else: 367 | print(ycb_data_folder, target_object, viewpoint_camera, viewpoint_angle) 368 | 369 | try: 370 | if not os.path.exists(os.path.join(ycb_data_folder, target_object, "clouds")): 371 | os.makedirs(os.path.join(ycb_data_folder, target_object, "clouds")) 372 | 373 | basename = "{0}_{1}".format(viewpoint_camera, viewpoint_angle) 374 | depthFilename = os.path.join(ycb_data_folder, target_object, basename + ".h5") 375 | rgbFilename = os.path.join(ycb_data_folder, target_object, basename + ".jpg") 376 | pbmFilename = os.path.join(ycb_data_folder, target_object, 'masks', basename + "_mask.pbm") 377 | 378 | calibrationFilename = os.path.join(ycb_data_folder, target_object, "calibration.h5") 379 | objFromrefFilename = os.path.join(ycb_data_folder, target_object, 'poses', '{0}_{1}_pose.h5'.format(referenceCamera, viewpoint_angle)) 380 | calibration = h5.File(calibrationFilename, 'r') 381 | objFromref = h5.File(objFromrefFilename, 'r')['H_table_from_reference_camera'][:] 382 | 383 | if not os.path.isfile(rgbFilename): 384 | print("The rgbd data is not available for the target object \"%s\". Please download the data first." % target_object) 385 | exit(1) 386 | rgbImage = imread(rgbFilename) 387 | pbmImage = imread(pbmFilename) 388 | depthK = calibration["{0}_depth_K".format(viewpoint_camera)][:] # use depth instead of ir 389 | rgbK = calibration["{0}_rgb_K".format(viewpoint_camera)][:] 390 | depthScale = np.array(calibration["{0}_ir_depth_scale".format(viewpoint_camera)]) * .0001 # 100um to meters 391 | H_RGBFromDepth, refFromRGB = getRGBFromDepthTransform(calibration, viewpoint_camera, referenceCamera) 392 | 393 | unregisteredDepthMap = h5.File(depthFilename, 'r')["depth"][:] 394 | unregisteredDepthMap = filterDiscontinuities(unregisteredDepthMap) * depthScale 395 | 396 | registeredDepthMap = registerDepthMap(unregisteredDepthMap, 397 | rgbImage, 398 | depthK, 399 | rgbK, 400 | H_RGBFromDepth) 401 | # apply mask 402 | registeredDepthMap[pbmImage == 255] = 0 403 | 404 | pointCloud = registeredDepthMapToPointCloud(registeredDepthMap, rgbImage, rgbK, refFromRGB, objFromref) 405 | 406 | writePLY(ply_fname, pointCloud) 407 | writePCD(pcd_fname, pointCloud) 408 | np.save(npy_fname, pointCloud[: ,:, :3].reshape(-1, 3)) 409 | except: 410 | f = open('exception.txt', 'a') 411 | f.write('/'.join(path) + '\n') 412 | f.close() 413 | print(ycb_data_folder, target_object, viewpoint_camera, viewpoint_angle, 'failed') 414 | return 415 | 416 | 417 | def main(): 418 | #获取当前文件所在目录的绝对路径 419 | path=os.path.dirname(os.path.abspath(__file__)) 420 | #更改（确保）当前所在目录是工作目录 421 | os.chdir(path) 422 | print(os.getcwd()) 423 | home_dir = os.environ['HOME'] 424 | file_dir = home_dir + "/dataset/PointNetGPD/ycb_rgbd/0*/*.jpg" #获取模型的路径 425 | #遍历（相对路径,相对于工作目录的路径）data/ycb_rgbd中，所有的以0开头的文件夹，并再次遍历该文件夹中的jpg图片 426 | fl = np.array(glob.glob(file_dir)) 427 | print(fl[0]) 428 | np.random.shuffle(fl) 429 | #获取cpu核心数 430 | cores = mp.cpu_count() 431 | #构建cpu的线程数 432 | pool = mp.Pool(processes=1) 433 | pool.map(generate, fl) 434 | 435 | if __name__ == "__main__": 436 | main() 437 | -------------------------------------------------------------------------------- /data/1v_500_2class_ourway2sample.model: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Hymwgk/PointNetGPD/e33b37f8da0b6326976bb339faa976f724980d82/data/1v_500_2class_ourway2sample.model -------------------------------------------------------------------------------- /data/1v_750_2class_ourway2sample.model: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Hymwgk/PointNetGPD/e33b37f8da0b6326976bb339faa976f724980d82/data/1v_750_2class_ourway2sample.model -------------------------------------------------------------------------------- /data/1v_750_2class_ourway2sample0.model: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Hymwgk/PointNetGPD/e33b37f8da0b6326976bb339faa976f724980d82/data/1v_750_2class_ourway2sample0.model -------------------------------------------------------------------------------- /data/image-20210311111731964.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Hymwgk/PointNetGPD/e33b37f8da0b6326976bb339faa976f724980d82/data/image-20210311111731964.png -------------------------------------------------------------------------------- /data/image-20210311162402868.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Hymwgk/PointNetGPD/e33b37f8da0b6326976bb339faa976f724980d82/data/image-20210311162402868.png -------------------------------------------------------------------------------- /data/在线检测时的夹爪各项参数定义.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Hymwgk/PointNetGPD/e33b37f8da0b6326976bb339faa976f724980d82/data/在线检测时的夹爪各项参数定义.png -------------------------------------------------------------------------------- /data/在线检测时的夹爪数学模型各点以及夹爪坐标系定义.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Hymwgk/PointNetGPD/e33b37f8da0b6326976bb339faa976f724980d82/data/在线检测时的夹爪数学模型各点以及夹爪坐标系定义.png -------------------------------------------------------------------------------- /main_1v.py: -------------------------------------------------------------------------------- 1 | #在python3中训练 2 | import argparse 3 | import os 4 | import time 5 | import pickle 6 | import sys 7 | 8 | print(sys.path) 9 | #sys.path.append('/home/wgk/code/PointNetGPD/model') 10 | 11 | 12 | import torch 13 | import torch.utils.data 14 | import torch.nn as nn 15 | import torch.nn.functional as F 16 | import torch.optim as optim 17 | import numpy as np 18 | #from tensorboardX import SummaryWriter 19 | #现在pytorch已经集成了tensorboard了 20 | from torch.utils.tensorboard import SummaryWriter 21 | from torch.optim.lr_scheduler import StepLR 22 | 23 | ##################/home/wgk/code/PointNetGPD/PointNetGPD/model 包########################### 24 | 25 | #需要安装pcl 26 | from model.dataset import * 27 | #pointnet的具体网络结构 28 | from model.pointnet import PointNetCls, DualPointNetCls 29 | 30 | 31 | import argparse 32 | 33 | #创建一个argparse模块中的ArgumentParser类对象 34 | # description只是解释一下，这个类对象是做什么的 35 | parser = argparse.ArgumentParser(description='pointnetGPD') 36 | #添加一个参数 tag 这个可以自定义，主要用于区分指定训练结果的存放文件夹等 37 | parser.add_argument('--tag', type=str, default='default') 38 | #默认epoch训练的轮数，大小为200 39 | parser.add_argument('--epoch', type=int, default=200) 40 | #添加模式参数，指定从train或者test中选择，且是必选项 41 | parser.add_argument('--mode', choices=['train', 'test'], required=True) 42 | #设定batch-size，默认是1 43 | parser.add_argument('--batch-size', type=int, default=1) 44 | #意思是，当命令中出现 --cuda的时候，就把 object.cuda的值设置为真 45 | parser.add_argument('--cuda', action='store_true') 46 | #添加参数 选择使用的gpu编号，默认使用0号gpu 47 | parser.add_argument('--gpu', type=int, default=0) 48 | parser.add_argument('--lr', type=float, default=0.005) 49 | parser.add_argument('--load-model', type=str, default='') 50 | 51 | 52 | parser.add_argument('--load-epoch', type=int, default=-1) 53 | parser.add_argument('--model-path', type=str, default='./assets/learned_models', 54 | help='pre-trained model path') 55 | parser.add_argument('--data-path', type=str, default=os.environ['HOME']+'/dataset/PointNetGPD', help='data path') 56 | #设置每隔多少次迭代，就打印出来一次loss还有训练进度 57 | parser.add_argument('--log-interval', type=int, default=10) 58 | #保存间隔，设置每隔几个epoch就保存一次模型，比如epoch=10就是训练10轮之后就保存当前的训练模型 59 | parser.add_argument('--save-interval', type=int, default=1) 60 | 61 | 62 | #此时，使用对象parser中的解析函数.parse_args()，读取命令行给出的参数，集合成namspace 返回给args 63 | #例子： python main_1v.py --epoch 200 --mode train --batch-size 3 64 | args = parser.parse_args() 65 | 66 | #查看是否安装好了cuda 67 | args.cuda = args.cuda if torch.cuda.is_available else False 68 | #为当前的gpu设置随机种子；这样在以后运行该程序的时候，随机数都是相同的，不会每次运行都变化一次 69 | #作用主要是为了固定随机初始化的权重值，这样就可以在每次重新从头训练网络的时候，权重值虽然是随机的 70 | #但是都是固定的，不会每次都在变化 71 | #其中的seed值，可以随便写 72 | """ 73 | https://blog.csdn.net/weixin_43002433/article/details/104706950?utm_ 74 | medium=distribute.pc_relevant.none-task-blog-BlogCommendFromMachineLearnPai2-1 75 | .add_param_isCf&depth_1-utm_source=distribute.pc_relevant.none-task-blog-Blog 76 | CommendFromMachineLearnPai2-1.add_param_isCf 77 | """ 78 | if args.cuda: 79 | torch.cuda.manual_seed(1) 80 | 81 | #获取当前文件所在目录（文件夹）的绝对路径 82 | path=os.path.dirname(os.path.abspath(__file__)) 83 | #更改（确保）当前所在目录是工作目录 84 | os.chdir(path) 85 | print(os.getcwd()) 86 | #获取当前的时间，年月日 时间 87 | current_time = time.strftime("%Y-%m-%dT%H:%M", time.localtime()) 88 | logger = SummaryWriter(os.path.join('./assets/log/', current_time)) 89 | 90 | #设置numpy的随机数种子，但是注意，这里的随机数种子是随着系统时间变化的，因此每次运行出现的随机数都是不同的 91 | np.random.seed(int(time.time())) 92 | 93 | """ 94 | 使用单线程进行数据集导入时候，有时候比较慢，会阻碍到计算的过程，于是考虑用多个线程进行数据的导入 95 | 如果参数num_workers设置的大于1，就是指定了多个线程，于是考虑使用多个线程导入数据，防止影响计算 96 | 每个线程叫一个"worker"，这个 worker_init_fn就是定义每个线程worker初始化的时候，需要执行哪些操作 97 | 其中，pid就是子线程的线程号，直接写成这样就行 98 | """ 99 | def worker_init_fn(pid): 100 | """ 101 | 为后台工作进程设置唯一但具有确定性的随机种子，只需要对numpy设置种子就行了， 102 | pytorch和python的随机数生成器会自己管理自己的子线程？ 103 | 设置torch.initial_seed() % (2**31-1)取余数，是因为numpy的种子范围是0到2**31-1 104 | """ 105 | np.random.seed(torch.initial_seed() % (2**31-1)) 106 | 107 | 108 | """ 109 | 将样本采样器返回的list中对应的样构建成一个minibatch的tensor，自定义的tensor的话 110 | 对于map型的数据集，且有多个读取进程的情况，采样器会先将样本集分成一个个batch，返回每个batch的样本的索引的 list， 111 | 再根据my_collate函数，说明如何把这些索引指定的样本整合成一个data Tensor和lable Tensor 112 | """ 113 | def my_collate(batch): 114 | batch = list(filter(lambda x:x is not None, batch)) 115 | return torch.utils.data.dataloader.default_collate(batch) 116 | 117 | # 118 | grasp_points_num=750 119 | thresh_good=0.6 120 | thresh_bad=0.6 121 | #设置点是只有xyz ？ 122 | point_channel=3 123 | 124 | 125 | """ 126 | 数据加载器,主要实现数据加载到网络中的相关作用，核心类是torch.utils.data.DataLoader 127 | DataLoader(dataset, batch_size=1, shuffle=False, sampler=None, 128 | batch_sampler=None, num_workers=0, collate_fn=None, 129 | pin_memory=False, drop_last=False, timeout=0, 130 | worker_init_fn=None) 131 | 132 | dataset：加载的数据集，继承自torch.utils.data.Dataset类的自定义类的对象，或者就是torch.utils.data.Dataset类的对象 133 | batch_size：batch size，就是batchs ize 134 | shuffle:：是否将数据打乱 135 | sampler： 样本抽样，后续会详细介绍 136 | batch_sampler，从注释可以看出，其和batch_size、shuffle等参数是互斥的，一般采用默认 137 | num_workers：使用多进程加载的进程数，0代表不使用多进程 138 | collate_fn： 如何将多个样本数据拼接成一个batch，一般使用默认的拼接方式即可 139 | pin_memory：是否将数据保存在pin memory区，pin memory中的数据转到GPU会快一些 140 | drop_last：dataset中的数据个数可能不是batch_size的整数倍，drop_last为True会将多出来不足一个batch的数据丢弃 141 | 142 | """ 143 | train_loader = torch.utils.data.DataLoader( 144 | #设置dataset，传入的是继承了torch.utils.data.Dataset类的子类的实例对象 class PointGraspOneViewDataset(torch.utils.data.Dataset): 145 | PointGraspOneViewDataset( 146 | #设置夹爪内部的点数量最少要有多少个 147 | grasp_points_num=grasp_points_num, 148 | path=args.data_path, 149 | prefix='panda', 150 | tag='train', 151 | grasp_amount_per_file=16800, #每个物体已经生成的抓取点云个数，140×120 （单物体的生成140个不同抓取姿态×单物体共有120个不同视角点云） 152 | thresh_good=thresh_good, 153 | thresh_bad=thresh_bad, 154 | ), 155 | #设置batch_size 156 | batch_size=args.batch_size, 157 | #设置使用多少个子线程来导入数据，如果设置为0，那就是直接把数据导入到主线程中 158 | num_workers=32, 159 | #如果设置为True，就使得Tensor数据最开始存放于内存的锁页内存中，这样将内存Tensor转移到GPU显存就会快一些 160 | # 当计算机内存充足的时候，选择True，如果不充足，可能使用到虚拟内存的时候，就写为False 161 | pin_memory=True, 162 | #如果设置为True，会默认构建一个乱序采样器（为False时，构建顺序采样器）；每次epoch之后都会把数据集打乱 163 | shuffle=True, 164 | 165 | # 166 | worker_init_fn=worker_init_fn, 167 | #collate_fn函数，将sampler返回的数据list合并成为一个整体的tensor，作为一个mini-batch 168 | collate_fn=my_collate, 169 | #设置，如何处理训练到最后，数据集长度不足一个batch_size时的数据，True就抛弃，否则保留 170 | drop_last=True, 171 | ) 172 | """ 173 | 测试时候用的数据集加载器；这里的test数据集是用来测试训练好的网络的准确率 174 | """ 175 | test_loader = torch.utils.data.DataLoader( 176 | PointGraspOneViewDataset( 177 | grasp_points_num=grasp_points_num, 178 | path=args.data_path, 179 | prefix='panda', 180 | #设置标签为test 181 | tag='test', 182 | grasp_amount_per_file=500, # 183 | thresh_good=thresh_good, 184 | thresh_bad=thresh_bad, 185 | with_obj=True, 186 | ), 187 | batch_size=args.batch_size, 188 | num_workers=32, 189 | pin_memory=True, 190 | shuffle=True, 191 | worker_init_fn=worker_init_fn, 192 | collate_fn=my_collate, 193 | ) 194 | 195 | is_resume = 0 196 | if args.load_model and args.load_epoch != -1: 197 | is_resume = 1 198 | #如果是测试模式 199 | if is_resume or args.mode == 'test': 200 | #加载网络结构和参数，加载到命令行 指定的gpu中 201 | model = torch.load(args.load_model, map_location='cuda:{}'.format(args.gpu)) 202 | 203 | model.device_ids = [args.gpu] 204 | print('load model {}'.format(args.load_model)) 205 | #如果是训练模式，就加载模型 206 | else: 207 | model = PointNetCls(num_points=grasp_points_num, input_chann=point_channel, k=2) 208 | 209 | 210 | 211 | #如果命令行出现了cuda字眼（args.cuda将会自动设置为True） 212 | if args.cuda: 213 | if args.gpu != -1: 214 | #设置为指定编号的gpu，使用指定编号的gpu进行训练 215 | torch.cuda.set_device(args.gpu) 216 | #将网络模型的所有参数等都转移到gpu中（刚指定的那个） 217 | model = model.cuda() 218 | else: 219 | #如果args.gpu=-1 220 | device_id = [0] 221 | #在这里选择使用哪个gpu 222 | torch.cuda.set_device(device_id[0]) 223 | #这句话，使得原有的model（网络），重构成为一个新的model，这个模型能够调用多个gpu同时运行 224 | model = nn.DataParallel(model,device_id).cuda() 225 | #选择adam优化器 226 | optimizer = optim.Adam(model.parameters(), lr=args.lr) 227 | """ 228 | from torch.optim.lr_scheduler import StepLR 229 | StepLR类主要用来调整学习率，lr=learning rate，让学习率随着epoch变化 230 | 构造函数 输入数据optimizer选择优化器；选择step_size=30, gamma=0.5 231 | 成员函数StepLR.step()是用来 232 | """ 233 | scheduler = StepLR(optimizer, step_size=30, gamma=0.5) 234 | 235 | 236 | #训练函数，仅仅是一个epoch，其中包含了很多batch 237 | def train(model, loader, epoch): 238 | #每一次epoch之前，都更新一下学习率 239 | scheduler.step() 240 | #调用train函数训练；这个model是提前设置好的pointnet 241 | model.train() 242 | 243 | torch.set_grad_enabled(True) 244 | correct = 0 245 | dataset_size = 0 246 | """ 247 | 注意，如果使用了多线程导入数据集的情况下，在当调用enumerate时，将会在此时在后台创建多线程导入数据 248 | loader是一个可迭代对象，索引是batch_idx，对象是(data, target) 249 | 这里实现的效果是：一个一个batch的训练网络，直到一个epoch训练完毕 250 | """ 251 | for batch_idx, (data, target) in enumerate(loader): 252 | #print(len(data),"data len is") 253 | """ 254 | 在实际的实验过程中发现，当某一个batch(比如剩下来的某一个batch)中间只含有一个sample 255 | 那么很有可能会报错，这里判断一下本次的batch中是不是只有一个样本，如果只有一个样本 256 | 那就跳过这个batch 257 | """ 258 | if len(data) <=1: 259 | continue 260 | 261 | dataset_size += data.shape[0] 262 | data, target = data.float(), target.long().squeeze() 263 | #如果使用cuda的话 264 | if args.cuda: 265 | data, target = data.cuda(), target.cuda() 266 | optimizer.zero_grad() 267 | #这里进行前向传播，并输出结果 但是这个output是什么 268 | output, _ = model(data) 269 | #计算loss 270 | loss = F.nll_loss(output, target) 271 | #反向传播，更新权重 272 | loss.backward() 273 | #进行学习率的优化 274 | optimizer.step() 275 | 276 | pred = output.data.max(1, keepdim=True)[1] 277 | correct += pred.eq(target.view_as(pred)).long().cpu().sum() 278 | #设置每隔多少次迭代（batch）打印一次loss，这里默认是10次迭代 279 | if batch_idx % args.log_interval == 0: 280 | print('Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}\t{}'.format( 281 | #当前第几个epoch， 当前训练到该epoch中的第几个batch，数据集总共有多大 282 | epoch, batch_idx * args.batch_size, len(loader.dataset), 283 | #计算当前训练百分比 284 | 100. * batch_idx * args.batch_size / len(loader.dataset), loss.item(), args.tag)) 285 | 286 | """在tensorboard上面绘制折线图 287 | def add_scalar(self, tag, scalar_value, global_step=None, walltime=None): 288 | 参数含义: 289 | tag (string): 显示什么名称 290 | scalar_value (float or string/blobname): 需要打印出来（同时会保存下来）的变量 291 | global_step (int): 就是坐标系横轴的下标显示什么变量 292 | walltime (float): Optional override default walltime (time.time()) 293 | with seconds after epoch of event 294 | """ 295 | logger.add_scalar('train_loss', 296 | loss.cpu().item(), #这里的loss.cpu是不是利用到了cpu？ 做了什么计算？ 297 | batch_idx + epoch * len(loader))#横坐标下标是迭代次数，训练了几次epoch 298 | """len(loader)和len(loader.dataset)区别 299 | len(loader.dataset) 返回的是数据集中的sample的数量 300 | len(loader) 返回的是 len(loader.dataset)/batch_size向上取整 就是一个epoch中包含有多少个batch（一个epoch迭代多少次） 301 | """ 302 | return float(correct)/float(dataset_size) 303 | 304 | #测试训练好的网络 305 | def test(model, loader): 306 | model.eval() 307 | torch.set_grad_enabled(False) 308 | test_loss = 0 309 | correct = 0 310 | dataset_size = 0 311 | da = {} 312 | db = {} 313 | res = [] 314 | for data, target, obj_name in loader: 315 | dataset_size += data.shape[0] 316 | data, target = data.float(), target.long().squeeze() 317 | if args.cuda: 318 | data, target = data.cuda(), target.cuda() 319 | output, _ = model(data) # N*C 320 | test_loss += F.nll_loss(output, target, size_average=False).cpu().item() 321 | pred = output.data.max(1, keepdim=True)[1] 322 | correct += pred.eq(target.view_as(pred)).long().cpu().sum() 323 | for i, j, k in zip(obj_name, pred.data.cpu().numpy(), target.data.cpu().numpy()): 324 | res.append((i, j[0], k)) 325 | 326 | test_loss /= len(loader.dataset) 327 | acc = float(correct)/float(dataset_size) 328 | return acc, test_loss 329 | 330 | 331 | def main(): 332 | # 333 | if args.mode == 'train': 334 | for epoch_i in range(is_resume*args.load_epoch, args.epoch): 335 | #训练一个epoch，把网络模型、数据集、当前是第几轮epoch的编号，都写进去 336 | acc_train = train(model, train_loader, epoch_i) 337 | #训练完毕i，精度等于多少 338 | print('Train done, acc={}'.format(acc_train)) 339 | #使用测试数据集进行测试 340 | acc, loss = test(model, test_loader) 341 | #打印测试出的精度和loss 342 | print('Test done, acc={}, loss={}'.format(acc, loss)) 343 | #然后把训练的结果测试的结果，存放在logger中，之后使用折线图打印出来 344 | logger.add_scalar('train_acc', acc_train, epoch_i) 345 | logger.add_scalar('test_acc', acc, epoch_i) 346 | logger.add_scalar('test_loss', loss, epoch_i) 347 | if epoch_i % args.save_interval == 0: 348 | path = os.path.join(args.model_path, current_time + '_{}.model'.format(epoch_i)) 349 | #如果满足要求了，就保存下来，注意这里，需要选定_use_new_zipfile_serialization=False 350 | torch.save(model, path,_use_new_zipfile_serialization=False) 351 | print('Save model @ {}'.format(path)) 352 | else: 353 | print('testing...') 354 | acc, loss = test(model, test_loader) 355 | print('Test done, acc={}, loss={}'.format(acc, loss)) 356 | 357 | if __name__ == "__main__": 358 | main() 359 | -------------------------------------------------------------------------------- /main_1v_gpd.py: -------------------------------------------------------------------------------- 1 | import argparse 2 | import os 3 | import time 4 | import pickle 5 | 6 | import torch 7 | import torch.utils.data 8 | import torch.nn as nn 9 | import torch.nn.functional as F 10 | import torch.optim as optim 11 | import numpy as np 12 | from tensorboardX import SummaryWriter 13 | from torch.optim.lr_scheduler import StepLR 14 | 15 | from model.dataset import * 16 | from model.gpd import * 17 | 18 | parser = argparse.ArgumentParser(description='pointnetGPD, gpd baseline') 19 | parser.add_argument('--tag', type=str, default='default') 20 | parser.add_argument('--epoch', type=int, default=200) 21 | parser.add_argument('--mode', choices=['train', 'test'], required=True) 22 | parser.add_argument('--batch-size', type=int, default=1) 23 | parser.add_argument('--cuda', action='store_true') 24 | parser.add_argument('--gpu', type=int, default=0) 25 | parser.add_argument('--lr', type=float, default=0.001) 26 | parser.add_argument('--load-model', type=str, default='') 27 | parser.add_argument('--load-epoch', type=int, default=-1) 28 | parser.add_argument('--model-path', type=str, default='./assets/learned_models', 29 | help='pre-trained model path') 30 | parser.add_argument('--data-path', type=str, default='./data', help='data path') 31 | parser.add_argument('--log-interval', type=int, default=10) 32 | parser.add_argument('--save-interval', type=int, default=1) 33 | 34 | args = parser.parse_args() 35 | 36 | args.cuda = args.cuda if torch.cuda.is_available else False 37 | 38 | if args.cuda: 39 | torch.cuda.manual_seed(1) 40 | 41 | logger = SummaryWriter(os.path.join('./assets/log/', 'GPD_{}'.format(args.tag))) 42 | np.random.seed(int(time.time())) 43 | 44 | def worker_init_fn(pid): 45 | np.random.seed(torch.initial_seed() % (2**31-1)) 46 | 47 | def my_collate(batch): 48 | batch = list(filter(lambda x:x is not None, batch)) 49 | return torch.utils.data.dataloader.default_collate(batch) 50 | 51 | grasp_points_num=750 52 | thresh_good=0.6 53 | thresh_bad=0.6 54 | 55 | input_size=60 56 | input_chann=3 # 3 # 12 57 | 58 | train_loader = torch.utils.data.DataLoader( 59 | PointGraspOneViewDataset( 60 | grasp_points_num=grasp_points_num, 61 | path=args.data_path, 62 | tag='train', 63 | grasp_amount_per_file=6500, 64 | thresh_good=thresh_good, 65 | thresh_bad=thresh_bad, 66 | projection=True, 67 | project_chann=input_chann, 68 | project_size=input_size, 69 | ), 70 | batch_size=args.batch_size, 71 | num_workers=32, 72 | pin_memory=True, 73 | shuffle=True, 74 | worker_init_fn=worker_init_fn, 75 | collate_fn=my_collate, 76 | ) 77 | 78 | test_loader = torch.utils.data.DataLoader( 79 | PointGraspOneViewDataset( 80 | grasp_points_num=grasp_points_num, 81 | path=args.data_path, 82 | tag='test', 83 | grasp_amount_per_file=500, 84 | thresh_good=thresh_good, 85 | thresh_bad=thresh_bad, 86 | with_obj=True, 87 | projection=True, 88 | project_chann=input_chann, 89 | project_size=input_size, 90 | ), 91 | batch_size=args.batch_size, 92 | num_workers=32, 93 | pin_memory=True, 94 | shuffle=True, 95 | worker_init_fn=worker_init_fn, 96 | collate_fn=my_collate, 97 | ) 98 | 99 | is_resume = 0 100 | if args.load_model and args.load_epoch != -1: 101 | is_resume = 1 102 | 103 | if is_resume or args.mode == 'test': 104 | model = torch.load(args.load_model, map_location='cuda:{}'.format(args.gpu)) 105 | model.device_ids = [args.gpu] 106 | print('load model {}'.format(args.load_model)) 107 | else: 108 | model = GPDClassifier(input_chann) 109 | if args.cuda: 110 | if args.gpu != -1: 111 | torch.cuda.set_device(args.gpu) 112 | model = model.cuda() 113 | else: 114 | device_id = [0,1] 115 | torch.cuda.set_device(device_id[0]) 116 | model = nn.DataParallel(model, device_ids=device_id).cuda() 117 | optimizer = optim.Adam(model.parameters(), lr=args.lr) 118 | scheduler = StepLR(optimizer, step_size=30, gamma=0.1) 119 | 120 | def train(model, loader, epoch): 121 | scheduler.step() 122 | model.train() 123 | torch.set_grad_enabled(True) 124 | correct = 0 125 | dataset_size = 0 126 | for batch_idx, (data, target) in enumerate(loader): 127 | dataset_size += data.shape[0] 128 | data, target = data.float(), target.long().squeeze() 129 | if args.cuda: 130 | data, target = data.cuda(), target.cuda() 131 | optimizer.zero_grad() 132 | output = model(data) 133 | loss = F.nll_loss(output, target) 134 | loss.backward() 135 | optimizer.step() 136 | pred = output.data.max(1, keepdim=True)[1] 137 | correct += pred.eq(target.view_as(pred)).long().cpu().sum() 138 | if batch_idx % args.log_interval == 0: 139 | print('Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}\t{}'.format( 140 | epoch, batch_idx * len(data), len(loader.dataset), 141 | 100. * batch_idx * len(data) / len(loader.dataset), loss.item(), args.tag)) 142 | logger.add_scalar('train_loss', loss.cpu().item(), 143 | batch_idx + epoch * len(loader)) 144 | return float(correct)/float(dataset_size) 145 | 146 | 147 | def test(model, loader): 148 | model.eval() 149 | torch.set_grad_enabled(False) 150 | test_loss = 0 151 | correct = 0 152 | dataset_size = 0 153 | da = {} 154 | db = {} 155 | res = [] 156 | for batch_idx, (data, target, obj_name) in enumerate(loader): 157 | dataset_size += data.shape[0] 158 | data, target = data.float(), target.long().squeeze() 159 | if args.cuda: 160 | data, target = data.cuda(), target.cuda() 161 | output = model(data) # N*C 162 | test_loss += F.nll_loss(output, target, size_average=False).cpu().item() 163 | pred = output.data.max(1, keepdim=True)[1] 164 | correct += pred.eq(target.view_as(pred)).long().cpu().sum() 165 | for i, j, k in zip(obj_name, pred.data.cpu().numpy(), target.data.cpu().numpy()): 166 | res.append((i, j[0], k)) 167 | 168 | test_loss /= len(loader.dataset) 169 | acc = float(correct)/float(dataset_size) 170 | return acc, test_loss 171 | 172 | 173 | def main(): 174 | if args.mode == 'train': 175 | for epoch in range(is_resume*args.load_epoch, args.epoch): 176 | acc_train = train(model, train_loader, epoch) 177 | print('Train done, acc={}'.format(acc_train)) 178 | acc, loss = test(model, test_loader) 179 | print('Test done, acc={}, loss={}'.format(acc, loss)) 180 | logger.add_scalar('train_acc', acc_train, epoch) 181 | logger.add_scalar('test_acc', acc, epoch) 182 | logger.add_scalar('test_loss', loss, epoch) 183 | if epoch % args.save_interval == 0: 184 | path = os.path.join(args.model_path, args.tag + '_{}.model'.format(epoch)) 185 | torch.save(model, path) 186 | print('Save model @ {}'.format(path)) 187 | else: 188 | print('testing...') 189 | acc, loss = test(model, test_loader) 190 | print('Test done, acc={}, loss={}'.format(acc, loss)) 191 | 192 | if __name__ == "__main__": 193 | main() 194 | -------------------------------------------------------------------------------- /main_1v_mc.py: -------------------------------------------------------------------------------- 1 | import argparse 2 | import os 3 | import time 4 | import pickle 5 | 6 | import torch 7 | import torch.utils.data 8 | import torch.nn as nn 9 | import torch.nn.functional as F 10 | import torch.optim as optim 11 | import numpy as np 12 | from tensorboardX import SummaryWriter 13 | from torch.optim.lr_scheduler import StepLR 14 | 15 | from model.dataset import * 16 | from model.pointnet import PointNetCls, DualPointNetCls 17 | 18 | parser = argparse.ArgumentParser(description='pointnetGPD') 19 | parser.add_argument('--tag', type=str, default='default') 20 | parser.add_argument('--epoch', type=int, default=200) 21 | parser.add_argument('--mode', choices=['train', 'test'], required=True) 22 | parser.add_argument('--batch-size', type=int, default=1) 23 | parser.add_argument('--cuda', action='store_true') 24 | parser.add_argument('--gpu', type=int, default=0) 25 | parser.add_argument('--lr', type=float, default=0.005) 26 | parser.add_argument('--load-model', type=str, default='') 27 | parser.add_argument('--load-epoch', type=int, default=-1) 28 | parser.add_argument('--model-path', type=str, default='./assets/learned_models', 29 | help='pre-trained model path') 30 | parser.add_argument('--data-path', type=str, default='./data', help='data path') 31 | parser.add_argument('--log-interval', type=int, default=10) 32 | parser.add_argument('--save-interval', type=int, default=1) 33 | 34 | args = parser.parse_args() 35 | 36 | args.cuda = args.cuda if torch.cuda.is_available else False 37 | 38 | if args.cuda: 39 | torch.cuda.manual_seed(1) 40 | 41 | logger = SummaryWriter(os.path.join('./assets/log/', args.tag)) 42 | np.random.seed(int(time.time())) 43 | 44 | def worker_init_fn(pid): 45 | np.random.seed(torch.initial_seed() % (2**31-1)) 46 | 47 | def my_collate(batch): 48 | batch = list(filter(lambda x:x is not None, batch)) 49 | return torch.utils.data.dataloader.default_collate(batch) 50 | 51 | grasp_points_num=750 52 | point_channel=3 53 | # FCs for dataset 54 | # fc_list_sub1 = np.arange(2.0, 0.75, -0.4) 55 | # fc_list_sub2 = np.arange(0.5, 0.36, -0.05) 56 | # ->[0.4, 0.5] class 1 (0.4, 0.45, 0.5) 57 | # ->(0.5, 1.2] class 2 (0.8, 1.2) 58 | # ->(1.2, 2.0] class 3 (1.6, 2.0) 59 | thresh_bad=1.2 60 | thresh_good=0.5 61 | 62 | train_loader = torch.utils.data.DataLoader( 63 | PointGraspOneViewMultiClassDataset( 64 | grasp_points_num=grasp_points_num, 65 | path=args.data_path, 66 | tag='train', 67 | grasp_amount_per_file=6500, 68 | thresh_good=thresh_good, 69 | thresh_bad=thresh_bad, 70 | ), 71 | batch_size=args.batch_size, 72 | num_workers=32, 73 | pin_memory=True, 74 | shuffle=True, 75 | worker_init_fn=worker_init_fn, 76 | collate_fn=my_collate, 77 | ) 78 | 79 | test_loader = torch.utils.data.DataLoader( 80 | PointGraspOneViewMultiClassDataset( 81 | grasp_points_num=grasp_points_num, 82 | path=args.data_path, 83 | tag='test', 84 | grasp_amount_per_file=500, 85 | thresh_good=thresh_good, 86 | thresh_bad=thresh_bad, 87 | with_obj=True, 88 | ), 89 | batch_size=args.batch_size, 90 | num_workers=32, 91 | pin_memory=True, 92 | shuffle=True, 93 | worker_init_fn=worker_init_fn, 94 | collate_fn=my_collate, 95 | ) 96 | 97 | is_resume = 0 98 | if args.load_model and args.load_epoch != -1: 99 | is_resume = 1 100 | 101 | if is_resume or args.mode == 'test': 102 | model = torch.load(args.load_model, map_location='cuda:{}'.format(args.gpu)) 103 | model.device_ids = [args.gpu] 104 | print('load model {}'.format(args.load_model)) 105 | else: 106 | model = PointNetCls(num_points=grasp_points_num, input_chann=point_channel, k=3) 107 | if args.cuda: 108 | if args.gpu != -1: 109 | torch.cuda.set_device(args.gpu) 110 | model = model.cuda() 111 | else: 112 | device_id = [0,1,2,3] 113 | torch.cuda.set_device(device_id[0]) 114 | model = nn.DataParallel(model, device_ids=device_id).cuda() 115 | optimizer = optim.Adam(model.parameters(), lr=args.lr) 116 | scheduler = StepLR(optimizer, step_size=30, gamma=0.5) 117 | 118 | def train(model, loader, epoch): 119 | scheduler.step() 120 | model.train() 121 | torch.set_grad_enabled(True) 122 | correct = 0 123 | dataset_size = 0 124 | for batch_idx, (data, target) in enumerate(loader): 125 | dataset_size += data.shape[0] 126 | data, target = data.float(), target.long().squeeze() 127 | if args.cuda: 128 | data, target = data.cuda(), target.cuda() 129 | optimizer.zero_grad() 130 | output, _ = model(data) 131 | loss = F.nll_loss(output, target) 132 | loss.backward() 133 | optimizer.step() 134 | pred = output.data.max(1, keepdim=True)[1] 135 | correct += pred.eq(target.view_as(pred)).long().cpu().sum() 136 | if batch_idx % args.log_interval == 0: 137 | print('Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}\t{}'.format( 138 | epoch, batch_idx * args.batch_size, len(loader.dataset), 139 | 100. * batch_idx * args.batch_size / len(loader.dataset), loss.item(), args.tag)) 140 | logger.add_scalar('train_loss', loss.cpu().item(), 141 | batch_idx + epoch * len(loader)) 142 | return float(correct)/float(dataset_size) 143 | 144 | 145 | def test(model, loader): 146 | model.eval() 147 | torch.set_grad_enabled(False) 148 | test_loss = 0 149 | correct = 0 150 | dataset_size = 0 151 | da = {} 152 | db = {} 153 | res = [] 154 | for data, target, obj_name in loader: 155 | dataset_size += data.shape[0] 156 | data, target = data.float(), target.long().squeeze() 157 | if args.cuda: 158 | data, target = data.cuda(), target.cuda() 159 | output, _ = model(data) # N*C 160 | test_loss += F.nll_loss(output, target, size_average=False).cpu().item() 161 | pred = output.data.max(1, keepdim=True)[1] 162 | correct += pred.eq(target.view_as(pred)).long().cpu().sum() 163 | for i, j, k in zip(obj_name, pred.data.cpu().numpy(), target.data.cpu().numpy()): 164 | res.append((i, j[0], k)) 165 | 166 | test_loss /= len(loader.dataset) 167 | acc = float(correct)/float(dataset_size) 168 | return acc, test_loss 169 | 170 | 171 | def main(): 172 | if args.mode == 'train': 173 | for epoch in range(is_resume*args.load_epoch, args.epoch): 174 | acc_train = train(model, train_loader, epoch) 175 | print('Train done, acc={}'.format(acc_train)) 176 | acc, loss = test(model, test_loader) 177 | print('Test done, acc={}, loss={}'.format(acc, loss)) 178 | logger.add_scalar('train_acc', acc_train, epoch) 179 | logger.add_scalar('test_acc', acc, epoch) 180 | logger.add_scalar('test_loss', loss, epoch) 181 | if epoch % args.save_interval == 0: 182 | path = os.path.join(args.model_path, args.tag + '_{}.model'.format(epoch)) 183 | torch.save(model, path) 184 | print('Save model @ {}'.format(path)) 185 | else: 186 | print('testing...') 187 | acc, loss = test(model, test_loader) 188 | print('Test done, acc={}, loss={}'.format(acc, loss)) 189 | 190 | if __name__ == "__main__": 191 | main() 192 | -------------------------------------------------------------------------------- /main_fullv.py: -------------------------------------------------------------------------------- 1 | import argparse 2 | import os 3 | import time 4 | import pickle 5 | 6 | import torch 7 | import torch.utils.data 8 | import torch.nn as nn 9 | import torch.nn.functional as F 10 | import torch.optim as optim 11 | import numpy as np 12 | from tensorboardX import SummaryWriter 13 | from torch.optim.lr_scheduler import StepLR 14 | 15 | from model.dataset import PointGraspDataset, PointGraspMultiViewDataset 16 | from model.pointnet import PointNetCls, DualPointNetCls 17 | 18 | parser = argparse.ArgumentParser(description='pointnetGPD') 19 | parser.add_argument('--tag', type=str, default='default') 20 | parser.add_argument('--epoch', type=int, default=200) 21 | parser.add_argument('--mode', choices=['train', 'test'], required=True) 22 | parser.add_argument('--batch-size', type=int, default=1) 23 | parser.add_argument('--cuda', action='store_true') 24 | parser.add_argument('--gpu', type=int, default=0) 25 | parser.add_argument('--lr', type=float, default=0.005) 26 | parser.add_argument('--load-model', type=str, default='') 27 | parser.add_argument('--load-epoch', type=int, default=-1) 28 | parser.add_argument('--model-path', type=str, default='./assets/learned_models', 29 | help='pre-trained model path') 30 | parser.add_argument('--data-path', type=str, default='./data', help='data path') 31 | parser.add_argument('--log-interval', type=int, default=10) 32 | parser.add_argument('--save-interval', type=int, default=1) 33 | 34 | args = parser.parse_args() 35 | 36 | args.cuda = args.cuda if torch.cuda.is_available else False 37 | 38 | if args.cuda: 39 | torch.cuda.manual_seed(1) 40 | 41 | logger = SummaryWriter(os.path.join('./assets/log/', args.tag)) 42 | np.random.seed(int(time.time())) 43 | 44 | def worker_init_fn(pid): 45 | np.random.seed(torch.initial_seed() % (2**31-1)) 46 | 47 | def my_collate(batch): 48 | batch = list(filter(lambda x:x is not None, batch)) 49 | return torch.utils.data.dataloader.default_collate(batch) 50 | 51 | grasp_points_num=1000 52 | obj_points_num=50000 53 | pc_file_used_num=20 54 | thresh_good=0.6 55 | thresh_bad=0.6 56 | point_channel=3 57 | 58 | train_loader = torch.utils.data.DataLoader( 59 | PointGraspDataset( 60 | obj_points_num=obj_points_num, 61 | grasp_points_num=grasp_points_num, 62 | pc_file_used_num=pc_file_used_num, 63 | path=args.data_path, 64 | tag='train', 65 | grasp_amount_per_file=6500, 66 | thresh_good=thresh_good, 67 | thresh_bad=thresh_bad, 68 | ), 69 | batch_size=args.batch_size, 70 | num_workers=32, 71 | pin_memory=True, 72 | shuffle=True, 73 | worker_init_fn=worker_init_fn, 74 | collate_fn=my_collate, 75 | ) 76 | 77 | test_loader = torch.utils.data.DataLoader( 78 | PointGraspDataset( 79 | obj_points_num=obj_points_num, 80 | grasp_points_num=grasp_points_num, 81 | pc_file_used_num=pc_file_used_num, 82 | path=args.data_path, 83 | tag='test', 84 | grasp_amount_per_file=500, 85 | thresh_good=thresh_good, 86 | thresh_bad=thresh_bad, 87 | with_obj=True, 88 | ), 89 | batch_size=args.batch_size, 90 | num_workers=32, 91 | pin_memory=True, 92 | shuffle=True, 93 | worker_init_fn=worker_init_fn, 94 | collate_fn=my_collate, 95 | ) 96 | 97 | is_resume = 0 98 | if args.load_model and args.load_epoch != -1: 99 | is_resume = 1 100 | 101 | if is_resume or args.mode == 'test': 102 | model = torch.load(args.load_model, map_location='cuda:{}'.format(args.gpu)) 103 | model.device_ids = [args.gpu] 104 | print('load model {}'.format(args.load_model)) 105 | else: 106 | model = PointNetCls(num_points=grasp_points_num, input_chann=point_channel, k=2) 107 | if args.cuda: 108 | if args.gpu != -1: 109 | torch.cuda.set_device(args.gpu) 110 | model = model.cuda() 111 | else: 112 | device_id = [0,1,2,3] 113 | torch.cuda.set_device(device_id[0]) 114 | model = nn.DataParallel(model, device_ids=device_id).cuda() 115 | optimizer = optim.Adam(model.parameters(), lr=args.lr) 116 | scheduler = StepLR(optimizer, step_size=30, gamma=0.5) 117 | 118 | def train(model, loader, epoch): 119 | scheduler.step() 120 | model.train() 121 | torch.set_grad_enabled(True) 122 | correct = 0 123 | dataset_size = 0 124 | for batch_idx, (data, target) in enumerate(loader): 125 | dataset_size += data.shape[0] 126 | data, target = data.float(), target.long().squeeze() 127 | if args.cuda: 128 | data, target = data.cuda(), target.cuda() 129 | optimizer.zero_grad() 130 | output, _ = model(data) 131 | loss = F.nll_loss(output, target) 132 | loss.backward() 133 | optimizer.step() 134 | pred = output.data.max(1, keepdim=True)[1] 135 | correct += pred.eq(target.view_as(pred)).long().cpu().sum() 136 | if batch_idx % args.log_interval == 0: 137 | print('Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}\t{}'.format( 138 | epoch, batch_idx * args.batch_size, len(loader.dataset), 139 | 100. * batch_idx * args.batch_size / len(loader.dataset), loss.item(), args.tag)) 140 | logger.add_scalar('train_loss', loss.cpu().item(), 141 | batch_idx + epoch * len(loader)) 142 | return float(correct)/float(dataset_size) 143 | 144 | 145 | def test(model, loader): 146 | model.eval() 147 | torch.set_grad_enabled(False) 148 | test_loss = 0 149 | correct = 0 150 | dataset_size = 0 151 | da = {} 152 | db = {} 153 | res = [] 154 | for data, target, obj_name in loader: 155 | dataset_size += data.shape[0] 156 | data, target = data.float(), target.long().squeeze() 157 | if args.cuda: 158 | data, target = data.cuda(), target.cuda() 159 | output, _ = model(data) # N*C 160 | test_loss += F.nll_loss(output, target, size_average=False).cpu().item() 161 | pred = output.data.max(1, keepdim=True)[1] 162 | correct += pred.eq(target.view_as(pred)).long().cpu().sum() 163 | for i, j, k in zip(obj_name, pred.data.cpu().numpy(), target.data.cpu().numpy()): 164 | res.append((i, j[0], k)) 165 | 166 | test_loss /= len(loader.dataset) 167 | acc = float(correct)/float(dataset_size) 168 | return acc, test_loss 169 | 170 | 171 | def main(): 172 | if args.mode == 'train': 173 | for epoch in range(is_resume*args.load_epoch, args.epoch): 174 | acc_train = train(model, train_loader, epoch) 175 | print('Train done, acc={}'.format(acc_train)) 176 | acc, loss = test(model, test_loader) 177 | print('Test done, acc={}, loss={}'.format(acc, loss)) 178 | logger.add_scalar('train_acc', acc_train, epoch) 179 | logger.add_scalar('test_acc', acc, epoch) 180 | logger.add_scalar('test_loss', loss, epoch) 181 | if epoch % args.save_interval == 0: 182 | path = os.path.join(args.model_path, args.tag + '_{}.model'.format(epoch)) 183 | torch.save(model, path) 184 | print('Save model @ {}'.format(path)) 185 | else: 186 | print('testing...') 187 | acc, loss = test(model, test_loader) 188 | print('Test done, acc={}, loss={}'.format(acc, loss)) 189 | 190 | if __name__ == "__main__": 191 | main() 192 | -------------------------------------------------------------------------------- /main_fullv_gpd.py: -------------------------------------------------------------------------------- 1 | import argparse 2 | import os 3 | import time 4 | import pickle 5 | 6 | import torch 7 | import torch.utils.data 8 | import torch.nn as nn 9 | import torch.nn.functional as F 10 | import torch.optim as optim 11 | import numpy as np 12 | from tensorboardX import SummaryWriter 13 | from torch.optim.lr_scheduler import StepLR 14 | 15 | from model.dataset import PointGraspDataset 16 | from model.gpd import * 17 | 18 | parser = argparse.ArgumentParser(description='pointnetGPD, gpd baseline') 19 | parser.add_argument('--tag', type=str, default='default') 20 | parser.add_argument('--epoch', type=int, default=200) 21 | parser.add_argument('--mode', choices=['train', 'test'], required=True) 22 | parser.add_argument('--batch-size', type=int, default=1) 23 | parser.add_argument('--cuda', action='store_true') 24 | parser.add_argument('--gpu', type=int, default=0) 25 | parser.add_argument('--lr', type=float, default=0.001) 26 | parser.add_argument('--load-model', type=str, default='') 27 | parser.add_argument('--load-epoch', type=int, default=-1) 28 | parser.add_argument('--model-path', type=str, default='./assets/learned_models', 29 | help='pre-trained model path') 30 | parser.add_argument('--data-path', type=str, default='./data', help='data path') 31 | parser.add_argument('--log-interval', type=int, default=10) 32 | parser.add_argument('--save-interval', type=int, default=1) 33 | 34 | args = parser.parse_args() 35 | 36 | args.cuda = args.cuda if torch.cuda.is_available else False 37 | 38 | if args.cuda: 39 | torch.cuda.manual_seed(1) 40 | 41 | logger = SummaryWriter(os.path.join('./assets/log/', 'GPD_{}'.format(args.tag))) 42 | np.random.seed(int(time.time())) 43 | 44 | def worker_init_fn(pid): 45 | np.random.seed(torch.initial_seed() % (2**31-1)) 46 | 47 | def my_collate(batch): 48 | batch = list(filter(lambda x:x is not None, batch)) 49 | return torch.utils.data.dataloader.default_collate(batch) 50 | 51 | grasp_points_num=1000 52 | obj_points_num=50000 53 | pc_file_used_num=20 54 | thresh_good=0.6 55 | thresh_bad=0.6 56 | 57 | input_size=60 58 | input_chann=3 # 3 # 12 59 | 60 | train_loader = torch.utils.data.DataLoader( 61 | PointGraspDataset( 62 | obj_points_num=obj_points_num, 63 | grasp_points_num=grasp_points_num, 64 | pc_file_used_num=pc_file_used_num, 65 | path=args.data_path, 66 | tag='train', 67 | grasp_amount_per_file=6500, 68 | thresh_good=thresh_good, 69 | thresh_bad=thresh_bad, 70 | projection=True, 71 | project_chann=input_chann, 72 | project_size=input_size, 73 | ), 74 | batch_size=args.batch_size, 75 | num_workers=32, 76 | pin_memory=True, 77 | shuffle=True, 78 | worker_init_fn=worker_init_fn, 79 | collate_fn=my_collate, 80 | ) 81 | 82 | test_loader = torch.utils.data.DataLoader( 83 | PointGraspDataset( 84 | obj_points_num=obj_points_num, 85 | grasp_points_num=grasp_points_num, 86 | pc_file_used_num=pc_file_used_num, 87 | path=args.data_path, 88 | tag='test', 89 | grasp_amount_per_file=500, 90 | thresh_good=thresh_good, 91 | thresh_bad=thresh_bad, 92 | with_obj=True, 93 | projection=True, 94 | project_chann=input_chann, 95 | project_size=input_size, 96 | ), 97 | batch_size=args.batch_size, 98 | num_workers=32, 99 | pin_memory=True, 100 | shuffle=True, 101 | worker_init_fn=worker_init_fn, 102 | collate_fn=my_collate, 103 | ) 104 | 105 | is_resume = 0 106 | if args.load_model and args.load_epoch != -1: 107 | is_resume = 1 108 | 109 | if is_resume or args.mode == 'test': 110 | model = torch.load(args.load_model, map_location='cuda:{}'.format(args.gpu)) 111 | model.device_ids = [args.gpu] 112 | print('load model {}'.format(args.load_model)) 113 | else: 114 | model = GPDClassifier(input_chann) 115 | if args.cuda: 116 | if args.gpu != -1: 117 | torch.cuda.set_device(args.gpu) 118 | model = model.cuda() 119 | else: 120 | device_id = [0,1] 121 | torch.cuda.set_device(device_id[0]) 122 | model = nn.DataParallel(model, device_ids=device_id).cuda() 123 | optimizer = optim.Adam(model.parameters(), lr=args.lr) 124 | scheduler = StepLR(optimizer, step_size=30, gamma=0.1) 125 | 126 | def train(model, loader, epoch): 127 | scheduler.step() 128 | model.train() 129 | torch.set_grad_enabled(True) 130 | correct = 0 131 | dataset_size = 0 132 | for batch_idx, (data, target) in enumerate(loader): 133 | dataset_size += data.shape[0] 134 | data, target = data.float(), target.long().squeeze() 135 | if args.cuda: 136 | data, target = data.cuda(), target.cuda() 137 | optimizer.zero_grad() 138 | output = model(data) 139 | loss = F.nll_loss(output, target) 140 | loss.backward() 141 | optimizer.step() 142 | pred = output.data.max(1, keepdim=True)[1] 143 | correct += pred.eq(target.view_as(pred)).long().cpu().sum() 144 | if batch_idx % args.log_interval == 0: 145 | print('Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}\t{}'.format( 146 | epoch, batch_idx * len(data), len(loader.dataset), 147 | 100. * batch_idx * len(data) / len(loader.dataset), loss.item(), args.tag)) 148 | logger.add_scalar('train_loss', loss.cpu().item(), 149 | batch_idx + epoch * len(loader)) 150 | return float(correct)/float(dataset_size) 151 | 152 | 153 | def test(model, loader): 154 | model.eval() 155 | torch.set_grad_enabled(False) 156 | test_loss = 0 157 | correct = 0 158 | dataset_size = 0 159 | da = {} 160 | db = {} 161 | res = [] 162 | for batch_idx, (data, target, obj_name) in enumerate(loader): 163 | dataset_size += data.shape[0] 164 | data, target = data.float(), target.long().squeeze() 165 | if args.cuda: 166 | data, target = data.cuda(), target.cuda() 167 | output = model(data) # N*C 168 | test_loss += F.nll_loss(output, target, size_average=False).cpu().item() 169 | pred = output.data.max(1, keepdim=True)[1] 170 | correct += pred.eq(target.view_as(pred)).long().cpu().sum() 171 | for i, j, k in zip(obj_name, pred.data.cpu().numpy(), target.data.cpu().numpy()): 172 | res.append((i, j[0], k)) 173 | 174 | test_loss /= len(loader.dataset) 175 | acc = float(correct)/float(dataset_size) 176 | return acc, test_loss 177 | 178 | 179 | def main(): 180 | if args.mode == 'train': 181 | for epoch in range(is_resume*args.load_epoch, args.epoch): 182 | acc_train = train(model, train_loader, epoch) 183 | print('Train done, acc={}'.format(acc_train)) 184 | acc, loss = test(model, test_loader) 185 | print('Test done, acc={}, loss={}'.format(acc, loss)) 186 | logger.add_scalar('train_acc', acc_train, epoch) 187 | logger.add_scalar('test_acc', acc, epoch) 188 | logger.add_scalar('test_loss', loss, epoch) 189 | if epoch % args.save_interval == 0: 190 | path = os.path.join(args.model_path, args.tag + '_{}.model'.format(epoch)) 191 | torch.save(model, path) 192 | print('Save model @ {}'.format(path)) 193 | else: 194 | print('testing...') 195 | acc, loss = test(model, test_loader) 196 | print('Test done, acc={}, loss={}'.format(acc, loss)) 197 | 198 | if __name__ == "__main__": 199 | main() 200 | -------------------------------------------------------------------------------- /main_fullv_mc.py: -------------------------------------------------------------------------------- 1 | import argparse 2 | import os 3 | import time 4 | import pickle 5 | 6 | import torch 7 | import torch.utils.data 8 | import torch.nn as nn 9 | import torch.nn.functional as F 10 | import torch.optim as optim 11 | import numpy as np 12 | from tensorboardX import SummaryWriter 13 | from torch.optim.lr_scheduler import StepLR 14 | 15 | from model.dataset import PointGraspDataset, PointGraspMultiClassDataset 16 | from model.pointnet import PointNetCls, DualPointNetCls 17 | 18 | parser = argparse.ArgumentParser(description='pointnetGPD') 19 | parser.add_argument('--tag', type=str, default='default') 20 | parser.add_argument('--epoch', type=int, default=200) 21 | parser.add_argument('--mode', choices=['train', 'test'], required=True) 22 | parser.add_argument('--batch-size', type=int, default=1) 23 | parser.add_argument('--cuda', action='store_true') 24 | parser.add_argument('--gpu', type=int, default=0) 25 | parser.add_argument('--lr', type=float, default=0.005) 26 | parser.add_argument('--load-model', type=str, default='') 27 | parser.add_argument('--load-epoch', type=int, default=-1) 28 | parser.add_argument('--model-path', type=str, default='./assets/learned_models', 29 | help='pre-trained model path') 30 | parser.add_argument('--data-path', type=str, default='./data', help='data path') 31 | parser.add_argument('--log-interval', type=int, default=10) 32 | parser.add_argument('--save-interval', type=int, default=1) 33 | 34 | args = parser.parse_args() 35 | 36 | args.cuda = args.cuda if torch.cuda.is_available else False 37 | 38 | if args.cuda: 39 | torch.cuda.manual_seed(1) 40 | 41 | logger = SummaryWriter(os.path.join('./assets/log/', args.tag)) 42 | np.random.seed(int(time.time())) 43 | 44 | def worker_init_fn(pid): 45 | np.random.seed(torch.initial_seed() % (2**31-1)) 46 | 47 | def my_collate(batch): 48 | batch = list(filter(lambda x:x is not None, batch)) 49 | return torch.utils.data.dataloader.default_collate(batch) 50 | 51 | grasp_points_num=1000 52 | obj_points_num=50000 53 | pc_file_used_num=20 54 | point_channel=3 55 | # FCs for dataset 56 | # fc_list_sub1 = np.arange(2.0, 0.75, -0.4) 57 | # fc_list_sub2 = np.arange(0.5, 0.36, -0.05) 58 | # ->[0.4, 0.5] class 1 (0.4, 0.45, 0.5) 59 | # ->(0.5, 1.2] class 2 (0.8, 1.2) 60 | # ->(1.2, 2.0] class 3 (1.6, 2.0) 61 | thresh_bad=1.2 62 | thresh_good=0.5 63 | 64 | train_loader = torch.utils.data.DataLoader( 65 | PointGraspMultiClassDataset( 66 | obj_points_num=obj_points_num, 67 | grasp_points_num=grasp_points_num, 68 | pc_file_used_num=pc_file_used_num, 69 | path=args.data_path, 70 | tag='train', 71 | grasp_amount_per_file=6500, 72 | thresh_good=thresh_good, 73 | thresh_bad=thresh_bad, 74 | ), 75 | batch_size=args.batch_size, 76 | num_workers=32, 77 | pin_memory=True, 78 | shuffle=True, 79 | worker_init_fn=worker_init_fn, 80 | collate_fn=my_collate, 81 | ) 82 | 83 | test_loader = torch.utils.data.DataLoader( 84 | PointGraspMultiClassDataset( 85 | obj_points_num=obj_points_num, 86 | grasp_points_num=grasp_points_num, 87 | pc_file_used_num=pc_file_used_num, 88 | path=args.data_path, 89 | tag='test', 90 | grasp_amount_per_file=500, 91 | thresh_good=thresh_good, 92 | thresh_bad=thresh_bad, 93 | with_obj=True, 94 | ), 95 | batch_size=args.batch_size, 96 | num_workers=32, 97 | pin_memory=True, 98 | shuffle=True, 99 | worker_init_fn=worker_init_fn, 100 | collate_fn=my_collate, 101 | ) 102 | 103 | is_resume = 0 104 | if args.load_model and args.load_epoch != -1: 105 | is_resume = 1 106 | 107 | if is_resume or args.mode == 'test': 108 | model = torch.load(args.load_model, map_location='cuda:{}'.format(args.gpu)) 109 | model.device_ids = [args.gpu] 110 | print('load model {}'.format(args.load_model)) 111 | else: 112 | model = PointNetCls(num_points=grasp_points_num, input_chann=point_channel, k=3) 113 | if args.cuda: 114 | if args.gpu != -1: 115 | torch.cuda.set_device(args.gpu) 116 | model = model.cuda() 117 | else: 118 | device_id = [0,1,2,3] 119 | torch.cuda.set_device(device_id[0]) 120 | model = nn.DataParallel(model, device_ids=device_id).cuda() 121 | optimizer = optim.Adam(model.parameters(), lr=args.lr) 122 | scheduler = StepLR(optimizer, step_size=30, gamma=0.5) 123 | 124 | def train(model, loader, epoch): 125 | scheduler.step() 126 | model.train() 127 | torch.set_grad_enabled(True) 128 | correct = 0 129 | dataset_size = 0 130 | for batch_idx, (data, target) in enumerate(loader): 131 | dataset_size += data.shape[0] 132 | data, target = data.float(), target.long().squeeze() 133 | if args.cuda: 134 | data, target = data.cuda(), target.cuda() 135 | optimizer.zero_grad() 136 | output, _ = model(data) 137 | loss = F.nll_loss(output, target) 138 | loss.backward() 139 | optimizer.step() 140 | pred = output.data.max(1, keepdim=True)[1] 141 | correct += pred.eq(target.view_as(pred)).long().cpu().sum() 142 | if batch_idx % args.log_interval == 0: 143 | print('Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}\t{}'.format( 144 | epoch, batch_idx * args.batch_size, len(loader.dataset), 145 | 100. * batch_idx * args.batch_size / len(loader.dataset), loss.item(), args.tag)) 146 | logger.add_scalar('train_loss', loss.cpu().item(), 147 | batch_idx + epoch * len(loader)) 148 | return float(correct)/float(dataset_size) 149 | 150 | 151 | def test(model, loader): 152 | model.eval() 153 | torch.set_grad_enabled(False) 154 | test_loss = 0 155 | correct = 0 156 | dataset_size = 0 157 | da = {} 158 | db = {} 159 | res = [] 160 | for data, target, obj_name in loader: 161 | dataset_size += data.shape[0] 162 | data, target = data.float(), target.long().squeeze() 163 | if args.cuda: 164 | data, target = data.cuda(), target.cuda() 165 | output, _ = model(data) # N*C 166 | test_loss += F.nll_loss(output, target, size_average=False).cpu().item() 167 | pred = output.data.max(1, keepdim=True)[1] 168 | correct += pred.eq(target.view_as(pred)).long().cpu().sum() 169 | for i, j, k in zip(obj_name, pred.data.cpu().numpy(), target.data.cpu().numpy()): 170 | res.append((i, j[0], k)) 171 | 172 | test_loss /= len(loader.dataset) 173 | acc = float(correct)/float(dataset_size) 174 | return acc, test_loss 175 | 176 | 177 | def main(): 178 | if args.mode == 'train': 179 | for epoch in range(is_resume*args.load_epoch, args.epoch): 180 | acc_train = train(model, train_loader, epoch) 181 | print('Train done, acc={}'.format(acc_train)) 182 | acc, loss = test(model, test_loader) 183 | print('Test done, acc={}, loss={}'.format(acc, loss)) 184 | logger.add_scalar('train_acc', acc_train, epoch) 185 | logger.add_scalar('test_acc', acc, epoch) 186 | logger.add_scalar('test_loss', loss, epoch) 187 | if epoch % args.save_interval == 0: 188 | path = os.path.join(args.model_path, args.tag + '_{}.model'.format(epoch)) 189 | torch.save(model, path) 190 | print('Save model @ {}'.format(path)) 191 | else: 192 | print('testing...') 193 | acc, loss = test(model, test_loader) 194 | print('Test done, acc={}, loss={}'.format(acc, loss)) 195 | 196 | if __name__ == "__main__": 197 | main() 198 | -------------------------------------------------------------------------------- /main_test.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python 2 | # -*- coding: utf-8 -*- 3 | from __future__ import print_function 4 | import argparse 5 | import os 6 | import time 7 | import torch 8 | import torch.utils.data 9 | import torch.nn as nn 10 | import numpy as np 11 | from scipy.stats import mode 12 | import sys 13 | from os import path 14 | 15 | sys.path.append(path.dirname(path.dirname(path.abspath("__file__")))) 16 | parser = argparse.ArgumentParser(description="pointnetGPD") 17 | #使用本py文件的路径作为基准 18 | os.chdir(os.path.dirname(os.path.abspath(__file__))) 19 | print(os.getcwd()) 20 | #是否使用cuda 21 | parser.add_argument("--cuda", action="store_true", default=False) 22 | #指定使用的gpu编号 23 | parser.add_argument("--gpu", type=int, default=0) 24 | 25 | #加载训练好的pointnet 26 | parser.add_argument("--load-model", type=str, 27 | default="../data/pointgpd_3class.model") 28 | #是否显示最终的抓取结果 29 | parser.add_argument("--show_final_grasp", action="store_true", default=False) 30 | #是否是托盘抓取？就是待抓取物品有没有托盘 31 | parser.add_argument("--tray_grasp", action="store_true", default=False) 32 | # 33 | parser.add_argument("--using_mp", action="store_true", default=False) 34 | 35 | parser.add_argument("--model_type", type=str) 36 | 37 | args = parser.parse_args() 38 | 39 | args.cuda = args.cuda if torch.cuda.is_available else False 40 | 41 | if args.cuda: 42 | torch.cuda.manual_seed(1) 43 | 44 | np.random.seed(int(time.time())) 45 | 46 | 47 | 48 | if args.model_type == "500": 49 | args.load_model = "./data/1v_500_2class_ourway2sample.model" 50 | elif args.model_type == "750": 51 | args.load_model = "./data/1v_750_2class_ourway2sample.model" 52 | elif args.model_type == "3class": # input points number is 500 53 | args.load_model = "./data/pointgpd_3class.model" 54 | else: 55 | print("Using default model file") 56 | model = torch.load(args.load_model, map_location="cpu") 57 | model.device_ids = [args.gpu] 58 | print("load model {}".format(args.load_model)) 59 | 60 | if args.cuda: 61 | model = torch.load(args.load_model, map_location="cuda:{}".format(args.gpu)) 62 | if args.gpu != -1: 63 | torch.cuda.set_device(args.gpu) 64 | model = model.cuda() 65 | else: 66 | device_id = [0, 1] 67 | torch.cuda.set_device(device_id[0]) 68 | model = nn.DataParallel(model, device_ids=device_id).cuda() 69 | if isinstance(model, torch.nn.DataParallel): 70 | model = model.module 71 | 72 | 73 | def test_network(model_, local_pc): 74 | local_pc = local_pc.T 75 | local_pc = local_pc[np.newaxis, ...] 76 | local_pc = torch.FloatTensor(local_pc) 77 | if args.cuda: 78 | local_pc = local_pc.cuda() 79 | output, _ = model_(local_pc) # N*C 80 | output = output.softmax(1) 81 | pred = output.data.max(1, keepdim=True)[1] 82 | output = output.cpu() 83 | return pred[0], output.data.numpy() 84 | 85 | 86 | def main(): 87 | repeat = 10 88 | num_point = 500 89 | model.eval() 90 | torch.set_grad_enabled(False) 91 | 92 | # load pc(should be in local gripper coordinate) 93 | # local_pc: (N, 3) 94 | # local_pc = np.load("test.npy") 95 | local_pc = np.random.random([500, 3]) # test only 96 | predict = [] 97 | for _ in range(repeat): 98 | if len(local_pc) >= num_point: 99 | local_pc = local_pc[np.random.choice(len(local_pc), num_point, replace=False)] 100 | else: 101 | local_pc = local_pc[np.random.choice(len(local_pc), num_point, replace=True)] 102 | 103 | # run model 104 | predict.append(test_network(model, local_pc)[0]) 105 | print("voting: ", predict) 106 | predict = mode(predict).mode[0] 107 | 108 | # output 109 | print("Test result:", predict) 110 | 111 | 112 | if __name__ == "__main__": 113 | main() 114 | -------------------------------------------------------------------------------- /model/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Hymwgk/PointNetGPD/e33b37f8da0b6326976bb339faa976f724980d82/model/__init__.py -------------------------------------------------------------------------------- /model/gpd.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | import torch.nn.functional as F 4 | 5 | class GPDClassifier(nn.Module): 6 | """ 7 | Input: (batch_size, input_chann, 60, 60) 8 | """ 9 | def __init__(self, input_chann, dropout=False): 10 | super(GPDClassifier, self).__init__() 11 | self.conv1 = nn.Conv2d(input_chann, 20, 5) 12 | self.pool1 = nn.MaxPool2d(2, stride=2) 13 | self.conv2 = nn.Conv2d(20, 50, 5) 14 | self.pool2 = nn.MaxPool2d(2, stride=2) 15 | self.fc1 = nn.Linear(12*12*50, 500) 16 | self.dp = nn.Dropout2d(p=0.5, inplace=False) 17 | self.relu = nn.ReLU() 18 | self.fc2 = nn.Linear(500, 2) 19 | self.if_dropout = dropout 20 | 21 | def forward(self, x): 22 | x = self.pool1(self.conv1(x)) 23 | x = self.pool2(self.conv2(x)) 24 | x = x.view(-1, 7200) 25 | x = self.relu(self.fc1(x)) 26 | if self.if_dropout: 27 | x = self.dp(x) 28 | x = self.fc2(x) 29 | x = F.log_softmax(x, dim=-1) 30 | 31 | return x 32 | -------------------------------------------------------------------------------- /model/pointnet.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import torch 3 | import torch.nn as nn 4 | from torch.autograd import Variable 5 | import torch.nn.functional as F 6 | 7 | 8 | class STN3d(nn.Module): 9 | def __init__(self, num_points = 2500, input_chann = 3): 10 | super(STN3d, self).__init__() 11 | self.num_points = num_points 12 | self.conv1 = torch.nn.Conv1d(input_chann, 64, 1) 13 | self.conv2 = torch.nn.Conv1d(64, 128, 1) 14 | self.conv3 = torch.nn.Conv1d(128, 1024, 1) 15 | self.mp1 = torch.nn.MaxPool1d(num_points) 16 | self.fc1 = nn.Linear(1024, 512) 17 | self.fc2 = nn.Linear(512, 256) 18 | self.fc3 = nn.Linear(256, 9) 19 | self.relu = nn.ReLU() 20 | 21 | self.bn1 = nn.BatchNorm1d(64) 22 | self.bn2 = nn.BatchNorm1d(128) 23 | self.bn3 = nn.BatchNorm1d(1024) 24 | self.bn4 = nn.BatchNorm1d(512) 25 | self.bn5 = nn.BatchNorm1d(256) 26 | 27 | 28 | def forward(self, x): 29 | batchsize = x.size()[0] 30 | x = F.relu(self.bn1(self.conv1(x))) 31 | x = F.relu(self.bn2(self.conv2(x))) 32 | x = F.relu(self.bn3(self.conv3(x))) 33 | x = self.mp1(x) 34 | x = x.view(-1, 1024) 35 | 36 | x = F.relu(self.bn4(self.fc1(x))) 37 | x = F.relu(self.bn5(self.fc2(x))) 38 | x = self.fc3(x) 39 | 40 | iden = Variable(torch.from_numpy(np.array([1,0,0,0,1,0,0,0,1]).astype(np.float32))).view(1,9).repeat(batchsize,1) 41 | if x.is_cuda: 42 | iden = iden.cuda() 43 | x = x + iden 44 | x = x.view(-1, 3, 3) 45 | return x 46 | 47 | 48 | class SimpleSTN3d(nn.Module): 49 | def __init__(self, num_points = 2500, input_chann = 3): 50 | super(SimpleSTN3d, self).__init__() 51 | self.num_points = num_points 52 | self.conv1 = torch.nn.Conv1d(input_chann, 64, 1) 53 | self.conv2 = torch.nn.Conv1d(64, 128, 1) 54 | self.conv3 = torch.nn.Conv1d(128, 256, 1) 55 | self.mp1 = torch.nn.MaxPool1d(num_points) 56 | self.fc1 = nn.Linear(256, 128) 57 | self.fc2 = nn.Linear(128, 64) 58 | self.fc3 = nn.Linear(64, 9) 59 | self.relu = nn.ReLU() 60 | 61 | self.bn1 = nn.BatchNorm1d(64) 62 | self.bn2 = nn.BatchNorm1d(128) 63 | self.bn3 = nn.BatchNorm1d(256) 64 | self.bn4 = nn.BatchNorm1d(128) 65 | self.bn5 = nn.BatchNorm1d(64) 66 | 67 | 68 | def forward(self, x): 69 | batchsize = x.size()[0] 70 | x = F.relu(self.bn1(self.conv1(x))) 71 | x = F.relu(self.bn2(self.conv2(x))) 72 | x = F.relu(self.bn3(self.conv3(x))) 73 | x = self.mp1(x) 74 | x = x.view(-1, 256) 75 | 76 | x = F.relu(self.bn4(self.fc1(x))) 77 | x = F.relu(self.bn5(self.fc2(x))) 78 | x = self.fc3(x) 79 | 80 | iden = Variable(torch.from_numpy(np.array([1,0,0,0,1,0,0,0,1]).astype(np.float32))).view(1,9).repeat(batchsize,1) 81 | if x.is_cuda: 82 | iden = iden.cuda() 83 | x = x + iden 84 | x = x.view(-1, 3, 3) 85 | return x 86 | 87 | 88 | class DualPointNetfeat(nn.Module): 89 | def __init__(self, num_points = 2500, input_chann = 6, global_feat = True): 90 | super(DualPointNetfeat, self).__init__() 91 | self.stn1 = SimpleSTN3d(num_points = num_points, input_chann = input_chann//2) 92 | self.stn2 = SimpleSTN3d(num_points = num_points, input_chann = input_chann//2) 93 | self.conv1 = torch.nn.Conv1d(input_chann, 64, 1) 94 | self.conv2 = torch.nn.Conv1d(64, 128, 1) 95 | self.conv3 = torch.nn.Conv1d(128, 1024, 1) 96 | self.bn1 = nn.BatchNorm1d(64) 97 | self.bn2 = nn.BatchNorm1d(128) 98 | self.bn3 = nn.BatchNorm1d(1024) 99 | self.mp1 = torch.nn.MaxPool1d(num_points) 100 | self.num_points = num_points 101 | self.global_feat = global_feat 102 | def forward(self, x): 103 | batchsize = x.size()[0] 104 | trans1 = self.stn1(x[:, 0:3, :]) 105 | trans2 = self.stn2(x[:, 3:6, :]) 106 | x = x.transpose(2,1) 107 | x = torch.cat([torch.bmm(x[..., 0:3], trans1), torch.bmm(x[..., 3:6], trans2)], dim=-1) 108 | x = x.transpose(2,1) 109 | x = F.relu(self.bn1(self.conv1(x))) 110 | pointfeat = x 111 | x = F.relu(self.bn2(self.conv2(x))) 112 | x = self.bn3(self.conv3(x)) 113 | x = self.mp1(x) 114 | x = x.view(-1, 1024) 115 | if self.global_feat: 116 | return x, trans1 + trans2 117 | else: 118 | x = x.view(-1, 1024, 1).repeat(1, 1, self.num_points) 119 | return torch.cat([x, pointfeat], 1), trans1 + trans2 120 | 121 | 122 | class PointNetfeat(nn.Module): 123 | def __init__(self, num_points = 2500, input_chann = 3, global_feat = True): 124 | super(PointNetfeat, self).__init__() 125 | self.stn = STN3d(num_points = num_points, input_chann = input_chann) 126 | self.conv1 = torch.nn.Conv1d(input_chann, 64, 1) 127 | self.conv2 = torch.nn.Conv1d(64, 128, 1) 128 | self.conv3 = torch.nn.Conv1d(128, 1024, 1) 129 | self.bn1 = nn.BatchNorm1d(64) 130 | self.bn2 = nn.BatchNorm1d(128) 131 | self.bn3 = nn.BatchNorm1d(1024) 132 | self.mp1 = torch.nn.MaxPool1d(num_points) 133 | self.num_points = num_points 134 | self.global_feat = global_feat 135 | def forward(self, x): 136 | batchsize = x.size()[0] 137 | trans = self.stn(x) 138 | x = x.transpose(2,1) 139 | x = torch.bmm(x, trans) 140 | # x = torch.cat([torch.bmm(x[..., 0:3], trans), torch.bmm(x[..., 3:6], trans)], dim=-1) 141 | x = x.transpose(2,1) 142 | x = F.relu(self.bn1(self.conv1(x))) 143 | pointfeat = x 144 | x = F.relu(self.bn2(self.conv2(x))) 145 | x = self.bn3(self.conv3(x)) 146 | x = self.mp1(x) 147 | x = x.view(-1, 1024) 148 | if self.global_feat: 149 | return x, trans 150 | else: 151 | x = x.view(-1, 1024, 1).repeat(1, 1, self.num_points) 152 | return torch.cat([x, pointfeat], 1), trans 153 | 154 | 155 | class DualPointNetCls(nn.Module): 156 | def __init__(self, num_points = 2500, input_chann = 3, k = 2): 157 | super(DualPointNetCls, self).__init__() 158 | self.num_points = num_points 159 | self.feat = DualPointNetfeat(num_points, input_chann=input_chann, global_feat=True) 160 | self.fc1 = nn.Linear(1024, 512) 161 | self.fc2 = nn.Linear(512, 256) 162 | self.fc3 = nn.Linear(256, k) 163 | self.bn1 = nn.BatchNorm1d(512) 164 | self.bn2 = nn.BatchNorm1d(256) 165 | self.relu = nn.ReLU() 166 | def forward(self, x): 167 | x, trans = self.feat(x) 168 | x = F.relu(self.bn1(self.fc1(x))) 169 | x = F.relu(self.bn2(self.fc2(x))) 170 | x = self.fc3(x) 171 | return F.log_softmax(x, dim=-1), trans 172 | 173 | 174 | class PointNetCls(nn.Module): 175 | def __init__(self, num_points = 2500, input_chann = 3, k = 2): 176 | super(PointNetCls, self).__init__() 177 | self.num_points = num_points 178 | self.feat = PointNetfeat(num_points, input_chann=input_chann, global_feat=True) 179 | self.fc1 = nn.Linear(1024, 512) 180 | self.fc2 = nn.Linear(512, 256) 181 | self.fc3 = nn.Linear(256, k) 182 | self.bn1 = nn.BatchNorm1d(512) 183 | self.bn2 = nn.BatchNorm1d(256) 184 | self.relu = nn.ReLU() 185 | def forward(self, x): 186 | x, trans = self.feat(x) 187 | x = F.relu(self.bn1(self.fc1(x))) 188 | x = F.relu(self.bn2(self.fc2(x))) 189 | x = self.fc3(x) 190 | return F.log_softmax(x, dim=-1), trans 191 | 192 | class PointNetDenseCls(nn.Module): 193 | def __init__(self, num_points = 2500, input_chann = 3, k = 2): 194 | super(PointNetDenseCls, self).__init__() 195 | self.num_points = num_points 196 | self.k = k 197 | self.feat = PointNetfeat(num_points, input_chann=input_chann, global_feat=False) 198 | self.conv1 = torch.nn.Conv1d(1088, 512, 1) 199 | self.conv2 = torch.nn.Conv1d(512, 256, 1) 200 | self.conv3 = torch.nn.Conv1d(256, 128, 1) 201 | self.conv4 = torch.nn.Conv1d(128, self.k, 1) 202 | self.bn1 = nn.BatchNorm1d(512) 203 | self.bn2 = nn.BatchNorm1d(256) 204 | self.bn3 = nn.BatchNorm1d(128) 205 | 206 | def forward(self, x): 207 | batchsize = x.size()[0] 208 | x, trans = self.feat(x) 209 | x = F.relu(self.bn1(self.conv1(x))) 210 | x = F.relu(self.bn2(self.conv2(x))) 211 | x = F.relu(self.bn3(self.conv3(x))) 212 | x = self.conv4(x) 213 | x = x.transpose(2,1).contiguous() 214 | x = F.log_softmax(x.view(-1,self.k), dim=-1) 215 | x = x.view(batchsize, self.num_points, self.k) 216 | return x, trans 217 | 218 | 219 | if __name__ == '__main__': 220 | pass 221 | -------------------------------------------------------------------------------- /requirements.txt: -------------------------------------------------------------------------------- 1 | networkx==2.2 2 | scipy 3 | h5py 4 | numpy 5 | NearPy 6 | Pillow 7 | python-dateutil 8 | PyYAML 9 | scikit-image 10 | scikit-learn 11 | six 12 | trimesh 13 | pyhull 14 | cvxopt 15 | pandas 16 | visualization 17 | pyquaternion 18 | numba 19 | autolab-core 20 | autolab-perception 21 | shapely 22 | 23 | -------------------------------------------------------------------------------- /requirements2.txt: -------------------------------------------------------------------------------- 1 | opencv-python==4.2.0.32 2 | trimesh 3 | ffmpeg-python==0.1.5 4 | networkx==2.2 5 | scipy 6 | h5py 7 | numpy 8 | NearPy 9 | Pillow 10 | python-dateutil 11 | PyYAML 12 | scikit-image 13 | scikit-learn 14 | six 15 | trimesh 16 | pyhull 17 | cvxopt 18 | pandas 19 | pyquaternion 20 | numba 21 | autolab-core 22 | autolab-perception 23 | shapely 24 | --------------------------------------------------------------------------------