Low PSNR and SSIM Values

[mahiratmis]

Hello, thank you very much for the code. I applied your model on SHIQ Dataset and the metric scores except for MSE are really low
Average PSNR: 22.78
Average MSE: 0.0070
Average SSIM: 0.8810
Am I missing something?
The code I used is as follows

import torch
from PIL import Image
import os 

base_folder = os.path.join("datasets", "SHIQ_data_10825", "test")
output_folder = "output_shiq"
os.makedirs(output_folder, exist_ok=True)

for filename in os.listdir(base_folder):
    if filename.endswith("_A.png"):
        input_image = Image.open(os.path.join(base_folder, filename))
        predictor = torch.hub.load("Stable-X/StableDelight", "StableDelight_turbo", trust_repo=True)
        delight_image = predictor(input_image)
        delight_image.save(os.path.join(output_folder, filename))

[mahiratmis]

The evaluation code is as follows:

import os
import torch
import numpy as np
from PIL import Image
from torchvision import transforms
from skimage.metrics import structural_similarity as compare_ssim  # for SSIM

def psnr(pred, target, pixel_max_cnt=1):
    """
    Compute PSNR between pred and target images.
    Assumes input tensors are in the [0,1] range.
    """
    mse_val = torch.mean((target - pred) ** 2).item()
    rmse_avg = mse_val ** 0.5
    p = 20 * np.log10(pixel_max_cnt / rmse_avg)
    return p

def mse(pred, target):
    """
    Compute Mean Squared Error between pred and target images.
    """
    return torch.mean((target - pred) ** 2).item()

def ssim(pred, target):
    """
    Compute SSIM between pred and target images.
    This function expects tensors with shape [batch, channels, height, width].
    """
    # Rearrange tensors from [batch, channels, height, width] to [batch, height, width, channels]
    pred_np = pred.clone().data.permute(0, 2, 3, 1).cpu().numpy()
    target_np = target.clone().data.permute(0, 2, 3, 1).cpu().numpy()
    # Take the first image in the batch
    pred_np = pred_np[0]
    target_np = target_np[0]
    # print("target.shape, pred.shape:", target_np.shape, pred_np.shape)
    ssim_val = compare_ssim(target_np, pred_np, channel_axis=2, data_range=1.0)
    return ssim_val

def main(ground_truth_folder, prediction_folder):
    # List filenames in the ground truth folder that end with _D.png
    gt_filenames = [f for f in os.listdir(ground_truth_folder) if f.endswith('_D.png')]

    # List filenames in the prediction folder that end with _A.png
    pred_filenames = [f for f in os.listdir(prediction_folder) if f.endswith('_A.png')]
    
    # Check that both folders have the same number of images
    if len(gt_filenames) != len(pred_filenames):
        print("Error: The number of images in the ground truth folder and prediction folder do not match.")
        print(f"Ground truth images: {len(gt_filenames)}, Prediction images: {len(pred_filenames)}")
        return  
    
    
    # Define a transform to convert images to torch tensors in the [0,1] range.
    transform = transforms.ToTensor()
    
    psnr_scores = []
    mse_scores = []
    ssim_scores = []
    
    for gt_filename in gt_filenames:
        gt_path = os.path.join(ground_truth_folder, gt_filename)
        # Construct corresponding prediction filename by replacing _D.png with _A.png
        pred_filename = gt_filename.replace('_D.png', '_A.png')
        pred_path = os.path.join(prediction_folder, pred_filename)
        
        if not os.path.exists(pred_path):
            print(f"Prediction for {gt_filename} (expected as {pred_filename}) not found, skipping.")
            continue
        
        # Open and convert images to RGB for consistency
        gt_image = Image.open(gt_path).convert('RGB')
        pred_image = Image.open(pred_path).convert('RGB')
        
        # Convert images to torch tensors and add a batch dimension
        gt_tensor = transform(gt_image).unsqueeze(0)
        pred_tensor = transform(pred_image).unsqueeze(0)
        
        # Compute the metrics
        psnr_val = psnr(pred_tensor, gt_tensor)
        mse_val = mse(pred_tensor, gt_tensor)
        ssim_val = ssim(pred_tensor, gt_tensor)
        
        psnr_scores.append(psnr_val)
        mse_scores.append(mse_val)
        ssim_scores.append(ssim_val)
        
        print(f"{gt_filename} vs {pred_filename}: PSNR = {psnr_val:.2f}, MSE = {mse_val:.4f}, SSIM = {ssim_val:.4f}")
    
    # Compute and print average scores over all images
    if psnr_scores:
        avg_psnr = np.mean(psnr_scores)
        avg_mse = np.mean(mse_scores)
        avg_ssim = np.mean(ssim_scores)
        print("\nAverage Metrics:")
        print("Average PSNR: {:.2f}".format(avg_psnr))
        print("Average MSE: {:.4f}".format(avg_mse))
        print("Average SSIM: {:.4f}".format(avg_ssim))

if __name__ == '__main__':
    # Update these paths to point to your ground truth and prediction image folders.
    ground_truth_folder = '../datasets/SHIQ_data_10825/test'
    prediction_folder = 'output_shiq'
    
    main(ground_truth_folder, prediction_folder)