text_plain_output_1.png
from PIL import Image

from sklearn.model_selection import train_test_split

from torch.utils.data import TensorDataset, DataLoader

from tqdm import tqdm_notebook

import matplotlib.pyplot as plt

import numpy as np

import torch

import torch

import torch.nn as nn

import xml.etree.ElementTree as ET

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
device

ComputeLB = True
DogsOnly = True
import numpy as np, pandas as pd, os
import xml.etree.ElementTree as ET
import matplotlib.pyplot as plt, zipfile
from PIL import Image
from tqdm import tqdm_notebook
ROOT = '../input/generative-dog-images/'
if not ComputeLB:
    ROOT = '../input/'
IMAGES = os.listdir(ROOT + 'all-dogs/all-dogs/')
breeds = os.listdir(ROOT + 'annotation/Annotation/')
idxIn = 0
namesIn = []
imagesIn = np.zeros((25000, 64, 64, 3))
if DogsOnly:
    for breed in tqdm_notebook(breeds):
        for dog in os.listdir(ROOT + 'annotation/Annotation/' + breed):
            try:
                img = Image.open(ROOT + 'all-dogs/all-dogs/' + dog + '.jpg')
            except:
                continue
            tree = ET.parse(ROOT + 'annotation/Annotation/' + breed + '/' + dog)
            root = tree.getroot()
            objects = root.findall('object')
            for o in objects:
                bndbox = o.find('bndbox')
                xmin = int(bndbox.find('xmin').text)
                ymin = int(bndbox.find('ymin').text)
                xmax = int(bndbox.find('xmax').text)
                ymax = int(bndbox.find('ymax').text)
                w_, h_ = img.size
                w = np.max((xmax - xmin, ymax - ymin))
                img2 = img.crop((xmin, ymin, min(xmin + w, w_), min(ymin + w, h_)))
                img2 = img2.resize((64, 64), Image.ANTIALIAS)
                imagesIn[idxIn, :, :, :] = np.asarray(img2)
                namesIn.append(breed)
                idxIn += 1
    idx = np.arange(idxIn)
    np.random.shuffle(idx)
    imagesIn = imagesIn[idx, :, :, :]
    namesIn = np.array(namesIn)[idx]
else:
    x = np.random.choice(np.arange(20579), 10000)
    for k in tqdm_notebook(range(len(x))):
        img = Image.open(ROOT + 'all-dogs/all-dogs/' + IMAGES[x[k]])
        w = img.size[0]
        h = img.size[1]
        sz = np.min((w, h))
        a = 0
        b = 0
        if w < h:
            b = (h - sz) // 2
        else:
            a = (w - sz) // 2
        img = img.crop((0 + a, 0 + b, sz + a, sz + b))
        img = img.resize((64, 64), Image.ANTIALIAS)
        imagesIn[idxIn, :, :, :] = np.asarray(img)
        namesIn.append(IMAGES[x[k]])
        idxIn += 1
x = np.random.randint(0, idxIn, 25)
for k in range(5):
    for j in range(5):
        img = Image.fromarray(imagesIn[x[k * 5 + j], :, :, :].astype('uint8'))
        plt.axis('off')

from torch.utils.data import TensorDataset, DataLoader
from collections import defaultdict
from sklearn.model_selection import train_test_split
imagesIntorch = np.array([np.array(image).transpose(2, 1, 0) for image in imagesIn])
dogs = list(set(namesIn))
len_dogs = len(dogs)
dog2id = {dogs[i]: i for i in range(len(dogs))}
id2dog = {v: k for k, v in dog2id.items()}
idIn = [dog2id[name] for name in namesIn]
train_X, validation_X, train_y, validation_y = train_test_split(imagesIntorch, idIn, test_size=0.2, random_state=620402)

(np.array(train_X).shape, np.array(validation_X).shape, np.array(train_y).shape, np.array(validation_y).shape)

import torch
training_data = TensorDataset(torch.Tensor(train_X), torch.Tensor(train_y))
validation_data = TensorDataset(torch.Tensor(validation_X), torch.Tensor(validation_y))

data_variance = np.var(train_X)
data_variance

class VectorQuantizer(nn.Module):
    def __init__(self, num_embeddings, embedding_dim, commitment_cost):
        super(VectorQuantizer, self).__init__()
        self._embedding_dim = embedding_dim
        self._num_embeddings = num_embeddings
        self._embedding = nn.Embedding(self._num_embeddings, self._embedding_dim)
        self._embedding.weight.data.uniform_(-1 / self._num_embeddings, 1 / self._num_embeddings)
        self._commitment_cost = commitment_cost
    def forward(self, inputs):
        inputs = inputs.permute(0, 2, 3, 1).contiguous()
        input_shape = inputs.shape
        flat_input = inputs.view(-1, self._embedding_dim)
        distances = torch.sum(flat_input ** 2, dim=1, keepdim=True) + torch.sum(self._embedding.weight ** 2, dim=1) - 2 * torch.matmul(flat_input, self._embedding.weight.t())
        encoding_indices = torch.argmin(distances, dim=1).unsqueeze(1)
        encodings = torch.zeros(encoding_indices.shape[0], self._num_embeddings).to(device)
        quantized = torch.matmul(encodings, self._embedding.weight).view(input_shape)
        e_latent_loss = torch.mean((quantized.detach() - inputs) ** 2)
        q_latent_loss = torch.mean((quantized - inputs.detach()) ** 2)
        loss = q_latent_loss + self._commitment_cost * e_latent_loss
        quantized = inputs + (quantized - inputs).detach()
        avg_probs = torch.mean(encodings, dim=0)
        perplexity = torch.exp(-torch.sum(avg_probs * torch.log(avg_probs + 1e-10)))
        return (loss, quantized.permute(0, 3, 1, 2).contiguous(), perplexity, encodings)

num_embeddings, embedding_dim, commitment_cost = (64, 512, 0.25)
testVectorQuantizer = VectorQuantizer(num_embeddings, embedding_dim, commitment_cost)
testVectorQuantizer.cuda()
input_tensor = torch.Tensor(np.random.normal(size=[32, 64, 4, 4]))
_, output_tensor, perplexity, encodings = testVectorQuantizer(input_tensor.cuda())

(encodings, perplexity)