def estimate_alpha_sm(
    image,
    trimap,
    return_foreground_background=False,
    trimap_expansion_radius=10,
    trimap_expansion_threshold=0.02,
    sample_gathering_angles=4,
    sample_gathering_weights=(3.0, 2.0, 1.0, 4.0),
    sample_gathering_Np_radius=1,
    sample_refinement_radius=5,
    local_smoothing_radius1=5,
    local_smoothing_radius2=5,
    local_smoothing_radius3=5,
    local_smoothing_sigma_sq1=None,
    local_smoothing_sigma_sq2=0.1,
    local_smoothing_sigma_sq3=None,

):
    """
    Estimate alpha from an input image and an input trimap using Shared Matting as proposed by :cite:`GastalOliveira2010SharedMatting`.

    Parameters
    ----------
    image: numpy.ndarray
        Image with shape :math:`h \\times  w \\times d` for which the alpha matte should be estimated
    trimap: numpy.ndarray
        Trimap with shape :math:`h \\times  w` of the image
    return_foreground_background: numpy.ndarray
        Whether to return foreground and background estimate. They will be computed either way
    trimap_expansion_radius: int
        How much to expand trimap.
    trimap_expansion_threshold: float
        Which pixel colors are similar enough to expand trimap into
    sample_gathering_angles: int
        In how many directions to search for new samples.
    sample_gathering_weights: Tuple[float, float, float, float]
        Weights for various cost functions
    sample_gathering_Np_radius: int
        Radius of Np function
    sample_refinement_radius: int
        Search region for better neighboring samples
    local_smoothing_radius1: int
        Radius for foreground/background smoothing
    local_smoothing_radius2: int
        Radius for confidence computation
    local_smoothing_radius3: int
        Radius for low frequency alpha computation
    local_smoothing_sigma_sq1: float
        Squared sigma value for foreground/background smoothing
        Defaults to :code:`(2 * local_smoothing_radius1 + 1)**2 / (9 * pi)` if not given
    local_smoothing_sigma_sq2: float
        Squared sigma value for confidence computation
    local_smoothing_sigma_sq3: float
        Squared sigma value for low frequency alpha computation
        Defaults to :code:`(2 * local_smoothing_radius3 + 1)**2 / (9 * pi)` if not given

    Returns
    -------
    alpha: numpy.ndarray
        Estimated alpha matte
    foreground: numpy.ndarray
        Estimated foreground
    background: numpy.ndarray
        Estimated background

    Example
    -------
    >>> from pymatting import *
    >>> image = load_image("data/lemur/lemur.png", "RGB")
    >>> trimap = load_image("data/lemur/lemur_trimap.png", "GRAY")
    >>> alpha, foreground, background = estimate_alpha_sm(
    ...     image,
    ...     trimap,
    ...     return_foreground_background=True,
    ...     sample_gathering_angles=4)
    """
    assert image.dtype in [np.float32, np.float64], f"image.dtype should be float32 or float64, but is {image.dtype}"
    assert trimap.dtype in [np.float32, np.float64], f"trimap.dtype should be float32 or float64, but is {trimap.dtype}"
    assert trimap.shape == image.shape[:2], f"image height and width should match trimap height and width, but image shape is {image.shape} and trimap shape is {trimap.shape}"
    assert len(image.shape) == 3 and image.shape[2] == 3, f"image should be RGB, but shape is {image.shape}"

    if local_smoothing_sigma_sq1 is None:
        local_smoothing_sigma_sq1 = (2 * local_smoothing_radius1 + 1)**2 / (9.0 * np.pi)

    if local_smoothing_sigma_sq3 is None:
        local_smoothing_sigma_sq3 = (2 * local_smoothing_radius1 + 1)**2 / (9.0 * np.pi)

    # Convert to float32 if float64. Precision should be sufficient.
    image = image.astype(np.float32)
    trimap = trimap.astype(np.float32)

    expanded_trimap = trimap.copy()

    if trimap_expansion_radius > 0 and trimap_expansion_threshold > 0.0:
        expand_trimap(
            expanded_trimap,
            trimap,
            image,
            k_i=trimap_expansion_radius,
            k_c=trimap_expansion_threshold)

    avg_F = image[expanded_trimap == 1.0].mean(axis=0)
    avg_B = image[expanded_trimap == 0.0].mean(axis=0)

    gathering_F = np.zeros_like(image) + avg_F
    gathering_B = np.zeros_like(image) + avg_B

    gathering_alpha = np.zeros_like(trimap)

    eN, eA, ef, eb = sample_gathering_weights

    sample_gathering(
        gathering_F,
        gathering_B,
        gathering_alpha,
        image,
        expanded_trimap,
        num_angles=sample_gathering_angles,
        eN=eN,
        eA=eA,
        ef=ef,
        eb=eb,
        Np_radius=sample_gathering_Np_radius)

    refined_F = np.zeros_like(image)
    refined_B = np.zeros_like(image)
    refined_alpha = np.zeros_like(trimap)

    sample_refinement(
        refined_F,
        refined_B,
        refined_alpha,
        gathering_F,
        gathering_B,
        image,
        expanded_trimap,
        radius=sample_refinement_radius)

    final_F = np.zeros_like(image)
    final_B = np.zeros_like(image)
    final_alpha = np.zeros_like(trimap)

    local_smoothing(
        final_F,
        final_B,
        final_alpha,
        refined_F,
        refined_B,
        refined_alpha,
        image,
        expanded_trimap,
        radius1=local_smoothing_radius1,
        radius2=local_smoothing_radius2,
        radius3=local_smoothing_radius3,
        sigma_sq1=local_smoothing_sigma_sq1,
        sigma_sq2=local_smoothing_sigma_sq2,
        sigma_sq3=local_smoothing_sigma_sq3)

    if return_foreground_background:
        return final_alpha, final_F, final_B
    else:
        return final_alpha

    Estimate alpha from an input image and an input trimap using Shared Matting as proposed by :cite:`GastalOliveira2010SharedMatting`.

    Parameters
    ----------
    image: numpy.ndarray
        Image with shape :math:`h \times  w \times d` for which the alpha matte should be estimated
    trimap: numpy.ndarray
        Trimap with shape :math:`h \times  w` of the image
    return_foreground_background: numpy.ndarray
        Whether to return foreground and background estimate. They will be computed either way
    trimap_expansion_radius: int
        How much to expand trimap.
    trimap_expansion_threshold: float
        Which pixel colors are similar enough to expand trimap into
    sample_gathering_angles: int
        In how many directions to search for new samples.
    sample_gathering_weights: Tuple[float, float, float, float]
        Weights for various cost functions
    sample_gathering_Np_radius: int
        Radius of Np function
    sample_refinement_radius: int
        Search region for better neighboring samples
    local_smoothing_radius1: int
        Radius for foreground/background smoothing
    local_smoothing_radius2: int
        Radius for confidence computation
    local_smoothing_radius3: int
        Radius for low frequency alpha computation
    local_smoothing_sigma_sq1: float
        Squared sigma value for foreground/background smoothing
        Defaults to :code:`(2 * local_smoothing_radius1 + 1)**2 / (9 * pi)` if not given
    local_smoothing_sigma_sq2: float
        Squared sigma value for confidence computation
    local_smoothing_sigma_sq3: float
        Squared sigma value for low frequency alpha computation
        Defaults to :code:`(2 * local_smoothing_radius3 + 1)**2 / (9 * pi)` if not given

    Returns
    -------
    alpha: numpy.ndarray
        Estimated alpha matte
    foreground: numpy.ndarray
        Estimated foreground
    background: numpy.ndarray
        Estimated background

    Example
    -------
    >>> from pymatting import *
    >>> image = load_image("data/lemur/lemur.png", "RGB")
    >>> trimap = load_image("data/lemur/lemur_trimap.png", "GRAY")
    >>> alpha, foreground, background = estimate_alpha_sm(
    ...     image,
    ...     trimap,
    ...     return_foreground_background=True,
    ...     sample_gathering_angles=4)
    
pymatting/alpha/estimate_alpha_sm.py