|
|
| import os
|
| import cv2
|
| import math
|
| import random
|
| import numpy as np
|
|
|
| def get_mask_boxes(mask):
|
| y_coords, x_coords = np.nonzero(mask)
|
| x_min = x_coords.min()
|
| x_max = x_coords.max()
|
| y_min = y_coords.min()
|
| y_max = y_coords.max()
|
| bbox = np.array([x_min, y_min, x_max, y_max]).astype(np.int32)
|
| return bbox
|
|
|
|
|
| def get_aug_mask(body_mask, w_len=10, h_len=20):
|
| body_bbox = get_mask_boxes(body_mask)
|
|
|
| bbox_wh = body_bbox[2:4] - body_bbox[0:2]
|
| w_slice = np.int32(bbox_wh[0] / w_len)
|
| h_slice = np.int32(bbox_wh[1] / h_len)
|
|
|
| for each_w in range(body_bbox[0], body_bbox[2], w_slice):
|
| w_start = min(each_w, body_bbox[2])
|
| w_end = min((each_w + w_slice), body_bbox[2])
|
| for each_h in range(body_bbox[1], body_bbox[3], h_slice):
|
| h_start = min(each_h, body_bbox[3])
|
| h_end = min((each_h + h_slice), body_bbox[3])
|
| if body_mask[h_start:h_end, w_start:w_end].sum() > 0:
|
| body_mask[h_start:h_end, w_start:w_end] = 1
|
|
|
| return body_mask
|
|
|
| def get_mask_body_img(img_copy, hand_mask, k=7, iterations=1):
|
| kernel = np.ones((k, k), np.uint8)
|
| dilation = cv2.dilate(hand_mask, kernel, iterations=iterations)
|
| mask_hand_img = img_copy * (1 - dilation[:, :, None])
|
|
|
| return mask_hand_img, dilation
|
|
|
|
|
| def get_face_bboxes(kp2ds, scale, image_shape, ratio_aug):
|
| h, w = image_shape
|
| kp2ds_face = kp2ds.copy()[23:91, :2]
|
|
|
| min_x, min_y = np.min(kp2ds_face, axis=0)
|
| max_x, max_y = np.max(kp2ds_face, axis=0)
|
|
|
|
|
| initial_width = max_x - min_x
|
| initial_height = max_y - min_y
|
|
|
| initial_area = initial_width * initial_height
|
|
|
| expanded_area = initial_area * scale
|
|
|
| new_width = np.sqrt(expanded_area * (initial_width / initial_height))
|
| new_height = np.sqrt(expanded_area * (initial_height / initial_width))
|
|
|
| delta_width = (new_width - initial_width) / 2
|
| delta_height = (new_height - initial_height) / 4
|
|
|
| if ratio_aug:
|
| if random.random() > 0.5:
|
| delta_width += random.uniform(0, initial_width // 10)
|
| else:
|
| delta_height += random.uniform(0, initial_height // 10)
|
|
|
| expanded_min_x = max(min_x - delta_width, 0)
|
| expanded_max_x = min(max_x + delta_width, w)
|
| expanded_min_y = max(min_y - 3 * delta_height, 0)
|
| expanded_max_y = min(max_y + delta_height, h)
|
|
|
| return [int(expanded_min_x), int(expanded_max_x), int(expanded_min_y), int(expanded_max_y)]
|
|
|
|
|
| def calculate_new_size(orig_w, orig_h, target_area, divisor=64):
|
|
|
| target_ratio = orig_w / orig_h
|
|
|
| def check_valid(w, h):
|
|
|
| if w <= 0 or h <= 0:
|
| return False
|
| return (w * h <= target_area and
|
| w % divisor == 0 and
|
| h % divisor == 0)
|
|
|
| def get_ratio_diff(w, h):
|
|
|
| return abs(w / h - target_ratio)
|
|
|
| def round_to_64(value, round_up=False, divisor=64):
|
|
|
| if round_up:
|
| return divisor * ((value + (divisor - 1)) // divisor)
|
| return divisor * (value // divisor)
|
|
|
| possible_sizes = []
|
|
|
| max_area_h = int(np.sqrt(target_area / target_ratio))
|
| max_area_w = int(max_area_h * target_ratio)
|
|
|
| max_h = round_to_64(max_area_h, round_up=True, divisor=divisor)
|
| max_w = round_to_64(max_area_w, round_up=True, divisor=divisor)
|
|
|
| for h in range(divisor, max_h + divisor, divisor):
|
| ideal_w = h * target_ratio
|
|
|
| w_down = round_to_64(ideal_w)
|
| w_up = round_to_64(ideal_w, round_up=True)
|
|
|
| for w in [w_down, w_up]:
|
| if check_valid(w, h, divisor):
|
| possible_sizes.append((w, h, get_ratio_diff(w, h)))
|
|
|
| if not possible_sizes:
|
| raise ValueError("Can not find suitable size")
|
|
|
| possible_sizes.sort(key=lambda x: (-x[0] * x[1], x[2]))
|
|
|
| best_w, best_h, _ = possible_sizes[0]
|
| return int(best_w), int(best_h)
|
|
|
|
|
| def resize_by_area(image, target_area, keep_aspect_ratio=True, divisor=64, padding_color=(0, 0, 0)):
|
| h, w = image.shape[:2]
|
| try:
|
| new_w, new_h = calculate_new_size(w, h, target_area, divisor)
|
| except:
|
| aspect_ratio = w / h
|
|
|
| if keep_aspect_ratio:
|
| new_h = math.sqrt(target_area / aspect_ratio)
|
| new_w = target_area / new_h
|
| else:
|
| new_w = new_h = math.sqrt(target_area)
|
|
|
| new_w, new_h = int((new_w // divisor) * divisor), int((new_h // divisor) * divisor)
|
|
|
| interpolation = cv2.INTER_AREA if (new_w * new_h < w * h) else cv2.INTER_LINEAR
|
|
|
| resized_image = padding_resize(image, height=new_h, width=new_w, padding_color=padding_color,
|
| interpolation=interpolation)
|
| return resized_image
|
|
|
|
|
| def padding_resize(img_ori, height=512, width=512, padding_color=(0, 0, 0), interpolation=cv2.INTER_LINEAR):
|
| ori_height = img_ori.shape[0]
|
| ori_width = img_ori.shape[1]
|
| channel = img_ori.shape[2]
|
|
|
| img_pad = np.zeros((height, width, channel), dtype=img_ori.dtype)
|
| if channel == 1:
|
| img_pad[:, :, 0] = padding_color[0]
|
| else:
|
| img_pad[:, :, 0] = padding_color[0]
|
| img_pad[:, :, 1] = padding_color[1]
|
| img_pad[:, :, 2] = padding_color[2]
|
|
|
| if (ori_height / ori_width) > (height / width):
|
| new_width = int(height / ori_height * ori_width)
|
| img = cv2.resize(img_ori, (new_width, height), interpolation=interpolation)
|
| padding = int((width - new_width) / 2)
|
| if len(img.shape) == 2:
|
| img = img[:, :, np.newaxis]
|
| img_pad[:, padding: padding + new_width, :] = img
|
| else:
|
| new_height = int(width / ori_width * ori_height)
|
| img = cv2.resize(img_ori, (width, new_height), interpolation=interpolation)
|
| padding = int((height - new_height) / 2)
|
| if len(img.shape) == 2:
|
| img = img[:, :, np.newaxis]
|
| img_pad[padding: padding + new_height, :, :] = img
|
|
|
| return img_pad
|
|
|
| def resize_to_bounds(img_ori, height=512, width=512, padding_color=(0, 0, 0), interpolation=cv2.INTER_LINEAR, extra_padding=64, crop_target_image=None):
|
|
|
| if crop_target_image is not None:
|
| ref = crop_target_image
|
| if ref.ndim == 2:
|
| mask = ref > 0
|
| else:
|
| mask = np.any(ref != 0, axis=2)
|
| coords = np.argwhere(mask)
|
| if coords.size == 0:
|
|
|
| y0, x0 = 0, 0
|
| y1, x1 = img_ori.shape[0], img_ori.shape[1]
|
| else:
|
| y0, x0 = coords.min(axis=0)
|
| y1, x1 = coords.max(axis=0) + 1
|
|
|
| pad_y0 = y0 - extra_padding
|
| pad_x0 = x0 - extra_padding
|
| pad_y1 = y1 + extra_padding
|
| pad_x1 = x1 + extra_padding
|
|
|
| crop_y0 = max(pad_y0, 0)
|
| crop_x0 = max(pad_x0, 0)
|
| crop_y1 = min(pad_y1, img_ori.shape[0])
|
| crop_x1 = min(pad_x1, img_ori.shape[1])
|
| crop_img = img_ori[crop_y0:crop_y1, crop_x0:crop_x1]
|
|
|
| pad_top = crop_y0 - pad_y0
|
| pad_left = crop_x0 - pad_x0
|
| pad_bottom = pad_y1 - crop_y1
|
| pad_right = pad_x1 - crop_x1
|
| if any([pad_top, pad_left, pad_bottom, pad_right]):
|
| channel = crop_img.shape[2] if crop_img.ndim == 3 else 1
|
| crop_img = np.pad(
|
| crop_img,
|
| ((pad_top, pad_bottom), (pad_left, pad_right)) + ((0, 0),) if channel > 1 else ((pad_top, pad_bottom), (pad_left, pad_right)),
|
| mode='constant', constant_values=0
|
| )
|
| else:
|
| if img_ori.ndim == 2:
|
| mask = img_ori > 0
|
| else:
|
| mask = np.any(img_ori != 0, axis=2)
|
| coords = np.argwhere(mask)
|
| if coords.size == 0:
|
|
|
| crop_img = img_ori
|
| else:
|
| y0, x0 = coords.min(axis=0)
|
| y1, x1 = coords.max(axis=0) + 1
|
| pad_y0 = y0 - extra_padding
|
| pad_x0 = x0 - extra_padding
|
| pad_y1 = y1 + extra_padding
|
| pad_x1 = x1 + extra_padding
|
| crop_y0 = max(pad_y0, 0)
|
| crop_x0 = max(pad_x0, 0)
|
| crop_y1 = min(pad_y1, img_ori.shape[0])
|
| crop_x1 = min(pad_x1, img_ori.shape[1])
|
| crop_img = img_ori[crop_y0:crop_y1, crop_x0:crop_x1]
|
| pad_top = crop_y0 - pad_y0
|
| pad_left = crop_x0 - pad_x0
|
| pad_bottom = pad_y1 - crop_y1
|
| pad_right = pad_x1 - crop_x1
|
| if any([pad_top, pad_left, pad_bottom, pad_right]):
|
| channel = crop_img.shape[2] if crop_img.ndim == 3 else 1
|
| crop_img = np.pad(
|
| crop_img,
|
| ((pad_top, pad_bottom), (pad_left, pad_right)) + ((0, 0),) if channel > 1 else ((pad_top, pad_bottom), (pad_left, pad_right)),
|
| mode='constant', constant_values=0
|
| )
|
|
|
| ori_height = crop_img.shape[0]
|
| ori_width = crop_img.shape[1]
|
| channel = crop_img.shape[2] if crop_img.ndim == 3 else 1
|
|
|
| img_pad = np.zeros((height, width, channel), dtype=crop_img.dtype)
|
| if channel == 1:
|
| img_pad[:, :, 0] = padding_color[0]
|
| else:
|
| for c in range(channel):
|
| img_pad[:, :, c] = padding_color[c % len(padding_color)]
|
|
|
|
|
| crop_aspect = ori_width / ori_height
|
| target_aspect = width / height
|
| if crop_aspect > target_aspect:
|
| new_width = width
|
| new_height = int(width / crop_aspect)
|
| else:
|
| new_height = height
|
| new_width = int(height * crop_aspect)
|
| img = cv2.resize(crop_img, (new_width, new_height), interpolation=interpolation)
|
| if img.ndim == 2:
|
| img = img[:, :, np.newaxis]
|
| y_pad = (height - new_height) // 2
|
| x_pad = (width - new_width) // 2
|
| img_pad[y_pad:y_pad + new_height, x_pad:x_pad + new_width, :] = img
|
|
|
| return img_pad
|
|
|
|
|
| def get_frame_indices(frame_num, video_fps, clip_length, train_fps):
|
|
|
| start_frame = 0
|
| times = np.arange(0, clip_length) / train_fps
|
| frame_indices = start_frame + np.round(times * video_fps).astype(int)
|
| frame_indices = np.clip(frame_indices, 0, frame_num - 1)
|
|
|
| return frame_indices.tolist()
|
|
|
|
|
| def get_face_bboxes(kp2ds, scale, image_shape):
|
| h, w = image_shape
|
| kp2ds_face = kp2ds.copy()[1:] * (w, h)
|
|
|
| min_x, min_y = np.min(kp2ds_face, axis=0)
|
| max_x, max_y = np.max(kp2ds_face, axis=0)
|
|
|
| initial_width = max_x - min_x
|
| initial_height = max_y - min_y
|
|
|
| initial_area = initial_width * initial_height
|
|
|
| expanded_area = initial_area * scale
|
|
|
| new_width = np.sqrt(expanded_area * (initial_width / initial_height))
|
| new_height = np.sqrt(expanded_area * (initial_height / initial_width))
|
|
|
| delta_width = (new_width - initial_width) / 2
|
| delta_height = (new_height - initial_height) / 4
|
|
|
| expanded_min_x = max(min_x - delta_width, 0)
|
| expanded_max_x = min(max_x + delta_width, w)
|
| expanded_min_y = max(min_y - 3 * delta_height, 0)
|
| expanded_max_y = min(max_y + delta_height, h)
|
|
|
| return [int(expanded_min_x), int(expanded_max_x), int(expanded_min_y), int(expanded_max_y)] |
