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

DISTS는 올해(2020년) 발표된 FR-IQA 방법입니다. FR 방법이지만 훈련과정이 필요한 알고리즘입니다. IQA_pytorch 패키지는 IQA 데이터베이스 중 하나인 KADID에 훈련된 DISTS의 파이썬 코드를 제공합니다. 파이토치 기반의 코드입니다. 

 

DISTS로 이미지의 품질을 평가하기 위해서는 먼저 다음 세가지의 패키지를 설치해주셔야 합니다. 

 

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

2. torchvision: pip install torchvision

3. IQA_pytorch: pip install IQA_pytorch

 

저는 미리 훈련된 DISTS로 다음 왜곡 이미지의 품질을 평가하겠습니다.

 

 

참고로 DISTS는 좋은 품질일수록 0에 가까운 작은 점수를, 나쁜 품질일수록 큰 점수를 반환합니다. 

 

DISTS를 이용해서 품질을 예측하기 위해 필요한 파이썬 코드는 다음과 같습니다. 

 

from IQA_pytorch import DISTS, 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 = DISTS().to(device)
 
score = model(dist, ref, as_loss=False)
print('score: %.4f' % score.item())

 

코드를 실행했더니, 0.3347의 점수가 나왔습니다. 

 

 

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

 

 

이번에는 훨씬 작은 점수가 나왔습니다. 아까 말씀드렸듯이 DISTS는 품질이 좋을 수록 작은 점수를 반환합니다. 

 

 

b스카이비전의 추천글

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