1+ import cv2
2+ import numpy as np
3+ import math
4+ from inference import show_single
5+ from shapely .geometry import LineString
6+
7+
8+ def lsd ():
9+ src = cv2 .imread ("/home/wangbowen/DATA/Facade/translated_data/images/32052284_477d66a5ae_o.png" )
10+ gray = cv2 .cvtColor (src , cv2 .COLOR_BGR2GRAY )
11+ src = cv2 .cvtColor (src , cv2 .COLOR_BGR2RGB )
12+ gray = cv2 .GaussianBlur (gray , (5 , 5 ), 5 )
13+ gray = cv2 .GaussianBlur (gray , (3 , 3 ), 5 )
14+
15+ LSD = cv2 .createLineSegmentDetector (0 )
16+ dlines = LSD .detect (gray )
17+
18+ line_record = []
19+
20+ for dline in dlines [0 ]:
21+ x0 = int (round (dline [0 ][0 ]))
22+ y0 = int (round (dline [0 ][1 ]))
23+ x1 = int (round (dline [0 ][2 ]))
24+ y1 = int (round (dline [0 ][3 ]))
25+ line_record .append ([x0 , y0 , x1 , y1 ])
26+
27+ return line_record , src
28+
29+
30+ def calc_abc_from_line_2d (x0 , y0 , x1 , y1 ):
31+ a = y0 - y1
32+ b = x1 - x0
33+ c = x0 * y1 - x1 * y0
34+ return a , b , c
35+
36+
37+ def get_line_cross_point (line1 , line2 ):
38+ a0 , b0 , c0 = calc_abc_from_line_2d (* line1 )
39+ a1 , b1 , c1 = calc_abc_from_line_2d (* line2 )
40+ D = a0 * b1 - a1 * b0
41+ if D == 0 :
42+ return None
43+ x = (b0 * c1 - b1 * c0 )/ D
44+ y = (a1 * c0 - a0 * c1 )/ D
45+ return x , y
46+
47+
48+ def combine (lines ):
49+ def get_line (j ):
50+ if index [j ] == 1 :
51+ current_line = [lines [j ][2 ], lines [j ][3 ], lines [j ][4 ], lines [j ][5 ]]
52+ else :
53+ if j == 0 or j == 2 :
54+ current_line = [5 , lines [j ][0 ], 20 , lines [j ][0 ]]
55+ else :
56+ current_line = [lines [j ][0 ], 20 , lines [j ][0 ], 40 ]
57+ return current_line
58+
59+ index = []
60+ for i in range (len (lines )):
61+ if len (lines [i ]) == 1 :
62+ index .append (0 )
63+ else :
64+ index .append (1 )
65+
66+ if np .array (index ).sum () < 2 :
67+ return None
68+
69+ cross_record = []
70+ start_line = None
71+ for s in range (len (index )):
72+ p_current_line = get_line (s )
73+ if s == 0 :
74+ start_line = p_current_line
75+
76+ if s == 3 :
77+ p_next_line = start_line
78+ else :
79+ p_next_line = get_line (s + 1 )
80+ cross_point = get_line_cross_point (p_current_line , p_next_line )
81+ cross_record .append (cross_point )
82+ return cross_record
83+
84+
85+ def line_search (enhance_stats , lines ):
86+ x , y , w , h = enhance_stats
87+ top = LineString ([(x , y ), (x + w , y )])
88+ bottom = LineString ([(x , y + h ), (x + w , y + h )])
89+ left = LineString ([(x , y ), (x , y + h )])
90+ right = LineString ([(x + w , y ), (x + w , y + h )])
91+
92+ line_list = {"top" : top , "left" : left , "bottom" : bottom , "right" : right }
93+ distance_thresh = 10
94+ degree_thresh = 0.1
95+ max_selection = 1
96+ record = {"top" : [], "left" : [], "bottom" : [], "right" : []}
97+
98+ for key , value in line_list .items ():
99+ # print(key)
100+ for (x0 , y0 , x1 , y1 ) in lines :
101+ current_line = LineString ([(x0 , y0 ), (x1 , y1 )])
102+ current_degree = math .atan2 (y0 - y1 , x0 - x1 )
103+ current_dis = value .distance (current_line )
104+ line_len = (x0 - x1 )** 2 + (y0 - y1 )** 2
105+
106+ if current_dis > distance_thresh :
107+ continue
108+
109+ # print([current_dis, current_degree, x0, y0, x1, y1])
110+
111+ if key == "top" or key == "bottom" :
112+ if math .pi * 1 / 5 < abs (current_degree ) < math .pi * 4 / 5 :
113+ continue
114+ if line_len > w ** 2 * 1.5 or line_len < w ** 2 / 3 :
115+ continue
116+ else :
117+ if abs (current_degree ) < math .pi * 1 / 3 or abs (current_degree ) > math .pi * 2 / 3 :
118+ continue
119+ if line_len > h ** 2 * 1.5 or line_len < h ** 2 / 3 :
120+ continue
121+
122+ status = True
123+ for i in range (len (record [key ])):
124+ if abs (abs (abs (current_degree ) - math .pi / 2 ) - abs (abs (record [key ][i ][1 ]) - math .pi / 2 )) < degree_thresh :
125+ if record [key ][i ][0 ] > current_dis :
126+ record [key ][i ] = [current_dis , current_degree , x0 , y0 , x1 , y1 ]
127+ status = False
128+
129+ if status :
130+ record [key ].append ([current_dis , current_degree , x0 , y0 , x1 , y1 ])
131+
132+ final_line = []
133+ for key2 , value2 in record .items ():
134+ value2 .sort (key = lambda s : s [0 ], reverse = False )
135+ num = min (len (value2 ), max_selection )
136+ if num == 0 :
137+ if key2 == "top" :
138+ final_line .append ([y ])
139+ elif key2 == "bottom" :
140+ final_line .append ([y + h ])
141+ elif key2 == "left" :
142+ final_line .append ([x ])
143+ else :
144+ final_line .append ([x + w ])
145+ continue
146+
147+ for j in range (num ):
148+ final_line .append (value2 [j ])
149+
150+ return final_line
151+
152+
153+ def revision ():
154+ img = cv2 .imread ("../demo/test.png" )
155+ img = cv2 .resize (img , (2048 , 1152 ))
156+ img_orl = img
157+
158+ gray = cv2 .cvtColor (img , cv2 .COLOR_BGR2GRAY )
159+ ret , binary = cv2 .threshold (gray , 0 , 255 , cv2 .THRESH_BINARY | cv2 .THRESH_OTSU )
160+ kernel1 = cv2 .getStructuringElement (cv2 .MORPH_RECT , (5 , 5 ))
161+ bin_clo = cv2 .erode (binary , kernel1 , iterations = 1 )
162+ num_labels , labels , stats , centroids = cv2 .connectedComponentsWithStats (bin_clo , connectivity = 8 )
163+
164+ # print('num_labels = ', num_labels)
165+ # print('labels = ', labels)
166+ # # 不同的连通域赋予不同的颜色
167+ lines , scr = lsd ()
168+
169+ for i in range (1 , num_labels ):
170+ # if i < 10:
171+ # continue
172+ mask = labels == i
173+ # 连通域的信息：对应各个轮廓的x、y、width、height和面积
174+ if stats [i ][4 ] < 100 :
175+ continue
176+ # # 连通域的信息：对应各个轮廓的x、y、width、height和面积
177+ # print('stats = ', stats[i])
178+ # # 连通域的中心点
179+ # print('centroids = ', centroids[i])
180+
181+ current_patch = np .zeros ((img .shape [0 ], img .shape [1 ]), np .uint8 )
182+ current_patch [mask ] = 255
183+ # show_single(current_patch)
184+
185+ # kernel2 = cv2.getStructuringElement(cv2.MORPH_RECT, (7, 7))
186+ # current_patch = cv2.dilate(current_patch, kernel2, iterations=1)
187+ # show_single(current_patch)
188+
189+ x , y , w , h = cv2 .boundingRect (current_patch )
190+ # print(x, y, w, h)
191+ # cv2.rectangle(current_patch, (x, y), (x + w, y + h), (225, 0, 255), 2)
192+ # show_single(current_patch)
193+
194+ detect_lines = line_search ([x , y , w , h ], lines )
195+ final_point = combine (detect_lines )
196+
197+ if final_point is None :
198+ continue
199+
200+ start_x , start_y = None , None
201+ for w in range (len (final_point )):
202+ x0 , y0 = round (final_point [w ][0 ]), round (final_point [w ][1 ])
203+ if w == 0 :
204+ start_x , start_y = x0 , y0
205+ if w == 3 :
206+ x1 , y1 = start_x , start_y
207+ else :
208+ x1 , y1 = round (final_point [w + 1 ][0 ]), round (final_point [w + 1 ][1 ])
209+ cv2 .line (scr , (x0 , y0 ), (x1 , y1 ), 255 , 2 , cv2 .LINE_AA )
210+
211+ show_single (scr , save = True , name = "lines_detect.png" )
212+
213+
214+ if __name__ == '__main__' :
215+ revision ()
0 commit comments