# AUTOGENERATED! DO NOT EDIT! File to edit: ../../nbs/14g_segmask.vis.ipynb.

# %% auto 0
__all__ = ['overlay_array', 'overlay']

# %% ../../nbs/14g_segmask.vis.ipynb 3
from fastcore.all import *
import PIL
import numpy as np
import polvo as pv
import polvo.segmask as ps
import matplotlib as mpl

# %% ../../nbs/14g_segmask.vis.ipynb 4
# def overlay(
#     image:PIL.Image,
#     mask:ps.SegMask,
# #     nclasses:int,
#     fill_background:bool=True,
#     cmap:mpl.colors.Colormap=mpl.colormaps['gist_ncar'],
#     alpha:float=0.5,
# ):
#     mask_arr = np.array(mask.mask)
#     if not len(mask_arr.shape) == 2: raise ValueError(f'Mask bust be 2d, but got {mask.shape}')
    
#     colored_mask = np.zeros((*mask_arr.shape, 3), dtype=np.uint8)
#     for class_idx in np.unique(mask_arr):
#         if not fill_background and mask.class_map.id2name[class_idx] == 'background': continue
#         mask_idxs = mask_arr == class_idx
#         colored_mask[mask_idxs] = np.array(cmap(class_idx/len(mask.class_map))[:3]) * 255

#     return PIL.Image.blend(image, PIL.Image.fromarray(colored_mask), alpha)

# %% ../../nbs/14g_segmask.vis.ipynb 5
def overlay_array(
    image:PIL.Image,
    mask:Union[PIL.Image.Image, np.ndarray],
    nclasses:int,
    ignore_idxs:Sequence=None,
    cmap:mpl.colors.Colormap=mpl.colormaps['gist_ncar'],
    alpha:float=0.5,
):
    mask = np.asarray(mask)
    ignore_idxs = ignore_idxs or []
    if not len(mask.shape) == 2: raise ValueError(f'Mask bust be 2d, but got {mask.shape}')
    
    colored_mask = np.zeros((*mask.shape, 3), dtype=np.uint8)
    for class_idx in np.unique(mask):
        if class_idx in ignore_idxs: continue
        mask_idxs = mask == class_idx
        colored_mask[mask_idxs] = np.array(cmap(class_idx/nclasses)[:3]) * 255

    return PIL.Image.blend(image, PIL.Image.fromarray(colored_mask), alpha)

# %% ../../nbs/14g_segmask.vis.ipynb 7
def overlay(
    image:PIL.Image,
    mask:ps.SegMask,
    fill_background:bool=True,
    ignore_idxs=None,
    cmap:mpl.colors.Colormap=mpl.colormaps['gist_ncar'],
    alpha:float=0.5,
):
    ignore_idxs = ignore_idxs or []
    if not fill_background: ignore_idxs.append(mask.class_map.name2id['background'])
    return overlay_array(image=image, mask=mask.mask, nclasses=len(mask.class_map), ignore_idxs=ignore_idxs, cmap=cmap, alpha=alpha)
