File size: 606 Bytes
62dbcfb |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 |
import cv2
from cv2.typing import MatLike
from traiNNer.metrics.psnr_ssim import calculate_ssim
def open_image(file_path: str) -> MatLike:
img = cv2.imread(file_path, cv2.IMREAD_UNCHANGED)
if len(img.shape) == 2:
img = cv2.cvtColor(img, cv2.COLOR_GRAY2RGB)
else:
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
return img
def main() -> None:
img1 = open_image(r"./img1.png")
img2 = open_image(r"./img2.png")
ssim_value = calculate_ssim(img1, img2, crop_border=4)
print(f"SSIM: {ssim_value}")
if __name__ == "__main__":
main()
|