DataGenerator

class DataGenerator(tf.keras.utils.Sequence):
    #'Generates data for Keras'

    def __init__(self,
                 x = None,
                 y = None,
                 batch_size = 20,
                 target_shape = (100, 100, 3),
                 sample_std = False,
                 feature_std = False,
                 proj_parameters = None,
                 blur_parameters = None,
                 nois_parameters = None,
                 flip_parameters = None,
                 gamm_parameters = None):

        #'Initialization'
        self.x = np.array(x)
        self.y = np.array(y)
        self.batch_size = batch_size
        self.target_shape = target_shape
        self.sample_std = sample_std
        self.feature_std = feature_std
        self.proj_parameters = proj_parameters
        self.blur_parameters = blur_parameters
        self.nois_parameters = nois_parameters
        self.flip_parameters = flip_parameters
        self.gamm_parameters = gamm_parameters

    def __len__(self):
        #'Denotes the number of batches per epoch'
        return len(self.x) // self.batch_size

    def __getitem__(self, index):
        #'Generate one batch of data'
        # Generate indexes of the batch
        indexes = np.random.randint(len(self.x), size = self.batch_size)

        # Generate data
        x, y = self.__data_generation(self.x[indexes], self.y[indexes])

        return x, y

    def on_epoch_end(self):
        pass

    def __data_generation(self, imgs, key_points):
        x = []
        y = []

        # Generate data
        for img, points in zip(imgs, key_points):
            img, points = self.__preprocess_input(np.copy(img), np.copy(points))
            x += [img]
            y += [points]

        x = np.array(x)
        if self.sample_std:
            batch_x = self.smap_standardize(x)

        y = np.array(y)
        return x, y

    def __preprocess_input(self, image, key_points):

        if self.proj_parameters is not None:
            image, key_points = self.__projection(image, key_points, proj_parameters)#0.05, 0.35, 25, 0.1)

        if image.shape != self.target_shape:
            image, key_points = self.__resize(image, key_points, self.target_shape)

        if self.flip_parameters is not None:
            image, key_points = self.__flip(image, key_points, self.flip_parameters)#0.5)
        if self.gamm_parameters is not None:
            image = self.__gamma(image, self.gamm_parameters)#1/2, 1.7)
        if self.blur_parameters is not None:
            image = self.__gaussian_blur(image, self.blur_parameters)#1.8)
        if self.nois_parameters is not None:
            image = self.__gaussian_noise(image, self.nois_parameters)#0, 0.06)
        if self.feature_std:
            image = self.__feat_standardize(image)
        return image, key_points

    def __gaussian_noise(self, image, params):
        mean, sigma = params
        image = image + np.random.normal(mean * random(), sigma * random(), image.shape)
        image = np.clip(image, 0.0, 1.0)
        return image

    def __gaussian_blur(self, image, params):
        sigma = params
        image = skimage.filters.gaussian(image, sigma = sigma * random(), multichannel = True)
        return image

    def __gamma(self, image, params):
        min_gamma, max_gamma = params
        d = max_gamma - min_gamma
        image = skimage.exposure.adjust_gamma(image, min_gamma + d * random())
        return image

    def __flip(self, image, key_points, params):
        probability = params
        if probability < random():
            image, key_points = mirr(image, key_points)
        return image, key_points

    def __projection(self, image, key_points, params):
        low, high, max_angle, alpha = params
        #print(key_points)
        try_count = 200
        for i in range(try_count):
            height, width = image.shape[0], image.shape[1]
            corners = np.array([[0, 0],
                                [height, 0],
                                [height, width],
                                [0, width]])
            d_h = low * height + high * height
            d_w = low * width + high * width
            s_h = -high * height
            s_w = -high * width
            edges =   np.array([[         s_h + d_h * random(),         s_w + d_w * random()],
                                [height - s_h - d_h * random(),         s_w + d_w * random()],
                                [height - s_h - d_h * random(), width - s_w - d_w * random()],
                                [         s_h + d_h * random(), width - s_w - d_w * random()]])
            theta = np.radians(max_angle * random())
            if i == try_count - 1:
                edges = np.copy(corners)
                theta = 0

            halflen = corners.shape[0]

            ax = np.array([-corners[..., 0], -corners[..., 1], -np.ones(halflen),
                           np.zeros(halflen), np.zeros(halflen), np.zeros(halflen),
                           edges[..., 0] * corners[..., 0], edges[..., 0] * corners[..., 1], edges[..., 0]]).T

            ay = np.array([np.zeros(halflen), np.zeros(halflen), np.zeros(halflen),
                           -corners[..., 0], -corners[..., 1], -np.ones(halflen),
                           edges[..., 1] * corners[..., 0], edges[..., 1] * corners[..., 1], edges[..., 1]]).T
            #print(ax.shape)
            A = np.zeros((halflen * 2, 9))
            A[::2, ...] = ax
            A[1::2, ...] = ay

            u, s, v = np.linalg.svd(A)

            H = v[-1, ...].reshape((3,3))[(1, 0, 2), :][:, (1, 0, 2)]
            #print(H)

            c, s = np.cos(theta), np.sin(theta)
            R = np.array(((c,-s,0), (s, c,0), (0,0,1)))

            #print(H)
            H = H @ R
            #print(H)
            if np.linalg.det(H) == 0:
                continue

            points = np.zeros((key_points.shape[0] // 2, 2))

            points[..., 0] = key_points[::2]
            points[..., 1] = key_points[1::2]
            test_points = np.zeros(key_points.shape)
            #print(key_points)
            points = skimage.transform.ProjectiveTransform(H)(points)
            test_points[::2]  = points[..., 0]
            test_points[1::2] = points[..., 1]
            #print(key_points)
            low_keys = np.ones(test_points.shape)
            high_keys = np.ones(test_points.shape)
            low_keys[::2] = alpha * height * low_keys[::2]
            low_keys[1::2] = alpha * width * low_keys[1::2]
            high_keys[::2] = (1 - alpha) * height * high_keys[::2]
            high_keys[1::2] = (1 - alpha) * width * high_keys[1::2]
            #print(low,high,test_points, np.greater(key_points, low).all(), np.less(key_points, high).all())
            if np.greater(test_points, low_keys).all() and np.less(test_points, high_keys).all():
                break

        key_points = test_points
        inv = np.linalg.inv(H)
        transf = skimage.transform.ProjectiveTransform(inv)
        image = skimage.transform.warp(image, transf, output_shape = image.shape, mode='edge')
        return image, key_points

    def __resize(self, image, key_points, shape):
        h, w = image.shape[1], image.shape[0]
        image = skimage.transform.resize(image, shape)
        points = np.zeros(key_points.shape, dtype=np.int64)
        points[ ::2] = (key_points[ ::2] / h * 100).round()
        points[1::2] = (key_points[1::2] / w * 100).round()
        return image, points

    def fit(self, imgs, passes):
        imgs = np.array(imgs)
        #i = np.copy(imgs)

        for i in range(passes):
            for j, img in enumerate(imgs):
                imgs[j], _ = self.preprocess_input(img, np.zeros(28))

            imgs = np.array(imgs)
            self.std = np.std(imgs, axis = 0)
            self.mean = np.mean(imgs, axis = 0)

            for j, img in enumerate(imgs):
                imgs[j] -= (self.mean)
                imgs[j] /= (self.std + 1e-6)
                low = np.min(imgs[j])
                high = np.max(imgs[j] - low)
                imgs[j] = (imgs[j] - low) / high
                #imgs[j] = np.clip(imgs[j], 0.0, 1.0)

        self.fited = True

    def feat_standardize(self, image):
        if self.fited:
            #print("lel")
            image -= self.mean
            image /= (self.std + 1e-6)
            low = np.min(image)
            high = np.max(image - low)
            image = (image - low) / high #np.clip(image, 0.0, 1.0)

        return image

    def smap_standardize(self, images):
        images -= np.mean(images, keepdims=True)
        images /= (np.std(images, keepdims=True) + 1e-6)
        return images