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