things

src/benchmark/__init__.py

@@ -1,15 +1,22 @@
+import lpips
+from pytorch_msssim import ssim
 from torchmetrics.image import (
-    LearnedPerceptualImagePatchSimilarity,
     PeakSignalNoiseRatio,
     StructuralSimilarityIndexMeasure,
 )
 
 psnr = PeakSignalNoiseRatio(1.0, reduction="none", dim=(1, 2, 3))
-ssim = StructuralSimilarityIndexMeasure(data_range=1.0, reduction="none")
-lpips = LearnedPerceptualImagePatchSimilarity(
-    net_type="alex", reduction="none", normalize=True
-)
+lp = lpips.LPIPS(net="alex")
 
 
 def benchmark(image1, image2):
-    return psnr(image1, image2), ssim(image1, image2), lpips(image1, image2)
+    return (
+        psnr(image1, image2),
+        ssim(
+            image1,
+            image2,
+            data_range=1.0,
+            size_average=False,
+        ),
+        lp(image1 * 2 - 1, image2 * 2 - 1),
+    )