更改了桌面定位算法，仍有问题

.vscode/launch.json

@@ -9,7 +9,8 @@
             "type": "python",
             "request": "launch",
             "program": "${file}",
-            "console": "integratedTerminal"
+            "console": "integratedTerminal",
+            "args": ["13"]
         }
     ]
 }

Kmeans.py

@@ -0,0 +1,33 @@
+def Cluster(lines, var1, var2):
+    '''
+    数轴上两类聚类
+    输入：初始line，两个估计聚集点
+    输出：两个聚类后的数组
+    '''
+    while True:
+        line1 = []
+        line2 = []
+        for line in lines:
+            for rho, _ in line:
+                if abs(abs(rho) - var1) <= abs(abs(rho) - var2):
+                    line1.append(line)
+                else:
+                    line2.append(line)
+        nvar1 = getMeans(line1)
+        nvar2 = getMeans(line2)
+        # 若新聚集点误差大于原聚集点的百分之一则继续
+        if abs(nvar1 - var1) > var1 / 100 or abs(nvar2 - var2) > var2 / 100:
+            var1 = nvar1
+            var2 = nvar2
+        else:
+            return line1, line2
+
+def getMeans(lines):
+    '''
+    计算平均值，line为数组
+    '''
+    sum = 0
+    for line in lines:
+        for rho, _ in line:
+            sum += abs(rho)
+    return (sum / len(lines)) if len(lines) else 0

QT.py

@@ -14,6 +14,8 @@
 from yolo6D.Predict import predict, predict_thread, draw_predict
 from camera import Camera
 from corner import square_desk
+import numpy as np
+import math
 
 RECORD_LENGTH = 18
 
@@ -23,6 +25,48 @@ def make_directories(folder):
     if not os.path.exists(folder+"depth/"):
         os.makedirs(folder+"depth/")
 
+def ReadData():
+    fp = open('data1.txt', 'r')
+    lines = fp.readlines()
+    options = dict()
+    x = np.zeros((6, 1))
+    for line in lines:
+        line = line.strip()
+        if line == '':
+            continue
+        key,value = line.split('=')
+        key = key.strip()
+        value = value.strip()
+        options[key] = value
+    x[0] = options['red_low']
+    x[1] = options['red_high']
+    x[2] = options['green_low']
+    x[3] = options['green_high']
+    x[4] = options['blue_low']
+    x[5] = options['blue_high']
+    fp.close()
+
+    ca = np.zeros((3, 1))
+    ho = np.zeros((3, 1))
+    fp = open('data2.txt', 'r')
+    lines = fp.readlines()
+    for line in lines:
+        line = line.strip()
+        if line == '':
+            continue
+        key,value = line.split('=')
+        key = key.strip()
+        value = value.strip()
+        options[key] = value
+    ca[0] = options['canny_threshold_1']
+    ca[1] = options['canny_threshold_2']
+    ca[2] = options['sobel']
+    ho[0] = options['hough_rho']
+    ho[1] = int(options['hough_theta']) * math.pi/180
+    ho[2] = options['hough_threshold']
+    fp.close()
+    return x, ca, ho
+
 class Ui_MainWindow(object):
     count = -1
     objectnum = 2
@@ -231,10 +275,11 @@ def setupUi(self, MainWindow):
         self.type.clicked.connect(lambda:self.changetype(round))
         self.type.setGeometry(QtCore.QRect(1100, 550, 100, 100))
         self.timer_camera = QtCore.QTimer()
-        print(" 0% 开始初始化相机")
+        self.x, self.ca, self.ho = ReadData()
+        print(" 2% 开始初始化相机")
         camera = Camera()
         camera.init()
-        print(" 5% \033[0;32m相机初始化完成\033[0m")
+        print(" 6% \033[0;32m相机初始化完成\033[0m")
         self.thread_init()
         print("10% \033[0;32m多线程初始化完成\033[0m")
         self.START.clicked.connect(lambda:self.open_camera(camera))
@@ -339,7 +384,13 @@ def capture_camera(self, camera):
             writer.write(f, zgray2list)
 
         print("    \033[0;32m%s.jpg已拍摄\033[0m" % self.count)
-
+        print("    \033[0;34m定位图片%s.jpg...\033[0m" % self.count)
+        lined = square_desk(self.count, self.x, self.ca, self.ho)
+        print("    \033[0;32mmarked%s.jpg已保存\033[0m" % self.count)
+        # marked_img = 'marked' + str(self.count) + '.jpg'
+        # marked_img = 'corner.jpg'
+        # marked = cv2.imread(marked_img,1)
+        # self.show(marked)
         # 预测
         print("    \033[0;34m预测图片%s.jpg...\033[0m" % self.count)
         threads = []
@@ -362,18 +413,12 @@ def capture_camera(self, camera):
                 if num_done is self.objectnum:
                     break
 
-        draw_predict(bs, ret, c, self.count)
+        draw_predict(bs, ret, lined, self.count)
         print("        \033[0;34m用时%s秒\033[0m" % (time.time() - starttime))
         print("    \033[0;32m%s.jpg已保存\033[0m" % self.count)
-        # predicted = cv2.imread('JPEGImages/%s.jpg' % self.count, 1)
-        # self.show(predicted)
-        print("    \033[0;34m定位图片%s.jpg...\033[0m" % self.count)
-        square_desk(self.count)
-        print("    \033[0;32mmarked%s.jpg已保存\033[0m" % self.count)
-        # marked_img = 'marked' + str(self.count) + '.jpg'
-        marked_img = 'corner.jpg'
-        marked = cv2.imread(marked_img,1)
-        self.show(marked)
+        predicted = cv2.imread('JPEGImages/predict%s.jpg' % self.count, 1)
+        self.show(predicted)
+
 
     def close_camera(self, camera):
         if self.timer_camera.isActive():

corner.py

@@ -5,6 +5,8 @@
 import sys
 from yolo6D.utils import get_camera_intrinsic
 from yolo6D.Predict import draw
+from Kmeans import Cluster
+import copy
 
 # def draw(img, corner, imgpts):
 #     img = cv2.line(img, corner, tuple(imgpts[0].ravel()), (255,0,0), 3)
@@ -55,32 +57,55 @@ def getcrosspoint(rho1,theta1,rho2,theta2):
     print(corss_x,cross_y)
     return (corss_x,cross_y)
 
-def square_desk(num, canny, hough):
-    '''
-    输入：图片编号
-    输出：坐标系
-    导出：带坐标系图片
-    '''
-    img_path = 'JPEGImages/' + str(num) + '.jpg'
-    img = cv2.imread(img_path,1)
+# 二值化
+def thres(img, x):
+    red_low = x[0]
+    red_high = x[1]
+    green_low = x[2]
+    green_high = x[3]
+    blue_low = x[4]
+    blue_high = x[5]
+    height, width = img.shape[:2]
+    im_new = np.zeros((height, width, 1), np.uint8)
+    for i in range(height):
+        for j in range(width):
+            judger = 1
+            temp = img[i, j]
+            if temp[0] < red_low or temp[0] > red_high:
+                judger = 0
+            if temp[1] < green_low or temp[1] > green_high:
+                judger = 0
+            if temp[2] < blue_low or temp[2] > blue_high:
+                judger = 0
+            if judger:
+                im_new[i, j] = 255
+    return im_new
+
+def square_canny(img, canny):
     img = cv2.resize(img,(640,480), interpolation = cv2.INTER_CUBIC)
-
-    internal_calibration = get_camera_intrinsic()
-    internal_calibration = np.array(internal_calibration, dtype='float32')
-    distCoeffs = np.zeros((8, 1), dtype='float32')
-
-    gray = cv2.cvtColor(img,cv2.COLOR_BGR2GRAY)
+    # gray = cv2.cvtColor(img,cv2.COLOR_BGR2GRAY)
     # threshold
-    lower_gray = np.array([140])
-    upper_gray = np.array([200])
-    test_gray = cv2.inRange(gray,lower_gray,upper_gray)
-    kernel = np.ones((3,3),np.uint8)
-    test_gray = cv2.morphologyEx(test_gray, cv2.MORPH_CLOSE, kernel)
+    # lower_gray = np.array([140])
+    # upper_gray = np.array([200])
+    # test_gray = cv2.inRange(gray,lower_gray,upper_gray)
+    kernel = np.ones((9,9),np.uint8)
+    test_gray = cv2.morphologyEx(img, cv2.MORPH_CLOSE, kernel)
     test_gray = cv2.morphologyEx(test_gray, cv2.MORPH_OPEN, kernel)
-    kernel2 = np.ones((5,5),np.uint8)
+    kernel2 = np.ones((9,9),np.uint8)
     test_gray = cv2.morphologyEx(test_gray, cv2.MORPH_CLOSE, kernel2)
     edges = cv2.Canny(test_gray, canny[0], canny[1], canny[2])
+    return edges
+
+def square_line(origin, edges, hough):
+    internal_calibration = get_camera_intrinsic()
+    internal_calibration = np.array(internal_calibration, dtype='float32')    
+    distCoeffs = np.zeros((8, 1), dtype='float32')
+    img = copy.copy(origin)
     lines = cv2.HoughLines(edges, hough[0], hough[1], hough[2])
+
+    # rho：ρ，图片左上角向直线所做垂线的长度
+    # theta：Θ，图片左上角向直线所做垂线与顶边夹角
+    # 垂足高于原点时，ρ为负，Θ为垂线关于原点的对称线与顶边的夹角
     top_line_theta = []
     top_line_rho = []
     left_line_theta = []
@@ -89,29 +114,32 @@ def square_desk(num, canny, hough):
     right_line_rho = []
     bottom_line_theta = []
     bottom_line_rho = []
+    horizon = []
     summ = 0
     final_lines = np.zeros((4,2))
-    if lines is None:
+    if len(lines) < 4:
         print("    \033[0;31m未检测到方桌！\033[0m")
+        return edges
     else:
         for line in lines:
             for rho,theta in line:
-
-                if theta > (math.pi*(2/3)):
-                    if abs(rho) <320:
-                        right_line_theta.append(theta)
-                        right_line_rho.append(rho)
-                elif theta > (math.pi/4):
-                    if abs(rho) <320:
-                        top_line_theta.append(theta)
-                        top_line_rho.append(rho)
-                    else:
-                        bottom_line_theta.append(theta)
-                        bottom_line_rho.append(rho)
-                else:
+                if (theta > math.pi / 3 and theta < math.pi * 2 / 3): # 横线
+                    horizon.append(line)
+                elif rho < 0: # 右边
+                    right_line_rho.append(rho)
+                    right_line_theta.append(theta)
+                else: # 左边
                     left_line_theta.append(theta)
                     left_line_rho.append(rho)
-
+        top, bottom = Cluster(horizon, 120, 360) # 将横线依据abs(rho)分为上下
+        for line in top:
+            for rho, theta in line:
+                top_line_rho.append(rho)
+                top_line_theta.append(theta)
+        for line in bottom:
+            for rho, theta in line:
+                bottom_line_rho.append(rho)
+                bottom_line_theta.append(theta)
 
         for i in right_line_theta:
             summ +=i
@@ -197,6 +225,27 @@ def square_desk(num, canny, hough):
         _,rvector,tvector=cv2.solvePnP(Table_3D,Table_2D,internal_calibration,distCoeffs)
         axis = np.float32([[55,0,0], [0,55,0], [0,0,-20]]).reshape(-1,3)
         imgpts, _ = cv2.projectPoints(axis, rvector, tvector,internal_calibration,distCoeffs,)
-        img = draw(img,(left_top_point_x, left_top_point_y),imgpts)
+        lined = draw(img,(left_top_point_x, left_top_point_y),imgpts)
+    return lined
+
+def square_desk(num, x, canny, hough):
+    '''
+    输入：图片编号
+    输出：坐标系
+    导出：带坐标系图片
+    '''
+    img_path = 'JPEGImages/' + str(num) + '.jpg'
+    img = cv2.imread(img_path,1)
+    internal_calibration = get_camera_intrinsic()
+    internal_calibration = np.array(internal_calibration, dtype='float32')
+    distCoeffs = np.zeros((8, 1), dtype='float32')
+
+    thresed = thres(img, x)
 
-    cv2.imwrite('corner.jpg',img)
+    edges = square_canny(thresed, canny)
+    lined = square_line(img, edges, hough)
+    # affine_table_2D = np.float32([[0,0],[0,550],[550,0],[550,550]])
+    # M = cv2.getPerspectiveTransform(Table_2D,affine_table_2D)
+    # marked = cv2.warpPerspective(lined,M,(550,550)) # Perspective_Transformation
+    cv2.imwrite('JPEGImages/marked' + str(num) + '.jpg',lined)
+    return lined

data1.txt

@@ -0,0 +1,6 @@
+red_low = 150
+red_high = 195
+green_low = 160
+green_high = 205
+blue_low = 160
+blue_high = 205

data2.txt

@@ -0,0 +1,6 @@
+canny_threshold_1 = 100
+canny_threshold_2 = 150
+sobel = 3
+hough_rho = 1
+hough_theta = 1
+hough_threshold = 20