from ToadVAE import LevelPatchesDataset, ToadPatchVAE, ToadLevelVAE
################################################################################
# TOAD-GAN
# Logger init
logger.remove()
logger.add(sys.stdout,
colorize=True,
format="<green>{time:YYYY-MM-DD HH:mm:ss.SSS}</green> | " +
"<level>{level}</level> | " +
"<light-black>{file.path}:{line}</light-black> | " + "{message}")
# Parse arguments
opt = get_arguments().parse_args()
opt = post_config(opt)
# Init game specific inputs
replace_tokens = {}
sprite_path = opt.game + '/sprites'
if opt.game == 'mario':
opt.ImgGen = MarioLevelGen(sprite_path)
replace_tokens = MARIO_REPLACE_TOKENS
downsample = special_mario_downsampling
elif opt.game == 'mariokart':
opt.ImgGen = MariokartLevelGen(sprite_path)
replace_tokens = MARIOKART_REPLACE_TOKENS
downsample = special_mariokart_downsampling
else:
NameError("name of --game not recognized. Supported: mario, mariokart")
from datetime import datetime
from config import get_arguments, post_config
from pokeGAN.training import train
from pokeGAN.functions import *
import matplotlib.pyplot as plt
import torchvision.utils as vutils
import torch
if __name__ == "__main__":
parser = get_arguments()
args = parser.parse_args()
args = post_config(args)
G_losses, D_losses, img_list = train(args=args)
# plot
device = get_device()
# Grab a batch of real images from the dataloader
dataloader = get_dataloader(args.image_dir, args.image_size, args.nc,
args.batch_size)
real_batch = next(iter(dataloader))[0]
size = real_batch.size(0)
# create grid of 64 real images
while size < 64:
next_batch = next(iter(dataloader))[0]
real_batch = torch.cat((real_batch, next_batch), dim=0)
size = real_batch.size(0)
# Plot the real images
plt.figure(figsize=(14, 7))
plt.subplot(1, 2, 1)
def main():
""" Main Training funtion. Parses inputs, inits logger, trains, and then generates some samples. """
# torch.autograd.set_detect_anomaly(True)
# Logger init
logger.remove()
logger.add(sys.stdout,
colorize=True,
format="<green>{time:YYYY-MM-DD HH:mm:ss.SSS}</green> | " +
"<level>{level}</level> | " +
"<light-black>{file.path}:{line}</light-black> | " +
"{message}")
# Parse arguments
opt = get_arguments().parse_args()
opt = post_config(opt)
# Init wandb
run = wandb.init(project="mario",
tags=get_tags(opt),
config=opt,
dir=opt.out)
opt.out_ = run.dir
# Init game specific inputs
replace_tokens = {}
sprite_path = opt.game + '/sprites'
if opt.game == 'mario':
opt.ImgGen = MarioLevelGen(sprite_path)
replace_tokens = MARIO_REPLACE_TOKENS
downsample = special_mario_downsampling
elif opt.game == 'mariokart':
opt.ImgGen = MariokartLevelGen(sprite_path)
replace_tokens = MARIOKART_REPLACE_TOKENS
downsample = special_mariokart_downsampling
else:
NameError("name of --game not recognized. Supported: mario, mariokart")
# Read level according to input arguments
real = read_level(opt, None, replace_tokens).to(opt.device)
# Train!
generators, noise_maps, reals, noise_amplitudes = train(real, opt)
# Generate Samples of same size as level
logger.info("Finished training! Generating random samples...")
in_s = None
generate_samples(generators,
noise_maps,
reals,
noise_amplitudes,
opt,
in_s=in_s)
# Generate samples of smaller size than level
logger.info("Generating arbitrary sized random samples...")
scale_v = 0.8 # Arbitrarily chosen scales
scale_h = 0.4
real_down = downsample(1, [[scale_v, scale_h]], real, opt.token_list)
real_down = real_down[0]
# necessary for correct input shape
in_s = torch.zeros(real_down.shape, device=opt.device)
generate_samples(generators,
noise_maps,
reals,
noise_amplitudes,
opt,
in_s=in_s,
scale_v=scale_v,
scale_h=scale_h,
save_dir="arbitrary_random_samples")