{"cells":[{"metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"cell_type":"code","source":"# This Python 3 environment comes with many helpful analytics libraries installed\n# It is defined by the kaggle/python docker image: https://github.com/kaggle/docker-python\n# For example, here's several helpful packages to load in \n\nimport numpy as np # linear algebra\nimport pandas as pd # data processing, CSV file I/O (e.g. pd.read_csv)\n\n# Input data files are available in the \"../input/\" directory.\n# For example, running this (by clicking run or pressing Shift+Enter) will list the files in the input directory\n\nimport os\nprint(os.listdir(\"../input\"))\n\n# Any results you write to the current directory are saved as output.","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"import matplotlib.pyplot as plt\nimport numpy as np\nimport torch\nfrom torch import nn, optim\nimport torch.nn.functional as F\nfrom torchvision import datasets, transforms\nfrom torchvision.utils import save_image\nimport os\nfrom torch import nn, optim\nimport torch.nn.functional as F\nfrom torchvision import datasets, transforms\nfrom torchvision.utils import save_image\nfrom torch.utils.data import Dataset, DataLoader\nimport os\nfrom tqdm import tqdm_notebook as tqdm\nfrom time import time\nfrom PIL import Image\nimport torch.nn.parallel\nimport torch.optim as optim\nimport torch.utils.data\nimport torchvision\nimport torchvision.datasets as dset\nimport torchvision.transforms as transforms\nimport torchvision.utils as vutils\nfrom torch.autograd import Variable\nfrom torch import nn, optim\nimport torch.nn.functional as F\nfrom torchvision import datasets, transforms\nfrom torchvision.utils import save_image\nimport matplotlib.image as mpimg\nimport torchvision\nimport torchvision.datasets as dset\nfrom torch.utils.data import Dataset, DataLoader\nimport torchvision.transforms as transforms\nimport torchvision.utils as vutils\nfrom torch.autograd import Variable\nimport xml.etree.ElementTree as ET\nimport random\nfrom torch.nn.utils import spectral_norm\nfrom scipy.stats import truncnorm\nimport torch as th\n\n","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"start = time()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def seed_everything(seed=4242):\n    random.seed(seed)\n    os.environ['PYTHONHASHSEED'] = str(seed)\n    np.random.seed(seed)\n    torch.manual_seed(seed)\n    torch.cuda.manual_seed(seed)\n    torch.backends.cudnn.deterministic = True\nseed_everything()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"class DataGenerator(Dataset):\n    def __init__(self, directory, transform=None, n_samples=np.inf):\n        self.directory = directory\n        self.transform = transform\n        self.n_samples = n_samples\n\n        self.samples = self._load_subfolders_images(directory)\n        if len(self.samples) == 0:\n            raise RuntimeError(\"Found 0 files in subfolders of: {}\".format(directory))\n\n    def _load_subfolders_images(self, root):\n        IMG_EXTENSIONS = (\n        '.jpg', '.jpeg', '.png', '.ppm', '.bmp', '.pgm', '.tif', '.tiff', '.webp')\n\n        def is_valid_file(x):\n            return torchvision.datasets.folder.has_file_allowed_extension(x, IMG_EXTENSIONS)\n\n        required_transforms = torchvision.transforms.Compose([\n                torchvision.transforms.Resize(64),\n                torchvision.transforms.CenterCrop(64),\n        ])\n\n        imgs = []\n        paths = []\n        for root, _, fnames in sorted(os.walk(root)):\n            for fname in sorted(fnames)[:min(self.n_samples, 999999999999999)]:\n                path = os.path.join(root, fname)\n                paths.append(path)\n\n        for path in paths:\n            if is_valid_file(path):\n                # Load image\n                img = dset.folder.default_loader(path)\n\n                # Get bounding boxes\n                annotation_basename = os.path.splitext(os.path.basename(path))[0]\n                annotation_dirname = next(\n                        dirname for dirname in os.listdir('../input/annotation/Annotation/') if\n                        dirname.startswith(annotation_basename.split('_')[0]))\n                annotation_filename = os.path.join('../input/annotation/Annotation/',\n                                                   annotation_dirname, annotation_basename)\n                tree = ET.parse(annotation_filename)\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                    \n                    w = np.min((xmax - xmin, ymax - ymin))\n                    bbox = (xmin, ymin, xmin+w, ymin+w)\n                    object_img = required_transforms(img.crop(bbox))\n                    #object_img = object_img.resize((64,64), Image.ANTIALIAS)\n                    imgs.append(object_img)\n        return imgs\n\n    def __getitem__(self, index):\n        sample = self.samples[index]\n        \n        if self.transform is not None:\n            sample = self.transform(sample)\n            \n        return np.asarray(sample)\n    \n    def __len__(self):\n        return len(self.samples)\n","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"%%time\ndatabase = '../input/all-dogs/all-dogs/'\n\ntransform = transforms.Compose([transforms.RandomHorizontalFlip(p=0.1),\n                                transforms.ToTensor(),\n                                transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])\n\ntrain_data = DataGenerator(database, transform=transform,n_samples=25000)\n\n\ntrain_loader = torch.utils.data.DataLoader(train_data, shuffle=True,batch_size=16, num_workers = 4)\n","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# ----------------------------------------------------------------------------\n# Pixelwise feature vector normalization.\n# reference: https://github.com/tkarras/progressive_growing_of_gans/blob/master/networks.py#L120\n# ----------------------------------------------------------------------------\nclass PixelwiseNorm(nn.Module):\n    def __init__(self):\n        super(PixelwiseNorm, self).__init__()\n\n    def forward(self, x, alpha=1e-8):\n        \"\"\"\n        forward pass of the module\n        :param x: input activations volume\n        :param alpha: small number for numerical stability\n        :return: y => pixel normalized activations\n        \"\"\"\n        y = x.pow(2.).mean(dim=1, keepdim=True).add(alpha).sqrt()  # [N1HW]\n        y = x / y  # normalize the input x volume\n        return y","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def show_generated_img_all():\n    gen_z = torch.randn(32, nz, 1, 1, device=device)\n    gen_images = netG(gen_z).to(\"cpu\").clone().detach()\n    gen_images = gen_images.numpy().transpose(0, 2, 3, 1)\n    gen_images = (gen_images+1.0)/2.0\n    fig = plt.figure(figsize=(25, 16))\n    for ii, img in enumerate(gen_images):\n        ax = fig.add_subplot(4, 8, ii + 1, xticks=[], yticks=[])\n        plt.imshow(img)\n    #plt.savefig(filename)  ","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"### This is to show one sample image for iteration of chosing\ndef show_generated_img():\n    noise = torch.randn(1, nz, 1, 1, device=device)\n    gen_image = netG(noise).to(\"cpu\").clone().detach().squeeze(0)\n    gen_image = gen_image.numpy().transpose(1, 2, 0)\n    gen_image = ((gen_image+1.0)/2.0)\n    plt.imshow(gen_image)\n    plt.show()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"class MinibatchStdDev(th.nn.Module):\n    \"\"\"\n    Minibatch standard deviation layer for the discriminator\n    \"\"\"\n\n    def __init__(self):\n        \"\"\"\n        derived class constructor\n        \"\"\"\n        super(MinibatchStdDev, self).__init__()\n\n    def forward(self, x, alpha=1e-8):\n        \"\"\"\n        forward pass of the layer\n        :param x: input activation volume\n        :param alpha: small number for numerical stability\n        :return: y => x appended with standard deviation constant map\n        \"\"\"\n        batch_size, _, height, width = x.shape\n        # [B x C x H x W] Subtract mean over batch.\n        y = x - x.mean(dim=0, keepdim=True)\n        # [1 x C x H x W]  Calc standard deviation over batch\n        y = th.sqrt(y.pow(2.).mean(dim=0, keepdim=False) + alpha)\n\n        # [1]  Take average over feature_maps and pixels.\n        y = y.mean().view(1, 1, 1, 1)\n\n        # [B x 1 x H x W]  Replicate over group and pixels.\n        y = y.repeat(batch_size,1, height, width)\n\n        # [B x C x H x W]  Append as new feature_map.\n        y = th.cat([x, y], 1)\n        # return the computed values:\n        return y\n","execution_count":null,"outputs":[]},{"metadata":{"_cell_guid":"79c7e3d0-c299-4dcb-8224-4455121ee9b0","_uuid":"d629ff2d2480ee46fbb7e2d37f6b5fab8052498a","trusted":true,"_kg_hide-input":true},"cell_type":"code","source":"\nclass Generator(nn.Module):\n    def __init__(self, nz, nfeats, nchannels):\n        super(Generator, self).__init__()\n\n        # input is Z, going into a convolution\n        self.conv1 = spectral_norm(nn.ConvTranspose2d(nz, nfeats * 8, 4, 1, 0, bias=False))\n        #self.bn1 = nn.BatchNorm2d(nfeats * 8)\n        # state size. (nfeats*8) x 4 x 4\n        \n        self.conv2 = spectral_norm(nn.ConvTranspose2d(nfeats * 8, nfeats * 8, 4, 2, 1, bias=False))\n        #self.bn2 = nn.BatchNorm2d(nfeats * 8)\n        # state size. (nfeats*8) x 8 x 8\n        \n        self.conv3 = spectral_norm(nn.ConvTranspose2d(nfeats * 8, nfeats * 4, 4, 2, 1, bias=False))\n        #self.bn3 = nn.BatchNorm2d(nfeats * 4)\n        # state size. (nfeats*4) x 16 x 16\n        \n        self.conv4 = spectral_norm(nn.ConvTranspose2d(nfeats * 4, nfeats * 2, 4, 2, 1, bias=False))\n        #self.bn4 = nn.BatchNorm2d(nfeats * 2)\n        # state size. (nfeats * 2) x 32 x 32\n        \n        self.conv5 = spectral_norm(nn.ConvTranspose2d(nfeats * 2, nfeats, 4, 2, 1, bias=False))\n        #self.bn5 = nn.BatchNorm2d(nfeats)\n        # state size. (nfeats) x 64 x 64\n        \n        self.conv6 = spectral_norm(nn.ConvTranspose2d(nfeats, nchannels, 3, 1, 1, bias=False))\n        # state size. (nchannels) x 64 x 64\n        self.pixnorm = PixelwiseNorm()\n    def forward(self, x):\n        #x = F.leaky_relu(self.bn1(self.conv1(x)))\n        #x = F.leaky_relu(self.bn2(self.conv2(x)))\n        #x = F.leaky_relu(self.bn3(self.conv3(x)))\n        #x = F.leaky_relu(self.bn4(self.conv4(x)))\n        #x = F.leaky_relu(self.bn5(self.conv5(x)))\n        x = F.leaky_relu(self.conv1(x))\n        x = F.leaky_relu(self.conv2(x))\n        x = self.pixnorm(x)\n        x = F.leaky_relu(self.conv3(x))\n        x = self.pixnorm(x)\n        x = F.leaky_relu(self.conv4(x))\n        x = self.pixnorm(x)\n        x = F.leaky_relu(self.conv5(x))\n        x = self.pixnorm(x)\n        x = torch.tanh(self.conv6(x))\n        \n        return x\n\n\n\nclass Discriminator(nn.Module):\n    def __init__(self, nchannels, nfeats):\n        super(Discriminator, self).__init__()\n\n        # input is (nchannels) x 64 x 64\n        self.conv1 = nn.Conv2d(nchannels, nfeats, 4, 2, 1, bias=False)\n        # state size. (nfeats) x 32 x 32\n        \n        self.conv2 = spectral_norm(nn.Conv2d(nfeats, nfeats * 2, 4, 2, 1, bias=False))\n        self.bn2 = nn.BatchNorm2d(nfeats * 2)\n        # state size. (nfeats*2) x 16 x 16\n        \n        self.conv3 = spectral_norm(nn.Conv2d(nfeats * 2, nfeats * 4, 4, 2, 1, bias=False))\n        self.bn3 = nn.BatchNorm2d(nfeats * 4)\n        # state size. (nfeats*4) x 8 x 8\n       \n        self.conv4 = spectral_norm(nn.Conv2d(nfeats * 4, nfeats * 8, 4, 2, 1, bias=False))\n        self.bn4 = nn.MaxPool2d(2)\n        # state size. (nfeats*8) x 4 x 4\n        self.batch_discriminator = MinibatchStdDev()\n        \n        self.conv5 = spectral_norm(nn.Conv2d(nfeats * 8 +1, 1, 2, 1, 0, bias=False))\n        # state size. 1 x 1 x 1\n        \n    def forward(self, x):\n        x = F.leaky_relu(self.conv1(x), 0.2)\n        x = F.leaky_relu(self.bn2(self.conv2(x)), 0.2)\n        x = F.leaky_relu(self.bn3(self.conv3(x)), 0.2)\n        x = F.leaky_relu(self.bn4(self.conv4(x)), 0.2)\n        x = self.batch_discriminator(x)\n        x = torch.sigmoid(self.conv5(x))\n        #x= self.conv5(x)\n        return x.view(-1, 1)\n\n\ndevice = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n\nbatch_size = 16\nlr = 0.0003\nlr_d = 0.0001\nbeta1 = 0.5\nepochs = 800\nnetG = Generator(100, 32, 3).to(device)\nnetD = Discriminator(3, 48).to(device)\n\ncriterion = nn.BCELoss()\n#criterion = nn.MSELoss()\n\noptimizerD = optim.Adam(netD.parameters(), lr=lr, betas=(beta1, 0.999))\noptimizerG = optim.Adam(netG.parameters(), lr=lr_d, betas=(beta1, 0.999))\nlr_schedulerG = torch.optim.lr_scheduler.CosineAnnealingWarmRestarts(optimizerG,\n                                                                     T_0=epochs//200, eta_min=0.00005)\nlr_schedulerD = torch.optim.lr_scheduler.CosineAnnealingWarmRestarts(optimizerD,\n                                                                     T_0=epochs//200, eta_min=0.00005)\n\nnz = 100\nfixed_noise = torch.randn(25, nz, 1, 1, device=device)\n\nreal_label = 0.8\nfake_label = 0.0\nbatch_size = train_loader.batch_size\n\n\n\n### training here\n\n\nstep = 0\nfor epoch in range(epochs):\n    for ii, (real_images) in tqdm(enumerate(train_loader), total=len(train_loader)):\n        end = time()\n        if (end -start) > 30000 :\n            break\n        ############################\n        # (1) Update D network: maximize log(D(x)) + log(1 - D(G(z)))\n        ###########################\n        # train with real\n        netD.zero_grad()\n        real_images = real_images.to(device)\n        batch_size = real_images.size(0)\n        labels = torch.full((batch_size, 1), real_label, device=device) +  np.random.uniform(-0.1, 0.1)\n\n        output = netD(real_images)\n        errD_real = criterion(output, labels)\n        errD_real.backward()\n        D_x = output.mean().item()\n\n        # train with fake\n        noise = torch.randn(batch_size, nz, 1, 1, device=device)\n        fake = netG(noise)\n        labels.fill_(fake_label) + np.random.uniform(0, 0.2)\n        output = netD(fake.detach())\n        errD_fake = criterion(output, labels)\n        errD_fake.backward()\n        D_G_z1 = output.mean().item()\n        errD = errD_real + errD_fake\n        optimizerD.step()\n\n        ############################\n        # (2) Update G network: maximize log(D(G(z)))\n        ###########################\n        netG.zero_grad()\n        labels.fill_(real_label)  # fake labels are real for generator cost\n        output = netD(fake)\n        errG = criterion(output, labels)\n        errG.backward()\n        D_G_z2 = output.mean().item()\n        optimizerG.step()\n        \n        if step % 500 == 0:\n            print('[%d/%d][%d/%d] Loss_D: %.4f Loss_G: %.4f D(x): %.4f D(G(z)): %.4f / %.4f'\n                  % (epoch + 1, epochs, ii, len(train_loader),\n                     errD.item(), errG.item(), D_x, D_G_z1, D_G_z2))\n            \n            valid_image = netG(fixed_noise)\n        step += 1\n        lr_schedulerG.step(epoch)\n        lr_schedulerD.step(epoch)\n\n    if epoch % 5 == 0:\n        show_generated_img()\n        \n# torch.save(netG.state_dict(), 'generator.pth')\n# torch.save(netD.state_dict(), 'discriminator.pth')\n\ndef truncated_normal(size, threshold=1):\n    values = truncnorm.rvs(-threshold, threshold, size=size)\n    return values\n\nif not os.path.exists('../output_images'):\n    os.mkdir('../output_images')\nim_batch_size = 50\nn_images=10000\nfor i_batch in range(0, n_images, im_batch_size):\n    z = truncated_normal((im_batch_size, 100, 1, 1), threshold=1)\n    gen_z = torch.from_numpy(z).float().to(device)    \n    #gen_z = torch.randn(im_batch_size, 100, 1, 1, device=device)\n    gen_images = netG(gen_z)\n    images = gen_images.to(\"cpu\").clone().detach()\n    images = images.numpy().transpose(0, 2, 3, 1)\n    for i_image in range(gen_images.size(0)):\n        save_image((gen_images[i_image, :, :, :] +1.0)/2.0, os.path.join('../output_images', f'image_{i_batch+i_image:05d}.png'))\n\n\nimport shutil\nshutil.make_archive('images', 'zip', '../output_images')","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"show_generated_img_all()","execution_count":null,"outputs":[]}],"metadata":{"kernelspec":{"display_name":"Python 3","language":"python","name":"python3"},"language_info":{"name":"python","version":"3.6.4","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"}},"nbformat":4,"nbformat_minor":1}