├── CMakeLists.txt ├── README.md ├── config └── config_pose.yaml ├── example ├── 1 │ ├── FastNet-V1.png │ ├── FastNet-V2.png │ ├── VGG-16.png │ └── VGG-19.png ├── 2 │ ├── FastNet-V1.png │ ├── FastNet-V2.png │ ├── VGG-16.png │ └── VGG-19.png ├── 3 │ ├── FastNet-V1.png │ ├── FastNet-V2.png │ ├── VGG-16.png │ └── VGG-19.png ├── 4 │ ├── FastNet-V1.png │ ├── FastNet-V2.png │ ├── VGG-16.png │ └── VGG-19.png ├── 5 │ ├── FastNet-V1.png │ ├── FastNet-V2.png │ ├── Vgg-16.png │ └── Vgg-19.png ├── 6 │ ├── FastNet-V1.png │ ├── FastNet-V2.png │ ├── VGG-16.png │ └── VGG-19.png ├── camera_2020128113847.mp4 └── main.mp4 ├── license.md ├── src ├── dope.py └── inference │ ├── cuboid.py │ ├── cuboid_pnp_solver.py │ ├── detector_FastNet-V1.py │ └── detector_FastNet-V2.py └── weight └── Readme.txt /CMakeLists.txt: -------------------------------------------------------------------------------- 1 | cmake_minimum_required(VERSION 2.8.3) 2 | project(dope) 3 | 4 | ## Compile as C++11, supported in ROS Kinetic and newer 5 | # add_compile_options(-std=c++11) 6 | 7 | ## Find catkin macros and libraries 8 | ## if COMPONENTS list like find_package(catkin REQUIRED COMPONENTS xyz) 9 | ## is used, also find other catkin packages 10 | find_package(catkin REQUIRED COMPONENTS 11 | rospy 12 | std_msgs 13 | tf2 14 | ) 15 | 16 | ## System dependencies are found with CMake's conventions 17 | # find_package(Boost REQUIRED COMPONENTS system) 18 | 19 | 20 | ## Uncomment this if the package has a setup.py. This macro ensures 21 | ## modules and global scripts declared therein get installed 22 | ## See http://ros.org/doc/api/catkin/html/user_guide/setup_dot_py.html 23 | # catkin_python_setup() 24 | 25 | ################################################ 26 | ## Declare ROS messages, services and actions ## 27 | ################################################ 28 | 29 | ## To declare and build messages, services or actions from within this 30 | ## package, follow these steps: 31 | ## * Let MSG_DEP_SET be the set of packages whose message types you use in 32 | ## your messages/services/actions (e.g. std_msgs, actionlib_msgs, ...). 33 | ## * In the file package.xml: 34 | ## * add a build_depend tag for "message_generation" 35 | ## * add a build_depend and a exec_depend tag for each package in MSG_DEP_SET 36 | ## * If MSG_DEP_SET isn't empty the following dependency has been pulled in 37 | ## but can be declared for certainty nonetheless: 38 | ## * add a exec_depend tag for "message_runtime" 39 | ## * In this file (CMakeLists.txt): 40 | ## * add "message_generation" and every package in MSG_DEP_SET to 41 | ## find_package(catkin REQUIRED COMPONENTS ...) 42 | ## * add "message_runtime" and every package in MSG_DEP_SET to 43 | ## catkin_package(CATKIN_DEPENDS ...) 44 | ## * uncomment the add_*_files sections below as needed 45 | ## and list every .msg/.srv/.action file to be processed 46 | ## * uncomment the generate_messages entry below 47 | ## * add every package in MSG_DEP_SET to generate_messages(DEPENDENCIES ...) 48 | 49 | ## Generate messages in the 'msg' folder 50 | # add_message_files( 51 | # FILES 52 | # Message1.msg 53 | # Message2.msg 54 | # ) 55 | 56 | ## Generate services in the 'srv' folder 57 | # add_service_files( 58 | # FILES 59 | # Service1.srv 60 | # Service2.srv 61 | # ) 62 | 63 | ## Generate actions in the 'action' folder 64 | # add_action_files( 65 | # FILES 66 | # Action1.action 67 | # Action2.action 68 | # ) 69 | 70 | ## Generate added messages and services with any dependencies listed here 71 | # generate_messages( 72 | # DEPENDENCIES 73 | # std_msgs 74 | # ) 75 | 76 | ################################################ 77 | ## Declare ROS dynamic reconfigure parameters ## 78 | ################################################ 79 | 80 | ## To declare and build dynamic reconfigure parameters within this 81 | ## package, follow these steps: 82 | ## * In the file package.xml: 83 | ## * add a build_depend and a exec_depend tag for "dynamic_reconfigure" 84 | ## * In this file (CMakeLists.txt): 85 | ## * add "dynamic_reconfigure" to 86 | ## find_package(catkin REQUIRED COMPONENTS ...) 87 | ## * uncomment the "generate_dynamic_reconfigure_options" section below 88 | ## and list every .cfg file to be processed 89 | 90 | ## Generate dynamic reconfigure parameters in the 'cfg' folder 91 | # generate_dynamic_reconfigure_options( 92 | # cfg/DynReconf1.cfg 93 | # cfg/DynReconf2.cfg 94 | # ) 95 | 96 | ################################### 97 | ## catkin specific configuration ## 98 | ################################### 99 | ## The catkin_package macro generates cmake config files for your package 100 | ## Declare things to be passed to dependent projects 101 | ## INCLUDE_DIRS: uncomment this if your package contains header files 102 | ## LIBRARIES: libraries you create in this project that dependent projects also need 103 | ## CATKIN_DEPENDS: catkin_packages dependent projects also need 104 | ## DEPENDS: system dependencies of this project that dependent projects also need 105 | catkin_package( 106 | # INCLUDE_DIRS include 107 | # LIBRARIES dope 108 | # CATKIN_DEPENDS rospy std_msgs tf2 109 | # DEPENDS system_lib 110 | ) 111 | 112 | ########### 113 | ## Build ## 114 | ########### 115 | 116 | ## Specify additional locations of header files 117 | ## Your package locations should be listed before other locations 118 | include_directories( 119 | # include 120 | ${catkin_INCLUDE_DIRS} 121 | ) 122 | 123 | ## Declare a C++ library 124 | # add_library(${PROJECT_NAME} 125 | # src/${PROJECT_NAME}/dope/dope.cpp 126 | # ) 127 | 128 | ## Add cmake target dependencies of the library 129 | ## as an example, code may need to be generated before libraries 130 | ## either from message generation or dynamic reconfigure 131 | # add_dependencies(${PROJECT_NAME} ${${PROJECT_NAME}_EXPORTED_TARGETS} ${catkin_EXPORTED_TARGETS}) 132 | 133 | ## Declare a C++ executable 134 | ## With catkin_make all packages are built within a single CMake context 135 | ## The recommended prefix ensures that target names across packages don't collide 136 | # add_executable(${PROJECT_NAME}_node src/dope_vis_node.cpp) 137 | 138 | ## Rename C++ executable without prefix 139 | ## The above recommended prefix causes long target names, the following renames the 140 | ## target back to the shorter version for ease of user use 141 | ## e.g. "rosrun someones_pkg node" instead of "rosrun someones_pkg someones_pkg_node" 142 | # set_target_properties(${PROJECT_NAME}_node PROPERTIES OUTPUT_NAME node PREFIX "") 143 | 144 | ## Add cmake target dependencies of the executable 145 | ## same as for the library above 146 | # add_dependencies(${PROJECT_NAME}_node ${${PROJECT_NAME}_EXPORTED_TARGETS} ${catkin_EXPORTED_TARGETS}) 147 | 148 | ## Specify libraries to link a library or executable target against 149 | # target_link_libraries(${PROJECT_NAME}_node 150 | # ${catkin_LIBRARIES} 151 | # ) 152 | 153 | ############# 154 | ## Install ## 155 | ############# 156 | 157 | # all install targets should use catkin DESTINATION variables 158 | # See http://ros.org/doc/api/catkin/html/adv_user_guide/variables.html 159 | 160 | ## Mark executable scripts (Python etc.) for installation 161 | ## in contrast to setup.py, you can choose the destination 162 | # install(PROGRAMS 163 | # scripts/my_python_script 164 | # DESTINATION ${CATKIN_PACKAGE_BIN_DESTINATION} 165 | # ) 166 | 167 | ## Mark executables and/or libraries for installation 168 | # install(TARGETS ${PROJECT_NAME} ${PROJECT_NAME}_node 169 | # ARCHIVE DESTINATION ${CATKIN_PACKAGE_LIB_DESTINATION} 170 | # LIBRARY DESTINATION ${CATKIN_PACKAGE_LIB_DESTINATION} 171 | # RUNTIME DESTINATION ${CATKIN_PACKAGE_BIN_DESTINATION} 172 | # ) 173 | 174 | ## Mark cpp header files for installation 175 | # install(DIRECTORY include/${PROJECT_NAME}/ 176 | # DESTINATION ${CATKIN_PACKAGE_INCLUDE_DESTINATION} 177 | # FILES_MATCHING PATTERN "*.h" 178 | # PATTERN ".svn" EXCLUDE 179 | # ) 180 | 181 | ## Mark other files for installation (e.g. launch and bag files, etc.) 182 | # install(FILES 183 | # # myfile1 184 | # # myfile2 185 | # DESTINATION ${CATKIN_PACKAGE_SHARE_DESTINATION} 186 | # ) 187 | 188 | ############# 189 | ## Testing ## 190 | ############# 191 | 192 | ## Add gtest based cpp test target and link libraries 193 | # catkin_add_gtest(${PROJECT_NAME}-test test/test_dope_vis.cpp) 194 | # if(TARGET ${PROJECT_NAME}-test) 195 | # target_link_libraries(${PROJECT_NAME}-test ${PROJECT_NAME}) 196 | # endif() 197 | 198 | ## Add folders to be run by python nosetests 199 | # catkin_add_nosetests(test) 200 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # Practical_Robotic_Grasping_Method 2 | 3 | ## Prepare 4 | ### ROS packages 5 | >> 1. realsense2_camera 6 | >> 2. ROOT: git clone: https://github.com/zhanghui-hunan/Practical_Robotic_Grasping_Method.git 7 | 8 | ### Environments 9 | >> 1. Python2.7 10 | >> 2. torch==1.4.0 11 | 12 | ### Download weight 13 | >>
Download from Google Driver: https://drive.google.com/drive/folders/1jVd4fTcAuZC0nf5Vy5niOmiEvnarSzh8 14 | 15 | ## Run 16 |
> $ roslaunch realsense2_camera rs_camera.launch 17 |
> $ rosrun Practical_Robotic_Grasping_Method dope.py 18 |
> $ rqt_image_view 19 |
> PS: rostopic: rgb_dope_points 20 | # Practical_Robotic_Grasping_Method 21 | -------------------------------------------------------------------------------- /config/config_pose.yaml: -------------------------------------------------------------------------------- 1 | topic_camera: "/camera/color/image_raw" 2 | topic_publishing: "dope" 3 | frame_id: "/dope" 4 | 5 | # Comment any of these lines to prevent detection / pose estimation of that object 6 | weights: { 7 | #"soap":"VGG-19.pth" 8 | #"soap": "VGG-16.pth" 9 | #"soap":"FastNet-V2.pth" 10 | "soap":"FastNet-V1.pth" 11 | } 12 | 13 | # Cuboid dimension in cm x,y,z 14 | dimensions: { 15 | "soap": [9.3999996185302734, 5.8000001907348633, 3.2000000476837158] 16 | } 17 | 18 | draw_colors: { 19 | "soap": [13, 255, 128] # green 20 | } 21 | 22 | # Camera intrinsics (Logitech C920) 23 | camera_settings: { 24 | "name": "realsense-d435", 25 | "width": 640, 26 | "height": 480, 27 | "fx": 611.028564453125, 28 | "fy": 611.33837890625, 29 | "cx": 321.5130615234375, 30 | "cy": 240.90025329589844 31 | } 32 | 33 | # Config params for DOPE 34 | thresh_angle: 0.5 35 | thresh_map: 0.01 36 | sigma: 3 37 | thresh_points: 0.1 38 | -------------------------------------------------------------------------------- /example/1/FastNet-V1.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/zhanghui-HNU/Practical_Robotic_Grasping_Method/e33a4abe68627b8209c4d71246bc192eea075aa2/example/1/FastNet-V1.png -------------------------------------------------------------------------------- /example/1/FastNet-V2.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/zhanghui-HNU/Practical_Robotic_Grasping_Method/e33a4abe68627b8209c4d71246bc192eea075aa2/example/1/FastNet-V2.png -------------------------------------------------------------------------------- /example/1/VGG-16.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/zhanghui-HNU/Practical_Robotic_Grasping_Method/e33a4abe68627b8209c4d71246bc192eea075aa2/example/1/VGG-16.png -------------------------------------------------------------------------------- /example/1/VGG-19.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/zhanghui-HNU/Practical_Robotic_Grasping_Method/e33a4abe68627b8209c4d71246bc192eea075aa2/example/1/VGG-19.png -------------------------------------------------------------------------------- /example/2/FastNet-V1.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/zhanghui-HNU/Practical_Robotic_Grasping_Method/e33a4abe68627b8209c4d71246bc192eea075aa2/example/2/FastNet-V1.png -------------------------------------------------------------------------------- /example/2/FastNet-V2.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/zhanghui-HNU/Practical_Robotic_Grasping_Method/e33a4abe68627b8209c4d71246bc192eea075aa2/example/2/FastNet-V2.png -------------------------------------------------------------------------------- /example/2/VGG-16.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/zhanghui-HNU/Practical_Robotic_Grasping_Method/e33a4abe68627b8209c4d71246bc192eea075aa2/example/2/VGG-16.png -------------------------------------------------------------------------------- /example/2/VGG-19.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/zhanghui-HNU/Practical_Robotic_Grasping_Method/e33a4abe68627b8209c4d71246bc192eea075aa2/example/2/VGG-19.png -------------------------------------------------------------------------------- /example/3/FastNet-V1.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/zhanghui-HNU/Practical_Robotic_Grasping_Method/e33a4abe68627b8209c4d71246bc192eea075aa2/example/3/FastNet-V1.png -------------------------------------------------------------------------------- /example/3/FastNet-V2.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/zhanghui-HNU/Practical_Robotic_Grasping_Method/e33a4abe68627b8209c4d71246bc192eea075aa2/example/3/FastNet-V2.png -------------------------------------------------------------------------------- /example/3/VGG-16.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/zhanghui-HNU/Practical_Robotic_Grasping_Method/e33a4abe68627b8209c4d71246bc192eea075aa2/example/3/VGG-16.png -------------------------------------------------------------------------------- /example/3/VGG-19.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/zhanghui-HNU/Practical_Robotic_Grasping_Method/e33a4abe68627b8209c4d71246bc192eea075aa2/example/3/VGG-19.png -------------------------------------------------------------------------------- /example/4/FastNet-V1.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/zhanghui-HNU/Practical_Robotic_Grasping_Method/e33a4abe68627b8209c4d71246bc192eea075aa2/example/4/FastNet-V1.png -------------------------------------------------------------------------------- /example/4/FastNet-V2.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/zhanghui-HNU/Practical_Robotic_Grasping_Method/e33a4abe68627b8209c4d71246bc192eea075aa2/example/4/FastNet-V2.png -------------------------------------------------------------------------------- /example/4/VGG-16.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/zhanghui-HNU/Practical_Robotic_Grasping_Method/e33a4abe68627b8209c4d71246bc192eea075aa2/example/4/VGG-16.png -------------------------------------------------------------------------------- /example/4/VGG-19.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/zhanghui-HNU/Practical_Robotic_Grasping_Method/e33a4abe68627b8209c4d71246bc192eea075aa2/example/4/VGG-19.png -------------------------------------------------------------------------------- /example/5/FastNet-V1.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/zhanghui-HNU/Practical_Robotic_Grasping_Method/e33a4abe68627b8209c4d71246bc192eea075aa2/example/5/FastNet-V1.png -------------------------------------------------------------------------------- /example/5/FastNet-V2.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/zhanghui-HNU/Practical_Robotic_Grasping_Method/e33a4abe68627b8209c4d71246bc192eea075aa2/example/5/FastNet-V2.png -------------------------------------------------------------------------------- /example/5/Vgg-16.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/zhanghui-HNU/Practical_Robotic_Grasping_Method/e33a4abe68627b8209c4d71246bc192eea075aa2/example/5/Vgg-16.png -------------------------------------------------------------------------------- /example/5/Vgg-19.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/zhanghui-HNU/Practical_Robotic_Grasping_Method/e33a4abe68627b8209c4d71246bc192eea075aa2/example/5/Vgg-19.png -------------------------------------------------------------------------------- /example/6/FastNet-V1.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/zhanghui-HNU/Practical_Robotic_Grasping_Method/e33a4abe68627b8209c4d71246bc192eea075aa2/example/6/FastNet-V1.png -------------------------------------------------------------------------------- /example/6/FastNet-V2.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/zhanghui-HNU/Practical_Robotic_Grasping_Method/e33a4abe68627b8209c4d71246bc192eea075aa2/example/6/FastNet-V2.png -------------------------------------------------------------------------------- /example/6/VGG-16.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/zhanghui-HNU/Practical_Robotic_Grasping_Method/e33a4abe68627b8209c4d71246bc192eea075aa2/example/6/VGG-16.png -------------------------------------------------------------------------------- /example/6/VGG-19.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/zhanghui-HNU/Practical_Robotic_Grasping_Method/e33a4abe68627b8209c4d71246bc192eea075aa2/example/6/VGG-19.png -------------------------------------------------------------------------------- /example/camera_2020128113847.mp4: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/zhanghui-HNU/Practical_Robotic_Grasping_Method/e33a4abe68627b8209c4d71246bc192eea075aa2/example/camera_2020128113847.mp4 -------------------------------------------------------------------------------- /example/main.mp4: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/zhanghui-HNU/Practical_Robotic_Grasping_Method/e33a4abe68627b8209c4d71246bc192eea075aa2/example/main.mp4 -------------------------------------------------------------------------------- /license.md: -------------------------------------------------------------------------------- 1 | # Creative Commons 2 | 3 | ## [Attribution-NonCommercial-ShareAlike 4.0 International](https://creativecommons.org/licenses/by-nc-sa/4.0/legalcode) 4 | 5 | Creative Commons Corporation (“Creative Commons”) is not a law firm and does not provide legal services or legal advice. Distribution of Creative Commons public licenses does not create a lawyer-client or other relationship. Creative Commons makes its licenses and related information available on an “as-is” basis. Creative Commons gives no warranties regarding its licenses, any material licensed under their terms and conditions, or any related information. Creative Commons disclaims all liability for damages resulting from their use to the fullest extent possible. 6 | 7 | ### Using Creative Commons Public Licenses 8 | 9 | Creative Commons public licenses provide a standard set of terms and conditions that creators and other rights holders may use to share original works of authorship and other material subject to copyright and certain other rights specified in the public license below. The following considerations are for informational purposes only, are not exhaustive, and do not form part of our licenses. 10 | 11 | * __Considerations for licensors:__ Our public licenses are intended for use by those authorized to give the public permission to use material in ways otherwise restricted by copyright and certain other rights. Our licenses are irrevocable. Licensors should read and understand the terms and conditions of the license they choose before applying it. Licensors should also secure all rights necessary before applying our licenses so that the public can reuse the material as expected. Licensors should clearly mark any material not subject to the license. This includes other CC-licensed material, or material used under an exception or limitation to copyright. [More considerations for licensors](http://wiki.creativecommons.org/Considerations_for_licensors_and_licensees#Considerations_for_licensors). 12 | 13 | * __Considerations for the public:__ By using one of our public licenses, a licensor grants the public permission to use the licensed material under specified terms and conditions. If the licensor’s permission is not necessary for any reason–for example, because of any applicable exception or limitation to copyright–then that use is not regulated by the license. Our licenses grant only permissions under copyright and certain other rights that a licensor has authority to grant. Use of the licensed material may still be restricted for other reasons, including because others have copyright or other rights in the material. A licensor may make special requests, such as asking that all changes be marked or described. Although not required by our licenses, you are encouraged to respect those requests where reasonable. [More considerations for the public](http://wiki.creativecommons.org/Considerations_for_licensors_and_licensees#Considerations_for_licensees). 14 | 15 | ## Creative Commons Attribution-NonCommercial-ShareAlike 4.0 International Public License 16 | 17 | By exercising the Licensed Rights (defined below), You accept and agree to be bound by the terms and conditions of this Creative Commons Attribution-NonCommercial-ShareAlike 4.0 International Public License ("Public License"). To the extent this Public License may be interpreted as a contract, You are granted the Licensed Rights in consideration of Your acceptance of these terms and conditions, and the Licensor grants You such rights in consideration of benefits the Licensor receives from making the Licensed Material available under these terms and conditions. 18 | 19 | ### Section 1 – Definitions. 20 | 21 | a. __Adapted Material__ means material subject to Copyright and Similar Rights that is derived from or based upon the Licensed Material and in which the Licensed Material is translated, altered, arranged, transformed, or otherwise modified in a manner requiring permission under the Copyright and Similar Rights held by the Licensor. For purposes of this Public License, where the Licensed Material is a musical work, performance, or sound recording, Adapted Material is always produced where the Licensed Material is synched in timed relation with a moving image. 22 | 23 | b. __Adapter's License__ means the license You apply to Your Copyright and Similar Rights in Your contributions to Adapted Material in accordance with the terms and conditions of this Public License. 24 | 25 | c. __BY-NC-SA Compatible License__ means a license listed at [creativecommons.org/compatiblelicenses](http://creativecommons.org/compatiblelicenses), approved by Creative Commons as essentially the equivalent of this Public License. 26 | 27 | d. __Copyright and Similar Rights__ means copyright and/or similar rights closely related to copyright including, without limitation, performance, broadcast, sound recording, and Sui Generis Database Rights, without regard to how the rights are labeled or categorized. For purposes of this Public License, the rights specified in Section 2(b)(1)-(2) are not Copyright and Similar Rights. 28 | 29 | e. __Effective Technological Measures__ means those measures that, in the absence of proper authority, may not be circumvented under laws fulfilling obligations under Article 11 of the WIPO Copyright Treaty adopted on December 20, 1996, and/or similar international agreements. 30 | 31 | f. __Exceptions and Limitations__ means fair use, fair dealing, and/or any other exception or limitation to Copyright and Similar Rights that applies to Your use of the Licensed Material. 32 | 33 | g. __License Elements__ means the license attributes listed in the name of a Creative Commons Public License. The License Elements of this Public License are Attribution, NonCommercial, and ShareAlike. 34 | 35 | h. __Licensed Material__ means the artistic or literary work, database, or other material to which the Licensor applied this Public License. 36 | 37 | i. __Licensed Rights__ means the rights granted to You subject to the terms and conditions of this Public License, which are limited to all Copyright and Similar Rights that apply to Your use of the Licensed Material and that the Licensor has authority to license. 38 | 39 | h. __Licensor__ means the individual(s) or entity(ies) granting rights under this Public License. 40 | 41 | i. __NonCommercial__ means not primarily intended for or directed towards commercial advantage or monetary compensation. For purposes of this Public License, the exchange of the Licensed Material for other material subject to Copyright and Similar Rights by digital file-sharing or similar means is NonCommercial provided there is no payment of monetary compensation in connection with the exchange. 42 | 43 | j. __Share__ means to provide material to the public by any means or process that requires permission under the Licensed Rights, such as reproduction, public display, public performance, distribution, dissemination, communication, or importation, and to make material available to the public including in ways that members of the public may access the material from a place and at a time individually chosen by them. 44 | 45 | k. __Sui Generis Database Rights__ means rights other than copyright resulting from Directive 96/9/EC of the European Parliament and of the Council of 11 March 1996 on the legal protection of databases, as amended and/or succeeded, as well as other essentially equivalent rights anywhere in the world. 46 | 47 | l. __You__ means the individual or entity exercising the Licensed Rights under this Public License. Your has a corresponding meaning. 48 | 49 | ### Section 2 – Scope. 50 | 51 | a. ___License grant.___ 52 | 53 | 1. Subject to the terms and conditions of this Public License, the Licensor hereby grants You a worldwide, royalty-free, non-sublicensable, non-exclusive, irrevocable license to exercise the Licensed Rights in the Licensed Material to: 54 | 55 | A. reproduce and Share the Licensed Material, in whole or in part, for NonCommercial purposes only; and 56 | 57 | B. produce, reproduce, and Share Adapted Material for NonCommercial purposes only. 58 | 59 | 2. __Exceptions and Limitations.__ For the avoidance of doubt, where Exceptions and Limitations apply to Your use, this Public License does not apply, and You do not need to comply with its terms and conditions. 60 | 61 | 3. __Term.__ The term of this Public License is specified in Section 6(a). 62 | 63 | 4. __Media and formats; technical modifications allowed.__ The Licensor authorizes You to exercise the Licensed Rights in all media and formats whether now known or hereafter created, and to make technical modifications necessary to do so. The Licensor waives and/or agrees not to assert any right or authority to forbid You from making technical modifications necessary to exercise the Licensed Rights, including technical modifications necessary to circumvent Effective Technological Measures. For purposes of this Public License, simply making modifications authorized by this Section 2(a)(4) never produces Adapted Material. 64 | 65 | 5. __Downstream recipients.__ 66 | 67 | A. __Offer from the Licensor – Licensed Material.__ Every recipient of the Licensed Material automatically receives an offer from the Licensor to exercise the Licensed Rights under the terms and conditions of this Public License. 68 | 69 | B. __Additional offer from the Licensor – Adapted Material.__ Every recipient of Adapted Material from You automatically receives an offer from the Licensor to exercise the Licensed Rights in the Adapted Material under the conditions of the Adapter’s License You apply. 70 | 71 | C. __No downstream restrictions.__ You may not offer or impose any additional or different terms or conditions on, or apply any Effective Technological Measures to, the Licensed Material if doing so restricts exercise of the Licensed Rights by any recipient of the Licensed Material. 72 | 73 | 6. __No endorsement.__ Nothing in this Public License constitutes or may be construed as permission to assert or imply that You are, or that Your use of the Licensed Material is, connected with, or sponsored, endorsed, or granted official status by, the Licensor or others designated to receive attribution as provided in Section 3(a)(1)(A)(i). 74 | 75 | b. ___Other rights.___ 76 | 77 | 1. Moral rights, such as the right of integrity, are not licensed under this Public License, nor are publicity, privacy, and/or other similar personality rights; however, to the extent possible, the Licensor waives and/or agrees not to assert any such rights held by the Licensor to the limited extent necessary to allow You to exercise the Licensed Rights, but not otherwise. 78 | 79 | 2. Patent and trademark rights are not licensed under this Public License. 80 | 81 | 3. To the extent possible, the Licensor waives any right to collect royalties from You for the exercise of the Licensed Rights, whether directly or through a collecting society under any voluntary or waivable statutory or compulsory licensing scheme. In all other cases the Licensor expressly reserves any right to collect such royalties, including when the Licensed Material is used other than for NonCommercial purposes. 82 | 83 | ### Section 3 – License Conditions. 84 | 85 | Your exercise of the Licensed Rights is expressly made subject to the following conditions. 86 | 87 | a. ___Attribution.___ 88 | 89 | 1. If You Share the Licensed Material (including in modified form), You must: 90 | 91 | A. retain the following if it is supplied by the Licensor with the Licensed Material: 92 | 93 | i. identification of the creator(s) of the Licensed Material and any others designated to receive attribution, in any reasonable manner requested by the Licensor (including by pseudonym if designated); 94 | 95 | ii. a copyright notice; 96 | 97 | iii. a notice that refers to this Public License; 98 | 99 | iv. a notice that refers to the disclaimer of warranties; 100 | 101 | v. a URI or hyperlink to the Licensed Material to the extent reasonably practicable; 102 | 103 | B. indicate if You modified the Licensed Material and retain an indication of any previous modifications; and 104 | 105 | C. indicate the Licensed Material is licensed under this Public License, and include the text of, or the URI or hyperlink to, this Public License. 106 | 107 | 2. You may satisfy the conditions in Section 3(a)(1) in any reasonable manner based on the medium, means, and context in which You Share the Licensed Material. For example, it may be reasonable to satisfy the conditions by providing a URI or hyperlink to a resource that includes the required information. 108 | 109 | 3. If requested by the Licensor, You must remove any of the information required by Section 3(a)(1)(A) to the extent reasonably practicable. 110 | 111 | b. ___ShareAlike.___ 112 | 113 | In addition to the conditions in Section 3(a), if You Share Adapted Material You produce, the following conditions also apply. 114 | 115 | 1. The Adapter’s License You apply must be a Creative Commons license with the same License Elements, this version or later, or a BY-NC-SA Compatible License. 116 | 117 | 2. You must include the text of, or the URI or hyperlink to, the Adapter's License You apply. You may satisfy this condition in any reasonable manner based on the medium, means, and context in which You Share Adapted Material. 118 | 119 | 3. You may not offer or impose any additional or different terms or conditions on, or apply any Effective Technological Measures to, Adapted Material that restrict exercise of the rights granted under the Adapter's License You apply. 120 | 121 | ### Section 4 – Sui Generis Database Rights. 122 | 123 | Where the Licensed Rights include Sui Generis Database Rights that apply to Your use of the Licensed Material: 124 | 125 | a. for the avoidance of doubt, Section 2(a)(1) grants You the right to extract, reuse, reproduce, and Share all or a substantial portion of the contents of the database for NonCommercial purposes only; 126 | 127 | b. if You include all or a substantial portion of the database contents in a database in which You have Sui Generis Database Rights, then the database in which You have Sui Generis Database Rights (but not its individual contents) is Adapted Material, including for purposes of Section 3(b); and 128 | 129 | c. You must comply with the conditions in Section 3(a) if You Share all or a substantial portion of the contents of the database. 130 | 131 | For the avoidance of doubt, this Section 4 supplements and does not replace Your obligations under this Public License where the Licensed Rights include other Copyright and Similar Rights. 132 | 133 | ### Section 5 – Disclaimer of Warranties and Limitation of Liability. 134 | 135 | a. __Unless otherwise separately undertaken by the Licensor, to the extent possible, the Licensor offers the Licensed Material as-is and as-available, and makes no representations or warranties of any kind concerning the Licensed Material, whether express, implied, statutory, or other. This includes, without limitation, warranties of title, merchantability, fitness for a particular purpose, non-infringement, absence of latent or other defects, accuracy, or the presence or absence of errors, whether or not known or discoverable. Where disclaimers of warranties are not allowed in full or in part, this disclaimer may not apply to You.__ 136 | 137 | b. __To the extent possible, in no event will the Licensor be liable to You on any legal theory (including, without limitation, negligence) or otherwise for any direct, special, indirect, incidental, consequential, punitive, exemplary, or other losses, costs, expenses, or damages arising out of this Public License or use of the Licensed Material, even if the Licensor has been advised of the possibility of such losses, costs, expenses, or damages. Where a limitation of liability is not allowed in full or in part, this limitation may not apply to You.__ 138 | 139 | c. The disclaimer of warranties and limitation of liability provided above shall be interpreted in a manner that, to the extent possible, most closely approximates an absolute disclaimer and waiver of all liability. 140 | 141 | ### Section 6 – Term and Termination. 142 | 143 | a. This Public License applies for the term of the Copyright and Similar Rights licensed here. However, if You fail to comply with this Public License, then Your rights under this Public License terminate automatically. 144 | 145 | b. Where Your right to use the Licensed Material has terminated under Section 6(a), it reinstates: 146 | 147 | 1. automatically as of the date the violation is cured, provided it is cured within 30 days of Your discovery of the violation; or 148 | 149 | 2. automatically as of the date the violation is cured, provided it is cured within 30 days of Your discovery of the violation; or 150 | 151 | For the avoidance of doubt, this Section 6(b) does not affect any right the Licensor may have to seek remedies for Your violations of this Public License. 152 | 153 | c. For the avoidance of doubt, the Licensor may also offer the Licensed Material under separate terms or conditions or stop distributing the Licensed Material at any time; however, doing so will not terminate this Public License. 154 | 155 | d. Sections 1, 5, 6, 7, and 8 survive termination of this Public License. 156 | 157 | ### Section 7 – Other Terms and Conditions. 158 | 159 | a. The Licensor shall not be bound by any additional or different terms or conditions communicated by You unless expressly agreed. 160 | 161 | b. Any arrangements, understandings, or agreements regarding the Licensed Material not stated herein are separate from and independent of the terms and conditions of this Public License. 162 | 163 | ### Section 8 – Interpretation. 164 | 165 | a. For the avoidance of doubt, this Public License does not, and shall not be interpreted to, reduce, limit, restrict, or impose conditions on any use of the Licensed Material that could lawfully be made without permission under this Public License. 166 | 167 | b. To the extent possible, if any provision of this Public License is deemed unenforceable, it shall be automatically reformed to the minimum extent necessary to make it enforceable. If the provision cannot be reformed, it shall be severed from this Public License without affecting the enforceability of the remaining terms and conditions. 168 | 169 | c. No term or condition of this Public License will be waived and no failure to comply consented to unless expressly agreed to by the Licensor. 170 | 171 | d. Nothing in this Public License constitutes or may be interpreted as a limitation upon, or waiver of, any privileges and immunities that apply to the Licensor or You, including from the legal processes of any jurisdiction or authority. 172 | 173 | ``` 174 | Creative Commons is not a party to its public licenses. Notwithstanding, Creative Commons may elect to apply one of its public licenses to material it publishes and in those instances will be considered the “Licensor.” Except for the limited purpose of indicating that material is shared under a Creative Commons public license or as otherwise permitted by the Creative Commons policies published at [creativecommons.org/policies](http://creativecommons.org/policies), Creative Commons does not authorize the use of the trademark “Creative Commons” or any other trademark or logo of Creative Commons without its prior written consent including, without limitation, in connection with any unauthorized modifications to any of its public licenses or any other arrangements, understandings, or agreements concerning use of licensed material. For the avoidance of doubt, this paragraph does not form part of the public licenses. 175 | 176 | Creative Commons may be contacted at [creativecommons.org](http://creativecommons.org/). 177 | ``` 178 | -------------------------------------------------------------------------------- /src/dope.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python 2 | 3 | # Copyright (c) 2018 NVIDIA Corporation. All rights reserved. 4 | # This work is licensed under a Creative Commons Attribution-NonCommercial-ShareAlike 4.0 International License. 5 | # https://creativecommons.org/licenses/by-nc-sa/4.0/legalcode 6 | 7 | """ 8 | This file starts a ROS node to run DOPE, 9 | listening to an image topic and publishing poses. 10 | """ 11 | 12 | from __future__ import print_function 13 | import yaml 14 | import sys 15 | 16 | import numpy as np 17 | import cv2 18 | 19 | import rospy 20 | import rospkg 21 | import tf 22 | from std_msgs.msg import String, Empty 23 | from cv_bridge import CvBridge, CvBridgeError 24 | from sensor_msgs.msg import Image as ImageSensor_msg 25 | from geometry_msgs.msg import PoseStamped 26 | 27 | from PIL import Image 28 | from PIL import ImageDraw 29 | 30 | # Import DOPE code 31 | rospack = rospkg.RosPack() 32 | g_path2package = rospack.get_path('dope') 33 | sys.path.append("{}/src/inference".format(g_path2package)) 34 | from cuboid import * 35 | #from detector_FastNet-V2 import * 36 | from detector_FastNet-V1 import * 37 | 38 | ### Global Variables 39 | g_bridge = CvBridge() 40 | g_img = None 41 | g_draw = None 42 | 43 | 44 | ### Basic functions 45 | def __image_callback(msg): 46 | '''Image callback''' 47 | global g_img 48 | g_img = g_bridge.imgmsg_to_cv2(msg, "rgb8") 49 | # cv2.imwrite('img.png', cv2.cvtColor(g_img, cv2.COLOR_BGR2RGB)) # for debugging 50 | 51 | 52 | ### Code to visualize the neural network output 53 | 54 | def DrawLine(point1, point2, lineColor, lineWidth): 55 | '''Draws line on image''' 56 | global g_draw 57 | if not point1 is None and point2 is not None: 58 | g_draw.line([point1,point2], fill=lineColor, width=lineWidth) 59 | 60 | def DrawDot(point, pointColor, pointRadius): 61 | '''Draws dot (filled circle) on image''' 62 | global g_draw 63 | if point is not None: 64 | xy = [ 65 | point[0]-pointRadius, 66 | point[1]-pointRadius, 67 | point[0]+pointRadius, 68 | point[1]+pointRadius 69 | ] 70 | g_draw.ellipse(xy, 71 | fill=pointColor, 72 | outline=pointColor 73 | ) 74 | 75 | def DrawCube(points, color=(255, 0, 0)): 76 | ''' 77 | Draws cube with a thick solid line across 78 | the front top edge and an X on the top face. 79 | ''' 80 | 81 | lineWidthForDrawing = 2 82 | 83 | # draw front 84 | DrawLine(points[0], points[1], color, lineWidthForDrawing) 85 | DrawLine(points[1], points[2], color, lineWidthForDrawing) 86 | DrawLine(points[3], points[2], color, lineWidthForDrawing) 87 | DrawLine(points[3], points[0], color, lineWidthForDrawing) 88 | 89 | # draw back 90 | DrawLine(points[4], points[5], color, lineWidthForDrawing) 91 | DrawLine(points[6], points[5], color, lineWidthForDrawing) 92 | DrawLine(points[6], points[7], color, lineWidthForDrawing) 93 | DrawLine(points[4], points[7], color, lineWidthForDrawing) 94 | 95 | # draw sides 96 | DrawLine(points[0], points[4], color, lineWidthForDrawing) 97 | DrawLine(points[7], points[3], color, lineWidthForDrawing) 98 | DrawLine(points[5], points[1], color, lineWidthForDrawing) 99 | DrawLine(points[2], points[6], color, lineWidthForDrawing) 100 | 101 | # draw dots 102 | DrawDot(points[0], pointColor=color, pointRadius = 4) 103 | DrawDot(points[1], pointColor=color, pointRadius = 4) 104 | 105 | # Axis calculation 106 | x_left = (points[3][0] + points[4][0]) / 2 107 | y_left = (points[3][1] + points[4][1]) / 2 108 | 109 | x_right = (points[2][0] + points[5][0]) / 2 110 | y_right = (points[2][1] + points[5][1]) / 2 111 | 112 | x_top = (points[0][0] + points[5][0]) / 2 113 | y_top = (points[0][1] + points[5][1]) / 2 114 | 115 | x_down = (points[3][0] + points[6][0]) / 2 116 | y_down = (points[3][1] + points[6][1]) / 2 117 | 118 | x_center = (x_left + x_right) / 2 119 | y_center = (y_top + y_down) / 2 120 | 121 | print(",%f,%f" % (x_center, y_center)) 122 | 123 | x_front = (points[0][0] + points[2][0]) / 2 124 | y_front = (points[0][1] + points[2][1]) / 2 125 | 126 | # draw x on the top 127 | DrawLine(points[0], points[5], color, lineWidthForDrawing) 128 | DrawLine(points[1], points[4], color, lineWidthForDrawing) 129 | 130 | # x-red 131 | DrawLine((x_left, y_left), (x_center, y_center), (255, 0, 0), 5) 132 | 133 | # y-green 134 | DrawLine((x_down, y_down), (x_center, y_center), (0, 255, 0), 5) 135 | 136 | # z-blue 137 | DrawLine((x_front, y_front), (x_center, y_center), (0, 0, 255), 5) 138 | 139 | 140 | def run_dope_node(params, freq=5): 141 | '''Starts ROS node to listen to image topic, run DOPE, and publish DOPE results''' 142 | 143 | global g_img 144 | global g_draw 145 | 146 | pubs = {} 147 | models = {} 148 | pnp_solvers = {} 149 | pub_dimension = {} 150 | draw_colors = {} 151 | 152 | # Initialize parameters 153 | matrix_camera = np.zeros((3,3)) 154 | matrix_camera[0,0] = params["camera_settings"]['fx'] 155 | matrix_camera[1,1] = params["camera_settings"]['fy'] 156 | matrix_camera[0,2] = params["camera_settings"]['cx'] 157 | matrix_camera[1,2] = params["camera_settings"]['cy'] 158 | matrix_camera[2,2] = 1 159 | dist_coeffs = np.zeros((4,1)) 160 | 161 | if "dist_coeffs" in params["camera_settings"]: 162 | dist_coeffs = np.array(params["camera_settings"]['dist_coeffs']) 163 | config_detect = lambda: None 164 | config_detect.mask_edges = 1 165 | config_detect.mask_faces = 1 166 | config_detect.vertex = 1 167 | config_detect.threshold = 0.5 168 | config_detect.softmax = 1000 169 | config_detect.thresh_angle = params['thresh_angle'] 170 | config_detect.thresh_map = params['thresh_map'] 171 | config_detect.sigma = params['sigma'] 172 | config_detect.thresh_points = params["thresh_points"] 173 | 174 | # For each object to detect, load network model, create PNP solver, and start ROS publishers 175 | for model in params['weights']: 176 | models[model] =\ 177 | ModelData( 178 | model, 179 | g_path2package + "/weights/" + params['weights'][model] 180 | ) 181 | models[model].load_net_model() 182 | 183 | draw_colors[model] = \ 184 | tuple(params["draw_colors"][model]) 185 | pnp_solvers[model] = \ 186 | CuboidPNPSolver( 187 | model, 188 | matrix_camera, 189 | Cuboid3d(params['dimensions'][model]), 190 | dist_coeffs=dist_coeffs 191 | ) 192 | pubs[model] = \ 193 | rospy.Publisher( 194 | '{}/pose_{}'.format(params['topic_publishing'], model), 195 | PoseStamped, 196 | queue_size=10 197 | ) 198 | pub_dimension[model] = \ 199 | rospy.Publisher( 200 | '{}/dimension_{}'.format(params['topic_publishing'], model), 201 | String, 202 | queue_size=10 203 | ) 204 | 205 | # Start ROS publisher 206 | pub_rgb_dope_points = \ 207 | rospy.Publisher( 208 | params['topic_publishing']+"/rgb_points", 209 | ImageSensor_msg, 210 | queue_size=10 211 | ) 212 | 213 | #publish pose 7 value 214 | pub = rospy.Publisher('pose_pos_ori', String, queue_size=10) 215 | 216 | # Starts ROS listener 217 | rospy.Subscriber( 218 | topic_cam, 219 | ImageSensor_msg, 220 | __image_callback 221 | ) 222 | 223 | # Initialize ROS node 224 | rospy.init_node('dope_vis', anonymous=True) 225 | rate = rospy.Rate(freq) 226 | 227 | print ("Running DOPE... (Listening to camera topic: '{}')".format(topic_cam)) 228 | print ("Ctrl-C to stop") 229 | 230 | while not rospy.is_shutdown(): 231 | if g_img is not None: 232 | # Copy and draw image 233 | img_copy = g_img.copy() 234 | im = Image.fromarray(img_copy) 235 | g_draw = ImageDraw.Draw(im) 236 | 237 | for m in models: 238 | # Detect object 239 | results = ObjectDetector.detect_object_in_image( 240 | models[m].net, 241 | pnp_solvers[m], 242 | g_img, 243 | config_detect 244 | ) 245 | 246 | # Publish pose and overlay cube on image 247 | for i_r, result in enumerate(results): 248 | if result["location"] is None: 249 | continue 250 | loc = result["location"] 251 | ori = result["quaternion"] 252 | msg = PoseStamped() 253 | msg.header.frame_id = params["frame_id"] 254 | msg.header.stamp = rospy.Time.now() 255 | CONVERT_SCALE_CM_TO_METERS = 100 256 | msg.pose.position.x = loc[0] / CONVERT_SCALE_CM_TO_METERS 257 | msg.pose.position.y = loc[1] / CONVERT_SCALE_CM_TO_METERS 258 | msg.pose.position.z = loc[2] / CONVERT_SCALE_CM_TO_METERS 259 | msg.pose.orientation.x = ori[0] 260 | msg.pose.orientation.y = ori[1] 261 | msg.pose.orientation.z = ori[2] 262 | msg.pose.orientation.w = ori[3] 263 | 264 | ######################### 265 | # output translation X,Y,Z to_topic:translation 266 | x = msg.pose.position.x 267 | y = msg.pose.position.y 268 | z = msg.pose.position.z 269 | 270 | # translation = ("%f,%f,%f" % (x, y, z)) 271 | # talkerT(translation) 272 | 273 | # output rotate rx,ry,rz to_topic:rotate 274 | rx = msg.pose.orientation.x 275 | ry = msg.pose.orientation.y 276 | rz = msg.pose.orientation.z 277 | rw = msg.pose.orientation.w 278 | 279 | a = [rx, ry, rz, rw] 280 | q = tf.transformations.euler_from_quaternion(a) 281 | c = q[1] + 1.57 282 | b = tf.transformations.quaternion_from_euler(q[0], c, q[2]) 283 | 284 | # quaternion_to_Euler-angle 285 | #rx = math.atan2(2*(x1*w1 + y1*z1), 1-2*((x1)**2 + (y1)**2)) 286 | #ry = math.asin(2*(w1*y1 - z1*x1)) 287 | #rz = math.atan2(2*(w1*x1+x1*y1), 1-2*(y1**2 + z1**2)) 288 | 289 | # pose_pos_ori = ("%f,%f,%f,%f,%f,%f,%f" % (x, y, z, b[0], b[1], b[2], b[3])) 290 | #pose_pos_ori = ("x:%f,y:%f,z:%f,rx:%f,ry:%f,rz:%f,rw:%f" % (x, y, z, rx, ry, rz, rw)) 291 | pose_pos_ori = ("%f,%f,%f,%f,%f,%f,%f" % (x, y, z, rx, ry, rz, rw)) 292 | rospy.loginfo(pose_pos_ori) 293 | pub.publish(pose_pos_ori) 294 | ########################## 295 | 296 | # Publish 297 | pubs[m].publish(msg) 298 | pub_dimension[m].publish(str(params['dimensions'][m])) 299 | 300 | # Draw the cube 301 | if None not in result['projected_points']: 302 | points2d = [] 303 | for pair in result['projected_points']: 304 | points2d.append(tuple(pair)) 305 | DrawCube(points2d, draw_colors[m]) 306 | 307 | # Publish the image with results overlaid 308 | pub_rgb_dope_points.publish( 309 | CvBridge().cv2_to_imgmsg( 310 | np.array(im)[..., ::-1], 311 | "bgr8" 312 | ) 313 | ) 314 | rate.sleep() 315 | 316 | 317 | if __name__ == "__main__": 318 | '''Main routine to run DOPE''' 319 | 320 | if len(sys.argv) > 1: 321 | config_name = sys.argv[1] 322 | else: 323 | config_name = "config_pose.yaml" 324 | rospack = rospkg.RosPack() 325 | params = None 326 | yaml_path = g_path2package + '/config/{}'.format(config_name) 327 | with open(yaml_path, 'r') as stream: 328 | try: 329 | print("Loading DOPE parameters from '{}'...".format(yaml_path)) 330 | params = yaml.load(stream) 331 | print(' Parameters loaded.') 332 | except yaml.YAMLError as exc: 333 | print(exc) 334 | 335 | topic_cam = params['topic_camera'] 336 | 337 | try : 338 | run_dope_node(params) 339 | except rospy.ROSInterruptException: 340 | pass 341 | -------------------------------------------------------------------------------- /src/inference/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 License. 3 | # https://creativecommons.org/licenses/by-nc-sa/4.0/legalcode 4 | 5 | from enum import IntEnum, unique 6 | import numpy as np 7 | import cv2 8 | from pyrr import Quaternion, Matrix44, Vector3, euler 9 | 10 | # Related to the object's local coordinate system 11 | # @unique 12 | class CuboidVertexType(IntEnum): 13 | FrontTopRight = 0 14 | FrontTopLeft = 1 15 | FrontBottomLeft = 2 16 | FrontBottomRight = 3 17 | RearTopRight = 4 18 | RearTopLeft = 5 19 | RearBottomLeft = 6 20 | RearBottomRight = 7 21 | Center = 8 22 | TotalCornerVertexCount = 8 # Corner vertexes doesn't include the center point 23 | TotalVertexCount = 9 24 | 25 | # List of the vertex indexes in each line edges of the cuboid 26 | CuboidLineIndexes = [ 27 | # Front face 28 | [ CuboidVertexType.FrontTopLeft, CuboidVertexType.FrontTopRight ], 29 | [ CuboidVertexType.FrontTopRight, CuboidVertexType.FrontBottomRight ], 30 | [ CuboidVertexType.FrontBottomRight, CuboidVertexType.FrontBottomLeft ], 31 | [ CuboidVertexType.FrontBottomLeft, CuboidVertexType.FrontTopLeft ], 32 | # Back face 33 | [ CuboidVertexType.RearTopLeft, CuboidVertexType.RearTopRight ], 34 | [ CuboidVertexType.RearTopRight, CuboidVertexType.RearBottomRight ], 35 | [ CuboidVertexType.RearBottomRight, CuboidVertexType.RearBottomLeft ], 36 | [ CuboidVertexType.RearBottomLeft, CuboidVertexType.RearTopLeft ], 37 | # Left face 38 | [ CuboidVertexType.FrontBottomLeft, CuboidVertexType.RearBottomLeft ], 39 | [ CuboidVertexType.FrontTopLeft, CuboidVertexType.RearTopLeft ], 40 | # Right face 41 | [ CuboidVertexType.FrontBottomRight, CuboidVertexType.RearBottomRight ], 42 | [ CuboidVertexType.FrontTopRight, CuboidVertexType.RearTopRight ], 43 | ] 44 | 45 | 46 | # ========================= Cuboid3d ========================= 47 | class Cuboid3d(): 48 | '''This class contains a 3D cuboid.''' 49 | 50 | # Create a box with a certain size 51 | def __init__(self, size3d = [1.0, 1.0, 1.0], center_location = [0, 0, 0], 52 | coord_system = None, parent_object = None): 53 | 54 | # NOTE: This local coordinate system is similar 55 | # to the intrinsic transform matrix of a 3d object 56 | self.center_location = center_location 57 | self.coord_system = coord_system 58 | self.size3d = size3d 59 | self._vertices = [0, 0, 0] * CuboidVertexType.TotalVertexCount 60 | 61 | self.generate_vertexes() 62 | 63 | def get_vertex(self, vertex_type): 64 | """Returns the location of a vertex. 65 | 66 | Args: 67 | vertex_type: enum of type CuboidVertexType 68 | 69 | Returns: 70 | Numpy array(3) - Location of the vertex type in the cuboid 71 | """ 72 | return self._vertices[vertex_type] 73 | 74 | def get_vertices(self): 75 | return self._vertices 76 | 77 | def generate_vertexes(self): 78 | width, height, depth = self.size3d 79 | 80 | # By default just use the normal OpenCV coordinate system 81 | if (self.coord_system is None): 82 | cx, cy, cz = self.center_location 83 | # X axis point to the right 84 | right = cx + width / 2.0 85 | left = cx - width / 2.0 86 | # Y axis point downward 87 | top = cy - height / 2.0 88 | bottom = cy + height / 2.0 89 | # Z axis point forward 90 | front = cz + depth / 2.0 91 | rear = cz - depth / 2.0 92 | 93 | # List of 8 vertices of the box 94 | self._vertices = [ 95 | [right, top, front], # Front Top Right 96 | [left, top, front], # Front Top Left 97 | [left, bottom, front], # Front Bottom Left 98 | [right, bottom, front], # Front Bottom Right 99 | [right, top, rear], # Rear Top Right 100 | [left, top, rear], # Rear Top Left 101 | [left, bottom, rear], # Rear Bottom Left 102 | [right, bottom, rear], # Rear Bottom Right 103 | self.center_location, # Center 104 | ] 105 | else: 106 | sx, sy, sz = self.size3d 107 | forward = np.array(self.coord_system.forward, dtype=float) * sy * 0.5 108 | up = np.array(self.coord_system.up, dtype=float) * sz * 0.5 109 | right = np.array(self.coord_system.right, dtype=float) * sx * 0.5 110 | center = np.array(self.center_location, dtype=float) 111 | self._vertices = [ 112 | center + forward + up + right, # Front Top Right 113 | center + forward + up - right, # Front Top Left 114 | center + forward - up - right, # Front Bottom Left 115 | center + forward - up + right, # Front Bottom Right 116 | center - forward + up + right, # Rear Top Right 117 | center - forward + up - right, # Rear Top Left 118 | center - forward - up - right, # Rear Bottom Left 119 | center - forward - up + right, # Rear Bottom Right 120 | self.center_location, # Center 121 | ] 122 | 123 | def get_projected_cuboid2d(self, cuboid_transform, camera_intrinsic_matrix): 124 | """ 125 | Projects the cuboid into the image plane using camera intrinsics. 126 | 127 | Args: 128 | cuboid_transform: the world transform of the cuboid 129 | camera_intrinsic_matrix: camera intrinsic matrix 130 | 131 | Returns: 132 | Cuboid2d - the projected cuboid points 133 | """ 134 | 135 | world_transform_matrix = cuboid_transform 136 | rvec = [0, 0, 0] 137 | tvec = [0, 0, 0] 138 | dist_coeffs = np.zeros((4, 1)) 139 | 140 | transformed_vertices = [0, 0, 0] * CuboidVertexType.TotalVertexCount 141 | for vertex_index in range(CuboidVertexType.TotalVertexCount): 142 | vertex3d = self._vertices[vertex_index] 143 | transformed_vertices[vertex_index] = world_transform_matrix * vertex3d 144 | 145 | projected_vertices = cv2.projectPoints(transformed_vertices, rvec, tvec, 146 | camera_intrinsic_matrix, dist_coeffs) 147 | 148 | return Cuboid2d(projected_vertices) 149 | -------------------------------------------------------------------------------- /src/inference/cuboid_pnp_solver.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 License. 3 | # https://creativecommons.org/licenses/by-nc-sa/4.0/legalcode 4 | 5 | import numpy as np 6 | from pyrr import Quaternion 7 | import cv2 8 | from cuboid import * 9 | 10 | class CuboidPNPSolver(object): 11 | """ 12 | This class is used to find the 6-DoF pose of a cuboid given its projected vertices. 13 | 14 | Runs perspective-n-point (PNP) algorithm. 15 | """ 16 | 17 | # Class variables 18 | cv2version = cv2.__version__.split('.') 19 | cv2majorversion = int(cv2version[0]) 20 | 21 | def __init__(self, object_name="", camera_intrinsic_matrix = None, cuboid3d = None, 22 | dist_coeffs = np.zeros((4, 1))): 23 | self.object_name = object_name 24 | if (not camera_intrinsic_matrix is None): 25 | self._camera_intrinsic_matrix = camera_intrinsic_matrix 26 | else: 27 | camera_intrinsic_matrix = np.array([ 28 | [0, 0, 0], 29 | [0, 0, 0], 30 | [0, 0, 0] 31 | ]) 32 | self._cuboid3d = cuboid3d 33 | 34 | self._dist_coeffs = dist_coeffs 35 | 36 | def set_camera_intrinsic_matrix(self, new_intrinsic_matrix): 37 | '''Sets the camera intrinsic matrix''' 38 | self._camera_intrinsic_matrix = new_intrinsic_matrix 39 | 40 | def solve_pnp(self, cuboid2d_points, pnp_algorithm = None): 41 | """ 42 | Detects the rotation and traslation 43 | of a cuboid object from its vertexes' 44 | 2D location in the image 45 | """ 46 | 47 | # Fallback to default PNP algorithm base on OpenCV version 48 | if pnp_algorithm is None: 49 | if CuboidPNPSolver.cv2majorversion == 2: 50 | pnp_algorithm = cv2.CV_ITERATIVE 51 | elif CuboidPNPSolver.cv2majorversion == 3: 52 | pnp_algorithm = cv2.SOLVEPNP_ITERATIVE 53 | # Alternative algorithms: 54 | # pnp_algorithm = SOLVE_PNP_P3P 55 | # pnp_algorithm = SOLVE_PNP_EPNP 56 | 57 | location = None 58 | quaternion = None 59 | projected_points = cuboid2d_points 60 | 61 | cuboid3d_points = np.array(self._cuboid3d.get_vertices()) 62 | obj_2d_points = [] 63 | obj_3d_points = [] 64 | 65 | for i in range(CuboidVertexType.TotalVertexCount): 66 | check_point_2d = cuboid2d_points[i] 67 | # Ignore invalid points 68 | if (check_point_2d is None): 69 | continue 70 | obj_2d_points.append(check_point_2d) 71 | obj_3d_points.append(cuboid3d_points[i]) 72 | 73 | obj_2d_points = np.array(obj_2d_points, dtype=float) 74 | obj_3d_points = np.array(obj_3d_points, dtype=float) 75 | 76 | valid_point_count = len(obj_2d_points) 77 | 78 | # Can only do PNP if we have more than 3 valid points 79 | is_points_valid = valid_point_count >= 4 80 | 81 | if is_points_valid: 82 | 83 | ret, rvec, tvec = cv2.solvePnP( 84 | obj_3d_points, 85 | obj_2d_points, 86 | self._camera_intrinsic_matrix, 87 | self._dist_coeffs, 88 | flags=pnp_algorithm 89 | ) 90 | 91 | if ret: 92 | location = list(x[0] for x in tvec) 93 | quaternion = self.convert_rvec_to_quaternion(rvec) 94 | 95 | projected_points, _ = cv2.projectPoints(cuboid3d_points, rvec, tvec, self._camera_intrinsic_matrix, self._dist_coeffs) 96 | projected_points = np.squeeze(projected_points) 97 | 98 | # If the location.Z is negative or object is behind the camera then flip both location and rotation 99 | x, y, z = location 100 | if z < 0: 101 | # Get the opposite location 102 | location = [-x, -y, -z] 103 | 104 | # Change the rotation by 180 degree 105 | rotate_angle = np.pi 106 | rotate_quaternion = Quaternion.from_axis_rotation(location, rotate_angle) 107 | quaternion = rotate_quaternion.cross(quaternion) 108 | 109 | return location, quaternion, projected_points 110 | 111 | def convert_rvec_to_quaternion(self, rvec): 112 | '''Convert rvec (which is log quaternion) to quaternion''' 113 | theta = np.sqrt(rvec[0] * rvec[0] + rvec[1] * rvec[1] + rvec[2] * rvec[2]) # in radians 114 | raxis = [rvec[0] / theta, rvec[1] / theta, rvec[2] / theta] 115 | 116 | # pyrr's Quaternion (order is XYZW), https://pyrr.readthedocs.io/en/latest/oo_api_quaternion.html 117 | return Quaternion.from_axis_rotation(raxis, theta) 118 | 119 | # Alternatively: pyquaternion 120 | # return Quaternion(axis=raxis, radians=theta) # uses OpenCV's Quaternion (order is WXYZ) 121 | 122 | def project_points(self, rvec, tvec): 123 | '''Project points from model onto image using rotation, translation''' 124 | output_points, tmp = cv2.projectPoints( 125 | self.__object_vertex_coordinates, 126 | rvec, 127 | tvec, 128 | self.__camera_intrinsic_matrix, 129 | self.__dist_coeffs) 130 | 131 | output_points = np.squeeze(output_points) 132 | return output_points -------------------------------------------------------------------------------- /src/inference/detector_FastNet-V1.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 License. 3 | # https://creativecommons.org/licenses/by-nc-sa/4.0/legalcode 4 | 5 | ''' 6 | Contains the following classes: 7 | - ModelData - High level information encapsulation 8 | - ObjectDetector - Greedy algorithm to build cuboids from belief maps 9 | ''' 10 | 11 | import time 12 | import json 13 | import os, shutil 14 | import sys 15 | import traceback 16 | from os import path 17 | import threading 18 | from threading import Thread 19 | 20 | import numpy as np 21 | import cv2 22 | 23 | import torch 24 | import torch.nn as nn 25 | import torchvision.transforms as transforms 26 | from torch.autograd import Variable 27 | import torchvision.models as models 28 | 29 | #torch Dict 30 | from collections import OrderedDict 31 | 32 | from scipy import ndimage 33 | import scipy 34 | import scipy.ndimage as ndimage 35 | import scipy.ndimage.filters as filters 36 | from scipy.ndimage.filters import gaussian_filter 37 | 38 | # Import the definition of the neural network model and cuboids 39 | from cuboid_pnp_solver import * 40 | 41 | #global transform for image input 42 | transform = transforms.Compose([ 43 | # transforms.Scale(IMAGE_SIZE), 44 | # transforms.CenterCrop((imagesize,imagesize)), 45 | transforms.ToTensor(), 46 | transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)), 47 | ]) 48 | 49 | #ResNet-50/101 50 | class Bottleneck(nn.Module): 51 | expansion = 4 52 | 53 | def __init__(self, inplanes, planes, stride=1, downsample=None): 54 | super(Bottleneck, self).__init__() 55 | self.conv1 = nn.Conv2d(inplanes, planes, kernel_size=1, bias=False) 56 | self.bn1 = nn.BatchNorm2d(planes) 57 | self.conv2 = nn.Conv2d(planes, planes, kernel_size=3, stride=stride, padding=1, bias=False) 58 | self.bn2 = nn.BatchNorm2d(planes) 59 | self.conv3 = nn.Conv2d(planes, planes * 4, kernel_size=1, bias=False) 60 | self.bn3 = nn.BatchNorm2d(planes * 4) 61 | self.relu = nn.ReLU(inplace=True) 62 | self.downsample = downsample 63 | self.stride = stride 64 | 65 | def forward(self, x): 66 | residual = x 67 | 68 | out = self.conv1(x) 69 | out = self.bn1(out) 70 | out = self.relu(out) 71 | 72 | out = self.conv2(out) 73 | out = self.bn2(out) 74 | out = self.relu(out) 75 | 76 | out = self.conv3(out) 77 | out = self.bn3(out) 78 | 79 | if self.downsample is not None: 80 | residual = self.downsample(x) 81 | 82 | out += residual 83 | out = self.relu(out) 84 | 85 | return out 86 | #ResNet-18/34 87 | class BasicBlock(nn.Module): 88 | expansion = 1 89 | 90 | def __init__(self, inplanes, planes, stride=1, downsample=None): 91 | super(BasicBlock, self).__init__() 92 | self.conv1 = conv3x3(inplanes, planes, stride) 93 | self.bn1 = nn.BatchNorm2d(planes) 94 | self.relu = nn.ReLU(inplace=True) 95 | self.conv2 = conv3x3(planes, planes) 96 | self.bn2 = nn.BatchNorm2d(planes) 97 | self.downsample = downsample 98 | self.stride = stride 99 | 100 | def forward(self, x): 101 | residual = x 102 | 103 | out = self.conv1(x) 104 | out = self.bn1(out) 105 | out = self.relu(out) 106 | 107 | out = self.conv2(out) 108 | out = self.bn2(out) 109 | 110 | if self.downsample is not None: 111 | residual = self.downsample(x) 112 | 113 | out += residual 114 | out = self.relu(out) 115 | 116 | return out 117 | 118 | def conv3x3(in_planes, out_planes, stride=1): 119 | """3x3 convolution with padding""" 120 | return nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride, 121 | padding=1, bias=False) 122 | 123 | #================================ Models ================================ 124 | 125 | 126 | class DopeNetwork(nn.Module): 127 | def __init__( 128 | self, 129 | numBeliefMap=9, 130 | numAffinity=16, 131 | stop_at_stage=5 # number of stages to process (if less than total number of stages) 132 | ): 133 | global block 134 | block = BasicBlock 135 | global layers 136 | layers = [3, 4, 6, 2] 137 | self.inplanes = 64 138 | 139 | super(DopeNetwork, self).__init__() 140 | 141 | self.stop_at_stage = stop_at_stage 142 | 143 | self.vgg = nn.Sequential(OrderedDict([ 144 | ('0', nn.Conv2d(3, 64, kernel_size=7, stride=2, padding=3, bias=False)), 145 | ('1', nn.BatchNorm2d(64)), 146 | ('2', nn.ReLU(inplace=True)), 147 | ('3', nn.MaxPool2d(kernel_size=3, stride=2, padding=0)), 148 | ##block 149 | ('4', self._make_layer(block, 64, layers[0], stride=1)), 150 | ('5', self._make_layer(block, 128, layers[1], stride=1)), 151 | ('6', self._make_layer(block, 256, layers[2], stride=2)), 152 | ('7', self._make_layer(block, 512, layers[3], stride=1)) 153 | ])) 154 | for i_layer in range(8): 155 | self.vgg.add_module(str(i_layer), self.vgg[i_layer]) 156 | 157 | # Add some layers 158 | i_layer = 7 159 | self.vgg.add_module(str(i_layer), nn.Conv2d(256, 128, kernel_size=3, stride=1, padding=1)) 160 | self.vgg.add_module(str(i_layer+1), nn.ReLU(inplace=True)) 161 | 162 | # print('---Belief------------------------------------------------') 163 | # _2 are the belief map stages 164 | self.m1_2 = DopeNetwork.create_stage(128, numBeliefMap, True) 165 | self.m2_2 = DopeNetwork.create_stage(128 + numBeliefMap + numAffinity, 166 | numBeliefMap, False) 167 | self.m3_2 = DopeNetwork.create_stage(128 + numBeliefMap + numAffinity, 168 | numBeliefMap, False) 169 | self.m4_2 = DopeNetwork.create_stage(128 + numBeliefMap + numAffinity, 170 | numBeliefMap, False) 171 | self.m5_2 = DopeNetwork.create_stage(128 + numBeliefMap + numAffinity, 172 | numBeliefMap, False) 173 | self.m6_2 = DopeNetwork.create_stage(128 + numBeliefMap + numAffinity, 174 | numBeliefMap, False) 175 | 176 | # print('---Affinity----------------------------------------------') 177 | # _1 are the affinity map stages 178 | self.m1_1 = DopeNetwork.create_stage(128, numAffinity, True) 179 | self.m2_1 = DopeNetwork.create_stage(128 + numBeliefMap + numAffinity, 180 | numAffinity, False) 181 | self.m3_1 = DopeNetwork.create_stage(128 + numBeliefMap + numAffinity, 182 | numAffinity, False) 183 | self.m4_1 = DopeNetwork.create_stage(128 + numBeliefMap + numAffinity, 184 | numAffinity, False) 185 | self.m5_1 = DopeNetwork.create_stage(128 + numBeliefMap + numAffinity, 186 | numAffinity, False) 187 | self.m6_1 = DopeNetwork.create_stage(128 + numBeliefMap + numAffinity, 188 | numAffinity, False) 189 | 190 | def _make_layer(self, block, planes, blocks, stride=1): 191 | downsample = None 192 | if stride != 1 or self.inplanes != planes * block.expansion: 193 | downsample = nn.Sequential( 194 | nn.Conv2d(self.inplanes, planes * block.expansion, 195 | kernel_size=1, stride=stride, bias=False), 196 | nn.BatchNorm2d(planes * block.expansion), 197 | ) 198 | layers = [] 199 | layers.append(block(self.inplanes, planes, stride, downsample)) 200 | self.inplanes = planes * block.expansion 201 | for i in range(1, blocks): 202 | layers.append(block(self.inplanes, planes)) 203 | 204 | return nn.Sequential(*layers) 205 | 206 | def forward(self, x): 207 | '''Runs inference on the neural network''' 208 | # output from ResNet 209 | out1 = self.vgg(x) 210 | 211 | out1_2 = self.m1_2(out1) 212 | out1_1 = self.m1_1(out1) 213 | 214 | if self.stop_at_stage == 1: 215 | return [out1_2], \ 216 | [out1_1] 217 | 218 | out2 = torch.cat([out1_2, out1_1, out1], 1) 219 | out2_2 = self.m2_2(out2) 220 | out2_1 = self.m2_1(out2) 221 | 222 | if self.stop_at_stage == 2: 223 | return [out1_2, out2_2], \ 224 | [out1_1, out2_1] 225 | 226 | out3 = torch.cat([out2_2, out2_1, out1], 1) 227 | out3_2 = self.m3_2(out3) 228 | out3_1 = self.m3_1(out3) 229 | 230 | if self.stop_at_stage == 3: 231 | return [out1_2, out2_2, out3_2], \ 232 | [out1_1, out2_1, out3_1] 233 | 234 | out4 = torch.cat([out3_2, out3_1, out1], 1) 235 | out4_2 = self.m4_2(out4) 236 | out4_1 = self.m4_1(out4) 237 | 238 | if self.stop_at_stage == 4: 239 | return [out1_2, out2_2, out3_2, out4_2], \ 240 | [out1_1, out2_1, out3_1, out4_1] 241 | 242 | out5 = torch.cat([out4_2, out4_1, out1], 1) 243 | out5_2 = self.m5_2(out5) 244 | out5_1 = self.m5_1(out5) 245 | 246 | if self.stop_at_stage == 5: 247 | return [out1_2, out2_2, out3_2, out4_2, out5_2],\ 248 | [out1_1, out2_1, out3_1, out4_1, out5_1] 249 | 250 | @staticmethod 251 | def create_stage(in_channels, out_channels, first=False): 252 | '''Create the neural network layers for a single stage.''' 253 | 254 | model = nn.Sequential() 255 | mid_channels = 128 256 | if first: 257 | padding = 1 258 | kernel = 3 259 | count = 6 260 | final_channels = 512 261 | else: 262 | padding = 3 263 | kernel = 7 264 | count = 10 265 | final_channels = mid_channels 266 | 267 | # First convolution 268 | model.add_module("0", 269 | nn.Conv2d( 270 | in_channels, 271 | mid_channels, 272 | kernel_size=kernel, 273 | stride=1, 274 | padding=padding) 275 | ) 276 | 277 | # Middle convolutions 278 | i = 1 279 | while i < count - 1: 280 | model.add_module(str(i), nn.ReLU(inplace=True)) 281 | i += 1 282 | model.add_module(str(i), 283 | nn.Conv2d( 284 | mid_channels, 285 | mid_channels, 286 | kernel_size=kernel, 287 | stride=1, 288 | padding=padding)) 289 | i += 1 290 | 291 | # Penultimate convolution 292 | model.add_module(str(i), nn.ReLU(inplace=True)) 293 | i += 1 294 | model.add_module(str(i), nn.Conv2d(mid_channels, final_channels, kernel_size=1, stride=1)) 295 | i += 1 296 | 297 | # Last convolution 298 | model.add_module(str(i), nn.ReLU(inplace=True)) 299 | i += 1 300 | model.add_module(str(i), nn.Conv2d(final_channels, out_channels, kernel_size=1, stride=1)) 301 | i += 1 302 | 303 | return model 304 | 305 | 306 | 307 | class ModelData(object): 308 | '''This class contains methods for loading the neural network''' 309 | 310 | def __init__(self, name="", net_path="", gpu_id=0): 311 | self.name = name 312 | self.net_path = net_path # Path to trained network model 313 | self.net = None # Trained network 314 | self.gpu_id = gpu_id 315 | 316 | def get_net(self): 317 | '''Returns network''' 318 | if not self.net: 319 | self.load_net_model() 320 | return self.net 321 | 322 | def load_net_model(self): 323 | '''Loads network model from disk''' 324 | if not self.net and path.exists(self.net_path): 325 | self.net = self.load_net_model_path(self.net_path) 326 | if not path.exists(self.net_path): 327 | print("ERROR: Unable to find model weights: '{}'".format( 328 | self.net_path)) 329 | exit(0) 330 | 331 | def load_net_model_path(self, path): 332 | '''Loads network model from disk with given path''' 333 | model_loading_start_time = time.time() 334 | print("Loading DOPE model '{}'...".format(path)) 335 | net = DopeNetwork() 336 | net = torch.nn.DataParallel(net, [0]).cuda() 337 | net.load_state_dict(torch.load(path)) 338 | net.eval() 339 | print(' Model loaded in {} seconds.'.format( 340 | time.time() - model_loading_start_time)) 341 | return net 342 | 343 | def __str__(self): 344 | '''Converts to string''' 345 | return "{}: {}".format(self.name, self.net_path) 346 | 347 | 348 | #================================ ObjectDetector ================================ 349 | class ObjectDetector(object): 350 | '''This class contains methods for object detection''' 351 | 352 | @staticmethod 353 | def detect_object_in_image(net_model, pnp_solver, in_img, config): 354 | '''Detect objects in a image using a specific trained network model''' 355 | 356 | if in_img is None: 357 | return [] 358 | 359 | # Run network inference 360 | image_tensor = transform(in_img) 361 | image_torch = Variable(image_tensor).cuda().unsqueeze(0) 362 | out, seg = net_model(image_torch) 363 | vertex2 = out[-1][0] 364 | aff = seg[-1][0] 365 | 366 | # Find objects from network output 367 | detected_objects = ObjectDetector.find_object_poses(vertex2, aff, pnp_solver, config) 368 | return detected_objects 369 | 370 | @staticmethod 371 | def find_object_poses(vertex2, aff, pnp_solver, config): 372 | '''Detect objects given network output''' 373 | 374 | # Detect objects from belief maps and affinities 375 | objects, all_peaks = ObjectDetector.find_objects(vertex2, aff, config) 376 | detected_objects = [] 377 | obj_name = pnp_solver.object_name 378 | 379 | for obj in objects: 380 | # Run PNP 381 | points = obj[1] + [(obj[0][0]*8, obj[0][1]*8)] 382 | cuboid2d = np.copy(points) 383 | location, quaternion, projected_points = pnp_solver.solve_pnp(points) 384 | 385 | # Save results 386 | detected_objects.append({ 387 | 'name': obj_name, 388 | 'location': location, 389 | 'quaternion': quaternion, 390 | 'cuboid2d': cuboid2d, 391 | 'projected_points': projected_points, 392 | }) 393 | 394 | return detected_objects 395 | 396 | @staticmethod 397 | def find_objects(vertex2, aff, config, numvertex=8): 398 | '''Detects objects given network belief maps and affinities, using heuristic method''' 399 | 400 | all_peaks = [] 401 | peak_counter = 0 402 | for j in range(vertex2.size()[0]): 403 | belief = vertex2[j].clone() 404 | map_ori = belief.cpu().data.numpy() 405 | 406 | map = gaussian_filter(belief.cpu().data.numpy(), sigma=config.sigma) 407 | p = 1 408 | map_left = np.zeros(map.shape) 409 | map_left[p:,:] = map[:-p,:] 410 | map_right = np.zeros(map.shape) 411 | map_right[:-p,:] = map[p:,:] 412 | map_up = np.zeros(map.shape) 413 | map_up[:,p:] = map[:,:-p] 414 | map_down = np.zeros(map.shape) 415 | map_down[:,:-p] = map[:,p:] 416 | 417 | peaks_binary = np.logical_and.reduce( 418 | ( 419 | map >= map_left, 420 | map >= map_right, 421 | map >= map_up, 422 | map >= map_down, 423 | map > config.thresh_map) 424 | ) 425 | peaks = zip(np.nonzero(peaks_binary)[1], np.nonzero(peaks_binary)[0]) 426 | 427 | # Computing the weigthed average for localizing the peaks 428 | peaks = list(peaks) 429 | win = 5 430 | ran = win // 2 431 | peaks_avg = [] 432 | for p_value in range(len(peaks)): 433 | p = peaks[p_value] 434 | weights = np.zeros((win,win)) 435 | i_values = np.zeros((win,win)) 436 | j_values = np.zeros((win,win)) 437 | for i in range(-ran,ran+1): 438 | for j in range(-ran,ran+1): 439 | if p[1]+i < 0 \ 440 | or p[1]+i >= map_ori.shape[0] \ 441 | or p[0]+j < 0 \ 442 | or p[0]+j >= map_ori.shape[1]: 443 | continue 444 | 445 | i_values[j+ran, i+ran] = p[1] + i 446 | j_values[j+ran, i+ran] = p[0] + j 447 | 448 | weights[j+ran, i+ran] = (map_ori[p[1]+i, p[0]+j]) 449 | 450 | # if the weights are all zeros 451 | # then add the none continuous points 452 | OFFSET_DUE_TO_UPSAMPLING = 0.4395 453 | try: 454 | peaks_avg.append( 455 | (np.average(j_values, weights=weights) + OFFSET_DUE_TO_UPSAMPLING, \ 456 | np.average(i_values, weights=weights) + OFFSET_DUE_TO_UPSAMPLING)) 457 | except: 458 | peaks_avg.append((p[0] + OFFSET_DUE_TO_UPSAMPLING, p[1] + OFFSET_DUE_TO_UPSAMPLING)) 459 | # Note: Python3 doesn't support len for zip object 460 | peaks_len = min(len(np.nonzero(peaks_binary)[1]), len(np.nonzero(peaks_binary)[0])) 461 | 462 | peaks_with_score = [peaks_avg[x_] + (map_ori[peaks[x_][1], peaks[x_][0]],) for x_ in range(len(peaks))] 463 | 464 | id = range(peak_counter, peak_counter + peaks_len) 465 | 466 | peaks_with_score_and_id = [peaks_with_score[i] + (id[i],) for i in range(len(id))] 467 | 468 | all_peaks.append(peaks_with_score_and_id) 469 | peak_counter += peaks_len 470 | 471 | objects = [] 472 | 473 | # Check object centroid and build the objects if the centroid is found 474 | for nb_object in range(len(all_peaks[-1])): 475 | if all_peaks[-1][nb_object][2] > config.thresh_points: 476 | objects.append([ 477 | [all_peaks[-1][nb_object][:2][0],all_peaks[-1][nb_object][:2][1]], 478 | [None for i in range(numvertex)], 479 | [None for i in range(numvertex)], 480 | all_peaks[-1][nb_object][2] 481 | ]) 482 | 483 | # Working with an output that only has belief maps 484 | if aff is None: 485 | if len (objects) > 0 and len(all_peaks)>0 and len(all_peaks[0])>0: 486 | for i_points in range(8): 487 | if len(all_peaks[i_points])>0 and all_peaks[i_points][0][2] > config.threshold: 488 | objects[0][1][i_points] = (all_peaks[i_points][0][0], all_peaks[i_points][0][1]) 489 | else: 490 | # For all points found 491 | for i_lists in range(len(all_peaks[:-1])): 492 | lists = all_peaks[i_lists] 493 | 494 | for candidate in lists: 495 | if candidate[2] < config.thresh_points: 496 | continue 497 | 498 | i_best = -1 499 | best_dist = 10000 500 | best_angle = 100 501 | for i_obj in range(len(objects)): 502 | center = [objects[i_obj][0][0], objects[i_obj][0][1]] 503 | 504 | # integer is used to look into the affinity map, 505 | # but the float version is used to run 506 | point_int = [int(candidate[0]), int(candidate[1])] 507 | point = [candidate[0], candidate[1]] 508 | 509 | # look at the distance to the vector field. 510 | v_aff = np.array([ 511 | aff[i_lists*2, 512 | point_int[1], 513 | point_int[0]].data.item(), 514 | aff[i_lists*2+1, 515 | point_int[1], 516 | point_int[0]].data.item()]) * 10 517 | 518 | # normalize the vector 519 | xvec = v_aff[0] 520 | yvec = v_aff[1] 521 | 522 | norms = np.sqrt(xvec * xvec + yvec * yvec) 523 | 524 | xvec/=norms 525 | yvec/=norms 526 | 527 | v_aff = np.concatenate([[xvec], [yvec]]) 528 | 529 | v_center = np.array(center) - np.array(point) 530 | xvec = v_center[0] 531 | yvec = v_center[1] 532 | 533 | norms = np.sqrt(xvec * xvec + yvec * yvec) 534 | 535 | xvec /= norms 536 | yvec /= norms 537 | 538 | v_center = np.concatenate([[xvec], [yvec]]) 539 | 540 | # vector affinity 541 | dist_angle = np.linalg.norm(v_center - v_aff) 542 | 543 | # distance between vertexes 544 | dist_point = np.linalg.norm(np.array(point) - np.array(center)) 545 | 546 | if dist_angle < config.thresh_angle \ 547 | and best_dist > 1000 \ 548 | or dist_angle < config.thresh_angle \ 549 | and best_dist > dist_point: 550 | i_best = i_obj 551 | best_angle = dist_angle 552 | best_dist = dist_point 553 | 554 | if i_best is -1: 555 | continue 556 | 557 | if objects[i_best][1][i_lists] is None \ 558 | or best_angle < config.thresh_angle \ 559 | and best_dist < objects[i_best][2][i_lists][1]: 560 | objects[i_best][1][i_lists] = ((candidate[0])*8, (candidate[1])*8) 561 | objects[i_best][2][i_lists] = (best_angle, best_dist) 562 | 563 | return objects, all_peaks 564 | -------------------------------------------------------------------------------- /src/inference/detector_FastNet-V2.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 License. 3 | # https://creativecommons.org/licenses/by-nc-sa/4.0/legalcode 4 | 5 | ''' 6 | Contains the following classes: 7 | - ModelData - High level information encapsulation 8 | - ObjectDetector - Greedy algorithm to build cuboids from belief maps 9 | ''' 10 | 11 | import time 12 | import json 13 | import os, shutil 14 | import sys 15 | import traceback 16 | from os import path 17 | import threading 18 | from threading import Thread 19 | 20 | import numpy as np 21 | import cv2 22 | 23 | import torch 24 | import torch.nn as nn 25 | import torchvision.transforms as transforms 26 | from torch.autograd import Variable 27 | import torchvision.models as models 28 | 29 | #torch Dict 30 | from collections import OrderedDict 31 | 32 | from scipy import ndimage 33 | import scipy 34 | import scipy.ndimage as ndimage 35 | import scipy.ndimage.filters as filters 36 | from scipy.ndimage.filters import gaussian_filter 37 | 38 | # Import the definition of the neural network model and cuboids 39 | from cuboid_pnp_solver import * 40 | 41 | #global transform for image input 42 | transform = transforms.Compose([ 43 | # transforms.Scale(IMAGE_SIZE), 44 | # transforms.CenterCrop((imagesize,imagesize)), 45 | transforms.ToTensor(), 46 | transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)), 47 | ]) 48 | 49 | #ResNet-50/101 50 | class Bottleneck(nn.Module): 51 | expansion = 4 52 | 53 | def __init__(self, inplanes, planes, stride=1, downsample=None): 54 | super(Bottleneck, self).__init__() 55 | self.conv1 = nn.Conv2d(inplanes, planes, kernel_size=1, bias=False) 56 | self.bn1 = nn.BatchNorm2d(planes) 57 | self.conv2 = nn.Conv2d(planes, planes, kernel_size=3, stride=stride, padding=1, bias=False) 58 | self.bn2 = nn.BatchNorm2d(planes) 59 | self.conv3 = nn.Conv2d(planes, planes * 4, kernel_size=1, bias=False) 60 | self.bn3 = nn.BatchNorm2d(planes * 4) 61 | self.relu = nn.ReLU(inplace=True) 62 | self.downsample = downsample 63 | self.stride = stride 64 | 65 | def forward(self, x): 66 | residual = x 67 | 68 | out = self.conv1(x) 69 | out = self.bn1(out) 70 | out = self.relu(out) 71 | 72 | out = self.conv2(out) 73 | out = self.bn2(out) 74 | out = self.relu(out) 75 | 76 | out = self.conv3(out) 77 | out = self.bn3(out) 78 | 79 | if self.downsample is not None: 80 | residual = self.downsample(x) 81 | 82 | out += residual 83 | out = self.relu(out) 84 | 85 | return out 86 | #ResNet-18/34 87 | class BasicBlock(nn.Module): 88 | expansion = 1 89 | 90 | def __init__(self, inplanes, planes, stride=1, downsample=None): 91 | super(BasicBlock, self).__init__() 92 | self.conv1 = conv3x3(inplanes, planes, stride) 93 | self.bn1 = nn.BatchNorm2d(planes) 94 | self.relu = nn.ReLU(inplace=True) 95 | self.conv2 = conv3x3(planes, planes) 96 | self.bn2 = nn.BatchNorm2d(planes) 97 | self.downsample = downsample 98 | self.stride = stride 99 | 100 | def forward(self, x): 101 | residual = x 102 | 103 | out = self.conv1(x) 104 | out = self.bn1(out) 105 | out = self.relu(out) 106 | 107 | out = self.conv2(out) 108 | out = self.bn2(out) 109 | 110 | if self.downsample is not None: 111 | residual = self.downsample(x) 112 | 113 | out += residual 114 | out = self.relu(out) 115 | 116 | return out 117 | 118 | def conv3x3(in_planes, out_planes, stride=1): 119 | """3x3 convolution with padding""" 120 | return nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride, 121 | padding=1, bias=False) 122 | 123 | #================================ Models ================================ 124 | 125 | 126 | class DopeNetwork(nn.Module): 127 | def __init__( 128 | self, 129 | numBeliefMap=9, 130 | numAffinity=16, 131 | stop_at_stage=4 # number of stages to process (if less than total number of stages) 132 | ): 133 | global block 134 | block = BasicBlock 135 | global layers 136 | layers = [3, 4, 6, 2] 137 | self.inplanes = 64 138 | 139 | super(DopeNetwork, self).__init__() 140 | 141 | self.stop_at_stage = stop_at_stage 142 | 143 | self.vgg = nn.Sequential(OrderedDict([ 144 | ('0', nn.Conv2d(3, 64, kernel_size=7, stride=2, padding=3, bias=False)), 145 | ('1', nn.BatchNorm2d(64)), 146 | ('2', nn.ReLU(inplace=True)), 147 | ('3', nn.MaxPool2d(kernel_size=3, stride=2, padding=0)), 148 | ##block 149 | ('4', self._make_layer(block, 64, layers[0], stride=1)), 150 | ('5', self._make_layer(block, 128, layers[1], stride=1)), 151 | ('6', self._make_layer(block, 256, layers[2], stride=2)), 152 | ('7', self._make_layer(block, 512, layers[3], stride=1)) 153 | ])) 154 | for i_layer in range(8): 155 | self.vgg.add_module(str(i_layer), self.vgg[i_layer]) 156 | 157 | # Add some layers 158 | i_layer = 7 159 | self.vgg.add_module(str(i_layer), nn.Conv2d(256, 128, kernel_size=3, stride=1, padding=1)) 160 | self.vgg.add_module(str(i_layer+1), nn.ReLU(inplace=True)) 161 | 162 | # print('---Belief------------------------------------------------') 163 | # _2 are the belief map stages 164 | self.m1_2 = DopeNetwork.create_stage(128, numBeliefMap, True) 165 | self.m2_2 = DopeNetwork.create_stage(128 + numBeliefMap + numAffinity, 166 | numBeliefMap, False) 167 | self.m3_2 = DopeNetwork.create_stage(128 + numBeliefMap + numAffinity, 168 | numBeliefMap, False) 169 | self.m4_2 = DopeNetwork.create_stage(128 + numBeliefMap + numAffinity, 170 | numBeliefMap, False) 171 | self.m5_2 = DopeNetwork.create_stage(128 + numBeliefMap + numAffinity, 172 | numBeliefMap, False) 173 | self.m6_2 = DopeNetwork.create_stage(128 + numBeliefMap + numAffinity, 174 | numBeliefMap, False) 175 | 176 | # print('---Affinity----------------------------------------------') 177 | # _1 are the affinity map stages 178 | self.m1_1 = DopeNetwork.create_stage(128, numAffinity, True) 179 | self.m2_1 = DopeNetwork.create_stage(128 + numBeliefMap + numAffinity, 180 | numAffinity, False) 181 | self.m3_1 = DopeNetwork.create_stage(128 + numBeliefMap + numAffinity, 182 | numAffinity, False) 183 | self.m4_1 = DopeNetwork.create_stage(128 + numBeliefMap + numAffinity, 184 | numAffinity, False) 185 | self.m5_1 = DopeNetwork.create_stage(128 + numBeliefMap + numAffinity, 186 | numAffinity, False) 187 | self.m6_1 = DopeNetwork.create_stage(128 + numBeliefMap + numAffinity, 188 | numAffinity, False) 189 | 190 | def _make_layer(self, block, planes, blocks, stride=1): 191 | downsample = None 192 | if stride != 1 or self.inplanes != planes * block.expansion: 193 | downsample = nn.Sequential( 194 | nn.Conv2d(self.inplanes, planes * block.expansion, 195 | kernel_size=1, stride=stride, bias=False), 196 | nn.BatchNorm2d(planes * block.expansion), 197 | ) 198 | layers = [] 199 | layers.append(block(self.inplanes, planes, stride, downsample)) 200 | self.inplanes = planes * block.expansion 201 | for i in range(1, blocks): 202 | layers.append(block(self.inplanes, planes)) 203 | 204 | return nn.Sequential(*layers) 205 | 206 | def forward(self, x): 207 | '''Runs inference on the neural network''' 208 | # output from ResNet 209 | out1 = self.vgg(x) 210 | 211 | out1_2 = self.m1_2(out1) 212 | out1_1 = self.m1_1(out1) 213 | 214 | if self.stop_at_stage == 1: 215 | return [out1_2], \ 216 | [out1_1] 217 | 218 | out2 = torch.cat([out1_2, out1_1, out1], 1) 219 | out2_2 = self.m2_2(out2) 220 | out2_1 = self.m2_1(out2) 221 | 222 | if self.stop_at_stage == 2: 223 | return [out1_2, out2_2], \ 224 | [out1_1, out2_1] 225 | 226 | out3 = torch.cat([out2_2, out2_1, out1], 1) 227 | out3_2 = self.m3_2(out3) 228 | out3_1 = self.m3_1(out3) 229 | 230 | if self.stop_at_stage == 3: 231 | return [out1_2, out2_2, out3_2], \ 232 | [out1_1, out2_1, out3_1] 233 | 234 | out4 = torch.cat([out3_2, out3_1, out1], 1) 235 | out4_2 = self.m4_2(out4) 236 | out4_1 = self.m4_1(out4) 237 | 238 | if self.stop_at_stage == 4: 239 | return [out1_2, out2_2, out3_2, out4_2], \ 240 | [out1_1, out2_1, out3_1, out4_1] 241 | 242 | @staticmethod 243 | def create_stage(in_channels, out_channels, first=False): 244 | '''Create the neural network layers for a single stage.''' 245 | 246 | model = nn.Sequential() 247 | mid_channels = 128 248 | if first: 249 | padding = 3 250 | kernel = 7 251 | count = 6 252 | final_channels = 512 253 | else: 254 | padding = 3 255 | kernel = 7 256 | count = 6 257 | final_channels = mid_channels 258 | 259 | # First convolution 260 | model.add_module("0", 261 | nn.Conv2d( 262 | in_channels, 263 | mid_channels, 264 | kernel_size=kernel, 265 | stride=1, 266 | padding=padding) 267 | ) 268 | 269 | # Middle convolutions 270 | i = 1 271 | while i < count - 1: 272 | model.add_module(str(i), nn.ReLU(inplace=True)) 273 | i += 1 274 | model.add_module(str(i), 275 | nn.Conv2d( 276 | mid_channels, 277 | mid_channels, 278 | kernel_size=kernel, 279 | stride=1, 280 | padding=padding)) 281 | i += 1 282 | 283 | # Penultimate convolution 284 | model.add_module(str(i), nn.ReLU(inplace=True)) 285 | i += 1 286 | model.add_module(str(i), nn.Conv2d(mid_channels, final_channels, kernel_size=1, stride=1)) 287 | i += 1 288 | 289 | # Last convolution 290 | model.add_module(str(i), nn.ReLU(inplace=True)) 291 | i += 1 292 | model.add_module(str(i), nn.Conv2d(final_channels, out_channels, kernel_size=1, stride=1)) 293 | i += 1 294 | 295 | return model 296 | 297 | 298 | 299 | class ModelData(object): 300 | '''This class contains methods for loading the neural network''' 301 | 302 | def __init__(self, name="", net_path="", gpu_id=0): 303 | self.name = name 304 | self.net_path = net_path # Path to trained network model 305 | self.net = None # Trained network 306 | self.gpu_id = gpu_id 307 | 308 | def get_net(self): 309 | '''Returns network''' 310 | if not self.net: 311 | self.load_net_model() 312 | return self.net 313 | 314 | def load_net_model(self): 315 | '''Loads network model from disk''' 316 | if not self.net and path.exists(self.net_path): 317 | self.net = self.load_net_model_path(self.net_path) 318 | if not path.exists(self.net_path): 319 | print("ERROR: Unable to find model weights: '{}'".format( 320 | self.net_path)) 321 | exit(0) 322 | 323 | def load_net_model_path(self, path): 324 | '''Loads network model from disk with given path''' 325 | model_loading_start_time = time.time() 326 | print("Loading DOPE model '{}'...".format(path)) 327 | net = DopeNetwork() 328 | net = torch.nn.DataParallel(net, [0]).cuda() 329 | net.load_state_dict(torch.load(path)) 330 | net.eval() 331 | print(' Model loaded in {} seconds.'.format( 332 | time.time() - model_loading_start_time)) 333 | return net 334 | 335 | def __str__(self): 336 | '''Converts to string''' 337 | return "{}: {}".format(self.name, self.net_path) 338 | 339 | 340 | #================================ ObjectDetector ================================ 341 | class ObjectDetector(object): 342 | '''This class contains methods for object detection''' 343 | 344 | @staticmethod 345 | def detect_object_in_image(net_model, pnp_solver, in_img, config): 346 | '''Detect objects in a image using a specific trained network model''' 347 | 348 | if in_img is None: 349 | return [] 350 | 351 | # Run network inference 352 | image_tensor = transform(in_img) 353 | image_torch = Variable(image_tensor).cuda().unsqueeze(0) 354 | out, seg = net_model(image_torch) 355 | vertex2 = out[-1][0] 356 | aff = seg[-1][0] 357 | 358 | # Find objects from network output 359 | detected_objects = ObjectDetector.find_object_poses(vertex2, aff, pnp_solver, config) 360 | return detected_objects 361 | 362 | @staticmethod 363 | def find_object_poses(vertex2, aff, pnp_solver, config): 364 | '''Detect objects given network output''' 365 | 366 | # Detect objects from belief maps and affinities 367 | objects, all_peaks = ObjectDetector.find_objects(vertex2, aff, config) 368 | detected_objects = [] 369 | obj_name = pnp_solver.object_name 370 | 371 | for obj in objects: 372 | # Run PNP 373 | points = obj[1] + [(obj[0][0]*8, obj[0][1]*8)] 374 | cuboid2d = np.copy(points) 375 | location, quaternion, projected_points = pnp_solver.solve_pnp(points) 376 | 377 | # Save results 378 | detected_objects.append({ 379 | 'name': obj_name, 380 | 'location': location, 381 | 'quaternion': quaternion, 382 | 'cuboid2d': cuboid2d, 383 | 'projected_points': projected_points, 384 | }) 385 | 386 | return detected_objects 387 | 388 | @staticmethod 389 | def find_objects(vertex2, aff, config, numvertex=8): 390 | '''Detects objects given network belief maps and affinities, using heuristic method''' 391 | 392 | all_peaks = [] 393 | peak_counter = 0 394 | for j in range(vertex2.size()[0]): 395 | belief = vertex2[j].clone() 396 | map_ori = belief.cpu().data.numpy() 397 | 398 | map = gaussian_filter(belief.cpu().data.numpy(), sigma=config.sigma) 399 | p = 1 400 | map_left = np.zeros(map.shape) 401 | map_left[p:,:] = map[:-p,:] 402 | map_right = np.zeros(map.shape) 403 | map_right[:-p,:] = map[p:,:] 404 | map_up = np.zeros(map.shape) 405 | map_up[:,p:] = map[:,:-p] 406 | map_down = np.zeros(map.shape) 407 | map_down[:,:-p] = map[:,p:] 408 | 409 | peaks_binary = np.logical_and.reduce( 410 | ( 411 | map >= map_left, 412 | map >= map_right, 413 | map >= map_up, 414 | map >= map_down, 415 | map > config.thresh_map) 416 | ) 417 | peaks = zip(np.nonzero(peaks_binary)[1], np.nonzero(peaks_binary)[0]) 418 | 419 | # Computing the weigthed average for localizing the peaks 420 | peaks = list(peaks) 421 | win = 5 422 | ran = win // 2 423 | peaks_avg = [] 424 | for p_value in range(len(peaks)): 425 | p = peaks[p_value] 426 | weights = np.zeros((win,win)) 427 | i_values = np.zeros((win,win)) 428 | j_values = np.zeros((win,win)) 429 | for i in range(-ran,ran+1): 430 | for j in range(-ran,ran+1): 431 | if p[1]+i < 0 \ 432 | or p[1]+i >= map_ori.shape[0] \ 433 | or p[0]+j < 0 \ 434 | or p[0]+j >= map_ori.shape[1]: 435 | continue 436 | 437 | i_values[j+ran, i+ran] = p[1] + i 438 | j_values[j+ran, i+ran] = p[0] + j 439 | 440 | weights[j+ran, i+ran] = (map_ori[p[1]+i, p[0]+j]) 441 | 442 | # if the weights are all zeros 443 | # then add the none continuous points 444 | OFFSET_DUE_TO_UPSAMPLING = 0.4395 445 | try: 446 | peaks_avg.append( 447 | (np.average(j_values, weights=weights) + OFFSET_DUE_TO_UPSAMPLING, \ 448 | np.average(i_values, weights=weights) + OFFSET_DUE_TO_UPSAMPLING)) 449 | except: 450 | peaks_avg.append((p[0] + OFFSET_DUE_TO_UPSAMPLING, p[1] + OFFSET_DUE_TO_UPSAMPLING)) 451 | # Note: Python3 doesn't support len for zip object 452 | peaks_len = min(len(np.nonzero(peaks_binary)[1]), len(np.nonzero(peaks_binary)[0])) 453 | 454 | peaks_with_score = [peaks_avg[x_] + (map_ori[peaks[x_][1], peaks[x_][0]],) for x_ in range(len(peaks))] 455 | 456 | id = range(peak_counter, peak_counter + peaks_len) 457 | 458 | peaks_with_score_and_id = [peaks_with_score[i] + (id[i],) for i in range(len(id))] 459 | 460 | all_peaks.append(peaks_with_score_and_id) 461 | peak_counter += peaks_len 462 | 463 | objects = [] 464 | 465 | # Check object centroid and build the objects if the centroid is found 466 | for nb_object in range(len(all_peaks[-1])): 467 | if all_peaks[-1][nb_object][2] > config.thresh_points: 468 | objects.append([ 469 | [all_peaks[-1][nb_object][:2][0],all_peaks[-1][nb_object][:2][1]], 470 | [None for i in range(numvertex)], 471 | [None for i in range(numvertex)], 472 | all_peaks[-1][nb_object][2] 473 | ]) 474 | 475 | # Working with an output that only has belief maps 476 | if aff is None: 477 | if len (objects) > 0 and len(all_peaks)>0 and len(all_peaks[0])>0: 478 | for i_points in range(8): 479 | if len(all_peaks[i_points])>0 and all_peaks[i_points][0][2] > config.threshold: 480 | objects[0][1][i_points] = (all_peaks[i_points][0][0], all_peaks[i_points][0][1]) 481 | else: 482 | # For all points found 483 | for i_lists in range(len(all_peaks[:-1])): 484 | lists = all_peaks[i_lists] 485 | 486 | for candidate in lists: 487 | if candidate[2] < config.thresh_points: 488 | continue 489 | 490 | i_best = -1 491 | best_dist = 10000 492 | best_angle = 100 493 | for i_obj in range(len(objects)): 494 | center = [objects[i_obj][0][0], objects[i_obj][0][1]] 495 | 496 | # integer is used to look into the affinity map, 497 | # but the float version is used to run 498 | point_int = [int(candidate[0]), int(candidate[1])] 499 | point = [candidate[0], candidate[1]] 500 | 501 | # look at the distance to the vector field. 502 | v_aff = np.array([ 503 | aff[i_lists*2, 504 | point_int[1], 505 | point_int[0]].data.item(), 506 | aff[i_lists*2+1, 507 | point_int[1], 508 | point_int[0]].data.item()]) * 10 509 | 510 | # normalize the vector 511 | xvec = v_aff[0] 512 | yvec = v_aff[1] 513 | 514 | norms = np.sqrt(xvec * xvec + yvec * yvec) 515 | 516 | xvec/=norms 517 | yvec/=norms 518 | 519 | v_aff = np.concatenate([[xvec], [yvec]]) 520 | 521 | v_center = np.array(center) - np.array(point) 522 | xvec = v_center[0] 523 | yvec = v_center[1] 524 | 525 | norms = np.sqrt(xvec * xvec + yvec * yvec) 526 | 527 | xvec /= norms 528 | yvec /= norms 529 | 530 | v_center = np.concatenate([[xvec], [yvec]]) 531 | 532 | # vector affinity 533 | dist_angle = np.linalg.norm(v_center - v_aff) 534 | 535 | # distance between vertexes 536 | dist_point = np.linalg.norm(np.array(point) - np.array(center)) 537 | 538 | if dist_angle < config.thresh_angle \ 539 | and best_dist > 1000 \ 540 | or dist_angle < config.thresh_angle \ 541 | and best_dist > dist_point: 542 | i_best = i_obj 543 | best_angle = dist_angle 544 | best_dist = dist_point 545 | 546 | if i_best is -1: 547 | continue 548 | 549 | if objects[i_best][1][i_lists] is None \ 550 | or best_angle < config.thresh_angle \ 551 | and best_dist < objects[i_best][2][i_lists][1]: 552 | objects[i_best][1][i_lists] = ((candidate[0])*8, (candidate[1])*8) 553 | objects[i_best][2][i_lists] = (best_angle, best_dist) 554 | 555 | return objects, all_peaks 556 | -------------------------------------------------------------------------------- /weight/Readme.txt: -------------------------------------------------------------------------------- 1 | Download via provided web the weight and put here --------------------------------------------------------------------------------