R/rfSegmentation.R

In neuroconductor-devel/ANTsR: ANTs in R: Quantification Tools for Biomedical Images

#' A rfSegmentation function.
#'
#' Unsupervised image segmentation via random forests.  An example.
#'
#' @param featureMatrix input matrix of features matched to mask size
#' n predictors in rf
#' @param mask input antsImage mask
#' @param labelimg input antsImage labelimage, optional for supervised seg
#' @param ntrees number of rf trees
#' @param verbose boolean
#' @return list of n-probability images is output where n is number of classes
#' @author Tustison NJ, Avants BB
#' @seealso \code{\link{mrvnrfs}}
#' @examples
#' \dontrun{
#' if ( usePkg("randomForest") ) {
#' img<-antsImageRead( getANTsRData("r16") ) %>% iMath("Normalize")
#' mask<-getMask( img )
#' segs<-kmeansSegmentation( img, k=3, kmask = mask)
#' fmat0 = t( antsrimpute( getNeighborhoodInMask( img, mask, c(2,2) ) ) )
#' fmat1 = t( antsrimpute( getNeighborhoodInMask(
#'   segs$probabilityimages[[1]], mask, c(2,2) ) ) )
#' fmat2 = t( antsrimpute( getNeighborhoodInMask(
#'   segs$probabilityimages[[2]], mask, c(2,2) ) ) )
#' fmat3 = t( antsrimpute( getNeighborhoodInMask(
#'   segs$probabilityimages[[3]], mask, c(2,2) ) ) )
#' fmat = cbind( fmat0, fmat1, fmat2, fmat3 )
#' # produces proximity between all voxel pairs
#' rfsegs<-rfSegmentation( fmat, verbose=FALSE )
#' lrr = lowrankRowMatrix( rfsegs, 10, faster = TRUE )
#' nv = eanatSelect( inmat=lrr, mask=mask, selectorScale=1.2,  cthresh=50,
#'   verbose=T, smooth=1 )
#' ee = eanatDef( lrr, mask, nvecs=nv, smooth=0., cthresh=50,
#'   its=2, verbose=TRUE )
#' eseg = eigSeg( mask, ee )
#' plot( img, eseg )
#' }
#' }
#' @export rfSegmentation
rfSegmentation <- function( featureMatrix,
  mask, labelimg=NULL,
  ntrees = 100, verbose = FALSE) {
  if (nargs() == 0) {
    print("Usage:  probs<-rfSegmentation( x, x2 ) ")
    return(1)
  }
  haverf<-usePkg("randomForest")
  if ( !haverf ) {
    print("need randomForest package for this function")
    return(NA)
  }
  mask = check_ants(mask)
  haveLabels = !is.null( labelimg )
  if ( haveLabels )
    {
    labelimg = check_ants(labelimg)
    labels <- as.factor(labelimg[mask == 1])
    mydf <- data.frame(labels = labels, featureMatrix)
    myrf <- randomForest::randomForest(y=labels,x=featureMatrix, ntree = ntrees, type = "classification",
    importance = TRUE, na.action = na.omit, do.trace = verbose)
    }
  else
    {
    myrf <- randomForest::randomForest( x=featureMatrix, ntree = ntrees,
      type = "unsupervised",
      importance = TRUE, na.action = na.omit, do.trace = verbose)
    return( myrf$proximity )
    }
  if ( verbose ) print( myrf )
  probabilityimages <- predict(myrf, type = "prob")
  probabilityimages <- matrixToImages(t(probabilityimages), mask)
  segs <- antsImageClone(mask)
  segs[mask == 1] <- predict(myrf)
  myout <- list(segmentation = segs, probabilityimages = probabilityimages, rfModel = myrf)
  return(myout)
}

#' A rfSegmentationPredict function.
#'
#' Predict image segmentation via random forests.
#'
#' @param rfSegmentationModel input rf model
#' @param featureMatrix input feature matrix
#' @param mask antsImage mask
#' @param verbose bool
#' @return segmentation is output
#' @author Tustison NJ, Avants BB
#' @examples
#'
#' if ( usePkg('randomForest') ) {
#' img<-antsImageRead( getANTsRData("r16"))
#' mask<-getMask( img )
#' mask2<-getMask( img )
#' mask [ 129:255, 1:255 ]<-0
#' mask2[ 2:128, 1:255 ]<-0
#' segs<-kmeansSegmentation( img, k=3, kmask = mask)
#' fmat = t( antsrimpute( getNeighborhoodInMask( img, mask, c(2,2) ) ) )
#' rfsegs<-rfSegmentation(  fmat, mask, segs$segmentation, ntrees=100 )
#' fmat2 = t( antsrimpute( getNeighborhoodInMask( img, mask2, c(2,2) ) ) )
#' rfseg2<-rfSegmentationPredict(  rfsegs$rfModel , fmat2 , mask2 )
#' }
#'
#' @export rfSegmentationPredict
rfSegmentationPredict <- function(rfSegmentationModel, featureMatrix,
   mask, verbose = FALSE) {
  segs <- antsImageClone(mask)
  segs[mask == 1] <- predict(rfSegmentationModel, newdata = featureMatrix )
  return(segs)
}