# Copyright (c) OpenMMLab. All rights reserved. import copy as cp import pickle import numpy as np from mmcv.fileio import FileClient from scipy.stats import mode from ..builder import PIPELINES from .augmentations import Flip @PIPELINES.register_module() class UniformSampleFrames: """Uniformly sample frames from the video. To sample an n-frame clip from the video. UniformSampleFrames basically divide the video into n segments of equal length and randomly sample one frame from each segment. To make the testing results reproducible, a random seed is set during testing, to make the sampling results deterministic. Required keys are "total_frames", "start_index" , added or modified keys are "frame_inds", "clip_len", "frame_interval" and "num_clips". Args: clip_len (int): Frames of each sampled output clip. num_clips (int): Number of clips to be sampled. Default: 1. test_mode (bool): Store True when building test or validation dataset. Default: False. seed (int): The random seed used during test time. Default: 255. """ def __init__(self, clip_len, num_clips=1, test_mode=False, seed=255): self.clip_len = clip_len self.num_clips = num_clips self.test_mode = test_mode self.seed = seed def _get_train_clips(self, num_frames, clip_len): """Uniformly sample indices for training clips. Args: num_frames (int): The number of frames. clip_len (int): The length of the clip. """ assert self.num_clips == 1 if num_frames < clip_len: start = np.random.randint(0, num_frames) inds = np.arange(start, start + clip_len) elif clip_len <= num_frames < 2 * clip_len: basic = np.arange(clip_len) inds = np.random.choice( clip_len + 1, num_frames - clip_len, replace=False) offset = np.zeros(clip_len + 1, dtype=np.int64) offset[inds] = 1 offset = np.cumsum(offset) inds = basic + offset[:-1] else: bids = np.array( [i * num_frames // clip_len for i in range(clip_len + 1)]) bsize = np.diff(bids) bst = bids[:clip_len] offset = np.random.randint(bsize) inds = bst + offset return inds def _get_test_clips(self, num_frames, clip_len): """Uniformly sample indices for testing clips. Args: num_frames (int): The number of frames. clip_len (int): The length of the clip. """ np.random.seed(self.seed) if num_frames < clip_len: # Then we use a simple strategy if num_frames < self.num_clips: start_inds = list(range(self.num_clips)) else: start_inds = [ i * num_frames // self.num_clips for i in range(self.num_clips) ] inds = np.concatenate( [np.arange(i, i + clip_len) for i in start_inds]) elif clip_len <= num_frames < clip_len * 2: all_inds = [] for i in range(self.num_clips): basic = np.arange(clip_len) inds = np.random.choice( clip_len + 1, num_frames - clip_len, replace=False) offset = np.zeros(clip_len + 1, dtype=np.int64) offset[inds] = 1 offset = np.cumsum(offset) inds = basic + offset[:-1] all_inds.append(inds) inds = np.concatenate(all_inds) else: bids = np.array( [i * num_frames // clip_len for i in range(clip_len + 1)]) bsize = np.diff(bids) bst = bids[:clip_len] all_inds = [] for i in range(self.num_clips): offset = np.random.randint(bsize) all_inds.append(bst + offset) inds = np.concatenate(all_inds) return inds def __call__(self, results): num_frames = results['total_frames'] if self.test_mode: inds = self._get_test_clips(num_frames, self.clip_len) else: inds = self._get_train_clips(num_frames, self.clip_len) inds = np.mod(inds, num_frames) start_index = results['start_index'] inds = inds + start_index results['frame_inds'] = inds.astype(np.int64) results['clip_len'] = self.clip_len results['frame_interval'] = None results['num_clips'] = self.num_clips return results def __repr__(self): repr_str = (f'{self.__class__.__name__}(' f'clip_len={self.clip_len}, ' f'num_clips={self.num_clips}, ' f'test_mode={self.test_mode}, ' f'seed={self.seed})') return repr_str @PIPELINES.register_module() class PoseDecode: """Load and decode pose with given indices. Required keys are "keypoint", "frame_inds" (optional), "keypoint_score" (optional), added or modified keys are "keypoint", "keypoint_score" (if applicable). """ @staticmethod def _load_kp(kp, frame_inds): """Load keypoints given frame indices. Args: kp (np.ndarray): The keypoint coordinates. frame_inds (np.ndarray): The frame indices. """ return [x[frame_inds].astype(np.float32) for x in kp] @staticmethod def _load_kpscore(kpscore, frame_inds): """Load keypoint scores given frame indices. Args: kpscore (np.ndarray): The confidence scores of keypoints. frame_inds (np.ndarray): The frame indices. """ return [x[frame_inds].astype(np.float32) for x in kpscore] def __call__(self, results): if 'frame_inds' not in results: results['frame_inds'] = np.arange(results['total_frames']) if results['frame_inds'].ndim != 1: results['frame_inds'] = np.squeeze(results['frame_inds']) offset = results.get('offset', 0) frame_inds = results['frame_inds'] + offset if 'keypoint_score' in results: kpscore = results['keypoint_score'] results['keypoint_score'] = kpscore[:, frame_inds].astype(np.float32) if 'keypoint' in results: results['keypoint'] = results['keypoint'][:, frame_inds].astype( np.float32) return results def __repr__(self): repr_str = f'{self.__class__.__name__}()' return repr_str @PIPELINES.register_module() class LoadKineticsPose: """Load Kinetics Pose given filename (The format should be pickle) Required keys are "filename", "total_frames", "img_shape", "frame_inds", "anno_inds" (for mmpose source, optional), added or modified keys are "keypoint", "keypoint_score". Args: io_backend (str): IO backend where frames are stored. Default: 'disk'. squeeze (bool): Whether to remove frames with no human pose. Default: True. max_person (int): The max number of persons in a frame. Default: 10. keypoint_weight (dict): The weight of keypoints. We set the confidence score of a person as the weighted sum of confidence scores of each joint. Persons with low confidence scores are dropped (if exceed max_person). Default: dict(face=1, torso=2, limb=3). source (str): The sources of the keypoints used. Choices are 'mmpose' and 'openpose-18'. Default: 'mmpose'. kwargs (dict, optional): Arguments for FileClient. """ def __init__(self, io_backend='disk', squeeze=True, max_person=100, keypoint_weight=dict(face=1, torso=2, limb=3), source='mmpose', **kwargs): self.io_backend = io_backend self.squeeze = squeeze self.max_person = max_person self.keypoint_weight = cp.deepcopy(keypoint_weight) self.source = source if source == 'openpose-18': self.kpsubset = dict( face=[0, 14, 15, 16, 17], torso=[1, 2, 8, 5, 11], limb=[3, 4, 6, 7, 9, 10, 12, 13]) elif source == 'mmpose': self.kpsubset = dict( face=[0, 1, 2, 3, 4], torso=[5, 6, 11, 12], limb=[7, 8, 9, 10, 13, 14, 15, 16]) else: raise NotImplementedError('Unknown source of Kinetics Pose') self.kwargs = kwargs self.file_client = None def __call__(self, results): assert 'filename' in results filename = results.pop('filename') # only applicable to source == 'mmpose' anno_inds = None if 'anno_inds' in results: assert self.source == 'mmpose' anno_inds = results.pop('anno_inds') results.pop('box_score', None) if self.file_client is None: self.file_client = FileClient(self.io_backend, **self.kwargs) bytes = self.file_client.get(filename) # only the kp array is in the pickle file, each kp include x, y, score. kps = pickle.loads(bytes) total_frames = results['total_frames'] frame_inds = results.pop('frame_inds') if anno_inds is not None: kps = kps[anno_inds] frame_inds = frame_inds[anno_inds] frame_inds = list(frame_inds) def mapinds(inds): uni = np.unique(inds) map_ = {x: i for i, x in enumerate(uni)} inds = [map_[x] for x in inds] return np.array(inds, dtype=np.int16) if self.squeeze: frame_inds = mapinds(frame_inds) total_frames = np.max(frame_inds) + 1 # write it back results['total_frames'] = total_frames h, w = results['img_shape'] if self.source == 'openpose-18': kps[:, :, 0] *= w kps[:, :, 1] *= h num_kp = kps.shape[1] num_person = mode(frame_inds)[-1][0] new_kp = np.zeros([num_person, total_frames, num_kp, 2], dtype=np.float16) new_kpscore = np.zeros([num_person, total_frames, num_kp], dtype=np.float16) # 32768 is enough num_person_frame = np.zeros([total_frames], dtype=np.int16) for frame_ind, kp in zip(frame_inds, kps): person_ind = num_person_frame[frame_ind] new_kp[person_ind, frame_ind] = kp[:, :2] new_kpscore[person_ind, frame_ind] = kp[:, 2] num_person_frame[frame_ind] += 1 kpgrp = self.kpsubset weight = self.keypoint_weight results['num_person'] = num_person if num_person > self.max_person: for i in range(total_frames): np_frame = num_person_frame[i] val = new_kpscore[:np_frame, i] val = ( np.sum(val[:, kpgrp['face']], 1) * weight['face'] + np.sum(val[:, kpgrp['torso']], 1) * weight['torso'] + np.sum(val[:, kpgrp['limb']], 1) * weight['limb']) inds = sorted(range(np_frame), key=lambda x: -val[x]) new_kpscore[:np_frame, i] = new_kpscore[inds, i] new_kp[:np_frame, i] = new_kp[inds, i] results['num_person'] = self.max_person results['keypoint'] = new_kp[:self.max_person] results['keypoint_score'] = new_kpscore[:self.max_person] return results def __repr__(self): repr_str = (f'{self.__class__.__name__}(' f'io_backend={self.io_backend}, ' f'squeeze={self.squeeze}, ' f'max_person={self.max_person}, ' f'keypoint_weight={self.keypoint_weight}, ' f'source={self.source}, ' f'kwargs={self.kwargs})') return repr_str @PIPELINES.register_module() class GeneratePoseTarget: """Generate pseudo heatmaps based on joint coordinates and confidence. Required keys are "keypoint", "img_shape", "keypoint_score" (optional), added or modified keys are "imgs". Args: sigma (float): The sigma of the generated gaussian map. Default: 0.6. use_score (bool): Use the confidence score of keypoints as the maximum of the gaussian maps. Default: True. with_kp (bool): Generate pseudo heatmaps for keypoints. Default: True. with_limb (bool): Generate pseudo heatmaps for limbs. At least one of 'with_kp' and 'with_limb' should be True. Default: False. skeletons (tuple[tuple]): The definition of human skeletons. Default: ((0, 1), (0, 2), (1, 3), (2, 4), (0, 5), (5, 7), (7, 9), (0, 6), (6, 8), (8, 10), (5, 11), (11, 13), (13, 15), (6, 12), (12, 14), (14, 16), (11, 12)), which is the definition of COCO-17p skeletons. double (bool): Output both original heatmaps and flipped heatmaps. Default: False. left_kp (tuple[int]): Indexes of left keypoints, which is used when flipping heatmaps. Default: (1, 3, 5, 7, 9, 11, 13, 15), which is left keypoints in COCO-17p. right_kp (tuple[int]): Indexes of right keypoints, which is used when flipping heatmaps. Default: (2, 4, 6, 8, 10, 12, 14, 16), which is right keypoints in COCO-17p. """ def __init__(self, sigma=0.6, use_score=True, with_kp=True, with_limb=False, skeletons=((0, 1), (0, 2), (1, 3), (2, 4), (0, 5), (5, 7), (7, 9), (0, 6), (6, 8), (8, 10), (5, 11), (11, 13), (13, 15), (6, 12), (12, 14), (14, 16), (11, 12)), double=False, left_kp=(1, 3, 5, 7, 9, 11, 13, 15), right_kp=(2, 4, 6, 8, 10, 12, 14, 16)): self.sigma = sigma self.use_score = use_score self.with_kp = with_kp self.with_limb = with_limb self.double = double # an auxiliary const self.eps = 1e-4 assert self.with_kp or self.with_limb, ( 'At least one of "with_limb" ' 'and "with_kp" should be set as True.') self.left_kp = left_kp self.right_kp = right_kp self.skeletons = skeletons def generate_a_heatmap(self, img_h, img_w, centers, sigma, max_values): """Generate pseudo heatmap for one keypoint in one frame. Args: img_h (int): The height of the heatmap. img_w (int): The width of the heatmap. centers (np.ndarray): The coordinates of corresponding keypoints (of multiple persons). sigma (float): The sigma of generated gaussian. max_values (np.ndarray): The max values of each keypoint. Returns: np.ndarray: The generated pseudo heatmap. """ heatmap = np.zeros([img_h, img_w], dtype=np.float32) for center, max_value in zip(centers, max_values): mu_x, mu_y = center[0], center[1] if max_value < self.eps: continue st_x = max(int(mu_x - 3 * sigma), 0) ed_x = min(int(mu_x + 3 * sigma) + 1, img_w) st_y = max(int(mu_y - 3 * sigma), 0) ed_y = min(int(mu_y + 3 * sigma) + 1, img_h) x = np.arange(st_x, ed_x, 1, np.float32) y = np.arange(st_y, ed_y, 1, np.float32) # if the keypoint not in the heatmap coordinate system if not (len(x) and len(y)): continue y = y[:, None] patch = np.exp(-((x - mu_x)**2 + (y - mu_y)**2) / 2 / sigma**2) patch = patch * max_value heatmap[st_y:ed_y, st_x:ed_x] = np.maximum(heatmap[st_y:ed_y, st_x:ed_x], patch) return heatmap def generate_a_limb_heatmap(self, img_h, img_w, starts, ends, sigma, start_values, end_values): """Generate pseudo heatmap for one limb in one frame. Args: img_h (int): The height of the heatmap. img_w (int): The width of the heatmap. starts (np.ndarray): The coordinates of one keypoint in the corresponding limbs (of multiple persons). ends (np.ndarray): The coordinates of the other keypoint in the corresponding limbs (of multiple persons). sigma (float): The sigma of generated gaussian. start_values (np.ndarray): The max values of one keypoint in the corresponding limbs. end_values (np.ndarray): The max values of the other keypoint in the corresponding limbs. Returns: np.ndarray: The generated pseudo heatmap. """ heatmap = np.zeros([img_h, img_w], dtype=np.float32) for start, end, start_value, end_value in zip(starts, ends, start_values, end_values): value_coeff = min(start_value, end_value) if value_coeff < self.eps: continue min_x, max_x = min(start[0], end[0]), max(start[0], end[0]) min_y, max_y = min(start[1], end[1]), max(start[1], end[1]) min_x = max(int(min_x - 3 * sigma), 0) max_x = min(int(max_x + 3 * sigma) + 1, img_w) min_y = max(int(min_y - 3 * sigma), 0) max_y = min(int(max_y + 3 * sigma) + 1, img_h) x = np.arange(min_x, max_x, 1, np.float32) y = np.arange(min_y, max_y, 1, np.float32) if not (len(x) and len(y)): continue y = y[:, None] x_0 = np.zeros_like(x) y_0 = np.zeros_like(y) # distance to start keypoints d2_start = ((x - start[0])**2 + (y - start[1])**2) # distance to end keypoints d2_end = ((x - end[0])**2 + (y - end[1])**2) # the distance between start and end keypoints. d2_ab = ((start[0] - end[0])**2 + (start[1] - end[1])**2) if d2_ab < 1: full_map = self.generate_a_heatmap(img_h, img_w, [start], sigma, [start_value]) heatmap = np.maximum(heatmap, full_map) continue coeff = (d2_start - d2_end + d2_ab) / 2. / d2_ab a_dominate = coeff <= 0 b_dominate = coeff >= 1 seg_dominate = 1 - a_dominate - b_dominate position = np.stack([x + y_0, y + x_0], axis=-1) projection = start + np.stack([coeff, coeff], axis=-1) * ( end - start) d2_line = position - projection d2_line = d2_line[:, :, 0]**2 + d2_line[:, :, 1]**2 d2_seg = ( a_dominate * d2_start + b_dominate * d2_end + seg_dominate * d2_line) patch = np.exp(-d2_seg / 2. / sigma**2) patch = patch * value_coeff heatmap[min_y:max_y, min_x:max_x] = np.maximum( heatmap[min_y:max_y, min_x:max_x], patch) return heatmap def generate_heatmap(self, img_h, img_w, kps, sigma, max_values): """Generate pseudo heatmap for all keypoints and limbs in one frame (if needed). Args: img_h (int): The height of the heatmap. img_w (int): The width of the heatmap. kps (np.ndarray): The coordinates of keypoints in this frame. sigma (float): The sigma of generated gaussian. max_values (np.ndarray): The confidence score of each keypoint. Returns: np.ndarray: The generated pseudo heatmap. """ heatmaps = [] if self.with_kp: num_kp = kps.shape[1] for i in range(num_kp): heatmap = self.generate_a_heatmap(img_h, img_w, kps[:, i], sigma, max_values[:, i]) heatmaps.append(heatmap) if self.with_limb: for limb in self.skeletons: start_idx, end_idx = limb starts = kps[:, start_idx] ends = kps[:, end_idx] start_values = max_values[:, start_idx] end_values = max_values[:, end_idx] heatmap = self.generate_a_limb_heatmap(img_h, img_w, starts, ends, sigma, start_values, end_values) heatmaps.append(heatmap) return np.stack(heatmaps, axis=-1) def gen_an_aug(self, results): """Generate pseudo heatmaps for all frames. Args: results (dict): The dictionary that contains all info of a sample. Returns: list[np.ndarray]: The generated pseudo heatmaps. """ all_kps = results['keypoint'] kp_shape = all_kps.shape if 'keypoint_score' in results: all_kpscores = results['keypoint_score'] else: all_kpscores = np.ones(kp_shape[:-1], dtype=np.float32) img_h, img_w = results['img_shape'] num_frame = kp_shape[1] imgs = [] for i in range(num_frame): sigma = self.sigma kps = all_kps[:, i] kpscores = all_kpscores[:, i] max_values = np.ones(kpscores.shape, dtype=np.float32) if self.use_score: max_values = kpscores hmap = self.generate_heatmap(img_h, img_w, kps, sigma, max_values) imgs.append(hmap) return imgs def __call__(self, results): if not self.double: results['imgs'] = np.stack(self.gen_an_aug(results)) else: results_ = cp.deepcopy(results) flip = Flip( flip_ratio=1, left_kp=self.left_kp, right_kp=self.right_kp) results_ = flip(results_) results['imgs'] = np.concatenate( [self.gen_an_aug(results), self.gen_an_aug(results_)]) return results def __repr__(self): repr_str = (f'{self.__class__.__name__}(' f'sigma={self.sigma}, ' f'use_score={self.use_score}, ' f'with_kp={self.with_kp}, ' f'with_limb={self.with_limb}, ' f'skeletons={self.skeletons}, ' f'double={self.double}, ' f'left_kp={self.left_kp}, ' f'right_kp={self.right_kp})') return repr_str @PIPELINES.register_module() class PaddingWithLoop: """Sample frames from the video. To sample an n-frame clip from the video, PaddingWithLoop samples the frames from zero index, and loop the frames if the length of video frames is less than te value of 'clip_len'. Required keys are "total_frames", added or modified keys are "frame_inds", "clip_len", "frame_interval" and "num_clips". Args: clip_len (int): Frames of each sampled output clip. num_clips (int): Number of clips to be sampled. Default: 1. """ def __init__(self, clip_len, num_clips=1): self.clip_len = clip_len self.num_clips = num_clips def __call__(self, results): num_frames = results['total_frames'] start = 0 inds = np.arange(start, start + self.clip_len) inds = np.mod(inds, num_frames) results['frame_inds'] = inds.astype(np.int64) results['clip_len'] = self.clip_len results['frame_interval'] = None results['num_clips'] = self.num_clips return results @PIPELINES.register_module() class PoseNormalize: """Normalize the range of keypoint values to [-1,1]. Args: mean (list | tuple): The mean value of the keypoint values. min_value (list | tuple): The minimum value of the keypoint values. max_value (list | tuple): The maximum value of the keypoint values. """ def __init__(self, mean=(960., 540., 0.5), min_value=(0., 0., 0.), max_value=(1920, 1080, 1.)): self.mean = np.array(mean, dtype=np.float32).reshape(-1, 1, 1, 1) self.min_value = np.array( min_value, dtype=np.float32).reshape(-1, 1, 1, 1) self.max_value = np.array( max_value, dtype=np.float32).reshape(-1, 1, 1, 1) def __call__(self, results): keypoint = results['keypoint'] keypoint = (keypoint - self.mean) / (self.max_value - self.min_value) results['keypoint'] = keypoint results['keypoint_norm_cfg'] = dict( mean=self.mean, min_value=self.min_value, max_value=self.max_value) return results