Spaces:
Runtime error
Runtime error
| # 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 | |