Commit c8c220e2 authored by dmMaze's avatar dmMaze
Browse files

fix lines merging

parent 862a223f

utils/textblock.py

+24 −12
Original line number Diff line number Diff line
@@ -606,7 +606,7 @@ def examine_textblk(blk: TextBlock, im_w: int, im_h: int, sort: bool = False) -> 
    if sort:

        blk.sort_lines()



def try_merge_textline(blk: TextBlock, blk2: TextBlock, fntsize_tol=1.7, distance_tol=2) -> bool:

def try_merge_textline(blk: TextBlock, blk2: TextBlock, fntsize_tol=1.7, distance_tol=2, canvas=None) -> bool:

    if blk2.merged:

        return False

    fntsize_div = blk.font_size / blk2.font_size
@@ -619,29 +619,37 @@ def try_merge_textline(blk: TextBlock, blk2: TextBlock, fntsize_tol=1.7, distanc 
    # distance_p1 = np.linalg.norm(np.array(blk2.lines[-1][0]) - np.array(blk.lines[-1][0]))

    minrect1 = blk.min_rect(ids=[-1])[0]

    xyxy1 = [*minrect1[0], *minrect1[2]]

    minrect2 = blk2.min_rect(ids=[-1])[0]

    minrect2 = blk2.min_rect(ids=[0])[0]

    xyxy2 = [*minrect2[0], *minrect2[2]]

    distance_x = max(xyxy1[0], xyxy2[0]) - min(xyxy1[2], xyxy2[2])

    distance_y = max(xyxy1[1], xyxy2[1]) - min(xyxy1[3], xyxy2[3])

    w1 = xyxy1[2] - xyxy1[0]

    w2 = xyxy2[2] - xyxy2[0]



    l1, l2 = Polygon(blk.lines[-1]), Polygon(blk2.lines[-1])

    l1, l2 = Polygon(blk.lines[-1]), Polygon(blk2.lines[0])

    if not l1.intersects(l2):

        if blk.vertical:

            if distance_y > 0:

                return False

            if distance_x > fntsz_avg * 0.8:

                return False

        else:

            if distance_x > 0:

                return False

            fntsz_thr = 0.5

            if fntsz_avg < 24:

                fntsz_thr = 0.6

            if distance_y > fntsz_avg * fntsz_thr:

                return False

            if abs(distance_x) / min(w1, w2) < 0.3:

                return False

        if fntsize_div > fntsize_tol or 1 / fntsize_div > fntsize_tol:

            return False

        if abs(cos_vec) < 0.866:   # cos30

            return False

        # if distance > distance_tol * fntsz_avg:

        #     return False

        if blk.vertical and blk2.vertical and distance_x > fntsz_avg * 0.8:

            return False

        if not blk.vertical and distance_y > fntsz_avg * 0.5:

            return False



    # merge

    for line in blk2.lines:

        blk.lines.append(line)
@@ -655,16 +663,15 @@ def try_merge_textline(blk: TextBlock, blk2: TextBlock, fntsize_tol=1.7, distanc 
    blk2.merged = True

    return True



def merge_textlines(blk_list: List[TextBlock]) -> List[TextBlock]:

def merge_textlines(blk_list: List[TextBlock], canvas=None, fntsize_tol=1.7) -> List[TextBlock]:

    if len(blk_list) < 2:

        return blk_list

    blk_list.sort(key=lambda blk: blk.distance[0])

    merged_list = []

    for ii, current_blk in enumerate(blk_list):

        if current_blk.merged:

            continue

        for jj, blk in enumerate(blk_list[ii+1:]):

            try_merge_textline(current_blk, blk)

            try_merge_textline(current_blk, blk, canvas=canvas, fntsize_tol=fntsize_tol)

        merged_list.append(current_blk)

    for blk in merged_list:

        blk.adjust_bbox(with_bbox=False)
@@ -778,8 +785,13 @@ def group_output(blks, lines, im_w, im_h, mask=None, sort_blklist=True, canvas=N 
        else:

            _final_blk_list.append(blk)

    final_blk_list = _final_blk_list



    # step3: merge scattered lines, sort textblocks by "grid"

    final_blk_list += merge_textlines(scattered_lines['hor'])

    scattered_lines['ver'].sort(key=lambda blk: blk.center()[0])

    scattered_lines['hor'].sort(key=lambda blk: blk.center()[1])

    # c = visualize_textblocks(canvas, scattered_lines['hor'])

    # cv2.imwrite('local_tst.jpg', c)

    final_blk_list += merge_textlines(scattered_lines['hor'], canvas=canvas, fntsize_tol=2.0)

    final_blk_list += merge_textlines(scattered_lines['ver'])

    if sort_blklist:

        final_blk_list = sort_regions(final_blk_list, )
@@ -830,7 +842,7 @@ def visualize_textblocks(canvas, blk_list: List[TextBlock]): 
        cv2.polylines(canvas, [blk.min_rect()], True, (127,127,0), 2)

        center = [int((bx1 + bx2)/2), int((by1 + by2)/2)]

        cv2.putText(canvas, str(blk.angle), center, cv2.FONT_HERSHEY_SIMPLEX, 1, (127,127,255), 2)

        cv2.putText(canvas, str(ii), (bx1, by1 + lw + 2), 0, lw / 3, (255,127,127), max(lw-1, 1), cv2.LINE_AA)

        cv2.putText(canvas, str(ii), (bx1, by1 + lw + 2), 0, lw / 6, (255,127,127), max(lw-7, 1), cv2.LINE_AA)

    return canvas



def collect_textblock_regions(img: np.ndarray, textblk_lst: List[TextBlock], text_height=48, maxwidth=8100, split_textblk = False, seg_func: Callable = None):

utils/textlines_merge.py

+4 −1
Original line number Diff line number Diff line
@@ -303,12 +303,15 @@ def sort_pnts(pts: np.ndarray): 
    pairwise_vec = (pts[:, None] - pts[None]).reshape((16, -1))

    pairwise_vec_norm = np.linalg.norm(pairwise_vec, axis=1)

    long_side_ids = np.argsort(pairwise_vec_norm)[[8, 10]]

    pairwise_vec_norm_sorted = pairwise_vec_norm[long_side_ids]

    long_side_vecs = pairwise_vec[long_side_ids]

    inner_prod = (long_side_vecs[0] * long_side_vecs[1]).sum()

    if inner_prod < 0:

        long_side_vecs[0] = -long_side_vecs[0]

    struc_vec = np.abs(long_side_vecs.mean(axis=0))

    is_vertical = struc_vec[0] <= struc_vec[1]

    is_vertical = struc_vec[0] * 1.2 <= struc_vec[1]

    if len(set(pairwise_vec_norm_sorted[4: 12])) == 0:  # is square

        is_vertical = False



    if is_vertical:

        pts = pts[np.argsort(pts[:, 1])]