Spaces:
Sleeping
Sleeping
| from __future__ import absolute_import | |
| from __future__ import print_function | |
| from __future__ import division | |
| import os | |
| import os.path as osp | |
| from collections import OrderedDict | |
| import cv2 | |
| import numpy as np | |
| from skimage.filters import gaussian | |
| def get_transform(center, scale, res, rot=0): | |
| """Generate transformation matrix.""" | |
| # res: (height, width), (rows, cols) | |
| crop_aspect_ratio = res[0] / float(res[1]) | |
| h = 200 * scale | |
| w = h / crop_aspect_ratio | |
| t = np.zeros((3, 3)) | |
| t[0, 0] = float(res[1]) / w | |
| t[1, 1] = float(res[0]) / h | |
| t[0, 2] = res[1] * (-float(center[0]) / w + .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): | |
| """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) | |
| return np.array([round(new_pt[0]), round(new_pt[1])], dtype=int) + 1 | |
| def crop(img, center, scale, res): | |
| """ | |
| Crop image according to the supplied bounding box. | |
| res: [rows, cols] | |
| """ | |
| # Upper left point | |
| ul = np.array(transform([1, 1], center, scale, res, invert=1)) - 1 | |
| # Bottom right point | |
| br = np.array(transform([res[1] + 1, res[0] + 1], center, scale, res, invert=1)) - 1 | |
| 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, dtype=np.float32) | |
| # Range to fill new array | |
| new_x = max(0, -ul[0]), min(br[0], len(img[0])) - ul[0] | |
| new_y = max(0, -ul[1]), min(br[1], len(img)) - ul[1] | |
| # Range to sample from original image | |
| old_x = max(0, ul[0]), min(len(img[0]), br[0]) | |
| old_y = max(0, ul[1]), min(len(img), br[1]) | |
| try: | |
| 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]] | |
| except Exception as e: | |
| print(e) | |
| new_img = cv2.resize(new_img, (res[1], res[0])) # (cols, rows) | |
| return new_img, ul, br | |
| def process_image(orig_img_rgb, center, scale, crop_height=256, crop_width=192, blur=False, do_crop=True): | |
| """ | |
| Read image, do preprocessing and possibly crop it according to the bounding box. | |
| If there are bounding box annotations, use them to crop the image. | |
| If no bounding box is specified but openpose detections are available, use them to get the bounding box. | |
| """ | |
| if blur: | |
| # Blur image to avoid aliasing artifacts | |
| downsampling_factor = ((scale * 200 * 1.0) / crop_height) | |
| downsampling_factor = downsampling_factor / 2.0 | |
| if downsampling_factor > 1.1: | |
| orig_img_rgb = gaussian(orig_img_rgb, sigma=(downsampling_factor-1)/2, channel_axis=2, preserve_range=True) | |
| IMG_NORM_MEAN = [0.485, 0.456, 0.406] | |
| IMG_NORM_STD = [0.229, 0.224, 0.225] | |
| if do_crop: | |
| img, ul, br = crop(orig_img_rgb, center, scale, (crop_height, crop_width)) | |
| else: | |
| img = orig_img_rgb.copy() | |
| crop_img = img.copy() | |
| img = img / 255. | |
| mean = np.array(IMG_NORM_MEAN, dtype=np.float32) | |
| std = np.array(IMG_NORM_STD, dtype=np.float32) | |
| norm_img = (img - mean) / std | |
| norm_img = np.transpose(norm_img, (2, 0, 1)) | |
| return norm_img, crop_img |