Spaces:
Running
Running
| import numpy as np | |
| import cv2 | |
| from collections import Counter | |
| import lib_ip.ip_draw as draw | |
| from config.CONFIG_UIED import Config | |
| C = Config() | |
| # detect object(connected region) | |
| # def boundary_bfs_connected_area(img, x, y, mark): | |
| # def neighbor(img, x, y, mark, stack): | |
| # for i in range(x - 1, x + 2): | |
| # if i < 0 or i >= img.shape[0]: continue | |
| # for j in range(y - 1, y + 2): | |
| # if j < 0 or j >= img.shape[1]: continue | |
| # if img[i, j] == 255 and mark[i, j] == 0: | |
| # stack.append([i, j]) | |
| # mark[i, j] = 255 | |
| # | |
| # stack = [[x, y]] # points waiting for inspection | |
| # area = [[x, y]] # points of this area | |
| # mark[x, y] = 255 # drawing broad | |
| # | |
| # while len(stack) > 0: | |
| # point = stack.pop() | |
| # area.append(point) | |
| # neighbor(img, point[0], point[1], mark, stack) | |
| # return area | |
| # def line_check_perpendicular(lines_h, lines_v, max_thickness): | |
| # """ | |
| # lines: [line_h, line_v] | |
| # -> line_h: horizontal {'head':(column_min, row), 'end':(column_max, row), 'thickness':int) | |
| # -> line_v: vertical {'head':(column, row_min), 'end':(column, row_max), 'thickness':int} | |
| # """ | |
| # is_per_h = np.full(len(lines_h), False) | |
| # is_per_v = np.full(len(lines_v), False) | |
| # for i in range(len(lines_h)): | |
| # # save the intersection point of h | |
| # lines_h[i]['inter_point'] = set() | |
| # h = lines_h[i] | |
| # | |
| # for j in range(len(lines_v)): | |
| # # save the intersection point of v | |
| # if 'inter_point' not in lines_v[j]: lines_v[j]['inter_point'] = set() | |
| # v = lines_v[j] | |
| # | |
| # # if h is perpendicular to v in head of v | |
| # if abs(h['head'][1]-v['head'][1]) <= max_thickness: | |
| # if abs(h['head'][0] - v['head'][0]) <= max_thickness: | |
| # lines_h[i]['inter_point'].add('head') | |
| # lines_v[j]['inter_point'].add('head') | |
| # is_per_h[i] = True | |
| # is_per_v[j] = True | |
| # elif abs(h['end'][0] - v['head'][0]) <= max_thickness: | |
| # lines_h[i]['inter_point'].add('end') | |
| # lines_v[j]['inter_point'].add('head') | |
| # is_per_h[i] = True | |
| # is_per_v[j] = True | |
| # | |
| # # if h is perpendicular to v in end of v | |
| # elif abs(h['head'][1]-v['end'][1]) <= max_thickness: | |
| # if abs(h['head'][0] - v['head'][0]) <= max_thickness: | |
| # lines_h[i]['inter_point'].add('head') | |
| # lines_v[j]['inter_point'].add('end') | |
| # is_per_h[i] = True | |
| # is_per_v[j] = True | |
| # elif abs(h['end'][0] - v['head'][0]) <= max_thickness: | |
| # lines_h[i]['inter_point'].add('end') | |
| # lines_v[j]['inter_point'].add('end') | |
| # is_per_h[i] = True | |
| # is_per_v[j] = True | |
| # per_h = [] | |
| # per_v = [] | |
| # for i in range(len(is_per_h)): | |
| # if is_per_h[i]: | |
| # lines_h[i]['inter_point'] = list(lines_h[i]['inter_point']) | |
| # per_h.append(lines_h[i]) | |
| # for i in range(len(is_per_v)): | |
| # if is_per_v[i]: | |
| # lines_v[i]['inter_point'] = list(lines_v[i]['inter_point']) | |
| # per_v.append(lines_v[i]) | |
| # return per_h, per_v | |
| # def line_shrink_corners(corner, lines_h, lines_v): | |
| # """ | |
| # shrink the corner according to lines: | |
| # col_min_shrink: shrink right (increase) | |
| # col_max_shrink: shrink left (decrease) | |
| # row_min_shrink: shrink down (increase) | |
| # row_max_shrink: shrink up (decrease) | |
| # :param lines_h: horizontal {'head':(column_min, row), 'end':(column_max, row), 'thickness':int) | |
| # :param lines_v: vertical {'head':(column, row_min), 'end':(column, row_max), 'thickness':int} | |
| # :return: shrunken corner: (top_left, bottom_right) | |
| # """ | |
| # (col_min, row_min), (col_max, row_max) = corner | |
| # col_min_shrink, row_min_shrink = col_min, row_min | |
| # col_max_shrink, row_max_shrink = col_max, row_max | |
| # valid_frame = False | |
| # | |
| # for h in lines_h: | |
| # # ignore outer border | |
| # if len(h['inter_point']) == 2: | |
| # valid_frame = True | |
| # continue | |
| # # shrink right -> col_min move to end | |
| # if h['inter_point'][0] == 'head': | |
| # col_min_shrink = max(h['end'][0], col_min_shrink) | |
| # # shrink left -> col_max move to head | |
| # elif h['inter_point'][0] == 'end': | |
| # col_max_shrink = min(h['head'][0], col_max_shrink) | |
| # | |
| # for v in lines_v: | |
| # # ignore outer border | |
| # if len(v['inter_point']) == 2: | |
| # valid_frame = True | |
| # continue | |
| # # shrink down -> row_min move to end | |
| # if v['inter_point'][0] == 'head': | |
| # row_min_shrink = max(v['end'][1], row_min_shrink) | |
| # # shrink up -> row_max move to head | |
| # elif v['inter_point'][0] == 'end': | |
| # row_max_shrink = min(v['head'][1], row_max_shrink) | |
| # | |
| # # return the shrunken corner if only there is line intersecting with two other lines | |
| # if valid_frame: | |
| # return (col_min_shrink, row_min_shrink), (col_max_shrink, row_max_shrink) | |
| # return corner | |
| # def line_cvt_relative_position(col_min, row_min, lines_h, lines_v): | |
| # """ | |
| # convert the relative position of lines in the entire image | |
| # :param col_min: based column the img lines belong to | |
| # :param row_min: based row the img lines belong to | |
| # :param lines_h: horizontal {'head':(column_min, row), 'end':(column_max, row), 'thickness':int) | |
| # :param lines_v: vertical {'head':(column, row_min), 'end':(column, row_max), 'thickness':int} | |
| # :return: lines_h_cvt, lines_v_cvt | |
| # """ | |
| # for h in lines_h: | |
| # h['head'][0] += col_min | |
| # h['head'][1] += row_min | |
| # h['end'][0] += col_min | |
| # h['end'][1] += row_min | |
| # for v in lines_v: | |
| # v['head'][0] += col_min | |
| # v['head'][1] += row_min | |
| # v['end'][0] += col_min | |
| # v['end'][1] += row_min | |
| # | |
| # return lines_h, lines_v | |
| # check if an object is so slim | |
| # @boundary: [border_up, border_bottom, border_left, border_right] | |
| # -> up, bottom: (column_index, min/max row border) | |
| # -> left, right: (row_index, min/max column border) detect range of each row | |
| def clipping_by_line(boundary, boundary_rec, lines): | |
| boundary = boundary.copy() | |
| for orient in lines: | |
| # horizontal | |
| if orient == 'h': | |
| # column range of sub area | |
| r1, r2 = 0, 0 | |
| for line in lines[orient]: | |
| if line[0] == 0: | |
| r1 = line[1] | |
| continue | |
| r2 = line[0] | |
| b_top = [] | |
| b_bottom = [] | |
| for i in range(len(boundary[0])): | |
| if r2 > boundary[0][i][0] >= r1: | |
| b_top.append(boundary[0][i]) | |
| for i in range(len(boundary[1])): | |
| if r2 > boundary[1][i][0] >= r1: | |
| b_bottom.append(boundary[1][i]) | |
| b_left = [x for x in boundary[2]] # (row_index, min column border) | |
| for i in range(len(b_left)): | |
| if b_left[i][1] < r1: | |
| b_left[i][1] = r1 | |
| b_right = [x for x in boundary[3]] # (row_index, max column border) | |
| for i in range(len(b_right)): | |
| if b_right[i][1] > r2: | |
| b_right[i][1] = r2 | |
| boundary_rec.append([b_top, b_bottom, b_left, b_right]) | |
| r1 = line[1] | |
| # remove imgs that contain text | |
| # def rm_text(org, corners, compo_class, | |
| # max_text_height=C.THRESHOLD_TEXT_MAX_HEIGHT, max_text_width=C.THRESHOLD_TEXT_MAX_WIDTH, | |
| # ocr_padding=C.OCR_PADDING, ocr_min_word_area=C.OCR_MIN_WORD_AREA, show=False): | |
| # """ | |
| # Remove area that full of text | |
| # :param org: original image | |
| # :param corners: [(top_left, bottom_right)] | |
| # -> top_left: (column_min, row_min) | |
| # -> bottom_right: (column_max, row_max) | |
| # :param compo_class: classes of corners | |
| # :param max_text_height: Too large to be text | |
| # :param max_text_width: Too large to be text | |
| # :param ocr_padding: Padding for clipping | |
| # :param ocr_min_word_area: If too text area ratio is too large | |
| # :param show: Show or not | |
| # :return: corners without text objects | |
| # """ | |
| # new_corners = [] | |
| # new_class = [] | |
| # for i in range(len(corners)): | |
| # corner = corners[i] | |
| # (top_left, bottom_right) = corner | |
| # (col_min, row_min) = top_left | |
| # (col_max, row_max) = bottom_right | |
| # height = row_max - row_min | |
| # width = col_max - col_min | |
| # # highly likely to be block or img if too large | |
| # if height > max_text_height and width > max_text_width: | |
| # new_corners.append(corner) | |
| # new_class.append(compo_class[i]) | |
| # else: | |
| # row_min = row_min - ocr_padding if row_min - ocr_padding >= 0 else 0 | |
| # row_max = row_max + ocr_padding if row_max + ocr_padding < org.shape[0] else org.shape[0] | |
| # col_min = col_min - ocr_padding if col_min - ocr_padding >= 0 else 0 | |
| # col_max = col_max + ocr_padding if col_max + ocr_padding < org.shape[1] else org.shape[1] | |
| # # check if this area is text | |
| # clip = org[row_min: row_max, col_min: col_max] | |
| # if not ocr.is_text(clip, ocr_min_word_area, show=show): | |
| # new_corners.append(corner) | |
| # new_class.append(compo_class[i]) | |
| # return new_corners, new_class | |
| # def rm_img_in_compo(corners_img, corners_compo): | |
| # """ | |
| # Remove imgs in component | |
| # """ | |
| # corners_img_new = [] | |
| # for img in corners_img: | |
| # is_nested = False | |
| # for compo in corners_compo: | |
| # if util.corner_relation(img, compo) == -1: | |
| # is_nested = True | |
| # break | |
| # if not is_nested: | |
| # corners_img_new.append(img) | |
| # return corners_img_new | |
| # def block_or_compo(org, binary, corners, | |
| # max_thickness=C.THRESHOLD_BLOCK_MAX_BORDER_THICKNESS, max_block_cross_points=C.THRESHOLD_BLOCK_MAX_CROSS_POINT, | |
| # min_compo_w_h_ratio=C.THRESHOLD_UICOMPO_MIN_W_H_RATIO, max_compo_w_h_ratio=C.THRESHOLD_UICOMPO_MAX_W_H_RATIO, | |
| # min_block_edge=C.THRESHOLD_BLOCK_MIN_EDGE_LENGTH): | |
| # """ | |
| # Check if the objects are img components or just block | |
| # :param org: Original image | |
| # :param binary: Binary image from pre-processing | |
| # :param corners: [(top_left, bottom_right)] | |
| # -> top_left: (column_min, row_min) | |
| # -> bottom_right: (column_max, row_max) | |
| # :param max_thickness: The max thickness of border of blocks | |
| # :param max_block_cross_points: Ratio of point of interaction | |
| # :return: corners of blocks and imgs | |
| # """ | |
| # blocks = [] | |
| # imgs = [] | |
| # compos = [] | |
| # for corner in corners: | |
| # (top_left, bottom_right) = corner | |
| # (col_min, row_min) = top_left | |
| # (col_max, row_max) = bottom_right | |
| # height = row_max - row_min | |
| # width = col_max - col_min | |
| # | |
| # block = False | |
| # vacancy = [0, 0, 0, 0] | |
| # for i in range(1, max_thickness): | |
| # try: | |
| # # top to bottom | |
| # if vacancy[0] == 0 and (col_max - col_min - 2 * i) is not 0 and ( | |
| # np.sum(binary[row_min + i, col_min + i: col_max - i]) / 255) / (col_max - col_min - 2 * i) <= max_block_cross_points: | |
| # vacancy[0] = 1 | |
| # # bottom to top | |
| # if vacancy[1] == 0 and (col_max - col_min - 2 * i) is not 0 and ( | |
| # np.sum(binary[row_max - i, col_min + i: col_max - i]) / 255) / (col_max - col_min - 2 * i) <= max_block_cross_points: | |
| # vacancy[1] = 1 | |
| # # left to right | |
| # if vacancy[2] == 0 and (row_max - row_min - 2 * i) is not 0 and ( | |
| # np.sum(binary[row_min + i: row_max - i, col_min + i]) / 255) / (row_max - row_min - 2 * i) <= max_block_cross_points: | |
| # vacancy[2] = 1 | |
| # # right to left | |
| # if vacancy[3] == 0 and (row_max - row_min - 2 * i) is not 0 and ( | |
| # np.sum(binary[row_min + i: row_max - i, col_max - i]) / 255) / (row_max - row_min - 2 * i) <= max_block_cross_points: | |
| # vacancy[3] = 1 | |
| # if np.sum(vacancy) == 4: | |
| # block = True | |
| # except: | |
| # pass | |
| # | |
| # # too big to be UI components | |
| # if block: | |
| # if height > min_block_edge and width > min_block_edge: | |
| # blocks.append(corner) | |
| # else: | |
| # if min_compo_w_h_ratio < width / height < max_compo_w_h_ratio: | |
| # compos.append(corner) | |
| # # filter out small objects | |
| # else: | |
| # if height > min_block_edge: | |
| # imgs.append(corner) | |
| # else: | |
| # if min_compo_w_h_ratio < width / height < max_compo_w_h_ratio: | |
| # compos.append(corner) | |
| # return blocks, imgs, compos | |
| # def compo_on_img(processing, org, binary, clf, | |
| # compos_corner, compos_class): | |
| # """ | |
| # Detect potential UI components inner img; | |
| # Only leave non-img | |
| # """ | |
| # pad = 2 | |
| # for i in range(len(compos_corner)): | |
| # if compos_class[i] != 'img': | |
| # continue | |
| # ((col_min, row_min), (col_max, row_max)) = compos_corner[i] | |
| # col_min = max(col_min - pad, 0) | |
| # col_max = min(col_max + pad, org.shape[1]) | |
| # row_min = max(row_min - pad, 0) | |
| # row_max = min(row_max + pad, org.shape[0]) | |
| # area = (col_max - col_min) * (row_max - row_min) | |
| # if area < 600: | |
| # continue | |
| # | |
| # clip_org = org[row_min:row_max, col_min:col_max] | |
| # clip_bin_inv = pre.reverse_binary(binary[row_min:row_max, col_min:col_max]) | |
| # | |
| # compos_boundary_new, compos_corner_new, compos_class_new = processing(clip_org, clip_bin_inv, clf) | |
| # compos_corner_new = util.corner_cvt_relative_position(compos_corner_new, col_min, row_min) | |
| # | |
| # assert len(compos_corner_new) == len(compos_class_new) | |
| # | |
| # # only leave non-img elements | |
| # for i in range(len(compos_corner_new)): | |
| # ((col_min_new, row_min_new), (col_max_new, row_max_new)) = compos_corner_new[i] | |
| # area_new = (col_max_new - col_min_new) * (row_max_new - row_min_new) | |
| # if compos_class_new[i] != 'img' and area_new / area < 0.8: | |
| # compos_corner.append(compos_corner_new[i]) | |
| # compos_class.append(compos_class_new[i]) | |
| # | |
| # return compos_corner, compos_class | |
| # def strip_img(corners_compo, compos_class, corners_img): | |
| # """ | |
| # Separate img from other compos | |
| # :return: compos without img | |
| # """ | |
| # corners_compo_withuot_img = [] | |
| # compo_class_withuot_img = [] | |
| # for i in range(len(compos_class)): | |
| # if compos_class[i] == 'img': | |
| # corners_img.append(corners_compo[i]) | |
| # else: | |
| # corners_compo_withuot_img.append(corners_compo[i]) | |
| # compo_class_withuot_img.append(compos_class[i]) | |
| # return corners_compo_withuot_img, compo_class_withuot_img | |
| # def merge_corner(corners, compos_class, min_selected_IoU=C.THRESHOLD_MIN_IOU, is_merge_nested_same=True): | |
| # """ | |
| # Calculate the Intersection over Overlap (IoU) and merge corners according to the value of IoU | |
| # :param is_merge_nested_same: if true, merge the nested corners with same class whatever the IoU is | |
| # :param corners: corners: [(top_left, bottom_right)] | |
| # -> top_left: (column_min, row_min) | |
| # -> bottom_right: (column_max, row_max) | |
| # :return: new corners | |
| # """ | |
| # new_corners = [] | |
| # new_class = [] | |
| # for i in range(len(corners)): | |
| # is_intersected = False | |
| # for j in range(len(new_corners)): | |
| # r = util.corner_relation_nms(corners[i], new_corners[j], min_selected_IoU) | |
| # # r = util.corner_relation(corners[i], new_corners[j]) | |
| # if is_merge_nested_same: | |
| # if compos_class[i] == new_class[j]: | |
| # # if corners[i] is in new_corners[j], ignore corners[i] | |
| # if r == -1: | |
| # is_intersected = True | |
| # break | |
| # # if new_corners[j] is in corners[i], replace new_corners[j] with corners[i] | |
| # elif r == 1: | |
| # is_intersected = True | |
| # new_corners[j] = corners[i] | |
| # | |
| # # if above IoU threshold, and corners[i] is in new_corners[j], ignore corners[i] | |
| # if r == -2: | |
| # is_intersected = True | |
| # break | |
| # # if above IoU threshold, and new_corners[j] is in corners[i], replace new_corners[j] with corners[i] | |
| # elif r == 2: | |
| # is_intersected = True | |
| # new_corners[j] = corners[i] | |
| # new_class[j] = compos_class[i] | |
| # | |
| # # containing and too small | |
| # elif r == -3: | |
| # is_intersected = True | |
| # break | |
| # elif r == 3: | |
| # is_intersected = True | |
| # new_corners[j] = corners[i] | |
| # | |
| # # if [i] and [j] are overlapped but no containing relation, merge corners when same class | |
| # elif r == 4: | |
| # is_intersected = True | |
| # if compos_class[i] == new_class[j]: | |
| # new_corners[j] = util.corner_merge_two_corners(corners[i], new_corners[j]) | |
| # | |
| # if not is_intersected: | |
| # new_corners.append(corners[i]) | |
| # new_class.append(compos_class[i]) | |
| # return new_corners, new_class | |
| # def select_corner(corners, compos_class, class_name): | |
| # """ | |
| # Select corners in given compo type | |
| # """ | |
| # corners_wanted = [] | |
| # for i in range(len(compos_class)): | |
| # if compos_class[i] == class_name: | |
| # corners_wanted.append(corners[i]) | |
| # return corners_wanted | |
| # def flood_fill_bfs(img, x_start, y_start, mark, grad_thresh): | |
| # def neighbor(x, y): | |
| # for i in range(x - 1, x + 2): | |
| # if i < 0 or i >= img.shape[0]: continue | |
| # for j in range(y - 1, y + 2): | |
| # if j < 0 or j >= img.shape[1]: continue | |
| # if mark[i, j] == 0 and abs(img[i, j] - img[x, y]) < grad_thresh: | |
| # stack.append([i, j]) | |
| # mark[i, j] = 255 | |
| # | |
| # stack = [[x_start, y_start]] # points waiting for inspection | |
| # region = [[x_start, y_start]] # points of this connected region | |
| # mark[x_start, y_start] = 255 # drawing broad | |
| # while len(stack) > 0: | |
| # point = stack.pop() | |
| # region.append(point) | |
| # neighbor(point[0], point[1]) | |
| # return region |