MCP-1st-Birthday/Free-View_Expressive_Talking_Head_Video_Editing

1from __future__ import print_function2import os3import sys4import time5import torch6import math7import numpy as np8import cv29 10 11def _gaussian(12        size=3, sigma=0.25, amplitude=1, normalize=False, width=None,13        height=None, sigma_horz=None, sigma_vert=None, mean_horz=0.5,14        mean_vert=0.5):15    # handle some defaults16    if width is None:17        width = size18    if height is None:19        height = size20    if sigma_horz is None:21        sigma_horz = sigma22    if sigma_vert is None:23        sigma_vert = sigma24    center_x = mean_horz * width + 0.525    center_y = mean_vert * height + 0.526    gauss = np.empty((height, width), dtype=np.float32)27    # generate kernel28    for i in range(height):29        for j in range(width):30            gauss[i][j] = amplitude * math.exp(-(math.pow((j + 1 - center_x) / (31                sigma_horz * width), 2) / 2.0 + math.pow((i + 1 - center_y) / (sigma_vert * height), 2) / 2.0))32    if normalize:33        gauss = gauss / np.sum(gauss)34    return gauss35 36 37def draw_gaussian(image, point, sigma):38    # Check if the gaussian is inside39    ul = [math.floor(point[0] - 3 * sigma), math.floor(point[1] - 3 * sigma)]40    br = [math.floor(point[0] + 3 * sigma), math.floor(point[1] + 3 * sigma)]41    if (ul[0] > image.shape[1] or ul[1] > image.shape[0] or br[0] < 1 or br[1] < 1):42        return image43    size = 6 * sigma + 144    g = _gaussian(size)45    g_x = [int(max(1, -ul[0])), int(min(br[0], image.shape[1])) - int(max(1, ul[0])) + int(max(1, -ul[0]))]46    g_y = [int(max(1, -ul[1])), int(min(br[1], image.shape[0])) - int(max(1, ul[1])) + int(max(1, -ul[1]))]47    img_x = [int(max(1, ul[0])), int(min(br[0], image.shape[1]))]48    img_y = [int(max(1, ul[1])), int(min(br[1], image.shape[0]))]49    assert (g_x[0] > 0 and g_y[1] > 0)50    image[img_y[0] - 1:img_y[1], img_x[0] - 1:img_x[1]51          ] = image[img_y[0] - 1:img_y[1], img_x[0] - 1:img_x[1]] + g[g_y[0] - 1:g_y[1], g_x[0] - 1:g_x[1]]52    image[image > 1] = 153    return image54 55 56def transform(point, center, scale, resolution, invert=False):57    """Generate and affine transformation matrix.58 59    Given a set of points, a center, a scale and a targer resolution, the60    function generates and affine transformation matrix. If invert is ``True``61    it will produce the inverse transformation.62 63    Arguments:64        point {torch.tensor} -- the input 2D point65        center {torch.tensor or numpy.array} -- the center around which to perform the transformations66        scale {float} -- the scale of the face/object67        resolution {float} -- the output resolution68 69    Keyword Arguments:70        invert {bool} -- define wherever the function should produce the direct or the71        inverse transformation matrix (default: {False})72    """73    _pt = torch.ones(3)74    _pt[0] = point[0]75    _pt[1] = point[1]76 77    h = 200.0 * scale78    t = torch.eye(3)79    t[0, 0] = resolution / h80    t[1, 1] = resolution / h81    t[0, 2] = resolution * (-center[0] / h + 0.5)82    t[1, 2] = resolution * (-center[1] / h + 0.5)83 84    if invert:85        t = torch.inverse(t)86 87    new_point = (torch.matmul(t, _pt))[0:2]88 89    return new_point.int()90 91 92def crop(image, center, scale, resolution=256.0):93    """Center crops an image or set of heatmaps94 95    Arguments:96        image {numpy.array} -- an rgb image97        center {numpy.array} -- the center of the object, usually the same as of the bounding box98        scale {float} -- scale of the face99 100    Keyword Arguments:101        resolution {float} -- the size of the output cropped image (default: {256.0})102 103    Returns:104        [type] -- [description]105    """  # Crop around the center point106    """ Crops the image around the center. Input is expected to be an np.ndarray """107    ul = transform([1, 1], center, scale, resolution, True)108    br = transform([resolution, resolution], center, scale, resolution, True)109    # pad = math.ceil(torch.norm((ul - br).float()) / 2.0 - (br[0] - ul[0]) / 2.0)110    if image.ndim > 2:111        newDim = np.array([br[1] - ul[1], br[0] - ul[0],112                           image.shape[2]], dtype=np.int32)113        newImg = np.zeros(newDim, dtype=np.uint8)114    else:115        newDim = np.array([br[1] - ul[1], br[0] - ul[0]], dtype=np.int)116        newImg = np.zeros(newDim, dtype=np.uint8)117    ht = image.shape[0]118    wd = image.shape[1]119    newX = np.array(120        [max(1, -ul[0] + 1), min(br[0], wd) - ul[0]], dtype=np.int32)121    newY = np.array(122        [max(1, -ul[1] + 1), min(br[1], ht) - ul[1]], dtype=np.int32)123    oldX = np.array([max(1, ul[0] + 1), min(br[0], wd)], dtype=np.int32)124    oldY = np.array([max(1, ul[1] + 1), min(br[1], ht)], dtype=np.int32)125    newImg[newY[0] - 1:newY[1], newX[0] - 1:newX[1]126           ] = image[oldY[0] - 1:oldY[1], oldX[0] - 1:oldX[1], :]127    newImg = cv2.resize(newImg, dsize=(int(resolution), int(resolution)),128                        interpolation=cv2.INTER_LINEAR)129    return newImg130 131 132def get_preds_fromhm(hm, center=None, scale=None):133    """Obtain (x,y) coordinates given a set of N heatmaps. If the center134    and the scale is provided the function will return the points also in135    the original coordinate frame.136 137    Arguments:138        hm {torch.tensor} -- the predicted heatmaps, of shape [B, N, W, H]139 140    Keyword Arguments:141        center {torch.tensor} -- the center of the bounding box (default: {None})142        scale {float} -- face scale (default: {None})143    """144    max, idx = torch.max(145        hm.view(hm.size(0), hm.size(1), hm.size(2) * hm.size(3)), 2)146    idx += 1147    preds = idx.view(idx.size(0), idx.size(1), 1).repeat(1, 1, 2).float()148    preds[..., 0].apply_(lambda x: (x - 1) % hm.size(3) + 1)149    preds[..., 1].add_(-1).div_(hm.size(2)).floor_().add_(1)150 151    for i in range(preds.size(0)):152        for j in range(preds.size(1)):153            hm_ = hm[i, j, :]154            pX, pY = int(preds[i, j, 0]) - 1, int(preds[i, j, 1]) - 1155            if pX > 0 and pX < 63 and pY > 0 and pY < 63:156                diff = torch.FloatTensor(157                    [hm_[pY, pX + 1] - hm_[pY, pX - 1],158                     hm_[pY + 1, pX] - hm_[pY - 1, pX]])159                preds[i, j].add_(diff.sign_().mul_(.25))160 161    preds.add_(-.5)162 163    preds_orig = torch.zeros(preds.size())164    if center is not None and scale is not None:165        for i in range(hm.size(0)):166            for j in range(hm.size(1)):167                preds_orig[i, j] = transform(168                    preds[i, j], center, scale, hm.size(2), True)169 170    return preds, preds_orig171 172def get_preds_fromhm_batch(hm, centers=None, scales=None):173    """Obtain (x,y) coordinates given a set of N heatmaps. If the centers174    and the scales is provided the function will return the points also in175    the original coordinate frame.176 177    Arguments:178        hm {torch.tensor} -- the predicted heatmaps, of shape [B, N, W, H]179 180    Keyword Arguments:181        centers {torch.tensor} -- the centers of the bounding box (default: {None})182        scales {float} -- face scales (default: {None})183    """184    max, idx = torch.max(185        hm.view(hm.size(0), hm.size(1), hm.size(2) * hm.size(3)), 2)186    idx += 1187    preds = idx.view(idx.size(0), idx.size(1), 1).repeat(1, 1, 2).float()188    preds[..., 0].apply_(lambda x: (x - 1) % hm.size(3) + 1)189    preds[..., 1].add_(-1).div_(hm.size(2)).floor_().add_(1)190 191    for i in range(preds.size(0)):192        for j in range(preds.size(1)):193            hm_ = hm[i, j, :]194            pX, pY = int(preds[i, j, 0]) - 1, int(preds[i, j, 1]) - 1195            if pX > 0 and pX < 63 and pY > 0 and pY < 63:196                diff = torch.FloatTensor(197                    [hm_[pY, pX + 1] - hm_[pY, pX - 1],198                     hm_[pY + 1, pX] - hm_[pY - 1, pX]])199                preds[i, j].add_(diff.sign_().mul_(.25))200 201    preds.add_(-.5)202 203    preds_orig = torch.zeros(preds.size())204    if centers is not None and scales is not None:205        for i in range(hm.size(0)):206            for j in range(hm.size(1)):207                preds_orig[i, j] = transform(208                    preds[i, j], centers[i], scales[i], hm.size(2), True)209 210    return preds, preds_orig211 212def shuffle_lr(parts, pairs=None):213    """Shuffle the points left-right according to the axis of symmetry214    of the object.215 216    Arguments:217        parts {torch.tensor} -- a 3D or 4D object containing the218        heatmaps.219 220    Keyword Arguments:221        pairs {list of integers} -- [order of the flipped points] (default: {None})222    """223    if pairs is None:224        pairs = [16, 15, 14, 13, 12, 11, 10, 9, 8, 7, 6, 5, 4, 3, 2, 1, 0,225                 26, 25, 24, 23, 22, 21, 20, 19, 18, 17, 27, 28, 29, 30, 35,226                 34, 33, 32, 31, 45, 44, 43, 42, 47, 46, 39, 38, 37, 36, 41,227                 40, 54, 53, 52, 51, 50, 49, 48, 59, 58, 57, 56, 55, 64, 63,228                 62, 61, 60, 67, 66, 65]229    if parts.ndimension() == 3:230        parts = parts[pairs, ...]231    else:232        parts = parts[:, pairs, ...]233 234    return parts235 236 237def flip(tensor, is_label=False):238    """Flip an image or a set of heatmaps left-right239 240    Arguments:241        tensor {numpy.array or torch.tensor} -- [the input image or heatmaps]242 243    Keyword Arguments:244        is_label {bool} -- [denote wherever the input is an image or a set of heatmaps ] (default: {False})245    """246    if not torch.is_tensor(tensor):247        tensor = torch.from_numpy(tensor)248 249    if is_label:250        tensor = shuffle_lr(tensor).flip(tensor.ndimension() - 1)251    else:252        tensor = tensor.flip(tensor.ndimension() - 1)253 254    return tensor255 256# From pyzolib/paths.py (https://bitbucket.org/pyzo/pyzolib/src/tip/paths.py)257 258 259def appdata_dir(appname=None, roaming=False):260    """ appdata_dir(appname=None, roaming=False)261 262    Get the path to the application directory, where applications are allowed263    to write user specific files (e.g. configurations). For non-user specific264    data, consider using common_appdata_dir().265    If appname is given, a subdir is appended (and created if necessary).266    If roaming is True, will prefer a roaming directory (Windows Vista/7).267    """268 269    # Define default user directory270    userDir = os.getenv('FACEALIGNMENT_USERDIR', None)271    if userDir is None:272        userDir = os.path.expanduser('~')273        if not os.path.isdir(userDir):  # pragma: no cover274            userDir = '/var/tmp'  # issue #54275 276    # Get system app data dir277    path = None278    if sys.platform.startswith('win'):279        path1, path2 = os.getenv('LOCALAPPDATA'), os.getenv('APPDATA')280        path = (path2 or path1) if roaming else (path1 or path2)281    elif sys.platform.startswith('darwin'):282        path = os.path.join(userDir, 'Library', 'Application Support')283    # On Linux and as fallback284    if not (path and os.path.isdir(path)):285        path = userDir286 287    # Maybe we should store things local to the executable (in case of a288    # portable distro or a frozen application that wants to be portable)289    prefix = sys.prefix290    if getattr(sys, 'frozen', None):291        prefix = os.path.abspath(os.path.dirname(sys.executable))292    for reldir in ('settings', '../settings'):293        localpath = os.path.abspath(os.path.join(prefix, reldir))294        if os.path.isdir(localpath):  # pragma: no cover295            try:296                open(os.path.join(localpath, 'test.write'), 'wb').close()297                os.remove(os.path.join(localpath, 'test.write'))298            except IOError:299                pass  # We cannot write in this directory300            else:301                path = localpath302                break303 304    # Get path specific for this app305    if appname:306        if path == userDir:307            appname = '.' + appname.lstrip('.')  # Make it a hidden directory308        path = os.path.join(path, appname)309        if not os.path.isdir(path):  # pragma: no cover310            os.mkdir(path)311 312    # Done313    return path314