R/ssim.R

In juoe/spatialtools: Collection of functions for doing spatial stuff.

#' Calculate SSIM from two rasters
#'
#' adopted from package 'spatialcompare'
#'
#' @param img1 is a \code{raster} object to compare
#' @param img2 is a \code{raster} object to compare
#' @param w is the width of the neighbourhood in number of pixels out from centre cell
#' @param gauss_filter is a binary flag for whether a Gaussian filter should be applied as a smoothing function
#' @param edge is a binary flag for whether a torroidal edge correction should be applied
#' @param ks is a vector of length 2 which contains values for constants in the SSIM formula. If ignored default values will be used.
#' @param global_values a binary flag whether overall index values should be returned or else raster stack with cell-by-cell values
#' @return A list of length 4 or a raster stack with 4 bands. List objects / layer
#'   contain the composite index, luminance, contrast, and structure components respectively.
#' @export


ssim <- function(img1, img2, w, gauss_filter=TRUE, edge=FALSE, ks=c(0.01, 0.03), global_values = FALSE) {
  #set constants
  N <- FALSE
  library(raster)

  L <- max(cellStats(img1, max), cellStats(img2, max))
  globalMin <- abs(min(cellStats(img1, min), cellStats(img2, min)))
  L <- L - globalMin
  K <- ks
  C1 <-(K[1]*L)^2
  C2 <-(K[2]*L)^2
  C3 <-C2/2

  sigma = 1.5
  #create null filter, optionally replace with weighted version
  filterx=matrix(1, nrow=(w*2)+1, ncol=(w*2)+1) / sum(matrix(1, nrow=(w*2)+1, ncol=(w*2)+1))

  getGauss <- function(sigma, w) {
    w2 = (w*2) + 1
    gf1 <- matrix(nrow=w2, ncol=w2)
    gf2 <- matrix(nrow=w2, ncol=w2)
    for(i in 1:w2) {
      gf1[i,] <-  c(w:0, 1:w)
      gf2[,i] <-  c(w:0, 1:w)
    }
    gf <- (1 / (2 * pi * (sigma^2))) * exp(-(gf1^2 + gf2^2) / (2 * (sigma^2)))
    return(gf/sum(gf))
  }

  #get the Gaussian Kernel
  if(gauss_filter == TRUE) {
    filterx=getGauss(sigma, w)
  }
  #get mu
  mu1 <- focal(img1, filterx, fun=sum, na.rm=N)
  mu2 <- focal(img2, filterx, fun=sum, na.rm=N)
  img12 <- img1 * img2
  #square
  mu1mu2 <- mu1 * mu2
  mu1sq <- mu1 * mu1
  mu2sq <- mu2 * mu2
  #normalized sigma sq
  sigsq1<- focal(img1*img1,filterx,fun=sum, na.rm=N) - mu1sq
  sigsq2<- focal(img2*img2,filterx,fun=sum, na.rm=N) - mu2sq
  sig12 <- focal(img1*img2,filterx,fun=sum, na.rm=N) - mu1mu2
  #std dev
  sig1 <- sigsq1 ^ 0.5
  sig2 <- sigsq2 ^ 0.5
  #compute components
  L <- ((2*mu1mu2)+C1) / (mu1sq + mu2sq + C1)
  C <- (2*sig1*sig2+C2) / (sigsq1 + sigsq2 + C2)
  S <- (sig12 + C3) / (sig1 * sig2 + C3)
  #compute SSIMap
  SSIM2 <- L * C * S
  #compute SSIM
  num <- (2*mu1mu2+C1)*(2*sig12+C2)
  denom <- (mu1sq+mu2sq+C1) * (sigsq1+sigsq2+C2)

  #global values
  mSSIM <- cellStats((num / denom), mean)
  mSIM <- cellStats(L, mean)
  mSIV <- cellStats(C, mean)
  mSIP <- cellStats(S, mean)

  if(global_values) {
    out <- list(mSSIM, mSIM, mSIV, mSIP)
  } else {
    out <- raster::stack(SSIM2, L, C, S)
  }

  names(out) <- c("SSIM", "SIM", "SIV", "SIP")
  return(out)

}