CoolFace
Apppublic

iti/HandMesh

sourceHugging Faceupdated 5y agoView on Hugging Face
0likes
augmentation.py258 linesDownload Raw Back to utils
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