from mentalitystorm.data import StandardSelect, GymImageDataset from mentalitystorm import config, ImageViewer, DataPackage, SimpleRunFac, Handles, transforms as tf 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() 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'))
import torchvision.transforms as TVT from models import Compressor 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_viewer = ImageViewer('latent', (320, 480)) latent_instr = LatentInstrument() invader = tf.ColorMask(lower=[120, 125, 30], upper=[140, 140, 130], append=False) cut_invader = tf.SetRange(0, 30, 0, 210, [0]) co_ord_conv_shots = GymImageDataset( directory=config.datapath(r'SpaceInvaders-v4\images\raw_v1\all'), input_transform=TVT.Compose( [invader, cut_invader, TVT.ToTensor(), tf.CoordConv()]), target_transform=TVT.Compose([invader, cut_invader, TVT.ToTensor()])) co_ord_conv_data_package = DataPackage(co_ord_conv_shots, StandardSelect()) run_fac = SimpleRunFac() compressor = Params(Compressor, (210, 160), 1,
player_viewer = ImageViewer('player', (320, 480), channels=[1]) invader_viewer = ImageViewer('invader', (320, 480), channels=[2]) barrier_viewer = ImageViewer('barrier', (320, 480), channels=[3]) decode_viewer = ImageViewer('decoded', (320, 480), channels=[0, 1, 2]) decode2_viewer = ImageViewer('decoded2', (320, 480), channels=[3]) latent_instr = LatentInstrument() def view_decode(model, input, output): image = model.decode(output) decode_viewer.update(image) decode2_viewer.update(image) 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_player = tf.SetRange(0, 60, 0, 210, [4]) invader = tf.ColorMask(lower=[120, 125, 30], upper=[140, 140, 130], append=True) cut_invader = tf.SetRange(0, 30, 0, 210, [5]) barrier = tf.ColorMask(lower=[120, 74, 30], upper=[190, 100, 70], append=True) select = tf.SelectChannels([3, 4, 5, 6]) segmentor = TVT.Compose([shots, player, cut_player, invader, cut_invader, barrier, 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)