mask.py 文件源码

def main(args):
    if args.threshold is None:
        print("Please provide a binarization threshold")
        return 1
    data, hdr = read(args.input, inc_header=True)
    mask = data >= args.threshold
    if args.minvol is not None:
        mask = binary_volume_opening(mask, args.minvol)
    if args.fill:
        mask = binary_fill_holes(mask)
    if args.extend is not None and args.extend > 0:
        if args.relion:
            se = binary_sphere(args.extend, False)
            mask = binary_dilation(mask, structure=se, iterations=1)
        else:
            dt = distance_transform_edt(~mask)
            mask = mask | (dt <= args.edge_width)
    if args.close:
        se = binary_sphere(args.extend, False)
        mask = binary_closing(mask, structure=se, iterations=1)
    final = mask.astype(np.single)
    if args.edge_width is not None:
        dt = distance_transform_edt(~mask)  # Compute *outward* distance transform of mask.
        idx = (dt <= args.edge_width) & (dt > 0)  # Identify edge points by distance from mask.
        x = np.arange(1, args.edge_width + 1)  # Domain of the edge profile.
        if "sin" in args.edge_profile:
            y = np.sin(np.linspace(np.pi/2, 0, args.edge_width + 1))  # Range of the edge profile.
        f = interp1d(x, y[1:])
        final[idx] = f(dt[idx])  # Insert edge heights interpolated at distance transform values.
    write(args.output, final, psz=hdr["xlen"] / hdr["nx"])
    return 0