[python] SSIM 파이썬 코드 실행하기(IQA_pytorch 패키지)

SSIM은 대표적인 FR IQA(image quality assessment, 이미지 품질 평가) 방법입니다. SSIM의 저자들은 기본적으로 SSIM에 대한 Matlab 코드를 제공합니다. 그런데 점차 연구자들이 python으로 IQA 알고리즘들을 개발하기 시작하면서 기존 IQA 알고리즘의 Matlab 코드들도 python으로 새롭게 쓰여지고 있는 추세입니다. 오늘 소개해드리는 IQA_pytorch 패키지는 SSIM, MS-SSIM, CW-SSIM, FSIM, VSI, GMSD, VIF, LPIPS, DISTS 등 FR-IQA 알고리즘들의 python 코드들을 제공합니다.

 

오늘은 그 중에서 IQA_pytorch 패키지를 이용해서 SSIM 파이썬 코드를 실행하는 방법에 대해 다루도록 하겠습니다. 우선 다음과 같은 패키지들을 우선 설치해주셔야 합니다. 저는 아나콘다에서 가상환경을 하나 만든 후에 다음과 같은 명령어로 패키지들을 설치했습니다. 

 

1. pytorch 1.4버전 conda install pytorch==1.4

2. torchvision 설치 pip install torchvision

3. IQA_pytorch 설치 pip install IQA_pytorch

 

주의하셔야 할 것은 pytorch를 1.4버전으로 설치해야 합니다. 그래야 호환되는 torchvision을 설치할 수 있습니다. 설치를 마치셨다면 품질을 평가할 이미지를 준비해주세요. 원본 이미지 1장과 왜곡 이미지 1장이 필요합니다. 저는 아래와 같은 이미지의 품질을 평가하도록 하겠습니다. 

 

오른쪽에 있는 이미지를 보면 품질이 현저히 떨어지죠? SSIM은 왜곡 이미지가 원본 이미지와 비교했을 때 품질이 어떠한지를 알려줍니다. 1에 가까울 수록 높은 품질이고, 0에 가까울 수록 나쁜 품질입니다. 

 

SSIM 파이썬 코드를 실행하시려면 다음과 같은 코드를 실행하시면 됩니다. 

 

from IQA_pytorch import SSIM, utils
from PIL import Image
import torch
 
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 
ref_path  = 'r0.png'
dist_path = 'r1.png' 
 
ref = utils.prepare_image(Image.open(ref_path).convert("RGB")).to(device)
dist = utils.prepare_image(Image.open(dist_path).convert("RGB")).to(device)
 
model = SSIM(channels=3)
 
score = model(dist, ref, as_loss=False)
print('score: %.4f' % score.item())

 

실행했더니 위 왜곡 이미지의 품질 점수는 다음과 같이 나왔습니다. 0.6717

 

 

비교를 위해서 위 이미지보다는 좀 더 괜찮은 품질의 왜곡 이미지의 점수를 산출해보도록 하겠습니다. 

 

 

육안으로 보기에 두 이미지가 큰 차이가 없죠? 1에 가까운 값이 나와야 알고리즘이 제대로 작동하는 것입니다. 그러면 이미지 path만 바꾼 후에 위 코드를 다시 실행해보겠습니다. 

 

 

이번에는 아까보다 훨씬 높은 0.9541이 나왔습니다. 아마 첫번째 왜곡 이미지와 두번째 왜곡 이미지의 중간 정도 품질의 이미지를 테스트해보면 0.6717과 0.9541 중간 정도의 값이 산출될 것입니다. 

 

 

b스카이비전의 추천글

☞ [ubuntu+python] BRISQUE 모델로 이미지 품질 평가하기  

☞ [IQA] 2D 이미지 품질 평가에 구조변화를 반영하는 SSIM과 그의 변형들  

☞ [Anaconda+python] GMSD의 matlab 코드를 python으로 작성하기