| | import csv |
| | import torch |
| | import numpy as np |
| | from skimage.metrics import structural_similarity as compare_ssim |
| | import math |
| |
|
| | class ImageQualityMetrics: |
| | """Class to compute image quality metrics like SSIM, PSNR, and MSE.""" |
| |
|
| | @staticmethod |
| | @staticmethod |
| | def ssim_3d(img1, img2): |
| | """Calculate SSIM for each 2D slice in the 3D image and return the average.""" |
| | ssim_vals = [] |
| | for i in range(img1.shape[1]): |
| | slice1 = img1[0, i, :, :] |
| | slice2 = img2[0, i, :, :] |
| | ssim_val = compare_ssim(slice1, slice2, data_range=slice1.max() - slice1.min()) |
| | ssim_vals.append(ssim_val) |
| | return np.mean(ssim_vals) |
| | |
| | @staticmethod |
| | def psnr(img1, img2): |
| | """Calculate PSNR (Peak Signal-to-Noise Ratio) between two images.""" |
| | mse = torch.mean((img1 - img2) ** 2) |
| | if mse == 0: |
| | return math.inf |
| | return 20 * math.log10(img1.max() - img1.min()) - 10 * math.log10(mse) |
| |
|
| | @staticmethod |
| | def mse(img1, img2): |
| | """Calculate MSE (Mean Squared Error) between two images.""" |
| | return torch.mean((img1 - img2) ** 2) |
| |
|
| |
|
| | class ValidationRecorder: |
| | """Class to handle validation process and record the metrics.""" |
| |
|
| | def __init__(self, csv_file_path): |
| | """Initialize the recorder with the path to the CSV file.""" |
| | self.csv_file_path = csv_file_path |
| |
|
| | def initialize_csv(self): |
| | """Initialize the CSV file with headers.""" |
| | with open(self.csv_file_path, mode='w', newline='') as file: |
| | writer = csv.writer(file) |
| | writer.writerow(['Epoch', 'Loss', 'SSIM', 'PSNR', 'MSE']) |
| |
|
| | def validate_and_record(self, epoch, dataloader, device, generator, |
| | criterion_g): |
| | """Validate the model and record the metrics in the CSV file.""" |
| | generator.eval() |
| | total_loss, total_ssim, total_psnr, total_mse = 0.0, 0.0, 0.0, 0.0 |
| |
|
| | with torch.no_grad(): |
| | for _, (low_res, high_res) in enumerate(dataloader): |
| | low_res, high_res = low_res.to(device), high_res.to(device) |
| | fake_images = generator(low_res) |
| |
|
| | loss = criterion_g(fake_images, high_res) |
| | total_loss += loss.item() |
| |
|
| | for j in range(high_res.size(0)): |
| | ssim_val = ImageQualityMetrics.ssim_3d( |
| | high_res[j].cpu().numpy(), fake_images[j].cpu().numpy()) |
| | psnr_val = ImageQualityMetrics.psnr( |
| | high_res[j], fake_images[j]) |
| | mse_val = ImageQualityMetrics.mse( |
| | high_res[j], fake_images[j]) |
| | total_ssim += ssim_val |
| | total_psnr += psnr_val |
| | total_mse += mse_val.item() |
| |
|
| | avg_loss = total_loss / len(dataloader) |
| | avg_ssim = total_ssim / (len(dataloader) * dataloader.batch_size) |
| | avg_psnr = total_psnr / (len(dataloader) * dataloader.batch_size) |
| | avg_mse = total_mse / (len(dataloader) * dataloader.batch_size) |
| |
|
| | self._write_to_csv(epoch, avg_loss, avg_ssim, avg_psnr, avg_mse) |
| |
|
| | def _write_to_csv(self, epoch, avg_loss, avg_ssim, avg_psnr, avg_mse): |
| | """Write the validation metrics to the CSV file.""" |
| | with open(self.csv_file_path, mode='a', newline='') as file: |
| | writer = csv.writer(file) |
| | writer.writerow([epoch, avg_loss, avg_ssim, avg_psnr, avg_mse]) |
| |
|
