# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/01_vision.ipynb.

# %% auto 0
__all__ = ['image_extensions', 'remove_images', 'bgr2rgb', 'rgb2bgr', 'gray2rgb', 'rgb2gray', 'rgb2rgba', 'bgra2rgb', 'rgba2rgb',
           'rgb_read', 'c1_read', 'tensor_to_img', 'plt_show', 'get_hw', 'get_image_files', 'to_pil', 'add_alpha',
           'has_alpha', 'color_to_rgb', 'solid_color_img', 'solid_color_img_like', 'get_pos_factors', 'get_pos_',
           'get_pos', 'paste_img']

# %% ../nbs/01_vision.ipynb 3
from .imports import *
from .core import *

# %% ../nbs/01_vision.ipynb 4
image_extensions = {'.art','.bmp','.cdr','.cdt','.cpt','.cr2','.crw','.djv','.djvu','.erf','.gif','.ico',
                    '.ief','.jng','.jp2','.jpe','.jpeg','.jpf','.jpg','.jpg2','.jpm','.jpx','.nef','.orf',
                    '.pat','.pbm','.pcx','.pgm','.png','.pnm','.ppm','.psd','.ras','.rgb','.svg','.svgz',
                    '.tif','.tiff','.wbmp','.xbm','.xpm','.xwd','.webp'}

def remove_images(images_path, final_count=0, fn=lambda x:True):
    imgs = get_image_files(images_path)
    remove = max(0, len(imgs) - final_count)
    removed = 0
    for i in imgs:
        if removed >= remove:
            break
        if fn(i):
            os.remove(i)
            removed+=1

def bgr2rgb(img):
    return cv2.cvtColor(img,cv2.COLOR_BGR2RGB)

def rgb2bgr(img):
    return cv2.cvtColor(img,cv2.COLOR_RGB2BGR)

def gray2rgb(img):
    return cv2.cvtColor(img,cv2.COLOR_GRAY2RGB)

def rgb2gray(img):
    return cv2.cvtColor(img,cv2.COLOR_RGB2GRAY)

def rgb2rgba(img):
    return cv2.cvtColor(img,cv2.COLOR_RGB2RGBA)

def bgra2rgb(img):
    if len(img.shape) > 2 and img.shape[2] == 4:
        return cv2.cvtColor(img,cv2.COLOR_BGRA2RGB)

def rgba2rgb(img):
    if len(img.shape) > 2 and img.shape[2] == 4:
        return cv2.cvtColor(img, cv2.COLOR_BGRA2BGR)

def rgb_read(img, shape=None):
    if not path_or_str(img):
        return img
    img = bgr2rgb(cv2.imread(str(img)))
    if shape is not None:
        img = cv2.resize(img, (shape[1], shape[0]))
    return img

def c1_read(img):
    "Read an image from `img` and convert it to grayscale."
    return cv2.imread(str(img), 0)

def tensor_to_img(t):
    if t.dim() > 3:
        return [np.array(np.transpose(t_,(1,2,0))) for t_ in t]
    return np.array(np.transpose(t,(1,2,0)))

def plt_show(im, cmap=None, title='', figsize=(7,7)):
    if path_or_str(im):
        im = rgb_read(im)
    if is_tensor(im):
        im = tensor_to_img(im)
        if is_list(im): im = im[0]
    fig=plt.figure(figsize=figsize)
    plt.imshow(im, cmap=cmap)
    plt.title(title)
    plt.show()

def get_hw(x):
    "Return the height and width of `x`."
    if is_clip(x):
        return x.size[1], x.size[0]
    return np.array(x).shape[:2]

def get_image_files(path, recurse=True, folders=None,
                    map_fn=None, # A function to apply to each item.
                    sort_key=None, # A function to sort the items by.
                    reverse=False, shuffle=False):
    "Get image files in `path` recursively, only in `folders`, if specified."
    l = get_files(path, extensions=image_extensions, recurse=recurse, folders=folders)
    if sort_key is not None:
        l = sorted(l, key=sort_key)
    if reverse:
        l = list(l)[::-1]
    if shuffle:
        random.shuffle(l)
    if map_fn is not None:
        return list_map(l, map_fn)
    return l

def to_pil(x):
    if is_array(x):
        return Image.fromarray(x)
    return x

def add_alpha(img, alpha=None):
    "Add an alpha channel to `img`."
    if alpha is None:
        alpha = 255
    img = to_pil(img)
    img.putalpha(alpha)
    return np.array(img)

def has_alpha(img):
    return np.array(img).shape[-1] == 4

def color_to_rgb(color):
    "Convert str `color` to RGB."
    if type(color) == str:
        return list_map(np.ceil(colors.to_rgb(color)).astype(int)*255, int)
    return list_map(color, int)

def solid_color_img(shape=(300,300,3), color='black', alpha=None):
    if len(shape) < 3:
        shape = tuple(shape) + (3,)
    image = np.zeros(shape, np.uint8)
    color = color_to_rgb(color)
    image[:] = color[:3]
    if len(color) == 4:
        alpha = color[3]
    if alpha is not None:
        image = Image.fromarray(image)
        image.putalpha(alpha)
        image = np.array(image)
    return image

def solid_color_img_like(img, color='black', alpha=None):
    h,w = np.array(img).shape[:2]
    shape = (h,w,3)
    return solid_color_img(shape=shape, color=color, alpha=alpha)

def get_pos_factors(pos):
    pos_dict = {'center':0.5, 'left':0.15, 'right':0.85, 'top':0.15, 'bottom':0.85}
    x_factor, y_factor = [pos_dict.get(p,p) for p in pos]
    return x_factor, y_factor

def get_pos_(bh, bw, h, w, x_factor, y_factor):
    x = int((bw - w) * x_factor)
    y = int((bh - h) * y_factor)
    return x,y

def get_pos(img, bg,
            pos=('center', 'center'), # int for exact position, float for relative position, or str for position in ['center', 'left', 'right', 'top', 'bottom']
            relative=True):
    "Get the position (x,y) of `img` on `bg` with `pos` as reference."
    p0,p1 = copy.deepcopy(pos)
    if is_int(p0) and is_int(p1):
        relative = False
    if not relative:
        return pos
    x_factor, y_factor = get_pos_factors(pos)
    bg = np.array(bg)
    img = np.array(img)
    bh,bw = bg.shape[:2]
    try:
        h,w = img.shape[:2]
    except:
        h,w = img[:2]
    x,y = get_pos_(bh, bw, h, w, x_factor, y_factor)
    if is_int(p0):
        x = p0
    if is_int(p1):
        y = p1
        # print(x, y)
    return x,y

def paste_img(img, bg, pos=('center', 'center'), relative=True):
    "Paste `img` on `bg` at `pos`."
    bg = np.array(copy.deepcopy(bg))
    img = np.array(copy.deepcopy(img))
    if bg.shape[-1] != 4:
        bg = add_alpha(bg)
    if img.shape[-1] != 4:
        img = add_alpha(img)
    h,w = img.shape[:2]
    bgh, bgw = bg.shape[:2]
    x,y = get_pos(img, bg, pos, relative=relative)
    
    r1,r2 = y,y+h
    c1,c2 = x,x+w
   
    if x < 0 or y < 0 or x >= bgw or y >= bgh:
        # print(x,y,h,w,bgh,bgw)
        bgh2 = max(bgh,h)
        bgw = max(bgw,w)
        px, py = 0.,0.
        if x < 0:
            bgw2 = bgw + abs(x)
            x = 0
            px = 1.
        elif x >= bgw:
            bgw2 = x+w
        else:
            bgw2 = bgw
        if y < 0:
            bgh2 = bgh + abs(y)
            y = 0
            py = 1.
        elif y >= bgh:
            bgh2 = y+h
        else:
            bgh2 = bgh
        new_bg = solid_color_img((bgh2, bgw2, 3), alpha=0)
        bg = paste_img(bg, new_bg, pos=(px,py))
        bgh, bgw = bg.shape[:2]
        r1,r2 = y,y+h
        c1,c2 = x,x+w
    
    elif h>bgh and w>bgw:
        hf = (h-bgh)//2
        wf = (w-bgw)//2
        img = img[hf:hf+bgh, wf:wf+bgw]
        r1,r2 = 0,bgh
        c1,c2 = 0,bgw
    elif h>bgh:
        hf = (h-bgh)//2
        img = img[hf:hf+bgh]
        r1,r2 = 0,bgh
    elif w>bgw:
        wf = (w-bgw)//2
        img = img[:,wf:wf+bgw]
        c1,c2 = 0,bgw

    bg2 = copy.deepcopy(bg[r1:r2, c1:c2])
    bg2h, bg2w = get_hw(bg2)
    img = img[:bg2h, :bg2w]
    # print(bg.shape, bg2.shape, img.shape, r2-r1, c2-c1)
    bgh, bgw = bg2.shape[:2]
    bg_rgb = bg2[...,:3]
    img_rgb = img[...,:3]
    bg_alpha = bg2[...,3]/255.0
    img_alpha = img[...,3]/255.0
    # print(img_alpha.shape)
    if (bg_alpha == 1).all():
        bg_alpha = 1-img_alpha
    i_len = len(img_alpha)
    for i in range(i_len):
        j_len = len(img_alpha[i])
        for j in range(j_len):
            if img_alpha[i][j] == 1:
                # try:
                bg_alpha[i][j] = 0
                # except:
                    # pass
    out_alpha = bg_alpha + img_alpha*(1-bg_alpha)
    out_rgb = (bg_rgb*bg_alpha[...,np.newaxis] + img_rgb*img_alpha[...,np.newaxis]*(1-bg_alpha[...,np.newaxis])) / out_alpha[...,np.newaxis]
    out_rgba = np.dstack((out_rgb,out_alpha*255)).astype(np.uint8)
    bg[r1:r2,c1:c2] = out_rgba
    return bg
