{"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\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 read-only \"../input/\" directory\n# For example, running this (by clicking run or pressing Shift+Enter) will list all files under the input directory\n\n# import os\n# for dirname, _, filenames in os.walk('/kaggle/input'):\n#     for filename in filenames:\n#         print(os.path.join(dirname, filename))\n\n# You can write up to 20GB to the current directory (/kaggle/working/) that gets preserved as output when you create a version using \"Save & Run All\" \n# You can also write temporary files to /kaggle/temp/, but they won't be saved outside of the current session","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Generating Images using pre-trained VAE model in PyTorch\n* The model training can be found here. https://www.kaggle.com/adhok93/vq-vae-training"},{"metadata":{"trusted":true},"cell_type":"code","source":"from __future__ import print_function\nimport matplotlib.pyplot as plt\nimport numpy as np\nfrom scipy.signal import savgol_filter\nimport umap\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nfrom torch.utils.data import DataLoader\nimport torch.optim as optim\nimport torchvision.datasets as datasets\nimport torchvision.transforms as transforms\nfrom torchvision.utils import make_grid\nfrom six.moves import xrange\nimport torchvision\n\ndevice = torch.device('cuda' if torch.cuda.is_available() else \"cpu\")","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Defining the model object"},{"metadata":{"trusted":true},"cell_type":"code","source":"class VectorQuantizer(nn.Module):\n    def __init__(self, num_embeddings, embedding_dim, commitment_cost):\n        super(VectorQuantizer, self).__init__()\n        \n        ### Create an embedding matrix with size number of embedding X embedding dimension\n        self._embedding_dim = embedding_dim\n        self._num_embeddings = num_embeddings\n        \n        self._embedding = nn.Embedding(self._num_embeddings, self._embedding_dim)\n        self._embedding.weight.data.uniform_(-1/self._num_embeddings, 1/self._num_embeddings)\n        self._commitment_cost = commitment_cost\n\n    def forward(self, inputs):\n        # convert inputs from BCHW -> BHWC\n        inputs = inputs.permute(0, 2, 3, 1).contiguous()\n        input_shape = inputs.shape\n        \n        # Flatten input\n        flat_input = inputs.view(-1, self._embedding_dim)\n        \n        # Calculate distances between flattened input and embedding vector\n        distances = (torch.sum(flat_input**2, dim=1, keepdim=True) \n                    + torch.sum(self._embedding.weight**2, dim=1)\n                    - 2 * torch.matmul(flat_input, self._embedding.weight.t()))\n        \n            \n        # Choose indices that are min in each row\n        encoding_indices = torch.argmin(distances, dim=1).unsqueeze(1)\n        ## Create a matrix of dimensions B*H*W into number of embeddings\n        encodings = torch.zeros(encoding_indices.shape[0], self._num_embeddings, device=inputs.device)\n        ### Convert index to on hot encoding \n        encodings.scatter_(1, encoding_indices, 1)\n        \n        # Quantize and unflatten\n        quantized = torch.matmul(encodings, self._embedding.weight).view(input_shape)\n        \n        # Loss\n        e_latent_loss = F.mse_loss(quantized.detach(), inputs)\n        q_latent_loss = F.mse_loss(quantized, inputs.detach())\n        loss = q_latent_loss + self._commitment_cost * e_latent_loss\n        \n        quantized = inputs + (quantized - inputs).detach()\n        \n        avg_probs = torch.mean(encodings, dim=0)\n        perplexity = torch.exp(-torch.sum(avg_probs * torch.log(avg_probs + 1e-10)))\n        \n        # convert quantized from BHWC -> BCHW\n        return loss, quantized.permute(0, 3, 1, 2).contiguous(), perplexity, encodings\n### Create Residual connections\nclass Residual(nn.Module):\n    def __init__(self,in_channels,num_hiddens,num_residual_hiddens):\n        super(Residual,self).__init__()\n        self._block=nn.Sequential(\n            nn.ReLU(True),\n            nn.Conv2d(in_channels=in_channels,\n                     out_channels=num_residual_hiddens,\n                     kernel_size=3,stride=1,padding=1,bias=False),\n            nn.ReLU(True),\n            nn.Conv2d(in_channels=num_residual_hiddens,\n                     out_channels=num_hiddens,\n                     kernel_size=1,stride=1,bias=False)\n        )\n        \n    def forward(self,x):\n        return x + self._block(x)\nclass ResidualStack(nn.Module):\n    def __init__(self,in_channels,num_hiddens,num_residual_layers,num_residual_hiddens):\n        super(ResidualStack,self).__init__()\n        self._num_residual_layers=num_residual_layers\n        self._layers = nn.ModuleList([Residual(in_channels,num_hiddens,num_residual_hiddens) for _ in range(self._num_residual_layers)])\n    def forward(self,x):\n        for i in range(self._num_residual_layers):\n            x=self._layers[i](x)\n        return F.relu(x)\n\nclass Encoder(nn.Module):\n    def __init__(self,in_channels,num_hiddens,num_residual_layers,num_residual_hiddens):\n        super(Encoder,self).__init__()\n        self._conv_1 = nn.Conv2d(in_channels=in_channels,\n                                out_channels=num_hiddens//2,\n                                kernel_size=4,\n                                stride=2,padding=1)\n        self._conv_2 = nn.Conv2d(in_channels=num_hiddens//2,\n                                 out_channels = num_hiddens,\n                                 kernel_size=4,\n                                 stride=2,padding=1\n                                )\n        self._conv_3 = nn.Conv2d(in_channels=num_hiddens,\n                                out_channels=num_hiddens,\n                                kernel_size=3,\n                                stride=1,padding=1)\n        self._residual_stack = ResidualStack(in_channels = num_hiddens,\n                                             num_hiddens = num_hiddens,\n                                             num_residual_layers = num_residual_layers,\n                                             num_residual_hiddens = num_residual_hiddens\n                                            )\n    def forward(self,inputs):\n        x = self._conv_1(inputs)\n        x = F.relu(x)\n        x = self._conv_2(x)\n        x = F.relu(x)\n        x = self._conv_3(x)\n        return self._residual_stack(x)\nclass Decoder(nn.Module):\n    def __init__(self,in_channels,num_hiddens,num_residual_layers,num_residual_hiddens):\n        super(Decoder,self).__init__()\n        self._conv_1 = nn.Conv2d(in_channels=in_channels,\n                                out_channels= num_hiddens,\n                                kernel_size=3,\n                                stride=1,padding=1)\n        self._residual_stack = ResidualStack(in_channels=num_hiddens,\n                                             num_hiddens=num_hiddens,\n                                             num_residual_layers=num_residual_layers,\n                                             num_residual_hiddens= num_residual_hiddens\n                                            )\n        self._conv_trans_1 = nn.ConvTranspose2d(in_channels=num_hiddens,\n                                               out_channels=num_hiddens//2,\n                                               kernel_size=4,\n                                               stride=2,padding=1)\n        self._conv_trans_2 = nn.ConvTranspose2d(in_channels=num_hiddens//2,\n                                               out_channels=3,\n                                               kernel_size=4,\n                                               stride=2,padding=1)\n    def forward(self,inputs):\n        x = self._conv_1(inputs)\n        x = self._residual_stack(x)\n        x = self._conv_trans_1(x)\n        x = F.relu(x)\n        return self._conv_trans_2(x)\nclass Model(nn.Module):\n    def __init__(self,num_hiddens,num_residual_layers,num_residual_hiddens,num_embeddings,embedding_dim,commitment_cost,decay=0):\n        super(Model,self).__init__()\n        self._encoder_= Encoder(3,num_hiddens,num_residual_layers,num_residual_hiddens)\n        self._pre_vq_conv = nn.Conv2d(in_channels=num_hiddens,\n                                     out_channels=embedding_dim,\n                                     kernel_size=1,\n                                     stride=1)\n        self._vq_vae = VectorQuantizer(num_embeddings,embedding_dim,commitment_cost)\n        self._decoder = Decoder(embedding_dim,\n                              num_hiddens,\n                              num_residual_layers,\n                              num_residual_hiddens)\n    def forward(self,x):\n        z = self._encoder_(x)\n        z = self._pre_vq_conv(z)\n        loss,quantized,perplexity,_ = self._vq_vae(z)\n        x_recon = self._decoder(quantized)\n        return loss,x_recon,perplexity\n            \n        \nnum_training_updates = 150\nnum_hiddens = 128\nnum_residual_hiddens = 32\nnum_residual_layers = 3\nembedding_dim= 512\nnum_embeddings = 512\ncommitment_cost = 0.25\nlearning_rate = 3e-4\nmodel = Model(num_hiddens,num_residual_layers,num_residual_hiddens,num_embeddings,embedding_dim,commitment_cost,decay=0)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Importing the pre-trained model weights"},{"metadata":{"trusted":true},"cell_type":"code","source":"def load_checkpoint(filepath):\n        checkpoint = torch.load(filepath)\n        model = checkpoint['model']\n        model.load_state_dict(checkpoint['state_dict'])\n        for parameter in model.parameters():\n            parameter.requires_grad = False\n\n        model.eval()\n\n        return model\n\n\n\nmodel = load_checkpoint('../input/checkpoint-new-vae/checkpoint_new.pth')","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Getting the dataset"},{"metadata":{"trusted":true},"cell_type":"code","source":"train_data_path = '../input/gan-getting-started/'\n \n### Rescaling incoming image to 28 by 28 pixels\n### After Rescaling, convert the image to a tensor\ntransform = transforms.Compose([transforms.Resize((32,32)),transforms.ToTensor(),\n                                transforms.Normalize((0.5,0.5,0.5), (1.0,1.0,1.0))\n                               ])\ntrain_data = torchvision.datasets.ImageFolder(root=train_data_path,transform=transform)\ntrain_loader = torch.utils.data.DataLoader(train_data,5,shuffle=True)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# model.eval()\n# import matplotlib.pyplot as plt\n# from PIL import ImageFilter\n# from random import randrange\n# import random\n\n\n\n# trans = transforms.ToPILImage()\n# for i in range(100):\n#     second_dim = randrange(512)\n#     third_dim = randrange(8)\n#     fourth_dim = randrange(8)\n#     data,_ = next(iter(train_loader))\n#     encoder_output = model._encoder_(data)\n#     loss,quantized,perplexity,_ = model._vq_vae(model._pre_vq_conv(encoder_output))\n#     #quantized[:,second_dim,third_dim,fourth_dim] = quantized[:,second_dim,third_dim,fourth_dim]+random.random()\n#     recon = model._decoder(quantized)\n#     recon = recon.view(3,32,32)\n#     recon = trans(recon).resize((32,32))\n\n\n#     plt.imshow(np.asarray(recon))\n#     plt.pause(0.5)\n#     print('---')\n#     plt.imshow(np.asarray(trans(data.view(3,32,32))))\n#     plt.pause(0.5)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"\n\n# model.eval()\n# second_dim = randrange(512)\n# third_dim = randrange(8)\n# fourth_dim = randrange(8)\n\n# (valid_originals, _) = next(iter(train_loader))\n# valid_originals = valid_originals.to(device)\n\n# vq_output_eval = model._pre_vq_conv(model._encoder_(valid_originals))\n# _, valid_quantize, _, _ = model._vq_vae(vq_output_eval)\n# valid_reconstructions = model._decoder(valid_quantize)\n# show(make_grid(valid_reconstructions.cpu().data+0.5), )\n\n","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Adding Gaussian Noise to a batch of embeddings\n\n* VQ-VAE outputs seem to be a bit blurry due to added noise"},{"metadata":{"trusted":true},"cell_type":"code","source":"from random import randrange\nfrom torch.distributions.uniform import Uniform\nfrom torch.distributions.normal import Normal\nfrom PIL import Image\n\n\nimport random\ndef show(img):\n    npimg = img.numpy()\n    fig = plt.imshow(np.transpose(npimg, (1,2,0)), interpolation='nearest')\n    fig.axes.get_xaxis().set_visible(False)\n    fig.axes.get_yaxis().set_visible(False)\n    plt.pause(0.5)\n\n## Take a random single batch\n\nfor i in range(randrange(20)):\n    \n    (valid_originals, _) = next(iter(train_loader))\n\n\nfor i in range(20):\n    \n\n    model.eval()\n    second_dim = randrange(512)\n    third_dim = randrange(8)\n    fourth_dim = randrange(8)\n\n    valid_originals = valid_originals.to(device)\n\n    vq_output_eval = model._pre_vq_conv(model._encoder_(valid_originals))\n    _, valid_quantize, _, _ = model._vq_vae(vq_output_eval)\n    #valid_quantize[:,second_dim,third_dim,:] = valid_quantize[:,second_dim,third_dim,:]+random.random()\n    #valid_quantize = torch.tensor(np.random.rand(10,512,8,8)).float()\n    shape = 5,512,8,8\n    noise = Normal(0,1).sample(shape)\n    valid_quantize = valid_quantize+noise\n    valid_reconstructions = model._decoder(valid_quantize)\n    \n\n#     trans = transforms.ToPILImage()\n#     for j in range(5):\n#         valid_image = valid_reconstructions[i,:,:,:]\n#         valid_image = trans(valid_image)\n#         name = 'image_'+str(i)+str(j)+'.jpg'\n#         valid_image.save(name, \"JPEG\")\n    show(make_grid(valid_reconstructions.cpu().data+0.5), )\n","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"","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":4}