from torchmetrics.image import (
    LearnedPerceptualImagePatchSimilarity,
    PeakSignalNoiseRatio,
    StructuralSimilarityIndexMeasure,
)

psnr = PeakSignalNoiseRatio(1.0, reduction="none")
ssim = StructuralSimilarityIndexMeasure(data_range=1.0, reduction="none")
lpips = LearnedPerceptualImagePatchSimilarity(
    net_type="alex", reduction="none", normalize=True
)


def benchmark(image1, image2):
    return psnr(image1, image2), ssim(image1, image2), lpips(image1, image2)
