test_blur.py

import torch
import torch.nn.functional as F
import math

def gaussian_blur_old(img: torch.Tensor, kernel_size: int) -> torch.Tensor:
    """Old incorrect separable convolution"""
    if kernel_size <= 0:
        return img
    
    # Create Gaussian kernel
    sigma = (kernel_size - 1) / 6.0
    x = torch.arange(kernel_size, dtype=torch.float32, device=img.device) - (kernel_size-1)/2
    kernel_1d = torch.exp(-x**2 / (2*sigma**2))
    kernel_1d /= kernel_1d.sum()
    kernel_2d = torch.outer(kernel_1d, kernel_1d).view(1, 1, kernel_size, kernel_size)
    
    # Apply separable convolution for efficiency
    channels = img.shape[-3]
    kernel_2d = kernel_2d.expand(channels, 1, kernel_size, kernel_size)
    padding = kernel_size // 2
    
    return F.conv2d(
        F.conv2d(img, kernel_2d, padding=padding, groups=channels),
        kernel_2d.transpose(2, 3), 
        padding=padding, 
        groups=channels
    )

def gaussian_blur_new(img: torch.Tensor, kernel_size: int) -> torch.Tensor:
    """New correct separable convolution"""
    if kernel_size <= 0:
        return img
    
    # Create Gaussian kernel
    sigma = (kernel_size - 1) / 6.0
    x = torch.arange(kernel_size, dtype=torch.float32, device=img.device) - (kernel_size-1)/2
    kernel_1d = torch.exp(-x**2 / (2*sigma**2))
    kernel_1d /= kernel_1d.sum()
    
    # Apply separable convolution for efficiency
    channels = img.shape[-3]
    padding = kernel_size // 2
    
    # Horizontal pass
    kernel_h = kernel_1d.view(1, 1, 1, kernel_size).expand(channels, 1, 1, kernel_size)
    img = F.conv2d(img, kernel_h, padding=(0, padding), groups=channels)
    
    # Vertical pass
    kernel_v = kernel_1d.view(1, 1, kernel_size, 1).expand(channels, 1, kernel_size, 1)
    img = F.conv2d(img, kernel_v, padding=(padding, 0), groups=channels)
    
    return img

# Test
if __name__ == "__main__":
    # Create test image
    device = 'cuda' if torch.cuda.is_available() else 'cpu'
    img = torch.randn(1, 3, 64, 64).to(device)
    kernel_size = 5
    
    # Apply both blurs
    blurred_old = gaussian_blur_old(img, kernel_size)
    blurred_new = gaussian_blur_new(img, kernel_size)
    
    # Compare
    diff = torch.abs(blurred_old - blurred_new).mean()
    print(f"Mean difference between old and new blur: {diff.item()}")
    
    # Also compare with standard 2D convolution
    sigma = (kernel_size - 1) / 6.0
    x = torch.arange(kernel_size, dtype=torch.float32, device=img.device) - (kernel_size-1)/2
    kernel_1d = torch.exp(-x**2 / (2*sigma**2))
    kernel_1d /= kernel_1d.sum()
    kernel_2d = torch.outer(kernel_1d, kernel_1d).view(1, 1, kernel_size, kernel_size)
    channels = img.shape[-3]
    kernel_2d = kernel_2d.expand(channels, 1, kernel_size, kernel_size)
    padding = kernel_size // 2
    blurred_2d = F.conv2d(img, kernel_2d, padding=padding, groups=channels)
    
    diff_new_2d = torch.abs(blurred_new - blurred_2d).mean()
    print(f"Mean difference between new separable and 2D convolution: {diff_new_2d.item()}")
    
    print("Test completed. New blur should match 2D convolution closely.")