├── .gitattributes ├── .gitignore ├── LICENSE.txt ├── MANIFEST.in ├── NVDUIntro.png ├── nvdu ├── __init__.py ├── config │ └── object_settings │ │ ├── _ycb_aligned_cm.json │ │ ├── _ycb_aligned_m.json │ │ └── _ycb_original.json ├── core │ ├── __init__.py │ ├── box.py │ ├── camera.py │ ├── cuboid.py │ ├── mesh.py │ ├── nvdu_data.py │ ├── pivot_axis.py │ ├── scene_object.py │ ├── transform3d.py │ └── utils3d.py ├── tools │ ├── __init__.py │ ├── nvdu_ycb.py │ └── test_nvdu_visualizer.py └── viz │ ├── __init__.py │ ├── background_image.py │ ├── camera.py │ ├── cuboid.py │ ├── image_draw.py │ ├── mesh.py │ ├── nvdu_visualizer.py │ ├── nvdu_viz_window.py │ ├── pivot_axis.py │ ├── pointcloud.py │ ├── scene.py │ ├── scene_object.py │ ├── utils3d.py │ └── viewport.py ├── readme.md ├── setup.cfg └── setup.py /.gitattributes: -------------------------------------------------------------------------------- 1 | *.uasset filter=lfs diff=lfs merge=lfs -text 2 | *.umap filter=lfs diff=lfs merge=lfs -text 3 | *.fbx filter=lfs diff=lfs merge=lfs -text 4 | *.stl filter=lfs diff=lfs merge=lfs -text 5 | *.obj filter=lfs diff=lfs merge=lfs -text 6 | *.wrl filter=lfs diff=lfs merge=lfs -text 7 | *.a filter=lfs diff=lfs merge=lfs -text 8 | *.lib filter=lfs diff=lfs merge=lfs -text 9 | *.so filter=lfs diff=lfs merge=lfs -text 10 | *.dll filter=lfs diff=lfs merge=lfs -text 11 | *.exe filter=lfs diff=lfs merge=lfs -text 12 | *.pdb filter=lfs diff=lfs merge=lfs -text 13 | *.mp4 filter=lfs diff=lfs merge=lfs -text 14 | *.pdf filter=lfs diff=lfs merge=lfs -text 15 | -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | # See http://help.github.com/ignore-files/ for more about ignoring files. 2 | 3 | # compiled output 4 | /dist 5 | /tmp 6 | /out-tsc 7 | 8 | # dependencies 9 | /node_modules 10 | /.mypy_cache 11 | /__pycache__ 12 | /*/__pycache__ 13 | /virtual_env 14 | *.pyc 15 | 16 | # data 17 | /data 18 | /output 19 | /content 20 | 21 | # IDEs and editors 22 | /.idea 23 | .project 24 | .classpath 25 | .c9/ 26 | *.launch 27 | .settings/ 28 | *.sublime-workspace 29 | /.vscode 30 | 31 | # IDE - VSCode 32 | .vscode/* 33 | !.vscode/settings.json 34 | !.vscode/tasks.json 35 | !.vscode/launch.json 36 | !.vscode/extensions.json 37 | 38 | # misc 39 | /.sass-cache 40 | /connect.lock 41 | /coverage 42 | /libpeerconnection.log 43 | npm-debug.log 44 | testem.log 45 | /typings 46 | 47 | # e2e 48 | /e2e/*.js 49 | /e2e/*.map 50 | 51 | # System Files 52 | .DS_Store 53 | Thumbs.db 54 | 55 | *.egg-info 56 | nvdu/data/ 57 | -------------------------------------------------------------------------------- /LICENSE.txt: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NVIDIA/Dataset_Utilities/532b8c76e3d7946748a10af3398438b35383f157/LICENSE.txt -------------------------------------------------------------------------------- /MANIFEST.in: -------------------------------------------------------------------------------- 1 | recursive-include nvdu/data/ycb * -------------------------------------------------------------------------------- /NVDUIntro.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NVIDIA/Dataset_Utilities/532b8c76e3d7946748a10af3398438b35383f157/NVDUIntro.png -------------------------------------------------------------------------------- /nvdu/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NVIDIA/Dataset_Utilities/532b8c76e3d7946748a10af3398438b35383f157/nvdu/__init__.py -------------------------------------------------------------------------------- /nvdu/config/object_settings/_ycb_aligned_cm.json: -------------------------------------------------------------------------------- 1 | { 2 | "exported_object_classes": [ 3 | "002_master_chef_can", 4 | "003_cracker_box", 5 | "004_sugar_box", 6 | "005_tomato_soup_can", 7 | "006_mustard_bottle", 8 | "007_tuna_fish_can", 9 | "008_pudding_box", 10 | "009_gelatin_box", 11 | "010_potted_meat_can", 12 | "011_banana", 13 | "019_pitcher_base", 14 | "021_bleach_cleanser", 15 | "024_bowl", 16 | "025_mug", 17 | "035_power_drill", 18 | "036_wood_block", 19 | "037_scissors", 20 | "040_large_marker", 21 | "051_large_clamp", 22 | "052_extra_large_clamp", 23 | "061_foam_brick" 24 | ], 25 | "exported_objects": [ 26 | { 27 | "class": "002_master_chef_can", 28 | "segmentation_class_id": 12, 29 | "fixed_model_transform": [ 30 | [ 1, 0, 0, 0 ], 31 | [ 0, 1, 0, 0 ], 32 | [ 0, 0, 1, 0 ], 33 | [ 0, 0, 0, 1 ] 34 | ], 35 | "cuboid_dimensions": [ 10.240400314331055, 14.0177001953125, 10.230899810791016 ] 36 | }, 37 | { 38 | "class": "003_cracker_box", 39 | "segmentation_class_id": 24, 40 | "fixed_model_transform": [ 41 | [ 1, 0, 0, 0 ], 42 | [ 0, 1, 0, 0 ], 43 | [ 0, 0, 1, 0 ], 44 | [ 0, 0, 0, 1 ] 45 | ], 46 | "cuboid_dimensions": [ 16.403600692749023, 21.343700408935547, 7.179999828338623 ] 47 | }, 48 | { 49 | "class": "004_sugar_box", 50 | "segmentation_class_id": 36, 51 | "fixed_model_transform": [ 52 | [ 1, 0, 0, 0 ], 53 | [ 0, 1, 0, 0 ], 54 | [ 0, 0, 1, 0 ], 55 | [ 0, 0, 0, 1 ] 56 | ], 57 | "cuboid_dimensions": [ 9.2677001953125, 17.625299453735352, 4.5134000778198242 ] 58 | }, 59 | { 60 | "class": "005_tomato_soup_can", 61 | "segmentation_class_id": 48, 62 | "fixed_model_transform": [ 63 | [ 1, 0, 0, 0 ], 64 | [ 0, 1, 0, 0 ], 65 | [ 0, 0, 1, 0 ], 66 | [ 0, 0, 0, 1 ] 67 | ], 68 | "cuboid_dimensions": [ 6.7659001350402832, 10.185500144958496, 6.771399974822998 ] 69 | }, 70 | { 71 | "class": "006_mustard_bottle", 72 | "segmentation_class_id": 60, 73 | "fixed_model_transform": [ 74 | [ 1, 0, 0, 0 ], 75 | [ 0, 1, 0, 0 ], 76 | [ 0, 0, 1, 0 ], 77 | [ 0, 0, 0, 1 ] 78 | ], 79 | "cuboid_dimensions": [ 9.6023998260498047, 19.130100250244141, 5.8249001502990723 ] 80 | }, 81 | { 82 | "class": "007_tuna_fish_can", 83 | "segmentation_class_id": 72, 84 | "fixed_model_transform": [ 85 | [ 1, 0, 0, 0 ], 86 | [ 0, 1, 0, 0 ], 87 | [ 0, 0, 1, 0 ], 88 | [ 0, 0, 0, 1 ] 89 | ], 90 | "cuboid_dimensions": [ 8.5558996200561523, 3.3538000583648682, 8.5539999008178711 ] 91 | }, 92 | { 93 | "class": "008_pudding_box", 94 | "segmentation_class_id": 84, 95 | "fixed_model_transform": [ 96 | [ 1, 0, 0, 0 ], 97 | [ 0, 1, 0, 0 ], 98 | [ 0, 0, 1, 0 ], 99 | [ 0, 0, 0, 1 ] 100 | ], 101 | "cuboid_dimensions": [ 11.365400314331055, 8.9816999435424805, 3.8471000194549561 ] 102 | }, 103 | { 104 | "class": "009_gelatin_box", 105 | "segmentation_class_id": 96, 106 | "fixed_model_transform": [ 107 | [ 1, 0, 0, 0 ], 108 | [ 0, 1, 0, 0 ], 109 | [ 0, 0, 1, 0 ], 110 | [ 0, 0, 0, 1 ] 111 | ], 112 | "cuboid_dimensions": [ 8.918299674987793, 7.311500072479248, 2.9983000755310059 ] 113 | }, 114 | { 115 | "class": "010_potted_meat_can", 116 | "segmentation_class_id": 108, 117 | "fixed_model_transform": [ 118 | [ 1, 0, 0, 0 ], 119 | [ 0, 1, 0, 0 ], 120 | [ 0, 0, 1, 0 ], 121 | [ 0, 0, 0, 1 ] 122 | ], 123 | "cuboid_dimensions": [ 10.164699554443359, 8.3542995452880859, 5.7600998878479004 ] 124 | }, 125 | { 126 | "class": "011_banana", 127 | "segmentation_class_id": 120, 128 | "fixed_model_transform": [ 129 | [ 1, 0, 0, 0 ], 130 | [ 0, 1, 0, 0 ], 131 | [ 0, 0, 1, 0 ], 132 | [ 0, 0, 0, 1 ] 133 | ], 134 | "cuboid_dimensions": [ 19.717399597167969, 3.8649001121520996, 7.4065999984741211 ] 135 | }, 136 | { 137 | "class": "019_pitcher_base", 138 | "segmentation_class_id": 132, 139 | "fixed_model_transform": [ 140 | [ 1, 0, 0, 0 ], 141 | [ 0, 1, 0, 0 ], 142 | [ 0, 0, 1, 0 ], 143 | [ 0, 0, 0, 1 ] 144 | ], 145 | "cuboid_dimensions": [ 18.861799240112305, 24.238700866699219, 13.310600280761719 ] 146 | }, 147 | { 148 | "class": "021_bleach_cleanser", 149 | "segmentation_class_id": 144, 150 | "fixed_model_transform": [ 151 | [ 1, 0, 0, 0 ], 152 | [ 0, 1, 0, 0 ], 153 | [ 0, 0, 1, 0 ], 154 | [ 0, 0, 0, 1 ] 155 | ], 156 | "cuboid_dimensions": [ 10.243300437927246, 25.058000564575195, 6.769899845123291 ] 157 | }, 158 | { 159 | "class": "024_bowl", 160 | "segmentation_class_id": 156, 161 | "fixed_model_transform": [ 162 | [ 1, 0, 0, 0 ], 163 | [ 0, 1, 0, 0 ], 164 | [ 0, 0, 1, 0 ], 165 | [ 0, 0, 0, 1 ] 166 | ], 167 | "cuboid_dimensions": [ 16.116399765014648, 5.5008997917175293, 16.146299362182617 ] 168 | }, 169 | { 170 | "class": "025_mug", 171 | "segmentation_class_id": 168, 172 | "fixed_model_transform": [ 173 | [ 1, 0, 0, 0 ], 174 | [ 0, 1, 0, 0 ], 175 | [ 0, 0, 1, 0 ], 176 | [ 0, 0, 0, 1 ] 177 | ], 178 | "cuboid_dimensions": [ 11.699000358581543, 8.1302995681762695, 9.3087997436523438 ] 179 | }, 180 | { 181 | "class": "035_power_drill", 182 | "segmentation_class_id": 180, 183 | "fixed_model_transform": [ 184 | [ 1, 0, 0, 0 ], 185 | [ 0, 1, 0, 0 ], 186 | [ 0, 0, 1, 0 ], 187 | [ 0, 0, 0, 1 ] 188 | ], 189 | "cuboid_dimensions": [ 18.436100006103516, 18.683599472045898, 5.7192001342773438 ] 190 | }, 191 | { 192 | "class": "036_wood_block", 193 | "segmentation_class_id": 192, 194 | "fixed_model_transform": [ 195 | [ 1, 0, 0, 0 ], 196 | [ 0, 1, 0, 0 ], 197 | [ 0, 0, 1, 0 ], 198 | [ 0, 0, 0, 1 ] 199 | ], 200 | "cuboid_dimensions": [ 8.969599723815918, 20.609199523925781, 9.0780000686645508 ] 201 | }, 202 | { 203 | "class": "037_scissors", 204 | "segmentation_class_id": 204, 205 | "fixed_model_transform": [ 206 | [ 1, 0, 0, 0 ], 207 | [ 0, 1, 0, 0 ], 208 | [ 0, 0, 1, 0 ], 209 | [ 0, 0, 0, 1 ] 210 | ], 211 | "cuboid_dimensions": [ 20.259799957275391, 8.7530002593994141, 1.565500020980835 ] 212 | }, 213 | { 214 | "class": "040_large_marker", 215 | "segmentation_class_id": 216, 216 | "fixed_model_transform": [ 217 | [ 1, 0, 0, 0 ], 218 | [ 0, 1, 0, 0 ], 219 | [ 0, 0, 1, 0 ], 220 | [ 0, 0, 0, 1 ] 221 | ], 222 | "cuboid_dimensions": [ 1.8833999633789063, 12.088600158691406, 1.9476000070571899 ] 223 | }, 224 | { 225 | "class": "051_large_clamp", 226 | "segmentation_class_id": 228, 227 | "fixed_model_transform": [ 228 | [ 1, 0, 0, 0 ], 229 | [ 0, 1, 0, 0 ], 230 | [ 0, 0, 1, 0 ], 231 | [ 0, 0, 0, 1 ] 232 | ], 233 | "cuboid_dimensions": [ 16.494100570678711, 12.166600227355957, 3.6412999629974365 ] 234 | }, 235 | { 236 | "class": "052_extra_large_clamp", 237 | "segmentation_class_id": 240, 238 | "fixed_model_transform": [ 239 | [ 1, 0, 0, 0 ], 240 | [ 0, 1, 0, 0 ], 241 | [ 0, 0, 1, 0 ], 242 | [ 0, 0, 0, 1 ] 243 | ], 244 | "cuboid_dimensions": [ 20.259599685668945, 16.520299911499023, 3.6475999355316162 ] 245 | }, 246 | { 247 | "class": "061_foam_brick", 248 | "segmentation_class_id": 252, 249 | "fixed_model_transform": [ 250 | [ 1, 0, 0, 0 ], 251 | [ 0, 1, 0, 0 ], 252 | [ 0, 0, 1, 0 ], 253 | [ 0, 0, 0, 1 ] 254 | ], 255 | "cuboid_dimensions": [ 7.7873997688293457, 5.1192998886108398, 5.2560000419616699 ] 256 | } 257 | ] 258 | } -------------------------------------------------------------------------------- /nvdu/config/object_settings/_ycb_aligned_m.json: -------------------------------------------------------------------------------- 1 | { 2 | "exported_object_classes": [ 3 | "002_master_chef_can", 4 | "003_cracker_box", 5 | "004_sugar_box", 6 | "005_tomato_soup_can", 7 | "006_mustard_bottle", 8 | "007_tuna_fish_can", 9 | "008_pudding_box", 10 | "009_gelatin_box", 11 | "010_potted_meat_can", 12 | "011_banana", 13 | "019_pitcher_base", 14 | "021_bleach_cleanser", 15 | "024_bowl", 16 | "025_mug", 17 | "035_power_drill", 18 | "036_wood_block", 19 | "037_scissors", 20 | "040_large_marker", 21 | "051_large_clamp", 22 | "052_extra_large_clamp", 23 | "061_foam_brick" 24 | ], 25 | "exported_objects": [ 26 | { 27 | "class": "002_master_chef_can", 28 | "segmentation_class_id": 12, 29 | "fixed_model_transform": [ 30 | [ 1, 0, 0, 0 ], 31 | [ 0, 1, 0, 0 ], 32 | [ 0, 0, 1, 0 ], 33 | [ 0, 0, 0, 1 ] 34 | ], 35 | "cuboid_dimensions": [ 10.240400314331055, 14.0177001953125, 10.230899810791016 ] 36 | }, 37 | { 38 | "class": "003_cracker_box", 39 | "segmentation_class_id": 24, 40 | "fixed_model_transform": [ 41 | [ 1, 0, 0, 0 ], 42 | [ 0, 1, 0, 0 ], 43 | [ 0, 0, 1, 0 ], 44 | [ 0, 0, 0, 1 ] 45 | ], 46 | "cuboid_dimensions": [ 16.403600692749023, 21.343700408935547, 7.179999828338623 ] 47 | }, 48 | { 49 | "class": "004_sugar_box", 50 | "segmentation_class_id": 36, 51 | "fixed_model_transform": [ 52 | [ 1, 0, 0, 0 ], 53 | [ 0, 1, 0, 0 ], 54 | [ 0, 0, 1, 0 ], 55 | [ 0, 0, 0, 1 ] 56 | ], 57 | "cuboid_dimensions": [ 9.2677001953125, 17.625299453735352, 4.5134000778198242 ] 58 | }, 59 | { 60 | "class": "005_tomato_soup_can", 61 | "segmentation_class_id": 48, 62 | "fixed_model_transform": [ 63 | [ 1, 0, 0, 0 ], 64 | [ 0, 1, 0, 0 ], 65 | [ 0, 0, 1, 0 ], 66 | [ 0, 0, 0, 1 ] 67 | ], 68 | "cuboid_dimensions": [ 6.7659001350402832, 10.185500144958496, 6.771399974822998 ] 69 | }, 70 | { 71 | "class": "006_mustard_bottle", 72 | "segmentation_class_id": 60, 73 | "fixed_model_transform": [ 74 | [ 1, 0, 0, 0 ], 75 | [ 0, 1, 0, 0 ], 76 | [ 0, 0, 1, 0 ], 77 | [ 0, 0, 0, 1 ] 78 | ], 79 | "cuboid_dimensions": [ 9.6023998260498047, 19.130100250244141, 5.8249001502990723 ] 80 | }, 81 | { 82 | "class": "007_tuna_fish_can", 83 | "segmentation_class_id": 72, 84 | "fixed_model_transform": [ 85 | [ 1, 0, 0, 0 ], 86 | [ 0, 1, 0, 0 ], 87 | [ 0, 0, 1, 0 ], 88 | [ 0, 0, 0, 1 ] 89 | ], 90 | "cuboid_dimensions": [ 8.5558996200561523, 3.3538000583648682, 8.5539999008178711 ] 91 | }, 92 | { 93 | "class": "008_pudding_box", 94 | "segmentation_class_id": 84, 95 | "fixed_model_transform": [ 96 | [ 1, 0, 0, 0 ], 97 | [ 0, 1, 0, 0 ], 98 | [ 0, 0, 1, 0 ], 99 | [ 0, 0, 0, 1 ] 100 | ], 101 | "cuboid_dimensions": [ 11.365400314331055, 8.9816999435424805, 3.8471000194549561 ] 102 | }, 103 | { 104 | "class": "009_gelatin_box", 105 | "segmentation_class_id": 96, 106 | "fixed_model_transform": [ 107 | [ 1, 0, 0, 0 ], 108 | [ 0, 1, 0, 0 ], 109 | [ 0, 0, 1, 0 ], 110 | [ 0, 0, 0, 1 ] 111 | ], 112 | "cuboid_dimensions": [ 8.918299674987793, 7.311500072479248, 2.9983000755310059 ] 113 | }, 114 | { 115 | "class": "010_potted_meat_can", 116 | "segmentation_class_id": 108, 117 | "fixed_model_transform": [ 118 | [ 1, 0, 0, 0 ], 119 | [ 0, 1, 0, 0 ], 120 | [ 0, 0, 1, 0 ], 121 | [ 0, 0, 0, 1 ] 122 | ], 123 | "cuboid_dimensions": [ 10.164699554443359, 8.3542995452880859, 5.7600998878479004 ] 124 | }, 125 | { 126 | "class": "011_banana", 127 | "segmentation_class_id": 120, 128 | "fixed_model_transform": [ 129 | [ 1, 0, 0, 0 ], 130 | [ 0, 1, 0, 0 ], 131 | [ 0, 0, 1, 0 ], 132 | [ 0, 0, 0, 1 ] 133 | ], 134 | "cuboid_dimensions": [ 19.717399597167969, 3.8649001121520996, 7.4065999984741211 ] 135 | }, 136 | { 137 | "class": "019_pitcher_base", 138 | "segmentation_class_id": 132, 139 | "fixed_model_transform": [ 140 | [ 1, 0, 0, 0 ], 141 | [ 0, 1, 0, 0 ], 142 | [ 0, 0, 1, 0 ], 143 | [ 0, 0, 0, 1 ] 144 | ], 145 | "cuboid_dimensions": [ 18.861799240112305, 24.238700866699219, 13.310600280761719 ] 146 | }, 147 | { 148 | "class": "021_bleach_cleanser", 149 | "segmentation_class_id": 144, 150 | "fixed_model_transform": [ 151 | [ 1, 0, 0, 0 ], 152 | [ 0, 1, 0, 0 ], 153 | [ 0, 0, 1, 0 ], 154 | [ 0, 0, 0, 1 ] 155 | ], 156 | "cuboid_dimensions": [ 10.243300437927246, 25.058000564575195, 6.769899845123291 ] 157 | }, 158 | { 159 | "class": "024_bowl", 160 | "segmentation_class_id": 156, 161 | "fixed_model_transform": [ 162 | [ 1, 0, 0, 0 ], 163 | [ 0, 1, 0, 0 ], 164 | [ 0, 0, 1, 0 ], 165 | [ 0, 0, 0, 1 ] 166 | ], 167 | "cuboid_dimensions": [ 16.116399765014648, 5.5008997917175293, 16.146299362182617 ] 168 | }, 169 | { 170 | "class": "025_mug", 171 | "segmentation_class_id": 168, 172 | "fixed_model_transform": [ 173 | [ 1, 0, 0, 0 ], 174 | [ 0, 1, 0, 0 ], 175 | [ 0, 0, 1, 0 ], 176 | [ 0, 0, 0, 1 ] 177 | ], 178 | "cuboid_dimensions": [ 11.699000358581543, 8.1302995681762695, 9.3087997436523438 ] 179 | }, 180 | { 181 | "class": "035_power_drill", 182 | "segmentation_class_id": 180, 183 | "fixed_model_transform": [ 184 | [ 1, 0, 0, 0 ], 185 | [ 0, 1, 0, 0 ], 186 | [ 0, 0, 1, 0 ], 187 | [ 0, 0, 0, 1 ] 188 | ], 189 | "cuboid_dimensions": [ 18.436100006103516, 18.683599472045898, 5.7192001342773438 ] 190 | }, 191 | { 192 | "class": "036_wood_block", 193 | "segmentation_class_id": 192, 194 | "fixed_model_transform": [ 195 | [ 1, 0, 0, 0 ], 196 | [ 0, 1, 0, 0 ], 197 | [ 0, 0, 1, 0 ], 198 | [ 0, 0, 0, 1 ] 199 | ], 200 | "cuboid_dimensions": [ 8.969599723815918, 20.609199523925781, 9.0780000686645508 ] 201 | }, 202 | { 203 | "class": "037_scissors", 204 | "segmentation_class_id": 204, 205 | "fixed_model_transform": [ 206 | [ 1, 0, 0, 0 ], 207 | [ 0, 1, 0, 0 ], 208 | [ 0, 0, 1, 0 ], 209 | [ 0, 0, 0, 1 ] 210 | ], 211 | "cuboid_dimensions": [ 20.259799957275391, 8.7530002593994141, 1.565500020980835 ] 212 | }, 213 | { 214 | "class": "040_large_marker", 215 | "segmentation_class_id": 216, 216 | "fixed_model_transform": [ 217 | [ 1, 0, 0, 0 ], 218 | [ 0, 1, 0, 0 ], 219 | [ 0, 0, 1, 0 ], 220 | [ 0, 0, 0, 1 ] 221 | ], 222 | "cuboid_dimensions": [ 1.8833999633789063, 12.088600158691406, 1.9476000070571899 ] 223 | }, 224 | { 225 | "class": "051_large_clamp", 226 | "segmentation_class_id": 228, 227 | "fixed_model_transform": [ 228 | [ 1, 0, 0, 0 ], 229 | [ 0, 1, 0, 0 ], 230 | [ 0, 0, 1, 0 ], 231 | [ 0, 0, 0, 1 ] 232 | ], 233 | "cuboid_dimensions": [ 16.494100570678711, 12.166600227355957, 3.6412999629974365 ] 234 | }, 235 | { 236 | "class": "052_extra_large_clamp", 237 | "segmentation_class_id": 240, 238 | "fixed_model_transform": [ 239 | [ 1, 0, 0, 0 ], 240 | [ 0, 1, 0, 0 ], 241 | [ 0, 0, 1, 0 ], 242 | [ 0, 0, 0, 1 ] 243 | ], 244 | "cuboid_dimensions": [ 20.259599685668945, 16.520299911499023, 3.6475999355316162 ] 245 | }, 246 | { 247 | "class": "061_foam_brick", 248 | "segmentation_class_id": 252, 249 | "fixed_model_transform": [ 250 | [ 1, 0, 0, 0 ], 251 | [ 0, 1, 0, 0 ], 252 | [ 0, 0, 1, 0 ], 253 | [ 0, 0, 0, 1 ] 254 | ], 255 | "cuboid_dimensions": [ 7.7873997688293457, 5.1192998886108398, 5.2560000419616699 ] 256 | } 257 | ] 258 | } -------------------------------------------------------------------------------- /nvdu/config/object_settings/_ycb_original.json: -------------------------------------------------------------------------------- 1 | { 2 | "exported_object_classes": [ 3 | "002_master_chef_can_16k", 4 | "003_cracker_box_16k", 5 | "004_sugar_box_16k", 6 | "005_tomato_soup_can_16k", 7 | "006_mustard_bottle_16k", 8 | "007_tuna_fish_can_16k", 9 | "008_pudding_box_16k", 10 | "009_gelatin_box_16k", 11 | "010_potted_meat_can_16k", 12 | "011_banana_16k", 13 | "019_pitcher_base_16k", 14 | "021_bleach_cleanser_16k", 15 | "024_bowl_16k", 16 | "025_mug_16k", 17 | "035_power_drill_16k", 18 | "036_wood_block_16k", 19 | "037_scissors_16k", 20 | "040_large_marker_16k", 21 | "051_large_clamp_16k", 22 | "052_extra_large_clamp_16k", 23 | "061_foam_brick_16k" 24 | ], 25 | "exported_objects": [ 26 | { 27 | "class": "002_master_chef_can_16k", 28 | "segmentation_class_id": 12, 29 | "fixed_model_transform": [ 30 | [ 86.602500915527344, 0, 50, 0 ], 31 | [ -50, 0, 86.602500915527344, 0 ], 32 | [ 0, -100, 0, 0 ], 33 | [ 0.99080002307891846, 6.9902000427246094, 1.6902999877929688, 1 ] 34 | ], 35 | "cuboid_dimensions": [ 10.240400314331055, 14.0177001953125, 10.230899810791016 ] 36 | }, 37 | { 38 | "class": "003_cracker_box_16k", 39 | "segmentation_class_id": 24, 40 | "fixed_model_transform": [ 41 | [ 0, 0, 100, 0 ], 42 | [ -100, 0, 0, 0 ], 43 | [ 0, -100, 0, 0 ], 44 | [ -1.4141999483108521, 10.347499847412109, 1.2884999513626099, 1 ] 45 | ], 46 | "cuboid_dimensions": [ 16.403600692749023, 21.343700408935547, 7.179999828338623 ] 47 | }, 48 | { 49 | "class": "004_sugar_box_16k", 50 | "segmentation_class_id": 36, 51 | "fixed_model_transform": [ 52 | [ -3.4877998828887939, 3.4899001121520996, 99.878196716308594, 0 ], 53 | [ -99.926002502441406, -1.7441999912261963, -3.4284999370574951, 0 ], 54 | [ 1.6224000453948975, -99.923896789550781, 3.5481998920440674, 0 ], 55 | [ -1.795199990272522, 8.7579002380371094, 0.38839998841285706, 1 ] 56 | ], 57 | "cuboid_dimensions": [ 9.2677001953125, 17.625299453735352, 4.5134000778198242 ] 58 | }, 59 | { 60 | "class": "005_tomato_soup_can_16k", 61 | "segmentation_class_id": 48, 62 | "fixed_model_transform": [ 63 | [ 99.144500732421875, 0, -13.052599906921387, 0 ], 64 | [ 13.052599906921387, 0, 99.144500732421875, 0 ], 65 | [ 0, -100, 0, 0 ], 66 | [ -0.1793999969959259, 5.1006999015808105, -8.4443998336791992, 1 ] 67 | ], 68 | "cuboid_dimensions": [ 6.7659001350402832, 10.185500144958496, 6.771399974822998 ] 69 | }, 70 | { 71 | "class": "006_mustard_bottle_16k", 72 | "segmentation_class_id": 60, 73 | "fixed_model_transform": [ 74 | [ 92.050498962402344, 0, 39.073101043701172, 0 ], 75 | [ -39.073101043701172, 0, 92.050498962402344, 0 ], 76 | [ 0, -100, 0, 0 ], 77 | [ 0.49259999394416809, 9.2497997283935547, 2.7135999202728271, 1 ] 78 | ], 79 | "cuboid_dimensions": [ 9.6023998260498047, 19.130100250244141, 5.8249001502990723 ] 80 | }, 81 | { 82 | "class": "007_tuna_fish_can_16k", 83 | "segmentation_class_id": 72, 84 | "fixed_model_transform": [ 85 | [ 100, 0, 0, 0 ], 86 | [ 0, 0, 100, 0 ], 87 | [ 0, -100, 0, 0 ], 88 | [ 2.6048998832702637, 1.3551000356674194, 2.2132000923156738, 1 ] 89 | ], 90 | "cuboid_dimensions": [ 8.5558996200561523, 3.3538000583648682, 8.5539999008178711 ] 91 | }, 92 | { 93 | "class": "008_pudding_box_16k", 94 | "segmentation_class_id": 84, 95 | "fixed_model_transform": [ 96 | [ -88.264503479003906, 46.947200775146484, -2.3113000392913818, 0 ], 97 | [ -46.878200531005859, -88.281303405761719, -2.9734001159667969, 0 ], 98 | [ -3.4363000392913818, -1.5410000085830688, 99.929100036621094, 0 ], 99 | [ 1.010200023651123, 1.6993999481201172, -1.7572000026702881, 1 ] 100 | ], 101 | "cuboid_dimensions": [ 11.365400314331055, 8.9816999435424805, 3.8471000194549561 ] 102 | }, 103 | { 104 | "class": "009_gelatin_box_16k", 105 | "segmentation_class_id": 96, 106 | "fixed_model_transform": [ 107 | [ 22.494199752807617, 97.436996459960938, 0.19629999995231628, 0 ], 108 | [ -97.433296203613281, 22.495100021362305, -0.85030001401901245, 0 ], 109 | [ -0.87269997596740723, 0, 99.996200561523438, 0 ], 110 | [ -0.29069998860359192, 2.3998000621795654, -1.4543999433517456, 1 ] 111 | ], 112 | "cuboid_dimensions": [ 8.918299674987793, 7.311500072479248, 2.9983000755310059 ] 113 | }, 114 | { 115 | "class": "010_potted_meat_can_16k", 116 | "segmentation_class_id": 108, 117 | "fixed_model_transform": [ 118 | [ 99.862998962402344, 0, -5.2336001396179199, 0 ], 119 | [ 5.2336001396179199, 0, 99.862998962402344, 0 ], 120 | [ 0, -100, 0, 0 ], 121 | [ 3.4065999984741211, 3.8582999706268311, 2.4767000675201416, 1 ] 122 | ], 123 | "cuboid_dimensions": [ 10.164699554443359, 8.3542995452880859, 5.7600998878479004 ] 124 | }, 125 | { 126 | "class": "011_banana_16k", 127 | "segmentation_class_id": 120, 128 | "fixed_model_transform": [ 129 | [ -36.395401000976563, -17.364799499511719, 91.508697509765625, 0 ], 130 | [ -93.087699890136719, 3.4368999004364014, -36.371200561523438, 0 ], 131 | [ 3.1707000732421875, -98.420799255371094, -17.415399551391602, 0 ], 132 | [ 0.036600001156330109, 1.497499942779541, 0.44449999928474426, 1 ] 133 | ], 134 | "cuboid_dimensions": [ 19.717399597167969, 3.8649001121520996, 7.4065999984741211 ] 135 | }, 136 | { 137 | "class": "019_pitcher_base_16k", 138 | "segmentation_class_id": 132, 139 | "fixed_model_transform": [ 140 | [ 71.933998107910156, 0, -69.465797424316406, 0 ], 141 | [ 69.465797424316406, 0, 71.933998107910156, 0 ], 142 | [ 0, -100, 0, 0 ], 143 | [ -2.4330000877380371, 11.837200164794922, -3.410599946975708, 1 ] 144 | ], 145 | "cuboid_dimensions": [ 18.861799240112305, 24.238700866699219, 13.310600280761719 ] 146 | }, 147 | { 148 | "class": "021_bleach_cleanser_16k", 149 | "segmentation_class_id": 144, 150 | "fixed_model_transform": [ 151 | [ -100, 0, 0, 0 ], 152 | [ 0, 0, -100, 0 ], 153 | [ 0, -100, 0, 0 ], 154 | [ -2.1663999557495117, 12.483799934387207, 1.1710000038146973, 1 ] 155 | ], 156 | "cuboid_dimensions": [ 10.243300437927246, 25.058000564575195, 6.769899845123291 ] 157 | }, 158 | { 159 | "class": "024_bowl_16k", 160 | "segmentation_class_id": 156, 161 | "fixed_model_transform": [ 162 | [ 0, 0, 100, 0 ], 163 | [ -100, 0, 0, 0 ], 164 | [ 0, -100, 0, 0 ], 165 | [ -4.3688998222351074, 2.6974999904632568, 1.4839999675750732, 1 ] 166 | ], 167 | "cuboid_dimensions": [ 16.116399765014648, 5.5008997917175293, 16.146299362182617 ] 168 | }, 169 | { 170 | "class": "025_mug_16k", 171 | "segmentation_class_id": 168, 172 | "fixed_model_transform": [ 173 | [ 99.862899780273438, 0, 5.2336001396179199, 0 ], 174 | [ -5.2336001396179199, 0, 99.862899780273438, 0 ], 175 | [ 0, -100, 0, 0 ], 176 | [ 0.97670000791549683, 4.0118999481201172, -1.6074999570846558, 1 ] 177 | ], 178 | "cuboid_dimensions": [ 11.699000358581543, 8.1302995681762695, 9.3087997436523438 ] 179 | }, 180 | { 181 | "class": "035_power_drill_16k", 182 | "segmentation_class_id": 180, 183 | "fixed_model_transform": [ 184 | [ 99.923896789550781, 1.7452000379562378, -3.4893999099731445, 0 ], 185 | [ 1.6521999835968018, -99.95050048828125, -2.6770000457763672, 0 ], 186 | [ -3.5343999862670898, 2.6173000335693359, -99.9031982421875, 0 ], 187 | [ 4.5402998924255371, 1.1200000047683716, 2.2990999221801758, 1 ] 188 | ], 189 | "cuboid_dimensions": [ 18.436100006103516, 18.683599472045898, 5.7192001342773438 ] 190 | }, 191 | { 192 | "class": "036_wood_block_16k", 193 | "segmentation_class_id": 192, 194 | "fixed_model_transform": [ 195 | [ -22.494199752807617, 0.87269997596740723, 97.433296203613281, 0 ], 196 | [ -97.436996459960938, 0, -22.495100021362305, 0 ], 197 | [ -0.19629999995231628, -99.996200561523438, 0.85030001401901245, 0 ], 198 | [ -0.52619999647140503, 10.234100341796875, -2.6282000541687012, 1 ] 199 | ], 200 | "cuboid_dimensions": [ 8.969599723815918, 20.609199523925781, 9.0780000686645508 ] 201 | }, 202 | { 203 | "class": "037_scissors_16k", 204 | "segmentation_class_id": 204, 205 | "fixed_model_transform": [ 206 | [ -27.525999069213867, 96.126197814941406, -1.4426000118255615, 0 ], 207 | [ -96.136802673339844, -27.525999069213867, 0.20250000059604645, 0 ], 208 | [ -0.20250000059604645, 1.4426000118255615, 99.989402770996094, 0 ], 209 | [ 4.518700122833252, -1.2343000173568726, -0.71310001611709595, 1 ] 210 | ], 211 | "cuboid_dimensions": [ 20.259799957275391, 8.7530002593994141, 1.565500020980835 ] 212 | }, 213 | { 214 | "class": "040_large_marker_16k", 215 | "segmentation_class_id": 216, 216 | "fixed_model_transform": [ 217 | [ -6.1027998924255371, 2.6177000999450684, 99.779296875, 0 ], 218 | [ -0.15979999303817749, -99.9656982421875, 2.6127998828887939, 0 ], 219 | [ 99.813499450683594, 0, 6.1048998832702637, 0 ], 220 | [ -1.1348999738693237, -0.40759998559951782, 3.5237998962402344, 1 ] 221 | ], 222 | "cuboid_dimensions": [ 1.8833999633789063, 12.088600158691406, 1.9476000070571899 ] 223 | }, 224 | { 225 | "class": "051_large_clamp_16k", 226 | "segmentation_class_id": 228, 227 | "fixed_model_transform": [ 228 | [ -14.553999900817871, -98.32550048828125, 10.96720027923584, 0 ], 229 | [ 98.754302978515625, -15.107999801635742, -4.3980998992919922, 0 ], 230 | [ 5.9814000129699707, 10.190500259399414, 99.299400329589844, 0 ], 231 | [ 1.027400016784668, -0.8101000189781189, -1.802299976348877, 1 ] 232 | ], 233 | "cuboid_dimensions": [ 16.494100570678711, 12.166600227355957, 3.6412999629974365 ] 234 | }, 235 | { 236 | "class": "052_extra_large_clamp_16k", 237 | "segmentation_class_id": 240, 238 | "fixed_model_transform": [ 239 | [ 99.357200622558594, -11.320300102233887, 0, 0 ], 240 | [ 11.318599700927734, 99.342002868652344, -1.7452000379562378, 0 ], 241 | [ 0.19760000705718994, 1.7339999675750732, 99.98480224609375, 0 ], 242 | [ 2.7757999897003174, 3.1777999401092529, -1.8323999643325806, 1 ] 243 | ], 244 | "cuboid_dimensions": [ 20.259599685668945, 16.520299911499023, 3.6475999355316162 ] 245 | }, 246 | { 247 | "class": "061_foam_brick_16k", 248 | "segmentation_class_id": 252, 249 | "fixed_model_transform": [ 250 | [ 0, 0, 100, 0 ], 251 | [ -100, 0, 0, 0 ], 252 | [ 0, -100, 0, 0 ], 253 | [ 1.7200000286102295, 2.5250999927520752, 1.8260999917984009, 1 ] 254 | ], 255 | "cuboid_dimensions": [ 7.7873997688293457, 5.1192998886108398, 5.2560000419616699 ] 256 | } 257 | ] 258 | } -------------------------------------------------------------------------------- /nvdu/core/__init__.py: -------------------------------------------------------------------------------- 1 | from .utils3d import * 2 | from .transform3d import * 3 | from .scene_object import * 4 | from .cuboid import * 5 | from .camera import * 6 | from .box import * 7 | from .nvdu_data import * -------------------------------------------------------------------------------- /nvdu/core/box.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) 2018 NVIDIA Corporation. All rights reserved. 2 | # This work is licensed under a Creative Commons Attribution-NonCommercial-ShareAlike 4.0 International 3 | # License. (https://creativecommons.org/licenses/by-nc-sa/4.0/legalcode) 4 | 5 | from .scene_object import * 6 | 7 | class Box2d(SceneObject): 8 | # Create a box from its border 9 | def __init__(self, left, right, top, bottom): 10 | super(BoundingBox2d, self).__init__() 11 | 12 | self.left = left 13 | self.right = right 14 | self.top = top 15 | self.bottom = bottom 16 | 17 | self.generate_vertexes() 18 | 19 | def get_width(): 20 | return (self.right - self.left) 21 | 22 | def get_height(): 23 | return (self.bottom - self.top) 24 | 25 | def get_size(): 26 | return [self.get_width(), self.get_height()] -------------------------------------------------------------------------------- /nvdu/core/camera.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) 2018 NVIDIA Corporation. All rights reserved. 2 | # This work is licensed under a Creative Commons Attribution-NonCommercial-ShareAlike 4.0 International 3 | # License. (https://creativecommons.org/licenses/by-nc-sa/4.0/legalcode 4 | 5 | import numpy as np 6 | import json 7 | from os import path 8 | from .transform3d import * 9 | from .utils3d import * 10 | 11 | class CameraIntrinsicSettings(object): 12 | DEFAULT_ZNEAR = 1 13 | # DEFAULT_ZFAR = 100000.0 14 | DEFAULT_ZFAR = DEFAULT_ZNEAR 15 | 16 | def __init__(self, 17 | res_width = 640.0, res_height = 480.0, 18 | fx = 640.0, fy = 640.0, 19 | cx = 320.0, cy = 240.0, 20 | projection_matrix = None): 21 | self.res_width = res_width 22 | self.res_height = res_height 23 | self.fx = fx 24 | self.fy = fy 25 | self.cx = cx 26 | self.cy = cy 27 | 28 | self.znear = self.DEFAULT_ZNEAR 29 | self.zfar = self.DEFAULT_ZFAR 30 | 31 | self.projection_matrix = projection_matrix 32 | 33 | @staticmethod 34 | def from_perspective_fov_horizontal(res_width = 640.0, res_height = 480.0, hfov = 90.0): 35 | ''' 36 | Create camera intrinsics settings from 3d rendering horizontal field of view 37 | ''' 38 | cx = res_width / 2.0 39 | cy = res_height / 2.0 40 | fx = cx / np.tan(np.deg2rad(hfov) / 2.0) 41 | fy = fx 42 | 43 | # print("CameraIntrinsicSettings: res_width = {} - res_height = {} - hfov = {} - cx = {} - cy = {} - fx = {} - fy = {}".format( 44 | # res_width, res_height, hfov, cx, cy, fx, fy)) 45 | 46 | new_cam_instrinsics = CameraIntrinsicSettings(res_width, res_height, fx, fy, cx, cy) 47 | return new_cam_instrinsics 48 | 49 | @staticmethod 50 | def from_json_object(json_obj): 51 | intrinsic_settings = json_obj["intrinsic_settings"] if ("intrinsic_settings" in json_obj) else None 52 | if (intrinsic_settings is None): 53 | return None 54 | 55 | print("intrinsic_settings: {}".format(intrinsic_settings)) 56 | 57 | try: 58 | captured_image_size = json_obj['captured_image_size'] 59 | res_width = captured_image_size['width'] 60 | res_height = captured_image_size['height'] 61 | except KeyError: 62 | print("*** Error ***: 'captured_image_size' is not present in camera settings file. Using default 640 x 480.") 63 | res_width = 640 64 | res_height = 480 65 | 66 | fx = intrinsic_settings['fx'] if ('fx' in intrinsic_settings) else 640.0 67 | fy = intrinsic_settings['fy'] if ('fy' in intrinsic_settings) else 640.0 68 | cx = intrinsic_settings['cx'] if ('cx' in intrinsic_settings) else (res_width / 2.0) 69 | cy = intrinsic_settings['cy'] if ('cy' in intrinsic_settings) else (res_height / 2.0) 70 | 71 | projection_matrix_json = json_obj["cameraProjectionMatrix"] if ("cameraProjectionMatrix" in json_obj) else None 72 | projection_matrix = None 73 | if (not projection_matrix_json is None): 74 | projection_matrix = Matrix44(projection_matrix_json) 75 | projection_matrix[2, 0] = -projection_matrix[2, 0] 76 | projection_matrix[2, 1] = -projection_matrix[2, 1] 77 | projection_matrix[2, 3] = -projection_matrix[2, 3] 78 | projection_matrix[3, 2] = -projection_matrix[3, 2] 79 | 80 | # print("projection_matrix_json: {}".format(projection_matrix_json)) 81 | print("projection_matrix: {}".format(projection_matrix)) 82 | 83 | return CameraIntrinsicSettings(res_width, res_height, fx, fy, cx, cy, projection_matrix) 84 | 85 | @staticmethod 86 | def from_json_file(json_file_path): 87 | if (path.exists(json_file_path)): 88 | with open(json_file_path, 'r') as json_file: 89 | json_obj = json.load(json_file) 90 | if ('camera_settings' in json_obj): 91 | viewpoint_list = json_obj['camera_settings'] 92 | # TODO: Need to parse all the viewpoints information, right now we only parse the first viewpoint 93 | viewpoint_obj = viewpoint_list[0] 94 | return CameraIntrinsicSettings.from_json_object(viewpoint_obj) 95 | return None 96 | 97 | def get_intrinsic_matrix(self): 98 | """ 99 | Get the camera intrinsic matrix as numpy array 100 | """ 101 | intrinsic_matrix = np.array([ 102 | [self.fx, 0, self.cx], 103 | [0, self.fy, self.cy], 104 | [0, 0, 1.0] 105 | ], dtype='double') 106 | return intrinsic_matrix 107 | 108 | def get_projection_matrix(self): 109 | if (self.projection_matrix is None): 110 | self.calculate_projection_matrix() 111 | 112 | return self.projection_matrix 113 | 114 | def calculate_projection_matrix(self): 115 | zdiff = float(self.zfar - self.znear) 116 | a = (2.0 * self.fx) / float(self.res_width) 117 | b = (2.0 * self.fy) / float(self.res_height) 118 | # print('a: {} - b: {}'.format(a, b)) 119 | c = -self.znear / zdiff if (zdiff > 0) else 0 120 | d = (self.znear * self.zfar) / zdiff if (zdiff > 0) else (-self.znear) 121 | c1 = 1.0 - (2.0 * self.cx) / self.res_width 122 | c2 = (2.0 * self.cy) / self.res_height - 1.0 123 | 124 | self.projection_matrix = Matrix44([ 125 | [a, 0, 0, 0], 126 | [0, b, 0, 0], 127 | [c1, c2, c, d], 128 | [0, 0, -1.0, 0] 129 | ]) 130 | 131 | def str(): 132 | return "{}".format(self.get_intrinsic_matrix()) -------------------------------------------------------------------------------- /nvdu/core/cuboid.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) 2018 NVIDIA Corporation. All rights reserved. 2 | # This work is licensed under a Creative Commons Attribution-NonCommercial-ShareAlike 4.0 International 3 | # License. (https://creativecommons.org/licenses/by-nc-sa/4.0/legalcode 4 | 5 | # import future 6 | # from enum import IntEnum, unique 7 | import numpy as np 8 | import cv2 9 | from .scene_object import * 10 | 11 | # Related to the object's local coordinate system 12 | # @unique 13 | # class CuboidVertexType(IntEnum): 14 | class CuboidVertexType(): 15 | FrontTopRight = 0 16 | FrontTopLeft = 1 17 | FrontBottomLeft = 2 18 | FrontBottomRight = 3 19 | RearTopRight = 4 20 | RearTopLeft = 5 21 | RearBottomLeft = 6 22 | RearBottomRight = 7 23 | Center = 8 24 | TotalCornerVertexCount = 8 # Corner vertexes doesn't include the center point 25 | TotalVertexCount = 9 26 | 27 | # List of the vertex indexes in each line edges of the cuboid 28 | CuboidLineIndexes = [ 29 | # Front face 30 | [ CuboidVertexType.FrontTopLeft, CuboidVertexType.FrontTopRight ], 31 | [ CuboidVertexType.FrontTopRight, CuboidVertexType.FrontBottomRight ], 32 | [ CuboidVertexType.FrontBottomRight, CuboidVertexType.FrontBottomLeft ], 33 | [ CuboidVertexType.FrontBottomLeft, CuboidVertexType.FrontTopLeft ], 34 | # Back face 35 | [ CuboidVertexType.RearTopLeft, CuboidVertexType.RearTopRight ], 36 | [ CuboidVertexType.RearTopRight, CuboidVertexType.RearBottomRight ], 37 | [ CuboidVertexType.RearBottomRight, CuboidVertexType.RearBottomLeft ], 38 | [ CuboidVertexType.RearBottomLeft, CuboidVertexType.RearTopLeft ], 39 | # Left face 40 | [ CuboidVertexType.FrontBottomLeft, CuboidVertexType.RearBottomLeft ], 41 | [ CuboidVertexType.FrontTopLeft, CuboidVertexType.RearTopLeft ], 42 | # Right face 43 | [ CuboidVertexType.FrontBottomRight, CuboidVertexType.RearBottomRight ], 44 | [ CuboidVertexType.FrontTopRight, CuboidVertexType.RearTopRight ], 45 | ] 46 | 47 | # ========================= Cuboid2d ========================= 48 | class Cuboid2d(SceneObject): 49 | """Container for 2d projected points of a cuboid on an image""" 50 | def __init__(self, vertices=[]): 51 | """Create a cuboid 2d from a list of 2d points 52 | Args: 53 | vertices - numpy array ([8, 9] * 2) 54 | """ 55 | super(Cuboid2d, self).__init__() 56 | 57 | self._vertices = np.array(vertices) 58 | # print('Cuboid2d - vertices: {}'.format(self._vertices)) 59 | 60 | def get_vertex(self, vertex_type): 61 | """Get the location of a vertex 62 | Args: 63 | vertex_type: enum of type CuboidVertexType 64 | Return: 65 | Numpy array(3) - Location of the vertex type in the cuboid 66 | """ 67 | return self._vertices[vertex_type] 68 | 69 | def get_vertices(self): 70 | return self._vertices 71 | 72 | # ========================= Cuboid3d ========================= 73 | class Cuboid3d(SceneObject): 74 | # Create a box with a certain size 75 | # TODO: Instead of using center_location and coord_system, should pass in a Transform3d 76 | def __init__(self, size3d = [1.0, 1.0, 1.0], center_location = [0, 0, 0], 77 | coord_system = None, parent_object = None): 78 | super(Cuboid3d, self).__init__(parent_object) 79 | 80 | # NOTE: This local coordinate system is similar 81 | # to the intrinsic transform matrix of a 3d object 82 | self.center_location = center_location 83 | self.coord_system = coord_system 84 | self.size3d = size3d 85 | self._vertices = [0, 0, 0] * CuboidVertexType.TotalVertexCount 86 | 87 | self.generate_vertexes() 88 | 89 | def get_vertex(self, vertex_type): 90 | """Get the location of a vertex 91 | Args: 92 | vertex_type: enum of type CuboidVertexType 93 | Return: 94 | Numpy array(3) - Location of the vertex type in the cuboid 95 | """ 96 | return self._vertices[vertex_type] 97 | 98 | def get_vertices(self): 99 | return self._vertices 100 | 101 | def generate_vertexes(self): 102 | width, height, depth = self.size3d 103 | 104 | # By default just use the normal OpenCV coordinate system 105 | if (self.coord_system is None): 106 | cx, cy, cz = self.center_location 107 | # X axis point to the right 108 | right = cx + width / 2.0 109 | left = cx - width / 2.0 110 | # Y axis point downward 111 | top = cy - height / 2.0 112 | bottom = cy + height / 2.0 113 | # Z axis point forward 114 | front = cz + depth / 2.0 115 | rear = cz - depth / 2.0 116 | 117 | # List of 8 vertices of the box 118 | self._vertices = [ 119 | [right, top, front], # Front Top Right 120 | [left, top, front], # Front Top Left 121 | [left, bottom, front], # Front Bottom Left 122 | [right, bottom, front], # Front Bottom Right 123 | [right, top, rear], # Rear Top Right 124 | [left, top, rear], # Rear Top Left 125 | [left, bottom, rear], # Rear Bottom Left 126 | [right, bottom, rear], # Rear Bottom Right 127 | self.center_location, # Center 128 | ] 129 | else: 130 | # NOTE: should use quaternion for initial transform 131 | sx, sy, sz = self.size3d 132 | forward = np.array(self.coord_system.forward, dtype=float) * sy * 0.5 133 | up = np.array(self.coord_system.up, dtype=float) * sz * 0.5 134 | right = np.array(self.coord_system.right, dtype=float) * sx * 0.5 135 | center = np.array(self.center_location, dtype=float) 136 | self._vertices = [ 137 | center + forward + up + right, # Front Top Right 138 | center + forward + up - right, # Front Top Left 139 | center + forward - up - right, # Front Bottom Left 140 | center + forward - up + right, # Front Bottom Right 141 | center - forward + up + right, # Rear Top Right 142 | center - forward + up - right, # Rear Top Left 143 | center - forward - up - right, # Rear Bottom Left 144 | center - forward - up + right, # Rear Bottom Right 145 | self.center_location, # Center 146 | ] 147 | # print("cuboid3d - forward: {} - up: {} - right: {}".format(forward, up, right)) 148 | 149 | # print("cuboid3d - size3d: {}".format(self.size3d)) 150 | # print("cuboid3d - depth: {} - width: {} - height: {}".format(depth, width, height)) 151 | # print("cuboid3d - vertices: {}".format(self._vertices)) 152 | 153 | def get_projected_cuboid2d(self, cuboid_transform, camera_intrinsic_matrix): 154 | """ 155 | Project the cuboid into the projection plane using CameraIntrinsicSettings to get a cuboid 2d 156 | Args: 157 | cuboid_transform: the world transform of the cuboid 158 | camera_intrinsic_matrix: camera intrinsic matrix 159 | Return: 160 | Cuboid2d - the projected cuboid points 161 | """ 162 | 163 | # projected_vertices = [0, 0] * CuboidVertexType.TotalVertexCount 164 | # world_transform_matrix = self.get_world_transform_matrix() 165 | world_transform_matrix = cuboid_transform 166 | rvec = [0, 0, 0] 167 | tvec = [0, 0, 0] 168 | dist_coeffs = np.zeros((4, 1)) 169 | 170 | transformed_vertices = [0, 0, 0] * CuboidVertexType.TotalVertexCount 171 | for vertex_index in range(CuboidVertexType.TotalVertexCount): 172 | vertex3d = self._vertices[vertex_index] 173 | transformed_vertices[vertex_index] = world_transform_matrix * vertex3d 174 | 175 | projected_vertices = cv2.projectPoints(transformed_vertices, rvec, tvec, 176 | camera_intrinsic_matrix, dist_coeffs) 177 | 178 | return Cuboid2d(projected_vertices) 179 | -------------------------------------------------------------------------------- /nvdu/core/mesh.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) 2018 NVIDIA Corporation. All rights reserved. 2 | # This work is licensed under a Creative Commons Attribution-NonCommercial-ShareAlike 4.0 International 3 | # License. (https://creativecommons.org/licenses/by-nc-sa/4.0/legalcode 4 | 5 | from .scene_object import * 6 | 7 | # ========================= Cuboid2d ========================= 8 | class Mesh(SceneObject): 9 | """Container for a 3d model""" 10 | def __init__(self, mesh_file_path=None, in_parent_object = None): 11 | super(Mesh, self).__init__(in_parent_object) 12 | 13 | self.source_file_path = mesh_file_path 14 | 15 | def set_initial_matrix(self, new_initial_matrix): 16 | self._relative_transform.set_initial_matrix(new_initial_matrix) 17 | 18 | def get_initial_matrix(self): 19 | return self._relative_transform.initial_matrix 20 | -------------------------------------------------------------------------------- /nvdu/core/nvdu_data.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) 2018 NVIDIA Corporation. All rights reserved. 2 | # This work is licensed under a Creative Commons Attribution-NonCommercial-ShareAlike 4.0 International 3 | # License. (https://creativecommons.org/licenses/by-nc-sa/4.0/legalcode 4 | 5 | import future 6 | from os import listdir, path 7 | from ctypes import * 8 | import json 9 | # from typing import List, Dict, Tuple 10 | import fnmatch 11 | import glob 12 | 13 | import numpy as np 14 | import cv2 15 | import pyrr 16 | import copy 17 | from fuzzyfinder import fuzzyfinder 18 | 19 | from .utils3d import * 20 | from .transform3d import * 21 | from .cuboid import * 22 | from .pivot_axis import * 23 | from .mesh import * 24 | from .camera import * 25 | 26 | FrameDataExt = ".json" 27 | FrameImageExt = ".png" 28 | FrameImageExt_Depth = ".depth" 29 | FrameImageExt_Pls = ".pls" 30 | FrameImageExt_Pls_no = ".pls_no" 31 | 32 | FrameAspectDict = { 33 | 'main': { 34 | 'name': 'Main', 35 | 'ext': '' 36 | }, 37 | 'depth': { 38 | 'name': 'Depth', 39 | 'ext': FrameImageExt_Depth 40 | }, 41 | 'pls': { 42 | 'name': 'Pixel Level Segmentation', 43 | 'ext': FrameImageExt_Pls 44 | }, 45 | 'pls_no': { 46 | 'name': 'Pixel Level Segmentation No Occlusion', 47 | 'ext': FrameImageExt_Pls_no 48 | } 49 | } 50 | 51 | # =============================== Helper functions =============================== 52 | DEFAULT_MESH_NAME_FORMAT = "{}/google_16k/textured.obj" 53 | def get_mesh_file_path(mesh_folder_path, mesh_name, mesh_name_format=DEFAULT_MESH_NAME_FORMAT): 54 | # NOTE: The name of object in the object settings have postfix '_16k', we need to remove it 55 | if mesh_name.endswith('_16k'): 56 | mesh_name = mesh_name[:-4] 57 | return path.join(mesh_folder_path, mesh_name_format.format(mesh_name)) 58 | 59 | DEFAULT_FRAME_NAME_FORMAT = "{0:06d}" 60 | def get_frame_name(frame_index, frame_name_format=DEFAULT_FRAME_NAME_FORMAT): 61 | """Get the number part of a frame's name""" 62 | frame_number_name = frame_name_format.format(frame_index) 63 | return frame_number_name 64 | 65 | def is_frame_exists(data_dir_path, frame_index, frame_name_format=DEFAULT_FRAME_NAME_FORMAT): 66 | frame_data_file_name = get_frame_name(frame_index, frame_name_format) + FrameImageExt 67 | frame_data_file_path = path.join(data_dir_path, frame_data_file_name) 68 | # print("Checking frame {} at: {}".format(frame_index, frame_data_file_path)) 69 | return path.exists(frame_data_file_path) 70 | 71 | def is_frame_data_exists(data_dir_path, frame_index, frame_name_format=DEFAULT_FRAME_NAME_FORMAT): 72 | frame_data_file_name = get_frame_name(frame_index, frame_name_format) + FrameDataExt 73 | frame_data_file_path = path.join(data_dir_path, frame_data_file_name) 74 | # print("Checking frame {} at: {}".format(frame_index, frame_data_file_path)) 75 | return path.exists(frame_data_file_path) 76 | 77 | def get_dataset_setting_file_path(data_dir_path): 78 | return path.join(data_dir_path, '_settings.json') 79 | 80 | def get_dataset_object_setting_file_path(data_dir_path): 81 | return path.join(data_dir_path, '_object_settings.json') 82 | 83 | def get_frame_data_path(data_dir_path, frame_index, frame_name_format=DEFAULT_FRAME_NAME_FORMAT): 84 | frame_data_file_name = get_frame_name(frame_index, frame_name_format) + FrameDataExt 85 | frame_data_file_path = path.join(data_dir_path, frame_data_file_name) 86 | return frame_data_file_path 87 | 88 | def get_frame_image_path(data_dir_path, frame_index, frame_name_format=DEFAULT_FRAME_NAME_FORMAT, aspect_id='main'): 89 | frame_name = get_frame_name(frame_index, frame_name_format) 90 | frame_aspect_data = FrameAspectDict[aspect_id] 91 | aspect_ext = frame_aspect_data['ext'] 92 | frame_aspect_image_file = path.join(data_dir_path, frame_name) + aspect_ext + FrameImageExt 93 | return frame_aspect_image_file 94 | 95 | # Get the total number of frames in a data set 96 | # data_dir_path: path to the dataset's directory 97 | def get_number_of_frames(data_dir_path, frame_name_format=DEFAULT_FRAME_NAME_FORMAT): 98 | number_of_frame = 0 99 | min_index = 1 100 | # NOTE: We only count till 1000000 frames for now 101 | max_index = 1000000 102 | while (min_index <= max_index): 103 | check_index = int((min_index + max_index) / 2) 104 | frame_exists = is_frame_exists(data_dir_path, check_index, frame_name_format) 105 | if frame_exists: 106 | # NOTE: The frame start from 0 => the number of frame will be the last index + 1 107 | number_of_frame = check_index + 1 108 | min_index = check_index + 1 109 | else: 110 | max_index = check_index - 1 111 | 112 | return number_of_frame 113 | 114 | class NVDUDataset(object): 115 | DEFAULT_IMAGE_NAME_FILTERS = ["*.png"] 116 | 117 | def __init__(self, in_dataset_dir, 118 | in_annotation_dir = None, 119 | in_img_name_filters = None): 120 | self._dataset_dir = in_dataset_dir 121 | self._annotation_dr = in_annotation_dir if (not in_annotation_dir is None) else self._dataset_dir 122 | self._img_name_filters = in_img_name_filters if (not in_img_name_filters is None) else NVDUDataset.DEFAULT_IMAGE_NAME_FILTERS 123 | 124 | self._frame_names = [] 125 | self._frame_count = 0 126 | 127 | @property 128 | def frame_names(self): 129 | return self._frame_names 130 | 131 | @property 132 | def frame_count(self): 133 | return self._frame_count 134 | 135 | @property 136 | def camera_setting_file_path(self): 137 | return NVDUDataset.get_camera_setting_file_path(self._dataset_dir) 138 | 139 | @property 140 | def object_setting_file_path(self): 141 | return NVDUDataset.get_object_setting_file_path(self._dataset_dir) 142 | 143 | # Scane the dataset and return how many frames are in it 144 | def scan(self): 145 | self._frame_names = [] 146 | if not path.exists(self._dataset_dir): 147 | return 0 148 | 149 | print("scan - _dataset_dir: {} - _img_name_filters: {}".format(self._dataset_dir, self._img_name_filters)) 150 | for file_name in listdir(self._dataset_dir): 151 | check_file_path = path.join(self._dataset_dir, file_name) 152 | if path.isfile(check_file_path): 153 | is_name_match_filters = any(fnmatch.fnmatch(file_name, check_filter) for check_filter in self._img_name_filters) 154 | # is_name_match_filters = False 155 | # for check_filter in self._img_name_filters: 156 | # if fnmatch.fnmatch(file_name, check_filter): 157 | # is_name_match_filters = True 158 | # print("check_filter: {} - file_name: {} - is good".format(check_filter, file_name)) 159 | # break 160 | 161 | if (is_name_match_filters): 162 | # NOTE: Consider frame name as the file name without its extension 163 | frame_name = path.splitext(file_name)[0] 164 | # NOTE: Consider frame name as the first part of the string before the '.' 165 | # frame_name = file_name.split(".")[0] 166 | # Check if it have annotation data or not 167 | check_annotation_file_path = self.get_annotation_file_path_of_frame(frame_name) 168 | # print("file_name: {} - check_file_path: {} - frame_name: {} - check_annotation_file_path: {}".format(file_name, check_file_path, frame_name, check_annotation_file_path)) 169 | if path.exists(check_annotation_file_path): 170 | self._frame_names.append(frame_name) 171 | 172 | self._frame_names = sorted(self._frame_names) 173 | self._frame_count = len(self._frame_names) 174 | # print("_frame_names: {}".format(self._frame_names)) 175 | return self._frame_count 176 | 177 | def get_image_file_path_of_frame(self, in_frame_name): 178 | for existing_file in glob.glob(in_frame_name + '*'): 179 | for name_filter in self._img_name_filters: 180 | if fnmatch.fnmatch(existing_file, name_filter): 181 | return path.join(self._dataset_dir, existing_file) 182 | 183 | raise Exception('File not found: {}.*'.format(in_frame_name)) 184 | 185 | def get_annotation_file_path_of_frame(self, in_frame_name): 186 | return path.join(self._annotation_dr, in_frame_name + FrameDataExt) 187 | 188 | def get_frame_name_from_index(self, in_frame_index): 189 | return self._frame_names[in_frame_index] if ((in_frame_index >= 0) and (in_frame_index < len(self._frame_names))) else "" 190 | 191 | # Return the frame's image file path and annotation file path when know its index 192 | def get_frame_file_path_from_index(self, in_frame_index): 193 | frame_name = self.get_frame_name_from_index(in_frame_index) 194 | return self.get_frame_file_path_from_name(frame_name) 195 | 196 | # Return the frame's image file path and annotation file path when know its name 197 | def get_frame_file_path_from_name(self, in_frame_name): 198 | if (not in_frame_name): 199 | return ("", "") 200 | 201 | image_file_path = self.get_image_file_path_of_frame(in_frame_name) 202 | annotation_file_path = self.get_annotation_file_path_of_frame(in_frame_name) 203 | return (image_file_path, annotation_file_path) 204 | 205 | @staticmethod 206 | def get_default_camera_setting_file_path(data_dir_path): 207 | return path.join(data_dir_path, '_camera_settings.json') 208 | 209 | @staticmethod 210 | def get_default_object_setting_file_path(data_dir_path): 211 | return path.join(data_dir_path, '_object_settings.json') 212 | 213 | 214 | # =============================== Dataset Settings =============================== 215 | # Class represent the settings data for each exported object 216 | class ExportedObjectSettings(object): 217 | def __init__(self, name = '', mesh_file_path = '', initial_matrix = None, 218 | cuboid_dimension = Vector3([0, 0, 0]), cuboid_center = Vector3([0, 0, 0]), 219 | coord_system = CoordinateSystem(), obj_class_id = 0, obj_color = None): 220 | self.name = '' 221 | self.mesh_file_path = mesh_file_path 222 | self.initial_matrix = initial_matrix 223 | 224 | self.class_id = obj_class_id 225 | # If object's color not specified then calculate it automatically from the id 226 | if (obj_color is None): 227 | self.class_color = [0, 0, 0, 255] 228 | self.class_color[0] = int((self.class_id >> 5) / 7.0 * 255) 229 | self.class_color[1] = int(((self.class_id >> 2) & 7) / 7.0 * 255) 230 | self.class_color[2] = int((self.class_id & 3) / 3.0 * 255) 231 | else: 232 | self.class_color = obj_color 233 | 234 | self.cuboid_dimension = cuboid_dimension 235 | self.cuboid_center_local = cuboid_center 236 | self.coord_system = coord_system 237 | self.cuboid3d = Cuboid3d(self.cuboid_dimension, self.cuboid_center_local, self.coord_system) 238 | 239 | self.mesh_model = None 240 | self.pivot_axis = PivotAxis(np.array(self.cuboid_dimension)) 241 | 242 | def __str__(self): 243 | return "({} - {} - {})".format(self.name, self.mesh_file_path, self.initial_matrix) 244 | 245 | class ExporterSettings(object): 246 | def __init__(self): 247 | self.captured_image_size = [0, 0] 248 | 249 | @classmethod 250 | def parse_from_json_data(cls, json_data): 251 | parsed_exporter_settings = ExporterSettings() 252 | parsed_exporter_settings.captured_image_size = [json_data['camera_settings'][0]['captured_image_size']['width'], 253 | json_data['camera_settings'][0]['captured_image_size']['height']] 254 | print("parsed_exporter_settings.captured_image_size: {}".format(parsed_exporter_settings.captured_image_size)) 255 | 256 | return parsed_exporter_settings 257 | 258 | class DatasetSettings(): 259 | def __init__(self, mesh_dir_path=''): 260 | self.mesh_dir_path = mesh_dir_path 261 | self.obj_settings = {} 262 | self.exporter_settings = ExporterSettings() 263 | self.coord_system = CoordinateSystem() 264 | 265 | @classmethod 266 | def parse_from_json_data(cls, json_data, mesh_dir_path=''): 267 | parsed_settings = DatasetSettings(mesh_dir_path) 268 | 269 | coord_system = None 270 | parsed_settings.coord_system = coord_system 271 | 272 | for check_obj in json_data['exported_objects']: 273 | obj_class = check_obj['class'] 274 | obj_mesh_file_path = get_mesh_file_path(parsed_settings.mesh_dir_path, obj_class) 275 | obj_initial_matrix = Matrix44(check_obj['fixed_model_transform']) 276 | obj_cuboid_dimension = check_obj['cuboid_dimensions'] if ('cuboid_dimensions' in check_obj) else Vector3([0, 0, 0]) 277 | obj_cuboid_center = check_obj['cuboid_center_local'] if ('cuboid_center_local' in check_obj) else Vector3([0, 0, 0]) 278 | obj_class_id = int(check_obj['segmentation_class_id']) if ('segmentation_class_id' in check_obj) else 0 279 | obj_color = check_obj['color'] if ('color' in check_obj) else None 280 | 281 | new_obj_info = ExportedObjectSettings(obj_class, obj_mesh_file_path, 282 | obj_initial_matrix, obj_cuboid_dimension, obj_cuboid_center, coord_system, obj_class_id, obj_color) 283 | 284 | # print('scan_settings_data: {}'.format(new_obj_info)) 285 | # print('obj_mesh_file_path: {}'.format(obj_mesh_file_path)) 286 | parsed_settings.obj_settings[obj_class] = new_obj_info 287 | 288 | return parsed_settings 289 | 290 | @classmethod 291 | def parse_from_file(cls, setting_file_path, mesh_dir_path=''): 292 | parsed_settings = None 293 | 294 | if (path.exists(setting_file_path)): 295 | json_data = json.load(open(setting_file_path)) 296 | parsed_settings = cls.parse_from_json_data(json_data, mesh_dir_path) 297 | print('parse_from_file: setting_file_path: {} - mesh_dir_path: {} - parsed_settings: {}'.format( 298 | setting_file_path, mesh_dir_path, parsed_settings)) 299 | 300 | return parsed_settings 301 | 302 | @classmethod 303 | def parse_from_dataset(cls, dataset_dir_path, mesh_dir_path=''): 304 | setting_file_path = get_dataset_object_setting_file_path(dataset_dir_path) 305 | # print('parse_from_dataset: dataset_dir_path: {} - mesh_dir_path: {}'.format(dataset_dir_path, mesh_dir_path)) 306 | return cls.parse_from_file(setting_file_path, mesh_dir_path) 307 | 308 | # Get the settings info for a specified object class 309 | def get_object_settings(self, object_class): 310 | if (object_class in self.obj_settings): 311 | return self.obj_settings[object_class] 312 | # TODO: If there are no match object_class name then try to find the closest match using fuzzy find 313 | all_object_classes = list(self.obj_settings.keys()) 314 | fuzzy_object_classes = list(fuzzyfinder(object_class, all_object_classes)) 315 | if (len(fuzzy_object_classes) > 0): 316 | fuzzy_object_class = fuzzy_object_classes[0] 317 | # print("fuzzy_object_classes: {} - fuzzy_object_class: {}".format(fuzzy_object_classes, fuzzy_object_class)) 318 | return self.obj_settings[fuzzy_object_class] 319 | 320 | return None 321 | 322 | # =============================== AnnotatedObjectInfo =============================== 323 | # Class contain annotation data of each object in the scene 324 | class AnnotatedObjectInfo(SceneObject): 325 | def __init__(self, dataset_settings, obj_class = '', name = ''): 326 | super(AnnotatedObjectInfo, self).__init__() 327 | 328 | self.name = name 329 | self.obj_class = obj_class 330 | self.object_settings = dataset_settings.get_object_settings(obj_class) if not (dataset_settings is None) else None 331 | 332 | self.location = pyrr.Vector3() 333 | self.cuboid_center = None 334 | self.quaternion = pyrr.Quaternion([0.0, 0.0, 0.0, 1.0]) 335 | # self.bb2d = BoundingBox() 336 | self.cuboid2d = None 337 | self.keypoints = [] 338 | 339 | if not (self.object_settings is None): 340 | self.dimension = self.object_settings.cuboid_dimension 341 | self.cuboid3d = copy.deepcopy(self.object_settings.cuboid3d) if (not self.object_settings is None) else None 342 | self.mesh = Mesh(self.object_settings.mesh_file_path) 343 | self.mesh.set_initial_matrix(self.object_settings.initial_matrix) 344 | self.pivot_axis = self.object_settings.pivot_axis 345 | else: 346 | self.dimension = None 347 | self.cuboid3d = None 348 | self.mesh = None 349 | self.pivot_axis = None 350 | 351 | self.is_modified = False 352 | self.relative_transform = transform3d() 353 | 354 | # Parse and create an annotated object from a json object 355 | @classmethod 356 | def parse_from_json_object(self, dataset_settings, json_obj): 357 | try: 358 | obj_class = json_obj['class'] 359 | # print('parse_from_json_object: dataset_settings: {} - name: {} - class: {}'.format( 360 | # dataset_settings, obj_name, obj_class)) 361 | except KeyError: 362 | print("*** Error ***: 'class' is not present in annotation file. Using default '002_master_chef_can_16k'.") 363 | obj_class = '002_master_chef_can_16k' 364 | 365 | parsed_object = AnnotatedObjectInfo(dataset_settings, obj_class) 366 | if ('location' in json_obj): 367 | parsed_object.location = json_obj['location'] 368 | if ('quaternion_xyzw' in json_obj): 369 | parsed_object.quaternion = Quaternion(json_obj['quaternion_xyzw']) 370 | if ('cuboid_centroid' in json_obj): 371 | parsed_object.cuboid_center = json_obj['cuboid_centroid'] 372 | 373 | # TODO: Parse bounding box 2d 374 | # json_obj['bounding_rectangle_imagespace'] 375 | 376 | # Parse the cuboid in image space 377 | if ('projected_cuboid' in json_obj): 378 | img_width, img_height = dataset_settings.exporter_settings.captured_image_size 379 | cuboid2d_vertices_json_data = json_obj['projected_cuboid'] 380 | # Convert the fraction coordinate to absolute coordinate 381 | # cuboid2d_vertices = list([img_width * vertex['x'], img_height * vertex['y']] for vertex in cuboid2d_vertices_json_data) 382 | cuboid2d_vertices = cuboid2d_vertices_json_data 383 | # print('img_width: {} - img_height: {}'.format(img_width, img_height)) 384 | # print('cuboid2d_vertices: {}'.format(cuboid2d_vertices)) 385 | parsed_object.cuboid2d = Cuboid2d(cuboid2d_vertices) 386 | 387 | # Parse the keypoints 388 | if ('keypoints' in json_obj): 389 | annotated_keypoints_json_data = json_obj['keypoints'] 390 | for check_keypoint_json_obj in annotated_keypoints_json_data: 391 | parsed_object.keypoints.append(check_keypoint_json_obj) 392 | 393 | parsed_object.update_transform() 394 | 395 | return parsed_object 396 | 397 | def set_transform(self, new_location, new_quaternion): 398 | self.location = new_location 399 | self.quaternion = new_quaternion 400 | self.is_modified = True 401 | 402 | def set_location(self, new_location): 403 | self.location = new_location 404 | self.is_modified = True 405 | 406 | def set_quaternion(self, new_quaternion): 407 | self.quaternion = new_quaternion 408 | self.is_modified = True 409 | 410 | def update_transform(self): 411 | self.set_relative_transform(self.location, self.quaternion) 412 | 413 | # print('update_transform: location: {} - quaternion: {}'.format(self.location, self.quaternion)) 414 | should_show = not (self.location is None) and not (self.quaternion is None) 415 | if (not self.mesh is None) and should_show: 416 | self.mesh.set_relative_transform(self.location, self.quaternion) 417 | 418 | if (not self.cuboid3d is None) and should_show: 419 | cuboid_location = self.cuboid_center if (not self.cuboid_center is None) else self.location 420 | # self.cuboid3d.set_relative_transform(self.location, self.quaternion) 421 | self.cuboid3d.set_relative_transform(cuboid_location, self.quaternion) 422 | 423 | if (not self.pivot_axis is None) and should_show: 424 | self.pivot_axis.set_relative_transform(self.location, self.quaternion) 425 | 426 | self.is_modified = False 427 | 428 | # =============================== AnnotatedSceneInfo =============================== 429 | class AnnotatedSceneInfo(object): 430 | """Annotation data of a scene""" 431 | def __init__(self, dataset_settings): 432 | self.source_file_path = "" 433 | self.dataset_settings = dataset_settings 434 | self.objects = [] 435 | # Numpy array of pixel data 436 | self.image_data = None 437 | self.camera_intrinsics = None 438 | 439 | def get_object_info(self, object_class_name): 440 | found_objects = [] 441 | for check_object in self.objects: 442 | if not (check_object is None) and (check_object.obj_class == object_class_name): 443 | found_objects.append(check_object) 444 | return found_objects 445 | 446 | def set_image_data(self, new_image_numpy_data): 447 | self.image_data = new_image_numpy_data 448 | # print("set_image_data: {}".format(self.image_data.shape)) 449 | 450 | def get_scene_info_str(self): 451 | info_str = path.splitext(path.basename(self.source_file_path))[0] 452 | return info_str 453 | 454 | # Parse and create an annotated scene from a json object 455 | @classmethod 456 | def create_from_json_data(cls, dataset_settings, frame_json_data, image_data): 457 | parsed_scene = AnnotatedSceneInfo(dataset_settings) 458 | 459 | # self.camera_intrinsics = dataset_settings.camera_intrinsics 460 | if ('view_data' in frame_json_data): 461 | view_data = frame_json_data['view_data'] 462 | # TODO: Need to handle the randomized FOV in different frame 463 | # if ((not view_data is None) and ('fov' in view_data)): 464 | # camera_fovx = frame_json_data['view_data']['fov'] 465 | # parsed_scene.camera_fovx = camera_fovx 466 | # parsed_scene.camera_intrinsics = CameraIntrinsicSettings.from_perspective_fov_horizontal(hfov=camera_fovx) 467 | 468 | parsed_scene.objects = [] 469 | try: 470 | objects_data = frame_json_data['objects'] 471 | for check_obj_info in objects_data: 472 | new_obj = AnnotatedObjectInfo.parse_from_json_object(dataset_settings, check_obj_info) 473 | parsed_scene.objects.append(new_obj) 474 | except KeyError: 475 | print("*** Error ***: 'objects' is not present in annotation file. No annotations will be displayed.") 476 | 477 | # !TEST 478 | # obj_transform = new_obj.get_world_transform_matrix() 479 | # projected_cuboid = new_obj.cuboid3d.get_projected_cuboid2d(obj_transform, self.camera_intrinsics.get_intrinsic_matrix()) 480 | 481 | parsed_scene.image_data = image_data 482 | 483 | return parsed_scene 484 | 485 | # Parse and create an annotated scene from a json object 486 | @classmethod 487 | def create_from_file(cls, dataset_settings, frame_file_path, image_file_path=""): 488 | json_data = json.load(open(frame_file_path)) 489 | if (path.exists(image_file_path)): 490 | image_data = np.array(cv2.imread(image_file_path)) 491 | image_data = image_data[:,:,::-1] # Reorder color channels to be RGB 492 | else: 493 | image_data = None 494 | 495 | new_scene_info = cls.create_from_json_data(dataset_settings, json_data, image_data) 496 | new_scene_info.source_file_path = frame_file_path 497 | return new_scene_info 498 | -------------------------------------------------------------------------------- /nvdu/core/pivot_axis.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) 2018 NVIDIA Corporation. All rights reserved. 2 | # This work is licensed under a Creative Commons Attribution-NonCommercial-ShareAlike 4.0 International 3 | # License. (https://creativecommons.org/licenses/by-nc-sa/4.0/legalcode 4 | 5 | from .scene_object import * 6 | 7 | # ========================= Cuboid2d ========================= 8 | class PivotAxis(SceneObject): 9 | """Object represent a pivot axis in 3d space""" 10 | def __init__(self, in_size3d = [1.0, 1.0, 1.0], in_parent_object = None): 11 | super(PivotAxis, self).__init__(in_parent_object) 12 | 13 | self.size3d = in_size3d 14 | self.generate_vertexes() 15 | 16 | def generate_vertexes(self): 17 | self.origin_loc = [0.0, 0.0, 0.0] 18 | x, y, z = self.size3d 19 | self.x_axis = [x, 0.0, 0.0] 20 | self.y_axis = [0.0, y, 0.0] 21 | self.z_axis = [0.0, 0.0, z] 22 | 23 | -------------------------------------------------------------------------------- /nvdu/core/scene_object.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) 2018 NVIDIA Corporation. All rights reserved. 2 | # This work is licensed under a Creative Commons Attribution-NonCommercial-ShareAlike 4.0 International 3 | # License. (https://creativecommons.org/licenses/by-nc-sa/4.0/legalcode 4 | 5 | from pyrr import Quaternion, Matrix44, Vector3, euler 6 | import numpy as np 7 | 8 | from .utils3d import * 9 | from .transform3d import * 10 | 11 | class SceneObject(object): 12 | def __init__(self, in_parent_object = None): 13 | # Relative transform relate to the parent object 14 | self._relative_transform = transform3d() 15 | 16 | # List of child object attached to this scene object 17 | self.child_objects = [] 18 | # The object which this object is attached to 19 | self.parent_object = in_parent_object 20 | self.attach_to_object(in_parent_object) 21 | 22 | def attach_to_object(self, new_parent_object): 23 | if (self.parent_object): 24 | self.parent_object.remove_child_object(self) 25 | 26 | self.parent_object = new_parent_object 27 | if (self.parent_object): 28 | self.parent_object.child_objects.append(self) 29 | 30 | def remove_child_object(self, child_object_to_be_removed): 31 | print('***********************************') 32 | if (child_object_to_be_removed): 33 | try: 34 | self.child_objects.remove(child_object_to_be_removed) 35 | except ValueError: 36 | pass 37 | 38 | def set_relative_transform(self, new_location, new_quaternion): 39 | self._relative_transform.set_location(new_location) 40 | self._relative_transform.set_quaternion(new_quaternion) 41 | 42 | def get_relative_transform(self): 43 | return self._relative_transform 44 | 45 | def get_relative_transform_matrix(self): 46 | return self._relative_transform.to_matrix() 47 | 48 | def get_world_transform_matrix(self): 49 | """Get the World-to-Object transform matrix""" 50 | if (self.parent_object is None): 51 | return self.get_relative_transform_matrix() 52 | parent_world_matrix = self.parent_object.get_world_transform_matrix() 53 | world_transform_matrix = parent_world_matrix * self.get_relative_transform_matrix() 54 | return world_transform_matrix 55 | -------------------------------------------------------------------------------- /nvdu/core/transform3d.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) 2018 NVIDIA Corporation. All rights reserved. 2 | # This work is licensed under a Creative Commons Attribution-NonCommercial-ShareAlike 4.0 International 3 | # License. (https://creativecommons.org/licenses/by-nc-sa/4.0/legalcode 4 | 5 | from pyrr import Quaternion, Matrix44, Vector3, euler 6 | import numpy as np 7 | 8 | from .utils3d import * 9 | 10 | class transform3d(): 11 | def __init__(self): 12 | self.location = Vector3() 13 | self.scale = Vector3([1, 1, 1]) 14 | self.rotation = Rotator([0.0, 0.0, 0.0]) 15 | self.quaternion = Quaternion([0, 0, 0, 1]) 16 | self.initial_matrix = Matrix44.identity() 17 | 18 | self.transform_matrix = Matrix44.identity() 19 | # Flag indicate whether the transformation is modified or not 20 | self.is_changed = False 21 | 22 | def to_matrix(self): 23 | if self.is_changed: 24 | self.update_transform_matrix() 25 | 26 | return self.transform_matrix 27 | 28 | def update_transform_matrix(self): 29 | scale_matrix = Matrix44.from_scale(self.scale) 30 | translation_matrix = Matrix44.from_translation(self.location) 31 | 32 | # TODO: For some reason the qx, qy, qz part of the quaternion must be flipped 33 | # Need to understand why and fix it 34 | # The change need to be made together with the coordinate conversion in NDDS 35 | 36 | # rotation_matrix = Matrix44.from_quaternion(self.quaternion) 37 | qx, qy, qz, qw = self.quaternion 38 | test_quaternion = Quaternion([-qx, -qy, -qz, qw]) 39 | rotation_matrix = Matrix44.from_quaternion(test_quaternion) 40 | 41 | # print('update_transform_matrix: rotation_matrix = {}'.format(rotation_matrix)) 42 | 43 | relative_matrix = (translation_matrix * scale_matrix * rotation_matrix) 44 | # self.transform_matrix = relative_matrix * self.initial_matrix 45 | self.transform_matrix = relative_matrix 46 | # print('update_transform_matrix: transform_matrix = {}'.format(self.transform_matrix)) 47 | 48 | def mark_changed(self): 49 | self.is_changed = True 50 | 51 | # ======================== Rotation ======================== 52 | def set_euler_rotation(self, new_rotation): 53 | self.rotation = new_rotation 54 | new_quaternion = self.rotation.to_quaternion() 55 | self.set_quaternion(new_quaternion) 56 | 57 | def rotate(self, angle_rotator): 58 | # new_rotator = self.rotation + angle_rotator 59 | new_rotator = self.rotation.add(angle_rotator) 60 | # print("New rotation: {}".format(new_rotator)) 61 | self.set_euler_rotation(new_rotator) 62 | 63 | def set_quaternion(self, new_quaternion): 64 | # print("New quaternion: {}".format(new_quaternion)) 65 | self.quaternion = new_quaternion 66 | self.mark_changed() 67 | 68 | # ======================== Scale ======================== 69 | # Scale the mesh with the same amount between all the axis 70 | def set_scale_uniform(self, uniform_scale): 71 | self.scale *= uniform_scale 72 | self.mark_changed() 73 | 74 | def set_scale(self, new_scale): 75 | self.scale = new_scale 76 | self.mark_changed() 77 | 78 | # ======================== Translation ======================== 79 | def set_location(self, new_location): 80 | self.location = new_location 81 | self.mark_changed() 82 | 83 | def move(self, move_vector): 84 | new_location = self.location + move_vector 85 | self.set_location(new_location) 86 | 87 | # ======================== Others ======================== 88 | def set_initial_matrix(self, new_initial_matrix): 89 | self.initial_matrix = new_initial_matrix 90 | self.mark_changed() 91 | 92 | def reset_transform(self): 93 | self.set_location(Vector3()) 94 | # self.set_rotation(euler.create(0.0, 0.0, 0.0)) 95 | self.set_quaternion(Quaternion([0, 0, 0, 1])) 96 | self.mark_changed() 97 | -------------------------------------------------------------------------------- /nvdu/core/utils3d.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) 2018 NVIDIA Corporation. All rights reserved. 2 | # This work is licensed under a Creative Commons Attribution-NonCommercial-ShareAlike 4.0 International 3 | # License. (https://creativecommons.org/licenses/by-nc-sa/4.0/legalcode 4 | 5 | from pyrr import Quaternion, Matrix44, Vector3, euler 6 | import numpy as np 7 | 8 | # =============================== Data parsing =============================== 9 | 10 | # ========================= CoordinateSystem ========================= 11 | # TODO: May just want to use the Transform3d class 12 | class CoordinateSystem(object): 13 | """This class present a coordinate system with 3 main directions 14 | Each direction is represent by a vector in OpenCV coordinate system 15 | By default in OpenCV coordinate system: 16 | Forward: Z - [0, 0, 1] 17 | Right: X - [1, 0, 0] 18 | Up: -Y - [0, -1, 0] 19 | """ 20 | def __init__(self, 21 | forward = [0, 0, 1], 22 | right = [1, 0, 0], 23 | up = [0, -1, 0]): 24 | self.forward = forward 25 | self.right = right 26 | self.up = up 27 | 28 | # TODO: Build the transform matrix to convert from OpenCV to this coordinate system 29 | 30 | # ========================= Rotator ========================= 31 | class Rotator(): 32 | def __init__(self, angles = [0, 0, 0]): 33 | # Store the angle (in radian) rotated in each axis: X, Y, Z 34 | self.angles = angles 35 | 36 | # NOTE: All the calculation use the OpenCV coordinate system 37 | # http://homepages.inf.ed.ac.uk/rbf/CVonline/LOCAL_COPIES/OWENS/LECT9/node2.html 38 | 39 | def __str__(self): 40 | return "({}, {}, {})".format(self.angles[0], self.angles[1], self.angles[2]) 41 | 42 | @property 43 | def yaw(self): 44 | return self.angle[1] 45 | 46 | @property 47 | def pitch(self): 48 | return self.angle[0] 49 | 50 | @property 51 | def roll(self): 52 | return self.angle[2] 53 | 54 | @staticmethod 55 | def create_from_yaw_pitch_roll(yaw = 0, pitch = 0, roll = 0): 56 | return rotator([pitch, yaw, roll]) 57 | 58 | @staticmethod 59 | def create_from_yaw_pitch_roll_degree(yaw = 0, pitch = 0, roll = 0): 60 | return rotator([np.deg2rad(pitch), np.deg2rad(yaw), np.deg2rad(roll)]) 61 | 62 | def add(self, other_rotator): 63 | return rotator([ 64 | self.angles[0] + other_rotator.angles[0], 65 | self.angles[1] + other_rotator.angles[1], 66 | self.angles[2] + other_rotator.angles[2], 67 | ]) 68 | 69 | # https://en.wikipedia.org/wiki/Euler_angles#Rotation_matrix 70 | # NOTE: The Tait-Bryan angles result is reverted so instead of ZXY we must use the YXZ 71 | # NOTE: The rotation order is Yaw Pitch Roll or Y X Z 72 | # Output: 3x3 rotation row-major matrix 73 | def to_rotation_matrix(self): 74 | x, y, z = self.angles 75 | 76 | # Z1 X2 Y3 77 | c1 = np.cos(z) 78 | s1 = np.sin(z) 79 | c2 = np.cos(x) 80 | s2 = np.sin(x) 81 | c3 = np.cos(y) 82 | s3 = np.sin(y) 83 | 84 | return np.array( 85 | [ 86 | [ 87 | c1 * c3 - s1 * s2 * s3, 88 | -c2 * s1, 89 | c1 * s3 + c3 * s1 * s2 90 | ], 91 | [ 92 | c3 * s1 + c1 * s2 * s3, 93 | c1 * c2, 94 | s1 * s3 - c1 * c3 * s2 95 | ], 96 | [ 97 | -c2 * s3, 98 | s2, 99 | c2 * c3 100 | ] 101 | ] 102 | ) 103 | 104 | # https://en.wikipedia.org/wiki/Conversion_between_quaternions_and_Euler_angles#Euler_Angles_to_Quaternion_Conversion 105 | # NOTE: The rotation order is Yaw Pitch Roll or Y X Z 106 | def to_quaternion(self): 107 | x, y, z = self.angles 108 | 109 | halfX = x * 0.5 110 | halfY = y * 0.5 111 | halfZ = z * 0.5 112 | 113 | cX = np.cos(halfX) 114 | sX = np.sin(halfX) 115 | cY = np.cos(halfY) 116 | sY = np.sin(halfY) 117 | cZ = np.cos(halfZ) 118 | sZ = np.sin(halfZ) 119 | 120 | return np.array( 121 | [ 122 | -sX * cY * cZ - cX * sY * sZ, 123 | sX * cY * sZ - cX * sY * cZ, 124 | sX * sY * cZ - cX * cY * sZ, 125 | sX * sY * sZ + cX * cY * cZ 126 | ] 127 | ) -------------------------------------------------------------------------------- /nvdu/tools/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NVIDIA/Dataset_Utilities/532b8c76e3d7946748a10af3398438b35383f157/nvdu/tools/__init__.py -------------------------------------------------------------------------------- /nvdu/tools/nvdu_ycb.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) 2018 NVIDIA Corporation. All rights reserved. 2 | # This work is licensed under a Creative Commons Attribution-NonCommercial-ShareAlike 4.0 International 3 | # License. (https://creativecommons.org/licenses/by-nc-sa/4.0/legalcode 4 | 5 | import future 6 | import os 7 | from os import path 8 | import numpy as np 9 | import pyrr 10 | from pyrr import Quaternion, Matrix33, Matrix44, Vector4 11 | import urllib.request 12 | import tarfile 13 | import zipfile 14 | import shutil 15 | # from enum import IntEnum, unique 16 | import argparse 17 | 18 | import nvdu 19 | from nvdu.core.nvdu_data import * 20 | 21 | # =============================== Constant variables =============================== 22 | # YCB_DIR_ORIGINAL = "ycb/original" 23 | # YCB_DIR_ALIGNED = "ycb/aligned_m" 24 | # YCB_DIR_ALIGNED_SCALED = "ycb/aligned_cm" 25 | YCB_DATA_URL = "http://ycb-benchmarks.s3-website-us-east-1.amazonaws.com/data/" 26 | YCB_URL_POST_FIX = "_google_16k" # Only support the 16k meshes at the moment 27 | 28 | # @unique 29 | # class YCBModelType(IntEnum): 30 | class YCBModelType(): 31 | Original = 0 # Original model, no modifications 32 | AlignedOnly = 1 # The model is aligned but doesn't get scaled 33 | AlignedCm = 2 # The model is aligned and get scaled to centimeter unit 34 | 35 | YCB_DIR = [ 36 | path.join('ycb', 'original'), 37 | path.join('ycb', 'aligned_m'), 38 | path.join('ycb', 'aligned_cm'), 39 | ] 40 | 41 | YCB_OBJECT_SETTINGS = [ 42 | path.join('object_settings', '_ycb_original.json'), 43 | path.join('object_settings', '_ycb_aligned_m.json'), 44 | path.join('object_settings', '_ycb_aligned_cm.json'), 45 | ] 46 | 47 | # =============================== Helper functions =============================== 48 | def get_data_root_path(): 49 | nvdu_root_path = nvdu.__path__[0] 50 | data_root_path = path.join(nvdu_root_path, "data") 51 | return data_root_path 52 | 53 | def get_config_root_path(): 54 | nvdu_root_path = nvdu.__path__[0] 55 | config_root_path = path.join(nvdu_root_path, "config") 56 | return config_root_path 57 | 58 | def get_ycb_object_settings_path(ycb_model_type): 59 | config_root_path = get_config_root_path() 60 | ycb_object_settings_path = path.join(config_root_path, YCB_OBJECT_SETTINGS[ycb_model_type]) 61 | return ycb_object_settings_path 62 | 63 | def get_ycb_root_dir(ycb_model_type): 64 | return path.join(get_data_root_path(), YCB_DIR[ycb_model_type]) 65 | 66 | def get_ycb_object_url(ycb_obj_name): 67 | ycb_obj_full_url = YCB_DATA_URL + "google/" + ycb_obj_name + YCB_URL_POST_FIX + ".tgz" 68 | return ycb_obj_full_url 69 | 70 | def get_ycb_object_dir(ycb_obj_name, model_type): 71 | """ 72 | Get the directory path of an ycb object 73 | ycb_obj_name: name of the YCB object 74 | model_type: YCBModelType - type of the YCB model 75 | """ 76 | ycb_obj_dir = path.join(get_ycb_root_dir(model_type), ycb_obj_name) 77 | return ycb_obj_dir 78 | 79 | def get_ycb_model_path(ycb_obj_name, model_type): 80 | """ 81 | Get the path to the .obj model of an ycb object 82 | ycb_obj_name: name of the YCB object 83 | model_type: YCBModelType - type of the YCB model 84 | """ 85 | ycb_obj_dir = path.join(get_ycb_root_dir(model_type), ycb_obj_name) 86 | ycb_model_path = path.join(ycb_obj_dir, 'google_16k', 'textured.obj') 87 | return ycb_model_path 88 | 89 | def log_all_path_info(): 90 | ycb_dir_org = get_ycb_root_dir(YCBModelType.Original) 91 | ycb_dir_aligned_cm = get_ycb_root_dir(YCBModelType.AlignedCm) 92 | print("YCB original models: {}\nYCB aligned models in centimeter: {}".format(ycb_dir_org, ycb_dir_aligned_cm)) 93 | 94 | def log_path_info(ycb_obj_name): 95 | ycb_obj_dir_org = get_ycb_object_dir(ycb_obj_name, YCBModelType.Original) 96 | ycb_obj_dir_aligned_cm = get_ycb_object_dir(ycb_obj_name, YCBModelType.AlignedCm) 97 | print("YCB object: '{}'\nOriginal model: {}\nAligned model:{}".format(ycb_obj_name, ycb_obj_dir_org, ycb_obj_dir_aligned_cm)) 98 | if not (path.exists(ycb_obj_dir_org) and path.exists(ycb_obj_dir_aligned_cm)): 99 | print("WARNING: This YCB object model does not exist, please run 'nvdu_ycb --setup'") 100 | 101 | def log_all_object_names(): 102 | print("Supported YCB objects:") 103 | ycb_object_settings_org_path = get_ycb_object_settings_path(YCBModelType.Original) 104 | all_ycb_object_settings = DatasetSettings.parse_from_file(ycb_object_settings_org_path) 105 | 106 | for obj_name, obj_settings in all_ycb_object_settings.obj_settings.items(): 107 | # NOTE: The name of object in the object settings have postfix '_16k', we need to remove it 108 | if obj_name.endswith('_16k'): 109 | obj_name = obj_name[:-4] 110 | 111 | print("'{}'".format(obj_name)) 112 | 113 | # =============================== Mesh functions =============================== 114 | def transform_wavefront_file(src_file_path, dest_file_path, transform_matrix): 115 | dest_dir = path.dirname(dest_file_path) 116 | if (not path.exists(dest_dir)): 117 | os.makedirs(dest_dir) 118 | 119 | src_file = open(src_file_path, 'r') 120 | dest_file = open(dest_file_path, 'w') 121 | 122 | # Keep a separated non-translation matrix to use on the mesh's vertex normal 123 | non_translation_matrix = Matrix44.from_matrix33(transform_matrix.matrix33) 124 | # print("non_translation_matrix: {}".format(non_translation_matrix)) 125 | 126 | # Parse each lines in the original mesh file 127 | for line in src_file: 128 | line_args = line.split() 129 | if len(line_args): 130 | type = line_args[0] 131 | # Transform each vertex 132 | if (type == 'v'): 133 | src_vertex = pyrr.Vector4([float(line_args[1]), float(line_args[2]), float(line_args[3]), 1.0]) 134 | dest_vertex = transform_matrix * src_vertex 135 | dest_file.write("v {:.6f} {:.6f} {:.6f}\n".format(dest_vertex.x, dest_vertex.y, dest_vertex.z)) 136 | continue 137 | 138 | # Transform each vertex normal of the mesh 139 | elif (type == 'vn'): 140 | src_vertex = pyrr.Vector4([float(line_args[1]), float(line_args[2]), float(line_args[3]), 1.0]) 141 | dest_vertex = non_translation_matrix * src_vertex 142 | dest_vertex = pyrr.vector.normalize([dest_vertex.x, dest_vertex.y, dest_vertex.z]) 143 | dest_file.write("vn {:.6f} {:.6f} {:.6f}\n".format(dest_vertex[0], dest_vertex[1], dest_vertex[2])) 144 | continue 145 | dest_file.write(line) 146 | 147 | src_file.close() 148 | dest_file.close() 149 | 150 | def extract_ycb_model(ycb_obj_name): 151 | ycb_obj_dir = get_ycb_root_dir(YCBModelType.Original) 152 | 153 | # Extract the .tgz to .tar 154 | ycb_obj_local_file_name = ycb_obj_name + ".tgz" 155 | ycb_obj_local_path = path.join(ycb_obj_dir, ycb_obj_local_file_name) 156 | print("Extracting: '{}'".format(ycb_obj_local_path)) 157 | tar = tarfile.open(ycb_obj_local_path, 'r:gz') 158 | tar.extractall(path=ycb_obj_dir) 159 | tar.close() 160 | 161 | def download_ycb_model(ycb_obj_name, auto_extract=False): 162 | """ 163 | Download an ycb object's 3d models 164 | ycb_obj_name: string - name of the YCB object to download 165 | auto_extract: bool - if True then automatically extract the downloaded tgz 166 | """ 167 | ycb_obj_full_url = get_ycb_object_url(ycb_obj_name) 168 | ycb_obj_local_file_name = ycb_obj_name + ".tgz" 169 | ycb_obj_local_dir = get_ycb_root_dir(YCBModelType.Original) 170 | ycb_obj_local_path = path.join(ycb_obj_local_dir, ycb_obj_local_file_name) 171 | 172 | if (not path.exists(ycb_obj_local_dir)): 173 | os.makedirs(ycb_obj_local_dir) 174 | 175 | print("Downloading:\nURL: '{}'\nFile:'{}'".format(ycb_obj_full_url, ycb_obj_local_path)) 176 | urllib.request.urlretrieve(ycb_obj_full_url, ycb_obj_local_path) 177 | 178 | if (auto_extract): 179 | extract_ycb_model(ycb_obj_name) 180 | 181 | def align_ycb_model(ycb_obj_name, ycb_obj_settings=None): 182 | # Use the default object settings file if it's not specified 183 | if (ycb_obj_settings is None): 184 | ycb_object_settings_org_path = get_ycb_object_settings_path(YCBModelType.Original) 185 | all_ycb_object_settings = DatasetSettings.parse_from_file(ycb_object_settings_org_path) 186 | ycb_obj_settings = all_ycb_object_settings.get_object_settings(ycb_obj_name) 187 | if (ycb_obj_settings is None): 188 | print("Can't find settings of object: '{}'".format(ycb_obj_name)) 189 | return 190 | 191 | src_file_path = get_ycb_model_path(ycb_obj_name, YCBModelType.Original) 192 | dest_file_path = get_ycb_model_path(ycb_obj_name, YCBModelType.AlignedCm) 193 | # Transform the original models 194 | print("Align model:\nSource:{}\nTarget:{}".format(src_file_path, dest_file_path)) 195 | # Use the fixed transform matrix to align YCB models 196 | transform_wavefront_file(src_file_path, dest_file_path, ycb_obj_settings.initial_matrix) 197 | 198 | src_dir = path.dirname(src_file_path) 199 | dest_dir = path.dirname(dest_file_path) 200 | # Copy the material and texture to the new directory 201 | shutil.copy(path.join(src_dir, 'textured.mtl'), path.join(dest_dir, 'textured.mtl')) 202 | shutil.copy(path.join(src_dir, 'texture_map.png'), path.join(dest_dir, 'texture_map.png')) 203 | 204 | def setup_all_ycb_models(): 205 | """ 206 | Read the original YCB object settings 207 | For each object in the list: 208 | Download the 16k 3d model 209 | Extract the .tgz file 210 | Convert the original model into the aligned one 211 | """ 212 | ycb_object_settings_org_path = get_ycb_object_settings_path(YCBModelType.Original) 213 | all_ycb_object_settings = DatasetSettings.parse_from_file(ycb_object_settings_org_path) 214 | 215 | for obj_name, obj_settings in all_ycb_object_settings.obj_settings.items(): 216 | # NOTE: The name of object in the object settings have postfix '_16k', we need to remove it 217 | if obj_name.endswith('_16k'): 218 | obj_name = obj_name[:-4] 219 | print("Setting up object: '{}'".format(obj_name)) 220 | download_ycb_model(obj_name, True) 221 | align_ycb_model(obj_name, obj_settings) 222 | # break 223 | 224 | # =============================== Main =============================== 225 | def main(): 226 | parser = argparse.ArgumentParser(description='NVDU YCB models Support') 227 | parser.add_argument('ycb_object_name', type=str, nargs='?', 228 | help="Name of the YCB object to check info", default=None) 229 | parser.add_argument('-s', '--setup', action='store_true', help="Setup the YCB models for the FAT dataset", default=False) 230 | parser.add_argument('-l', '--list', action='store_true', help="List all the supported YCB objects", default=False) 231 | 232 | args = parser.parse_args() 233 | 234 | if (args.list): 235 | log_all_object_names() 236 | 237 | if (args.setup): 238 | setup_all_ycb_models() 239 | else: 240 | if (args.ycb_object_name): 241 | log_path_info(args.ycb_object_name) 242 | else: 243 | # Print the info 244 | log_all_path_info() 245 | 246 | 247 | if __name__ == "__main__": 248 | main() -------------------------------------------------------------------------------- /nvdu/tools/test_nvdu_visualizer.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) 2018 NVIDIA Corporation. All rights reserved. 2 | # This work is licensed under a Creative Commons Attribution-NonCommercial-ShareAlike 4.0 International 3 | # License. (https://creativecommons.org/licenses/by-nc-sa/4.0/legalcode 4 | 5 | #!/usr/bin/env python 6 | try: 7 | import pyglet 8 | except Exception as ex: 9 | print("Can't import pyglet module: {}".format(ex)) 10 | 11 | from pyrr import Quaternion, Matrix44, Vector3, euler 12 | import numpy as np 13 | 14 | from ctypes import * 15 | import argparse 16 | from os import path 17 | import sys 18 | 19 | import nvdu 20 | from nvdu.core import * 21 | from nvdu.viz.nvdu_visualizer import * 22 | from nvdu.viz.nvdu_viz_window import * 23 | 24 | from nvdu.tools.nvdu_ycb import * 25 | 26 | CAMERA_FOV_HORIZONTAL = 90 27 | 28 | # By default, just visualize the current directory 29 | DEFAULT_data_dir_path = '.' 30 | DEFAULT_output_dir_path = '' 31 | DEFAULT_CAMERA_SETTINGS_FILE_PATH = './config/camera_lov.json' 32 | 33 | data_dir_path = DEFAULT_data_dir_path 34 | # mesh_file_name = 'textured.obj' 35 | 36 | # ============================= MAIN ============================= 37 | def main(): 38 | DEFAULT_WINDOW_WIDTH = 0 39 | DEFAULT_WINDOW_HEIGHT = 0 40 | 41 | # Launch Arguments 42 | parser = argparse.ArgumentParser(description='NVDU Data Visualiser') 43 | parser.add_argument('dataset_dir', type=str, nargs='?', 44 | help="Dataset directory. This is where all the images (required) and annotation info (optional) are. Default is the current directory", default=DEFAULT_data_dir_path) 45 | parser.add_argument('-m', '--model_dir', type=str, help="Model directory. Default is /data/ycb/original/", default=get_ycb_root_dir(YCBModelType.Original)) 46 | parser.add_argument('-a', '--data_annot_dir', type=str, help="Directory path - where to find the annotation data. Default is the same directory as the dataset directory", default="") 47 | parser.add_argument('-s', '--size', type=int, nargs=2, metavar=('WIDTH', 'HEIGHT'), help="Window's size: [width, height]. If not specified then the window fit the resolution of the camera", default=[DEFAULT_WINDOW_WIDTH, DEFAULT_WINDOW_HEIGHT]) 48 | parser.add_argument('-o', '--object_settings_path', type=str, help="Object settings file path") 49 | parser.add_argument('-c', '--camera_settings_path', type=str, help="Camera settings file path", default=None) 50 | parser.add_argument('-n', '--name_filters', type=str, nargs='*', help="The name filter of each frame. e.g: *.png", default=["*.png"]) 51 | parser.add_argument('--fps', type=float, help="How fast do we want to automatically change frame", default=10) 52 | parser.add_argument('--auto_change', action='store_true', help="If specified, the visualizer will automatically change the frame", default=False) 53 | parser.add_argument('-e', '--export_dir', type=str, help="Directory path - where to store the visualized images. If specified, the script will automatically export the visualized image to the export directory. If not specified, the current directory will be used.", default='') 54 | parser.add_argument('--auto_export', action='store_true', help="If specified, the visualizer will automatically export the visualized frame to image file in the `export_dir` directory", default=False) 55 | parser.add_argument('--ignore_fixed_transform', action='store_true', help="If specified, the visualizer will not use the fixed transform matrix for the 3d model", default=False) 56 | # parser.add_argument('--gui', type=str, help="Show GUI window") 57 | 58 | # subparsers = parser.add_subparsers(help='sub-command help') 59 | # parser_export = subparsers.add_parser('export', help="Export the visualized frames to image files") 60 | # parser_export.add_argument('--out_dir', type=str, help="Directory path - where to store the visualized images. If this is set, the script will automatically export the visualized image to the export directory", default='viz') 61 | # parser_export.add_argument('--movie_name', type=str, help="Name of the movie", default='viz.mp4') 62 | # parser_export.add_argument('--movie_fps', type=float, help="Framerate of the generated movie", default=30) 63 | 64 | args = parser.parse_args() 65 | print("args: {}".format(args)) 66 | 67 | # TODO: May want to add auto_export as a launch arguments flag 68 | # auto_export = not (not args.export_dir) 69 | auto_export = args.auto_export 70 | 71 | dataset_dir_path = args.dataset_dir 72 | data_annot_dir_path = args.data_annot_dir if (args.data_annot_dir) else dataset_dir_path 73 | 74 | name_filters = args.name_filters 75 | print("name_filters: {}".format(name_filters)) 76 | viz_dataset = NVDUDataset(dataset_dir_path, data_annot_dir_path, name_filters) 77 | # frame_count = viz_dataset.scan() 78 | # print("Number of frames in the dataset: {}".format(frame_count)) 79 | 80 | # NOTE: Just use the YCB models path for now 81 | model_dir_path = args.model_dir 82 | 83 | object_settings_path = args.object_settings_path 84 | camera_settings_path = args.camera_settings_path 85 | # If not specified then use the default object and camera settings files from the dataset directory 86 | if (object_settings_path is None): 87 | object_settings_path = NVDUDataset.get_default_object_setting_file_path(dataset_dir_path) 88 | if (camera_settings_path is None): 89 | camera_settings_path = NVDUDataset.get_default_camera_setting_file_path(dataset_dir_path) 90 | 91 | dataset_settings = DatasetSettings.parse_from_file(object_settings_path, model_dir_path) 92 | if (dataset_settings is None): 93 | print("Error: Could not locate dataset settings at {}".format(object_settings_path)) 94 | else: 95 | if path.exists(camera_settings_path): 96 | camera_json_data = json.load(open(camera_settings_path)) 97 | dataset_settings.exporter_settings = ExporterSettings.parse_from_json_data(camera_json_data) 98 | 99 | camera_intrinsic_settings = CameraIntrinsicSettings.from_json_file(camera_settings_path) 100 | if (camera_intrinsic_settings is None): 101 | print("Error: Could not locate camera settings at {}".format(camera_settings_path)) 102 | 103 | if (dataset_settings is None or camera_intrinsic_settings is None): 104 | exit(1) 105 | 106 | # print("camera_intrinsic_settings: {} - {}".format(camera_settings_path, camera_intrinsic_settings)) 107 | 108 | # By default fit the window size to the resolution of the images 109 | # NOTE: Right now we don't really support scaling 110 | width, height = args.size 111 | if (width <= 0) or (height <= 0): 112 | width = camera_intrinsic_settings.res_width 113 | height = camera_intrinsic_settings.res_height 114 | 115 | main_window = NVDUVizWindow(width, height, 'NVDU Data Visualiser') 116 | main_window.visualizer.dataset_settings = dataset_settings 117 | main_window.visualizer.visualizer_settings.ignore_initial_matrix = args.ignore_fixed_transform 118 | main_window.dataset = viz_dataset 119 | main_window.set_auto_fps(args.fps) 120 | main_window.should_export = auto_export 121 | main_window.set_auto_change_frame(args.auto_change) 122 | main_window.export_dir = args.export_dir 123 | main_window.setup() 124 | 125 | main_window.set_camera_intrinsic_settings(camera_intrinsic_settings) 126 | 127 | pyglet.app.run() 128 | 129 | if __name__ == '__main__': 130 | main() 131 | -------------------------------------------------------------------------------- /nvdu/viz/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NVIDIA/Dataset_Utilities/532b8c76e3d7946748a10af3398438b35383f157/nvdu/viz/__init__.py -------------------------------------------------------------------------------- /nvdu/viz/background_image.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) 2018 NVIDIA Corporation. All rights reserved. 2 | # This work is licensed under a Creative Commons Attribution-NonCommercial-ShareAlike 4.0 International 3 | # License. (https://creativecommons.org/licenses/by-nc-sa/4.0/legalcode 4 | 5 | import cv2 6 | import numpy as np 7 | from pyglet.gl import * 8 | 9 | class BackgroundImage(object): 10 | def __init__(self, width = 0, height = 0): 11 | self.width = width 12 | self.height = height 13 | self.location = [0, 0, 0] 14 | self.vlist = pyglet.graphics.vertex_list(4, 15 | ('v2f', [0,0, width,0, 0,height, width,height]), 16 | ('t2f', [0,0, width,0, 0,height, width,height])) 17 | 18 | self.scale = [self.width, self.height, 1.0] 19 | 20 | @classmethod 21 | def create_from_numpy_image_data(cls, numpy_image_data, width = 0, height = 0): 22 | img_width = numpy_image_data.shape[1] if (width == 0) else width 23 | img_height = numpy_image_data.shape[0] if (height == 0) else height 24 | 25 | new_image = cls(img_width, img_height) 26 | new_image.load_image_data_from_numpy(numpy_image_data) 27 | return new_image 28 | 29 | @classmethod 30 | def create_from_file_path(cls, image_file_path, width = 0, height = 0): 31 | image_np = np.array(cv2.imread(image_file_path)) 32 | image_np = image_np[:,:,::-1] # Convert BGR to RGB format. Alternatively, use cv2.cvtColor(image, cv2.COLOR_BGR2RGB) 33 | return cls.create_from_numpy_image_data(image_np, width, height) 34 | 35 | def load_image_data_from_numpy(self, numpy_image_data): 36 | width = numpy_image_data.shape[1] 37 | height = numpy_image_data.shape[0] 38 | color_channel_count = numpy_image_data.shape[2] 39 | pitch = -width * color_channel_count 40 | # print('numpy_image_data.shape: {}'.format(numpy_image_data.shape)) 41 | img_data = numpy_image_data 42 | img_data = img_data.tostring() 43 | self.image = pyglet.image.ImageData(width, height, 'RGB', img_data, pitch) 44 | self.texture = self.image.get_texture(True, True) 45 | 46 | def load_image_from_file(self, image_file_path): 47 | self.image = pyglet.image.load(image_file_path) 48 | self.texture = self.image.get_texture(True, True) 49 | 50 | def load_new_image(self, image_file_path): 51 | self.image = pyglet.image.load(image_file_path) 52 | self.texture = self.image.get_texture(True, True) 53 | # print('Texture: {} - id: {} - target:{} - width: {} - height: {}'.format( 54 | # self.texture, self.texture.id, self.texture.target, self.texture.width, self.texture.height)) 55 | # print('GL_TEXTURE_RECTANGLE_ARB: {} - GL_TEXTURE_RECTANGLE_NV: {} - GL_TEXTURE_2D: {}'.format( 56 | # GL_TEXTURE_RECTANGLE_ARB, GL_TEXTURE_RECTANGLE_NV, GL_TEXTURE_2D)) 57 | 58 | def draw(self): 59 | glMatrixMode(GL_MODELVIEW) 60 | glPushMatrix() 61 | glLoadIdentity() 62 | x, y, z = self.location 63 | glTranslatef(x, y, z) 64 | glColor3f(1, 1, 1) 65 | 66 | texture_target = self.texture.target 67 | glEnable(texture_target) 68 | glBindTexture(texture_target, self.texture.id) 69 | self.vlist.draw(GL_TRIANGLE_STRIP) 70 | glDisable(texture_target) 71 | 72 | glPopMatrix() 73 | -------------------------------------------------------------------------------- /nvdu/viz/camera.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) 2018 NVIDIA Corporation. All rights reserved. 2 | # This work is licensed under a Creative Commons Attribution-NonCommercial-ShareAlike 4.0 International 3 | # License. (https://creativecommons.org/licenses/by-nc-sa/4.0/legalcode 4 | 5 | from pyrr import Quaternion, Matrix44, Vector3, euler 6 | from pyglet.gl import * 7 | from pyglet.gl.gl import c_float, c_double, c_int, glGetFloatv, GL_MODELVIEW_MATRIX 8 | from pyglet.gl.glu import * 9 | import numpy as np 10 | 11 | from .utils3d import * 12 | 13 | from nvdu.core import transform3d 14 | from nvdu.core import scene_object 15 | from nvdu.core.camera import * 16 | from .scene_object import * 17 | 18 | # A Camera handle the perspective and frustrum of a scene 19 | class Camera(SceneObjectViz3d): 20 | # DEFAULT_ZNEAR = 0.000001 21 | # DEFAULT_ZNEAR = 0.00001 22 | DEFAULT_ZNEAR = 1 23 | DEFAULT_ZFAR = 100000.0 24 | # Horizontal field of view of the camera (in degree) 25 | DEFAULT_FOVX = 90.0 26 | DEFAULT_ASPECT_RATIO_XY = 1920.0 / 1080.0 27 | 28 | def __init__(self, cam_intrinsic_settings = CameraIntrinsicSettings(), scene_object = None): 29 | super(Camera, self).__init__(scene_object) 30 | 31 | self.camera_matrix = Matrix44.identity() 32 | self.projection_matrix = Matrix44.identity() 33 | 34 | self.set_instrinsic_settings(cam_intrinsic_settings) 35 | 36 | def draw(self): 37 | # glClipControl(GL_LOWER_LEFT, GL_ZERO_TO_ONE) 38 | glMatrixMode(GL_PROJECTION) 39 | glLoadIdentity() 40 | # TODO: Need to invert the CameraToWorld matrix => WorldToCamera matrix 41 | # transform_matrix = self.relative_transform.to_matrix() 42 | # glMultMatrixf(get_opengl_matrixf(transform_matrix)) 43 | 44 | # self.build_perspective_projection_matrix() 45 | 46 | # print("on_draw - camera_fov: {}".format(self.camera_fov)) 47 | # print("camera - on_draw - projection_matrix: {}".format(self.projection_matrix)) 48 | 49 | glMultMatrixf(get_opengl_matrixf(self.projection_matrix)) 50 | 51 | # Scale up the scene so the objects really close to the camera doesn't get clipped by the near plane 52 | # glScalef(100.0, 100.0, 100.0) 53 | 54 | # Set up the camera using its intrinsic parameters: 55 | # (fx, fy): focal length 56 | # (cx, cy): optical center 57 | def set_instrinsic_settings(self, cam_intrinsic_settings): 58 | self.intrinsic_settings = cam_intrinsic_settings 59 | if (cam_intrinsic_settings): 60 | self.projection_matrix = cam_intrinsic_settings.get_projection_matrix() 61 | 62 | # print("set_instrinsic_settings: {} - projection matrix: {}".format(str(self.intrinsic_settings), self.projection_matrix)) 63 | 64 | def set_perspective_params(self, fovy, aspect_ratio_xy, znear = DEFAULT_ZNEAR, zfar = DEFAULT_ZFAR): 65 | self.fovy = np.deg2rad(fovy) 66 | self.aspect_ratio_xy = aspect_ratio_xy 67 | self.znear = znear 68 | self.zfar = zfar 69 | 70 | self.build_perspective_projection_matrix() 71 | 72 | def set_fovx(self, new_fovx): 73 | # new_fovy = convert_HFOV_to_VFOV(new_fovx, 1.0 / self.aspect_ratio_xy) 74 | # self.set_fovy(new_fovy) 75 | # self.build_perspective_projection_matrix() 76 | new_cam_instrinsics = CameraIntrinsicSettings.from_perspective_fov_horizontal( 77 | self.intrinsic_settings.res_width, self.intrinsic_settings.res_height, new_fovx) 78 | 79 | def build_perspective_projection_matrix(self): 80 | zdiff = float(self.znear - self.zfar) 81 | 82 | fovy_tan = np.tan(self.fovy / 2.0) 83 | # TODO: Handle fovy_tan = 0? 84 | a = 1.0 / (fovy_tan * self.aspect_ratio_xy) 85 | b = 1.0 / fovy_tan 86 | # print('a: {} - b: {}'.format(a, b)) 87 | c = (self.znear + self.zfar) / zdiff 88 | d = 2 * (self.znear * self.zfar) / zdiff 89 | 90 | self.projection_matrix = Matrix44([ 91 | [a, 0, 0, 0], 92 | [0, b, 0, 0], 93 | [0, 0, c, -1.0], 94 | [0, 0, d, 0] 95 | ]) 96 | # print('build_perspective_projection_matrix: {} - znear: {} - zfar: {} - aspect_ratio_xy: {} - fovy: {}'.format( 97 | # self.projection_matrix, self.znear, self.zfar, self.aspect_ratio_xy, self.fovy)) 98 | 99 | def set_viewport_size(self, viewport_size): 100 | self.viewport_size = viewport_size 101 | -------------------------------------------------------------------------------- /nvdu/viz/cuboid.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) 2018 NVIDIA Corporation. All rights reserved. 2 | # This work is licensed under a Creative Commons Attribution-NonCommercial-ShareAlike 4.0 International 3 | # License. (https://creativecommons.org/licenses/by-nc-sa/4.0/legalcode 4 | 5 | from pyrr import Quaternion, Matrix44, Vector3, euler 6 | import numpy as np 7 | from pyglet.gl import * 8 | from ctypes import * 9 | 10 | from .scene_object import * 11 | from nvdu.core.cuboid import * 12 | 13 | # ========================= Cuboid3d ========================= 14 | # TODO: Should merge Cuboid and Box3d 15 | class Cuboid3dViz(SceneObjectViz3d): 16 | # Create a box with a certain size 17 | def __init__(self, cuboid3d, in_color=None): 18 | super(Cuboid3dViz, self).__init__(cuboid3d) 19 | self.cuboid3d = cuboid3d 20 | self.vertices = self.cuboid3d.get_vertices() 21 | 22 | self.render_line = True 23 | self.color = in_color 24 | self.face_alpha = 128 25 | # self.render_line = False 26 | 27 | self.generate_vertex_buffer() 28 | 29 | def generate_vertex_buffer(self): 30 | self.vertices_gl = [] 31 | for i in range(0, CuboidVertexType.TotalCornerVertexCount): 32 | self.vertices_gl.append(self.vertices[i][0]) 33 | self.vertices_gl.append(self.vertices[i][1]) 34 | self.vertices_gl.append(self.vertices[i][2]) 35 | 36 | # List of color for each vertices of the box 37 | if (self.color is None): 38 | self.colors_gl = [ 39 | 0, 0, 255, 255, # Front Top Right 40 | 0, 0, 255, 255, # Front Top Left 41 | 255, 0, 255, 255, # Front Bottom Left 42 | 255, 0, 255, 255, # Front Bottom Right 43 | 0, 255, 0, 255, # Rear Top Right 44 | 0, 255, 0, 255, # Rear Top Left 45 | 255, 255, 0, 255, # Rear Bottom Left 46 | 255, 255, 0, 255, # Rear Bottom Right 47 | ] 48 | else: 49 | self.colors_gl = [] 50 | for i in range(0, CuboidVertexType.TotalCornerVertexCount): 51 | for color_channel in self.color: 52 | self.colors_gl.append(color_channel) 53 | 54 | # Reduce the alpha of the vertex colors when we rendering triangles 55 | self.colors_tri_gl = list(int(color / 4) for color in self.colors_gl) 56 | 57 | cvt = CuboidVertexType 58 | # Counter-Clockwise order triangle indices 59 | self.indices_tri_gl = [ 60 | # Front face 61 | cvt.FrontBottomLeft, cvt.FrontTopLeft, cvt.FrontTopRight, 62 | cvt.FrontTopRight, cvt.FrontBottomRight, cvt.FrontBottomLeft, 63 | # Right face 64 | cvt.FrontBottomRight, cvt.FrontTopRight, cvt.RearBottomRight, 65 | cvt.RearTopRight, cvt.RearBottomRight, cvt.FrontTopRight, 66 | # Back face 67 | cvt.RearBottomLeft, cvt.RearBottomRight, cvt.RearTopRight, 68 | cvt.RearTopRight, cvt.RearTopLeft, cvt.RearBottomLeft, 69 | # Left face 70 | cvt.FrontTopLeft, cvt.FrontBottomLeft, cvt.RearBottomLeft, 71 | cvt.RearBottomLeft, cvt.RearTopLeft, cvt.FrontTopLeft, 72 | # Top face 73 | cvt.RearTopLeft, cvt.RearTopRight, cvt.FrontTopRight, 74 | cvt.FrontTopRight, cvt.FrontTopLeft, cvt.RearTopLeft, 75 | # Bottom face 76 | cvt.RearBottomLeft, cvt.FrontBottomLeft, cvt.FrontBottomRight, 77 | cvt.FrontBottomRight, cvt.RearBottomRight, cvt.RearBottomLeft, 78 | ] 79 | self.indices_line_gl = np.array(CuboidLineIndexes).flatten() 80 | # print("indices_line_gl: {}".format(self.indices_line_gl)) 81 | 82 | self.indices_point_gl = list(range(0, len(self.vertices))) 83 | # print('indices_point_gl: {}'.format(self.indices_point_gl)) 84 | 85 | self.vertex_gl_array = (GLfloat* len(self.vertices_gl))(*self.vertices_gl) 86 | self.color_gl_array = (GLubyte* len(self.colors_gl))(*self.colors_gl) 87 | self.colors_tri_gl_array = (GLubyte* len(self.colors_tri_gl))(*self.colors_tri_gl) 88 | self.indices_tri_gl_array = (GLubyte* len(self.indices_tri_gl))(*self.indices_tri_gl) 89 | self.indices_line_gl_array = (GLubyte* len(self.indices_line_gl))(*self.indices_line_gl) 90 | self.indices_point_gl_array = (GLubyte* len(self.indices_point_gl))(*self.indices_point_gl) 91 | 92 | def on_draw(self): 93 | super(Cuboid3dViz, self).on_draw() 94 | # print('Cuboid3dViz - on_draw - vertices: {}'.format(self.vertices)) 95 | 96 | glEnableClientState(GL_VERTEX_ARRAY) 97 | glEnableClientState(GL_COLOR_ARRAY) 98 | # glEnable(GL_POLYGON_SMOOTH) 99 | 100 | glPolygonMode(GL_FRONT_AND_BACK, RenderMode.normal) 101 | 102 | glVertexPointer(3, GL_FLOAT, 0, self.vertex_gl_array) 103 | 104 | # Render each faces of the cuboid 105 | glColorPointer(4, GL_UNSIGNED_BYTE, 0, self.colors_tri_gl_array) 106 | glDrawElements(GL_TRIANGLES, len(self.indices_tri_gl), GL_UNSIGNED_BYTE, self.indices_tri_gl_array) 107 | 108 | # Render each edge lines 109 | glEnable(GL_LINE_SMOOTH) 110 | glLineWidth(3.0) 111 | glColorPointer(4, GL_UNSIGNED_BYTE, 0, self.color_gl_array) 112 | # TODO: May want to use GL_LINE_STRIP or GL_LINE_LOOP 113 | glDrawElements(GL_LINES, len(self.indices_line_gl), GL_UNSIGNED_BYTE, self.indices_line_gl_array) 114 | glDisable(GL_LINE_SMOOTH) 115 | 116 | # Render each corner vertices in POINTS mode 117 | glPointSize(10.0) 118 | glDrawElements(GL_POINTS, len(self.indices_point_gl_array), GL_UNSIGNED_BYTE, self.indices_point_gl_array) 119 | glPointSize(1.0) 120 | 121 | # Deactivate vertex arrays after drawing 122 | glDisableClientState(GL_VERTEX_ARRAY) 123 | glDisableClientState(GL_COLOR_ARRAY) 124 | # glDisable(GL_POLYGON_SMOOTH) 125 | 126 | # ========================= Cuboid2d ========================= 127 | class Cuboid2dViz(SceneObjectVizBase): 128 | # Create a box with a certain size 129 | def __init__(self, cuboid2d, in_color=None): 130 | super(Cuboid2dViz, self).__init__(cuboid2d) 131 | self.cuboid2d = cuboid2d 132 | self.color = in_color 133 | if not (self.cuboid2d is None): 134 | self.vertices = self.cuboid2d.get_vertices() 135 | self.generate_vertexes_buffer() 136 | # self.render_line = False 137 | 138 | def generate_vertexes_buffer(self): 139 | self.vertices_gl = [] 140 | max_vertex_count = min(CuboidVertexType.TotalVertexCount, len(self.vertices)) 141 | for i in range(0, max_vertex_count): 142 | vertex = self.vertices[i] 143 | if (not vertex is None): 144 | self.vertices_gl.append(self.vertices[i][0]) 145 | self.vertices_gl.append(self.vertices[i][1]) 146 | else: 147 | self.vertices_gl.append(0.0) 148 | self.vertices_gl.append(0.0) 149 | 150 | # List of color for each vertices of the box 151 | self.vertex_colors_gl = [ 152 | 0, 0, 255, 255, # Front Top Right 153 | 0, 0, 255, 255, # Front Top Left 154 | 255, 0, 255, 255, # Front Bottom Left 155 | 255, 0, 255, 255, # Front Bottom Right 156 | 0, 255, 0, 255, # Rear Top Right 157 | 0, 255, 0, 255, # Rear Top Left 158 | 255, 255, 0, 255, # Rear Bottom Left 159 | 255, 255, 0, 255, # Rear Bottom Right 160 | ] 161 | 162 | # List of color for each vertices of the box 163 | if (self.color is None): 164 | self.edge_colors_gl = self.vertex_colors_gl 165 | else: 166 | self.edge_colors_gl = [] 167 | for i in range(0, CuboidVertexType.TotalCornerVertexCount): 168 | for color_channel in self.color: 169 | self.edge_colors_gl.append(color_channel) 170 | 171 | # NOTE: Only add valid lines: 172 | self.indices_line_gl = [] 173 | for line in CuboidLineIndexes: 174 | vi0, vi1 = line 175 | v0 = self.vertices[vi0] 176 | v1 = self.vertices[vi1] 177 | if not (v0 is None) and not (v1 is None): 178 | self.indices_line_gl.append(vi0) 179 | self.indices_line_gl.append(vi1) 180 | # print('indices_line_gl: {}'.format(self.indices_line_gl)) 181 | 182 | # self.indices_line_gl = [ 183 | # # Front face 184 | # cvt.FrontTopLeft, cvt.FrontTopRight, 185 | # cvt.FrontTopRight, cvt.FrontBottomRight, 186 | # cvt.FrontBottomRight, cvt.FrontBottomLeft, 187 | # cvt.FrontBottomLeft, cvt.FrontTopLeft, 188 | # # Back face 189 | # cvt.RearTopLeft, cvt.RearTopRight, 190 | # cvt.RearTopRight, cvt.RearBottomRight, 191 | # cvt.RearBottomRight, cvt.RearBottomLeft, 192 | # cvt.RearBottomLeft, cvt.RearTopLeft, 193 | # # Left face 194 | # cvt.FrontBottomLeft, cvt.RearBottomLeft, 195 | # cvt.FrontTopLeft, cvt.RearTopLeft, 196 | # # Right face 197 | # cvt.FrontBottomRight, cvt.RearBottomRight, 198 | # cvt.FrontTopRight, cvt.RearTopRight, 199 | # ] 200 | 201 | # self.indices_point_gl = list(i for i in range(0, len(self.vertices))) 202 | # NOTE: Only add valid points: 203 | self.indices_point_gl = [] 204 | for i in range (len(self.vertices)): 205 | if (not self.vertices[i] is None): 206 | self.indices_point_gl.append(i) 207 | # print('indices_point_gl: {}'.format(self.indices_point_gl)) 208 | 209 | self.vertex_gl_array = (GLfloat* len(self.vertices_gl))(*self.vertices_gl) 210 | self.vertex_color_gl_array = (GLubyte* len(self.vertex_colors_gl))(*self.vertex_colors_gl) 211 | self.edge_colors_gl_array = (GLubyte* len(self.edge_colors_gl))(*self.edge_colors_gl) 212 | self.indices_line_gl_array = (GLubyte* len(self.indices_line_gl))(*self.indices_line_gl) 213 | self.indices_point_gl_array = (GLubyte* len(self.indices_point_gl))(*self.indices_point_gl) 214 | 215 | def on_draw(self): 216 | if (self.cuboid2d is None): 217 | return 218 | 219 | super(Cuboid2dViz, self).on_draw() 220 | # print('Cuboid2dViz - on_draw - vertices: {}'.format(self.vertices)) 221 | 222 | glEnableClientState(GL_VERTEX_ARRAY) 223 | glEnableClientState(GL_COLOR_ARRAY) 224 | # glEnable(GL_POLYGON_SMOOTH) 225 | 226 | glPolygonMode(GL_FRONT_AND_BACK, RenderMode.normal) 227 | 228 | glVertexPointer(2, GL_FLOAT, 0, self.vertex_gl_array) 229 | 230 | glEnable(GL_LINE_SMOOTH) 231 | glLineWidth(3.0) 232 | glColorPointer(4, GL_UNSIGNED_BYTE, 0, self.edge_colors_gl_array) 233 | # TODO: May want to use GL_LINE_STRIP or GL_LINE_LOOP 234 | glDrawElements(GL_LINES, len(self.indices_line_gl), GL_UNSIGNED_BYTE, self.indices_line_gl_array) 235 | glDisable(GL_LINE_SMOOTH) 236 | 237 | glColorPointer(4, GL_UNSIGNED_BYTE, 0, self.vertex_color_gl_array) 238 | glPointSize(10.0) 239 | glDrawElements(GL_POINTS, len(self.indices_point_gl_array), GL_UNSIGNED_BYTE, self.indices_point_gl_array) 240 | glPointSize(1.0) 241 | 242 | # Deactivate vertex arrays after drawing 243 | glDisableClientState(GL_VERTEX_ARRAY) 244 | glDisableClientState(GL_COLOR_ARRAY) 245 | # glDisable(GL_POLYGON_SMOOTH) 246 | -------------------------------------------------------------------------------- /nvdu/viz/image_draw.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) 2018 NVIDIA Corporation. All rights reserved. 2 | # This work is licensed under a Creative Commons Attribution-NonCommercial-ShareAlike 4.0 International 3 | # License. (https://creativecommons.org/licenses/by-nc-sa/4.0/legalcode 4 | 5 | import math 6 | import sys 7 | import cv2 8 | 9 | from nvdu.core.cuboid import * 10 | 11 | def is_point_valid(point): 12 | if (point is None): 13 | return False 14 | if (math.isnan(point[0]) or math.isnan(point[1])): 15 | return False 16 | # NOTE: sometime the value get too big and we run into error: 17 | # OverflowError: Python int too large to convert to C long 18 | # if (math.fabs(point[0]) >= sys.maxsize) or (math.fabs(point[1]) >= sys.maxsize): 19 | if (math.fabs(point[0]) >= 10000) or (math.fabs(point[1]) >= 10000): 20 | return False 21 | return True 22 | 23 | # This module contains all the functions related to drawing on image 24 | def draw_cuboid2d(image, cuboid2d, color, line_thickness=1, point_size=1): 25 | if (image is None) or (cuboid2d is None): 26 | return 27 | 28 | # print("image: {} - image.shape: {}".format(image, image.shape)) 29 | 30 | line_type = cv2.LINE_AA 31 | # Draw the lines edge of the cuboid 32 | for line in CuboidLineIndexes: 33 | vi0, vi1 = line 34 | v0 = cuboid2d.get_vertex(vi0) 35 | v1 = cuboid2d.get_vertex(vi1) 36 | # print("draw line - v0: {} - v1: {}".format(v0, v1)) 37 | if (is_point_valid(v0) and is_point_valid(v1)): 38 | v0 = (int(v0[0]), int(v0[1])) 39 | v1 = (int(v1[0]), int(v1[1])) 40 | # print("draw line - v0: {} - v1: {}".format(v0, v1)) 41 | 42 | cv2.line(image, v0, v1, color, line_thickness, line_type) 43 | 44 | # Draw circle at each corner vertices of the cuboid 45 | thickness = -1 46 | # TODO: Highlight the top front vertices 47 | for vertex_index in range(CuboidVertexType.TotalVertexCount): 48 | vertex = cuboid2d.get_vertex(vertex_index) 49 | if (not is_point_valid(vertex)): 50 | continue 51 | 52 | point = (int(vertex[0]), int(vertex[1])) 53 | cv2.circle(image, point, point_size, color, thickness, line_type) 54 | 55 | if (vertex_index == CuboidVertexType.FrontTopRight): 56 | cv2.circle(image, point, point_size, (0,0,0), int(point_size / 2), line_type) 57 | elif (vertex_index == CuboidVertexType.FrontTopLeft): 58 | cv2.circle(image, point, point_size, (0,0,0), 1, line_type) 59 | 60 | 61 | -------------------------------------------------------------------------------- /nvdu/viz/mesh.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) 2018 NVIDIA Corporation. All rights reserved. 2 | # This work is licensed under a Creative Commons Attribution-NonCommercial-ShareAlike 4.0 International 3 | # License. (https://creativecommons.org/licenses/by-nc-sa/4.0/legalcode 4 | 5 | import os 6 | from os import path 7 | # import asyncio 8 | from pyrr import Quaternion, Matrix44, Vector3, euler 9 | import numpy as np 10 | from pyglet.gl import * 11 | from pyglet.gl.gl import * 12 | from pyglet.gl.glu import * 13 | from ctypes import * 14 | import pywavefront 15 | import pywavefront.visualization 16 | 17 | from nvdu.core.mesh import * 18 | from .utils3d import * 19 | from .scene_object import * 20 | from .pivot_axis import * 21 | 22 | class Model3dManager(object): 23 | def __init__(self): 24 | self.model_map = {} 25 | 26 | def get_model(self, model_path, auto_load = True): 27 | if (model_path in self.model_map): 28 | return self.model_map[model_path] 29 | 30 | if (auto_load): 31 | self.load_model(model_path) 32 | 33 | return None 34 | 35 | def load_model(self, model_path): 36 | new_model = self.load_model_from_file(model_path) 37 | self.model_map[model_path] = new_model 38 | return new_model 39 | 40 | # def load_model_list(self, model_paths): 41 | # # print("load_model_list: {}".format(model_paths)) 42 | # tasks = [] 43 | # for check_path in model_paths: 44 | # if not (check_path in self.model_map): 45 | # tasks.append(asyncio.ensure_future(self.load_model(check_path))) 46 | # if (len(tasks) > 0): 47 | # # print("Load all the models: {}".format(model_paths)) 48 | # self.loop.run_until_complete(asyncio.wait(tasks)) 49 | 50 | def load_model_from_file(self, model_file_path): 51 | if (path.exists(model_file_path)): 52 | print("Model3dManager::load_model_from_file: {}".format(model_file_path)) 53 | return pywavefront.Wavefront(model_file_path) 54 | else: 55 | print("Model3dManager::load_model_from_file - can NOT find 3d model: {}".format(model_file_path)) 56 | return None 57 | 58 | GlobalModelManager = Model3dManager() 59 | 60 | class MeshViz(SceneObjectViz3d): 61 | def __init__(self, mesh_obj): 62 | super(MeshViz, self).__init__(mesh_obj) 63 | 64 | self.mesh_obj = mesh_obj 65 | self.mesh_model = None 66 | 67 | # pivot_size = [10, 10, 10] 68 | # self.pivot_axis = PivotAxis(pivot_size) 69 | self.pivot_axis = None 70 | self.ignore_initial_matrix = False 71 | 72 | def on_draw(self): 73 | super(MeshViz, self).on_draw() 74 | 75 | if (self.mesh_model is None): 76 | self.mesh_model = GlobalModelManager.get_model(self.mesh_obj.source_file_path) 77 | 78 | if (self.mesh_model): 79 | if (self.pivot_axis): 80 | self.pivot_axis.draw() 81 | 82 | glPolygonMode(GL_FRONT_AND_BACK, self.render_mode) 83 | 84 | if (not self.ignore_initial_matrix): 85 | mesh_initial_matrix = self.mesh_obj.get_initial_matrix() 86 | # print("mesh_initial_matrix: {}".format(mesh_initial_matrix)) 87 | glMultMatrixf(get_opengl_matrixf(mesh_initial_matrix)) 88 | 89 | # TODO: Need to get the color from the object settings 90 | glColor4f(1.0, 1.0, 0.0, 0.5) 91 | self.mesh_model.draw() 92 | glColor4f(1.0, 1.0, 1.0, 1.0) 93 | glPolygonMode(GL_FRONT_AND_BACK, RenderMode.normal) 94 | -------------------------------------------------------------------------------- /nvdu/viz/nvdu_visualizer.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) 2018 NVIDIA Corporation. All rights reserved. 2 | # This work is licensed under a Creative Commons Attribution-NonCommercial-ShareAlike 4.0 International 3 | # License. (https://creativecommons.org/licenses/by-nc-sa/4.0/legalcode 4 | 5 | import future 6 | import pyrr 7 | import pywavefront 8 | import numpy as np 9 | from ctypes import * 10 | import json 11 | from os import listdir, path 12 | import pickle 13 | import cv2 14 | import pyglet 15 | 16 | from nvdu.core.nvdu_data import * 17 | from .camera import * 18 | from .cuboid import * 19 | from .viewport import * 20 | from .pointcloud import * 21 | # from .pivot_axis import * 22 | from .mesh import * 23 | from .background_image import * 24 | 25 | # =============================== Helper functions =============================== 26 | # =============================== Data parsing =============================== 27 | # =============================== Dataset Settings =============================== 28 | # =============================== AnnotatedObjectViz =============================== 29 | # Class contain annotation data of each object in the scene 30 | class AnnotatedObjectViz(object): 31 | def __init__(self, annotated_object_info): 32 | self.object_info = annotated_object_info 33 | object_settings = self.object_info.object_settings 34 | # self.bb2d = BoundingBox() 35 | self.cuboid2d = Cuboid2dViz(self.object_info.cuboid2d, object_settings.class_color) 36 | self.cuboid3d = Cuboid3dViz(self.object_info.cuboid3d, object_settings.class_color) 37 | 38 | self.pivot_axis = PivotAxis(self.object_info.pivot_axis) 39 | self.mesh = MeshViz(self.object_info.mesh) 40 | 41 | raw_keypoints = self.object_info.keypoints 42 | keypoint_locations = [] 43 | 44 | if not (raw_keypoints is None): 45 | for check_keypoint in raw_keypoints: 46 | # TODO: don't hardcode the key name here and use const instead 47 | if not (check_keypoint is None) and ('projected_location' in check_keypoint): 48 | check_keypoint_location = check_keypoint['projected_location'] 49 | keypoint_locations.append(check_keypoint_location) 50 | 51 | # print("raw_keypoints: {}".format(raw_keypoints)) 52 | # print("keypoint_locations: {}".format(keypoint_locations)) 53 | self.keypoint2d = PointCloud2d(keypoint_locations) 54 | 55 | self.is_modified = False 56 | 57 | def update_transform(self): 58 | # print('update_transform: location: {} - quaternion: {}'.format(self.location, self.quaternion)) 59 | should_show = not (self.location is None) and not (self.quaternion is None) 60 | if (self.mesh): 61 | self.mesh.set_visibility(should_show) 62 | 63 | if (self.cuboid3d): 64 | self.cuboid3d.set_visibility(should_show) 65 | 66 | if (self.pivot_axis): 67 | self.pivot_axis.set_visibility(should_show) 68 | 69 | self.is_modified = False 70 | 71 | def draw(self, visualizer_settings=None): 72 | if (self.is_modified): 73 | self.update_transform() 74 | 75 | glPolygonMode(GL_FRONT_AND_BACK, visualizer_settings.render_mode) 76 | if ((visualizer_settings is None) or visualizer_settings.show_mesh) and self.mesh: 77 | self.mesh.draw() 78 | 79 | if ((visualizer_settings is None) or visualizer_settings.show_cuboid3d) and self.cuboid3d: 80 | self.cuboid3d.draw() 81 | # self.cuboid2d.draw() 82 | 83 | if ((visualizer_settings is None) or visualizer_settings.show_pivot_axis) and self.pivot_axis: 84 | self.pivot_axis.draw() 85 | 86 | def update_settings(self, visualizer_settings=None): 87 | if (visualizer_settings is None): 88 | return 89 | 90 | have_valid_transform = not (self.object_info.location is None) and not (self.object_info.quaternion is None) 91 | 92 | if (self.is_modified): 93 | self.update_transform() 94 | 95 | if self.mesh: 96 | self.mesh.set_visibility(visualizer_settings.show_mesh and have_valid_transform) 97 | self.mesh.render_mode = visualizer_settings.render_mode 98 | 99 | if self.cuboid3d: 100 | self.cuboid3d.set_visibility(visualizer_settings.show_cuboid3d and have_valid_transform) 101 | 102 | if self.cuboid2d: 103 | # print('update_settings: self.cuboid2d = {} - visualizer_settings.show_cuboid2d = {}'.format(self.cuboid2d, visualizer_settings.show_cuboid2d)) 104 | self.cuboid2d.set_visibility(visualizer_settings.show_cuboid2d) 105 | 106 | if self.pivot_axis: 107 | self.pivot_axis.set_visibility(visualizer_settings.show_pivot_axis and have_valid_transform) 108 | 109 | if self.keypoint2d: 110 | # print("visualizer_settings.show_keypoint2d: {}".format(visualizer_settings.show_keypoint2d)) 111 | self.keypoint2d.set_visibility(visualizer_settings.show_keypoint2d) 112 | 113 | # =============================== AnnotatedSceneViz =============================== 114 | class AnnotatedSceneViz(object): 115 | """Class contain annotation data of a scene""" 116 | def __init__(self, annotated_scene_info): 117 | self._scene_info = annotated_scene_info 118 | self.camera_intrinsics = self._scene_info.camera_intrinsics 119 | self._object_vizs = [] 120 | for check_object in self._scene_info.objects: 121 | if not (check_object is None): 122 | new_object_viz = AnnotatedObjectViz(check_object) 123 | self._object_vizs.append(new_object_viz) 124 | 125 | dataset_settings = self._scene_info.dataset_settings 126 | if not (dataset_settings is None or dataset_settings.exporter_settings is None): 127 | img_width, img_height = dataset_settings.exporter_settings.captured_image_size 128 | # NOTE: Fallback to a default resolution. 129 | else: 130 | img_width = self.camera_intrinsics.res_width 131 | img_height = self.camera_intrinsics.res_height 132 | # img_width, img_height = 640, 480 133 | # print("Image size: {} x {}".format(img_width, img_height)) 134 | # print("AnnotatedSceneViz - dataset_settings: {} - dataset_settings.exporter_settings: {}".format( 135 | # dataset_settings, dataset_settings.exporter_settings)) 136 | # print("AnnotatedSceneViz - img_width: {} - img_height: {}".format(img_width, img_height)) 137 | # print("_scene_info.image_data = {}".format(self._scene_info.image_data)) 138 | # print("object_vizs = {}".format(self._object_vizs)) 139 | 140 | # print("AnnotatedSceneViz: img_width = {} - img_height = {} - image_data: {}".format(img_width, img_height, self._scene_info.image_data.shape)) 141 | 142 | if not (self._scene_info.image_data is None): 143 | self.background_image = BackgroundImage.create_from_numpy_image_data(self._scene_info.image_data, img_width, img_height) 144 | else: 145 | self.background_image = None 146 | 147 | info_str = self._scene_info.get_scene_info_str() 148 | # print("Scene info: {}".format(info_str)) 149 | self.info_text = pyglet.text.Label(info_str, 150 | font_size=16, 151 | x=0, y=0, 152 | # color=(255, 0, 0, 255), 153 | anchor_x='left', anchor_y='baseline') 154 | 155 | def set_image_data(self, new_image_numpy_data): 156 | img_width, img_height = self.dataset_settings.exporter_settings.captured_image_size 157 | print("set_image_data - img_width: {} - img_height: {}".format(img_width, img_height)) 158 | if (self.background_image is None): 159 | self.background_image = BackgroundImage.create_from_numpy_image_data(new_image_numpy_data, img_width, img_height) 160 | else: 161 | self.background_image.load_image_data_from_numpy(new_image_numpy_data) 162 | 163 | def draw(self, visualizer_settings=None): 164 | for obj_viz in self._object_vizs: 165 | # print('draw object: {}'.format(obj.source_file_path)) 166 | obj_viz.draw(visualizer_settings) 167 | 168 | def set_text_color(self, new_text_color): 169 | self.info_text.color = new_text_color 170 | 171 | def update_settings(self, visualizer_settings=None): 172 | if (visualizer_settings): 173 | for obj_viz in self._object_vizs: 174 | obj_viz.update_settings(visualizer_settings) 175 | 176 | # TODO: Create a new viz object to handle the text overlay 177 | # if (self.info_text): 178 | # self.info_text.set_visibility() 179 | 180 | # =============================== Visualizer =============================== 181 | class VisualizerSettings(object): 182 | def __init__(self): 183 | self.render_mode = RenderMode.normal 184 | self.show_mesh = True 185 | self.show_pivot_axis = True 186 | self.show_cuboid3d = True 187 | self.show_cuboid2d = True 188 | self.show_bb2d = True 189 | self.show_info_text = True 190 | self.show_keypoint2d = False 191 | self.ignore_initial_matrix = False 192 | 193 | # TODO: Find a way to use template for all these flags 194 | def toggle_mesh(self): 195 | self.show_mesh = not self.show_mesh 196 | 197 | def toggle_pivot_axis(self): 198 | self.show_pivot_axis = not self.show_pivot_axis 199 | 200 | def toggle_cuboid3d(self): 201 | self.show_cuboid3d = not self.show_cuboid3d 202 | 203 | def toggle_cuboid2d(self): 204 | self.show_cuboid2d = not self.show_cuboid2d 205 | 206 | def toggle_bb2d(self): 207 | self.show_bb2d = not self.show_bb2d 208 | 209 | def toggle_keypoint2d(self): 210 | print("toggle_keypoint2d==========================") 211 | self.show_keypoint2d = not self.show_keypoint2d 212 | 213 | def toggle_info_overlay(self): 214 | self.show_info_text = not self.show_info_text 215 | 216 | class NVDUVisualizer(): 217 | def __init__(self): 218 | self.render_mode = RenderMode.normal 219 | self.camera = Camera() 220 | 221 | self.dataset_settings = None 222 | self.visualizer_settings = VisualizerSettings() 223 | self.annotated_scene = None 224 | self.scene_viz = None 225 | 226 | self.viewport = Viewport(None) 227 | self.viewport.size = [512, 512] 228 | 229 | def draw(self): 230 | if (self.annotated_scene is None) or (self.scene_viz is None): 231 | return 232 | 233 | self.viewport.clear() 234 | 235 | # TODO: Should let the AnnotatedSceneViz handle all these draw logic 236 | self.viewport.scene_bg.add_object(self.scene_viz.background_image) 237 | 238 | self.viewport.scene3d.camera = self.camera 239 | if (not self.scene_viz.camera_intrinsics is None): 240 | self.viewport.scene3d.camera.set_instrinsic_settings(self.scene_viz.camera_intrinsics) 241 | # self.viewport.scene3d.camera.set_fovx(self.scene_viz.camera_fovx) 242 | 243 | #TODO: Move this code to a separated function 244 | # mesh_paths = [] 245 | # for obj_viz in self.scene_viz._object_vizs: 246 | # mesh_paths.append(obj_viz.mesh.mesh_obj.source_file_path) 247 | # GlobalModelManager.load_model_list(mesh_paths) 248 | # print("NVDUVisualizer - draw - mesh_paths: {}".format(mesh_paths)) 249 | 250 | for obj in self.scene_viz._object_vizs: 251 | if (obj.mesh): 252 | obj.mesh.ignore_initial_matrix = self.visualizer_settings.ignore_initial_matrix 253 | self.viewport.scene3d.add_object(obj.mesh) 254 | self.viewport.scene3d.add_object(obj.cuboid3d) 255 | self.viewport.scene3d.add_object(obj.pivot_axis) 256 | self.viewport.scene_overlay.add_object(obj.cuboid2d) 257 | self.viewport.scene_overlay.add_object(obj.keypoint2d) 258 | 259 | if (self.visualizer_settings.show_info_text): 260 | self.scene_viz.info_text.draw() 261 | 262 | self.scene_viz.update_settings(self.visualizer_settings) 263 | 264 | self.viewport.draw() 265 | 266 | # ========================== CONTROL ========================== 267 | def toggle_cuboid2d_overlay(self): 268 | self.visualizer_settings.toggle_cuboid2d() 269 | 270 | def toggle_cuboid3d_overlay(self): 271 | self.visualizer_settings.toggle_cuboid3d() 272 | 273 | def toggle_object_overlay(self): 274 | self.visualizer_settings.toggle_mesh() 275 | 276 | def toggle_pivot_axis(self): 277 | self.visualizer_settings.toggle_pivot_axis() 278 | 279 | def toggle_info_overlay(self): 280 | self.visualizer_settings.toggle_info_overlay() 281 | 282 | def toggle_keypoint2d_overlay(self): 283 | self.visualizer_settings.toggle_keypoint2d() 284 | 285 | def set_render_mode(self, new_render_mode): 286 | self.visualizer_settings.render_mode = new_render_mode 287 | 288 | def set_text_color(self, new_text_color): 289 | self.scene_viz.set_text_color(new_text_color) 290 | 291 | def visualize_dataset_frame(self, in_dataset, in_frame_index = 0): 292 | frame_image_file_path, frame_data_file_path = in_dataset.get_frame_file_path_from_index(in_frame_index) 293 | if not path.exists(frame_image_file_path): 294 | print("Can't find image file for frame: {} - {}".format(in_frame_index, frame_image_file_path)) 295 | return 296 | if not path.exists(frame_data_file_path): 297 | print("Can't find annotation file for frame: {} - {}".format(in_frame_index, frame_data_file_path)) 298 | return 299 | 300 | print("visualize_dataset_frame: frame_image_file_path: {} - frame_data_file_path: {}".format( 301 | frame_image_file_path, frame_data_file_path)) 302 | 303 | frame_scene_data = AnnotatedSceneInfo.create_from_file(self.dataset_settings, 304 | frame_data_file_path, frame_image_file_path) 305 | self.visualize_scene(frame_scene_data) 306 | 307 | def set_scene_data(self, new_scene_data): 308 | self.annotated_scene = new_scene_data 309 | self.scene_viz = AnnotatedSceneViz(self.annotated_scene) 310 | 311 | def visualize_scene(self, annotated_scene): 312 | self.set_scene_data(annotated_scene) 313 | self.draw() 314 | -------------------------------------------------------------------------------- /nvdu/viz/nvdu_viz_window.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) 2018 NVIDIA Corporation. All rights reserved. 2 | # This work is licensed under a Creative Commons Attribution-NonCommercial-ShareAlike 4.0 International 3 | # License. (https://creativecommons.org/licenses/by-nc-sa/4.0/legalcode 4 | 5 | import future 6 | import os 7 | from os import path 8 | import pyglet 9 | from pyglet.window import key 10 | 11 | from nvdu.viz.nvdu_visualizer import * 12 | from nvdu.core.nvdu_data import * 13 | 14 | class NVDUVizWindow(pyglet.window.Window): 15 | DEFAULT_EXPORT_DIR = "viz" 16 | 17 | def __init__(self, width, height, caption =''): 18 | super(NVDUVizWindow, self).__init__(width, height, caption) 19 | 20 | self._org_caption = caption 21 | print('Window created: width = {} - height = {} - title = {}'.format(self.width, self.height, self.caption)) 22 | # print('Window context: {} - config: {}'.format(self.context, self.context.config)) 23 | 24 | self.frame_index = 0 25 | self.visualizer = NVDUVisualizer() 26 | 27 | self.auto_change_frame = False 28 | self.auto_fps = 0 29 | 30 | self._dataset = None 31 | self.export_dir = "" 32 | self._should_export = False 33 | 34 | @property 35 | def dataset(self): 36 | return self._dataset 37 | 38 | @dataset.setter 39 | def dataset(self, new_dataset): 40 | self._dataset = new_dataset 41 | frame_count = self._dataset.scan() 42 | print("Number of frames in the dataset: {}".format(frame_count)) 43 | 44 | @property 45 | def should_export(self): 46 | # Can export if the export directory is valid 47 | # return not (not self.export_dir) 48 | return self._should_export and self.export_dir 49 | 50 | @should_export.setter 51 | def should_export(self, new_export): 52 | self._should_export = new_export 53 | 54 | def set_caption_postfix(self, postfix): 55 | self.set_caption(self._org_caption + postfix) 56 | 57 | def setup(self): 58 | glClearColor(0, 0, 0, 1) 59 | glEnable(GL_DEPTH_TEST) 60 | # glEnable(GL_DEPTH_CLAMP) 61 | glFrontFace(GL_CCW) 62 | # glFrontFace(GL_CW) 63 | glEnable(GL_BLEND) 64 | glBlendFunc(GL_SRC_ALPHA, GL_ONE_MINUS_SRC_ALPHA) 65 | 66 | self.visualize_current_frame() 67 | 68 | def on_draw(self): 69 | # Clear the current GL Window 70 | self.clear() 71 | 72 | self.update_text_color() 73 | 74 | if (self.visualizer): 75 | self.visualizer.draw() 76 | if (self.should_export): 77 | self.save_current_viz_frame() 78 | 79 | def on_resize(self, width, height): 80 | super(NVDUVizWindow, self).on_resize(width, height) 81 | # set the Viewport 82 | glViewport(0, 0, width, height) 83 | 84 | self.visualizer.viewport.size = [width, height] 85 | 86 | # new_cam_intrinsic_settings = CameraIntrinsicSettings.from_perspective_fov_horizontal(width, height, CAMERA_FOV_HORIZONTAL) 87 | # self.visualizer.camera.set_instrinsic_settings(new_cam_intrinsic_settings) 88 | 89 | def set_camera_intrinsic_settings(self, new_cam_intrinsic_settings): 90 | # print("set_camera_intrinsic_settings: {}".format(new_cam_intrinsic_settings)) 91 | self.visualizer.camera.set_instrinsic_settings(new_cam_intrinsic_settings) 92 | 93 | # Save the current screenshot to a file 94 | def save_screenshot(self, export_path): 95 | screen_image = pyglet.image.get_buffer_manager().get_color_buffer() 96 | screen_image.save(export_path) 97 | print("save_screenshot: {}".format(export_path)) 98 | 99 | def save_current_viz_frame(self): 100 | # TODO: Should ignore? if the visualized frame already exist 101 | current_frame_name = self.dataset.get_frame_name_from_index(self.frame_index) 102 | # TODO: May need to add config to control the viz postfix 103 | viz_frame_file_name = current_frame_name + "_viz.png" 104 | export_viz_path = path.join(self.export_dir, viz_frame_file_name) 105 | if not path.exists(self.export_dir): 106 | os.makedirs(self.export_dir) 107 | 108 | self.save_screenshot(export_viz_path) 109 | 110 | # ========================== DATA PROCESSING ========================== 111 | def visualize_current_frame(self): 112 | print('Visualizing frame: {}'.format(self.frame_index)) 113 | self.visualizer.visualize_dataset_frame(self.dataset, self.frame_index) 114 | 115 | def set_frame_index(self, new_frame_index): 116 | total_frame_count = self.dataset.frame_count 117 | if (new_frame_index < 0): 118 | new_frame_index += total_frame_count 119 | # TODO: May need to update the total frame count when it's invalid 120 | if (total_frame_count > 0): 121 | new_frame_index = new_frame_index % total_frame_count 122 | 123 | if (self.frame_index != new_frame_index): 124 | self.frame_index = new_frame_index 125 | self.visualize_current_frame() 126 | 127 | # ========================== INPUT CONTROL ========================== 128 | def on_key_press(self, symbol, modifiers): 129 | super(NVDUVizWindow, self).on_key_press(symbol, modifiers) 130 | if (symbol == key.F3): 131 | self.toggle_cuboid2d_overlay() 132 | if (symbol == key.F4): 133 | self.toggle_cuboid3d_overlay() 134 | elif (symbol == key.F5): 135 | self.toggle_object_overlay() 136 | elif (symbol == key.F6): 137 | self.toggle_pivot() 138 | elif (symbol == key.F7): 139 | self.toggle_info_overlay() 140 | elif (symbol == key.F8): 141 | self.toggle_keypoint2d_overlay() 142 | elif (symbol == key.F12): 143 | self.toggle_export_viz_frame() 144 | elif (symbol == key._1): 145 | self.visualizer.set_render_mode(RenderMode.normal) 146 | elif (symbol == key._2): 147 | self.visualizer.set_render_mode(RenderMode.wire_frame) 148 | elif (symbol == key._3): 149 | self.visualizer.set_render_mode(RenderMode.point) 150 | elif (symbol == key.SPACE): 151 | self.toggle_auto_change_frame() 152 | 153 | def on_text_motion(self, motion): 154 | if motion == key.LEFT: 155 | self.set_frame_index(self.frame_index - 1) 156 | elif motion == key.RIGHT: 157 | self.visualize_next_frame() 158 | elif motion == key.UP: 159 | self.set_frame_index(self.frame_index + 100) 160 | elif motion == key.DOWN: 161 | self.set_frame_index(self.frame_index - 100) 162 | 163 | def visualize_next_frame(self, dt=0): 164 | self.set_frame_index(self.frame_index + 1) 165 | 166 | def toggle_export_viz_frame(self): 167 | self._should_export = not self._should_export 168 | if (self._should_export and not self.export_dir): 169 | self.export_dir = NVDUVizWindow.DEFAULT_EXPORT_DIR 170 | 171 | self.update_text_color() 172 | 173 | def update_text_color(self): 174 | # Use different color when we are exporting visualized frame 175 | if (self.should_export): 176 | self.set_caption_postfix(" - Exporting ...") 177 | self.visualizer.set_text_color((255, 0, 0, 255)) 178 | else: 179 | self.set_caption_postfix("") 180 | self.visualizer.set_text_color((255, 255, 255, 255)) 181 | 182 | def toggle_cuboid2d_overlay(self): 183 | self.visualizer.toggle_cuboid2d_overlay() 184 | 185 | def toggle_cuboid3d_overlay(self): 186 | self.visualizer.toggle_cuboid3d_overlay() 187 | 188 | def toggle_object_overlay(self): 189 | self.visualizer.toggle_object_overlay() 190 | 191 | def toggle_pivot(self): 192 | self.visualizer.toggle_pivot_axis() 193 | 194 | def toggle_keypoint2d_overlay(self): 195 | self.visualizer.toggle_keypoint2d_overlay() 196 | 197 | def toggle_info_overlay(self): 198 | self.visualizer.toggle_info_overlay() 199 | 200 | def toggle_auto_change_frame(self): 201 | self.set_auto_change_frame(not self.auto_change_frame) 202 | 203 | def set_auto_fps(self, new_fps): 204 | self.auto_fps = new_fps 205 | if (new_fps <= 0): 206 | self.set_auto_change_frame(False) 207 | else: 208 | self.set_auto_change_frame(True) 209 | 210 | def set_auto_change_frame(self, new_bool): 211 | if (self.auto_change_frame == new_bool): 212 | return 213 | 214 | self.auto_change_frame = new_bool 215 | if (self.auto_change_frame): 216 | print("Start auto changing frame ...") 217 | wait_duration = 1.0 / self.auto_fps 218 | pyglet.clock.schedule_interval(self.visualize_next_frame, wait_duration) 219 | else: 220 | print("Stop auto changing frame ...") 221 | pyglet.clock.unschedule(self.visualize_next_frame) 222 | 223 | -------------------------------------------------------------------------------- /nvdu/viz/pivot_axis.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) 2018 NVIDIA Corporation. All rights reserved. 2 | # This work is licensed under a Creative Commons Attribution-NonCommercial-ShareAlike 4.0 International 3 | # License. (https://creativecommons.org/licenses/by-nc-sa/4.0/legalcode 4 | 5 | import numpy as np 6 | from pyrr import Quaternion, Matrix44, Vector3, euler 7 | from pyglet.gl import * 8 | from pyglet.gl.gl import * 9 | from pyglet.gl.glu import * 10 | from ctypes import * 11 | 12 | from .scene_object import * 13 | 14 | class PivotAxis(SceneObjectViz3d): 15 | # Create a pivot axis object with 3 axes each can have different length 16 | def __init__(self, in_pivot_axis_obj, in_line_width = 5.0): 17 | super(PivotAxis, self).__init__(in_pivot_axis_obj) 18 | self.pivot_obj = in_pivot_axis_obj 19 | self.origin_loc = self.pivot_obj.origin_loc 20 | self.line_width = in_line_width 21 | self.render_stipple_line = False 22 | # List of color for each vertices of the box 23 | self.colors = [ 24 | [255, 0, 0], 25 | [0, 255, 0], 26 | [0, 0, 255], 27 | ] 28 | 29 | def on_draw(self): 30 | super(PivotAxis, self).on_draw() 31 | 32 | glEnable(GL_LINE_SMOOTH) 33 | if (self.render_stipple_line): 34 | glEnable(GL_LINE_STIPPLE) 35 | line_pattern = 0x00ff 36 | line_tripple_factor = 1 37 | glLineStipple(line_tripple_factor, line_pattern); 38 | 39 | glLineWidth(self.line_width) 40 | glBegin(GL_LINES) 41 | # X axis 42 | glColor3ub(self.colors[0][0], self.colors[0][1], self.colors[0][2]) 43 | glVertex3f(self.origin_loc[0], self.origin_loc[1], self.origin_loc[2]) 44 | glVertex3f(self.pivot_obj.x_axis[0], self.pivot_obj.x_axis[1], self.pivot_obj.x_axis[2]) 45 | 46 | # Y axis 47 | glColor3ub(self.colors[1][0], self.colors[1][1], self.colors[1][2]) 48 | glVertex3f(self.origin_loc[0], self.origin_loc[1], self.origin_loc[2]) 49 | glVertex3f(self.pivot_obj.y_axis[0], self.pivot_obj.y_axis[1], self.pivot_obj.y_axis[2]) 50 | 51 | # Z axis 52 | glColor3ub(self.colors[2][0], self.colors[2][1], self.colors[2][2]) 53 | glVertex3f(self.origin_loc[0], self.origin_loc[1], self.origin_loc[2]) 54 | glVertex3f(self.pivot_obj.z_axis[0], self.pivot_obj.z_axis[1], self.pivot_obj.z_axis[2]) 55 | glEnd() 56 | 57 | glColor3ub(255, 255, 255) 58 | glLineWidth(1.0) 59 | 60 | if (self.render_stipple_line): 61 | glDisable(GL_LINE_STIPPLE) 62 | glDisable(GL_LINE_SMOOTH) -------------------------------------------------------------------------------- /nvdu/viz/pointcloud.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) 2018 NVIDIA Corporation. All rights reserved. 2 | # This work is licensed under a Creative Commons Attribution-NonCommercial-ShareAlike 4.0 International 3 | # License. (https://creativecommons.org/licenses/by-nc-sa/4.0/legalcode 4 | 5 | from pyrr import Quaternion, Matrix44, Vector3, euler 6 | import numpy as np 7 | from pyglet.gl import * 8 | from ctypes import * 9 | 10 | from .scene_object import * 11 | from nvdu.core.cuboid import * 12 | 13 | # ========================= PointCloud2d ========================= 14 | class PointCloud2d(SceneObjectVizBase): 15 | # Create a box with a certain size 16 | def __init__(self, point2d_list=[], in_color=None): 17 | super(PointCloud2d, self).__init__(point2d_list) 18 | self.vertices = point2d_list 19 | self.color = in_color 20 | self.generate_vertexes_buffer() 21 | 22 | def generate_vertexes_buffer(self): 23 | self.vertices_gl = [] 24 | vertex_count = len(self.vertices) 25 | for i in range(0, vertex_count): 26 | vertex = self.vertices[i] 27 | if (not vertex is None): 28 | self.vertices_gl.append(vertex[0]) 29 | self.vertices_gl.append(vertex[1]) 30 | 31 | self.indices_point_gl = [] 32 | for i in range (len(self.vertices)): 33 | if (not self.vertices[i] is None): 34 | self.indices_point_gl.append(i) 35 | 36 | self.vertex_gl_array = (GLfloat* len(self.vertices_gl))(*self.vertices_gl) 37 | self.indices_point_gl_array = (GLubyte* len(self.indices_point_gl))(*self.indices_point_gl) 38 | 39 | # print("PointCloud2d: {}".format(self.vertices_gl)) 40 | 41 | def on_draw(self): 42 | super(PointCloud2d, self).on_draw() 43 | 44 | # print("Drawing pointcloud: {}".format(self.vertices_gl)) 45 | 46 | glEnableClientState(GL_VERTEX_ARRAY) 47 | # glEnableClientState(GL_COLOR_ARRAY) 48 | 49 | glPolygonMode(GL_FRONT_AND_BACK, RenderMode.normal) 50 | 51 | glVertexPointer(2, GL_FLOAT, 0, self.vertex_gl_array) 52 | 53 | # glColorPointer(4, GL_UNSIGNED_BYTE, 0, self.vertex_color_gl_array) 54 | glPointSize(10.0) 55 | glDrawElements(GL_POINTS, len(self.indices_point_gl_array), GL_UNSIGNED_BYTE, self.indices_point_gl_array) 56 | glPointSize(1.0) 57 | 58 | # Deactivate vertex arrays after drawing 59 | glDisableClientState(GL_VERTEX_ARRAY) 60 | # glDisableClientState(GL_COLOR_ARRAY) 61 | -------------------------------------------------------------------------------- /nvdu/viz/scene.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) 2018 NVIDIA Corporation. All rights reserved. 2 | # This work is licensed under a Creative Commons Attribution-NonCommercial-ShareAlike 4.0 International 3 | # License. (https://creativecommons.org/licenses/by-nc-sa/4.0/legalcode 4 | 5 | from pyrr import Quaternion, Matrix44, Vector3, euler 6 | import numpy as np 7 | 8 | from .utils3d import * 9 | from nvdu.core import transform3d 10 | from .camera import * 11 | 12 | # A Scene manage all of the objects need to be rendered 13 | class Scene(object): 14 | def __init__(self, owner_viewport): 15 | self.viewport = owner_viewport 16 | # List of the objects in the scene - SceneObject type 17 | self.objects = [] 18 | # TODO: May need to derive from SceneObject 19 | # self._is_visible = True 20 | 21 | def clear(self): 22 | self.objects = [] 23 | 24 | def add_object(self, new_obj): 25 | # TODO: Make sure the new_obj is new and not already in the objects list? 26 | self.objects.append(new_obj) 27 | 28 | def draw(self): 29 | pass 30 | 31 | # ================================= Scene2d ================================= 32 | class Scene2d(Scene): 33 | def __init__(self, owner_viewport): 34 | super(Scene2d, self).__init__(owner_viewport) 35 | 36 | def draw(self): 37 | super(Scene2d, self).draw() 38 | 39 | viewport_width, viewport_height = self.viewport.size 40 | # print('Scene2d - {} - draw'.format(self)) 41 | 42 | # Disable depth when rendering 2d scene 43 | glPolygonMode(GL_FRONT_AND_BACK, GL_FILL) 44 | glDepthMask(False) 45 | # Use Orthographic camera in full viewport size for 2d scene 46 | glMatrixMode(GL_PROJECTION) 47 | glLoadIdentity() 48 | gluOrtho2D(0.0, viewport_width, 0.0, viewport_height) 49 | 50 | glMatrixMode(GL_MODELVIEW) 51 | glPushMatrix() 52 | glLoadIdentity() 53 | # NOTE: OpenCV 2d image coordinate system have Y going down 54 | # while OpenGL have Y going up => need to flip the Y axis 55 | # and add the viewport_height so the OpenCV coordinate appear right 56 | glTranslatef(0.0, viewport_height, 0.0) 57 | glMultMatrixf(get_opengl_matrixf(opencv_to_opengl_matrix)) 58 | 59 | # Render the objects in the scene 60 | for obj in self.objects: 61 | if (obj): 62 | obj.draw() 63 | 64 | glPopMatrix() 65 | glDepthMask(True) 66 | 67 | # ================================= Scene3d ================================= 68 | class Scene3d(Scene): 69 | def __init__(self, owner_viewport): 70 | super(Scene3d, self).__init__(owner_viewport) 71 | self.camera = None 72 | 73 | def draw(self): 74 | super(Scene3d, self).draw() 75 | 76 | # print('Scene3d - {} - draw'.format(self)) 77 | 78 | glPushMatrix() 79 | self.camera.draw() 80 | 81 | glMatrixMode(GL_MODELVIEW) 82 | glLoadIdentity() 83 | glPushMatrix() 84 | glMultMatrixf(get_opengl_matrixf(opencv_to_opengl_matrix)) 85 | 86 | # TODO: Sort the 3d objects in the scene from back to front (Z reducing) 87 | for child_object in self.objects: 88 | if child_object: 89 | child_object.draw() 90 | 91 | glPopMatrix() 92 | glPopMatrix() 93 | -------------------------------------------------------------------------------- /nvdu/viz/scene_object.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) 2018 NVIDIA Corporation. All rights reserved. 2 | # This work is licensed under a Creative Commons Attribution-NonCommercial-ShareAlike 4.0 International 3 | # License. (https://creativecommons.org/licenses/by-nc-sa/4.0/legalcode 4 | 5 | from pyrr import Quaternion, Matrix44, Vector3, euler 6 | import numpy as np 7 | 8 | from .utils3d import * 9 | 10 | from nvdu.core import transform3d 11 | from nvdu.core import scene_object 12 | 13 | class SceneObjectVizBase(object): 14 | def __init__(self, scene_object): 15 | self.scene_object = scene_object 16 | self.render_mode = RenderMode.normal 17 | self._is_visible = True 18 | 19 | def draw(self): 20 | if ((self.scene_object is None) or (not self.is_visible())): 21 | return 22 | 23 | self.on_draw() 24 | 25 | def on_draw(self): 26 | pass 27 | 28 | def is_visible(self): 29 | return self._is_visible 30 | 31 | def set_visibility(self, should_visible): 32 | self._is_visible = should_visible 33 | 34 | def hide(self): 35 | self._is_visible = False 36 | 37 | def show(self): 38 | self._is_visible = True 39 | 40 | def toggle_visibility(self): 41 | self._is_visible = not self._is_visible 42 | 43 | class SceneObjectViz3d(SceneObjectVizBase): 44 | def __init__(self, scene_object): 45 | super(SceneObjectViz3d, self).__init__(scene_object) 46 | 47 | def draw(self): 48 | if ((self.scene_object is None) or (not self.is_visible())): 49 | return 50 | 51 | # transform_matrix = self.scene_object.relative_transform.to_matrix() 52 | # glMultMatrixf(get_opengl_matrixf(transform_matrix)) 53 | # for child_object in self.child_objects: 54 | # if (child_object != None): 55 | # child_object.draw() 56 | 57 | glPushMatrix() 58 | 59 | world_transform_matrix = self.scene_object.get_world_transform_matrix() 60 | # print("{} - draw - world_transform_matrix: {}".format(self, world_transform_matrix)) 61 | glMultMatrixf(get_opengl_matrixf(world_transform_matrix)) 62 | 63 | self.on_draw() 64 | 65 | glPopMatrix() 66 | -------------------------------------------------------------------------------- /nvdu/viz/utils3d.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) 2018 NVIDIA Corporation. All rights reserved. 2 | # This work is licensed under a Creative Commons Attribution-NonCommercial-ShareAlike 4.0 International 3 | # License. (https://creativecommons.org/licenses/by-nc-sa/4.0/legalcode 4 | 5 | from pyrr import Quaternion, Matrix44, Vector3, euler 6 | import numpy as np 7 | from pyglet.gl.gl import * 8 | from ctypes import * 9 | 10 | # Get the openGL matrix (GLfloat* 16) from a Matrix44 type 11 | # NOTE: OpenGL use column major while OpenCV use row major 12 | def get_opengl_matrixf(in_mat44): 13 | return (GLfloat* 16)( 14 | in_mat44.m11, in_mat44.m12, in_mat44.m13, in_mat44.m14, 15 | in_mat44.m21, in_mat44.m22, in_mat44.m23, in_mat44.m24, 16 | in_mat44.m31, in_mat44.m32, in_mat44.m33, in_mat44.m34, 17 | in_mat44.m41, in_mat44.m42, in_mat44.m43, in_mat44.m44 18 | ) 19 | 20 | def convert_HFOV_to_VFOV(hfov, hw_ratio): 21 | # https://en.wikipedia.org/wiki/Field_of_view_in_video_games 22 | vfov = 2 * np.arctan(np.tan(np.deg2rad(hfov / 2)) * hw_ratio) 23 | 24 | return np.rad2deg(vfov) 25 | 26 | opencv_to_opengl_matrix = Matrix44([ 27 | [1.0, 0.0, 0.0, 0.0], 28 | [0.0, -1.0, 0.0, 0.0], 29 | [0.0, 0.0, -1.0, 0.0], 30 | [0.0, 0.0, 0.0, 1.0] 31 | ]) 32 | 33 | class RenderMode(): 34 | normal = GL_FILL 35 | wire_frame = GL_LINE 36 | point = GL_POINT 37 | -------------------------------------------------------------------------------- /nvdu/viz/viewport.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) 2018 NVIDIA Corporation. All rights reserved. 2 | # This work is licensed under a Creative Commons Attribution-NonCommercial-ShareAlike 4.0 International 3 | # License. (https://creativecommons.org/licenses/by-nc-sa/4.0/legalcode 4 | 5 | from .utils3d import * 6 | from .scene import * 7 | 8 | class Viewport(object): 9 | def __init__(self, context): 10 | self._context = context 11 | self._size = [0, 0] 12 | 13 | self.scene_bg = Scene2d(self) 14 | self.scene3d = Scene3d(self) 15 | self.scene_overlay = Scene2d(self) 16 | 17 | self.scenes = [ 18 | self.scene_bg, 19 | self.scene3d, 20 | self.scene_overlay 21 | ] 22 | 23 | @property 24 | def size(self): 25 | return self._size 26 | 27 | @size.setter 28 | def size(self, new_size): 29 | self._size[0] = new_size[0] 30 | self._size[1] = new_size[1] 31 | 32 | def clear(self): 33 | for scene in self.scenes: 34 | if (scene): 35 | scene.clear() 36 | 37 | def draw(self): 38 | for scene in self.scenes: 39 | if (scene): 40 | scene.draw() 41 | -------------------------------------------------------------------------------- /readme.md: -------------------------------------------------------------------------------- 1 | # Nvidia Dataset Utilities (NVDU) 2 | This project is a collection of Python scripts to help work with datasets for deep learning. For example, visualizing annotation data associated with captured sensor images generated by NVIDIA Deep learning Dataset Synthesizer (NDDS) https://github.com/NVIDIA/Dataset_Synthesizer. 3 | > **This module depends on OpenCV-python which currently doesn't work with Python 3.7; this module requires either Python 3.5 or 3.6.** 4 | 5 | ![](./NVDUIntro.png) 6 | 7 | *Example of a dataset frame visualized using NVDU, showing axes and 3D cuboids for annotated objects.* 8 | 9 | ## Table of Contents 10 | - [Nvidia Dataset Utilities (NVDU)](#nvidia-dataset-utilities-nvdu) 11 | - [Table of Contents](#table-of-contents) 12 | - [Install](#install) 13 | - [Install from pip:](#install-from-pip) 14 | - [Install from source code git repo:](#install-from-source-code-git-repo) 15 | - [nvdu_ycb](#nvdu_ycb) 16 | - [Usage](#usage) 17 | - [nvdu_viz](#nvdu_viz) 18 | - [Usage](#usage-1) 19 | - [Examples](#examples) 20 | - [Visualize a dataset generated by NDDS:](#visualize-a-dataset-generated-by-ndds) 21 | - [Visualize a set of images using different annotation data:](#visualize-a-set-of-images-using-different-annotation-data) 22 | - [Controls](#controls) 23 | - [Visualization options:](#visualization-options) 24 | - [Other:](#other) 25 | 26 | # Install 27 | ## Install from pip: 28 | `pip install nvdu` 29 | 30 | ## Install from source code git repo: 31 | **Clone the repo** 32 | 33 | _Using ssh path:_ 34 | ``` 35 | git clone ssh://git@github.com:12051/NVIDIA/Dataset_Utilities.git 36 | ``` 37 | _Using https path:_ 38 | ``` 39 | git clone https://github.com/NVIDIA/Dataset_Utilities.git 40 | ``` 41 | **Go inside the cloned repo's directory** 42 | ``` 43 | cd Dataset_Utilities 44 | ``` 45 | 46 | **Install locally** 47 | 48 | `pip install -e .` 49 | 50 | # nvdu_ycb 51 | _nvdu_ycb_ command help download, extract and align the YCB 3d models (which are used in the FAT dataset: http://research.nvidia.com/publication/2018-06_Falling-Things). 52 | ## Usage 53 | ``` 54 | usage: nvdu_ycb [-h] [-s] [-l] [ycb_object_name] 55 | 56 | NVDU YCB models Support 57 | 58 | positional arguments: 59 | ycb_object_name Name of the YCB object to check. 60 | 61 | optional arguments: 62 | -h, --help show this help message and exit 63 | -s, --setup Setup the YCB models for the FAT dataset 64 | -l, --list List all the supported YCB objects 65 | ``` 66 | 67 | *NOTE: If you don't run the `nvdu_ycb --setup` before trying to use nvdu_viz, the visualizer will not be able to find the 3d models of the YCB object to overlay.* 68 | 69 | # nvdu_viz 70 | _nvdu_viz_ command visualizes the annotated datasets using the NDDS format. 71 | ## Usage 72 | ``` 73 | nvdu_viz [-h] [-a DATA_ANNOT_DIR] [-s SIZE SIZE] 74 | [-o OBJECT_SETTINGS_PATH] [-c CAMERA_SETTINGS_PATH] 75 | [-m MODEL_DIR] [-n [NAME_FILTERS [NAME_FILTERS ...]]] 76 | [--fps FPS] [--auto_change] [-e EXPORT_DIR] [--auto_export] 77 | [--ignore_fixed_transform] 78 | [dataset_dir] 79 | 80 | NVDU Data Visualiser 81 | 82 | positional arguments: 83 | dataset_dir Dataset directory. This is where all the images 84 | (required) and annotation info (optional) are. 85 | Defaults to the current directory. 86 | 87 | optional arguments: 88 | -h, --help show this help message and exit. 89 | -a DATA_ANNOT_DIR, --data_annot_dir DATA_ANNOT_DIR 90 | Directory path - where to find the annotation data. 91 | Defaults to be the same directory as the dataset 92 | directory. 93 | -s SIZE SIZE, --size SIZE SIZE 94 | Window's size: [width, height]. If not specified then 95 | the window is sized to fit the resolution of the camera. 96 | -o OBJECT_SETTINGS_PATH, --object_settings_path OBJECT_SETTINGS_PATH 97 | Object settings file path. 98 | -c CAMERA_SETTINGS_PATH, --camera_settings_path CAMERA_SETTINGS_PATH 99 | Camera settings file path. 100 | -n [NAME_FILTERS [NAME_FILTERS ...]], --name_filters [NAME_FILTERS [NAME_FILTERS ...]] 101 | The name filter of each frame. e.g: *.png. 102 | --fps FPS How fast to automatically change frame. 103 | --auto_change When using this flag, the visualizer will automatically 104 | change the frame. 105 | -e EXPORT_DIR, --export_dir EXPORT_DIR 106 | Directory path - where to store the visualized images. 107 | If this is set, the script will automatically export 108 | the visualized image to the export directory. 109 | --auto_export When using this flag, the visualizer will automatically 110 | export the visualized frame to an image file in the 111 | `export_dir` directory. 112 | --ignore_fixed_transform 113 | When using this flag, the visualizer will not use the 114 | fixed transform matrix for the 3d model. 115 | ``` 116 | _NOTE: The `nvdu_viz` script can work from any directory_ 117 | 118 | ## Examples 119 | ### Visualize a dataset generated by NDDS: 120 | 1. Visualize the current directory: 121 | ``` 122 | nvdu_viz 123 | ``` 124 | 2. Visualize a relative path: 125 | ``` 126 | nvdu_viz ../a_dataset 127 | ``` 128 | 3. Visualize an absolute path: 129 | ``` 130 | nvdu_viz ~/data/dataset 131 | ``` 132 | 4. Visualize different aspect of a frame using a filter: 133 | ``` 134 | nvdu_viz dataset_path --name_filters *.left.png *.right.png 135 | ``` 136 | 137 | ### Visualize a set of images using different annotation data: 138 | 1. The camera and object settings files are in the image directory: 139 | ``` 140 | nvdu_viz image_directory_here -a annotation_directory_here 141 | ``` 142 | 2. The camera and object settings files are NOT in the image directory: 143 | ``` 144 | nvdu_viz image_directory_here -a annotation_directory_here -c camera_setting_path_here -o object_setting_path_here 145 | ``` 146 | 147 | ## Controls 148 | ### Visualization options: 149 | ``` 150 | F3 - Toggle the 2d cuboid 151 | F4 - Toggle the 3d cuboid 152 | F5 - Toggle the 3d models 153 | F6 - Toggle the axes 154 | F7 - Toggle the overlay frame name 155 | 1 - Render the 3d models normally 156 | 2 - Render the 3d models using only the edge lines 157 | 3 - Render the 3d models as point clouds 158 | ``` 159 | 160 | ### Other: 161 | ``` 162 | ESC - Quit the visualizer 163 | Right - Go to the next frame 164 | Left - Go to the previous frame 165 | Up - Go to the next 100 frame 166 | Down - Go to the previous 100 frame 167 | Space - Toggle frame auto-changing 168 | F12 - Toggle exporting the visualized frame to file 169 | ``` 170 | -------------------------------------------------------------------------------- /setup.cfg: -------------------------------------------------------------------------------- 1 | [bdist_wheel] 2 | universal=1 -------------------------------------------------------------------------------- /setup.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) 2018 NVIDIA Corporation. All rights reserved. 2 | # This work is licensed under a Creative Commons Attribution-NonCommercial-ShareAlike 4.0 International 3 | # License. (https://creativecommons.org/licenses/by-nc-sa/4.0/legalcode 4 | 5 | from setuptools import Command, find_packages, setup 6 | import os 7 | from os import path 8 | import glob 9 | 10 | __version_info__ = (1, 0, 0, 1) 11 | _ROOT = os.path.abspath(os.path.dirname(__file__)) 12 | 13 | # Utility function to read the README file. 14 | # Used for the long_description. It's nice, because now 1) we have a top level 15 | # README file and 2) it's easier to type in the README file than to put a raw 16 | # string in below ... 17 | def read(fname): 18 | return open(os.path.join(os.path.dirname(__file__), fname)).read() 19 | 20 | def get_all_files(find_dir): 21 | all_files = [] 22 | for check_path in os.listdir(find_dir): 23 | full_path = path.join(find_dir, check_path) 24 | if (path.isdir(full_path)): 25 | all_files.extend(get_all_files(full_path)) 26 | else: 27 | relative_path = path.relpath(full_path, _ROOT) 28 | all_files.append(relative_path) 29 | return all_files 30 | 31 | all_config_files = get_all_files(path.join(_ROOT, path.join('nvdu', 'config'))) 32 | print("all_config_files: {}".format(all_config_files)) 33 | 34 | __version__ = '.'.join(map(str, __version_info__)) 35 | 36 | setup( 37 | name = "nvdu", 38 | version = __version__, 39 | description = "Nvidia Dataset Utilities", 40 | long_description = read('readme.md'), 41 | long_description_content_type = 'text/markdown', 42 | url = "https://github.com/NVIDIA/Dataset_Utilities", 43 | author = "NVIDIA Corporation", 44 | author_email = "info@nvidia.com", 45 | maintainer = "Thang To", 46 | maintainer_email = "thangt@nvidia.com", 47 | license = "Creative Commons Attribution-NonCommercial-ShareAlike 4.0. https://creativecommons.org/licenses/by-nc-sa/4.0/", 48 | # See https://pypi.python.org/pypi?%3Aaction=list_classifiers 49 | classifiers = [ 50 | "Development Status :: 5 - Production/Stable", 51 | "Intended Audience :: Developers", 52 | "Natural Language :: English", 53 | "Operating System :: OS Independent", 54 | "Programming Language :: Python :: 3.6", 55 | "Topic :: Utilities", 56 | "Topic :: Software Development :: Libraries :: Python Modules", 57 | ], 58 | keywords = "nvdu, nvidia", 59 | packages=find_packages(), 60 | package_data={'': all_config_files}, 61 | include_package_data=True, 62 | install_requires = [ 63 | "numpy", 64 | "opencv-python", 65 | "pyrr", 66 | "PyWavefront==0.2.0", 67 | "pyglet", 68 | "fuzzyfinder" 69 | ], 70 | extras_require = { 71 | # "test": [ "pytest" ] 72 | }, 73 | entry_points = { 74 | "console_scripts": [ 75 | "nvdu_viz=nvdu.tools.test_nvdu_visualizer:main", 76 | "nvdu_ycb=nvdu.tools.nvdu_ycb:main", 77 | ] 78 | }, 79 | scripts=[], 80 | ) --------------------------------------------------------------------------------