提高点到线的识别效果和效率

format_convert/utils.py

@@ -1232,7 +1232,7 @@ class LineTable:
                 all_match_box_list.append(match_box_list)
 
             # print("match_box_list", all_match_box_list)
-            all_match_box_list.sort(key=lambda x:(x[0][2]+x[0][4])/2,reverse=sourceP_LB)
+            all_match_box_list.sort(key=lambda x:(round(x[0][2]+x[0][4])/2,0),reverse=sourceP_LB)
             for box_list in all_match_box_list:
                 for box in box_list:
                     _cell["text"] += re.sub("\s",'',box[0])

otr/table_line.py

@@ -308,8 +308,230 @@ def draw_pixel(pred, prob=0.2, is_test=1):
         plt.show()
         return
 
+def expansionAndShrinkage(pred,width=3):
+
+
+    pred_array = np.array(pred)
+    print("pred_array=====",pred_array.shape)
+    _array = pred_array[...,0]
+
+    _l = [_array]
+    for _i in range(width):
+        tmp_array = np.pad(_array[:-(_i+1),...],((_i+1,0),(0,0)))
+        _l.append(tmp_array)
+    for _i in range(width):
+        tmp_array = np.pad(_array[_i+1:,...],((0,_i+1),(0,0)))
+        _l.append(tmp_array)
+
+    for _i in range(width):
+        tmp_array = np.pad(_array[...,:-(_i+1)],((0,0),(_i+1,0)))
+        _l.append(tmp_array)
+    for _i in range(width):
+        tmp_array = np.pad(_array[...,_i+1:],((0,0),(0,_i+1)))
+        _l.append(tmp_array)
+
+
+    for _a in _l:
+        print(_a.shape)
+
+    h_array = np.stack(_l,axis=0)
+    h_array = np.max(h_array,axis=0,keepdims=False)
+
+    _array = pred_array[...,1]
+
+    _l = [_array]
+    for _i in range(width):
+        tmp_array = np.pad(_array[:-(_i+1),...],((_i+1,0),(0,0)))
+        _l.append(tmp_array)
+    for _i in range(width):
+        tmp_array = np.pad(_array[_i+1:,...],((0,_i+1),(0,0)))
+        _l.append(tmp_array)
+
+    for _i in range(width):
+        tmp_array = np.pad(_array[...,:-(_i+1)],((0,0),(_i+1,0)))
+        _l.append(tmp_array)
+    for _i in range(width):
+        tmp_array = np.pad(_array[...,_i+1:],((0,0),(0,_i+1)))
+        _l.append(tmp_array)
+
+    v_array = np.stack(_l,axis=0)
+    print("v_array=====",v_array.shape)
+    v_array = np.max(v_array,axis=0,keepdims=False)
+
+    print("h_array=====",h_array.shape)
+    print("v_array=====",v_array.shape)
+    last_array = np.stack([h_array,v_array],axis=-1)
+    print("pred_array=====",last_array.shape)
+    return last_array
+
+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:
+        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 0
+
+def lines_cluster(list_lines,line_width):
+    after_len = 0
+    log("len lines %d"%len(list_lines))
+    append_width = line_width//2
+    while 1:
+        c_lines = []
+        first_len = after_len
+
+        for _line in list_lines:
+            bbox = _line["bbox"]
+            _find = False
+            for c_l_i in range(len(c_lines)):
+                c_l = c_lines[len(c_lines)-c_l_i-1]
+                bbox1 = c_l["bbox"]
+                bboxa = [max(0,bbox[0]-append_width),max(0,bbox[1]-append_width),bbox[2]+append_width,bbox[3]+append_width]
+                bboxb = [max(0,bbox1[0]-append_width),max(0,bbox1[1]-append_width),bbox1[2]+append_width,bbox1[3]+append_width]
+
+                _iou = getIOU(bboxa,bboxb)
+                if _iou>0:
+                    new_bbox = [min(bbox[0],bbox[2],bbox1[0],bbox1[2]),min(bbox[1],bbox[3],bbox1[1],bbox1[3]),max(bbox[0],bbox[2],bbox1[0],bbox1[2]),max(bbox[1],bbox[3],bbox1[1],bbox1[3])]
+                    _find = True
+                    c_l["bbox"] = new_bbox
+                    break
+            if not _find:
+                c_lines.append(_line)
+        after_len = len(c_lines)
+        if first_len==after_len:
+            break
+        list_lines = c_lines
+    return c_lines
+
+
+
+def points2lines(pred,sourceP_LB=True, prob=0.2, line_width=8, padding=3, min_len=10,
+                  cell_width=13):
+    _time = time.time()
+
+    log("starting points2lines")
+    height = len(pred)
+    width = len(pred[0])
 
-def points2lines(pred, sourceP_LB=True, prob=0.2, line_width=7, padding=3, min_len=10,
+    _sum = list(np.sum(np.array((pred[...,0]>prob)).astype(int),axis=1))
+
+    h_index = -1
+    h_lines = []
+    v_lines = []
+    _step = line_width
+    while 1:
+        h_index += 1
+        if h_index>=height:
+            break
+        w_index = -1
+        if sourceP_LB:
+            h_i = height-1-h_index
+        else:
+            h_i = h_index
+        _start = None
+        if _sum[h_index]<min_len:
+            continue
+        while 1:
+            w_index += _step
+            if w_index>=width:
+                break
+            _h,_v = pred[h_i][w_index]
+            if _h>prob:
+                if _start is None:
+                    _start = w_index
+            else:
+                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
+                w_index -= _step//2
+
+    log("starting points2lines 1")
+    w_index = -1
+
+    _sum = list(np.sum(np.array((pred[...,1]>prob)).astype(int),axis=0))
+    _step = line_width
+    while 1:
+        w_index += 1
+        if w_index>=width:
+            break
+        if _sum[w_index]<min_len:
+            continue
+        h_index = -1
+        _start = None
+        while 1:
+            h_index += _step
+            if h_index>=height:
+                break
+
+
+            if sourceP_LB:
+                h_i = height-1-h_index
+            else:
+                h_i = h_index
+
+            _h,_v = pred[h_index][w_index]
+            if _v>prob:
+                if _start is None:
+                    _start = h_i
+            else:
+                if _start is not None:
+                    _end = last_h
+                    _bbox = [w_index,_start,w_index,_end]
+                    _dict = {"bbox":_bbox}
+                    v_lines.append(_dict)
+                    _start = None
+                h_index -= _step//2
+            last_h = h_i
+    log("starting points2lines 2")
+
+    for _line in h_lines:
+        _bbox = _line["bbox"]
+        _bbox = [max(_bbox[0]-2,0),(_bbox[1]+_bbox[3])/2,_bbox[2]+2,(_bbox[1]+_bbox[3])/2]
+        _line["bbox"] = _bbox
+
+    for _line in v_lines:
+        _bbox = _line["bbox"]
+        _bbox = [(_bbox[0]+_bbox[2])/2,max(_bbox[1]-2,0),(_bbox[0]+_bbox[2])/2,_bbox[3]+2]
+        _line["bbox"] = _bbox
+
+    h_lines = lines_cluster(h_lines,line_width=line_width)
+    v_lines = lines_cluster(v_lines,line_width=line_width)
+
+    list_line = []
+    for _line in h_lines:
+        _bbox = _line["bbox"]
+        _bbox = [max(_bbox[0]-1,0),(_bbox[1]+_bbox[3])/2,_bbox[2]+1,(_bbox[1]+_bbox[3])/2]
+        list_line.append(_bbox)
+    for _line in v_lines:
+        _bbox = _line["bbox"]
+        _bbox = [(_bbox[0]+_bbox[2])/2,max(_bbox[1]-1,0),(_bbox[0]+_bbox[2])/2,_bbox[3]+1]
+        list_line.append(_bbox)
+
+    log("points2lines cost %.2fs"%(time.time()-_time))
+
+    # import matplotlib.pyplot as plt
+    # plt.figure()
+    # for _line in list_line:
+    #     x0,y0,x1,y1 = _line
+    #     plt.plot([x0,x1],[y0,y1])
+    # for _line in list_line:
+    #     x0,y0,x1,y1 = _line.bbox
+    #     plt.plot([x0,x1],[y0,y1])
+    # for point in list_crosspoints:
+    #     plt.scatter(point.get("point")[0],point.get("point")[1])
+    # plt.show()
+
+    return list_line
+
+
+
+
+def points2lines_bak(pred, sourceP_LB=True, prob=0.2, line_width=7, padding=3, min_len=10,
                  cell_width=13):
     def inBbox(bbox,point,line_width):
         x,y = point
@@ -317,6 +539,7 @@ def points2lines(pred, sourceP_LB=True, prob=0.2, line_width=7, padding=3, min_l
             return True,[min(x,bbox[0]),min(y,bbox[1]),max(x,bbox[2]),max(y,bbox[3])]
         return False,None
     _time = time.time()
+
     height = len(pred)
     width = len(pred[0])
 
@@ -993,7 +1216,7 @@ def delete_no_cross_lines(list_lines):
     return new_list_lines
 
 
-def delete_short_lines(list_lines, image_shape, scale=40):
+def delete_short_lines(list_lines, image_shape, scale=100):
     x_min_len = max(5, int(image_shape[0] / scale))
     y_min_len = max(5, int(image_shape[1] / scale))
     new_list_lines = []