img_helper.py

import numpy as np
import math
import cv2
from skimage import transform as stf


def transform(data, center, output_size, scale, rotation):
    scale_ratio = float(output_size) / scale
    rot = float(rotation) * np.pi / 180.0
    #translation = (output_size/2-center[0]*scale_ratio, output_size/2-center[1]*scale_ratio)
    t1 = stf.SimilarityTransform(scale=scale_ratio)
    cx = center[0] * scale_ratio
    cy = center[1] * scale_ratio
    t2 = stf.SimilarityTransform(translation=(-1 * cx, -1 * cy))
    t3 = stf.SimilarityTransform(rotation=rot)
    t4 = stf.SimilarityTransform(translation=(output_size / 2,
                                              output_size / 2))
    t = t1 + t2 + t3 + t4
    trans = t.params[0:2]
    #print('M', scale, rotation, trans)
    cropped = cv2.warpAffine(data,
                             trans, (output_size, output_size),
                             borderValue=0.0)
    return cropped, trans


def transform_pt(pt, trans):
    new_pt = np.array([pt[0], pt[1], 1.]).T
    new_pt = np.dot(trans, new_pt)
    #print('new_pt', new_pt.shape, new_pt)
    return new_pt[:2]


def gaussian(img, pt, sigma):
    # Draw a 2D gaussian
    assert (sigma >= 0)
    if sigma == 0:
        img[pt[1], pt[0]] = 1.0
        return True
    #assert pt[0]<=img.shape[1]
    #assert pt[1]<=img.shape[0]

    # Check that any part of the gaussian is in-bounds
    ul = [int(pt[0] - 3 * sigma), int(pt[1] - 3 * sigma)]
    br = [int(pt[0] + 3 * sigma + 1), int(pt[1] + 3 * sigma + 1)]
    if (ul[0] > img.shape[1] or ul[1] >= img.shape[0] or br[0] < 0
            or br[1] < 0):
        # If not, just return the image as is
        #print('gaussian error')
        return False
        #return img

    # Generate gaussian
    size = 6 * sigma + 1
    x = np.arange(0, size, 1, float)
    y = x[:, np.newaxis]
    x0 = y0 = size // 2
    # The gaussian is not normalized, we want the center value to equal 1
    g = np.exp(-((x - x0)**2 + (y - y0)**2) / (2 * sigma**2))

    # Usable gaussian range
    g_x = max(0, -ul[0]), min(br[0], img.shape[1]) - ul[0]
    g_y = max(0, -ul[1]), min(br[1], img.shape[0]) - ul[1]
    # Image range
    img_x = max(0, ul[0]), min(br[0], img.shape[1])
    img_y = max(0, ul[1]), min(br[1], img.shape[0])

    img[img_y[0]:img_y[1], img_x[0]:img_x[1]] = g[g_y[0]:g_y[1], g_x[0]:g_x[1]]
    return True
    #return img


def estimate_trans_bbox(face, input_size, s=2.0):
    w = face[2] - face[0]
    h = face[3] - face[1]
    wc = int((face[2] + face[0]) / 2)
    hc = int((face[3] + face[1]) / 2)
    im_size = max(w, h)
    #size = int(im_size*1.2)
    scale = input_size / (max(w, h) * s)
    M = [
        [scale, 0, input_size / 2 - wc * scale],
        [0, scale, input_size / 2 - hc * scale],
    ]
    M = np.array(M)
    return M