utils.py

import os
import time

from collections import OrderedDict
import numpy as np
import pymatting
import cv2
from PIL import Image

SIZES = OrderedDict({
    "1 inch": {
        'physics': (25, 35),
        'pixels': (295, 413)
    },
    "1 inch smaller": {
        'physics': (22, 32),
        'pixels': (260, 378)
    },
    "1 inch larger": {
        'physics': (33, 48),
        'pixels': (390, 567)
    },
    "2 inches": {
        'physics': (35, 49),
        'pixels': (413, 579)
    },
    "2 inches smaller": {
        'physics': (35, 45),
        'pixels': (413, 531)
    },
    "2 inches larger": {
        'physics': (35, 53),
        'pixels': (413, 626)
    },
    "3 inches": {
        'physics': (55, 84),
        'pixels': (649, 991)
    },
    "4 inches": {
        'physics': (76, 102),
        'pixels': (898, 1205)
    },
    "5 inches": {
        'physics': (89, 127),
        'pixels': (1050, 1500)
    }
})

# R, G, B
COLOR_MAP = {
    'White': [255, 255, 255],
    'Blue': [0, 191, 243],
    'Red': [255, 0, 0]
}

# jpg compress ratio
SAVE_SIZE = {'Small': 50, 'Middle': 75, 'Large': 95}


def delete_result():
    """clear old result in `.temp`"""
    root = '.temp'
    results = sorted(os.listdir(root))
    for res in results:
        if int(time.time()) - int(os.path.splitext(res)[0]) > 10000:
            os.remove(os.path.join(root, res))


def clear_all():
    delete_result()
    return None, None, size_play()[0], 'White', 'Large', None


def size_play():
    sizes = []
    for k, v in SIZES.items():
        size = ''.join([
            k, '(', str(v['physics'][0]), 'x', str(v['physics'][1]), 'mm,',
            str(v['pixels'][0]), 'x', str(v['pixels'][1]), 'px)'
        ])
        sizes.append(size)
    return sizes


def bg_replace(img, alpha, bg_name):
    bg = COLOR_MAP[bg_name]
    bg = np.array(bg)[None, None, :]
    alpha = alpha / 255.
    pymatting.estimate_foreground_ml(img / 255., alpha) * 255
    alpha = alpha[:, :, None]
    res = alpha * img + (1 - alpha) * bg
    return res.astype('uint8')


def adjust_size(img, size_index):
    key = list(SIZES.keys())[size_index]
    w_o, h_o = SIZES[key]['pixels']

    # scale
    h_ori, w_ori = img.shape[:2]
    scale = max(w_o / w_ori, h_o / h_ori)
    if scale > 1:
        interpolation = cv2.INTER_CUBIC
    else:
        interpolation = cv2.INTER_AREA
    img_scale = cv2.resize(
        img, dsize=None, fx=scale, fy=scale, interpolation=interpolation)

    # crop
    h_scale, w_scale = img_scale.shape[:2]
    h_cen = h_scale // 2
    w_cen = w_scale // 2
    h_start = max(0, h_cen - h_o // 2)
    h_end = min(h_scale, h_start + h_o)
    w_start = max(0, w_cen - w_o // 2)
    w_end = min(w_scale, w_start + w_o)
    img_c = img_scale[h_start:h_end, w_start:w_end]

    return img_c


def download(img, size):
    q = SAVE_SIZE[size]
    while True:
        name = str(int(time.time()))
        tmp_name = './.temp/' + name + '.jpg'
        if not os.path.exists(tmp_name):
            break
        else:
            time.sleep(1)
    dir_name = os.path.dirname(tmp_name)
    if not os.path.exists(dir_name):
        os.makedirs(dir_name)

    im = Image.fromarray(img)
    im.save(tmp_name, 'jpeg', quality=q, dpi=(300, 300))
    return tmp_name