R/pspaceWatershed.R

In PathwaySpace: Spatial Projection of Network Signals along Geodesic Paths

#' @title Variation of the watershed algorithm for summit detection
#'
#' @description The \code{summitWatershed} function implements a segmentation
#' strategy to identify summits within a landscape image generated by the
#' PathwaySpace package. This function is entirely coded in R, which helps
#' alleviating users from the task of loading an excessive number of
#' dependencies. Nonetheless, while this novel implementation prevents
#' the burden a 'dependency heaviness', it still requires optimization
#' as it currently exhibits slower performance compared to well-established
#' implementations such as the watershed function from the EBImage package.
#' The \code{summitWatershed} maintain a certain level of compatibility
#' with the EBImage's watershed function, and both can be used in the
#' PathwaySpace package.
#'
#' @param x A 2D-numeric array in which each point represents the coordinates
#' of a signal in a landscape image.
#' @param tolerance The minimum signal intensity of a summit (in [0,1]),
#' representing a fraction of the maximum signal intensity.
#' @param ext Radius (in pixels) for detecting neighboring objects.
#' @return A matrix with labeled summits.
#' @author Mauro Castro.
#' @seealso \code{\link{summitMapping}}
#' @examples
#' # Load a demo landscape image
#' data('gimage', package = 'PathwaySpace')
#'
#' # Scale down the image for a quicker demonstration
#' gimage <- gimage[200:300, 200:300]
#' 
#' # Check signal range
#' range(gimage, na.rm = TRUE)
#' # [1] 0 1
#'
#' # Check image
#' \donttest{
#' image(gimage)
#' }
#'
#' # Threshold the signal intensity, for example:
#' gimage[gimage < 0.5] <- 0
#'
#' # Run summit segmentation
#' gmask <- summitWatershed(x = gimage)
#'
#' # Check resulting image mask
#' \donttest{
#' image(gimage)
#' }
#'
#' @aliases summitWatershed
#' @export
#'
summitWatershed <- function(x, tolerance = 0.1, ext = 1) {
    .validate.args("numeric_mtx", "x", x)
    .validate.args("singleNumber", "tolerance", tolerance)
    .validate.args("singleInteger", "ext", ext)
    xmask <- .summitWatershed(x = x, ext = ext)
    xmask <- .mergeSummits(x = x, xmask, tolerance)
    return(xmask)
}
.summitWatershed <- function(x, ext = 1) {
    nr <- nrow(x)
    nc <- ncol(x)
    nrc <- max(nr, nc)
    x[x == 0] <- NA
    # make indices
    x.idx <- x
    x.idx[, ] <- seq_len(prod(dim(x)))
    c.idx <- arrayInd(seq_len(prod(dim(x))), dim(x), useNames = TRUE)
    c.idx <- cbind(c.idx, val = as.numeric(x), key = seq_len(nrow(c.idx)),
        label = NA)
    seeds <- c.idx[!is.na(c.idx[, "val"]), , drop = FALSE]
    seeds <- seeds[sort.list(seeds[, "val"], decreasing = TRUE), , 
        drop = FALSE]
    # make kernel masks
    r <- ext
    km1 <- .kernelMask(r)
    km1 <- t(which(km1 == 1, arr.ind = TRUE)) - (r + 1)
    a <- (pi * r^2) * max(1, nrc * 0.01)
    rr <- ceiling(sqrt(a / pi))
    km2 <- .kernelMask(rr)
    km2 <- t(which(km2 == 1, arr.ind = TRUE)) - (rr + 1)
    # start flooding
    lcount <- 1
    for (i in seq_len(nrow(seeds))) {
        lab <- c.idx[seeds[i, "key"], "label"]
        if (is.na(lab)) {
            si <- seeds[i, ]
            idx.si <- .getSidx(si, km2, nr, nc, c.idx, x.idx)
            lab <- idx.si[!is.na(idx.si[, "label"]), , drop = FALSE]
            if (nrow(lab) == 0) {
                lab <- lcount
                lcount <- lcount + 1
            } else if (nrow(lab) == 1) {
                lab <- lab[, "label"]
            } else {
                label <- lab[, "label"]
                lb <- table(label)[1] / length(label) > 0.5
                if (lb) {
                    lab <- .mode(label)
                } else {
                    lab <- label[which.max(lab[, "val"])]
                }
            }
            idx.si <- .getSidx(si, km1, nr, nc, c.idx, x.idx)
            c.idx[idx.si[, "key"], "label"] <- lab
        }
    }
    xmask <- x
    xmask[c.idx[, c("row", "col")]] <- c.idx[, "label"]
    xmask[is.na(xmask)] <- 0
    # keep masks of at least 3 pixels
    mz <- table(as.numeric(xmask))
    mz <- as.numeric(names(mz[mz < 3]))
    for (m in mz) {
        xmask[xmask == m] <- 0
    }
    xmask <- .relabel(xmask)
    xmask <- .fillCavity(xmask)
    # prune dams
    ml <- .findOutlines(xmask)
    dl <- .findDams(x, ml, xmask)$dl
    dl <- .dilatePxEdges(dl, .kernelMask(1))
    xmask[xmask == 0] <- NA
    xmask[dl > 0] <- 0
    xmask <- .expandPartitions(x, xmask)
    return(xmask)
}
.getSidx <- function(si, km, nr, nc, c.idx, x.idx, si.val = TRUE) {
    ix <- t(km + si[c(1, 2)])
    bl <- (ix[, 1] > 0 & ix[, 1] <= nr) & (ix[, 2] > 0 & ix[, 2] <= nc)
    ix <- ix[bl, , drop = FALSE]
    idx.si <- c.idx[x.idx[ix], , drop = FALSE]
    idx.si <- idx.si[!is.na(idx.si[, "val"]), , drop = FALSE]
    if (si.val) {
        bl <- (si["val"] - idx.si[, "val"]) <= 0
        idx.si <- idx.si[bl, , drop = FALSE]
    }
    idx.si
}
.mode <- function(c) {
    uc <- unique(c)
    ic <- match(c, uc)
    uc[which.max(tabulate(ic))]
}
#-------------------------------------------------------------------------------
.mergeSummits <- function(x, xmask, tolerance = 0.1) {
    if (tolerance < 0 || tolerance > 1) {
        stop("'tolerance' should be in [0,1]", call. = FALSE)
    }
    rsignal <- range(x[x > 0], na.rm = TRUE)
    ml <- .findOutlines(xmask)
    dams <- .findDams(x, ml, xmask)
    xmask <- .openDams(x, xmask, dams, ml, rsignal, tolerance)
    xmask <- .fillCavity(xmask)
    xmask <- .relabel(xmask)
    return(xmask)
}
.findDams <- function(x, ml, xmask) {
    mlabs <- sort(unique(as.numeric(xmask)))
    mlabs <- mlabs[mlabs > 0]
    xx <- mml <- dl <- array(0, dim = dim(x) + 2)
    nr <- nrow(xx)
    nc <- ncol(xx)
    mml[-c(1, nr), -c(1, nc)] <- ml
    xx[-c(1, nr), -c(1, nc)] <- x
    ij <- c(-1, 0, 1)
    km <- .kernelMask(1, shape = "square")
    km[2, 2] <- 0
    km <- which(km == 1, arr.ind = TRUE) - 2
    sts <- data.frame(neighbor.id = NA, n.signal = NA)[-1, ]
    dlstat <- lapply(mlabs, function(l) {
        sts
    })
    for (l in mlabs) {
        lp <- which(mml == l, arr.ind = TRUE)
        if (nrow(lp) > 0) {
            for (i in seq_len(nrow(lp))) {
                tp <- t(t(km) + lp[i, ])
                idx <- mml[tp]
                idx <- idx > 0 & idx != l
                if (any(idx)) {
                    tp <- tp[idx, , drop = FALSE]
                    dl[tp] <- mml[tp]
                    tpp <- data.frame(neighbor.id = mml[tp], n.signal = xx[tp])
                    dlstat[[l]] <- rbind(dlstat[[l]], tpp)
                }
            }
        }
    }
    dl <- dl[-c(1, nr), -c(1, nc)]
    dams <- list(dl = dl, dlstat = dlstat)
    return(dams)
}
.openDams <- function(x, xmask, dams, ml, rsignal, tolerance) {
    mlabs <- sort(unique(as.numeric(xmask)))
    mlabs <- mlabs[mlabs > 0]
    nlabs <- mlabs
    for (mid in mlabs) {
        if (mid %in% nlabs) {
            dlstat <- dams$dlstat[[mid]]
            if (nrow(dlstat) > 0) {
                neighbor.ids <- sort(unique(dlstat$neighbor.id))
                neighbor.ids <- neighbor.ids[neighbor.ids %in% nlabs]
                while (length(neighbor.ids) > 0) {
                    nid <- neighbor.ids[1]
                    dam.n <- dlstat[dlstat$neighbor.id == nid, ]
                    dam.signal <- max(dam.n$n.signal, na.rm = TRUE)
                    top.summit <- max(x[xmask == mid], na.rm = TRUE)
                    bot.summit <- max(x[xmask == nid], na.rm = TRUE)
                    dam.t1 <- (bot.summit - dam.signal) / 
                        (rsignal[2] - rsignal[1])
                    dam.t2 <- 1 - (dam.signal / bot.summit)
                    if (dam.t1 < tolerance || dam.t2 < tolerance) {
                        nlabs[nlabs == nid] <- mid
                        xmask[xmask == nid] <- mid
                        dlstat <- dlstat[dlstat$neighbor.id != nid, ]
                        dlstat <- rbind(dlstat, dams$dlstat[[nid]])
                        neighbor.ids <- sort(unique(dlstat$neighbor.id))
                        neighbor.ids <- neighbor.ids[neighbor.ids %in% nlabs]
                    } else {
                        neighbor.ids <- neighbor.ids[-1]
                    }
                }
            }
        }
    }
    xmask <- .relabel(xmask)
    return(xmask)
}
.expandPartitions <- function(x, xmask) {
    nr <- nrow(x)
    nc <- ncol(x)
    seeds <- xmask
    ol <- .findOutlines(seeds)
    seeds[seeds != ol & seeds > 0] <- NA
    # make indices with seeds
    x.idx <- x
    x.idx[, ] <- seq_len(prod(dim(x)))
    c.idx <- arrayInd(seq_len(prod(dim(x))), dim(x), useNames = TRUE)
    c.idx <- cbind(c.idx, val = as.numeric(x), key = seq_len(nrow(c.idx)),
        label = NA)
    c.idx <- cbind(c.idx, seed = as.numeric(seeds))
    c.idx[is.na(c.idx[, "seed"]), "val"] <- NA
    c.idx[, "label"] <- c.idx[, "seed"]
    # sort seeds by signal intensity
    seeds <- c.idx[!is.na(c.idx[, "val"]), , drop = FALSE]
    if (nrow(seeds) == 0) {
        return(xmask)
    }
    ixd <- order(seeds[, "val"], seeds[, "seed"], decreasing = TRUE)
    seeds <- seeds[ixd, , drop = FALSE]
    seeds <- cbind(seeds, keyseed = seq_len(nrow(seeds)), count = 1)
    seeds[seeds[, "seed"] == 0, "count"] <- 0
    c.idx <- cbind(c.idx, keyseed = NA)
    c.idx[seeds[, "key"], "keyseed"] <- seeds[, "keyseed"]
    km <- .kernelMask(1, "square")
    km <- t(which(km == 1, arr.ind = TRUE)) - 2
    # start expanding seeds
    while (sum(seeds[, "count"]) > 0) {
        si <- seeds[match(1, seeds[, "count"]), ]
        idx.si <- .getSidx(si, km, nr, nc, c.idx, x.idx, si.val = FALSE)
        idx.si <- idx.si[idx.si[, "label"] == 0, , drop = FALSE]
        if (nrow(idx.si) > 0) {
            c.idx[idx.si[, "key"], "label"] <- si["label"]
            seeds[idx.si[, "keyseed"], "label"] <- si["label"]
            seeds[idx.si[, "keyseed"], "count"] <- 1
        }
        seeds[si["keyseed"], "count"] <- 0
    }
    c.idx <- c.idx[!is.na(c.idx[, "label"]), , drop = FALSE]
    xmask[c.idx[, c("row", "col")]] <- c.idx[, "label"]
    xmask[is.na(xmask)] <- 0
    # keep masks of at least 3 pixels
    mz <- table(as.numeric(xmask))
    mz <- as.numeric(names(mz[mz < 3]))
    for (m in mz) { xmask[xmask == m] <- 0 }
    xmask <- .relabel(xmask)
    return(xmask)
}