Add tests

image_similarity_measures/quality_metrics.py

@@ -39,12 +39,9 @@ def rmse(org_img: np.ndarray, pred_img: np.ndarray, max_p: int = 4095) -> float:
         org_img = np.expand_dims(org_img, axis=-1)
 
     rmse_bands = []
-    for i in range(org_img.shape[2]):
-        dif = np.subtract(org_img[:, :, i], pred_img[:, :, i])
-        m = np.mean(np.square(dif / max_p))
-        s = np.sqrt(m)
-        rmse_bands.append(s)
-
+    diff = org_img - pred_img
+    mse_bands = np.mean(np.square(diff / max_p), axis=(0, 1))
+    rmse_bands = np.sqrt(mse_bands)
     return np.mean(rmse_bands)
 
 
@@ -66,21 +63,19 @@ def psnr(org_img: np.ndarray, pred_img: np.ndarray, max_p: int = 4095) -> float:
     if org_img.ndim == 2:
         org_img = np.expand_dims(org_img, axis=-1)
 
-    mse_bands = []
-    for i in range(org_img.shape[2]):
-        mse_bands.append(np.mean(np.square(org_img[:, :, i] - pred_img[:, :, i])))
-
-    return 20 * np.log10(max_p) - 10.0 * np.log10(np.mean(mse_bands))
+    mse_bands = np.mean(np.square(org_img - pred_img), axis=(0, 1))
+    mse = np.mean(mse_bands)
+    return 20 * np.log10(max_p / np.sqrt(mse))
 
 
 def _similarity_measure(x: np.array, y: np.array, constant: float):
     """
     Calculate feature similarity measurement between two images
     """
-    numerator = 2 * x * y + constant
-    denominator = x**2 + y**2 + constant
+    numerator = 2 * np.multiply(x, y) + constant
+    denominator = np.add(np.square(x), np.square(y)) + constant
 
-    return numerator / denominator
+    return np.divide(numerator, denominator)
 
 
 def _gradient_magnitude(img: np.ndarray, img_depth: int):

pyproject.toml

@@ -13,6 +13,10 @@ classifiers=[
     "Intended Audience :: Science/Research",
     "Development Status :: 5 - Production/Stable",
 ]
+[tool.poetry.group.test.dependencies]
+pytest = "^7.3.0"
+rasterio = "^1.3.6"
+
 
 [project.urls]
 Source = 'https://github.com/up42/image-similarity-measures'