q.3.1-all

import os
from PIL import Image
import torch
import torch.nn as nn
import torch.optim as optim
import torch.nn.functional as F
from torchvision import transforms
from torch.utils.data import Dataset, DataLoader
from sklearn.manifold import TSNE
import numpy as np
import matplotlib.pyplot as plt


class CustomBrainMRIDataset(Dataset):
    def __init__(self, directory, transform=None):
        self.directory = directory
        self.image_files = [f for f in os.listdir(directory) if f.endswith('.png')]

        if not self.image_files:
            raise FileNotFoundError(f"Directory {directory} contains no PNG files.")

        print(f"Found {len(self.image_files)} PNG files in {directory}")
        self.transform = transform

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

    def __getitem__(self, index):
        img_path = os.path.join(self.directory, self.image_files[index])
        image = Image.open(img_path).convert('L')

        if self.transform:
            image = self.transform(image)

        return image


transform_pipeline = transforms.Compose([
    transforms.Resize((64, 64)),
    transforms.ToTensor(),
    transforms.Normalize((0.5,), (0.5,))
])

train_seg_dir = 'keras_png_slices_data/keras_png_slices_seg_train'
train_no_seg_dir = 'keras_png_slices_data/keras_png_slices_train'

test_seg_dir = 'keras_png_slices_data/keras_png_slices_seg_test'
test_no_seg_dir = 'keras_png_slices_data/keras_png_slices_test'

validate_seg_dir = 'keras_png_slices_data/keras_png_slices_seg_validate'
validate_no_seg_dir = 'keras_png_slices_data/keras_png_slices_validate'

train_seg_dataset = CustomBrainMRIDataset(train_seg_dir, transform=transform_pipeline)
train_no_seg_dataset = CustomBrainMRIDataset(train_no_seg_dir, transform=transform_pipeline)

test_seg_dataset = CustomBrainMRIDataset(test_seg_dir, transform=transform_pipeline)
test_no_seg_dataset = CustomBrainMRIDataset(test_no_seg_dir, transform=transform_pipeline)

validate_seg_dataset = CustomBrainMRIDataset(validate_seg_dir, transform=transform_pipeline)
validate_no_seg_dataset = CustomBrainMRIDataset(validate_no_seg_dir, transform=transform_pipeline)

train_seg_loader = DataLoader(train_seg_dataset, batch_size=32, shuffle=True)
train_no_seg_loader = DataLoader(train_no_seg_dataset, batch_size=32, shuffle=True)

test_seg_loader = DataLoader(test_seg_dataset, batch_size=32, shuffle=False)
test_no_seg_loader = DataLoader(test_no_seg_dataset, batch_size=32, shuffle=False)

validate_seg_loader = DataLoader(validate_seg_dataset, batch_size=32, shuffle=False)
validate_no_seg_loader = DataLoader(validate_no_seg_dataset, batch_size=32, shuffle=False)


class VAE(nn.Module):
    def __init__(self, latent_dim=128):
        super(VAE, self).__init__()

        self.encoder = nn.Sequential(
            nn.Conv2d(1, 32, 4, stride=2, padding=1),
            nn.ReLU(),
            nn.Conv2d(32, 64, 4, stride=2, padding=1),
            nn.ReLU(),
            nn.Conv2d(64, 128, 4, stride=2, padding=1),
            nn.ReLU()
        )

        self.fc_mu = nn.Linear(128 * 8 * 8, latent_dim)
        self.fc_logvar = nn.Linear(128 * 8 * 8, latent_dim)

        self.fc_decode = nn.Linear(latent_dim, 128 * 8 * 8)
        self.decoder = nn.Sequential(
            nn.ConvTranspose2d(128, 64, 4, stride=2, padding=1),
            nn.ReLU(),
            nn.ConvTranspose2d(64, 32, 4, stride=2, padding=1),
            nn.ReLU(),
            nn.ConvTranspose2d(32, 1, 4, stride=2, padding=1),
            nn.Sigmoid()
        )

    def encode(self, x):
        x = self.encoder(x)
        x = x.view(x.size(0), -1)
        mu = self.fc_mu(x)
        logvar = self.fc_logvar(x)
        return mu, logvar

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

    def decode(self, z):
        x = F.relu(self.fc_decode(z))
        x = x.view(x.size(0), 128, 8, 8)
        x = self.decoder(x)
        return x

    def forward(self, x):
        mu, logvar = self.encode(x)
        z = self.reparameterize(mu, logvar)
        recon_x = self.decode(z)
        return recon_x, mu, logvar


def loss_function(recon_x, x, mu, logvar):
    recon_loss = F.mse_loss(recon_x, x, reduction='sum')
    kl_div = -0.5 * torch.sum(1 + logvar - mu.pow(2) - logvar.exp())
    return recon_loss + kl_div


device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
vae = VAE(latent_dim=128).to(device)
optimizer = optim.Adam(vae.parameters(), lr=0.001)

num_epochs = 10

for epoch in range(num_epochs):
    vae.train()
    total_loss = 0.0

    for data in train_seg_loader:
        data = data.to(device)
        optimizer.zero_grad()
        recon_data, mu, logvar = vae(data)
        loss = loss_function(recon_data, data, mu, logvar)
        loss.backward()
        optimizer.step()
        total_loss += loss.item()

    avg_loss = total_loss / len(train_seg_loader.dataset)
    print(f'Epoch [{epoch + 1}/{num_epochs}], Loss: {avg_loss:.4f}')

vae.eval()
latent_vectors = []
for data in validate_seg_loader:
    data = data.to(device)
    mu, logvar = vae.encode(data)
    latent_vectors.append(mu.cpu().detach().numpy())

latent_vectors = np.vstack(latent_vectors)

tsne = TSNE(n_components=2, random_state=42)
tsne_embedding = tsne.fit_transform(latent_vectors)

plt.scatter(tsne_embedding[:, 0], tsne_embedding[:, 1], alpha=0.5)
plt.title('2D Manifold of Latent Space (t-SNE)')
plt.show()