points2lines完善

otr/table_line.py

@@ -368,16 +368,16 @@ def expansionAndShrinkage(pred,width=3):
 def getIOU(bbox0, bbox1):
     width = abs(max(bbox0[2],bbox1[2])-min(bbox0[0],bbox1[0]))-(abs(bbox0[2]-bbox0[0])+abs(bbox1[2]-bbox1[0]))
     height = abs(max(bbox0[3],bbox1[3])-min(bbox0[1],bbox1[1]))-(abs(bbox0[3]-bbox0[1])+abs(bbox1[3]-bbox1[1]))
-    if width < 0 and height < 0:
+    if width <= 0 and height <= 0:
         iou = abs(width*height/min(abs((bbox0[2]-bbox0[0])*(bbox0[3]-bbox0[1])),
                                    abs((bbox1[2]-bbox1[0])*(bbox1[3]-bbox1[1]))))
         # print("getIOU", iou)
-        return iou
+        return iou+0.1
     return 0
 
 def lines_cluster(list_lines,line_width):
     after_len = 0
-    log("len lines %d"%len(list_lines))
+    prelength = len(list_lines)
     append_width = line_width//2
     while 1:
         c_lines = []
@@ -404,6 +404,7 @@ def lines_cluster(list_lines,line_width):
         if first_len==after_len:
             break
         list_lines = c_lines
+    log("cluster lines from %d to %d"%(prelength,len(list_lines)))
     return c_lines
 
 
@@ -434,14 +435,21 @@ def points2lines(pred,sourceP_LB=True, prob=0.2, line_width=8, padding=3, min_le
         _start = None
         if _sum[h_index]<min_len:
             continue
+        last_back = 0
         while 1:
-            w_index += _step
             if w_index>=width:
+                if _start is not None:
+                    _end = w_index-1
+                    _bbox = [_start,h_i,_end,h_i]
+                    _dict = {"bbox":_bbox}
+                    h_lines.append(_dict)
+                    _start = None
                 break
             _h,_v = pred[h_i][w_index]
             if _h>prob:
                 if _start is None:
                     _start = w_index
+                w_index += _step
             else:
                 if _start is not None:
                     _end = w_index-1
@@ -449,7 +457,11 @@ def points2lines(pred,sourceP_LB=True, prob=0.2, line_width=8, padding=3, min_le
                     _dict = {"bbox":_bbox}
                     h_lines.append(_dict)
                     _start = None
+
                 w_index -= _step//2
+                if w_index<=last_back:
+                    w_index = last_back + _step//2
+                last_back = w_index
 
     log("starting points2lines 1")
     w_index = -1
@@ -464,21 +476,30 @@ def points2lines(pred,sourceP_LB=True, prob=0.2, line_width=8, padding=3, min_le
             continue
         h_index = -1
         _start = None
+        last_back = 0
+        list_test = []
+        list_lineprob = []
         while 1:
-            h_index += _step
             if h_index>=height:
+                if _start is not None:
+                    _end = last_h
+                    _bbox = [w_index,_start,w_index,_end]
+                    _dict = {"bbox":_bbox}
+                    v_lines.append(_dict)
+                    _start = None
+                    list_test.append(_dict)
                 break
-
-
             if sourceP_LB:
                 h_i = height-1-h_index
             else:
                 h_i = h_index
 
             _h,_v = pred[h_index][w_index]
+            list_lineprob.append((h_index,_v))
             if _v>prob:
                 if _start is None:
                     _start = h_i
+                h_index += _step
             else:
                 if _start is not None:
                     _end = last_h
@@ -486,7 +507,13 @@ def points2lines(pred,sourceP_LB=True, prob=0.2, line_width=8, padding=3, min_le
                     _dict = {"bbox":_bbox}
                     v_lines.append(_dict)
                     _start = None
+                    list_test.append(_dict)
+
                 h_index -= _step//2
+                if h_index<=last_back:
+                    h_index = last_back + _step//2
+                last_back = h_index
+
             last_h = h_i
     log("starting points2lines 2")