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