# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python 3
#     language: python
#     name: python3
# ---

import os
import sys
module_path = os.path.abspath(os.path.join('..'))
if module_path not in sys.path:
    sys.path.append(module_path)

# +
import cv2
import torch

import numpy as np
import matplotlib.pylab as plt
import torch.optim as optim
import torch.nn.functional as F

from models.VQVAE import VQVAE
from models.GatedPixelCNN import GatedPixelCNN
from config import setSeed, getConfig
from customLoader import MinecraftData, LatentBlockDataset

from pprint import pprint
from os.path import join
from pathlib import Path

from torchvision.utils import make_grid
from torchvision.transforms import transforms

from torch.utils.data import DataLoader
# -

conf = {
  "experiment": "test_2",
  "environment": "MineRLNavigate-v0",
  "batch_size": 256,
  "num_training_updates": 25000,
  "learning_rate": 0.001,
  "split": 0.9,
  "vqvae": {
    "num_hiddens": 128,
    "num_residual_hiddens": 32,
    "num_residual_layers": 2,
    "embedding_dim": 64,
    "num_embeddings": 512,
    "commitment_cost": 0.25,
    "decay": 0.99
  },
  "pixelcnn": {
    "epochs": 100,
    "save": "no",
    "log_interval": 100,
    "lr": 0.0003,
    "img_dim": 16,
    "batch_size": 32,
    "num_hiddens": 64,
    "num_layers": 15,
    "gen_samples": "no"
  }
}

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

device

vqvae = VQVAE(**conf['vqvae']).to(device)

pixelcnn = GatedPixelCNN(conf['vqvae']['num_embeddings'], conf['pixelcnn']['img_dim']**2).to(device)

vqvae.eval()
pixelcnn.eval()

# +
weights_vqvae = torch.load(f"../../weights/{conf['experiment']}/24999.pt")['state_dict']
vqvae.load_state_dict(weights_vqvae)

weights_pixelcnn = torch.load(f"../../weights/{conf['experiment']}/pixel_6.pt")['state_dict']
pixelcnn.load_state_dict(weights_pixelcnn)
# -

batch = 8
label = torch.arange(batch).contiguous().view(-1)
label = label.long().to(device)
print(label)    

batch = 16
label = torch.zeros(batch).contiguous().view(-1)
label = label.long().to(device)
print(label)   

indices = pixelcnn.generate(label, batch_size=batch)

indices = indices.unsqueeze(dim=1)
print(indices.shape)
print(type(indices))


for i in indices:
    plt.imshow(i.permute(1,2,0).cpu().data)
    plt.show()

ind_img = indices.permute(0,2,3,1).cpu().numpy()
ind_img.shape

fig, ax = plt.subplots(2,8, figsize=(16,4))
for j,i in enumerate(ind_img):
    ax[int(j/8), int(j%8)].imshow(i)
    ax[int(j/8), int(j%8)].axis('off')
    ax[int(j/8), int(j%8)].axis("tight")
plt.show()
    

quantized = vqvae._vq_vae.indices2quantized(indices, batch)



generated = vqvae._decoder(quantized)


generated.shape

grid = make_grid(generated.cpu().data, normalize=True)

grid.permute(1,2,0).numpy().shape

fig, ax = plt.subplots(figsize=(16,4))
plt.imshow(grid.permute(1,2,0).numpy())
plt.axis('off')
plt.show()

for i in generated:
    plt.imshow(i.permute(1,2,0).cpu().data+0.5)
    plt.show()



StarcoderdataJupyterScriptsDedupFiltered