Spaces:
Runtime error
Runtime error
File size: 4,511 Bytes
753fd9a |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 |
# Modified from:
# https://github.com/anibali/pytorch-stacked-hourglass
# https://github.com/bearpaw/pytorch-pose
import numpy as np
import torch
from .imutils import im_to_numpy, im_to_torch
from .misc import to_torch
from .pilutil import imresize, imrotate
def color_normalize(x, mean, std):
if x.size(0) == 1:
x = x.repeat(3, 1, 1)
for t, m, s in zip(x, mean, std):
t.sub_(m)
return x
def flip_back(flip_output, hflip_indices):
"""flip and rearrange output maps"""
return fliplr(flip_output)[:, hflip_indices]
def shufflelr(x, width, hflip_indices):
"""flip and rearrange coords"""
# Flip horizontal
x[:, 0] = width - x[:, 0]
# Change left-right parts
x = x[hflip_indices]
return x
def fliplr(x):
"""Flip images horizontally."""
if torch.is_tensor(x):
return torch.flip(x, [-1])
else:
return np.ascontiguousarray(np.flip(x, -1))
def get_transform(center, scale, res, rot=0):
"""
General image processing functions
"""
# Generate transformation matrix
h = 200 * scale
t = np.zeros((3, 3))
t[0, 0] = float(res[1]) / h
t[1, 1] = float(res[0]) / h
t[0, 2] = res[1] * (-float(center[0]) / h + .5)
t[1, 2] = res[0] * (-float(center[1]) / h + .5)
t[2, 2] = 1
if not rot == 0:
rot = -rot # To match direction of rotation from cropping
rot_mat = np.zeros((3,3))
rot_rad = rot * np.pi / 180
sn,cs = np.sin(rot_rad), np.cos(rot_rad)
rot_mat[0,:2] = [cs, -sn]
rot_mat[1,:2] = [sn, cs]
rot_mat[2,2] = 1
# Need to rotate around center
t_mat = np.eye(3)
t_mat[0,2] = -res[1]/2
t_mat[1,2] = -res[0]/2
t_inv = t_mat.copy()
t_inv[:2,2] *= -1
t = np.dot(t_inv,np.dot(rot_mat,np.dot(t_mat,t)))
return t
def transform(pt, center, scale, res, invert=0, rot=0, as_int=True):
# Transform pixel location to different reference
t = get_transform(center, scale, res, rot=rot)
if invert:
t = np.linalg.inv(t)
new_pt = np.array([pt[0] - 1, pt[1] - 1, 1.]).T
new_pt = np.dot(t, new_pt)
if as_int:
return new_pt[:2].astype(int) + 1
else:
return new_pt[:2] + 1
def transform_preds(coords, center, scale, res):
# size = coords.size()
# coords = coords.view(-1, coords.size(-1))
# print(coords.size())
for p in range(coords.size(0)):
coords[p, 0:2] = to_torch(transform(coords[p, 0:2], center, scale, res, 1, 0))
return coords
def crop(img, center, scale, res, rot=0, interp='bilinear'):
# import pdb; pdb.set_trace()
# mode = 'F'
img = im_to_numpy(img)
# Preprocessing for efficient cropping
ht, wd = img.shape[0], img.shape[1]
sf = scale * 200.0 / res[0]
if sf < 2:
sf = 1
else:
new_size = int(np.math.floor(max(ht, wd) / sf))
new_ht = int(np.math.floor(ht / sf))
new_wd = int(np.math.floor(wd / sf))
if new_size < 2:
return torch.zeros(res[0], res[1], img.shape[2]) \
if len(img.shape) > 2 else torch.zeros(res[0], res[1])
else:
img = imresize(img, [new_ht, new_wd], interp=interp) # , mode=mode)
center = center * 1.0 / sf
scale = scale / sf
# Upper left point
ul = np.array(transform([0, 0], center, scale, res, invert=1))
# Bottom right point
br = np.array(transform(res, center, scale, res, invert=1))
# Padding so that when rotated proper amount of context is included
pad = int(np.linalg.norm(br - ul) / 2 - float(br[1] - ul[1]) / 2)
if not rot == 0:
ul -= pad
br += pad
new_shape = [br[1] - ul[1], br[0] - ul[0]]
if len(img.shape) > 2:
new_shape += [img.shape[2]]
new_img = np.zeros(new_shape)
# Range to fill new array
new_x = max(0, -ul[0]), min(br[0], img.shape[1]) - ul[0]
new_y = max(0, -ul[1]), min(br[1], img.shape[0]) - ul[1]
# Range to sample from original image
old_x = max(0, ul[0]), min(img.shape[1], br[0])
old_y = max(0, ul[1]), min(img.shape[0], br[1])
new_img[new_y[0]:new_y[1], new_x[0]:new_x[1]] = img[old_y[0]:old_y[1], old_x[0]:old_x[1]]
if not rot == 0:
# Remove padding
new_img = imrotate(new_img, rot, interp=interp) # , mode=mode)
new_img = new_img[pad:-pad, pad:-pad]
new_img = im_to_torch(imresize(new_img, res, interp=interp)) #, mode=mode))
return new_img
|