Commit 8aa8a00a authored by dmMaze's avatar dmMaze
Browse files

allow rectool to use existing mask #563

parent 00a2ff38

ui/drawingpanel.py

+11 −4
Original line number Diff line number Diff line
@@ -7,7 +7,7 @@ import numpy as np 
import cv2



from utils.imgproc_utils import enlarge_window

from utils.textblock_mask import canny_flood, connected_canny_flood

from utils.textblock_mask import canny_flood, connected_canny_flood, existing_mask

from utils.logger import logger



from .module_manager import ModuleManager
@@ -201,7 +201,11 @@ class RectPanel(Widget): 
        self.methodComboBox = QComboBox()

        self.methodComboBox.setFixedHeight(CONFIG_COMBOBOX_HEIGHT)

        self.methodComboBox.setFixedWidth(CONFIG_COMBOBOX_SHORT)

        self.methodComboBox.addItems([self.tr('method 1'), self.tr('method 2')])

        self.methodComboBox.addItems([

            self.tr('method 1'), 

            self.tr('method 2'),

            self.tr('Use Existing Mask')

        ])

        self.autoChecker = QCheckBox(self.tr("Auto"))

        self.autoChecker.setToolTip(self.tr("run inpainting automatically."))

        self.autoChecker.stateChanged.connect(self.on_auto_changed)
@@ -243,8 +247,10 @@ class RectPanel(Widget): 
    def get_maskseg_method(self):

        if self.methodComboBox.currentIndex() == 0:

            return canny_flood

        else:

        elif self.methodComboBox.currentIndex() == 1:

            return connected_canny_flood

        elif self.methodComboBox.currentIndex() == 2:

            return existing_mask



    def on_auto_changed(self):

        if self.autoChecker.isChecked():
@@ -775,7 +781,8 @@ class DrawingPanel(Widget): 
            if mode == 0:

                im = np.copy(img[y1: y2, x1: x2])

                maskseg_method = self.rectPanel.get_maskseg_method()

                inpaint_mask_array, ballon_mask, bub_dict = maskseg_method(im)

                mask = self.canvas.imgtrans_proj.mask_array[y1: y2, x1: x2]

                inpaint_mask_array, ballon_mask, bub_dict = maskseg_method(im, mask=mask)

                mask = self.rectPanel.post_process_mask(inpaint_mask_array)



                bground_bgr = bub_dict['bground_bgr']

utils/textblock_mask.py

+7 −2
Original line number Diff line number Diff line
@@ -94,7 +94,7 @@ def bground_calculator(buble_img, back_ground_mask, dilate=True): 
    return bground_aver, bground_region, sd



# 输入：文本块roi，分割出文本mask，根据mask计算文本bgr均值和标准差，决定纯色覆盖/inpaint修复

def canny_flood(img, show_process=False, inpaint_sdthresh=10):

def canny_flood(img, show_process=False, inpaint_sdthresh=10, **kwargs):

    # cv2.setNumThreads(4)

    WHITE = (255, 255, 255)

    BLACK = (0, 0, 0)
@@ -207,7 +207,7 @@ def canny_flood(img, show_process=False, inpaint_sdthresh=10): 
    return mask, ballon_mask, bub_dict



# 输入：文本块roi，分割出文本mask，根据mask计算文本bgr均值和标准差，决定纯色覆盖/inpaint修复

def connected_canny_flood(img, show_process=False, inpaint_sdthresh=10, apply_strokewidth_check=0):

def connected_canny_flood(img, show_process=False, inpaint_sdthresh=10, apply_strokewidth_check=0, **kwargs):



    # 寻找最可能是气泡的外轮廓mask

    def find_outermask(img):
@@ -341,6 +341,11 @@ def connected_canny_flood(img, show_process=False, inpaint_sdthresh=10, apply_st 
    return mask, ballon_mask, bub_dict





def existing_mask(img, mask: np.ndarray):

    bub_dict = {"bgr": [0, 0, 0],"bground_bgr": [255, 255, 255],"need_inpaint": True}

    return mask, mask, bub_dict





def extract_ballon_mask(img: np.ndarray, mask: np.ndarray) -> Tuple[np.ndarray, np.ndarray]:

    '''

    Given original img and text mask (cropped)