ScreenCoder_UI2Code/UIED/detect_compo/deprecated/Block.py

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)