{"cells":[{"metadata":{"colab":{},"colab_type":"code","id":"KEzXRu_kWzbO","trusted":true},"cell_type":"code","source":"import torch\nimport torch.nn as nn\nimport torch.optim as optim\nimport torch.nn.functional as F\nimport numpy as np\nimport matplotlib.pyplot as plt\nfrom torchvision import datasets\nfrom torchvision.transforms import transforms\nfrom torch.utils.data import DataLoader\n%matplotlib inline\nplt.rcParams['image.interpolation'] = 'nearest'\nfrom IPython.display import clear_output","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"import os\nimport PIL\nimport xml.etree.ElementTree as ET\n   \ndef get_all_paths(root_directory):\n    list_all_paths = []\n    for root, dirs,filenames in os.walk(root_directory):\n        for f in filenames:\n            list_all_paths.append(os.path.abspath(os.path.join(root, f)))\n    return list_all_paths\n\ndef make_folders(path_folder):\n    if not os.path.exists(path_folder):\n        os.makedirs(path_folder)\n        print(\"Folder: \" + path_folder + \" created\")\n        \ndef dogs_only_from_annotations(path_to_data_folder, path_to_annotation_folder, path_to_dogs_folder):\n    exceptions = {}\n    annotation_directories = get_sub_directories(path_to_annotation_folder)\n    for subdir in annotation_directories:\n        clear_output()\n        print(\"Processing Directory: \" + subdir)\n        make_folders(os.path.join(path_to_dogs_folder, subdir))\n        files = get_files_in_sub_directory(path_to_annotation_folder, subdir)\n        for f in files:\n            # No .xml extention (here plain text), good for open image project.\n            basename =  os.path.splitext(f)[0] \n            try:\n                objects = get_all_objects(os.path.join(path_to_annotation_folder, subdir, basename))\n                for i, obj in enumerate(objects):\n                    xmin, ymin, xmax, ymax = obj\n                    image = PIL.Image.open(os.path.join(path_to_data_folder, basename + \".jpg\"))\n                    cropped = image.crop((xmin, ymin, xmax, \n                                          ymax)).save(os.path.join(path_to_dogs_folder, subdir, \n                                                                   \"cropped_\" + basename + \n                                                                   str(i) + \".jpg\"),\n                                                      \"JPEG\")\n            except Exception as e:\n                exceptions[str(e)] = os.path.join(path_to_annotation_folder, subdir, basename)\n    return exceptions\n\ndef parse_annotation_file(file_name):\n    root = ET.parse(file_name).getroot()\n    for child in root:\n        print(child.tag, child.attrib)\n        \ndef print_annotation_file(file_name):\n    root = ET.parse(file_name).getroot()        \n    print(ET.tostring(root, encoding='utf8').decode('utf8'))\n    \ndef get_all_objects(file_path):\n    bbxs = []\n    root = ET.parse(file_path).getroot()\n    for obj in root.findall(\"object\"):\n        bndbox = obj.find(\"bndbox\")\n        bbxs.append([int(it.text) for it in bndbox])\n    return bbxs\n\ndef get_sub_directories(directory):\n    return sorted([name for name in os.listdir(directory) if os.path.isdir(os.path.join(directory, name))])\n\ndef get_files_in_sub_directory(parent_directory, sub_directory):\n    return os.listdir(os.path.join(parent_directory, sub_directory))\n\ndef get_images(imgs_folder, i, j):\n    imgs = os.listdir(imgs_folder)[i:j]\n    f, axes = plt.subplots(nrows=3, ncols=len(imgs) // 3 + 1, figsize=(20, 5))\n    axes = axes.ravel()\n    for ax in axes:\n        ax.axis('off')\n    for i, title in enumerate(imgs):\n        image = PIL.Image.open(os.path.join(imgs_folder, title))\n        axes[i].imshow(image, cmap='gray')\n        axes[i].set_title(title)  \n    plt.tight_layout()\n    plt.show()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"data_folder = \"../input/all-dogs/all-dogs\"\nannotation_folder = \"../input/annotation/Annotation\"","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"get_images(data_folder, 600, 620)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"dogs_folder = \"../Dogs-Cropped-Images\"\nout_folder = \"../output_images\"\nmake_folders(dogs_folder)\nmake_folders(out_folder)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"os.listdir(\"../\")","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"exceptions = dogs_only_from_annotations(data_folder, annotation_folder, dogs_folder)","execution_count":null,"outputs":[]},{"metadata":{"colab":{},"colab_type":"code","id":"qXAMKb1OMQ9P","trusted":true},"cell_type":"code","source":"batch_size = 128\nz_size = 100","execution_count":null,"outputs":[]},{"metadata":{"colab":{},"colab_type":"code","id":"O-f-ZLpfJ08W","trusted":true},"cell_type":"code","source":"def get_dataloader(batch_size, image_size, data_dir='../Dogs-Cropped-Images'):\n    \"\"\"\n    Batch the neural network data using DataLoader\n    :param batch_size: The size of each batch; the number of images in a batch\n    :param img_size: The square size of the image data (x, y)\n    :param data_dir: Directory where image data is located\n    :return: DataLoader with batched data\n    \"\"\"\n    transform = transforms.Compose([transforms.Resize((image_size, image_size)), \n                                transforms.ToTensor(), \n                                transforms.Normalize([0.5, 0.5, 0.5],\n                                                    [0.5, 0.5, 0.5])])\n                              \n                                \n    dataset=datasets.ImageFolder(data_dir,transform)\n    dataloader=torch.utils.data.DataLoader(dataset=dataset,batch_size=batch_size)\n    # TODO: Implement function and return a dataloader\n    \n    return dataloader\n","execution_count":null,"outputs":[]},{"metadata":{"colab":{},"colab_type":"code","id":"h0DP-XgmMX70","trusted":true},"cell_type":"code","source":"# Define function hyperparameters\n \nimg_size = 64\n\n\"\"\"\nDON'T MODIFY ANYTHING IN THIS CELL THAT IS BELOW THIS LINE\n\"\"\"\n# Call your function and get a dataloader\nceleba_train_loader = get_dataloader(batch_size, img_size)\n","execution_count":null,"outputs":[]},{"metadata":{"colab":{"base_uri":"https://localhost:8080/","height":267},"colab_type":"code","id":"i0OGWC8AMnE2","outputId":"6c3179c2-bd5c-4919-98ec-fbb9174410c8","scrolled":true,"trusted":true},"cell_type":"code","source":"# helper display function\ndef imshow(img):\n    npimg = img.numpy()\n    npimg = npimg *0.5\n    npimg +=0.5\n    plt.imshow(np.transpose(npimg, (1, 2, 0)))\n\n\"\"\"\nDON'T MODIFY ANYTHING IN THIS CELL THAT IS BELOW THIS LINE\n\"\"\"\n# obtain one batch of training images\ndataiter = iter(celeba_train_loader)\nimages, _ = dataiter.next() # _ for no labels\n\n# plot the images in the batch, along with the corresponding labels\nfig = plt.figure(figsize=(20, 4))\nplot_size=20\nfor idx in np.arange(plot_size):\n    ax = fig.add_subplot(2, plot_size/2, idx+1, xticks=[], yticks=[])\n    imshow(images[idx])","execution_count":null,"outputs":[]},{"metadata":{"colab":{},"colab_type":"code","id":"t3JKkRVXrhnq","trusted":true},"cell_type":"code","source":"def init_weight(m):\n  classname = m.__class__.__name__\n  if classname.find('Conv') != -1:\n    m.weight.data.normal_(0.0, 0.02)\n  elif classname.find('BatchNorm') != -1:\n    m.weight.data.normal_(0.1, 0.02)\n    m.bias.data.fill_(0)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def mila(input,beta=1.0):\n    return input * torch.tanh(F.softplus(input+beta))","execution_count":null,"outputs":[]},{"metadata":{"colab":{},"colab_type":"code","id":"BFXGkMeHXje-","trusted":true},"cell_type":"code","source":"# Discriminator\n\n# Probably a VGG16 or VGG19 for Simple Image Classification pretrained on ImageNet\n\nclass Discriminator(nn.Module):\n    \n    def __init__(self, inhw, c1_channels=64, c2_channels=128, c3_channels=256,\n                 c4_channels=512, i_channels_in_2=True):\n        '''\n        The constructor method for the Discriminator class\n        \n        Arguments:\n        - inhw : The number of \n        - c1_channels : the number of output channels from the\n                        first Convolutional Layer [Default - 128]\n                        \n        - c2_channels : the number of output channels from the\n                        second Convolutional Layer [Default - 256]\n                        \n        - c3_channels : the number of output channels from the\n                        third Convolutional Layer [Default - 512]\n        \n        - i_channels_in_2 : Increase the number of channels by 2\n                        in each layer.\n        '''\n        \n        super().__init__()\n        \n        # Define the class variables\n        self.c1_channels = c1_channels\n        \n        if i_channels_in_2:\n            self.c2_channels = self.c1_channels * 2\n            self.c3_channels = self.c2_channels * 2\n            self.c4_channels = self.c3_channels * 2\n        else:\n            self.c2_channels = c2_channels\n            self.c3_channels = c3_channels\n            self.c4_channels = c4_channels\n        \n        self.conv1 = nn.Conv2d(in_channels=3,\n                               out_channels=self.c1_channels,\n                               kernel_size=4,\n                               stride=2,\n                               padding=1,\n                               bias=False)\n        \n        self.conv2 = nn.Conv2d(in_channels=self.c1_channels,\n                               out_channels=self.c2_channels,\n                               kernel_size=4,\n                               stride=2,\n                               padding=1,\n                               bias=False)\n        \n        self.bnorm2 = nn.BatchNorm2d(num_features=self.c2_channels)\n        \n        self.conv3 = nn.Conv2d(in_channels=self.c2_channels,\n                               out_channels=self.c3_channels,\n                               kernel_size=4,\n                               stride=2,\n                               padding=1,\n                               bias=False)\n        \n        self.bnorm3 = nn.BatchNorm2d(num_features=self.c3_channels)\n        \n        self.conv4 = nn.Conv2d(in_channels=self.c3_channels,\n                               out_channels=self.c4_channels,\n                               kernel_size=4,\n                               stride=2,\n                               padding=1,\n                               bias=False)\n        \n        self.bnorm4 = nn.BatchNorm2d(num_features=self.c4_channels)\n        \n        self.conv5 = nn.Conv2d(in_channels=self.c4_channels,\n                               out_channels=1,\n                               kernel_size=4,\n                               padding=0,\n                               stride=1,\n                               bias=False)\n        \n        self.lrelu = nn.LeakyReLU(negative_slope=0.2)\n        \n        self.sigmoid = nn.Sigmoid()\n        \n        \n    def forward(self, img):\n        '''\n        The method for the forward pass in the network\n        \n        Arguments;\n        - img : a torch.tensor that is of the shape N x C x H x W\n                where, N - the batch_size\n                       C - the number of channels\n                       H - the height\n                       W - the width\n       \n       Returns:\n       - out : the output of the Discriminator \n               whether the passed image is real /fake\n        '''\n        \n        #print (img.shape)\n        \n        batch_size = img.shape[0]\n        \n        x = self.lrelu(self.conv1(img))\n        x = self.lrelu(self.bnorm2(self.conv2(x)))\n        x = self.lrelu(self.bnorm3(self.conv3(x)))\n        x = self.lrelu(self.bnorm4(self.conv4(x)))\n        x = self.conv5(x)\n        \n        x = self.sigmoid(x)\n        \n        return x.view(-1, 1).squeeze()\n      \n    def out_shape(self, inp_dim, kernel_size=4, padding=1, stride=2):\n        return ((inp_dim - kernel_size + (2 * padding)) // stride) + 1","execution_count":null,"outputs":[]},{"metadata":{"colab":{},"colab_type":"code","id":"WBtQGV_SoZx4","trusted":true},"cell_type":"code","source":"class Generator(nn.Module):\n    def __init__(self, ct1_channels=512, ct2_channels=256,\n                 ct3_channels=128, ct4_channels=64, d_channels_in_2=False):\n        \n        '''\n        The contructor class for the Generator\n        \n        Arguments:\n        - zin_channels: ###\n        \n        - ct1_channels: The number of output channels for the\n                        first ConvTranspose Layer. [Default - 1024]\n        \n        - ct2_channels: The number of putput channels for the\n                        second ConvTranspose Layer. [Default - 512]\n                        \n        - ct3_channels: The number of putput channels for the\n                        third ConvTranspose Layer. [Default - 256]\n                        \n        - ct4_channels: The number of putput channels for the\n                        fourth ConvTranspose Layer. [Default - 128]\n                        \n        - d_channnels_in_2 : Decrease the number of channels \n                        by 2 times in each layer.\n                        \n        '''\n        super().__init__()\n        \n        # Define the class variables\n        self.ct1_channels = ct1_channels\n        self.pheight = 4\n        self.pwidth = 4\n        \n        if d_channels_in_2:\n            self.ct2_channels = self.ct1_channels // 2\n            self.ct3_channels = self.ct2_channels // 2\n            self.ct4_channels = self.ct3_channels // 2\n        else:\n            self.ct2_channels = ct2_channels\n            self.ct3_channels = ct3_channels\n            self.ct4_channels = ct4_channels\n        \n        self.convt_0 = nn.ConvTranspose2d(in_channels=z_size,\n                                          out_channels=self.ct1_channels,\n                                          kernel_size=4,\n                                          padding=0,\n                                          stride=1,\n                                          bias=False)\n        \n        self.bnorm0 = nn.BatchNorm2d(self.ct1_channels)\n        \n        self.convt_1 = nn.ConvTranspose2d(in_channels=self.ct1_channels,\n                                          out_channels=self.ct2_channels,\n                                          kernel_size=4,\n                                          stride=2,\n                                          padding=1,\n                                          bias=False)\n        \n        self.bnorm1 = nn.BatchNorm2d(num_features=self.ct2_channels)\n        \n        self.convt_2 = nn.ConvTranspose2d(in_channels=self.ct2_channels,\n                                          out_channels=self.ct3_channels,\n                                          kernel_size=4,\n                                          stride=2,\n                                          padding=1,\n                                          bias=False)\n        \n        self.bnorm2 = nn.BatchNorm2d(num_features=self.ct3_channels)\n        \n        self.convt_3 = nn.ConvTranspose2d(in_channels=self.ct3_channels,\n                                          out_channels=self.ct4_channels,\n                                          kernel_size=4,\n                                          stride=2,\n                                          padding=1,\n                                          bias=False)\n        \n        self.bnorm3 = nn.BatchNorm2d(num_features=self.ct4_channels)\n        \n        self.convt_4 = nn.ConvTranspose2d(in_channels=self.ct4_channels,\n                                          out_channels=3,\n                                          kernel_size=4,\n                                          stride=2,\n                                          padding=1,\n                                          bias=False)\n        \n        self.relu = nn.ReLU()\n        self.tanh = nn.Tanh()\n        \n    def forward(self, z):\n        '''\n        The method for the forward pass for the Generator\n        \n        Arguments:\n        - z : the input random uniform vector sampled from uniform distribution\n        \n        Returns:\n        - out : The output of the forward pass through the network\n        '''\n        \n        # Project the input z and reshape\n        x = self.relu(self.bnorm0(self.convt_0(z)))\n        #print (x.shape)\n        x = mila(self.bnorm1(self.convt_1(x)))\n        x = mila(self.bnorm2(self.convt_2(x)))\n        x = mila(self.bnorm3(self.convt_3(x)))\n        out = self.tanh(self.convt_4(x))\n        \n        return out","execution_count":null,"outputs":[]},{"metadata":{"colab":{"base_uri":"https://localhost:8080/","height":238},"colab_type":"code","id":"SO8yTD5SPw8V","outputId":"77971385-5029-4772-93f7-2c21446e2ac8","trusted":true},"cell_type":"code","source":"dis = Discriminator(64).cuda()\ndis.apply(init_weight)\n\ngen = Generator().cuda()\ngen.apply(init_weight)","execution_count":null,"outputs":[]},{"metadata":{"colab":{"base_uri":"https://localhost:8080/","height":459},"colab_type":"code","id":"dgwPSQhaYAsi","outputId":"71356f78-7d6a-4d81-8a2d-89033405ec9c","trusted":true},"cell_type":"code","source":"print (dis)\nprint ()\nprint (gen)","execution_count":null,"outputs":[]},{"metadata":{"colab":{},"colab_type":"code","id":"nMxNayODNO4f","trusted":true},"cell_type":"code","source":"criterion = nn.BCELoss()","execution_count":null,"outputs":[]},{"metadata":{"colab":{},"colab_type":"code","id":"_bs13NRPP3pk","trusted":true},"cell_type":"code","source":"# Optimizers\ncriterion = nn.BCELoss()\nd_lr = 0.0002\ng_lr = 0.0002\n\nd_opt = optim.Adam(dis.parameters(), lr=d_lr, betas=[0.5, 0.999])\ng_opt = optim.Adam(gen.parameters(), lr=g_lr, betas=[0.5, 0.999])","execution_count":null,"outputs":[]},{"metadata":{"colab":{},"colab_type":"code","id":"sWyEye9vNUaz","trusted":true},"cell_type":"code","source":"device=[\"cuda\"]","execution_count":null,"outputs":[]},{"metadata":{"colab":{"base_uri":"https://localhost:8080/","height":9129},"colab_type":"code","id":"-WkGmEcL7Oec","outputId":"e0b0de05-b966-4409-9f4d-eb722b1a0063","scrolled":true,"trusted":true},"cell_type":"code","source":"# Train loop\n\np_every = 300\nt_every = 1\ne_every = 1\ns_every = 1\nepochs = 200\n\nreal_label = 0.9\nfake_label = 0.1\n\ntrain_losses = []\neval_losses = []\nsamples=[]\nsample_size=16\nfixed_z = np.random.uniform(-1, 1, size=(sample_size, z_size))\nfixed_z = torch.from_numpy(fixed_z).float()\nfixed_z = fixed_z.cuda()\n\nfor e in range(epochs):\n    \n    td_loss = 0\n    tg_loss = 0\n    \n    for batch_i, (real_images, _) in enumerate(celeba_train_loader):\n        \n        real_images = real_images.cuda()\n        \n        batch_size = real_images.size(0)\n\n        #### Train the Discriminator ####\n\n        d_opt.zero_grad()\n\t\t\n        d_real = dis(real_images)\n        \n        label = torch.full((batch_size,), real_label, device='cuda')\n        \n        r_loss = criterion(d_real,label)\n        r_loss.backward()\n\n\n        z = torch.randn(batch_size, z_size, 1, 1, device='cuda')\n\n        fake_images = gen(z)\n        \n        label.fill_(fake_label)\n        \n        d_fake = dis(fake_images.detach())\n        \n        f_loss = criterion(d_fake,label)\n        f_loss.backward()\n\n        d_loss = r_loss + f_loss\n\n        d_opt.step()\n\n\n        #### Train the Generator ####\n        g_opt.zero_grad()\n        \n        label.fill_(real_label)\n        \n        d_fake2 = dis(fake_images)\n        \n        g_loss = criterion(d_fake2, label)\n        g_loss.backward()\n\t\t\n        g_opt.step()\n        \n        if batch_i % p_every == 0:\n          noise = torch.randn(1, 100, 1, 1, device='cuda')\n          out = gen(noise)\n          out = out.detach().cpu().squeeze(0).transpose(0, 1).transpose(1, 2).numpy()\n          out = out * (0.5, 0.5, 0.5)\n          out += (0.5, 0.5, 0.5)\n          plt.axis('off')\n          plt.imshow(out)\n          plt.show()\n          print ('Epoch [{:5d} / {:5d}] | d_loss: {:6.4f} | g_loss: {:6.4f}'. \\\n                    format(e+1, epochs, d_loss, g_loss))\n            \ntrain_losses.append([d_loss, g_loss])\n    \n\n  \n        \nprint ('[INFO] Training Completed successfully!')\n    \n    # finally return losses\nreturn train_losses\n\n    # Save training generator samples\n \n                \n","execution_count":null,"outputs":[]},{"metadata":{"colab_type":"text","id":"YaWKOLE1NikK"},"cell_type":"markdown","source":"## Load the pretrained model"},{"metadata":{"colab":{},"colab_type":"code","id":"fybPVWhr3Yn3","trusted":true},"cell_type":"code","source":"#torch.save(gen.state_dict(),'dog_generator.pt')","execution_count":null,"outputs":[]},{"metadata":{"colab":{},"colab_type":"code","id":"TStShI983Ygr","trusted":true},"cell_type":"code","source":"#gen.load_state_dict(torch.load('dog_generator.pt'))","execution_count":null,"outputs":[]},{"metadata":{"colab_type":"text","id":"bknAvhJsLOdM"},"cell_type":"markdown","source":"## Use the Generator to create some random images"},{"metadata":{"colab":{"base_uri":"https://localhost:8080/","height":268},"colab_type":"code","id":"86kb9681atSk","outputId":"320094e0-6695-492e-9320-38089c161cb3","scrolled":true,"trusted":true},"cell_type":"code","source":"noise = torch.randn(1, 100, 1, 1, device='cuda')\nout = gen(noise)\nout = out.detach().cpu().squeeze(0).transpose(0, 1).transpose(1, 2).numpy()\nout = out * (0.5, 0.5, 0.5)\nout += (0.5, 0.5, 0.5)\nplt.axis('off')\nplt.imshow(out)\nplt.show()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"from torchvision.utils import save_image\nn_images=10000\nim_batch_size=50\n\nfor i_batch in range(0, n_images):\n    gen_z = torch.randn(1, 100, 1, 1, device='cuda')\n    gen_images = gen(gen_z)\n    gen_images = gen_images.to(\"cpu\").clone().detach()\n    gen_images.numpy().transpose(0, 2, 3, 1)\n    gen_images = gen_images * (0.5)\n    gen_images += (0.5)\n    for i_image in range(gen_images.size(0)):\n        save_image(gen_images, \n                   os.path.join(out_folder, f'image_{i_batch+i_image:05d}.png'))","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"get_images(out_folder, 0, 25)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"len(os.listdir(out_folder))","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"import shutil\nshutil.make_archive('images', 'zip', out_folder)","execution_count":null,"outputs":[]}],"metadata":{"accelerator":"GPU","colab":{"collapsed_sections":[],"name":"DCGAN.ipynb","provenance":[],"toc_visible":true,"version":"0.3.2"},"kernelspec":{"display_name":"Python 3","language":"python","name":"python3"},"language_info":{"codemirror_mode":{"name":"ipython","version":3},"file_extension":".py","mimetype":"text/x-python","name":"python","nbconvert_exporter":"python","pygments_lexer":"ipython3","version":"3.7.3"}},"nbformat":4,"nbformat_minor":1}