summaries.py 文件源码

def postprocess_colormap(cls, postprocess=True):
    """Create a colormap out of the classes and postprocess the face."""
    batch = vs.apply_colormap(cls, vmin=0, vmax=21, cmap=CMAP)
    cmap = vs.apply_colormap(np.array(range(22), dtype='uint8'),
                             vmin=0, vmax=21, cmap=CMAP)
    COLSET = cmap[18:22]
    FCOL = cmap[11]
    if postprocess:
        kernel = np.ones((2, 2), dtype=np.uint8)
        for im in batch:
            for col in COLSET:
                # Extract the map of the matching color.
                colmap = np.all(im == col, axis=2).astype(np.uint8)
                # Erode.
                while np.sum(colmap) > 10:
                    colmap = cv2.erode(colmap, kernel)
                # Prepare the original map for remapping.
                im[np.all(im == col, axis=2)] = FCOL
                # Backproject.
                im[colmap == 1] = col
    return batch[:, :, :, :3]