R/makeChips.R

In geodl: Geospatial Semantic Segmentation with Torch and Terra

#' makeChips
#'
#' Generate image chips from images and associated raster masks
#'
#' This function generates image and mask chips from an input image and
#' associated raster mask. The chips are written into the defined directory.
#' The number of rows and columns of pixels in each chip are equal to the
#' size argument. If a stride_x and/or stride_y is used that is different from
#' the size argument, resulting chips will either overlap or have gaps between
#' them. In order to not have overlap or gaps, the stride_x and stride_y arguments
#' should be the same as the size argument. Both the image chips and associated
#' masks are written to TIFF format (".tif"). Input data are not limited to three
#' band images. This function is specifically for a binary classification where
#' the positive case is indicated with a cell value of 1 and the background or
#' negative case is indicated with a cell value of 0. If an irregular shaped raster
#' grid is provided, only chips and masks that contain no NA or NoDATA cells will
#' be produced.
#'
#' Three modes are available. If "All" is used, all image chips are
#' generated even if they do not contain pixels mapped to the positive case.
#' Within the provided directory, image chips will be written to an "images" folder
#' and masks will be written to a "masks" folder. If "Positive" is used, only chips
#' that have at least 1 pixel mapped to the positive class will be produced. Background-
#' only chips will not be generated. Within the provided directory, image chips will
#' be written to an "images" folder and masks will be written to a "masks" folder.
#' Lastly, if the "Divided" method is used, separate "positive" and "background"
#' folders will be created with "images" and "masks" subfolders. Any chip that has
#' at least 1 pixel mapped to the positive class will be written to the "positive"
#' folder while any chip having only background pixels will be written to the
#' "background" folder.
#'
#' @param image SpatRaster object or path to input image. Function will generate a SpatRaster object
#' internally. The image and mask must have the same extent, number of rows and
#' columns of pixels, cell size, and coordinate reference system.
#' @param mask SpatRaster object or path to single-band mask. Function will generate a SpatRaster
#' object internally. The image and mask must have the same extent, number of
#' rows and columns of pixels, cell size, and coordinate reference system.
#' @param n_channels Number of channels in the input image. Default is 3.
#' @param size Size of image chips as number of rows and columns of pixels.
#' Default is 256.
#' @param stride_x Stride in the x (columns) direction. Default is 256.
#' @param stride_y Stride in the y (rows) direction. Default is 256.
#' @param outDir Full or relative path to the current working directory where you
#' want to write the chips to. Subfolders in this directory will be generated by
#' the function if useExistingDir = FALSE. You must include the final forward slash
#' in the file path (e.g., "C:/data/chips/").
#' @param mode Either "All", "Positive", or "Divided". Please see the explanations
#' provided above. The default is "All".
#' @param useExistingDir TRUE or FALSE. Write chips into an existing directory
#' with subfolders already defined as opposed to using a new directory. This can be used
#' if you want to add chips to an existing set of chips. However, the "mode" should
#' be the same as that used to generated the original chips. Default is FALSE.
#' @return Image and mask files written to disk in TIFF format. No R object is returned.
#' @examples
#' \dontrun{
#' makeChips(image = "INPUT IMAGE FILE NAME AND PATH",
#'           mask = "INPUT RASTER MASK FILE NAME AND PATH",
#'           n_channels = 3,
#'           size = 256,
#'           stride_x = 256,
#'           stride_y = 256,
#'           outDir = "OUTPUT DIRECTY IN WHICH TO SAVE CHIPS",
#'           mode = "Positive",
#'           useExistingDir=FALSE)
#' }
#' @export
#' @importFrom utils stack
makeChips <- function(image,
                      mask,
                      n_channels = 3,
                      size = 256,
                      stride_x = 256,
                      stride_y = 256,
                      outDir,
                      mode = "All",
                      useExistingDir = FALSE){
  if(mode == "All"){
    img1 <- terra::rast(image)
    mask1 <- terra::rast(mask)

    fName = basename(image)

    if(useExistingDir == FALSE){
      dir.create(paste0(outDir, "/images"))
      dir.create(paste0(outDir, "/masks"))
    }

    across_cnt = terra::ncol(img1)
    down_cnt = terra::nrow(img1)
    tile_size_across = size
    tile_size_down = size
    overlap_across = stride_x
    overlap_down = stride_y
    across <- ceiling(across_cnt/overlap_across)
    down <- ceiling(down_cnt/overlap_down)
    across_add <- (across*overlap_across)-across_cnt
    across_seq <- seq(0, across-1, by=1)
    down_seq <- seq(0, down-1, by=1)
    across_seq2 <- (across_seq*overlap_across)+1
    down_seq2 <- (down_seq*overlap_down)+1

    #Loop through row/column combinations to make predictions for entire image
    for (c in across_seq2){
      for (r in down_seq2){
        c1 <- c
        r1 <- r
        c2 <- c + (size-1)
        r2 <- r + (size-1)
        if(c2 <= across_cnt && r2 <= down_cnt){ #Full chip
          chip_data <- img1[r1:r2, c1:c2, 1:n_channels]
          mask_data <- mask1[r1:r2, c1:c2, 1]
        }else if(c2 > across_cnt && r2 <= down_cnt){ # Last column
          c1b <- across_cnt - (size-1)
          c2b <- across_cnt
          chip_data <- img1[r1:r2, c1b:c2b, 1:n_channels]
          mask_data <- mask1[r1:r2, c1b:c2b, 1]
        }else if(c2 <= across_cnt && r2 > down_cnt){ #Last row
          r1b <- down_cnt - (size-1)
          r2b <- down_cnt
          chip_data <- img1[r1b:r2b, c1:c2, 1:n_channels]
          mask_data <- mask1[r1b:r2b, c1:c2, 1]
        }else{ # Last row, last column
          c1b <- across_cnt - (size -1)
          c2b <- across_cnt
          r1b <- down_cnt - (size -1)
          r2b <- down_cnt
          chip_data <- img1[r1b:r2b, c1b:c2b, 1:n_channels]
          mask_data <- mask1[r1b:r2b, c1b:c2b, 1]
        }
        if(sum(is.na(chip_data)) == 0){
          chip_data2 <- c(stack(chip_data)[,1])
          chip_array <- array(chip_data2, c(size,size,n_channels))
          image1 <- terra::rast(chip_array)
          terra::writeRaster(image1,
                             paste0(outDir,
                                    "/images/",
                                    substr(fName, 1, nchar(fName)-4), "_", c1, "_", r1, ".tif"))
          names(mask_data) <- c("C")
          Cx <- as.vector(mask_data$C)
          mask_array <- array(Cx, c(size,size,1))
          msk1 <-terra::rast(mask_array)
          terra::writeRaster(msk1,
                             paste0(outDir,
                                    "/masks/",
                                    substr(fName, 1, nchar(fName)-4), "_", c1, "_", r1, ".tif"))
        }
      }
    }
  }else if(mode == "Positive"){
    img1 <- terra::rast(image)
    mask1 <- terra::rast(mask)

    fName = basename(image)

    if(useExistingDir == FALSE){
      dir.create(paste0(outDir, "/images"))
      dir.create(paste0(outDir, "/masks"))
    }

    across_cnt = terra::ncol(img1)
    down_cnt = terra::nrow(img1)
    tile_size_across = size
    tile_size_down = size
    overlap_across = stride_x
    overlap_down = stride_y
    across <- ceiling(across_cnt/overlap_across)
    down <- ceiling(down_cnt/overlap_down)
    across_add <- (across*overlap_across)-across_cnt
    across_seq <- seq(0, across-1, by=1)
    down_seq <- seq(0, down-1, by=1)
    across_seq2 <- (across_seq*overlap_across)+1
    down_seq2 <- (down_seq*overlap_down)+1

    #Loop through row/column combinations to make predictions for entire image
    for (c in across_seq2){
      for (r in down_seq2){
        c1 <- c
        r1 <- r
        c2 <- c + (size-1)
        r2 <- r + (size-1)
        if(c2 <= across_cnt && r2 <= down_cnt){ #Full chip
          chip_data <- img1[r1:r2, c1:c2, 1:n_channels]
          mask_data <- mask1[r1:r2, c1:c2, 1]
        }else if(c2 > across_cnt && r2 <= down_cnt){ # Last column
          c1b <- across_cnt - (size-1)
          c2b <- across_cnt
          chip_data <- img1[r1:r2, c1b:c2b, 1:n_channels]
          mask_data <- mask1[r1:r2, c1b:c2b, 1]
        }else if(c2 <= across_cnt && r2 > down_cnt){ #Last row
          r1b <- down_cnt - (size-1)
          r2b <- down_cnt
          chip_data <- img1[r1b:r2b, c1:c2, 1:n_channels]
          mask_data <- mask1[r1b:r2b, c1:c2, 1]
        }else{ # Last row, last column
          c1b <- across_cnt - (size -1)
          c2b <- across_cnt
          r1b <- down_cnt - (size -1)
          r2b <- down_cnt
          chip_data <- img1[r1b:r2b, c1b:c2b, 1:n_channels]
          mask_data <- mask1[r1b:r2b, c1b:c2b, 1]
        }
        if(sum(is.na(chip_data)) == 0){
          chip_data2 <- c(stack(chip_data)[,1])
          chip_array <- array(chip_data2, c(size,size,n_channels))
          image1 <- terra::rast(chip_array)
          names(mask_data) <- c("C")
          Cx <- as.vector(mask_data$C)
          mask_array <- array(Cx, c(size,size,1))
          msk1 <-terra::rast(mask_array)
          if(max(mask_array) > 0){
            terra::writeRaster(image1,
                               paste0(outDir,
                                      "/images/",
                                      substr(fName, 1, nchar(fName)-4), "_", c1, "_", r1, ".tif"))
            terra::writeRaster(msk1,
                               paste0(outDir,
                                      "/masks/",
                                      substr(fName, 1, nchar(fName)-4), "_", c1, "_", r1, ".tif"))
          }
        }
      }
    }
  }else if(mode == "Divided") {
    img1 <- terra::rast(image)
    mask1 <- terra::rast(mask)

    fName = basename(image)

    if(useExistingDir == FALSE){

      dir.create(paste0(outDir, "/images"))
      dir.create(paste0(outDir, "/masks"))

      dir.create(paste0(outDir, "/images/positive"))
      dir.create(paste0(outDir, "/images/background"))
      dir.create(paste0(outDir, "/masks/positive"))
      dir.create(paste0(outDir, "/masks/background"))

    }

    across_cnt <- terra::ncol(img1)
    down_cnt <- terra::nrow(img1)
    tile_size_across <- size
    tile_size_down <- size
    overlap_across <- stride_x
    overlap_down <- stride_y
    across <- ceiling(across_cnt/overlap_across)
    down <- ceiling(down_cnt/overlap_down)
    across_add <- (across*overlap_across)-across_cnt
    across_seq <- seq(0, across-1, by=1)
    down_seq <- seq(0, down-1, by=1)
    across_seq2 <- (across_seq*overlap_across)+1
    down_seq2 <- (down_seq*overlap_down)+1

    #Loop through row/column combinations to make predictions for entire image
    for (c in across_seq2){
      for (r in down_seq2){
        c1 <- c
        r1 <- r
        c2 <- c + (size-1)
        r2 <- r + (size-1)
        if(c2 <= across_cnt && r2 <= down_cnt){ #Full chip
          chip_data <- img1[r1:r2, c1:c2, 1:n_channels]
          mask_data <- mask1[r1:r2, c1:c2, 1]
        }else if(c2 > across_cnt && r2 <= down_cnt){ # Last column
          c1b <- across_cnt - (size-1)
          c2b <- across_cnt
          chip_data <- img1[r1:r2, c1b:c2b, 1:n_channels]
          mask_data <- mask1[r1:r2, c1b:c2b, 1]
        }else if(c2 <= across_cnt && r2 > down_cnt){ #Last row
          r1b <- down_cnt - (size-1)
          r2b <- down_cnt
          chip_data <- img1[r1b:r2b, c1:c2, 1:n_channels]
          mask_data <- mask1[r1b:r2b, c1:c2, 1]
        }else{ # Last row, last column
          c1b <- across_cnt - (size -1)
          c2b <- across_cnt
          r1b <- down_cnt - (size -1)
          r2b <- down_cnt
          chip_data <- img1[r1b:r2b, c1b:c2b, 1:n_channels]
          mask_data <- mask1[r1b:r2b, c1b:c2b, 1]
        }
        if(sum(is.na(chip_data)) == 0){
          chip_data2 <- c(stack(chip_data)[,1])
          chip_array <- array(chip_data2, c(size,size,n_channels))
          image1 <- terra::rast(chip_array)
          names(mask_data) <- c("C")
          Cx <- as.vector(mask_data$C)
          mask_array <- array(Cx, c(size,size,1))
          msk1 <- terra::rast(mask_array)
          if(max(mask_array) > 0){
            terra::writeRaster(image1,
                               paste0(outDir,
                                      "/images/positive/",
                                      substr(fName, 1, nchar(fName)-4), "_", c1, "_", r1, ".tif"))
            terra::writeRaster(msk1,
                               paste0(outDir,
                                      "/masks/positive/",
                                      substr(fName, 1, nchar(fName)-4), "_", c1, "_", r1, ".tif"))
          }else{
            terra::writeRaster(image1,
                               paste0(outDir,
                                      "/images/background/",
                                      substr(fName, 1, nchar(fName)-4), "_", c1, "_", r1, ".tif"))
            terra::writeRaster(msk1,
                               paste0(outDir,
                                      "/masks/background/",
                                      substr(fName, 1, nchar(fName)-4), "_", c1, "_", r1, ".tif"))
          }
        }
      }
    }
  } else {
    message("Invalid Mode Provided.")
  }
}