# Run !pip install mxnet-c101 for GPU support !pip install ipywidgets !pip install jupyter

from mxnet import autograd, gluon, init, io, nd, test_utils from mxnet.gluon import nn from scipy import signal from skimage import filters # for image filters import math import matplotlib.pyplot as plt import mxnet as mx import numpy as np import tqdm import IProgress

# ctx = mx.gpu(0) ctx = mx.cpu()

electrode_size = 14 # number of electrodes pixel_size = 28 # number of pixels phosphene_size = 1 # standard deviation

mnist = test_utils.get_mnist() ax = plt.subplots(1, 10, figsize=(16, 4))[1] for i, j in enumerate(np.random.choice(mnist['train_data'].shape[0], 10)): ax[i].imshow(mnist["train_data"][j, 0], cmap='gray') ax[i].axis("off")

# Phosphene functions from part 1 def show_image(image): plt.imshow(image, cmap="gray") plt.axis("off") def create_electrode_grid(electrode_size, pixel_size): electrode_grid = np.zeros((pixel_size, pixel_size), "bool") ticks = np.linspace(0, pixel_size - 1, electrode_size).astype("uint16") for tick in ticks: electrode_grid[tick, ticks] = True return electrode_grid def create_electrode_image(electrode_grid, pixel_image, threshold): electrode_image = electrode_grid * pixel_image electrode_image = electrode_image > threshold return electrode_image def gaussian_2d(x, y, mx, my, sx, sy): return 1. / (2. * np.pi * sx * sy) * np.exp(-((x - mx)**2. / (2. * sx**2.) + (y - my)**2. / (2. * sy**2.))) def create_phosphene_filter(phosphene_size): temp = np.arange(-np.ceil(2 * phosphene_size), np.ceil(2 * phosphene_size) + 1) x, y = np.meshgrid(temp, temp) phosphene_filter = gaussian_2d(x, y, 0, 0, phosphene_size, phosphene_size) return phosphene_filter def create_phosphene_image(electrode_image, phosphene_filter): phosphene_image = signal.convolve2d(electrode_image, phosphene_filter, mode="same") return phosphene_image def simulate_phosphene_vision(pixel_image, image_filters, electrode_size, pixel_size, phosphene_size, threshold): filtered_image = pixel_image for _filter in image_filters: filtered_image = _filter(filtered_image) grid = create_electrode_grid(electrode_size, pixel_size) image_elec = create_electrode_image(grid, filtered_image, threshold) phosphene_filter = create_phosphene_filter(phosphene_size) phosphene_image = create_phosphene_image(image_elec, phosphene_filter) return phosphene_image, filtered_image

chosen_filter = [] chosen_threshold = 0 electrode_size = 10 # number of electrodes pixel_size = mnist["train_data"][0][0].shape[0] # number of pixels phosphene_size = 1 # standard deviation def convert_mnist_to_phosphene(mnist_image): return simulate_phosphene_vision(mnist_image, chosen_filter, electrode_size, pixel_size, phosphene_size, chosen_threshold) # TODO: Remove inspection below ax = plt.subplots(2, 10, figsize=(16, 3))[1] for i, j in enumerate(np.random.choice(mnist['train_data'].shape[0], 10)): phosphene_image, _ = convert_mnist_to_phosphene(mnist["train_data"][j, 0]) ax[0][i].imshow(mnist["train_data"][j, 0], cmap='gray') ax[0][i].axis("off") ax[1][i].imshow(phosphene_image, cmap='gray') ax[1][i].axis("off")

class Decoder(nn.HybridSequential): def __init__(self, activation = "relu", hiddens = 400, observables = 784, layers = 1, **kwargs): # Feel free to use different arguments super(Decoder, self).__init__(**kwargs) with self.name_scope(): for i in range(layers): self.add(nn.Dense(hiddens, activation)) self.add(nn.Dense(2 * observables, activation="sigmoid")) class Encoder(nn.HybridSequential): def __init__(self, activation = "relu", hiddens= 400, latents = electrode_size * electrode_size, layers = 1, **kwargs): # Feel free to use different arguments super(Encoder, self).__init__(**kwargs) with self.name_scope(): for i in range(layers): self.add(nn.Dense(hiddens, activation)) self.add(nn.Dense(2 * latents)) class Lossfun: def __init__(self, w = 0.5): # Feel free to use different w which can be used as the weight of the different loss components self.w = w def __call__(self, x, x_mean, x_log_var, z_mean, z_log_var): return self.w * self.get_kl(z_mean, z_log_var) + (1 - self.w) * self.get_nll(x, x_mean, x_log_var) def get_kl(self, z_mean, z_log_var): return 0.5 * nd.mean(((z_mean ** 2) + nd.exp(z_log_var) - z_log_var - 1)) def get_nll(self, x, x_mean, x_log_var): return 0.5 * nd.mean(((x_log_var + math.log(2 * math.pi)) + ((x - x_mean) ** 2) * nd.exp(-x_log_var))) class VariationalAutoencoder(gluon.HybridBlock): def __init__(self, activation = "relu", hiddens= 400, latents = electrode_size * electrode_size, observables = 784, layers = 1, w = 0.5, **kwargs): # Feel free to use different arguments super(VariationalAutoencoder, self).__init__() with self.name_scope(): self.decoder = Encoder(activation, hiddens, observables, layers, **kwargs) self.encoder = Decoder(activation, hiddens, latents, layers, **kwargs) self.latents = latents self.lossfun = Lossfun(w) def hybrid_forward(self, F, x): z_mean, z_log_var = F.split(self.encoder(x), 2, 1) z = F.exp(0.5 * z_log_var) * F.random_normal(0, 1, (x.shape[0], self.latents), ctx) + z_mean x_mean, x_log_var = nd.split(self.decoder(z), 2, 1) return x, x_mean, x_log_var, z_mean, z_log_var

training_iter = io.NDArrayIter({"data": mnist["train_data"].reshape(-1, 784)}, {"label": mnist["train_label"]}, 30, True) test_iter = io.NDArrayIter({"data": mnist["test_data"].reshape(-1, 784)}, {"label": mnist["test_label"]}, 30) model = VariationalAutoencoder() model.collect_params().initialize(init.Xavier(), ctx) # variationalAutoencoder.hybridize() trainer = gluon.Trainer(model.collect_params(), "adam", {"learning_rate": 0.001})

idx = np.random.choice(mnist['test_data'].shape[0], 10) # Inputs ax = plt.subplots(1, 10, figsize=(16, 4))[1] for i, j in enumerate(idx): ax[i].imshow(mnist["test_data"][j, 0]) ax[i].axis("off")

# Reconstructions ax = plt.subplots(1, 10, figsize=(16, 4))[1] for i, j in enumerate(idx): ax[i].imshow(model(nd.array(mnist["test_data"][j, 0].reshape(1, 784)).as_in_context(ctx))[1].asnumpy().reshape(28, 28) > 0.5) ax[i].axis("off")

# Latents ax = plt.subplots(1, 10, figsize=(16, 4))[1] for i, j in enumerate(idx): ax[i].imshow(model(nd.array(mnist["test_data"][j, 0].reshape(1, 784)).as_in_context(ctx))[3].asnumpy().reshape(electrode_size, electrode_size)) ax[i].axis("off")

# Thresholded latents ax = plt.subplots(1, 10, figsize=(16, 4))[1] for i, j in enumerate(idx): ax[i].imshow(model(nd.array(mnist["test_data"][j, 0].reshape(1, 784)).as_in_context(ctx))[3].asnumpy().reshape(electrode_size, electrode_size) > 0.5) ax[i].axis("off")