Spaces:
Running
Running
File size: 4,322 Bytes
a0d91d3 |
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 |
import math
import random
import mmcv
import numpy as np
from ..builder import PIPELINES
import torch
from typing import Optional, Tuple, Union
@PIPELINES.register_module()
class Crop(object):
r"""Crop motion sequences.
Args:
crop_size (int): The size of the cropped motion sequence.
"""
def __init__(self,
crop_size: Optional[Union[int, None]] = None):
self.crop_size = crop_size
assert self.crop_size is not None
def __call__(self, results):
motion = results['motion']
length = len(motion)
if length >= self.crop_size:
idx = random.randint(0, length - self.crop_size)
motion = motion[idx: idx + self.crop_size]
results['motion_length'] = self.crop_size
else:
padding_length = self.crop_size - length
D = motion.shape[1:]
padding_zeros = np.zeros((padding_length, *D), dtype=np.float32)
motion = np.concatenate([motion, padding_zeros], axis=0)
results['motion_length'] = length
assert len(motion) == self.crop_size
results['motion'] = motion
results['motion_shape'] = motion.shape
if length >= self.crop_size:
results['motion_mask'] = torch.ones(self.crop_size).numpy()
else:
results['motion_mask'] = torch.cat(
(torch.ones(length), torch.zeros(self.crop_size - length))).numpy()
return results
def __repr__(self):
repr_str = self.__class__.__name__ + f'(crop_size={self.crop_size})'
return repr_str
@PIPELINES.register_module()
class RandomCrop(object):
r"""Random crop motion sequences. Each sequence will be padded with zeros to the maximum length.
Args:
min_size (int or None): The minimum size of the cropped motion sequence (inclusive).
max_size (int or None): The maximum size of the cropped motion sequence (inclusive).
"""
def __init__(self,
min_size: Optional[Union[int, None]] = None,
max_size: Optional[Union[int, None]] = None):
self.min_size = min_size
self.max_size = max_size
assert self.min_size is not None
assert self.max_size is not None
def __call__(self, results):
motion = results['motion']
length = len(motion)
crop_size = random.randint(self.min_size, self.max_size)
if length > crop_size:
idx = random.randint(0, length - crop_size)
motion = motion[idx: idx + crop_size]
results['motion_length'] = crop_size
else:
results['motion_length'] = length
padding_length = self.max_size - min(crop_size, length)
if padding_length > 0:
D = motion.shape[1:]
padding_zeros = np.zeros((padding_length, *D), dtype=np.float32)
motion = np.concatenate([motion, padding_zeros], axis=0)
results['motion'] = motion
results['motion_shape'] = motion.shape
if length >= self.max_size and crop_size == self.max_size:
results['motion_mask'] = torch.ones(self.max_size).numpy()
else:
results['motion_mask'] = torch.cat((
torch.ones(min(length, crop_size)),
torch.zeros(self.max_size - min(length, crop_size))), dim=0).numpy()
assert len(motion) == self.max_size
return results
def __repr__(self):
repr_str = self.__class__.__name__ + f'(min_size={self.min_size}'
repr_str += f', max_size={self.max_size})'
return repr_str
@PIPELINES.register_module()
class Normalize(object):
"""Normalize motion sequences.
Args:
mean_path (str): Path of mean file.
std_path (str): Path of std file.
"""
def __init__(self, mean_path, std_path, eps=1e-9):
self.mean = np.load(mean_path)
self.std = np.load(std_path)
self.eps = eps
def __call__(self, results):
motion = results['motion']
motion = (motion - self.mean) / (self.std + self.eps)
results['motion'] = motion
results['motion_norm_mean'] = self.mean
results['motion_norm_std'] = self.std
return results
|