denoising_model.py

import os
import torch
import torch.nn as nn
from PIL import Image
from torch.utils.data import Dataset, DataLoader
from torchvision.transforms import ToTensor
import time
from datetime import datetime
import multiprocessing

def get_optimal_threads():
    """Calculate optimal number of threads based on CPU cores"""
    return max(1, multiprocessing.cpu_count() - 1)  # Leave one core free for system

# Simple UNet-style denoising model
class DenoisingModel(nn.Module):
    def __init__(self):
        super(DenoisingModel, self).__init__()
        # Encoder
        self.enc1 = nn.Sequential(
            nn.Conv2d(3, 64, 3, padding=1),
            nn.ReLU(),
            nn.Conv2d(64, 64, 3, padding=1),
            nn.ReLU()
        )
        self.pool1 = nn.MaxPool2d(2, 2)

        # Decoder
        self.up1 = nn.ConvTranspose2d(64, 64, 2, stride=2)
        self.dec1 = nn.Sequential(
            nn.Conv2d(64, 64, 3, padding=1),
            nn.ReLU(),
            nn.Conv2d(64, 3, 3, padding=1)
        )

    def forward(self, x):
        # Encoder
        e1 = self.enc1(x)
        p1 = self.pool1(e1)

        # Decoder
        u1 = self.up1(p1)
        d1 = self.dec1(u1)
        return d1

class DenoiseDataset(Dataset):
    def __init__(self, noisy_folder, target_folder, patch_size=256):
        self.noisy_folder = noisy_folder
        self.target_folder = target_folder
        self.patch_size = patch_size
        self.image_pairs = [
            (os.path.join(noisy_folder, f), os.path.join(target_folder, f.replace("_noisy", "_target")))
            for f in os.listdir(noisy_folder) if "_noisy" in f
        ]
        self.transform = ToTensor()

        print(f"Dataset initialization:")
        print(f"- Noisy folder: {noisy_folder}")
        print(f"- Target folder: {target_folder}")
        print(f"- Patch size: {patch_size}")
        print(f"- Found {len(self.image_pairs)} image pairs")

        if not self.image_pairs:
            raise ValueError("No image pairs found. Check if noisy and target images are correctly named.")

        # Precalculate number of patches per image for better performance
        self.patches_per_image = {}
        for noisy_path, _ in self.image_pairs:
            try:
                self.patches_per_image[noisy_path] = self._get_num_patches_per_image(noisy_path)
            except Exception as e:
                print(f"Error calculating patches for {noisy_path}: {e}. Skipping this image pair.")
                self.image_pairs = [(n, t) for n, t in self.image_pairs if n != noisy_path]

        self.total_patches = sum(self.patches_per_image.values())

    def __len__(self):
        return self.total_patches

    def __getitem__(self, idx):
        image_idx = 0
        cumulative_patches = 0

        for i, (noisy_path, _) in enumerate(self.image_pairs):
            num_patches = self.patches_per_image[noisy_path]
            if cumulative_patches + num_patches > idx:
                image_idx = i
                break
            cumulative_patches += num_patches

        patch_idx = idx - cumulative_patches
        noisy_path, target_path = self.image_pairs[image_idx]

        try:
            noisy_image = self._load_image(noisy_path)
            target_image = self._load_image(target_path)
        except Exception as e:
            print(f"Error loading image pair ({noisy_path}, {target_path}): {e}. Returning default values.")
            return torch.zeros((3, self.patch_size, self.patch_size)), torch.zeros((3, self.patch_size, self.patch_size))

        try:
            noisy_patch = self._get_patch(noisy_image, patch_idx)
            target_patch = self._get_patch(target_image, patch_idx)
        except Exception as e:
            print(f"Error getting patch from image pair ({noisy_path}, {target_path}): {e}. Returning default values.")
            return torch.zeros((3, self.patch_size, self.patch_size)), torch.zeros((3, self.patch_size, self.patch_size))

        return noisy_patch, target_patch

    def _load_image(self, image_path):
        try:
            image = Image.open(image_path).convert("RGB")
            return self.transform(image)
        except Exception as e:
            raise Exception(f"Error loading image {image_path}: {e}")

    def _get_num_patches_per_image(self, image_path):
        try:
            image = Image.open(image_path)
            width, height = image.size
            num_patches = (width // self.patch_size) * (height // self.patch_size)
            return num_patches
        except Exception as e:
            raise Exception(f"Error calculating patches for {image_path}: {e}")

    def _get_patch(self, image, patch_idx):
        width, height = image.shape[2], image.shape[1]
        patches_per_row = width // self.patch_size
        row = patch_idx // patches_per_row
        col = patch_idx % patches_per_row

        x_start = col * self.patch_size
        y_start = row * self.patch_size
        return image[:, y_start:y_start+self.patch_size, x_start:x_start+self.patch_size]


def train_model(noisy_dir, target_dir, epochs, batch_size, learning_rate, save_interval, num_workers):
    # Set up CUDA if available
    if torch.cuda.is_available():
        torch.backends.cudnn.benchmark = True  # Enable cuDNN auto-tuner
        device = torch.device("cuda")
        print(f"\nUsing GPU: {torch.cuda.get_device_name(0)}")
        print(f"CUDA version: {torch.version.cuda}")
    else:
        device = torch.device("cpu")
        print("\nNo GPU detected, using CPU")

    # Create output directory for models
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    output_dir = f"model_checkpoints_{timestamp}"
    os.makedirs(output_dir, exist_ok=True)

    print("\nTraining Configuration:")
    print(f"- Number of epochs: {epochs}")
    print(f"- Batch size: {batch_size}")
    print(f"- Learning rate: {learning_rate}")
    print(f"- Number of worker threads: {num_workers}")
    print(f"- Model checkpoint directory: {output_dir}")

    # Initialize dataset and dataloader with specified number of workers
    dataset = DenoiseDataset(noisy_dir, target_dir)
    dataloader = DataLoader(
        dataset,
        batch_size=batch_size,
        shuffle=True,
        num_workers=num_workers,
        pin_memory=True if torch.cuda.is_available() else False
    )

    # Initialize model, loss function, and optimizer
    model = DenoisingModel().to(device)
    if torch.cuda.device_count() > 1:
        print(f"Using {torch.cuda.device_count()} GPUs!")
        model = nn.DataParallel(model)

    criterion = nn.MSELoss()
    optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)

    # Training loop
    total_batches = len(dataloader)
    start_time = time.time()

    print("\nStarting training...")
    for epoch in range(epochs):
        epoch_loss = 0.0
        for batch_idx, (noisy_patches, target_patches) in enumerate(dataloader):
            # Move data to device
            noisy_patches = noisy_patches.to(device, non_blocking=True)
            target_patches = target_patches.to(device, non_blocking=True)

            # Forward pass
            outputs = model(noisy_patches)
            loss = criterion(outputs, target_patches)

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

            # Update epoch loss
            epoch_loss += loss.item()

            # Print progress
            if (batch_idx + 1) % 100 == 0:
                elapsed_time = time.time() - start_time
                print(f"Epoch [{epoch+1}/{epochs}], "
                      f"Batch [{batch_idx+1}/{total_batches}], "
                      f"Loss: {loss.item():.6f}, "
                      f"Time: {elapsed_time:.2f}s")

            # Save model checkpoint
            if (batch_idx + 1) % save_interval == 0:
                checkpoint_path = os.path.join(output_dir,
                    f"denoising_model_epoch{epoch+1}_batch{batch_idx+1}.pth")
                torch.save({
                    'epoch': epoch,
                    'batch': batch_idx,
                    'model_state_dict': model.state_dict(),
                    'optimizer_state_dict': optimizer.state_dict(),
                    'loss': loss.item(),
                }, checkpoint_path)
                print(f"\nCheckpoint saved: {checkpoint_path}")

        # End of epoch summary
        avg_epoch_loss = epoch_loss / total_batches
        print(f"\nEpoch [{epoch+1}/{epochs}] completed. "
              f"Average loss: {avg_epoch_loss:.6f}")

        # Save epoch checkpoint
        checkpoint_path = os.path.join(output_dir, f"denoising_model_epoch{epoch+1}.pth")
        torch.save({
            'epoch': epoch,
            'model_state_dict': model.state_dict(),
            'optimizer_state_dict': optimizer.state_dict(),
            'loss': avg_epoch_loss,
        }, checkpoint_path)
        print(f"Epoch checkpoint saved: {checkpoint_path}")

    print("\nTraining completed!")
    print(f"Total training time: {time.time() - start_time:.2f} seconds")

    # Save final model
    final_model_path = os.path.join(output_dir, "denoising_model_final.pth")
    torch.save(model.state_dict(), final_model_path)
    print(f"Final model saved: {final_model_path}")

def main():
    noisy_dir = 'noisy_images'
    target_dir = 'target_images'
    epochs = 10
    batch_size = 4
    learning_rate = 0.001
    save_interval = 1000
    num_workers = get_optimal_threads()

    train_model(noisy_dir, target_dir, epochs, batch_size, learning_rate, save_interval, num_workers)

if __name__ == "__main__":
    main()