densenet_implementation.py

# -*- coding: utf-8 -*-
"""DenseNet Implementation.ipynb

Automatically generated by Colaboratory.

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

import h5py

def explore_hdf5_file(hdf5_path):
    with h5py.File(hdf5_path, 'r') as file:
        print("Contents of the HDF5 file:")
        file.visititems(print_name_and_shape)

def print_name_and_shape(name, node):
    if isinstance(node, h5py.Dataset):
        print(f"Dataset Name: {name}")
        print(f"Dataset Shape: {node.shape}")
        print(f"Dataset Dtype: {node.dtype}")
        print("-" * 40)

# Path to your HDF5 file
hdf5_path = '/content/drive/MyDrive/Course Project - Img and Video/Data/dataset.hdf5'

# Explore the file
explore_hdf5_file(hdf5_path)

import h5py
import torch
from torch.utils.data import Dataset, DataLoader
import numpy as np
class MRIDataset(Dataset):
    def __init__(self, file_path, org_dataset, csm_dataset, mask_dataset):
        self.file = h5py.File(file_path, 'r')
        self.org_data = self.file[org_dataset]
        self.csm_data = self.file[csm_dataset]
        self.mask_data = self.file[mask_dataset]

    def __len__(self):
        return self.org_data.shape[0]

    def __getitem__(self, idx):
        #print("Now loading the Data from the drive")
        org = torch.tensor(self.org_data[idx]).float()
        csm = torch.tensor(self.csm_data[idx]).float()
        mask = torch.tensor(self.mask_data[idx]).float()
        org_fft = torch.fft.fft2(org)
        #print("Creating the Undersampeled Data")
        undersampled = org_fft * mask
        undersampled_ifft = torch.fft.ifft2(undersampled)  # Inverse FFT to get undersampled image in spatial domain
        # Separate real and imaginary parts for model input
        undersampled_real = undersampled.real
        undersampled_imag = undersampled.imag
        return undersampled_real, undersampled_imag, org, csm, mask

# Example usage
train_dataset = MRIDataset('/content/drive/MyDrive/Course Project - Img and Video/Data/dataset.hdf5', 'trnOrg', 'trnCsm', 'trnMask')
train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True)

import matplotlib.pyplot as plt

# Load a single batch or image
train_dataset = MRIDataset('/content/drive/MyDrive/Course Project - Img and Video/Data/dataset.hdf5', 'trnOrg', 'trnCsm', 'trnMask')
train_loader = DataLoader(train_dataset, batch_size=1, shuffle=True)

# Get one sample
undersampled_real, undersampled_imag, org, csm, mask = next(iter(train_loader))

# Convert complex data to absolute values for visualization
org_abs = torch.abs(org)
csm_abs = torch.abs(csm)

# Display the original image
plt.figure(figsize=(8, 8))
plt.subplot(1, 2, 1)
plt.imshow(org_abs[0].numpy(), cmap='gray')
plt.title('Original Image')

# Display each coil image
for i in range(csm_abs.shape[1]):
    plt.subplot(4, 3, i + 1)  # Adjust subplot grid as needed
    plt.imshow(csm_abs[0, i].numpy(), cmap='gray')
    plt.title(f'Coil {i + 1}')

plt.show()

# Display the original image
plt.figure(figsize=(8, 8))
plt.subplot(1, 2, 1)
plt.imshow(org_abs[0].numpy(), cmap='gray')
plt.title('Original Image')

import matplotlib.pyplot as plt

# Get one sample
undersampled_real, undersampled_imag, org, csm, maskk = next(iter(train_loader))

# Assuming the images are complex, display the absolute values
undersampled = torch.fft.ifft2(torch.fft.fft2(org[0])*mask[0])

undersampled_abs = torch.abs(undersampled)
original_abs = torch.abs(org[0])

# Display the images
plt.figure(figsize=(12, 6))
plt.subplot(1, 2, 1)
plt.imshow(undersampled_abs.numpy(), cmap='gray')
plt.title('Undersampled Image')

plt.subplot(1, 2, 2)
plt.imshow(original_abs.numpy(), cmap='gray')
plt.title('Original Image')
plt.show()

np.shape(undersampled.numpy())
np.shape(undersampled_abs)

np.shape(org)

class DenseNetMRI(nn.Module):
    def __init__(self, input_channels=2, output_channels=1):
        super(DenseNetMRI, self).__init__()

        # Choose deeper DenseNet architecture
        self.features = models.densenet169(pretrained=True).features

        # Replace first convolutional layer for desired input channels
        self.features.conv0 = nn.Conv2d(input_channels, 64, kernel_size=(7, 7), stride=(2, 2), padding=(3, 3), bias=False)

        # Residual connections in additional layers
        self.additional_layers = nn.Sequential(
            nn.Conv2d(1024, 512, kernel_size=3, padding=1),
            nn.ReLU(),
            nn.Upsample(size=(128, 116)),
            nn.Conv2d(512, 256, kernel_size=3, padding=1),
            nn.ReLU(),
            nn.Conv2d(256, output_channels, kernel_size=1),
            nn.Upsample(size=(256, 232))
        )

        # Add attention mechanism (e.g., channel attention)
        self.channel_attention = nn.Sequential(
            nn.AdaptiveAvgPool2d((1, 1)),
            nn.Conv2d(1024, 512, kernel_size=1),
            nn.ReLU(),
            nn.Conv2d(512, output_channels, kernel_size=1),
            nn.Sigmoid()
        )

    def forward(self, x):
        features = self.features(x)
        x = self.additional_layers(features)
        # Apply channel attention (weighted sum)
        x = x * self.channel_attention(features)
        return x

import torch.nn as nn
import torchvision.models as models
class DenseNetMRI(nn.Module):
    def __init__(self):
        super(DenseNetMRI, self).__init__()
        # Initialize a pre-trained DenseNet
        densenet = models.densenet121(pretrained=True)

        # Replace the first convolutional layer to accept 2-channel input
        # Original first conv layer: nn.Conv2d(3, 64, kernel_size=(7, 7), stride=(2, 2), padding=(3, 3), bias=False)
        densenet.features.conv0 = nn.Conv2d(2, 64, kernel_size=(7, 7), stride=(2, 2), padding=(3, 3), bias=False)

        # Use the features from DenseNet and add additional layers
        self.features = densenet.features
        self.additional_layers = nn.Sequential(
            nn.Conv2d(1024, 512, kernel_size=3, padding=1),
            nn.ReLU(),
            nn.Upsample(size=(128, 116)),  # Adjust size as needed
            nn.Conv2d(512, 256, kernel_size=3, padding=1),
            nn.ReLU(),
            nn.Upsample(size=(256, 232))  # Adjust size as needed
        )

    def forward(self, x):
        x = self.features(x)
        x = self.additional_layers(x)
        return x

# Check if GPU is available and set the device accordingly
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
print("Using device:", device)

# Move the model to the specified device
model = DenseNetMRI().to(device)

# You'll need to specify the number of input and output channels
# For MRI images, this might be 1 (for grayscale) or 2 (for complex-valued)
# model = DenseNetMRI(input_channels=1, output_channels=1)
criterion = nn.MSELoss()
optimizer = torch.optim.Adam(model.parameters(), lr=0.01)

def calculate_psnr(output, target):
    mse = torch.mean((output - target) ** 2)
    if mse == 0:
        return float('inf')
    max_pixel = 1.0  # Assuming your images are scaled between 0 and 1
    psnr = 20 * torch.log10(max_pixel / torch.sqrt(mse))
    return psnr

num_epochs = 10
for epoch in range(num_epochs):
    total_loss = 0.0
    total_psnr = 0.0
    num_batches = 0

    for undersampled_real, undersampled_imag, org, csm , mask in train_loader:
        # Concatenate real and imaginary parts along channel dimension
        undersampled_combined = torch.cat([undersampled_real.unsqueeze(1).to(device), undersampled_imag.unsqueeze(1).to(device)], dim=1)

        # Forward pass
        outputs = model(undersampled_combined)

        # Ensure 'org' has a channel dimension
        org = org.unsqueeze(1).to(device)  # Adds a channel dimension

        # Compute loss
        loss = criterion(outputs, org)
        total_loss += loss.item()

        # Compute PSNR
        psnr = calculate_psnr(outputs, org)
        total_psnr += psnr.item()

        # Backward pass and optimize
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        num_batches += 1

    # Calculate average loss and PSNR for the epoch
    avg_loss = total_loss / num_batches
    avg_psnr = total_psnr / num_batches

    print(f'Epoch [{epoch + 1}/{num_epochs}], Loss: {avg_loss:.4f}, Average PSNR: {avg_psnr:.2f} dB')

torch.save(model.state_dict(), 'densenet_model.pth')
print("Model saved as densenet_model.pth")

model = DenseNetMRI()

model.load_state_dict(torch.load('densenet_model.pth'))

# Move model to the correct device (GPU or CPU)
model = model.to(device)

test_dataset = MRIDataset('/content/drive/MyDrive/Course Project - Img and Video/Data/dataset.hdf5', 'tstOrg', 'tstCsm', 'tstMask')
test_loader = DataLoader(test_dataset, batch_size=2, shuffle=False)  # No need to shuffle for testing

def calculate_psnr(output, target):
    # Check for channel mismatch and handle accordingly
    if output.dim() == 4 and target.dim() == 3:
        # If the output has an extra dimension and it's not singleton
        if output.size(1) != 1:
            # Average across the channels or select a specific channel
            # Here, we choose to average
            output = output.mean(1)
        else:
            # If it's a singleton channel, squeeze it out
            output = output.squeeze(1)

    # Check for spatial dimension alignment
    if output.shape[1:] != target.shape[1:]:
        # Resize output to match target spatial dimensions
        output = torch.nn.functional.interpolate(output, size=(target.size(1), target.size(2)))

    mse = torch.mean((output - target) ** 2)
    if mse == 0:
        psnr = torch.tensor(float('inf'))
    else:
        max_pixel = 1.0
        psnr = 20 * torch.log10(max_pixel / torch.sqrt(mse))
    return psnr




def evaluate_model(model, test_loader, device):
    model.eval()  # Set the model to evaluation mode
    total_psnr_original = 0.0
    total_psnr_reconstructed = 0.0
    num_samples = 0
    saved_originals = []
    saved_reconstructed = []

    with torch.no_grad():  # No need to track gradients
        for undersampled_real, undersampled_imag, org, csm, mask in test_loader:
            org = org.to(device)

            # Correctly combine real and imaginary parts
            # Ensure the concatenated tensor has shape [batch_size, 2, height, width]
            undersampled_combined = torch.cat([undersampled_real.unsqueeze(1).to(device), undersampled_imag.unsqueeze(1).to(device)], dim=1)

            # Forward pass through the model
            reconstructed = model(undersampled_combined)
            print(f"reconstructed shape:{reconstructed.shape}")
            print(f"Org Shape:{org.shape}")
            # Calculate PSNR for original and reconstructed images
            psnr_original = calculate_psnr(org, org)  # PSNR with itself will be max
            psnr_reconstructed = calculate_psnr(reconstructed, org)
            saved_originals.append(org.to(device))
            reconstructed_fft = torch.fft.ifft2(reconstructed)
            saved_reconstructed.append(reconstructed.to(device))

            total_psnr_original += psnr_original.item()
            total_psnr_reconstructed += psnr_reconstructed.item()
            num_samples += 1

    avg_psnr_original = total_psnr_original / num_samples
    avg_psnr_reconstructed = total_psnr_reconstructed / num_samples

    return saved_originals, saved_reconstructed,avg_psnr_original, avg_psnr_reconstructed

saved_originals, saved_reconstructed,avg_psnr_original, avg_psnr_reconstructed = evaluate_model(model, test_loader, device)
print(f'Average PSNR for Original Images: {avg_psnr_original:.2f} dB')
print(f'Average PSNR for Reconstructed Images: {avg_psnr_reconstructed:.2f} dB')

import matplotlib.pyplot as plt

def plot_saved_images(saved_originals, saved_reconstructed, num_samples=2):
    fig, axs = plt.subplots(num_samples, 2, figsize=(10, 4 * num_samples))

    for i in range(num_samples):
        # Original image - move to CPU, convert to NumPy array, and select first slice
        original_img = saved_originals[i][0].squeeze().cpu().numpy()
        # Assuming the image is 3D and you want to plot the first 2D slice
        original_img_slice = original_img[0, :, :] if original_img.ndim == 3 else original_img

        axs[i, 0].imshow(original_img_slice, cmap='gray')
        axs[i, 0].set_title(f"Original Image {i+1}")
        axs[i, 0].axis('off')

        # Reconstructed image - move to CPU, convert to NumPy array, and select first slice
        reconstructed_img = saved_reconstructed[i][0].squeeze().cpu().numpy()
        # Assuming the image is 3D and you want to plot the first 2D slice
        reconstructed_img_slice = reconstructed_img[0, :, :] if reconstructed_img.ndim == 3 else reconstructed_img

        axs[i, 1].imshow(reconstructed_img_slice, cmap='gray')
        axs[i, 1].set_title(f"Reconstructed Image {i+1}")
        axs[i, 1].axis('off')

    plt.tight_layout()
    plt.show()

# Plot the images
plot_saved_images(saved_originals, saved_reconstructed, num_samples=2)