├── LICENSE ├── README.md ├── cube_detect_class.py ├── detection.py ├── input_image.jpg └── results.png /LICENSE: -------------------------------------------------------------------------------- 1 | Apache License 2 | Version 2.0, January 2004 3 | http://www.apache.org/licenses/ 4 | 5 | TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION 6 | 7 | 1. Definitions. 8 | 9 | "License" shall mean the terms and conditions for use, reproduction, 10 | and distribution as defined by Sections 1 through 9 of this document. 11 | 12 | "Licensor" shall mean the copyright owner or entity authorized by 13 | the copyright owner that is granting the License. 14 | 15 | "Legal Entity" shall mean the union of the acting entity and all 16 | other entities that control, are controlled by, or are under common 17 | control with that entity. For the purposes of this definition, 18 | "control" means (i) the power, direct or indirect, to cause the 19 | direction or management of such entity, whether by contract or 20 | otherwise, or (ii) ownership of fifty percent (50%) or more of the 21 | outstanding shares, or (iii) beneficial ownership of such entity. 22 | 23 | "You" (or "Your") shall mean an individual or Legal Entity 24 | exercising permissions granted by this License. 25 | 26 | "Source" form shall mean the preferred form for making modifications, 27 | including but not limited to software source code, documentation 28 | source, and configuration files. 29 | 30 | "Object" form shall mean any form resulting from mechanical 31 | transformation or translation of a Source form, including but 32 | not limited to compiled object code, generated documentation, 33 | and conversions to other media types. 34 | 35 | "Work" shall mean the work of authorship, whether in Source or 36 | Object form, made available under the License, as indicated by a 37 | copyright notice that is included in or attached to the work 38 | (an example is provided in the Appendix below). 39 | 40 | "Derivative Works" shall mean any work, whether in Source or Object 41 | form, that is based on (or derived from) the Work and for which the 42 | editorial revisions, annotations, elaborations, or other modifications 43 | represent, as a whole, an original work of authorship. For the purposes 44 | of this License, Derivative Works shall not include works that remain 45 | separable from, or merely link (or bind by name) to the interfaces of, 46 | the Work and Derivative Works thereof. 47 | 48 | "Contribution" shall mean any work of authorship, including 49 | the original version of the Work and any modifications or additions 50 | to that Work or Derivative Works thereof, that is intentionally 51 | submitted to Licensor for inclusion in the Work by the copyright owner 52 | or by an individual or Legal Entity authorized to submit on behalf of 53 | the copyright owner. For the purposes of this definition, "submitted" 54 | means any form of electronic, verbal, or written communication sent 55 | to the Licensor or its representatives, including but not limited to 56 | communication on electronic mailing lists, source code control systems, 57 | and issue tracking systems that are managed by, or on behalf of, the 58 | Licensor for the purpose of discussing and improving the Work, but 59 | excluding communication that is conspicuously marked or otherwise 60 | designated in writing by the copyright owner as "Not a Contribution." 61 | 62 | "Contributor" shall mean Licensor and any individual or Legal Entity 63 | on behalf of whom a Contribution has been received by Licensor and 64 | subsequently incorporated within the Work. 65 | 66 | 2. Grant of Copyright License. Subject to the terms and conditions of 67 | this License, each Contributor hereby grants to You a perpetual, 68 | worldwide, non-exclusive, no-charge, royalty-free, irrevocable 69 | copyright license to reproduce, prepare Derivative Works of, 70 | publicly display, publicly perform, sublicense, and distribute the 71 | Work and such Derivative Works in Source or Object form. 72 | 73 | 3. Grant of Patent License. Subject to the terms and conditions of 74 | this License, each Contributor hereby grants to You a perpetual, 75 | worldwide, non-exclusive, no-charge, royalty-free, irrevocable 76 | (except as stated in this section) patent license to make, have made, 77 | use, offer to sell, sell, import, and otherwise transfer the Work, 78 | where such license applies only to those patent claims licensable 79 | by such Contributor that are necessarily infringed by their 80 | Contribution(s) alone or by combination of their Contribution(s) 81 | with the Work to which such Contribution(s) was submitted. If You 82 | institute patent litigation against any entity (including a 83 | cross-claim or counterclaim in a lawsuit) alleging that the Work 84 | or a Contribution incorporated within the Work constitutes direct 85 | or contributory patent infringement, then any patent licenses 86 | granted to You under this License for that Work shall terminate 87 | as of the date such litigation is filed. 88 | 89 | 4. Redistribution. You may reproduce and distribute copies of the 90 | Work or Derivative Works thereof in any medium, with or without 91 | modifications, and in Source or Object form, provided that You 92 | meet the following conditions: 93 | 94 | (a) You must give any other recipients of the Work or 95 | Derivative Works a copy of this License; and 96 | 97 | (b) You must cause any modified files to carry prominent notices 98 | stating that You changed the files; and 99 | 100 | (c) You must retain, in the Source form of any Derivative Works 101 | that You distribute, all copyright, patent, trademark, and 102 | attribution notices from the Source form of the Work, 103 | excluding those notices that do not pertain to any part of 104 | the Derivative Works; and 105 | 106 | (d) If the Work includes a "NOTICE" text file as part of its 107 | distribution, then any Derivative Works that You distribute must 108 | include a readable copy of the attribution notices contained 109 | within such NOTICE file, excluding those notices that do not 110 | pertain to any part of the Derivative Works, in at least one 111 | of the following places: within a NOTICE text file distributed 112 | as part of the Derivative Works; within the Source form or 113 | documentation, if provided along with the Derivative Works; or, 114 | within a display generated by the Derivative Works, if and 115 | wherever such third-party notices normally appear. The contents 116 | of the NOTICE file are for informational purposes only and 117 | do not modify the License. You may add Your own attribution 118 | notices within Derivative Works that You distribute, alongside 119 | or as an addendum to the NOTICE text from the Work, provided 120 | that such additional attribution notices cannot be construed 121 | as modifying the License. 122 | 123 | You may add Your own copyright statement to Your modifications and 124 | may provide additional or different license terms and conditions 125 | for use, reproduction, or distribution of Your modifications, or 126 | for any such Derivative Works as a whole, provided Your use, 127 | reproduction, and distribution of the Work otherwise complies with 128 | the conditions stated in this License. 129 | 130 | 5. Submission of Contributions. Unless You explicitly state otherwise, 131 | any Contribution intentionally submitted for inclusion in the Work 132 | by You to the Licensor shall be under the terms and conditions of 133 | this License, without any additional terms or conditions. 134 | Notwithstanding the above, nothing herein shall supersede or modify 135 | the terms of any separate license agreement you may have executed 136 | with Licensor regarding such Contributions. 137 | 138 | 6. Trademarks. This License does not grant permission to use the trade 139 | names, trademarks, service marks, or product names of the Licensor, 140 | except as required for reasonable and customary use in describing the 141 | origin of the Work and reproducing the content of the NOTICE file. 142 | 143 | 7. Disclaimer of Warranty. Unless required by applicable law or 144 | agreed to in writing, Licensor provides the Work (and each 145 | Contributor provides its Contributions) on an "AS IS" BASIS, 146 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or 147 | implied, including, without limitation, any warranties or conditions 148 | of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A 149 | PARTICULAR PURPOSE. You are solely responsible for determining the 150 | appropriateness of using or redistributing the Work and assume any 151 | risks associated with Your exercise of permissions under this License. 152 | 153 | 8. Limitation of Liability. In no event and under no legal theory, 154 | whether in tort (including negligence), contract, or otherwise, 155 | unless required by applicable law (such as deliberate and grossly 156 | negligent acts) or agreed to in writing, shall any Contributor be 157 | liable to You for damages, including any direct, indirect, special, 158 | incidental, or consequential damages of any character arising as a 159 | result of this License or out of the use or inability to use the 160 | Work (including but not limited to damages for loss of goodwill, 161 | work stoppage, computer failure or malfunction, or any and all 162 | other commercial damages or losses), even if such Contributor 163 | has been advised of the possibility of such damages. 164 | 165 | 9. Accepting Warranty or Additional Liability. While redistributing 166 | the Work or Derivative Works thereof, You may choose to offer, 167 | and charge a fee for, acceptance of support, warranty, indemnity, 168 | or other liability obligations and/or rights consistent with this 169 | License. However, in accepting such obligations, You may act only 170 | on Your own behalf and on Your sole responsibility, not on behalf 171 | of any other Contributor, and only if You agree to indemnify, 172 | defend, and hold each Contributor harmless for any liability 173 | incurred by, or claims asserted against, such Contributor by reason 174 | of your accepting any such warranty or additional liability. 175 | 176 | END OF TERMS AND CONDITIONS 177 | 178 | APPENDIX: How to apply the Apache License to your work. 179 | 180 | To apply the Apache License to your work, attach the following 181 | boilerplate notice, with the fields enclosed by brackets "[]" 182 | replaced with your own identifying information. (Don't include 183 | the brackets!) The text should be enclosed in the appropriate 184 | comment syntax for the file format. We also recommend that a 185 | file or class name and description of purpose be included on the 186 | same "printed page" as the copyright notice for easier 187 | identification within third-party archives. 188 | 189 | Copyright [yyyy] [name of copyright owner] 190 | 191 | Licensed under the Apache License, Version 2.0 (the "License"); 192 | you may not use this file except in compliance with the License. 193 | You may obtain a copy of the License at 194 | 195 | http://www.apache.org/licenses/LICENSE-2.0 196 | 197 | Unless required by applicable law or agreed to in writing, software 198 | distributed under the License is distributed on an "AS IS" BASIS, 199 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 200 | See the License for the specific language governing permissions and 201 | limitations under the License. 202 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # OpenCV-real-world-red-cube-detection 2 | OpenCV-real-world-red-cube-detection-真实场景红色物块三维坐标检测 3 | 4 | 这个代码可是我花了两周，给别人提供解决方案的时候值1.5k的呢，现在只需要github一个star就能拥有！ 5 | 6 | b站的讲解视频连接： 7 | https://www.bilibili.com/video/BV1SK4y1Y7tZ 8 | 9 | 输入图 10 | ![输入](https://github.com/kaixindelele/OpenCV-real-world-red-cube-detection--/blob/main/input_image.jpg) 11 | 12 | 13 | 结果图 14 | ![结果](https://github.com/kaixindelele/OpenCV-real-world-red-cube-detection--/blob/main/results.png) 15 | 16 | 我的博客，欢迎关注： 17 | https://blog.csdn.net/hehedadaq/article/details/110097481 18 | -------------------------------------------------------------------------------- /cube_detect_class.py: -------------------------------------------------------------------------------- 1 | import os 2 | import math 3 | import numpy as np 4 | import cv2 5 | import time 6 | from detection import extract_red, extract_yellow, extract_gray, extract_orange, extract_black, extract_white 7 | import pygame 8 | from pygame.locals import QUIT 9 | 10 | 11 | class Cube: 12 | def __init__(self, 13 | input_image, 14 | color='red', 15 | show_flag=False, 16 | points_num=6, 17 | ): 18 | self.origin_img = input_image.copy() 19 | self.color = color 20 | self.points_num = points_num 21 | self.show_flag = show_flag 22 | self.surface_points_init 23 | canny_image = self.img_process_rgb2canny(input_image) 24 | self.canny_image = canny_image.copy() 25 | max_contour = self.img_process_canny2max_contour(canny_image) 26 | dark_with_poly = self.approx_polygone(max_contour) 27 | poly_6_points = self.get_poly_6_points(dark_with_poly, 28 | points_num=self.points_num, 29 | minDistance=50) 30 | self.surface_points_update(poly_6_points) 31 | 32 | if self.show_flag: 33 | print("poly_6_points:", poly_6_points) 34 | dark = self.origin_img.copy() 35 | for p in self.top_surface: 36 | cv2.circle(dark, (p[0], p[1]), 2, (200,0,0),-1) 37 | for p in self.buttom_surface: 38 | cv2.circle(dark, (p[0], p[1]), 2, (200,0,0),-1) 39 | cv2.imshow('poly_6_points', dark) 40 | cv2.waitKey(0) 41 | 42 | cut_rate = 0.1 43 | new_six_lines_list = self.divide_6_lines_and_contours(poly_6_points, cut_rate) 44 | cross_6_points = self.update_6_corners_from_new_6_line_list(new_six_lines_list) 45 | if self.show_flag: 46 | print("cross_6_points:", cross_6_points) 47 | # dark = self.origin_img.copy() 48 | for p in cross_6_points: 49 | cv2.circle(dark, (p[0], p[1]), 2, (200,200,0),-1) 50 | self.surface_points_update(cross_6_points) 51 | for p in self.top_surface: 52 | cv2.circle(dark, (p[0], p[1]), 2, (200,200,0),-1) 53 | cv2.imshow('cross_6_points', dark) 54 | cv2.waitKey(0) 55 | # self.surface_points_update(cross_6_points) 56 | 57 | 58 | def draw_two_points_line(self, input_img, point1, point2): 59 | # 起点和终点的坐标 60 | img = input_img.copy() 61 | ptStart = (point1[0], point1[1]) 62 | ptEnd = (point2[0], point2[1]) 63 | point_color = (0, 255, 0) # BGR 64 | thickness = 1 65 | lineType = 4 66 | cv2.line(img, ptStart, ptEnd, point_color, thickness, lineType) 67 | cv2.imshow('image', img) 68 | cv2.waitKey(0) 69 | 70 | def draw_one_point_line(self, input_img, point, line_ab): 71 | img = input_img.copy() 72 | k,b = line_ab[0], line_ab[1] 73 | x0, y0 = point[0], point[1] 74 | img = np.zeros((320, 320, 3), np.uint8) #生成一个空灰度图像 75 | # 起点和终点的坐标 76 | ptStart = (x0, y0) 77 | 78 | ptEnd = (x0+10, int(k*(x0+10)+b)) 79 | point_color = (0, 255, 0) # BGR 80 | thickness = 1 81 | lineType = 4 82 | cv2.line(img, ptStart, ptEnd, point_color, thickness, lineType) 83 | 84 | cv2.imshow('image', img) 85 | cv2.waitKey(0) 86 | 87 | def calculate_dist_point2line(self, point, line_ab): 88 | k, b = line_ab[0], line_ab[1] 89 | x0, y0 = point[0], point[1] 90 | dist = (k*x0-y0+b)/np.sqrt(k*k+1) 91 | return dist 92 | 93 | def cross_point(self, line1_ab, line2_ab):#计算交点函数 94 | k1, b1 = line1_ab[0], line1_ab[1] 95 | k2, b2 = line2_ab[0], line2_ab[1] 96 | x=(b2-b1)*1.0/(k1-k2) 97 | y=k1*x*1.0+b1*1.0 98 | return [int(x), int(y)] 99 | 100 | def calculate_from_two_points_to_line(self, point1, point2): 101 | x_list = [point1[0], point2[0]] 102 | y_list = [point1[1], point2[1]] 103 | k,b = self.Least_squares(x_list, y_list) 104 | return k, b 105 | 106 | def calculate_from_points_list_to_line(self, points_list): 107 | x_list = [p[0] for p in points_list] 108 | y_list = [p[1] for p in points_list] 109 | k,b = self.Least_squares(x_list, y_list) 110 | return k, b 111 | 112 | def Least_squares(self, x, y): 113 | x_ = np.array(x).mean() 114 | y_ = np.array(y).mean() 115 | m = np.zeros(1) 116 | n = np.zeros(1) 117 | k = np.zeros(1) 118 | p = np.zeros(1) 119 | for i in np.arange(len(x)): 120 | k = (x[i]-x_)* (y[i]-y_) 121 | m += k 122 | p = np.square( x[i]-x_ ) 123 | n = n + p 124 | a = m/n 125 | b = y_ - a* x_ 126 | return a[0],b[0] 127 | 128 | def img_process_rgb2canny(self, input_image): 129 | img = input_image 130 | if self.show_flag: 131 | cv2.imshow("img", img) 132 | cv2.waitKey(0) 133 | 134 | gray_img = cv2.cvtColor(input_image, cv2.COLOR_RGB2GRAY) 135 | gray_img = cv2.Canny(gray_img, 50, 150, apertureSize = 3) 136 | if self.show_flag: 137 | cv2.imshow("gray_img", gray_img) 138 | cv2.waitKey(0) 139 | 140 | origin_img = img.copy() 141 | red = extract_red(img) 142 | if self.show_flag: 143 | cv2.imshow("red", red) 144 | cv2.waitKey(0) 145 | # 去除一些噪点 146 | red = cv2.medianBlur(red, 5) 147 | if self.show_flag: 148 | cv2.imshow("medianBlur", red) 149 | cv2.waitKey(0) 150 | gf_time = time.time() 151 | red = cv2.Canny(red, 50, 150, apertureSize = 3) 152 | if self.show_flag: 153 | cv2.imshow("Canny", red) 154 | cv2.waitKey(0) 155 | 156 | return red 157 | 158 | def img_process_canny2max_contour(self, canny_image): 159 | contours, hierarchy = cv2.findContours(canny_image, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE) 160 | max_cont = max([len(c) for c in contours]) 161 | contours_corners = [c for c in contours if len(c)==max_cont][0] 162 | contours_corners_array = np.array(contours_corners) 163 | corners = contours_corners_array 164 | corners = np.int0(corners) 165 | max_contour = corners.reshape(corners.shape[0], 2) 166 | if self.show_flag: 167 | dark = self.origin_img.copy() 168 | for i in corners: 169 | x,y = i.ravel() 170 | cv2.circle(dark, (x, y), 1, (0, 200, 200), -1) 171 | cv2.imshow("max_contour", dark) 172 | cv2.waitKey(0) 173 | 174 | points_num=6 175 | level=0.1 176 | minDistance=50 177 | corners = cv2.goodFeaturesToTrack(canny_image, 178 | points_num, 179 | level, 180 | minDistance, 181 | ) 182 | corners = np.int0(corners) 183 | corners = corners.reshape(corners.shape[0], 2) 184 | six_corners = corners[:, :] 185 | dark = self.origin_img.copy() 186 | for i in corners: 187 | x,y = i.ravel() 188 | cv2.circle(dark, (x, y), 3, (0, 200, 200), -1) 189 | cv2.imshow("canny contours features", dark) 190 | cv2.waitKey(0) 191 | 192 | return max_contour 193 | 194 | def approx_polygone(self, max_contour, 195 | points_num=6, 196 | level=0.1, 197 | minDistance=50): 198 | # approx polygone from contours! 199 | polygonelist = [] 200 | corners = max_contour 201 | perimeter = cv2.arcLength(corners, True) 202 | epsilon = 0.005*cv2.arcLength(corners, True) 203 | approx = cv2.approxPolyDP(corners, epsilon, True) 204 | 205 | polygonelist.append(approx) 206 | dark = np.zeros(self.origin_img.shape) 207 | cv2.drawContours(dark, polygonelist, -1, (255, 255, 255), 1) 208 | if self.show_flag: 209 | cv2.imshow("approx_corners", dark) 210 | cv2.waitKey(0) 211 | dark_with_poly = dark.copy() 212 | return dark_with_poly 213 | 214 | def get_poly_6_points(self, dark_with_poly, 215 | points_num=6, 216 | minDistance=50): 217 | # poly features 218 | dark = dark_with_poly.astype('uint8') 219 | dark = cv2.cvtColor(dark, cv2.COLOR_RGB2GRAY) 220 | # get 6 points!!!! 221 | level=0.1 222 | corners = cv2.goodFeaturesToTrack(dark, 223 | points_num, 224 | level, 225 | minDistance, 226 | ) 227 | corners = np.int0(corners) 228 | corners = corners.reshape(corners.shape[0], 2) 229 | six_corners = corners[:, :] 230 | if self.show_flag: 231 | dark = self.origin_img.copy() 232 | for i in corners: 233 | x,y = i.ravel() 234 | cv2.circle(dark, (x, y), 3, (0, 200, 200), -1) 235 | cv2.imshow("poly contours features", dark) 236 | cv2.waitKey(0) 237 | return six_corners 238 | 239 | def surface_points_update(self, poly_6_points): 240 | corners = poly_6_points 241 | data = corners[corners[:,1].argsort()] 242 | self.top_upper_point = data[0,:] 243 | if data[1,0] < data[2,0]: 244 | self.top_left_point = data[1,:] 245 | self.top_right_point = data[2,:] 246 | else: 247 | self.top_left_point = data[2,:] 248 | self.top_right_point = data[1,:] 249 | self.top_down_point = self.top_left_point + self.top_right_point - self.top_upper_point 250 | # 算出上表面中心点，虽然有畸变，但是应该够用了 251 | self.top_center_point = ((self.top_left_point + self.top_right_point)/2).astype(np.int16) 252 | 253 | # buttom surface points update 254 | data = corners[corners[:,1].argsort()] 255 | self.buttom_down_point = data[-1,:] 256 | # 哪个小，哪个是左 257 | if data[-3,0] < data[-2,0]: 258 | self.buttom_left_point = data[-3,:] 259 | self.buttom_right_point = data[-2,:] 260 | else: 261 | self.buttom_left_point = data[-2,:] 262 | self.buttom_right_point = data[-3,:] 263 | self.buttom_upper_point = self.buttom_left_point + self.buttom_right_point - self.buttom_down_point 264 | # 算出上表面中心点，虽然有畸变，但是应该够用了 265 | self.buttom_center_point = ((self.buttom_left_point + self.buttom_right_point)/2).astype(np.int16) 266 | 267 | def divide_6_lines_and_contours(self, poly_6_points, cut_rate = 0.1): 268 | other_conners = np.where(self.canny_image>0) 269 | other_conners = np.array([other_conners[1], other_conners[0]]) 270 | other_conners = other_conners.transpose() 271 | # print("other_corners:", other_conners) 272 | corners = other_conners 273 | top_left_points_list = [] 274 | top_right_points_list = [] 275 | mid_left_points_list = [] 276 | mid_right_points_list = [] 277 | buttom_left_points_list = [] 278 | buttom_right_points_list = [] 279 | # y_sorted_data = corners[corners[:,1].argsort()] 280 | y_sorted_data = corners 281 | 282 | for p in y_sorted_data: 283 | if p[0] < self.top_upper_point[0] and p[1] < self.top_left_point[1]: 284 | top_left_points_list.append(p) 285 | elif p[0] > self.top_upper_point[0] and p[1] < self.top_right_point[1]: 286 | top_right_points_list.append(p) 287 | 288 | elif p[0] < self.top_upper_point[0] and p[1] > self.top_left_point[1] and p[1] < self.buttom_left_point[1]: 289 | mid_left_points_list.append(p) 290 | elif p[0] > self.top_upper_point[0] and p[1] > self.top_right_point[1] and p[1] < self.buttom_right_point[1]: 291 | mid_right_points_list.append(p) 292 | 293 | elif p[0] < self.buttom_down_point[0] and p[1] > self.buttom_left_point[1] and p[1] < self.buttom_down_point[1]: 294 | buttom_left_points_list.append(p) 295 | elif p[0] > self.buttom_down_point[0] and p[1] > self.buttom_right_point[1] and p[1] < self.buttom_down_point[1]: 296 | buttom_right_points_list.append(p) 297 | 298 | if self.show_flag: 299 | dark = self.origin_img.copy() 300 | for p in top_left_points_list: 301 | cv2.circle(dark, (p[0], p[1]), 1, (200,0,0),-1) 302 | cv2.imshow('left', dark) 303 | cv2.waitKey(0) 304 | dark = self.origin_img.copy() 305 | for p in top_right_points_list: 306 | cv2.circle(dark, (p[0], p[1]), 2, (200,100,0),-1) 307 | cv2.imshow('right_list', dark) 308 | cv2.waitKey(0) 309 | 310 | dark = self.origin_img.copy() 311 | for p in mid_left_points_list: 312 | cv2.circle(dark, (p[0], p[1]), 2, (200,100,0),-1) 313 | cv2.imshow('mid_left', dark) 314 | cv2.waitKey(0) 315 | dark = self.origin_img.copy() 316 | for p in mid_right_points_list: 317 | cv2.circle(dark, (p[0], p[1]), 2, (200,100,0),-1) 318 | cv2.imshow('mid_right', dark) 319 | cv2.waitKey(0) 320 | 321 | dark = self.origin_img.copy() 322 | for p in buttom_left_points_list: 323 | cv2.circle(dark, (p[0], p[1]), 2, (200,100,0),-1) 324 | cv2.imshow('buttom_left', dark) 325 | cv2.waitKey(0) 326 | dark = self.origin_img.copy() 327 | for p in buttom_right_points_list: 328 | cv2.circle(dark, (p[0], p[1]), 2, (200,100,0),-1) 329 | cv2.imshow('buttom_right', dark) 330 | cv2.waitKey(0) 331 | 332 | six_lines_list = [top_left_points_list, 333 | top_right_points_list, 334 | mid_left_points_list, 335 | mid_right_points_list, 336 | buttom_left_points_list, 337 | buttom_right_points_list,] 338 | tem_six_lines_list = [] 339 | for line_list in six_lines_list: 340 | points_num = len(line_list) 341 | 342 | tem_line_list = line_list[int(cut_rate*points_num):-int(cut_rate*points_num)] 343 | tem_six_lines_list.append(tem_line_list) 344 | if self.show_flag: 345 | dark = self.origin_img.copy() 346 | for p in tem_line_list: 347 | cv2.circle(dark, (p[0], p[1]), 1, (200,0,0),-1) 348 | cv2.imshow('cut_line', dark) 349 | cv2.waitKey(0) 350 | 351 | return tem_six_lines_list 352 | 353 | def update_6_corners_from_new_6_line_list(self, new_6_lines): 354 | six_corners = [] 355 | lines_kb_list = [] 356 | for points_list in new_6_lines: 357 | k, b = self.calculate_from_points_list_to_line(points_list) 358 | lines_kb_list.append([k, b]) 359 | six_corners.append(self.cross_point(lines_kb_list[0], lines_kb_list[1])) 360 | six_corners.append(self.cross_point(lines_kb_list[0], lines_kb_list[2])) 361 | six_corners.append(self.cross_point(lines_kb_list[1], lines_kb_list[3])) 362 | six_corners.append(self.cross_point(lines_kb_list[2], lines_kb_list[4])) 363 | six_corners.append(self.cross_point(lines_kb_list[3], lines_kb_list[5])) 364 | six_corners.append(self.cross_point(lines_kb_list[4], lines_kb_list[5])) 365 | six_corners = np.array(six_corners) 366 | return six_corners 367 | 368 | def surface_points_init(self,): 369 | self.top_upper_point = [] 370 | self.top_down_point = [] 371 | self.top_left_point = [] 372 | self.top_right_point = [] 373 | self.top_center_point = [] 374 | self.buttom_upper_point = [] 375 | self.buttom_down_point = [] 376 | self.buttom_left_point = [] 377 | self.buttom_right_point = [] 378 | self.buttom_center_point = [] 379 | @property 380 | def top_surface(self,): 381 | return [self.top_upper_point, 382 | self.top_down_point, 383 | self.top_left_point, 384 | self.top_right_point, 385 | self.top_center_point,] 386 | @property 387 | def buttom_surface(self,): 388 | return [self.buttom_upper_point, 389 | self.buttom_down_point, 390 | self.buttom_left_point, 391 | self.buttom_right_point, 392 | self.buttom_center_point,] 393 | 394 | 395 | 396 | def main(): 397 | img = cv2.imread("input_image.jpg") 398 | st = time.time() 399 | cube_class = Cube(img, show_flag=True) 400 | print("st:", time.time()-st) 401 | cv2.destroyAllWindows() 402 | 403 | if __name__=="__main__": 404 | main() -------------------------------------------------------------------------------- /detection.py: -------------------------------------------------------------------------------- 1 | # -*- encoding: utf-8 -*- 2 | ''' 3 | @File : detection.py 4 | @Time : 2019/11/10 01:01:20 5 | @Author : Li Xiang 6 | @Version : 1.0 7 | @Contact : lixiang6632@outlook.com 8 | @License : (C)Copyright 2019-2020, Song-Group-IIM-R&A 9 | @Describe: this module is used to detect the cube and calculate the coordinates of cube 10 | under robot's base coordinate system 11 | functions: 12 | extract_red, extract_yellow, extract_gray, extract_background，find_two_points， 13 | calculate_centroid，calculate_distance，calculate_coordinate_of_robot 14 | ''' 15 | 16 | 17 | from math import sqrt 18 | # from cv2 import cv2 as cv 19 | import cv2 20 | import numpy as np 21 | from sympy import symbols, solve 22 | 23 | 24 | def extract_orange(img): 25 | """ 26 | extract green object from the input image 27 | parameter: image 28 | return: green object of the image 29 | """ 30 | img_hsv = cv2.cvtColor(img, cv2.COLOR_BGR2HSV) 31 | # extract the first area 32 | lower_green=np.array([11,43,46]) 33 | upper_green=np.array([25,255,255]) 34 | 35 | mask0 = cv2.inRange(img_hsv, lower_green, upper_green) 36 | # splice two areas 37 | mask = mask0 38 | return mask 39 | 40 | 41 | def extract_red(img, show_flag=False): 42 | """ 43 | extract red object from the input image 44 | parameter: image 45 | return: red object of the image 46 | """ 47 | img_hsv = cv2.cvtColor(img, cv2.COLOR_BGR2HSV) 48 | lower_red = np.array([0, 43, 46]) 49 | upper_red = np.array([10, 255, 255]) 50 | mask0 = cv2.inRange(img_hsv, lower_red, upper_red) 51 | cv2.imshow('mask0', mask0) 52 | cv2.waitKey(0) 53 | # extract the second area 54 | # lower_red = np.array([156, 43, 46]) 55 | lower_red = np.array([156, 43, 46]) 56 | upper_red = np.array([180, 255, 255]) 57 | mask1 = cv2.inRange(img_hsv, lower_red, upper_red) 58 | cv2.imshow('mask1', mask1) 59 | cv2.waitKey(0) 60 | # splice two areas 61 | mask = mask0 + mask1 62 | cv2.imshow('mask', mask) 63 | cv2.waitKey(0) 64 | if show_flag: 65 | return img_hsv 66 | else: 67 | return mask 68 | 69 | def extract_white(img): 70 | """ 71 | extract green object from the input image 72 | parameter: image 73 | return: green object of the image 74 | """ 75 | img_hsv = cv2.cvtColor(img, cv2.COLOR_BGR2HSV) 76 | # extract the first area 77 | lower_green=np.array([0,0,0]) 78 | upper_green=np.array([180,50,255]) 79 | mask0 = cv2.inRange(img_hsv, lower_green, upper_green) 80 | # splice two areas 81 | mask = mask0 82 | cv2.imshow("black", mask0) 83 | cv2.waitKey(0) 84 | return mask 85 | 86 | def extract_black(img): 87 | """ 88 | extract green object from the input image 89 | parameter: image 90 | return: green object of the image 91 | """ 92 | img_hsv = cv2.cvtColor(img, cv2.COLOR_BGR2HSV) 93 | # extract the first area 94 | lower_green=np.array([0,0,0]) 95 | upper_green=np.array([180,255,46]) 96 | mask0 = cv2.inRange(img_hsv, lower_green, upper_green) 97 | # splice two areas 98 | mask = mask0 99 | cv2.imshow("black", mask0) 100 | cv2.waitKey(0) 101 | return mask 102 | 103 | def extract_green(img): 104 | """ 105 | extract green object from the input image 106 | parameter: image 107 | return: green object of the image 108 | """ 109 | img_hsv = cv2.cvtColor(img, cv2.COLOR_BGR2HSV) 110 | # extract the first area 111 | lower_green=np.array([35,43,46]) 112 | upper_green=np.array([77,255,255]) 113 | 114 | mask0 = cv2.inRange(img_hsv, lower_green, upper_green) 115 | # splice two areas 116 | mask = mask0 117 | return mask 118 | 119 | def extract_cyan(img): 120 | """ 121 | extract green object from the input image 122 | parameter: image 123 | return: green object of the image 124 | """ 125 | img_hsv = cv2.cvtColor(img, cv2.COLOR_BGR2HSV) 126 | 127 | # extract the first area 128 | lower_green=np.array([78,100,46]) 129 | upper_green=np.array([110,255,255]) 130 | 131 | mask0 = cv2.inRange(img_hsv, lower_green, upper_green) 132 | # splice two areas 133 | mask = mask0 134 | return mask 135 | 136 | 137 | def extract_yellow(img): 138 | """ 139 | extract yellow object from the input image 140 | parameter: image 141 | return: yellow object of the image 142 | """ 143 | img_hsv = cv2.cvtColor(img, cv2.COLOR_BGR2HSV) 144 | # extract the yellow area 145 | lower_yellow = np.array([26, 43, 46]) 146 | upper_yellow = np.array([34, 255, 255]) 147 | mask = cv2.inRange(img_hsv, lower_yellow, upper_yellow) 148 | return mask 149 | 150 | 151 | def extract_gray(img): 152 | """ 153 | extract gray object from the input image 154 | parameter: image 155 | return: gray object of the image 156 | """ 157 | img_hsv = cv2.cvtColor(img, cv2.COLOR_BGR2HSV) 158 | # extract the gray area 159 | lower_gray = np.array([0, 0, 46]) 160 | upper_gray = np.array([180, 43, 220]) 161 | mask = cv2.inRange(img_hsv, lower_gray, upper_gray) 162 | return mask 163 | 164 | def extract_blue(img): 165 | """ 166 | extract gray object from the input image 167 | parameter: image 168 | return: gray object of the image 169 | """ 170 | img_hsv = cv2.cvtColor(img, cv2.COLOR_BGR2HSV) 171 | # extract the gray area 172 | lower_gray = np.array([90, 43, 46]) 173 | upper_gray = np.array([110, 255, 255]) 174 | mask = cv2.inRange(img_hsv, lower_gray, upper_gray) 175 | return mask 176 | 177 | def extract_background(img, background_points): 178 | """ 179 | extract background from the input image 180 | parameters: 181 | image 182 | background_points: an array of points 183 | return: background of the image 184 | """ 185 | four_points = np.array(background_points, dtype=np.int32) 186 | mask = np.zeros((img.shape[0], img.shape[1])) 187 | 188 | cv2.fillConvexPoly(mask, four_points, 1) 189 | mask = mask.astype(np.bool) 190 | 191 | res = np.zeros_like(img) 192 | res[mask] = img[mask] 193 | return res 194 | 195 | 196 | def find_two_points(img): 197 | """ 198 | obtain a pair of vertices of the cube from the object extracted image 199 | parameter: object extracted image 200 | return: a numpy array of two points 201 | """ 202 | row_sum = list(img.sum(1)) 203 | row_useful = [x for x in row_sum if x > 0] 204 | index_row_first = row_sum.index(row_useful[0]) 205 | index_row_second = index_row_first + len(row_useful) - 1 206 | index_col_first = list(img[index_row_first]).index(255) + row_useful[0] // 255 // 2 207 | index_col_second = list(img[index_row_second]).index(255) + row_useful[-1] // 255 // 2 208 | return np.array([[index_row_first, index_col_first], [index_row_second, index_col_second]]) 209 | 210 | 211 | def calculate_centroid(point_a, point_b): 212 | """ 213 | calculate the midpoint of the two input points 214 | parameter: two points of type: array 215 | return: midpoint of type: array 216 | """ 217 | return [(point_a[i] + point_b[i]) / 2 for i in range(3)] 218 | 219 | 220 | def calculate_distance(point_a, point_b): 221 | """ 222 | calculate European distance of the two points in space 223 | parameter: two points of type: array 224 | return: European distance of tpye: float 225 | """ 226 | dis = sqrt((point_a[0] - point_b[0]) ** 2 + (point_a[1] - point_b[1]) ** 2 + \ 227 | (point_a[2] - point_b[2]) ** 2) 228 | return dis 229 | 230 | 231 | def calculate_coordinate_of_robot(point_a, radius_a, point_b, radius_b, delta_h, side_length): 232 | """ 233 | calculate the intersection of two spheres, then find the position according to 234 | the height(delta_h + side_length / 2) 235 | parameters: 236 | point_a: center point of the first sphere 237 | radius_a: radius of the first sphere 238 | point_b: center point of the second sphere 239 | radius_b: radius of the second sphere 240 | delta_h: the difference between the heights of object desk and the robot desk 241 | side_length: side length of the cube 242 | return: the coordinate in robot's base coordinate system 243 | """ 244 | # import math 245 | 246 | # if math.isnan(side_length): 247 | # print("side_length:", side_length) 248 | # side_length = 48 249 | 250 | coordinate_x, coordinate_y, coordinate_z = symbols('coordinate_x coordinate_y coordinate_z') 251 | res = solve([(coordinate_x - point_a[0]) ** 2 + (coordinate_y - point_a[1]) ** 2 + \ 252 | (coordinate_z - point_a[2]) ** 2 - radius_a ** 2, (coordinate_x - point_b[0]) ** 2 + \ 253 | (coordinate_y - point_b[1]) ** 2 + (coordinate_z - point_b[2]) ** 2 - radius_b ** 2, \ 254 | coordinate_z - delta_h - side_length / 2], [coordinate_x, coordinate_y, coordinate_z]) 255 | res = [coordinate_x for coordinate_x in res if coordinate_x[1] < -40] 256 | return list(res[0]) 257 | 258 | 259 | def test(): 260 | """ 261 | just for test 262 | """ 263 | point_dest_a = [400, -400, 0] 264 | point_dest_b = [-400, -400, 0] 265 | radius_a, radius_b = 800, 800 266 | delta_h, side_length = 0.5, 4 267 | res = calculate_coordinate_of_robot(point_dest_a, radius_a, \ 268 | point_dest_b, radius_b, delta_h, side_length) 269 | print(res) 270 | 271 | 272 | if __name__ == "__main__": 273 | test() 274 | -------------------------------------------------------------------------------- /input_image.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/kaixindelele/OpenCV-real-world-red-cube-detection/bf5cd42fd8a82759bc98e06c362c0669f316d868/input_image.jpg -------------------------------------------------------------------------------- /results.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/kaixindelele/OpenCV-real-world-red-cube-detection/bf5cd42fd8a82759bc98e06c362c0669f316d868/results.png --------------------------------------------------------------------------------