{"cells":[{"metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"cell_type":"code","source":"EPOCHS = 574\nBATCH = 32\nSIZE = 64\n\nLRG = 0.0004\nLRD = 0.0008\n\nNZ = 64\nFEATURES_G = 80\nFEATURES_D = 112\n\nLRFACTOR_G = 0.9\nLRFACTOR_D = 0.9\n\nREAL_LABEL = [0.9, 0.9]\nFAKE_LABEL = [0.00, 0.00]\nREAL_FAKE_LABEL = 0.9\n\nSEED = 1234\nRATIO = 1.50\n\nPATH_ANNOTATIONS = '../input/annotation/Annotation'\nPATH_ALL_DOGS = '../input/all-dogs/all-dogs'\nPATH_DOGS = '../cropped-dogs'\n","execution_count":null,"outputs":[]},{"metadata":{"_uuid":"d629ff2d2480ee46fbb7e2d37f6b5fab8052498a","_cell_guid":"79c7e3d0-c299-4dcb-8224-4455121ee9b0","trusted":true},"cell_type":"code","source":"import os\nimport sys\nimport time\nimport math\nimport random\nimport shutil\nimport numpy as np\nimport tensorflow as tf\nimport matplotlib.pyplot as plt\nfrom tqdm import tqdm as tqdm\n\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nfrom torch.utils.data import Dataset\nfrom torch import nn, optim\nimport torch.optim as op\n\nfrom torchvision import datasets, transforms\nfrom torchvision.utils import make_grid, save_image\nfrom PIL import Image\n\nDEVICE = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n\ndef random_seed(seed):\n    torch.backends.cudnn.deterministic = True\n    torch.backends.cudnn.benchmark = False\n    random.seed(seed)\n    np.random.seed(seed)\n    torch.manual_seed(seed)\n    tf.set_random_seed(seed)\n\ndef R4(f): return F'{f:0.4}'.ljust(6, '0').ljust(8, ' ')\ndef R6(f): return F'{f:0.6}'.ljust(8, '0').ljust(8, ' ')\n    \nrandom_seed(SEED) \n","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"%%time\n\nimport xml.etree.ElementTree as ET\n\n'''\n    Cropped Images\n'''\ndef enumerate_images(path):\n    for fn in os.listdir(path):\n        with Image.open(os.path.join(path, fn)) as image:\n            yield image.copy()\n\ndef enumerate_annotation_images(path_annotation, path_dogs):\n    breed_id = 0\n    breeds = os.listdir(path_annotation)\n    for breed in breeds:\n        dog_id = 0\n        breed_id += 1\n        for dog in os.listdir(os.path.join(path_annotation, F'{breed}')):\n            try:\n                image = Image.open(os.path.join(path_dogs, F'{dog}.jpg'))\n            except :\n                print(dog, 'not found!')\n                continue\n\n            tree = ET.parse(os.path.join(path_annotation, F'{breed}', F'{dog}'))\n            root = tree.getroot()\n            objects = root.findall('object')\n            for o in objects:\n                dog_id += 1\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                img = image.crop((xmin, ymin, xmax, ymax))\n                yield img, breed_id, dog_id\n\ndef save_cropped_images(path_annotation, path_dogs, path_target, size=64):\n    if not os.path.exists(path_target):\n        os.mkdir(path_target)\n    n = 0\n    current_id = 0\n    with tqdm(enumerate_annotation_images(path_annotation, path_dogs)) as t:\n        for image, breed_id, dog_id in t:\n            if current_id != breed_id:\n                current_id = breed_id\n                t.set_postfix({'id': breed_id})\n            w, h = image.size\n            if w < size or h < size: continue\n            ratio = w / h\n            if h > w: ratio = h / w\n            sz = math.ceil(size * ratio)\n            if ratio < RATIO:\n                image.thumbnail((sz, sz), Image.ANTIALIAS)\n                image.save(os.path.join(path_target, F'{breed_id:03}_{dog_id:03}.png'))\n                n += 1\n    return n\n\nsave_cropped_images(PATH_ANNOTATIONS, PATH_ALL_DOGS, PATH_DOGS)\n","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"import torch\nimport torch.nn as nn\nimport torch.nn.functional as F\n\n'''\n    GENERATOR\n'''\nclass Generator(nn.Module):\n    def __init__(self, nz, nfeats):\n        super(Generator, self).__init__()\n\n        self.conv1 = nn.ConvTranspose2d(nz, nfeats * 8, 4, 1, 0, bias=False)\n        self.bn1 = nn.BatchNorm2d(nfeats * 8)\n        self.conv2 = nn.ConvTranspose2d(nfeats * 8, nfeats * 8, 4, 2, 1, bias=False)\n        self.bn2 = nn.BatchNorm2d(nfeats * 8)\n        self.conv3 = nn.ConvTranspose2d(nfeats * 8, nfeats * 4, 4, 2, 1, bias=False)\n        self.bn3 = nn.BatchNorm2d(nfeats * 4)\n        self.conv4 = nn.ConvTranspose2d(nfeats * 4, nfeats * 2, 4, 2, 1, bias=False)\n        self.bn4 = nn.BatchNorm2d(nfeats * 2)\n        self.conv5 = nn.ConvTranspose2d(nfeats * 2, nfeats, 4, 2, 1, bias=False)\n        self.bn5 = nn.BatchNorm2d(nfeats)\n        self.conv6 = nn.ConvTranspose2d(nfeats, 3, 3, 1, 1, bias=False)\n\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 = torch.tanh(self.conv6(x))\n        return x\n\n    def load(self, path, device=None):\n        self.load_state_dict(torch.load(path, map_location=device))\n\n    def save(self, path):\n        torch.save(self.state_dict(), path)\n\n'''\n    DISCRIMINATOR\n'''\nclass Discriminator(nn.Module):\n    def __init__(self, nfeats):\n        super(Discriminator, self).__init__()\n\n        self.conv1 = nn.Conv2d(3, nfeats, 4, 2, 1, bias=False)\n        self.conv2 = nn.Conv2d(nfeats, nfeats * 2, 4, 2, 1, bias=False)\n        self.bn2 = nn.BatchNorm2d(nfeats * 2)\n        self.conv3 = nn.Conv2d(nfeats * 2, nfeats * 4, 4, 2, 1, bias=False)\n        self.bn3 = nn.BatchNorm2d(nfeats * 4)\n        self.conv4 = nn.Conv2d(nfeats * 4, nfeats * 8, 4, 2, 1, bias=False)\n        self.bn4 = nn.BatchNorm2d(nfeats * 8)\n        self.conv5 = nn.Conv2d(nfeats * 8, 1, 4, 1, 0, bias=False)\n        \n    def forward(self, x):\n        x = F.leaky_relu(self.conv1(x), 0.1)\n        x = F.leaky_relu(self.bn2(self.conv2(x)), 0.1)\n        x = F.leaky_relu(self.bn3(self.conv3(x)), 0.1)\n        x = F.leaky_relu(self.bn4(self.conv4(x)), 0.1)\n        x = torch.sigmoid(self.conv5(x))\n        x = x.view(-1, 1)\n        return x\n\n    def load(self, path, device=None):\n        self.load_state_dict(torch.load(path, map_location=device))\n\n    def save(self, path):\n        torch.save(self.state_dict(), path)\n","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"import os\nimport random\nimport numpy as np\nimport matplotlib.pyplot as plt\n\nfrom PIL import Image\nfrom torch.utils.data import Dataset\n\nimport cv2\n\n'''\n    DATASET\n'''\ndef image_resize(image, width=None, height=None, inter=cv2.INTER_AREA):\n    dim = None\n    (h, w) = image.shape[:2]\n    if width is None and height is None:\n        return image\n    if width is None:\n        r = height / float(h)\n        width = max(height, int(w * r))\n    else:\n        r = width / float(w)\n        height = max(width, int(h * r))\n    dim = (width, height)\n    resized = cv2.resize(image, dim, interpolation=inter)\n    return resized\n\ndef random_crop(arr_image, width, height):\n    x = random.randint(0, arr_image.shape[1] - width)\n    y = random.randint(0, arr_image.shape[0] - height)\n    arr_image = arr_image[y:y + height, x:x + width]\n    return arr_image\n\ndef random_fliplr(arr_image):\n    r = random.randint(0, 2)\n    if r == 1:\n        return np.fliplr(arr_image)\n    return arr_image\n\nclass ArrayDataset(Dataset):\n    def __init__(self, path, size=64, margin=0, mean=None, std=None, count=None):\n        self.size = size\n        self.margin = margin\n        files = os.listdir(path)\n        if count is not None:\n            files = files[:count]\n        self.arrange(path, files)\n        if mean is None or std is None:\n            m, s = self.calc_norm()\n            if mean is None: mean = m\n            if std is None: std = s\n        self.mean = np.array(mean).reshape(-1, 1, 1)\n        self.std = np.array(std).reshape(-1, 1, 1)\n\n    def arrange(self, path, files):\n        self.images = []\n        self.labels = []\n        sizem = self.size + self.margin\n        for fn in files:\n            image = cv2.imread(os.path.join(path, fn))\n            if image.shape[0] > image.shape[1]:\n                image = image_resize(image, width=sizem)\n            else:\n                image = image_resize(image, height=sizem)\n            image = image[..., ::-1].copy()\n            self.images.append(image)\n            lbl = fn[:3]\n            if str.isdigit(lbl):\n                self.labels.append(int(lbl))\n            else:\n                self.labels.append(0)\n\n    def calc_norm(self):\n        mean = np.zeros(3)\n        std = np.zeros(3)\n        for x in self.images:\n            x = np.transpose(x, (2, 0, 1))\n            x = x.reshape(3, -1)\n            mean += x.mean(1)\n            std += x.std(1)\n        m = (mean / len(self.images)) / 255.0\n        s = (std / len(self.images)) / 255.0\n        return m, s\n\n    def __len__(self):\n        return len(self.images)\n\n    def __getitem__(self, idx):\n        x = self.images[idx]\n        x = random_crop(x, self.size, self.size)\n        x = random_fliplr(x).copy()\n        x = np.transpose(x, (2, 0, 1))\n        x = x / 255.0\n        x = x - self.mean\n        x = x / self.std\n        x = x.astype(np.float32)\n        return x, self.labels[idx]\n","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"import numpy as np\nimport torch\n\nfrom tqdm import tqdm as tqdm\n\nDEVICE = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n\n'''\n    TRAINER\n'''\nclass Trainer(object):\n    def __init__(self, netG, netD, optimizerG, optimizerD, criterion, mean, std):\n        self.netG = netG\n        self.netD = netD\n        self.optimizerG = optimizerG\n        self.optimizerD = optimizerD\n        self.criterion = criterion\n        self.mean = mean\n        self.std = std\n\n    def next_step(self, data_loader, noise_gen, label_gen):\n        for step, (real_images, real_labels) in enumerate(data_loader):\n            batch_size = real_images.size(0)\n            labels_true = label_gen(batch_size, 'real').to(DEVICE)\n            labels_false = label_gen(batch_size, 'fake').to(DEVICE)\n            labels_true_false = label_gen(batch_size, 'real-fake').to(DEVICE)\n\n            # train with real\n            self.netD.zero_grad()\n            real_images = real_images.to(DEVICE)\n            output = self.netD(real_images)\n            errD_real = self.criterion(output, labels_true)\n            errD_real.backward()\n            D_x = output.mean().item()\n\n            # train with fake\n            noise = noise_gen(batch_size).to(DEVICE)\n            fake = self.netG(noise)\n            fake = fake / 2 + 0.5\n            fake = fake - self.mean\n            fake = fake / self.std\n\n            output = self.netD(fake.detach())\n            errD_fake = self.criterion(output, labels_false)\n            errD_fake.backward()\n            D_G_z1 = output.mean().item()\n            errD = errD_real + errD_fake\n            self.optimizerD.step()\n        \n            # Update G network\n            self.netG.zero_grad()\n            output = self.netD(fake)\n            errG = self.criterion(output, labels_true_false)\n            errG.backward()\n            D_G_z2 = output.mean().item()\n            self.optimizerG.step()\n\n        return errD.item(), errG.item(), D_x, D_G_z1, D_G_z2\n\n    def stabilize(self, count, batch_size, noise_gen):\n        self.netG.train()\n        self.netD.train()\n        with torch.no_grad():\n            for n in range(count // batch_size):\n                gen_z = noise_gen(batch_size).to(DEVICE)\n                _ = self.netG(gen_z)\n\n    def generate(self, count, batch_size, noise_gen):\n        self.netG.eval()\n        images_set = []\n        with torch.no_grad():\n            for n in range(count // batch_size):\n                gen_z = noise_gen(batch_size).to(DEVICE)\n                gen_images = self.netG(gen_z)\n                images = gen_images.detach().cpu()\n                images_set.append(images)\n        images = torch.cat(images_set)\n        self.netG.train()\n        return images\n\n    def generate_numpy(self, count, batch_size, noise_gen):\n        self.netG.eval()\n        images_set = []\n        with torch.no_grad():\n            for n in range(count // batch_size):\n                gen_z = noise_gen(batch_size).to(DEVICE)\n                gen_images = self.netG(gen_z)\n                images = gen_images.detach().cpu()\n                images = images.numpy().transpose(0, 2, 3, 1)\n                images = (images + 1.0) * 0.5 * 255\n                images = images.astype(int)\n                images_set.append(images)\n        images = np.concatenate(images_set)\n        self.netG.train()\n        return images\n","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"%%time\n'''\n    PREPARE\n'''\ntrain_data = ArrayDataset(PATH_DOGS, size=SIZE, margin=0)\ntrain_loader = torch.utils.data.DataLoader(train_data, shuffle=True, batch_size=BATCH, num_workers=4, drop_last=True)\n\nnetG = Generator(NZ, FEATURES_G).to(DEVICE)\nnetD = Discriminator(FEATURES_D).to(DEVICE)\n\ncriterion = nn.BCELoss()\n\noptimizerG = optim.Adam(netG.parameters(), lr=LRG, betas=(0.5, 0.999))\noptimizerD = optim.Adam(netD.parameters(), lr=LRD, betas=(0.5, 0.999))\n\nschedulerG = op.lr_scheduler.LambdaLR(optimizerG, lambda epoch: LRFACTOR_G ** epoch)\nschedulerD = op.lr_scheduler.LambdaLR(optimizerD, lambda epoch: LRFACTOR_D ** epoch)\n\ndef noise_generator(batch_size):\n    return torch.randn(batch_size, NZ, 1, 1)\n\ndef label_generator(batch_size, mode):\n    if mode == 'real':\n        label = np.random.uniform(*REAL_LABEL, batch_size).reshape(batch_size, 1)\n        return torch.tensor(label, dtype=torch.float32)\n    if mode == 'fake':\n        label = np.random.uniform(*FAKE_LABEL, batch_size).reshape(batch_size, 1)\n        return torch.tensor(label, dtype=torch.float32)\n    if mode == 'real-fake':\n        return torch.full((batch_size, 1), REAL_FAKE_LABEL)\n    raise Exception('Invalid parameter mode')\n    ","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def needs_scheduler(epoch):\n    if epoch > 0:\n        if epoch % 50 == 0:\n            print('Scheduler update')\n            return True\n    return False\n","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"%%time\n'''\n    TRAIN\n'''\nhist_metrics = []\nlast_metrics = 0\n\nmean = train_data.mean\nmean = np.expand_dims(mean, axis=0)\nmean = np.repeat(mean, BATCH, 0)\nmean = torch.tensor(mean).float().to(DEVICE)\n\nstd = train_data.std\nstd = np.expand_dims(std, axis=0)\nstd = np.repeat(std, BATCH, 0)\nstd = torch.tensor(std).float().to(DEVICE)\n\ntrainer = Trainer(netG, netD, optimizerG, optimizerD, criterion, mean, std)\n\nmetrics = 0.0\nt0 = time.time()\nprint('Training...')\nfor epoch in range(EPOCHS):        \n    random_seed(SEED + epoch)\n    errD, errG, dx, dgz1, dgz2 = trainer.next_step(train_loader, noise_generator, label_generator)\n    \n    if needs_scheduler(epoch):\n        schedulerG.step()\n        schedulerD.step()\n","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"torch.save(netG.state_dict(), 'generator.pth')\ntorch.save(netD.state_dict(), 'discriminator.pth')\n","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"random_seed(1234)\nimages = trainer.generate(32, 32, noise_generator)\nimage = (images + 1.0) * 0.5\n\nsave_image(images, 'samples.jpg', normalize=True)\nimg = Image.open('samples.jpg')\nimg = img.resize((img.size[0] * 2, img.size[1] * 2), Image.ANTIALIAS)\nimg.save('samples.jpg')\n\nplt.figure(figsize=(16, 16))\nplt.imshow(np.transpose(make_grid(images, padding=2, normalize=True), (1, 2, 0)))\nplt.show()\n","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"SEEDS = [\n1525, 491, 1117, 1509, 1538, 1478, 16, 1251, 1126, 81, 254, 704, 1740, 2081, 79, 980, 1770, 795, 1317, 1471, 2198, 153, 1488, 1456, 454, 573, 690, 2289, 1753, 2075, 772, 301, 1240, 83, 686, 932, 826, 27, 24, 1744, 880, 2029, 1421, 1490, 667, 1387, 1948, 2130, 15, 133, 2151, 602, 76, 2090, 1277, 717, 2175, 2168, 1313, 1562, 1526, 1092, 800, 1804, 1605, 1120, 2320, 2372, 2080, 460, 126, 2302, 1940, 983, 668, 773, 1788, 565, 1540, 2337, 140, 2173, 2010, 1784, 1248, 1106, 1722, 208, 1757, 2371, 449, 481, 1914, 2217, 437, 1161, 169, 693, 807, 1839, 217, 1520, 1102, 30, 1199, 239, 603, 730, 569, 2309, 1434, 1679, 3, 313, 532, 35, 825, 824, 388, 255, 1758, 763, 87, 751, 2073, 585, 407, 1510, 365, 739, 417, 463, 1412, 1636, 1237, 1439, 278, 227, 1916, 2261, 154, 1622, 1833, 2213, 1738, 74, 892, 1203, 2001, 1227, 815, 872, 1769, 699, 2374, 1050, 1073, 620, 202, 1420, 2211, 840, 568, 1715, 1337, 1651, 1591, 878, 1445, 1100, 1702, 551, 72, 1917, 1256, 1479, 152, 2192, 559, 2031, 1920, 23, 2304, 2153, 1848, 2032, 1976, 578, 650, 1464, 1706, 820, 2232, 762, 33, 1627, 1141, 2262, 781, 174,\n]\n","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"state_dict = torch.load('generator.pth')\n\ndef load_model(model):\n    model.load_state_dict(state_dict)\n\ndef generate_numpy_images(netG, count, batch_size, noise_gen):\n    netG.eval()\n    images_set = []\n    with torch.no_grad():\n        for n in range(count // batch_size):\n            gen_z = noise_gen(batch_size).to(DEVICE)\n            gen_images = netG(gen_z)\n            aimages = gen_images.detach().cpu()\n            aimages = (aimages + 1.0) * 0.5 * 255\n            images = aimages.numpy().transpose(0, 2, 3, 1)\n            images = images.astype(int)\n            images_set.append(images)\n    images = np.concatenate(images_set)\n    return images\n\ndef generate_tensors(netG, count, batch_size, noise_gen):\n    netG.eval()\n    images_set = []\n    with torch.no_grad():\n        for n in range(count // batch_size):\n            gen_z = noise_gen(batch_size).to(DEVICE)\n            gen_images = netG(gen_z)\n            images = gen_images.detach().cpu()\n            images_set.append(images)\n    images = torch.cat(images_set)\n    netG.train()\n    return images\n\ndef generate_images(count, batch, seed):\n    netG = Generator(NZ, FEATURES_G).to(DEVICE)\n    load_model(netG)\n    random_seed(seed)\n    images = generate_tensors(netG, count, batch, noise_generator)\n    return images\n","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"%%time\n\nPATH_GENERATED = '../output_images'\n\nif not os.path.exists(PATH_GENERATED):\n    os.mkdir(PATH_GENERATED)\n\nn = 0\nfor seed in SEEDS:\n    images = generate_images(50, 50, seed)\n    images = (images + 1.0) * 0.5\n    for ix in range(50):\n        image = images[ix]\n        save_image(image, os.path.join(PATH_GENERATED, f'image_{n:05d}.png'), normalize=False)\n        n += 1\n\nprint(n)\n\nimport shutil\nshutil.make_archive('images', 'zip', PATH_GENERATED)    \n","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}