简体   繁体   中英

Comparing two images with ms_ssim on python

I want to compare two images but ms_ssim wants 4D tensor

https://pypi.org/project/pytorch-msssim/

I Tried

from PIL import Image
from tqdm import tqdm
from pytorch_msssim import ssim, ms_ssim, SSIM, MS_SSIM
import torchvision
import numpy as np

topil=torchvision.transforms.ToPILImage()
totensor=torchvision.transforms.ToTensor()

def ssimcompare(path1:str,path2:str)->float:
    image1 = Image.open(path1)
    image2 = Image.open(path2) 
    #it1=np.expand_dims(totensor(topil(np.array(image1))), axis=0)
    #it2=np.expand_dims(totensor(topil(np.array(image2))), axis=0)
    #it1=totensor(np.expand_dims(np.array(image1), axis=0))
    #it2=totensor(np.expand_dims(np.array(image2), axis=0))
    it1=totensor(np.array(image1))
    it2=totensor(np.array(image2))
    valor=ms_ssim( it1 , it2, data_range=255, size_average=False )
    return valor

But I get different errors

ValueError: Input images must be 4-d tensors.
TypeError: pic should be Tensor or ndarray. Got <class 'PIL.JpegImagePlugin.JpegImageFile'>.
AttributeError: 'numpy.ndarray' object has no attribute 'type'

The problem is that all these functions (and classes) requires batches of images as input. But, since an image is 3D, a batch is 4D.

When you have only one image tensor you can "unsqueeze" it into a one-item batch with

it = it.unsqueeze(0)

However, I don't recommend the pytorch_msssim package. You should consider the piqa or piq IQA packages, as they are well documented and have faster implementations.

For example,

pip install piqa

Then, your function becomes

from piqa import ssim

def msssim_compare(path1: str, path2: str) -> float:
    image1 = Image.open(path1)
    image2 = Image.open(path2)
    it1 = totensor(image1).unsqueeze(0)
    it2 = totensor(image2).unsqueeze(0)

    return ssim.msssim(it1, it2).squeeze(0)

The technical post webpages of this site follow the CC BY-SA 4.0 protocol. If you need to reprint, please indicate the site URL or the original address.Any question please contact:yoyou2525@163.com.

 
粤ICP备18138465号  © 2020-2024 STACKOOM.COM