# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python 3
#     language: python
#     name: python3
# ---

from importlib import reload
import make_models;reload(make_models)
del make_models;from make_models import *

encoder, decoder, vae = get_MRI_VAE_3D(input_shape=(64,64,64,1),
                   latent_dim=2,
                   batch_size = 32,
                   disentangle=False,
                   gamma=1,
                   kernel_size = 3,
                   filters = 16,
                   intermediate_dim = 128,
                   opt=None)

import numpy as np
data_batch = np.random.rand(32,64,64,64)

batch_loss = vae.train_on_batch(data_batch)

# +
from matplotlib import pyplot as plt
from tqdm import tqdm

loss = list()
#for i in range(25):
for i in tqdm(range(25)):
    
    batch_loss = vae.train_on_batch(data_batch)
    loss.append(batch_loss)
    
plt.plot(loss)
# -

predictions= vae.predict(data_batch)
predictions.shape

# +
plt.figure(figsize=(10,5))

plt.subplot(1,2,1);
plt.imshow(data_batch[0,:,:,32])
plt.title('input')

plt.subplot(1,2,2);
plt.imshow(predictions[0,:,:,32,0])
plt.title('prediction/reconstruction')


# -

class VAE_DataLoader:
    """A simple example dataloader"""
    
    def __init__(self):
        self.batch_id = 0 # count batches

    def get_batch(self):
        self.data_batch = np.random.rand(32,64,64,64)  
        self.batch_id += 1 # count batches 
        return self.data_batch



data_loader = VAE_DataLoader()

print(data_loader.batch_id)

# +
for i in range(10):
    a = data_loader.get_batch()
    
print(data_loader.batch_id)

# +
latent_representation = encoder.predict(data_batch)
print(len(latent_representation))

prediction = decoder.predict(latent_representation[2])
print(prediction.shape)
# -



StarcoderdataJupyterScriptsDedupFiltered