'''
This code is based on https://github.com/jrieke/shape-detection/
'''

import matplotlib.pyplot as plt
import matplotlib
import numpy as np
import tensorflow as tf
import datetime

class JriekeBboxDataset:
    def generate(self):
        print('Generating...')

        self.WIDTH = 8
        self.HEIGHT = 8
        num_imgs = 50000
        min_object_size = 1
        max_object_size = 4
        num_objects = 1

        self.bboxes = np.zeros((num_imgs, num_objects, 4))
        self.imgs = np.zeros((num_imgs, self.WIDTH, self.HEIGHT))  # set background to 0

        for i_img in range(num_imgs):
            for i_object in range(num_objects):
                w, h = np.random.randint(min_object_size, max_object_size, size=2)
                x = np.random.randint(0, self.WIDTH - w)
                y = np.random.randint(0, self.HEIGHT - h)
                self.imgs[i_img, y:y+h, x:x+w] = 1.  # set rectangle to 1
                self.bboxes[i_img, i_object] = [x, y, w, h]

        print("Shapes: imgs ", self.imgs.shape, " bboxes ", self.bboxes.shape)

        #why this?
        # X = (self.imgs.reshape(num_imgs, -1) - np.mean(self.imgs)) / np.std(self.imgs)
        X = self.imgs

        y = self.bboxes.reshape(num_imgs, -1) / self.WIDTH

        # Split training and test.
        i = int(0.8 * num_imgs)
        train_X = X[:i] #80% for training
        test_X = X[i:]
        train_y = y[:i]
        test_y = y[i:]
        self.test_imgs = self.imgs[i:]
        self.test_bboxes = self.bboxes[i:]

        return train_X, train_y, test_X, test_y

    def check_dataset_image_compability(self, test_X_sample, test_imgs_sample):
        fig = plt.figure(figsize=(12, 3))
        fig.suptitle('check if the generated imgs match to the test_X slice image')
        fig.subplots_adjust(top=0.85)

        plt.subplot(1, 2, 1)
        plt.gca().set_title('Returned by the dataset class: used for training')
        plt.imshow(test_X_sample, cmap='Greys', interpolation='none', origin='lower', extent=[0, self.WIDTH, 0, self.HEIGHT])

        plt.subplot(1, 2, 2)
        plt.gca().set_title('Global image holder: used for plotting.')
        plt.imshow(test_imgs_sample, cmap='Greys', interpolation='none', origin='lower', extent=[0, self.WIDTH, 0, self.HEIGHT])
        plt.show()
        print('compare:',TMP,test_imgs_sample)

    def IOU(self,bbox1, bbox2):
        '''Calculate overlap between two bounding boxes [x, y, w, h] as the area of intersection over the area of unity'''
        x1, y1, w1, h1 = bbox1[0], bbox1[1], bbox1[2], bbox1[3]
        x2, y2, w2, h2 = bbox2[0], bbox2[1], bbox2[2], bbox2[3]

        w_I = min(x1 + w1, x2 + w2) - max(x1, x2)
        h_I = min(y1 + h1, y2 + h2) - max(y1, y2)
        if w_I <= 0 or h_I <= 0:  # no overlap
            return 0.
        I = w_I * h_I

        U = w1 * h1 + w2 * h2 - I

        return I / U

    def convertDefaultAnnotToCoord(self, annot):
        '''
        annot -> [x, y, w, h]
        '''

        w = annot[2] * self.WIDTH
        h = annot[3] * self.HEIGHT

        x = annot[0] * self.HEIGHT
        y = annot[1] * self.HEIGHT

        return [x,y,w,h]

    def convertYoloAnnotToCoord(self, yolo_annot):
        '''
        yolo_annot -> [x, y, w, h]
        '''

        w = yolo_annot[2] * self.WIDTH
        h = yolo_annot[3] * self.HEIGHT

        x = (yolo_annot[0] * self.WIDTH) - (w/2)
        y = (yolo_annot[1] * self.HEIGHT) - (h/2)

        return [x,y,w,h]

    def show_generated(self, i=0):
        fig = plt.figure()
        fig.subplots_adjust(top=0.85)
        fig.suptitle('Generated image sample + GT')
        plt.imshow(self.imgs[i], cmap='Greys', interpolation='none', origin='lower', extent=[0, self.WIDTH, 0, self.HEIGHT])
        for bbox in self.bboxes[i]:
            plt.gca().add_patch(matplotlib.patches.Rectangle((bbox[0], bbox[1]), bbox[2], bbox[3], ec='r', fc='none'))
            plt.gca().legend(['GT'])
        plt.show()

    def plot_rectangle(self, img, bbox):

        fig = plt.figure()
        fig.suptitle('Plotting rectangle.')
        fig.subplots_adjust(top=0.85)

        plt.subplot(1, 1, 1)
        plt.imshow(img, cmap='Greys', interpolation='none', origin='lower', extent=[0, self.WIDTH, 0, self.HEIGHT])
        plt.gca().add_patch(matplotlib.patches.Rectangle((bbox[0], bbox[1]), bbox[2], bbox[3], ec='r', fc='none'))
        plt.show()

    def check_dataset_image_compability(self, test_X_sample, test_imgs_sample):
        fig = plt.figure(figsize=(12, 3))
        fig.suptitle('check if the generated imgs match to the test_X slice image')
        fig.subplots_adjust(top=0.85)

        plt.subplot(1, 2, 1)
        plt.gca().set_title('Returned by the dataset class: used for training')
        plt.imshow(test_X_sample, cmap='Greys', interpolation='none', origin='lower', extent=[0, self.WIDTH, 0, self.HEIGHT])

        plt.subplot(1, 2, 2)
        plt.gca().set_title('Global image holder: used for plotting.')
        plt.imshow(test_imgs_sample, cmap='Greys', interpolation='none', origin='lower', extent=[0, self.WIDTH, 0, self.HEIGHT])
        plt.show()
        print('compare:',test_X_sample,test_imgs_sample)

    def show_predicted(self, pred_bboxes):
        # Show a few images and predicted bounding boxes from the test dataset.

        fig = plt.figure(figsize=(12, 3))
        fig.subplots_adjust(top=0.85)
        fig.suptitle('Prediction demonstration. Random samples.')
        legend_plotted = False

        for i_subplot in range(1, 11):
            plt.subplot(1, 10, i_subplot)
            i = np.random.randint(len(pred_bboxes))
            plt.imshow(self.test_imgs[i], cmap='Greys', interpolation='none', origin='lower', extent=[0, self.WIDTH, 0, self.HEIGHT])
            for pred_bbox, exp_bbox in zip(pred_bboxes[i], self.test_bboxes[i]):
                # print('before convertion: pred',pred_bbox, 'gt',exp_bbox)
                pred_bbox = self.convertDefaultAnnotToCoord(pred_bbox)
                # exp_bbox = self.convertDefaultAnnotToCoord(exp_bbox)
                print('after convertion: pred',pred_bbox, 'gt',exp_bbox)
                plt.gca().add_patch(matplotlib.patches.Rectangle((pred_bbox[0], pred_bbox[1]), pred_bbox[2], pred_bbox[3], ec='r', fc='none'))
                #gt
                plt.gca().add_patch(matplotlib.patches.Rectangle((exp_bbox[0], exp_bbox[1]), exp_bbox[2], exp_bbox[3], ec='b', fc='none'))
                plt.annotate('IOU: {:.2f}'.format(self.IOU(pred_bbox, exp_bbox)), (pred_bbox[0], pred_bbox[1]+pred_bbox[3]+0.2), color='r')
                if not legend_plotted:
                    legend_plotted = True
                    plt.gca().legend(['Pred','GT'],loc='upper center', bbox_to_anchor=(0.5, -0.5), fancybox=True)
        plt.show()
        # plt.savefig('plots/bw-single-rectangle_prediction_{0:%Y-%m-%d%H:%M:%S}.png'.format(datetime.datetime.now()), dpi=300)

[((22, 22, 22, 58), 'numpy.zeros', 'np.zeros', ({(22, 31, 22, 57): '(num_imgs, num_objects, 4)'}, {}), '((num_imgs, num_objects, 4))', True, 'import numpy as np\n'), ((23, 20, 23, 65), 'numpy.zeros', 'np.zeros', ({(23, 29, 23, 64): '(num_imgs, self.WIDTH, self.HEIGHT)'}, {}), '((num_imgs, self.WIDTH, self.HEIGHT))', True, 'import numpy as np\n'), ((53, 14, 53, 41), 'matplotlib.pyplot.figure', 'plt.figure', (), '', True, 'import matplotlib.pyplot as plt\n'), ((57, 8, 57, 28), 'matplotlib.pyplot.subplot', 'plt.subplot', ({(57, 20, 57, 21): '(1)', (57, 23, 57, 24): '(2)', (57, 26, 57, 27): '(1)'}, {}), '(1, 2, 1)', True, 'import matplotlib.pyplot as plt\n'), ((59, 8, 59, 125), 'matplotlib.pyplot.imshow', 'plt.imshow', (), '', True, 'import matplotlib.pyplot as plt\n'), ((61, 8, 61, 28), 'matplotlib.pyplot.subplot', 'plt.subplot', ({(61, 20, 61, 21): '(1)', (61, 23, 61, 24): '(2)', (61, 26, 61, 27): '(2)'}, {}), '(1, 2, 2)', True, 'import matplotlib.pyplot as plt\n'), ((63, 8, 63, 128), 'matplotlib.pyplot.imshow', 'plt.imshow', (), '', True, 'import matplotlib.pyplot as plt\n'), ((64, 8, 64, 18), 'matplotlib.pyplot.show', 'plt.show', ({}, {}), '()', True, 'import matplotlib.pyplot as plt\n'), ((109, 14, 109, 26), 'matplotlib.pyplot.figure', 'plt.figure', ({}, {}), '()', True, 'import matplotlib.pyplot as plt\n'), ((112, 8, 112, 124), 'matplotlib.pyplot.imshow', 'plt.imshow', (), '', True, 'import matplotlib.pyplot as plt\n'), ((116, 8, 116, 18), 'matplotlib.pyplot.show', 'plt.show', ({}, {}), '()', True, 'import matplotlib.pyplot as plt\n'), ((120, 14, 120, 26), 'matplotlib.pyplot.figure', 'plt.figure', ({}, {}), '()', True, 'import matplotlib.pyplot as plt\n'), ((124, 8, 124, 28), 'matplotlib.pyplot.subplot', 'plt.subplot', ({(124, 20, 124, 21): '(1)', (124, 23, 124, 24): '(1)', (124, 26, 124, 27): '(1)'}, {}), '(1, 1, 1)', True, 'import matplotlib.pyplot as plt\n'), ((125, 8, 125, 115), 'matplotlib.pyplot.imshow', 'plt.imshow', (), '', True, 'import matplotlib.pyplot as plt\n'), ((127, 8, 127, 18), 'matplotlib.pyplot.show', 'plt.show', ({}, {}), '()', True, 'import matplotlib.pyplot as plt\n'), ((130, 14, 130, 41), 'matplotlib.pyplot.figure', 'plt.figure', (), '', True, 'import matplotlib.pyplot as plt\n'), ((134, 8, 134, 28), 'matplotlib.pyplot.subplot', 'plt.subplot', ({(134, 20, 134, 21): '(1)', (134, 23, 134, 24): '(2)', (134, 26, 134, 27): '(1)'}, {}), '(1, 2, 1)', True, 'import matplotlib.pyplot as plt\n'), ((136, 8, 136, 125), 'matplotlib.pyplot.imshow', 'plt.imshow', (), '', True, 'import matplotlib.pyplot as plt\n'), ((138, 8, 138, 28), 'matplotlib.pyplot.subplot', 'plt.subplot', ({(138, 20, 138, 21): '(1)', (138, 23, 138, 24): '(2)', (138, 26, 138, 27): '(2)'}, {}), '(1, 2, 2)', True, 'import matplotlib.pyplot as plt\n'), ((140, 8, 140, 128), 'matplotlib.pyplot.imshow', 'plt.imshow', (), '', True, 'import matplotlib.pyplot as plt\n'), ((141, 8, 141, 18), 'matplotlib.pyplot.show', 'plt.show', ({}, {}), '()', True, 'import matplotlib.pyplot as plt\n'), ((147, 14, 147, 41), 'matplotlib.pyplot.figure', 'plt.figure', (), '', True, 'import matplotlib.pyplot as plt\n'), ((168, 8, 168, 18), 'matplotlib.pyplot.show', 'plt.show', ({}, {}), '()', True, 'import matplotlib.pyplot as plt\n'), ((126, 28, 126, 113), 'matplotlib.patches.Rectangle', 'matplotlib.patches.Rectangle', (), '', False, 'import matplotlib\n'), ((153, 12, 153, 41), 'matplotlib.pyplot.subplot', 'plt.subplot', ({(153, 24, 153, 25): '(1)', (153, 27, 153, 29): '(10)', (153, 31, 153, 40): 'i_subplot'}, {}), '(1, 10, i_subplot)', True, 'import matplotlib.pyplot as plt\n'), ((155, 12, 155, 133), 'matplotlib.pyplot.imshow', 'plt.imshow', (), '', True, 'import matplotlib.pyplot as plt\n'), ((27, 23, 27, 82), 'numpy.random.randint', 'np.random.randint', (), '', True, 'import numpy as np\n'), ((28, 20, 28, 56), 'numpy.random.randint', 'np.random.randint', ({(28, 38, 28, 39): '0', (28, 41, 28, 55): 'self.WIDTH - w'}, {}), '(0, self.WIDTH - w)', True, 'import numpy as np\n'), ((29, 20, 29, 57), 'numpy.random.randint', 'np.random.randint', ({(29, 38, 29, 39): '0', (29, 41, 29, 56): 'self.HEIGHT - h'}, {}), '(0, self.HEIGHT - h)', True, 'import numpy as np\n'), ((58, 8, 58, 17), 'matplotlib.pyplot.gca', 'plt.gca', ({}, {}), '()', True, 'import matplotlib.pyplot as plt\n'), ((62, 8, 62, 17), 'matplotlib.pyplot.gca', 'plt.gca', ({}, {}), '()', True, 'import matplotlib.pyplot as plt\n'), ((114, 32, 114, 117), 'matplotlib.patches.Rectangle', 'matplotlib.patches.Rectangle', (), '', False, 'import matplotlib\n'), ((126, 8, 126, 17), 'matplotlib.pyplot.gca', 'plt.gca', ({}, {}), '()', True, 'import matplotlib.pyplot as plt\n'), ((135, 8, 135, 17), 'matplotlib.pyplot.gca', 'plt.gca', ({}, {}), '()', True, 'import matplotlib.pyplot as plt\n'), ((139, 8, 139, 17), 'matplotlib.pyplot.gca', 'plt.gca', ({}, {}), '()', True, 'import matplotlib.pyplot as plt\n'), ((114, 12, 114, 21), 'matplotlib.pyplot.gca', 'plt.gca', ({}, {}), '()', True, 'import matplotlib.pyplot as plt\n'), ((115, 12, 115, 21), 'matplotlib.pyplot.gca', 'plt.gca', ({}, {}), '()', True, 'import matplotlib.pyplot as plt\n'), ((161, 36, 161, 141), 'matplotlib.patches.Rectangle', 'matplotlib.patches.Rectangle', (), '', False, 'import matplotlib\n'), ((163, 36, 163, 137), 'matplotlib.patches.Rectangle', 'matplotlib.patches.Rectangle', (), '', False, 'import matplotlib\n'), ((161, 16, 161, 25), 'matplotlib.pyplot.gca', 'plt.gca', ({}, {}), '()', True, 'import matplotlib.pyplot as plt\n'), ((163, 16, 163, 25), 'matplotlib.pyplot.gca', 'plt.gca', ({}, {}), '()', True, 'import matplotlib.pyplot as plt\n'), ((167, 20, 167, 29), 'matplotlib.pyplot.gca', 'plt.gca', ({}, {}), '()', True, 'import matplotlib.pyplot as plt\n')]