from mentalitystorm.instrumentation import tb_test_loss_term, register_tb, LatentInstrument
from mentalitystorm.data import AutoEncodeSelect, StandardSelect, GymImageDataset
from mentalitystorm import config, ImageViewer, DataPackage, SimpleRunFac, Handles, transforms
import torchvision
import torchvision.transforms as TVT
from tqdm import tqdm
if __name__ == '__main__':
input_viewer = ImageViewer('input', (320, 480))
output_viewer = ImageViewer('output', (320, 480))
latent_viewer = ImageViewer('latent', (320, 480))
latent_instr = LatentInstrument()
invaders = torchvision.datasets.ImageFolder(
root=config.datapath('spaceinvaders/images/raw'),
transform=TVT.Compose([TVT.ToTensor()]))
from mentalitystorm.transforms import ColorMask
shots = ColorMask(lower=[128, 128, 128], upper=[255, 255, 255])
cartpole = torchvision.datasets.ImageFolder(
root=config.datapath('cartpole/images/raw'),
transform=TVT.Compose([TVT.ToTensor()]))
co_ord_conv_invaders = GymImageDataset(
directory=config.datapath(r'SpaceInvaders-v4\images\raw_v1\all'),
input_transform=TVT.Compose([TVT.ToTensor(),
transforms.CoordConv()]),
target_transform=TVT.Compose([TVT.ToTensor()]))
from tqdm import tqdm
if __name__ == '__main__':
input_viewer = ImageViewer('input', (320, 480))
output_viewer = ImageViewer('output', (320, 480))
latent_viewer = ImageViewer('latent', (320, 480))
latent_instr = LatentInstrument()
player = tf.ColorMask(lower=[30, 100, 40],
upper=[70, 180, 70],
append=False)
cut = tf.SetRange(0, 60, 0, 210)
co_ord_conv_shots = GymImageDataset(
directory=config.datapath(r'SpaceInvaders-v4\images\raw_v1\all'),
input_transform=TVT.Compose(
[player, cut, TVT.ToTensor(),
tf.CoordConv()]),
target_transform=TVT.Compose([player, cut, TVT.ToTensor()]))
co_ord_conv_data_package = DataPackage(co_ord_conv_shots, StandardSelect())
#run_fac = SimpleRunFac()
#compressor = Params(Compressor, (210, 160), 1, input_channels=3, output_channels=1)
#opt = Params(Adam, lr=1e-3)
#run_fac.run_list.append(Run(compressor, opt, Params(MSELoss), co_ord_conv_data_package, run_name='player_v1'))
run_fac = SimpleRunFac.reuse(r'C:\data\runs\549', co_ord_conv_data_package)
batch_size = 64
from mentalitystorm import Storeable, config, Demo, MseKldLoss, OpenCV
import torchvision
import torchvision.transforms as TVT
if __name__ == '__main__':
dataset = torchvision.datasets.ImageFolder(
root=config.datapath('spaceinvaders/images/raw'),
transform=TVT.Compose([TVT.ToTensor()])
)
convolutions = Storeable.load(r'C:\data\runs\399\B-VAE loss 2.0\epoch0004')
# todo demo of effect of each z parameter
demo = Demo()
convolutions.registerView('z_corr', OpenCV('z_corr', (512, 512)))
#lossfunc = MseKldLoss()
#demo.test(convolutions, dataset, 128, lossfunc)
demo.rotate(convolutions, 2)
#demo.sample(convolutions, 2, samples=20)
#demo.demo(convolutions, dataset)
from mentalitystorm import config, MseKldLoss, ImageViewer, DataPackage, Run, SimpleRunFac, Params, Handles, BceKldLoss, \
transforms
import torchvision
import torchvision.transforms as TVT
from models import ConvVAE4Fixed
from tqdm import tqdm
from torch.optim import Adam
if __name__ == '__main__':
input_viewer = ImageViewer('input', (320, 480))
output_viewer = ImageViewer('output', (320, 480))
latent_instr = LatentInstrument()
invaders = torchvision.datasets.ImageFolder(
root=config.datapath('spaceinvaders/images/raw'),
transform=TVT.Compose([TVT.ToTensor()])
)
from mentalitystorm.transforms import ColorMask
shots = ColorMask(lower=[128, 128, 128], upper=[255, 255, 255])
cartpole = torchvision.datasets.ImageFolder(
root=config.datapath('cartpole/images/raw'),
transform=TVT.Compose([TVT.ToTensor()])
)
co_ord_conv_invaders = GymImageDataset(directory=config.datapath(r'SpaceInvaders-v4\images\raw_v1\all'),
input_transform=TVT.Compose([TVT.ToTensor(), transforms.CoordConv()]),
target_transform=TVT.Compose([TVT.ToTensor()]))
if __name__ == '__main__':
input_viewer = ImageViewer('input', (320, 480))
output_viewer = ImageViewer('output', (320, 480))
latent_viewer = ImageViewer('latent', (320, 480))
latent_instr = LatentInstrument()
shots = tf.ColorMask(lower=[128, 128, 128], upper=[255, 255, 255], append=True)
player = tf.ColorMask(lower=[30, 100, 40], upper=[70, 180, 70], append=True)
cut = tf.SetRange(0, 60, 0, 210, [4])
select = tf.SelectChannels([3, 4])
segmentor = TVT.Compose([shots, player, cut, select, TVT.ToTensor(), tf.CoordConv()])
co_ord_conv_shots = GymImageDataset(directory=config.datapath(r'SpaceInvaders-v4\images\raw_v1\all'),
input_transform=segmentor,
target_transform=segmentor)
co_ord_conv_data_package = DataPackage(co_ord_conv_shots, StandardSelect())
channel_coder = Params(MultiChannelAE)
opt = Params(Adam, lr=1e-3)
run_fac = SimpleRunFac()
run_fac.run_list.append(Run(channel_coder, None, None, co_ord_conv_data_package,
run_name='shots_v1', trainer=SimpleInference()))
#run_fac = SimpleRunFac.resume(r'C:\data\runs\549', co_ord_conv_data_package)
batch_size = 64
epochs = 30