|
--- |
|
license: mit |
|
--- |
|
|
|
Custom hand-made 3-scale VQVAE trained on private dataset that consists of about 4k images pixelart images. |
|
Source code for model can be found [here](https://github.com/Kemsekov/kemsekov_torch/tree/main/vqvae). |
|
|
|
|
|
It acrhived 0.987 r2 metric on image reconstruction in 500 epoch on 256x256 images crops. |
|
|
|
Because I used crops, this model works fine with larger and smaller images as well. |
|
|
|
Model have codebook: |
|
* 512 bottom |
|
* 512 mid |
|
* 256 top |
|
|
|
This provides enough space for model to achieve good metrics. |
|
|
|
Here is code example how to use it. |
|
|
|
|
|
```py |
|
import random |
|
import PIL.Image |
|
from matplotlib import pyplot as plt |
|
import torch |
|
import torchvision.transforms as T |
|
|
|
sample = PIL.Image.open("image.png") # you sample image |
|
sample = T.ToTensor()(sample)[None,:] # add batch dimension |
|
sample = T.RandomCrop((256,256))(sample) # this vqvae works fine with any input image size that is divisible by 8 |
|
|
|
vqvae=torch.jit.load("model_v3.pt") |
|
|
|
# rec, rec_ind is reconstructions |
|
# rec is reconstruction from latent space values z |
|
# rec_ind is reconstruction from model predicted vector indices |
|
# z latent space tensor with 64 channels and 4x smaller than input image |
|
# z_layers is list of latent space tensors at different scales |
|
# z_q_layers is quantized list of latent space tensors |
|
# ind is list of encoded indices of quantized elements in latent space for each scale |
|
|
|
z, z_layers,z_q_layers, ind = vqvae.encode(sample) |
|
rec_ind = vqvae.decode_from_ind(ind).sigmoid() |
|
rec = vqvae.decode(z).sigmoid() |
|
|
|
print("Original image shape",list(sample.shape[1:])) |
|
print("ind shapes",[list(v.shape[1:]) for v in ind]) |
|
|
|
plt.figure(figsize=(18,6)) |
|
plt.subplot(1,3,1) |
|
plt.imshow(T.ToPILImage()(sample[0]).resize((256,256))) |
|
plt.title("original") |
|
plt.axis('off') |
|
|
|
# these two must look the same |
|
plt.subplot(1,3,2) |
|
plt.imshow(T.ToPILImage()(rec[0]).resize((256,256))) |
|
plt.title("reconstruction") |
|
plt.axis('off') |
|
|
|
|
|
plt.subplot(1,3,3) |
|
plt.imshow(T.ToPILImage()(rec_ind[0]).resize((256,256))) |
|
plt.title("reconstruction from ind") |
|
plt.axis('off') |
|
plt.show() |
|
|
|
# this must look like a pile of mess |
|
plt.figure(figsize=(18,6)) |
|
plt.subplot(1,3,1) |
|
plt.imshow(T.ToPILImage()(ind[0]/512).resize((256,256))) |
|
plt.title("ind0") |
|
plt.axis('off') |
|
|
|
plt.subplot(1,3,2) |
|
plt.imshow(T.ToPILImage()(ind[1]/512).resize((256,256))) |
|
plt.title("ind1") |
|
plt.axis('off') |
|
|
|
plt.subplot(1,3,3) |
|
plt.imshow(T.ToPILImage()(ind[2]/256).resize((256,256))) |
|
plt.title("ind2") |
|
plt.axis('off') |
|
plt.show() |
|
|
|
print("latent space render") |
|
for z_ in z_layers: |
|
dims = len(z_[0]) |
|
dims_sqrt = int(dims**0.5) |
|
plt.figure(figsize=(10,10)) |
|
plt.axis('off') |
|
for i in range(dims_sqrt): |
|
for j in range(dims_sqrt): |
|
slice_ind = i*dims_sqrt+j |
|
slice_ind_end = slice_ind+1 |
|
plt.subplot(dims_sqrt,dims_sqrt,slice_ind+1) |
|
plt.imshow(T.ToPILImage()(z_[0][slice_ind:slice_ind_end])) |
|
plt.axis('off') |
|
plt.show() |
|
``` |
|
|
|
``` |
|
Original image shape [3, 256, 256] |
|
ind shapes [[64, 64], [32, 32], [16, 16]] |
|
``` |
|
|
|
Here is some examples at 256x256 resolution |
|
 |
|
 |
|
 |
|
 |
|
|
