Pythonで学ぶ多段階ノイズ除去の実装

import torch
import torch.nn as nn
import torchvision.transforms as transforms
import torchvision.datasets as datasets
import matplotlib.pyplot as plt

# デバイスの設定
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# シンプルなCNNモデル定義
class SimpleCNN(nn.Module):
    def __init__(self):
        super(SimpleCNN, self).__init__()
        self.conv1 = nn.Conv2d(1, 16, kernel_size=3, padding=1)
        self.conv2 = nn.Conv2d(16, 32, kernel_size=3, padding=1)
        self.fc1 = nn.Linear(32 * 28 * 28, 128)
        self.fc2 = nn.Linear(128, 28 * 28)
    
    def forward(self, x):
        x = torch.relu(self.conv1(x))
        x = torch.relu(self.conv2(x))
        x = x.view(x.size(0), -1)
        x = torch.relu(self.fc1(x))
        x = self.fc2(x)
        x = x.view(-1, 1, 28, 28)
        return x

# データセットの準備
transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.5,), (0.5,))
])

mnist_train = datasets.MNIST(root='./data', train=True, transform=transform, download=True)

# ガウシアンノイズの追加
def add_gaussian_noise(img, noise_level):
    noise = torch.randn(*img.shape).to(device) * noise_level
    noisy_img = img + noise
    noisy_img = torch.clamp(noisy_img, -1.0, 1.0)
    return noisy_img

# モデルの初期化
model = SimpleCNN().to(device)
optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
criterion = nn.MSELoss()

# トレーニング
num_epochs = 5
initial_noise_level = 0.5
steps = 5  # ノイズ除去のステップ数

for epoch in range(num_epochs):
    for images, _ in train_loader:
        images = images.to(device)
        
        # 初期状態にノイズを追加
        noisy_images = add_gaussian_noise(images, initial_noise_level)
        
        # 多段階でノイズを除去
        for step in range(steps):
            current_noise_level = initial_noise_level * (1 - step / steps)
            noisy_images = add_gaussian_noise(noisy_images, current_noise_level)
            optimizer.zero_grad()
            outputs = model(noisy_images)
            loss = criterion(outputs, images)
            loss.backward()
            optimizer.step()
    
    print(f'Epoch [{epoch + 1}/{num_epochs}], Loss: {loss.item():.4f}')

# 結果の表示
def show_images(orig, noisy, denoised):
    orig = orig.cpu().numpy()
    noisy = noisy.cpu().numpy()
    denoised = denoised.cpu().detach().numpy()
    fig, axes = plt.subplots(1, 3, figsize=(12, 4))
    axes[0].imshow((orig[0][0] * 0.5 + 0.5), cmap='gray')
    axes[0].set_title("Original Image")
    axes[1].imshow((noisy[0][0] * 0.5 + 0.5), cmap='gray')
    axes[1].set_title("Noisy Image")
    axes[2].imshow((denoised[0][0] * 0.5 + 0.5), cmap='gray')
    axes[2].set_title("Denoised Image")
    plt.show()

# サンプル画像の取得と表示
sample_image, _ = mnist_train[0]
sample_image = sample_image.unsqueeze(0).to(device)

noisy_sample = sample_image
for step in range(steps):
    current_noise_level = initial_noise_level * (1 - step / steps)
    noisy_sample = add_gaussian_noise(noisy_sample, current_noise_level)

denoised_sample = model(noisy_sample)

show_images(sample_image, noisy_sample, denoised_sample)