SSIMLoss

ssim(x, y, kernel=ssim_kernel(x); peakval=1, crop=true, dims=:)

Return the structural similarity index measure (SSIM) between two signals. SSIM is computed via the mean of a sliding window of statistics computed between the two signals. By default, the sliding window is a Gaussian with side-length 11 in each signal dimension and σ=1.5. crop=false will pad x and y such that the sliding window computes statistics centered at every pixel of the input (via same-size convolution). ssim computes statistics independently over channel and batch dimensions. x and y may be 3D/4D/5D tensors with channel and batch-dimensions.

peakval=1 is the standard for image comparisons, but in practice should be set to the maximum value of your signal type.

dims determines which dimensions to average the computed statistics over. If dims=1:ndims(x)-1, SSIM will be computed for each batch-element separately.

The results of ssim are matched against those of ImageQualityIndexes for grayscale and RGB images (i.e. x, y both of size (N1, N2, 1, B) and (N1, N2, 3, B) for grayscale and color images, resp.).

See also ssim_loss, ssim_loss_fast.

ssim_kernel(T, N)

Return Gaussian kernel with σ=1.5 and side-length 11 for use in ssim. Returned kernel will be N-2 dimensional of type T.

ssim_kernel(x::AbstractArray{T, N}) where {T, N}

Return Gaussian kernel with σ=1.5 and side-length 11 for use in ssim. Returned array will be on the same device as x.

ssim_loss(x, y, kernel=ssim_kernel(x); peakval=1, crop=true, dims=:)

Computes 1 - ssim(x, y), suitable for use as a loss function with gradient descent. For faster training, it is recommended to store a kernel and reuse it, ex.,

kernel = ssim_kernel(Float32, 2) |> gpu
# or alternatively for faster computation
# kernel = ones(Float32, 5, 5, 1, num_channels) |> gpu

for (x, y) in dataloader
    x, y = (x, y) .|> gpu
    grads = gradient(model) do m
        x̂ = m(y)
        ssim_loss(x, x̂, kernel)
    end
    # update the model ...
end

See ssim for a detailed description of SSIM and the above arguments. See also ssim_loss_fast.

ssim_loss_fast(x, y; kernel_length=5, peakval=1, crop=true, dims=:)

Computes ssim_loss with an averaging kernel instead of a large Gaussian kernel for faster computation. kernel_length specifies the averaging kernel side-length in each signal dimension of x, y. See ssim for a detailed description of SSIM and the above arguments.

See also ssim_loss.