{"cells":[{"metadata":{"trusted":true},"cell_type":"code","source":"from keras.models import Sequential, Model, load_model\nfrom keras.layers import UpSampling2D, Conv2D, Activation, BatchNormalization, Reshape, Dense, Input, LeakyReLU, Dropout, Flatten, ZeroPadding2D\nfrom keras.optimizers import Adam\n\nimport glob\nfrom PIL import Image\nimport numpy as np\nimport os\nfrom ast import literal_eval\n\nfrom scipy.misc import imsave\nimport xml.etree.ElementTree as ET \nfrom tqdm import tqdm\nimport zipfile\nimport time","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"class DataLoader:\n    def __init__(self, root='../input/', image_size=(64, 64)):\n        self.root = root\n        self.breeds = os.listdir(self.root + 'annotation/Annotation/')\n        self.idxIn = 0\n        self.namesIn = []\n        self.image_size = image_size\n        self.imagesIn = np.zeros((25000, image_size[0], image_size[1], 3))\n    \n    def load_images(self):\n        start = time.time()\n        for breed in self.breeds:\n            for dog in os.listdir(self.root + 'annotation/Annotation/' + breed):\n                try:\n                    img = Image.open(self.root + 'all-dogs/all-dogs/' + dog + '.jpg')\n                except:\n                    continue\n                tree = ET.parse(self.root + 'annotation/Annotation/' + breed + '/' + dog)\n                root = tree.getroot()\n                objects = root.findall('object')\n                for o in objects:\n                    bndbox = o.find('bndbox') \n                    xmin = int(bndbox.find('xmin').text)\n                    ymin = int(bndbox.find('ymin').text)\n                    xmax = int(bndbox.find('xmax').text)\n                    ymax = int(bndbox.find('ymax').text)\n                    w = np.min((xmax - xmin, ymax - ymin))\n                    img2 = img.crop((xmin, ymin, xmin+w, ymin+w))\n                    img2 = img2.resize((self.image_size[0], self.image_size[1]), Image.ANTIALIAS)\n                    self.imagesIn[self.idxIn,:,:,:] = np.asarray(img2)\n                    if self.idxIn % 1000 == 0:\n                        print(self.idxIn)\n                    self.namesIn.append(breed)\n                    self.idxIn += 1\n        idx = np.arange(self.idxIn)\n        np.random.shuffle(idx)\n        self.imagesIn = self.imagesIn[idx,:,:,:]\n        self.namesIn = np.array(self.namesIn)[idx]\n\n        end = time.time()\n        print(end - start)\n\n        return self.imagesIn","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"class ImageZipper:\n    def __init__(self, folder='./', name='images.zip'):\n        self.folder = folder\n        self.z = zipfile.PyZipFile(self.folder + name, mode='w')\n    \n    def add_image(self, image, image_name):\n        file_name = self.folder + image_name + '.png'\n        imsave(file_name, image)\n        self.z.write(file_name)\n        os.remove(file_name)\n    \n    def close(self):\n        self.z.close()","execution_count":null,"outputs":[]},{"metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"cell_type":"code","source":"class DCGAN:\n    def __init__(self, discriminator_path, generator_path, output_directory, data_loader):\n        self.data_loader = data_loader\n        self.img_size = self.data_loader.image_size\n        self.upsample_layers = 5\n        self.starting_filters = 64\n        self.kernel_size = 3\n        self.channels = 3\n        self.discriminator_path = discriminator_path\n        self.generator_path = generator_path\n        self.output_directory = output_directory\n        self.generated_folder = f\"{self.output_directory}/generated_{self.img_size[0]}x{self.img_size[1]}\"\n\n    def build_generator(self):\n        noise_shape = (100,)\n\n        # This block of code can be a little daunting, but essentially it automatically calculates the required starting\n        # array size that will be correctly upscaled to our desired image size.\n        #\n        # We have 5 Upsample2D layers which each double the images width and height, so we can determine the starting\n        # x size by taking (x / 2^upsample_count) So for our target image size, 256x192, we do the following:\n        # x = (192 / 2^5), y = (256 / 2^5) [x and y are reversed within the model]\n        # We also need a 3rd dimension which is chosen relatively arbitrarily, in this case it's 64.\n        model = Sequential()\n        model.add(\n            Dense(self.starting_filters * (self.img_size[0] // (2 ** self.upsample_layers))  *  (self.img_size[1] // (2 ** self.upsample_layers)),\n                  activation=\"relu\", input_shape=noise_shape))\n        model.add(Reshape(((self.img_size[0] // (2 ** self.upsample_layers)),\n                           (self.img_size[1] // (2 ** self.upsample_layers)),\n                           self.starting_filters)))\n        model.add(BatchNormalization(momentum=0.8))\n\n        model.add(UpSampling2D())  # 6x8 -> 12x16\n        model.add(Conv2D(1024, kernel_size=self.kernel_size, padding=\"same\"))\n        model.add(Activation(\"relu\"))\n        model.add(BatchNormalization(momentum=0.8))\n\n        model.add(UpSampling2D())  # 12x16 -> 24x32\n        model.add(Conv2D(512, kernel_size=self.kernel_size, padding=\"same\"))\n        model.add(Activation(\"relu\"))\n        model.add(BatchNormalization(momentum=0.8))\n\n        model.add(UpSampling2D())  # 24x32 -> 48x64\n        model.add(Conv2D(256, kernel_size=self.kernel_size, padding=\"same\"))\n        model.add(Activation(\"relu\"))\n        model.add(BatchNormalization(momentum=0.8))\n\n        model.add(UpSampling2D())  # 48x64 -> 96x128\n        model.add(Conv2D(128, kernel_size=self.kernel_size, padding=\"same\"))\n        model.add(Activation(\"relu\"))\n        model.add(BatchNormalization(momentum=0.8))\n\n        model.add(UpSampling2D())  # 96x128 -> 192x256\n        model.add(Conv2D(64, kernel_size=self.kernel_size, padding=\"same\"))\n        model.add(Activation(\"relu\"))\n        model.add(BatchNormalization(momentum=0.8))\n\n        model.add(Conv2D(32, kernel_size=self.kernel_size, padding=\"same\"))\n        model.add(Activation(\"relu\"))\n        model.add(BatchNormalization(momentum=0.8))\n\n        model.add(Conv2D(self.channels, kernel_size=self.kernel_size, padding=\"same\"))\n        model.add(Activation(\"tanh\"))\n\n#         model.summary()\n\n        noise = Input(shape=noise_shape)\n        img = model(noise)\n\n        return Model(noise, img)\n\n    def build_discriminator(self):\n        img_shape = (self.img_size[0], self.img_size[1], self.channels)\n\n        model = Sequential()\n\n        model.add(Conv2D(32, kernel_size=self.kernel_size, strides=2, input_shape=img_shape, padding=\"same\"))  # 192x256 -> 96x128\n        model.add(LeakyReLU(alpha=0.2))\n        model.add(Dropout(0.25))\n\n        model.add(Conv2D(64, kernel_size=self.kernel_size, strides=2, padding=\"same\"))  # 96x128 -> 48x64\n        model.add(ZeroPadding2D(padding=((0, 1), (0, 1))))\n        model.add(LeakyReLU(alpha=0.2))\n        model.add(Dropout(0.25))\n        model.add(BatchNormalization(momentum=0.8))\n\n        model.add(Conv2D(128, kernel_size=self.kernel_size, strides=2, padding=\"same\"))  # 48x64 -> 24x32\n        model.add(LeakyReLU(alpha=0.2))\n        model.add(Dropout(0.25))\n        model.add(BatchNormalization(momentum=0.8))\n\n        model.add(Conv2D(256, kernel_size=self.kernel_size, strides=1, padding=\"same\"))  # 24x32 -> 12x16\n        model.add(LeakyReLU(alpha=0.2))\n        model.add(Dropout(0.25))\n\n        model.add(Conv2D(512, kernel_size=self.kernel_size, strides=1, padding=\"same\"))  # 12x16 -> 6x8\n        model.add(LeakyReLU(alpha=0.2))\n        model.add(Dropout(0.25))\n\n        model.add(Flatten())\n        model.add(Dense(1, activation='sigmoid'))\n\n#         model.summary()\n\n        img = Input(shape=img_shape)\n        validity = model(img)\n\n        return Model(img, validity)\n\n    def build_gan(self):\n        optimizer = Adam(0.0002, 0.5)\n\n        # See if the specified model paths exist, if they don't then we start training new models\n\n        if self.discriminator_path is not None and self.generator_path is not None and os.path.exists(self.discriminator_path) and os.path.exists(self.generator_path):\n            self.discriminator = load_model(self.discriminator_path)\n            self.generator = load_model(self.generator_path)\n            print(\"Loaded models...\")\n        else:\n            self.discriminator = self.build_discriminator()\n            self.discriminator.compile(loss='binary_crossentropy',\n                                       optimizer=optimizer,\n                                       metrics=['accuracy'])\n\n            self.generator = self.build_generator()\n            self.generator.compile(loss='binary_crossentropy', optimizer=optimizer)\n\n        # These next few lines setup the training for the GAN model\n        z = Input(shape=(100,))\n        img = self.generator(z)\n\n        self.discriminator.trainable = False\n\n        valid = self.discriminator(img)\n\n        self.combined = Model(z, valid)\n        self.combined.compile(loss='binary_crossentropy', optimizer=optimizer)\n\n    def train(self, epochs, batch_size=32, save_interval=50):\n        self.build_gan()\n        X_train = self.data_loader.load_images()\n        print(\"Training Data Shape: \", X_train.shape)\n\n        # Rescale images from -1 to 1\n        X_train = (X_train.astype(np.float32) - 127.5) / 127.5\n\n        half_batch = batch_size // 2\n\n        for epoch in range(epochs):\n\n\n            # Train Generator\n            noise = np.random.normal(0, 1, (batch_size, 100))\n            g_loss = self.combined.train_on_batch(noise, np.ones((batch_size, 1)))\n\n\n\n            # Train Discriminator\n            idx = np.random.randint(0, X_train.shape[0], half_batch)\n            imgs = X_train[idx]\n\n            # Sample noise and generate a half batch of new images\n            noise = np.random.normal(0, 1, (half_batch, 100))\n            gen_imgs = self.generator.predict(noise)\n\n            # Train the discriminator (real classified as ones and generated as zeros)\n            d_loss_real = self.discriminator.train_on_batch(imgs, np.ones((half_batch, 1)))\n            d_loss_fake = self.discriminator.train_on_batch(gen_imgs, np.zeros((half_batch, 1)))\n            d_loss = 0.5 * np.add(d_loss_real, d_loss_fake)\n\n            # If at save interval => save generated image samples, save model files\n            if epoch % (save_interval) == 0:\n                \n                # Print progress\n                print(f\"{epoch} [D loss: {d_loss[0]} | D Accuracy: {100 * d_loss[1]}] [G loss: {g_loss}]\")\n\n#                 save_path = self.output_directory\n#                 if not os.path.exists(save_path):\n#                     os.makedirs(save_path)\n#                 self.discriminator.save(save_path + \"/discrim.h5\")\n#                 self.generator.save(save_path + \"/generat.h5\")\n\n    def gene_imgs(self, count):\n        # Generate images from the currently loaded model\n        noise = np.random.normal(0, 1, (count, 100))\n        return self.generator.predict(noise)\n    \n    def generate_imgs(self, count, threshold):\n        # Generates (count) images from the model ensuring the discriminator scores them between the threshold values\n        # and saves them\n\n        imgs = []\n        for i in tqdm(range(count), desc='Generating images'):\n            score = [0]\n            while not(threshold[0] < score[0] < threshold[1]):\n                img = self.gene_imgs(1)\n                score = self.discriminator.predict(img)\n            imgs.append(img)\n\n        imgs = np.asarray(imgs).squeeze()\n        imgs = 0.5 * imgs + 0.5\n\n        return imgs","execution_count":null,"outputs":[]},{"metadata":{"_uuid":"d629ff2d2480ee46fbb7e2d37f6b5fab8052498a","_cell_guid":"79c7e3d0-c299-4dcb-8224-4455121ee9b0","trusted":true},"cell_type":"code","source":"# Number of epochs to train for\nepochs = 530000\n# Number of images to train on at once\nbatch_size = 24\n# How many epochs to go between saves/outputs\nsave_interval = 1000\n# Directoy to save weights and images to.\noutput_directory = './'\n\ndata_loader = DataLoader()\ndcgan = DCGAN(None, None, output_directory, data_loader)\nstart = time.time()\ndcgan.train(epochs=epochs, batch_size=batch_size, save_interval=save_interval)\nend = time.time()\nprint(end - start)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# Path to existing generator weights file\nload_generator = \"./generat.h5\"\n# Path to existing discriminator weights file\nload_discriminator = \"./discrim.h5\"\n\n# generate that many samples from existing model\nsample = 10000\n# The values between which a generated image must score from the discriminator\nsample_thresholds = (0.0, 1.0)\ndcgan = DCGAN(load_discriminator, load_generator, output_directory, data_loader)\ndcgan.build_gan()\n\nzipper = ImageZipper(folder=output_directory)\nfor i in range(0, sample, 100):\n    print('Images generated', i)\n    images = dcgan.generate_imgs(100, sample_thresholds)\n    for index, image in enumerate(images):\n        zipper.add_image(image, str(i + index))\n\nzipper.close()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"import os\nos.listdir('.')","execution_count":null,"outputs":[]}],"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"pygments_lexer":"ipython3","nbconvert_exporter":"python","version":"3.6.4","file_extension":".py","codemirror_mode":{"name":"ipython","version":3},"name":"python","mimetype":"text/x-python"}},"nbformat":4,"nbformat_minor":1}