Complete rewrite line finding code

RCS.py

@@ -2,77 +2,129 @@
 import numpy as np
 import math
 
+
+height = -1
+width = -1
+
+def standardizeLines(v):
+    rho = v[0][0]
+    theta = v[0][1]
+    if rho < 0:
+        rho *= -1
+        theta -= math.pi
+    return [rho, theta, -1]
+
+def extreme_xy_values(line):
+    a = math.cos(float(line[1]))
+    b = math.sin(float(line[1]))
+    x0 = a*float(line[0])
+    y0 = b*float(line[0])
+    if abs(b) < 10**-10:
+        return 0, height, x0, x0
+    if abs(a) < 10**-10:
+        return y0, y0, 0, width
+    slope = -a/b
+
+    y_left = y0 - slope * x0
+    y_right = y_left + slope * width
+    y_left = min(max(y_left, 0), height)
+    y_right = min(max(y_right, 0), height)
+    # x_bottom = y_left/float(slope)
+    x_bottom = (y_left-y0)/slope + x0
+    x_top = (y_right-y0)/slope + x0
+    x_bottom = min(max(x_bottom, 0), width)
+    x_top = min(max(x_top, 0), width)
+    # print str(y_left) + "|" + str(y_right) + "|" + str(x_bottom) + "|" + str(x_top) + "|" + str(a) + "|" + str(b) + "|" + str(slope) + "|"
+    return int(y_left), int(y_right), int(x_bottom), int(x_top)
+
+def are_lines_similar(line1, line2):
+    extreme_x = (0,1,2,width,width-1,width-2)
+    extreme_y = (0,1,2,height,height-1,height-2)
+    y1l, y1r, x1b, x1t = extreme_xy_values(line1)
+    y2l, y2r, x2b, x2t = extreme_xy_values(line2)
+    s_thres = 30
+    if abs(y2r-y1r) < s_thres and abs(y2l-y1l) < s_thres and x1b in extreme_x and x2b in extreme_x and x1t in extreme_x and x2t in extreme_x:
+        return True
+    if abs(x2b-x1b) < s_thres and abs(x2t-x1t) < s_thres and y1l in extreme_y and y2l in extreme_y and y1r in extreme_y and y2r in extreme_y:
+        return True
+    if abs(y2r-y1r) < s_thres and abs(y2l-y1l) < s_thres and abs(x2b-x1b) < s_thres and abs(x2t-x1t) < s_thres:
+        return True
+    return False
+
 def drawLine(img, rho, theta, color):
-    a = np.cos(theta)
-    b = np.sin(theta)
+    a = math.cos(theta)
+    b = math.sin(theta)
     x0 = a*rho
     y0 = b*rho
     x1 = int(x0 + 1000*(-b))
     y1 = int(y0 + 1000*(a))
     x2 = int(x0 - 1000*(-b))
     y2 = int(y0 - 1000*(a))
+    yl, yr, xb, xt = extreme_xy_values((rho,theta))
+    # cv2.line(img,(0,int(yl)),(width,int(yl)),(0,0,0),2)
+    # cv2.line(img,(0,int(yr)),(width,int(yr)),(0,0,0),2)
+    # cv2.line(img,(int(xb),0),(int(xb),height),(255,255,255),2)
+    # cv2.line(img,(int(xt),0),(int(xt),height),(255,255,255),2)
     cv2.line(img,(x1,y1),(x2,y2),color,2)
 
 def main():
+    global height, width
     cap = cv2.VideoCapture(0)
     thres = 100
     isBlacklist = True
-    blacklist = {}
-    cannyBlacklist
+    blacklist = []
     count = 0
-    eth = 20
+    eth = 40
     while True:
         ret, frame = cap.read()
         img = cv2.resize(frame, None,fx=0.5, fy=0.5, interpolation = cv2.INTER_CUBIC)
-        blur = cv2.GaussianBlur(img,(15,15),0)
+        if height == -1:
+            height, width = img.shape[:2]
+        black = np.zeros((height,width,3), np.uint8)
+        blur = cv2.GaussianBlur(img,(5,5),0)
         gray = cv2.cvtColor(blur,cv2.COLOR_BGR2GRAY)
         edges = cv2.Canny(gray,eth, eth*3,apertureSize = 3)
         lines = cv2.HoughLines(edges,1,np.pi/45,thres)
 
-        if lines is not None:
+        if lines is not None and len(lines) <= 100:
+            lines = map(standardizeLines, lines)
+            lineGroups = {}
+
+            if not isBlacklist:
+                if len(lines) < 25:
+                    thres -= 2
+                if len(lines) > 50:
+                    thres += 2
+
+            print str(thres) + "|" + str(len(lines))
 
-            linesByAngle = {}
+            for i in range(0, len(lines)):
+                # drawLine(img, lines[i][0], lines[i][1], (0,0,255))
 
-            for x in range(0, len(lines)):
-                for rho, theta in lines[x]:
-                    if (rho,theta) in blacklist:
+                if isBlacklist:
+                    blacklist.append((lines[i][0], lines[i][1]))
+                    continue
+
+                if (lines[i][0], lines[i][1]) in blacklist:
                         continue
-
-                    if isBlacklist:
-                        blacklist[(rho,theta)] = blacklist.get((rho,theta), 0) + 1
-
-                    drawLine(img, rho, theta, (0,0,255))
-                    angle = round(math.degrees(theta), -1)
-                    if not angle in linesByAngle:
-                        linesByAngle[angle] = [(rho,theta)]
-                    else:
-                        linesByAngle[angle].append((rho, theta))
-
-            lineGroups = {}#{int, list}
-
-            for approxAngle, values in linesByAngle.iteritems():
-                # sort by similar rhos
-                for i in range(0, len(values)):
-                    found = False
-                    for j in range(0, i):
-                        if abs(values[i][0] - values[j][0]) < 25 and len(lineGroups[values[j][2]]) <= 10:
-                            lineGroups[values[j][2]].append(values[i])
-                            values[i] = (values[i][0], values[i][1], values[j][2])
-                            found = True
-                            break
-                    if not found:
-                        values[i] = (values[i][0], values[i][1], len(lineGroups))
-                        lineGroups[len(lineGroups)] = [(values[i][0], values[i][1], len(lineGroups))]
-
-            if len(lineGroups) < 25:
-                thres -= 2
-            if len(lineGroups) > 30:
-                thres += 2
-
-            # print str(thres) + "|" + str(len(lineGroups)) + "|" + str(len(lines))
+
+                foundGroup = False
+
+                for j in range(0, i):
+                    if (lines[j][0], lines[j][1]) in blacklist:
+                            continue
+                    if i != 0 and are_lines_similar(lines[i], lines[j]):
+                        # print "Lines: " + str(extreme_xy_values(lines[i])) + " and " + str(extreme_xy_values(lines[j])) + " are similar"
+                        lines[i][2] = lines[j][2]
+                        lineGroups[lines[j][2]].append(lines[i])
+                        foundGroup = True
+                        break
+
+                if not foundGroup:
+                    lines[i][2] = len(lineGroups)
+                    lineGroups[len(lineGroups)] = [lines[i]]
 
             for index, lines in lineGroups.iteritems():
-                # print green on avg of those rhos and thetas
                 rho = 0
                 theta = 0
                 for l in lines:
@@ -82,17 +134,19 @@ def main():
                 theta /= len(lines)
                 drawLine(img, rho, theta, (0,255,0))
 
+        elif lines is not None and len(lines) > 100:
+            thres += 5
         elif thres > 2:
             thres -= 2
         else:
             print "Failed to detect anything....!"
 
-        cv2.imshow('edges',edges)
+        cv2.imshow('guess',edges)
         cv2.imshow('frame',img)
         count += 1
-        if count >= 100 and isBlacklist:
+        if count >= 10 and isBlacklist:
             isBlacklist = False
-            eth = 30
+            eth += 10
             print "Stopped recording blacklist!!"
         if cv2.waitKey(10) & 0xFF == ord('q'):
             break