Ensemble_Semantic_Seg/postProc.py at main · LIPAI-Org/Ensemble_Semantic_Seg · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
import numpy as np
from skimage.color import rgb2gray
from skimage.morphology import diamond, dilation, erosion, remove_small_objects
from scipy.ndimage import binary_fill_holes


def postProc(inputImage, orIm):
    #binarizar máscara
    if inputImage.ndim == 3:
        inputGray = rgb2gray(inputImage)
    else:
        inputGray = inputImage.astype(np.float32) / 255.0
    binaryIm = (inputGray > 0.5).astype(np.uint8)


    #Dilatação
    se = diamond(1)
    dilatedIm = dilation(binaryIm, se)
    finalMaskRGB = np.stack([dilatedIm*255]*3, axis=-1).astype(np.uint8)
    finalIm2 = orIm.copy()
    finalIm2[finalMaskRGB < 200] = 0

    #Preenchimento de buracos
    filledIm = binary_fill_holes(dilatedIm).astype(np.uint8)
    finalMaskRGB = np.stack([filledIm*255]*3, axis=-1).astype(np.uint8)
    finalIm3 = orIm.copy()
    finalIm3[finalMaskRGB < 200] = 0

    #Erosão
    erodedIm = erosion(filledIm, se)
    finalMaskRGB = np.stack([erodedIm*255]*3, axis=-1).astype(np.uint8)
    finalIm4 = orIm.copy()
    finalIm4[finalMaskRGB < 200] = 0

    #Remoção de objetos pequenos
    cleanedIm = remove_small_objects(erodedIm.astype(bool), 30).astype(np.uint8)
    finalMaskRGB = np.stack([cleanedIm*255]*3, axis=-1).astype(np.uint8)
    finalIm5 = orIm.copy()
    finalIm5[finalMaskRGB < 200] = 0

    return finalMaskRGB, finalIm5