better_vae copy.py

# -*- coding: utf-8 -*-
"""better_vae

Automatically generated by Colaboratory.

Original file is located at
    https://colab.research.google.com/drive/1gArgTJw2lZpmCH0CInG8m9FJuw5nhLOs
"""

import numpy as np
import pandas as pd
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import Dataset, DataLoader

# Load your data
frames_file_path = "data/mTBI/frame_pixels.npy"  # Add this line
df = pd.read_csv("data/mTBI/eye_motion_trace.csv")
print("loaded data")
# Dataset class
class LSTMDataset(Dataset):
    def __init__(self, dataframe, frames_file_path):
        self.dataframe = dataframe
        self.frames_file_path = frames_file_path  # This should be a string path

    def __len__(self):
        return len(self.dataframe)

    def __getitem__(self, idx):
        frame_data = self.load_frame_data(idx)
        x = self.dataframe.iloc[idx]['x']
        y = self.dataframe.iloc[idx]['y']
        inputs = torch.tensor([x, y], dtype=torch.float32)

        return inputs, torch.tensor(frame_data, dtype=torch.float32)

    def load_frame_data(self, idx):
        data = np.load(self.frames_file_path, mmap_mode='r')
        frame_data = data[idx]
        return frame_data

import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader, TensorDataset


# frames is a tensor of shape [num_frames, height, width] or [num_frames, channels, height, width]

# Hyperparameters
input_channels = frames.size(1) if len(frames.shape) == 4 else 1  # Assuming grayscale or RGB images
image_size = frames.size(2) if len(frames.shape) == 4 else frames.size(1)  # Assuming square images
latent_dim = 32
batch_size = 64
epochs = 20

# Flatten each frame and concatenate with eye movement shift data
frames = frames.view(frames.size(0), -1) if len(frames.shape) == 4 else frames.view(frames.size(0), -1, 1)
eye_movement_shifts_and_frames = torch.cat([eye_movement_shifts, frames], dim=1)

# Convert data to PyTorch tensors
eye_movement_shifts_and_frames = torch.FloatTensor(eye_movement_shifts_and_frames)

# Create a DataLoader
dataset = TensorDataset(eye_movement_shifts_and_frames, frames)
dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=True)

# VAE Model
class VAE(nn.Module):
    def __init__(self, input_size, latent_dim):
        super(VAE, self).__init__()

        # Encoder
        self.encoder = nn.Sequential(
            nn.Linear(input_size, 128),
            nn.ReLU(),
            nn.Linear(128, 64),
            nn.ReLU(),
        )
        self.z_mean = nn.Linear(64, latent_dim)
        self.z_log_var = nn.Linear(64, latent_dim)

        # Decoder
        self.decoder = nn.Sequential(
            nn.Linear(latent_dim, 64),
            nn.ReLU(),
            nn.Linear(64, 128),
            nn.ReLU(),
            nn.Linear(128, input_size),
            nn.Sigmoid()  # Assuming you want values between 0 and 1
        )

    def reparameterize(self, mu, log_var):
        std = torch.exp(0.5 * log_var)
        eps = torch.randn_like(std)
        return mu + eps * std

    def forward(self, x):
        # Encoder
        x = self.encoder(x)
        z_mean = self.z_mean(x)
        z_log_var = self.z_log_var(x)

        # Reparameterization trick
        z = self.reparameterize(z_mean, z_log_var)

        # Decoder
        x_hat = self.decoder(z)
        return x_hat, z_mean, z_log_var

# Instantiate the VAE model
input_size = eye_movement_shifts_and_frames.size(1)
model = VAE(input_size=input_size, latent_dim=latent_dim)

# Loss function
def loss_function(x_hat, x, z_mean, z_log_var):
    reconstruction_loss = nn.functional.mse_loss(x_hat, x, reduction='sum')
    kl_divergence = -0.5 * torch.sum(1 + z_log_var - z_mean.pow(2) - z_log_var.exp())
    return reconstruction_loss + kl_divergence

# Optimizer
optimizer = optim.Adam(model.parameters(), lr=1e-3)

# Training
for epoch in range(epochs):
    for batch in dataloader:
        x, _ = batch
        optimizer.zero_grad()
        x_hat, z_mean, z_log_var = model(x)
        loss = loss_function(x_hat, x, z_mean, z_log_var)
        loss.backward()
        optimizer.step()

    print(f'Epoch {epoch + 1}/{epochs}, Loss: {loss.item()}')

# Generate new videos
with torch.no_grad():
    # Assuming you have new eye movement shift data in new_eye_movement_shifts
    new_eye_movement_shifts = torch.FloatTensor(new_eye_movement_shifts)
    new_frames, _, _ = model(new_eye_movement_shifts)
    new_frames = new_frames.numpy()

# You can now use new_frames for further analysis or visualization

with torch.no_grad():
    # Assuming new_eye_movement_shifts_tensor is the tensor containing the new eye movement shifts
    new_eye_movement_shifts_tensor = torch.FloatTensor(new_eye_movement_shifts)
    generated_videos, _, _ = model(new_eye_movement_shifts_tensor)
    generated_videos = generated_videos.numpy()


import matplotlib.pyplot as plt
import numpy as np

# Assuming generated_videos is a NumPy array containing the generated videos
# Reshape the videos if necessary
generated_videos = generated_videos.reshape((-1, input_channels, image_size, image_size))

# Visualize a few frames from the generated video
num_frames_to_visualize = 5
for i in range(num_frames_to_visualize):
    plt.subplot(1, num_frames_to_visualize, i + 1)
    plt.imshow(np.transpose(generated_videos[i], (1, 2, 0)))
    plt.axis('off')

plt.show()