Image Quality Metrics

This example demonstrates the image quality metrics available in Kaira, including PSNR (Peak Signal-to-Noise Ratio), SSIM (Structural Similarity Index), MS-SSIM (Multi-Scale SSIM), and LPIPS (Learned Perceptual Image Patch Similarity).

These metrics are particularly useful for: * Evaluating image compression algorithms * Assessing deep learning-based image processing * Quality control in image transmission systems

import os

import matplotlib.pyplot as plt

First, let’s import the necessary modules

import torch
import torchvision.transforms as T
from PIL import Image

from kaira.data import ImageDataset
from kaira.metrics.image.lpips import LearnedPerceptualImagePatchSimilarity
from kaira.metrics.image.psnr import PeakSignalNoiseRatio
from kaira.metrics.image.ssim import MultiScaleSSIM, StructuralSimilarityIndexMeasure

# Load sample images using the new simplified dataset interface
dataset = ImageDataset(name="cifar10", size=(256, 256))

# Extract images from dataset
images = []
image_names = []
for i in range(4):  # Just use first 4 images
    image_tensor, label = dataset[i]  # ImageDataset returns (image, label)
    images.append(image_tensor)
    image_names.append(f"cifar10_image_{i}")

images = torch.stack(images)
print(f"Loaded {len(images)} test images: {image_names}")

  0%|          | 0.00/170M [00:00<?, ?B/s]
  0%|          | 426k/170M [00:00<00:40, 4.18MB/s]
  5%|▍         | 7.93M/170M [00:00<00:03, 45.5MB/s]
 11%|█▏        | 19.3M/170M [00:00<00:01, 76.6MB/s]
 18%|█▊        | 31.1M/170M [00:00<00:01, 92.5MB/s]
 25%|██▌       | 42.7M/170M [00:00<00:01, 101MB/s]
 32%|███▏      | 54.1M/170M [00:00<00:01, 105MB/s]
 39%|███▊      | 65.8M/170M [00:00<00:00, 109MB/s]
 45%|████▌     | 77.2M/170M [00:00<00:00, 111MB/s]
 52%|█████▏    | 88.7M/170M [00:00<00:00, 112MB/s]
 59%|█████▉    | 100M/170M [00:01<00:00, 113MB/s]
 66%|██████▌   | 112M/170M [00:01<00:00, 114MB/s]
 72%|███████▏  | 123M/170M [00:01<00:00, 114MB/s]
 79%|███████▉  | 135M/170M [00:01<00:00, 114MB/s]
 86%|████████▌ | 146M/170M [00:01<00:00, 114MB/s]
 93%|█████████▎| 158M/170M [00:01<00:00, 115MB/s]
 99%|█████████▉| 169M/170M [00:01<00:00, 115MB/s]
100%|██████████| 170M/170M [00:01<00:00, 106MB/s]
Loaded 4 test images: ['cifar10_image_0', 'cifar10_image_1', 'cifar10_image_2', 'cifar10_image_3']

Create different types of distortions

We’ll create different types of distortions to compare how various metrics assess them

def add_gaussian_noise(image, std=0.1):
    """Add Gaussian noise to image."""
    return image + torch.randn_like(image) * std


def add_salt_pepper_noise(image, prob=0.05):
    """Add salt and pepper noise to image."""
    mask = torch.rand_like(image)
    image = image.clone()
    image[mask < prob / 2] = 0  # salt
    image[mask > 1 - prob / 2] = 1  # pepper
    return image


def blur_image(image, kernel_size=3):
    """Apply Gaussian blur to image."""
    return T.GaussianBlur(kernel_size)(image)


def compress_image(image, quality=10):
    """Simulate JPEG compression artifacts."""
    to_pil = T.ToPILImage()
    to_tensor = T.ToTensor()
    pil_image = to_pil(image)
    # Create a temporary file for compression
    temp_file = "temp.jpg"
    pil_image.save(temp_file, quality=quality)
    try:
        compressed = Image.open(temp_file)
        return to_tensor(compressed)
    finally:
        if os.path.exists(temp_file):
            os.remove(temp_file)


# Create distorted versions
noisy_images = torch.stack([add_gaussian_noise(img) for img in images])
sp_noisy_images = torch.stack([add_salt_pepper_noise(img) for img in images])
blurred_images = torch.stack([blur_image(img) for img in images])
compressed_images = torch.stack([compress_image(img) for img in images])

Initialize metrics

We’ll create individual metrics directly without using the registry

# Initialize metrics manually
psnr = PeakSignalNoiseRatio(data_range=1.0, reduction="mean")  # or PSNR()
ssim = StructuralSimilarityIndexMeasure(data_range=1.0, reduction="mean")  # or SSIM()
ms_ssim = MultiScaleSSIM(data_range=1.0, reduction="mean")  # Add reduction parameter
lpips = LearnedPerceptualImagePatchSimilarity(net_type="alex")  # remove redundant reduction parameter

Downloading: "https://download.pytorch.org/models/alexnet-owt-7be5be79.pth" to /home/docs/.cache/torch/hub/checkpoints/alexnet-owt-7be5be79.pth

  0%|          | 0.00/233M [00:00<?, ?B/s]
 15%|█▍        | 34.4M/233M [00:00<00:00, 359MB/s]
 30%|███       | 70.4M/233M [00:00<00:00, 369MB/s]
 46%|████▌     | 106M/233M [00:00<00:00, 373MB/s]
 61%|██████    | 142M/233M [00:00<00:00, 374MB/s]
 76%|███████▋  | 178M/233M [00:00<00:00, 359MB/s]
 91%|█████████ | 212M/233M [00:00<00:00, 322MB/s]
100%|██████████| 233M/233M [00:00<00:00, 327MB/s]

Evaluate metrics on different distortions

def evaluate_all_metrics(original, distorted):
    """Evaluate all metrics between original and distorted images."""
    return {"PSNR": psnr(distorted, original), "SSIM": ssim(distorted, original), "MS-SSIM": ms_ssim(distorted, original), "LPIPS": lpips(distorted, original)}  # Now returns scalar mean  # Now returns scalar mean  # Now returns scalar mean  # Now returns scalar mean


# Evaluate metrics for each type of distortion
gaussian_metrics = evaluate_all_metrics(images, noisy_images)
sp_metrics = evaluate_all_metrics(images, sp_noisy_images)
blur_metrics = evaluate_all_metrics(images, blurred_images)
compress_metrics = evaluate_all_metrics(images, compressed_images)

Visualize results

First, let’s look at the distorted images

plt.figure(figsize=(15, 8))
titles = ["Original", "Gaussian Noise", "Salt & Pepper", "Blur", "Compressed"]
all_images = [images, noisy_images, sp_noisy_images, blurred_images, compressed_images]

for i, (title, imgs) in enumerate(zip(titles, all_images)):
    plt.subplot(2, 3, i + 1)
    plt.imshow(imgs[0].permute(1, 2, 0).clip(0, 1))
    plt.title(title)
    plt.axis("off")

plt.tight_layout()
plt.show()

Now let’s compare how different metrics evaluate each distortion type

all_metrics = [gaussian_metrics, sp_metrics, blur_metrics, compress_metrics]

labels = ["Gaussian", "Salt & Pepper", "Blur", "Compression"]

# Create a manual comparison visualization
plt.figure(figsize=(14, 10))

# Plot PSNR values
plt.subplot(2, 2, 1)
psnr_values = [metrics["PSNR"].item() for metrics in all_metrics]  # Convert tensor to Python scalar
plt.bar(labels, psnr_values, color="blue")
plt.xlabel("Distortion Type")
plt.ylabel("PSNR (dB)")
plt.title("PSNR Comparison")
plt.grid(axis="y", alpha=0.3)

# Plot SSIM values
plt.subplot(2, 2, 2)
ssim_values = [metrics["SSIM"].item() for metrics in all_metrics]  # Convert tensor to Python scalar
plt.bar(labels, ssim_values, color="green")
plt.xlabel("Distortion Type")
plt.ylabel("SSIM")
plt.title("SSIM Comparison")
plt.grid(axis="y", alpha=0.3)

# Plot MS-SSIM values
plt.subplot(2, 2, 3)
msssim_values = [metrics["MS-SSIM"].item() for metrics in all_metrics]  # Convert tensor to Python scalar
plt.bar(labels, msssim_values, color="purple")
plt.xlabel("Distortion Type")
plt.ylabel("MS-SSIM")
plt.title("MS-SSIM Comparison")
plt.grid(axis="y", alpha=0.3)

# Plot LPIPS values (lower is better)
plt.subplot(2, 2, 4)
lpips_values = [metrics["LPIPS"].item() for metrics in all_metrics]  # Convert tensor to Python scalar
plt.bar(labels, lpips_values, color="red")
plt.xlabel("Distortion Type")
plt.ylabel("LPIPS (lower is better)")
plt.title("LPIPS Comparison")
plt.grid(axis="y", alpha=0.3)

plt.tight_layout()
plt.show()

Interpreting the Results

The results show how different metrics capture various aspects of image quality:

• PSNR is a simple pixel-level metric that measures absolute differences * Higher values indicate better quality * More sensitive to noise than blurring * May not align well with human perception

• SSIM considers structural information * Values range from -1 to 1 (higher is better) * More tolerant of uniform changes * Better correlation with human perception than PSNR

• MS-SSIM evaluates structural similarity at multiple scales * Similar to SSIM but captures both local and global structures * Often preferred for high-resolution images * Better at detecting blur than basic SSIM

• LPIPS uses deep features to measure perceptual similarity * Lower values indicate better perceptual quality * Trained on human perceptual judgments * Often best matches human quality assessment

Different distortions affect these metrics differently:

• Gaussian noise heavily impacts PSNR but less so SSIM

• Blur might maintain good PSNR but show poor SSIM/MS-SSIM

• LPIPS often identifies perceptually significant distortions that other metrics might miss

For practical applications:

• Use multiple metrics for comprehensive evaluation

• Consider the specific requirements of your application

• LPIPS is recommended when perceptual quality is critical

• PSNR/SSIM are good for optimization objectives due to their mathematical properties