Python helpers

image_hlp.py

import io
import logging
import math

import cv2
import numpy as np
import matplotlib.pyplot as plt

from tqdm import tqdm

plt.style.use('ggplot')
plt.rcParams["figure.figsize"] = (8, 8)

import seaborn as sns
sns.set(style="darkgrid", color_codes=True)
sns.set_context("notebook", font_scale=1.1)

def load_image(file, flags=cv2.IMREAD_COLOR):
    if type(file) is io.BytesIO:
        a = np.frombuffer(buffer.getvalue(), np.uint8)
        img = cv2.imdecode(a, flags)
    elif isinstance(file, io.IOBase):
        a = np.fromfile(file, np.uint8)
        img = cv2.imdecode(a, flags)
    elif type(file) is bytes:
        a = np.frombuffer(file, np.uint8)
        img = cv2.imdecode(a, flags)
    else:
        img = cv2.imread(file, flags)
    if len(img.shape) == 3:
        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)    
    return img


def save_image(filename, image):
    if image.max() < 200:
        image = image * 255
    cv_img = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)  
    cv2.imwrite(filename, cv_img)

    
def plot_image(img, figsize=(15,15), cmap=None, norm=None, aspect=None):
    plt.figure(figsize=figsize)
    plt.imshow(img, cmap=cmap, norm=norm, aspect=aspect)
    plt.grid(False) 

    
def plot_images(images, figsize=(15,15), cmap=None, norm=None, aspect=None):
    plt.figure(figsize=figsize)

    r = math.ceil(math.sqrt(len(images)))
    c = math.ceil(len(images) / r)

    for i, img in tqdm(enumerate(images)):
        ax = plt.subplot(c, r, i + 1)
        ax.grid(False)
        ax.set_xticks([])
        ax.set_yticks([])  
        plt.imshow(img, cmap=cmap, norm=norm, aspect=aspect)
    plt.subplots_adjust(wspace=0, hspace=0)
    plt.show()
    
    
def hex_to_rgb(v):
    return np.array([int(v[i:i+2], 16) for i in range(0, len(v), 2)], dtype=np.uint8)


def rgb_to_hex(v):
    return ''.join(map(lambda c: hex(c)[2:], v))


def rgb_to_int(color):
    return (color[0].astype(np.int) << 16) + (color[1].astype(np.int) << 8) + color[2].astype(np.int)


def int_to_rgb(value):
    return np.array([(value >> 16) & 0xFF, (value >> 8) & 0xFF, (value) & 0xFF]).astype(np.uint8)


def most_frequent_colors(img, count=None):
    colors = rgb_to_int(img.reshape(-1, 3).transpose())
    counts = np.bincount(colors)
    index = np.nonzero(counts)
    
    colors_counts = np.vstack((index, counts[index])).T
    colors_counts = colors_counts[colors_counts[:,1].argsort()][::-1]
    colors_counts = colors_counts[:count]

    return (
        colors_counts[:,1].T.reshape(-1,1),
        int_to_rgb(colors_counts[:,0]).T
    )


def alpha_composition(src, dst, alpha):
    return src * alpha + dst * (1 - alpha)

def cv_rectangle(img, corner, size, color, thickness=1, bg_color=None):
    '''
        draw filled rectangle with supporting alpha channel
    '''
    x, y = corner
    w, h = size
    
    x1, y1 = max(x, 0), max(y, 0)
    x2, y2 = min(x + w, img.shape[1]), min(y + h, img.shape[0])
    
    if bg_color is not None:
        if bg_color.shape[0] == 4:
            alpha = bg_color[-1] / 255
            src = bg_color[:-1]
            dst = img[y1:y2,x1:x2,:]        
            img[y1:y2,x1:x2,:] = alpha_composition(src, dst, alpha).astype(dtype=np.uint8)
        else:
            img[y1:y2,x1:x2,:] = bg_color

    cv2.rectangle(img, (x1,y1), (x2, y2), color.tolist(), thickness)

    cv2.rectangle(img, (x1,y1), (x2, y2), color.tolist(), thickness)