nmf_using_different_dataset.py

# -*- coding: utf-8 -*-
"""nmf-using-different-dataset.ipynb

Automatically generated by Colab.

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

# This Python 3 environment comes with many helpful analytics libraries installed
# It is defined by the kaggle/python Docker image: https://github.com/kaggle/docker-python
# For example, here's several helpful packages to load

import numpy as np # linear algebra
import pandas as pd # data processing, CSV file I/O (e.g. pd.read_csv)

# Input data files are available in the read-only "../input/" directory
# For example, running this (by clicking run or pressing Shift+Enter) will list all files under the input directory

import os
for dirname, _, filenames in os.walk('/kaggle/input'):
    for filename in filenames:
        print(os.path.join(dirname, filename))

# You can write up to 20GB to the current directory (/kaggle/working/) that gets preserved as output when you create a version using "Save & Run All"
# You can also write temporary files to /kaggle/temp/, but they won't be saved outside of the current session

import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader, Dataset
from torchvision import transforms
from PIL import Image
import matplotlib.pyplot as plt
import os
from torchvision.models import squeezenet1_1
from torchmetrics.functional import structural_similarity_index_measure as ssim

# Define the Convolutional Autoencoder
class ConvAutoencoder(nn.Module):
    def __init__(self):
        super(ConvAutoencoder, self).__init__()

        # Encoder: Convolutional layers with downsampling
        self.encoder = nn.Sequential(
            nn.Conv2d(3, 64, kernel_size=3, stride=2, padding=1),  # 128x128 -> 64x64
            nn.ReLU(),
            nn.Conv2d(64, 128, kernel_size=3, stride=2, padding=1),  # 64x64 -> 32x32
            nn.ReLU(),
            nn.Conv2d(128, 256, kernel_size=3, stride=2, padding=1),  # 32x32 -> 16x16
            nn.ReLU(),
            nn.Conv2d(256, 512, kernel_size=3, stride=2, padding=1),  # 16x16 -> 8x8
            nn.ReLU(),
            nn.Conv2d(512, 1024, kernel_size=3, stride=2, padding=1),  # 10x10 -> 5x5
            nn.ReLU()
        )

        # Decoder: Transposed convolutional layers with upsampling
        self.decoder = nn.Sequential(
          nn.ConvTranspose2d(1024, 512, kernel_size=3, stride=2, padding=1, output_padding=1),  # 5x5 -> 10x10
            nn.ReLU(),
            nn.ConvTranspose2d(512, 256, kernel_size=3, stride=2, padding=1, output_padding=1),  # 10x10 -> 20x20
            nn.ReLU(),
            nn.ConvTranspose2d(256, 128, kernel_size=3, stride=2, padding=1, output_padding=1),  # 20x20 -> 40x40
            nn.ReLU(),
            nn.ConvTranspose2d(128, 64, kernel_size=3, stride=2, padding=1, output_padding=1),  # 40x40 -> 80x80
            nn.ReLU(),
            nn.ConvTranspose2d(64, 3, kernel_size=3, stride=2, padding=1, output_padding=1),  # 80x80 -> 160x160
            nn.Sigmoid()  # Use sigmoid to ensure the output is between 0 and 1
        )

    def forward(self, x):
        encoded = self.encoder(x)
        decoded = self.decoder(encoded)
        return decoded

# Custom dataset to load sharp and blurred images
class DeblurDataset(Dataset):
    def __init__(self, sharp_dir, blurred_dir, transform=None):
        self.sharp_dir = sharp_dir
        self.blurred_dir = blurred_dir
        self.transform = transform
        self.image_names = os.listdir(sharp_dir)

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

    def __getitem__(self, idx):
        img_name = self.image_names[idx]
        sharp_path = os.path.join(self.sharp_dir, img_name)
        blurred_path = os.path.join(self.blurred_dir, img_name)

        sharp_image = Image.open(sharp_path).convert('RGB')
        blurred_image = Image.open(blurred_path).convert('RGB')

        if self.transform:
            sharp_image = self.transform(sharp_image)
            blurred_image = self.transform(blurred_image)

        return blurred_image, sharp_image  # Input: blurred, Target: sharp

# Data transforms
transform = transforms.Compose([
    transforms.Resize((160, 160)),
    transforms.ToTensor()
])
#squeezenet = squeezenet1_1(pretrained=True).features[:8].cuda()

#def perceptual_loss(output, target):
   # output_features = squeezenet(output)
   # target_features = squeezenet(target)
   # return nn.MSELoss()(output_features, target_features)
def ssim_loss(output, target):
    return 1 - ssim(output, target)

# Instantiate the dataset and dataloader
sharp_dir = '/kaggle/input/sharp-data'  # Replace with the path to your sharp images
blurred_dir = '/kaggle/input/blurred'  # Replace with the path to your blurred images
print("Creating Data")
dataset = DeblurDataset(sharp_dir, blurred_dir, transform=transform)
print("Loading Data")
dataloader = DataLoader(dataset, batch_size=64, shuffle=True)
print("Loading Model")
# Model, loss function, and optimizer
model = ConvAutoencoder().cuda()  # Move the model to GPU if available
criterion = nn.MSELoss()  # Mean Squared Error loss for reconstruction
optimizer = optim.Adam(model.parameters(), lr=0.001)

# Training loop
print("Training Started")
num_epochs = 100
for epoch in range(num_epochs):
    running_loss = 0.0
    for data in dataloader:
        blurred_images, sharp_images = data
        blurred_images = blurred_images.cuda()
        sharp_images = sharp_images.cuda()

        # Forward pass
        outputs = model(blurred_images)
        loss_mse = nn.MSELoss()(outputs, sharp_images)
       # loss_perceptual = perceptual_loss(outputs, sharp_images)
        loss_ssim = ssim_loss(outputs, sharp_images)
        loss = loss_mse
        # Backward pass and optimization
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        running_loss += loss.item()

    print(f"Epoch [{epoch+1}/{num_epochs}], Loss: {running_loss/len(dataloader):.4f}")

print("Training completed!")
def display_images(model, dataloader, num_images=10):
    model.eval()  # Set the model to evaluation mode

    fig, axes = plt.subplots(num_images, 3, figsize=(10, num_images * 3))

    with torch.no_grad():  # Disable gradient calculation for inference
        for i, (blurred_images, sharp_images) in enumerate(dataloader):
            if i >= num_images:
                break

            # Move images to GPU if available
            blurred_images = blurred_images.cuda()
            sharp_images = sharp_images.cuda()

            # Forward pass to get reconstructed (deblurred) images
            reconstructed_images = model(blurred_images)

            # Move images back to CPU and detach from computation graph
            blurred_images = blurred_images.cpu().numpy().transpose(0, 2, 3, 1)
            sharp_images = sharp_images.cpu().numpy().transpose(0, 2, 3, 1)
            reconstructed_images = reconstructed_images.cpu().numpy().transpose(0, 2, 3, 1)

            # Display images
            for j in range(len(blurred_images)):
                if i * len(blurred_images) + j >= num_images:
                    break

                idx = i * len(blurred_images) + j

                # Blurred image
                axes[idx, 0].imshow(blurred_images[j])
                axes[idx, 0].set_title("Blurred")
                axes[idx, 0].axis("off")

                # Reconstructed (Deblurred) image
                axes[idx, 1].imshow(reconstructed_images[j])
                axes[idx, 1].set_title("Reconstructed")
                axes[idx, 1].axis("off")

                # Sharp (Original) image
                axes[idx, 2].imshow(sharp_images[j])
                axes[idx, 2].set_title("Sharp")
                axes[idx, 2].axis("off")

    plt.tight_layout()
    plt.show()

# Call the function to display the images
display_images(model, dataloader, num_images=10)


# Example function to save model weights after training
torch.save(model.state_dict(), 'deblur_autoencoder.pth')