import cv2 from os.path import join as pjoin import time import numpy as np from detect_compo.lib_ip.Component import Component from config.CONFIG_UIED import Config C = Config() class Block(Component): def __init__(self, region, image_shape): super().__init__(region, image_shape) self.category = 'Block' self.parent = None self.children = [] self.uicompo_ = None self.top_or_botm = None self.redundant = False def block_is_uicompo(self, image_shape, max_compo_scale): ''' Check the if the block is a ui component according to its relative size ''' row, column = image_shape[:2] # print(height, height / row, max_compo_scale[0], height / row > max_compo_scale[0]) # draw.draw_bounding_box(org, [corner], show=True) # ignore atomic components if self.bbox.height / row > max_compo_scale[0] or self.bbox.width / column > max_compo_scale[1]: return False return True def block_is_top_or_bottom_bar(self, image_shape, top_bottom_height): ''' Check if the block is top bar or bottom bar ''' height, width = image_shape[:2] (column_min, row_min, column_max, row_max) = self.bbox.put_bbox() if column_min < 5 and row_min < 5 and \ width - column_max < 5 and row_max < height * top_bottom_height[0]: self.uicompo_ = True return True if column_min < 5 and row_min > height * top_bottom_height[1] and \ width - column_max < 5 and height - row_max < 5: self.uicompo_ = True return True return False def block_erase_from_bin(self, binary, pad): (column_min, row_min, column_max, row_max) = self.put_bbox() column_min = max(column_min - pad, 0) column_max = min(column_max + pad, binary.shape[1]) row_min = max(row_min - pad, 0) row_max = min(row_max + pad, binary.shape[0]) cv2.rectangle(binary, (column_min, row_min), (column_max, row_max), (0), -1)