| import os |
| import glob |
| import torch |
| import torch.nn.functional as F |
| import math |
| from PIL import Image |
| from torchvision import transforms |
| from torch.utils.data import Dataset, DataLoader |
|
|
| |
| DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu") |
| _SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__)) |
| _SIGMAS = (15.0 / 255.0, 50.0 / 255.0, 75.0 / 255.0) |
|
|
|
|
| def _testset_root(name): |
| return os.path.join( |
| _SCRIPT_DIR, |
| "datasets", |
| "Test_Datasets", |
| "FFDNet-master", |
| "testsets", |
| name, |
| ) |
|
|
|
|
| class TestDataset(Dataset): |
| def __init__(self, root_dir, sigma): |
| self.sigma = sigma |
| self.image_paths = glob.glob(os.path.join(root_dir, "*.png")) + glob.glob( |
| os.path.join(root_dir, "*.jpg") |
| ) |
| self.transform = transforms.Compose( |
| [transforms.Grayscale(num_output_channels=1), transforms.ToTensor()] |
| ) |
|
|
| def __len__(self): |
| return len(self.image_paths) |
|
|
| def __getitem__(self, idx): |
| img = Image.open(self.image_paths[idx]) |
| clean = self.transform(img) |
| noisy = clean + torch.randn_like(clean) * self.sigma |
| return clean, torch.clamp(noisy, 0.0, 1.0) |
|
|
|
|
| def calculate_psnr(img1, img2): |
| mse = torch.mean((img1 - img2) ** 2) |
| if mse == 0: |
| return float("inf") |
| return 20 * math.log10(1.0 / math.sqrt(mse)) |
|
|
|
|
| def classical_telegraph_step(u_n, u_n_minus_1, tau=0.2, gamma=1.0): |
| kx = torch.tensor([[0, 0, 0], [-0.5, 0, 0.5], [0, 0, 0]], device=DEVICE).view( |
| 1, 1, 3, 3 |
| ) |
| ky = torch.tensor([[0, -0.5, 0], [0, 0, 0], [0, 0.5, 0]], device=DEVICE).view( |
| 1, 1, 3, 3 |
| ) |
|
|
| grad_x = F.conv2d(u_n, kx, padding=1) |
| grad_y = F.conv2d(u_n, ky, padding=1) |
| grad_mag = torch.sqrt(grad_x**2 + grad_y**2 + 1e-8) |
|
|
| c = 1.0 / (1.0 + (grad_mag / 0.1) ** 2) |
|
|
| divergence = F.conv2d(c * grad_x, kx, padding=1) + F.conv2d( |
| c * grad_y, ky, padding=1 |
| ) |
|
|
| alpha = (2 + gamma * tau) / (1 + gamma * tau) |
| beta = -1 / (1 + gamma * tau) |
| lam = (tau**2) / (1 + gamma * tau) |
|
|
| return alpha * u_n + beta * u_n_minus_1 + lam * divergence |
|
|
|
|
| def run_eval(dataset_name, sigma): |
| root = _testset_root(dataset_name) |
| dataset = TestDataset(root, sigma) |
| if len(dataset) == 0: |
| print(f"[!] Skip {dataset_name}: no images in {os.path.abspath(root)}") |
| return None |
|
|
| dataloader = DataLoader(dataset, batch_size=1, shuffle=False) |
| total_psnr = 0.0 |
| for clean, noisy in dataloader: |
| clean, noisy = clean.to(DEVICE), noisy.to(DEVICE) |
| u_n_minus_1, u_n = noisy.clone(), noisy.clone() |
|
|
| for _ in range(20): |
| u_next = classical_telegraph_step(u_n, u_n_minus_1) |
| u_n_minus_1, u_n = u_n, u_next |
|
|
| total_psnr += calculate_psnr(clean, u_n) |
|
|
| avg_psnr = total_psnr / len(dataset) |
| sigma_int = int(round(sigma * 255.0)) |
| print( |
| f"[+] {dataset_name} sigma={sigma_int}/255 PSNR: {avg_psnr:.2f} dB " |
| f"({len(dataset)} images)" |
| ) |
| return avg_psnr |
|
|
|
|
| def main(): |
| print("[*] Classical Majee 2020 baseline — Set12 & BSD68") |
| for dataset_name in ("Set12", "BSD68"): |
| for sigma in _SIGMAS: |
| run_eval(dataset_name, sigma) |
| print() |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
