Rotation, which changes shape about imgaug HOT 13 OPEN

class Rot90(Augmenter):
    def __init__(self, p=0, name=None, deterministic=False, random_state=None):
        super(Rot90, self).__init__(name=name, deterministic=deterministic, random_state=random_state)

        if ia.is_single_number(p):
            self.p = Binomial(p)
        elif isinstance(p, StochasticParameter):
            self.p = p
        else:
            raise Exception("Expected p to be int or float or StochasticParameter, got %s." % (type(p),))

    def _augment_images(self, images, random_state, parents, hooks):
        assert isinstance(image, list)
        nb_images = len(images)
        samples = self.p.draw_samples((nb_images,), random_state=random_state)
        for i in sm.xrange(nb_images):
            if samples[i] == 1:
                images[i] = np.rot90(images[i], k=1)
        return images

    def _augment_keypoints(self, keypoints_on_images, random_state, parents, hooks):
        nb_images = len(keypoints_on_images)
        samples = self.p.draw_samples((nb_images,), random_state=random_state)
        for i, keypoints_on_image in enumerate(keypoints_on_images):
            if samples[i] == 1:
                width, height = keypoints_on_image.shape[0:2]
                for keypoint in keypoints_on_image.keypoints:
                    keypoint.x, keypoint.y = keypoint.y, keypoint.x
                keypoints_on_image.shape = (width, height)
        return keypoints_on_images

    def get_parameters(self):
        return [self.p]

def _augment_keypoints(self, keypoints_on_images, random_state, parents, hooks):
        nb_images = len(keypoints_on_images)
        samples = self.p.draw_samples((nb_images,), random_state=random_state)
        for i, keypoints_on_image in enumerate(keypoints_on_images):
            if samples[i] == 1:
                width, height = keypoints_on_image.shape[0:2]
                for keypoint in keypoints_on_image.keypoints:
                    keypoint.x, keypoint.y = keypoint.y, width - keypoint.x
                keypoints_on_image.shape = (height, width)
        return keypoints_on_images