post_process.py 文件源码

def watershed_cut(depthImage, ssMask):
    ssMask = ssMask.astype(np.int32)
    resultImage = np.zeros(shape=ssMask.shape, dtype=np.float32)

    for semClass in CLASS_TO_CITYSCAPES.keys():
        csCode = CLASS_TO_CITYSCAPES[semClass]
        ssCode = CLASS_TO_SS[semClass]
        ssMaskClass = (ssMask == ssCode)

        ccImage = (depthImage > THRESHOLD[semClass]) * ssMaskClass
        ccImage = skimage.morphology.remove_small_objects(ccImage, min_size=MIN_SIZE[semClass])
        ccImage = skimage.morphology.remove_small_holes(ccImage)
        ccLabels = skimage.morphology.label(ccImage)

        ccIDs = np.unique(ccLabels)[1:]
        for ccID in ccIDs:
            ccIDMask = (ccLabels == ccID)
            ccIDMask = skimage.morphology.binary_dilation(ccIDMask, SELEM[THRESHOLD[semClass]])
            instanceID = 1000 * csCode + ccID
            resultImage[ccIDMask] = instanceID

    resultImage = resultImage.astype(np.uint16)
    return resultImage