alpha_matting_segmentation.py 文件源码

def create_scribble_mask(x, y, frame, pct_area_shrink=0.1, pct_area_grow=0.1, winsz=3):
    """
    Arguments:
        x               = x-coordinates of bbox
        y               = y-coordinates of bbox
        frame           = image to create scribble mask for
        pct_area_shrink = fraction of bbox area to shrink by: 
                          new area = old area * (1-pct_area_shrink)
        pct_area_grow   = fraction of bbox area to expand by: 
                          new area = old area * pct_area_grow
        winsz           = window size of the alpha matting algorithm

    Output:
      scribble_mask = mask for pixels. Contains 1 for definite foreground, 
                      -1 for definite background, and 0 for unknown pixels.
    """
    ih, iw = frame.shape[:2]

    # expand bbox by 'pct_area_grow' area
    xe, ye = expand_region(x, y, pct_area_grow, ih, iw, winsz)

    # shrink bbox by 'pct_area_shrink' area
    xs, ys = expand_region(x, y, 1-pct_area_shrink, ih, iw, winsz)

    # pixel indices (row, column) for expanded/contracted bboxes
    re, ce = polygon(xe, ye) # pixels for expanded polygon
    rs, cs = polygon(xs, ys) # pixels for contracted polygon

    # mask for outside of expanded bbox
    expanded_mask = np.zeros((ih, iw), dtype='bool')
    expanded_mask[ce, re] = True # mark inside expanded bbox
    expanded_mask = ~expanded_mask # invert to get outside

    # final scribble mask with -1 for outside expanded bbox, 1 inside, 0 otherwise
    scribble_mask = np.zeros((ih, iw), dtype='int')
    scribble_mask.flat[expanded_mask.ravel()] = -1
    scribble_mask[cs, rs] = 1

    return scribble_mask