iti/HandMesh
0
1import numpy as np2from numpy import random3from torchvision import transforms4import cv25 6 7class Compose(object):8 """Composes several augmentations together.9 Args:10 transforms (List[Transform]): list of transforms to compose.11 Example:12 >>> augmentations.Compose([13 >>> transforms.CenterCrop(10),14 >>> transforms.ToTensor(),15 >>> ])16 """17 18 def __init__(self, transforms):19 self.transforms = transforms20 21 def __call__(self, img):22 for t in self.transforms:23 img = t(img)24 return img25 26 27class ConvertFromInts(object):28 def __call__(self, image):29 return image.astype(np.float32)30 31 32class BaseTransform(object):33 def __init__(self, size, mean, std):34 self.mean = np.array(mean, dtype=np.float32)35 self.std = std36 self.size = size37 38 def __call__(self, image):39 image = cv2.resize(image, (self.size, self.size)).astype(np.float32)40 image -= self.mean41 image /= self.std42 image = image.transpose(2, 0, 1)43 44 return image45 46 47class RandomSaturation(object):48 def __init__(self, lower=0.5, upper=1.5):49 self.lower = lower50 self.upper = upper51 assert self.upper >= self.lower, "contrast upper must be >= lower."52 assert self.lower >= 0, "contrast lower must be non-negative."53 54 def __call__(self, image):55 if random.randint(2):56 image[:, :, 1] *= random.uniform(self.lower, self.upper)57 58 return image59 60 61class RandomHue(object):62 def __init__(self, delta=18.0):63 assert delta >= 0.0 and delta <= 360.064 self.delta = delta65 66 def __call__(self, image):67 if random.randint(2):68 image[:, :, 0] += random.uniform(-self.delta, self.delta)69 image[:, :, 0][image[:, :, 0] > 360.0] -= 360.070 image[:, :, 0][image[:, :, 0] < 0.0] += 360.071 return image72 73 74class RandomLightingNoise(object):75 def __init__(self):76 self.perms = ((0, 1, 2), (0, 2, 1),77 (1, 0, 2), (1, 2, 0),78 (2, 0, 1), (2, 1, 0))79 80 def __call__(self, image):81 if random.randint(2):82 swap = self.perms[random.randint(len(self.perms))]83 shuffle = SwapChannels(swap) # shuffle channels84 image = shuffle(image)85 return image86 87 88class ConvertColor(object):89 def __init__(self, current='RGB', transform='HSV'):90 self.transform = transform91 self.current = current92 93 def __call__(self, image):94 if self.current == 'RGB' and self.transform == 'HSV':95 image = cv2.cvtColor(image, cv2.COLOR_RGB2HSV)96 elif self.current == 'HSV' and self.transform == 'RGB':97 image = cv2.cvtColor(image, cv2.COLOR_HSV2RGB)98 else:99 raise NotImplementedError100 return image101 102 103class RandomContrast(object):104 def __init__(self, lower=0.5, upper=1.5):105 self.lower = lower106 self.upper = upper107 assert self.upper >= self.lower, "contrast upper must be >= lower."108 assert self.lower >= 0, "contrast lower must be non-negative."109 110 # expects float image111 def __call__(self, image):112 if random.randint(2):113 alpha = random.uniform(self.lower, self.upper)114 image *= alpha115 return image116 117 118class RandomBrightness(object):119 def __init__(self, delta=32):120 assert delta >= 0.0121 assert delta <= 255.0122 self.delta = delta123 124 def __call__(self, image):125 if random.randint(2):126 delta = random.uniform(-self.delta, self.delta)127 image += delta128 return image129 130 131class SwapChannels(object):132 """Transforms a tensorized image by swapping the channels in the order133 specified in the swap tuple.134 Args:135 swaps (int triple): final order of channels136 eg: (2, 1, 0)137 """138 139 def __init__(self, swaps):140 self.swaps = swaps141 142 def __call__(self, image):143 """144 Args:145 image (Tensor): image tensor to be transformed146 Return:147 a tensor with channels swapped according to swap148 """149 # if torch.is_tensor(image):150 # image = image.data.cpu().numpy()151 # else:152 # image = np.array(image)153 image = image[:, :, self.swaps]154 return image155 156 157class PhotometricDistort(object):158 def __init__(self):159 self.pd = [160 RandomContrast(),161 ConvertColor(transform='HSV'),162 RandomSaturation(),163 RandomHue(),164 ConvertColor(current='HSV', transform='RGB'),165 RandomContrast()166 ]167 self.rand_brightness = RandomBrightness()168 # self.rand_light_noise = RandomLightingNoise()169 170 def __call__(self, image):171 im = image.copy()172 im = self.rand_brightness(im)173 if random.randint(2):174 distort = Compose(self.pd[:-1])175 else:176 distort = Compose(self.pd[1:])177 im = distort(im)178 return im179 # return self.rand_light_noise(im)180 181 182class Augmentation(object):183 def __init__(self, size=224):184 # self.mean = mean185 # self.std = std186 self.size = size187 self.augment = Compose([188 ConvertFromInts(),189 PhotometricDistort(),190 #BaseTransform(self.size, self.mean, self.std)191 ])192 193 def __call__(self, img):194 return self.augment(img)195 196 197def crop_roi(img, bbox, out_sz, padding=(0, 0, 0)):198 bbox = [float(x) for x in bbox]199 a = (out_sz - 1) / (bbox[2] - bbox[0])200 b = (out_sz - 1) / (bbox[3] - bbox[1])201 c = -a * bbox[0]202 d = -b * bbox[1]203 mapping = np.array([[a, 0, c],204 [0, b, d]]).astype(np.float)205 crop = cv2.warpAffine(img, mapping, (out_sz, out_sz),206 borderMode=cv2.BORDER_CONSTANT,207 borderValue=padding)208 return crop209 210 211def crop_pad_im_from_bounding_rect(im, bb):212 """213 :param im: H x W x C214 :param bb: x, y, w, h (may exceed the image region)215 :return: cropped image216 """217 crop_im = im[max(0, bb[1]):min(bb[1] + bb[3], im.shape[0]), max(0, bb[0]):min(bb[0] + bb[2], im.shape[1]), :]218 219 if bb[1] < 0:220 crop_im = cv2.copyMakeBorder(crop_im, -bb[1], 0, 0, 0, # top, bottom, left, right, bb[3]-crop_im.shape[0]221 borderType=cv2.BORDER_CONSTANT, value=(0, 0, 0, 0))222 if bb[1] + bb[3] > im.shape[0]:223 crop_im = cv2.copyMakeBorder(crop_im, 0, bb[1] + bb[3] - im.shape[0], 0, 0,224 borderType=cv2.BORDER_CONSTANT, value=(0, 0, 0, 0))225 226 if bb[0] < 0:227 crop_im = cv2.copyMakeBorder(crop_im, 0, 0, -bb[0], 0, # top, bottom, left, right228 borderType=cv2.BORDER_CONSTANT, value=(0, 0, 0, 0))229 if bb[0] + bb[2] > im.shape[1]:230 crop_im = cv2.copyMakeBorder(crop_im, 0, 0, 0, bb[0] + bb[2] - im.shape[1],231 borderType=cv2.BORDER_CONSTANT, value=(0, 0, 0, 0))232 return crop_im233 234 235def rotate(img, mapping, padding=(0, 0, 0)):236 # mapping = cv2.getRotationMatrix2D((img.shape[1] // 2, img.shape[0] // 2), angle, 1.0) # 12237 rotated = cv2.warpAffine(img, mapping, (img.shape[1], img.shape[0]),238 borderMode=cv2.BORDER_CONSTANT,239 borderValue=padding)240 return rotated241 242 243def get_m1to1_gaussian_rand(scale):244 r = 2245 while r < -1 or r > 1:246 r = np.random.normal(scale=scale)247 248 return r249 250 251if __name__ == '__main__':252 img = cv2.imread('../data/FreiHAND/data/evaluation/rgb/00000001.jpg')253 img = crop_roi(img, (112-50*1.3, 112-50*1.3, 112+50*1.3, 112+50*1.3), 224)254 img, mapping = rotate(img, 30)255 cv2.imshow('test', img)256 print(mapping)257 cv2.waitKey(0)258 