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)