{"cells":[{"metadata":{},"cell_type":"markdown","source":"reference:\n- [GAN dogs starter](https://www.kaggle.com/wendykan/gan-dogs-starter)\n- [Demo MiFID metric for Dog image generation comp](https://www.kaggle.com/wendykan/demo-mifid-metric-for-dog-image-generation-comp)"},{"metadata":{"trusted":true},"cell_type":"code","source":"import os\nfrom pathlib import Path\nimport shutil\n\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","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"batch_size = 32\nepochs = 2\n\nMODEL_PATH = '../input/dog-face-generation-competition-kid-metric-input/classify_image_graph_def.pb'\nTRAIN_DIR = Path('../input/generative-dog-images/all-dogs/')\nOUT_DIR = Path('../output_images/')\nOUT_DIR.mkdir(exist_ok=True)\n\ndevice = torch.device('cuda')\n\nlr = 0.0003\nbeta1 = 0.5\n\nnz = 100\nfixed_noise = torch.randn(25, nz, 1, 1, device=device)\n\nreal_label = 0.9\nfake_label = 0","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"class Generator(nn.Module):\n\n    def __init__(self, nz, nfeats, nchannels):\n        super(Generator, self).__init__()\n\n        # input is Z, going into a convolution\n        self.conv1 = 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 = 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 = 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 = 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 = 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 = nn.ConvTranspose2d(nfeats, nchannels, 3, 1, 1, bias=False)\n        # state size. (nchannels) x 64 x 64\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        \n        return x\n\n\nclass Discriminator(nn.Module):\n\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 = 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 = 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 = nn.Conv2d(nfeats * 4, nfeats * 8, 4, 2, 1, bias=False)\n        self.bn4 = nn.BatchNorm2d(nfeats * 8)\n        # state size. (nfeats * 8) x 4 x 4\n        \n        self.conv5 = nn.Conv2d(nfeats * 8, 1, 4, 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 = torch.sigmoid(self.conv5(x))\n        \n        return x.view(-1, 1)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"transform = transforms.Compose([transforms.Resize(64),\n                                transforms.CenterCrop(64),\n                                transforms.ToTensor(),\n                                transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])\n\ntrain_data = datasets.ImageFolder(TRAIN_DIR, transform=transform)\ntrain_loader = torch.utils.data.DataLoader(train_data, shuffle=True, batch_size=batch_size)\n                                           \nimgs, label = next(iter(train_loader))\nimgs = imgs.numpy().transpose(0, 2, 3, 1)\n\nnet_g = Generator(100, 32, 3).to(device)\nnet_d = Discriminator(3, 48).to(device)\n\ncriterion = nn.BCELoss()\n\noptimizer_g = optim.Adam(net_g.parameters(), lr=lr, betas=(beta1, 0.999))\noptimizer_d = optim.Adam(net_d.parameters(), lr=lr, betas=(beta1, 0.999))","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## 1. Train"},{"metadata":{"trusted":true},"cell_type":"code","source":"for epoch in range(epochs):\n    for i, (real_images, train_labels) in enumerate(train_loader):\n        # --------------------------------------\n        # Update Discriminator network: maximize log(D(x)) + log(1 - D(G(z)))\n        # --------------------------------------\n        # train with real\n        net_d.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)\n\n        output = net_d(real_images)\n        err_d_real = criterion(output, labels)\n        err_d_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 = net_g(noise)\n        labels.fill_(fake_label)\n        output = net_d(fake.detach())\n        err_d_fake = criterion(output, labels)\n        err_d_fake.backward()\n        d_g_z1 = output.mean().item()\n        err_d = err_d_real + err_d_fake\n        optimizer_d.step()\n\n        # --------------------------------------\n        # Update Generator network: maximize log(D(G(z)))\n        # --------------------------------------\n        net_g.zero_grad()\n        labels.fill_(real_label)  # fake labels are real for generator cost\n        output = net_d(fake)\n        err_g = criterion(output, labels)\n        err_g.backward()\n        d_g_z2 = output.mean().item()\n        optimizer_g.step()\n        \n        if (i + 1) % (len(train_loader) // 2) == 0:\n            print('[{}/{}][{}/{}] Loss_d: {:.4f} Loss_g: {:.4f} D(x): {:.4f} D(G(z)): {:.4f} / {:.4f}'.format(\n                epoch + 1, epochs, i + 1, len(train_loader), err_d.item(), err_g.item(), d_x, d_g_z1, d_g_z2))","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## 2. Generate images"},{"metadata":{"trusted":true},"cell_type":"code","source":"im_batch_size = 50\nn_images = 10000\n\nfor i_batch in range(0, n_images, im_batch_size):\n    gen_z = torch.randn(im_batch_size, 100, 1, 1, device=device)\n    gen_images = net_g(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, :, :, :], OUT_DIR / f'image_{i_batch + i_image:05d}.png')","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## 3. Evaluate"},{"metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"cell_type":"code","source":"from dataclasses import dataclass\nfrom pathlib import Path\nimport warnings\n\nimport numpy as np\nfrom PIL import Image\nfrom scipy import linalg\nimport tensorflow as tf\n\n\nclass KernelEvalException(Exception):\n    pass\n\n\n@dataclass\nclass MiFIDEvaluator(object):\n    model_path: str\n    train_images_path: str\n    feature_path: str = None\n    imsize: int = 64\n    output_layer: str = 'Pretrained_Net/pool_3:0'\n    input_layer: str = 'Pretrained_Net/ExpandDims:0'\n    output_shape: int = 2048\n    cosine_distance_eps: float = 0.1\n    batch_size: int = 50\n    fid_epsilon: float = 1e-14\n    \n    def __post_init__(self):\n        tf.reset_default_graph()\n        self.create_model_graph()\n        with tf.Session() as sess:\n            if self.feature_path is None:\n                self.mu2, self.sigma2, self.features2 = self._handle_path_memorization(\n                    self.train_images_path, sess, is_checksize=False, is_check_png=False)\n            else:\n                with np.load(self.feature_path) as f:\n                    self.mu2, self.sigma2, self.features2 = f['m'], f['s'], f['features']\n    \n    def create_model_graph(self):\n        with tf.gfile.FastGFile(self.model_path, 'rb') as f:\n            graph_def = tf.GraphDef()\n            graph_def.ParseFromString(f.read())\n            _ = tf.import_graph_def(graph_def, name='Pretrained_Net')\n            \n    def img_read_checks(self, filename, is_checksize=False, is_check_png=False):\n        im = Image.open(str(filename))\n        if is_checksize and im.size != (self.imsize, self.imsize):\n            raise KernelEvalException(f'The images are not of size {check_imsize}')\n        if is_check_png and im.format != 'PNG':\n            raise KernelEvalException('Only PNG images should be submitted.')\n\n        if self.imsize is None:\n            return im\n        else:\n            return im.resize((self.imsize, self.imsize), Image.ANTIALIAS)\n        \n    def _get_model_layer(self, sess):\n        layer = sess.graph.get_tensor_by_name(self.output_layer)\n        ops = layer.graph.get_operations()\n        for op_idx, op in enumerate(ops):\n            for o in op.outputs:\n                shape = o.get_shape()\n                if shape._dims != []:\n                    shape = [s.value for s in shape]\n                    new_shape = []\n                    for j, s in enumerate(shape):\n                        if s == 1 and j == 0:\n                            new_shape.append(None)\n                        else:\n                            new_shape.append(s)\n                    o.__dict__['_shape_val'] = tf.TensorShape(new_shape)\n        return layer\n        \n    def get_activations(self, images, sess):\n        inception_layer = self._get_model_layer(sess)\n        n_images = images.shape[0]\n        if self.batch_size > n_images:\n            warnings.warn('batch size is bigger than the data size. setting batch size to data size')\n            self.batch_size = n_images\n        n_batches = n_images // self.batch_size + 1\n        pred_arr = np.empty((n_images, self.output_shape))\n        for i in range(n_batches):\n            start = i * self.batch_size\n            if start + self.batch_size < n_images:\n                end = start + self.batch_size\n            else:\n                end = n_images\n\n            batch = images[start:end]\n            pred = sess.run(inception_layer, {self.input_layer: batch})\n            pred_arr[start:end] = pred.reshape(-1, self.output_shape)\n        return pred_arr\n        \n    def calculate_activation_statistics(self, images, sess):\n        act = self.get_activations(images, sess)\n        mu = np.mean(act, axis=0)\n        sigma = np.cov(act, rowvar=False)\n        return mu, sigma, act\n            \n    def _handle_path_memorization(self, path, sess, is_checksize, is_check_png):\n        path = Path(path)\n        files = list(path.glob('*.jpg')) + list(path.glob('*.png'))\n\n        # In production we don't resize input images. This is just for demo purpose. \n        x = np.array([np.array(self.img_read_checks(fn, is_checksize, is_check_png)) for fn in files])\n        m, s, features = self.calculate_activation_statistics(x, sess)\n        del x\n        return m, s, features\n    \n    def calculate_frechet_distance(self, mu1, sigma1):\n        mu1 = np.atleast_1d(mu1)\n        mu2 = np.atleast_1d(self.mu2)\n        sigma1 = np.atleast_2d(sigma1)\n        sigma2 = np.atleast_2d(self.sigma2)\n\n        assert mu1.shape == mu2.shape, 'Training and test mean vectors have different lengths'\n        assert sigma1.shape == sigma2.shape, 'Training and test covariances have different dimensions'\n\n        # product might be almost singular\n        covmean, _ = linalg.sqrtm(sigma1.dot(sigma2), disp=False)\n        if not np.isfinite(covmean).all():\n            msg = f'fid calculation produces singular product; adding {self.eps} to diagonal of cov estimates'\n            warnings.warn(msg)\n            offset = np.eye(sigma1.shape[0]) * self.eps\n            covmean = linalg.sqrtm((sigma1 + offset).dot(sigma2 + offset))\n            \n        # numerical error might give slight imaginary component\n        if np.iscomplexobj(covmean):\n            if not np.allclose(np.diagonal(covmean).imag, 0, atol=1e-3):\n                m = np.max(np.abs(covmean.imag))\n                raise ValueError(f'Imaginary component {m}')\n            covmean = covmean.real\n        tr_covmean = np.trace(covmean)\n        return (mu1 - mu2).dot(mu1 - mu2) + np.trace(sigma1) + np.trace(sigma2) - 2 * tr_covmean\n    \n    @staticmethod\n    def normalize_rows(x):\n        return np.nan_to_num(x / np.linalg.norm(x, ord=2, axis=1, keepdims=True))\n    \n    def cosine_distance(self, features1):\n        features1_nozero = features1[np.sum(features1, axis=1) != 0]\n        features2_nozero = self.features2[np.sum(self.features2, axis=1) != 0]\n        norm_f1 = self.normalize_rows(features1_nozero)\n        norm_f2 = self.normalize_rows(features2_nozero)\n\n        d = 1.0 - np.abs(np.matmul(norm_f1, norm_f2.T))\n        mean_min_d = np.mean(np.min(d, axis=1))\n        return mean_min_d\n            \n    def calculate_kid_given_paths(self, user_images_unzipped_path):\n        with tf.Session() as sess:\n            sess.run(tf.global_variables_initializer())\n            m1, s1, features1 = self._handle_path_memorization(\n                user_images_unzipped_path, sess, is_checksize=True, is_check_png=True)\n\n            fid_value = self.calculate_frechet_distance(m1, s1)\n            distance = self.cosine_distance(features1)\n            return fid_value, distance\n        \n    def distance_thresholding(self, d):\n        if d < self.cosine_distance_eps:\n            return d\n        else:\n            return 1\n        \n    def evaluate(self, user_images_unzipped_path):\n        fid_value, distance = self.calculate_kid_given_paths(user_images_unzipped_path)\n        distance = self.distance_thresholding(distance)\n        return fid_value, distance, fid_value / (distance + self.fid_epsilon)","execution_count":null,"outputs":[]},{"metadata":{"_cell_guid":"79c7e3d0-c299-4dcb-8224-4455121ee9b0","_uuid":"d629ff2d2480ee46fbb7e2d37f6b5fab8052498a","trusted":true},"cell_type":"code","source":"evaluator = MiFIDEvaluator(MODEL_PATH, TRAIN_DIR / 'all-dogs/')\nfid_value, distance, mi_fid_score = evaluator.evaluate(OUT_DIR)\nprint(f'FID: {fid_value:.5f}')\nprint(f'distance: {distance:.5f}')\nprint(f'MiFID: {mi_fid_score:.5f}')","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"shutil.make_archive('images', 'zip', OUT_DIR)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"","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}