{"cells":[{"metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"cell_type":"code","source":"%load_ext tensorboard.notebook\n%matplotlib inline\n\nimport os\nimport shutil\nfrom dataclasses import dataclass\nfrom functools import partial\nimport pathlib\nimport time\nimport xml.etree.ElementTree as ET\nimport zipfile\n\nimport numpy as np\nimport pandas as pd\nimport skimage\nfrom PIL import Image\nimport cv2\n\nimport matplotlib.pyplot as plt\n\nfrom torch.utils.data import DataLoader, Dataset\nimport albumentations as albu\nfrom tqdm import tqdm_notebook as tqdm","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","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"@dataclass\nclass Example:\n    img: np.ndarray\n    category: str\n    difficult: int\n    transform: object = None\n\n    def read_img(self):\n        # read and crop\n        img = self.img\n\n        if self.transform is not None:\n            img = self.transform(image=img)['image']\n\n        # convert to ndarray and transpose HWC => CHW\n        img = np.array(img).transpose(2, 0, 1)\n        return img\n\n    def show(self):\n        img = (self.read_img() + 1) / 2\n        plt.imshow(img.transpose(1, 2, 0))\n        plt.title(self.category)\n\nclass DogDataset(Dataset):\n    def __init__(self,\n                 img_dir='../input/all-dogs/all-dogs/',\n                 anno_dir='../input/annotation/Annotation/',\n                 transform=None,\n                 examples=None):\n        self.img_dir = pathlib.Path(img_dir)\n        self.anno_dir = pathlib.Path(anno_dir)\n        self.transform = transform\n        self.preprocess = albu.Compose([\n            albu.SmallestMaxSize(64),\n        ])\n        \n        if examples is None:\n            self.examples = self._correct_examples()\n        else:\n            self.examples = examples\n\n        self.categories = sorted(set([e.category for e in self.examples]))\n        self.categ2id = dict(zip(self.categories, range(len(self.categories))))\n        \n    def _correct_examples(self):\n        examples = []\n        for anno in tqdm(list(self.anno_dir.glob('*/*'))):\n            tree = ET.parse(anno)\n            root = tree.getroot()\n\n            img_path = self.img_dir / f'{root.find(\"filename\").text}.jpg'\n            if not img_path.exists():\n                continue\n\n            objects = root.findall('object')\n            for obj in objects:\n                examples.append(self._create_example(img_path, obj))\n        return examples\n\n    def _create_example(self, img_path, obj):\n        # reading bound box\n        bbox = obj.find('bndbox')\n        # read and preprocess image\n        img = skimage.io.imread(img_path)\n        xmin=int(bbox.find('xmin').text)\n        ymin=int(bbox.find('ymin').text)\n        xmax=int(bbox.find('xmax').text)\n        ymax=int(bbox.find('ymax').text)\n        img = img[ymin:ymax, xmin:xmax]\n        img = self.preprocess(image=img)['image']\n\n        # add example\n        return Example(\n            img=img,\n            category=obj.find('name').text,\n            difficult=int(obj.find('difficult').text),\n            transform=self.transform,\n        )\n    \n    def __len__(self):\n        return len(self.examples)\n    \n    def __getitem__(self, i):\n        e = self.examples[i]\n        img = e.read_img()\n\n        categ_id = self.categ2id[e.category]\n        return img, categ_id\n\n    def show_examples(self, indices=None, n_cols=8):\n        if indices is None:\n            indices = np.random.randint(0, len(self), n_cols)\n\n        n_rows = (len(indices)-1) // n_cols + 1\n\n        fig = plt.figure(figsize=(n_cols*4, n_rows*4))\n        for i, idx in enumerate(indices, 1):\n            fig.add_subplot(n_rows, n_cols, i)\n            self.examples[idx].show()\n        plt.show()","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=42):\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":"import albumentations.pytorch","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"%%time\n\ntransform = albu.Compose([\n    albu.CenterCrop(64, 64),\n    albu.HorizontalFlip(p=0.5),\n    albu.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),\n])\n\ndataset = DogDataset(\n    transform=transform,\n)\n\nprint(dataset)\ndataset.show_examples()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"ls /opt/conda/lib/python3.6/site-packages/cv2/data/haarcascade_*.xml","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def crop_with_facebox(img, pos):\n    pos = pos[0]\n    left, right = pos[0], pos[0]+pos[2]\n    top, bottom = pos[1], pos[1]+pos[3]\n    img = img[top:bottom, left:right]\n    return img\n\ndef create_cropped_dataset(dataset):\n    rescale = albu.SmallestMaxSize(64)\n\n    cascade_file_path = '/opt/conda/lib/python3.6/site-packages/cv2/data/haarcascade_{}.xml'\n    \n    # classifiers are prioritized\n    ordered_classifiers = [\n        (cv2.CascadeClassifier(cascade_file_path.format('frontalcatface_extended')), crop_with_facebox),\n        (cv2.CascadeClassifier(cascade_file_path.format('frontalcatface')), crop_with_facebox),\n        (cv2.CascadeClassifier(cascade_file_path.format('frontalface_default')), crop_with_facebox),\n        (cv2.CascadeClassifier(cascade_file_path.format('frontalface_alt')), crop_with_facebox),\n        (cv2.CascadeClassifier(cascade_file_path.format('frontalface_alt2')), crop_with_facebox),\n    ]\n\n    cropped_examples = []\n    for i, e in enumerate(tqdm(dataset.examples)):\n        grayimg = cv2.cvtColor(e.img, cv2.COLOR_RGB2GRAY)\n        for clf, crop_fn in ordered_classifiers:\n            pos = clf.detectMultiScale(grayimg)\n            if len(pos) != 0:\n                break\n\n        if len(pos) == 0:\n            continue\n\n        img = crop_fn(e.img, pos)\n        if img is None:\n            continue\n\n        img = rescale(image=img)['image']\n        cropped_examples.append(Example(\n            img=img, category=e.category,\n            difficult=e.difficult,\n            transform=e.transform,\n        ))\n\n    return DogDataset(examples=cropped_examples)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"%%time\ncropped_dataset = create_cropped_dataset(dataset)\nprint(len(cropped_dataset))\ncropped_dataset.show_examples()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"batch_size = 32\ntrain_loader = torch.utils.data.DataLoader(cropped_dataset, shuffle=True,batch_size=batch_size, num_workers = 4)","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","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"class 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        self.pixnorm = PixelwiseNorm()\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 = self.pixnorm(x)\n        x = F.leaky_relu(self.bn3(self.conv3(x)), 0.2)\n       # x = self.pixnorm(x)\n        x = F.leaky_relu(self.bn4(self.conv4(x)), 0.2)\n       # x = self.pixnorm(x)\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\nlr = 0.0003\nlr_d = 0.0001\nbeta1 = 0.5\n#epochs = 900\nepochs = 14000\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.7\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) > 31800:\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 % 200 == 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 = 100\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":{"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}