R/makeChipsMultiClass.R

In geodl: Geospatial Semantic Segmentation with Torch and Terra

#' makeChipsMultiClass
#'
#' Generate image chips from images and associated raster masks for multiclass classification
#'
#' This function generates image and mask chips from an input image and
#' associated raster mask. The chips will be written into the defined directory.
#' The number of rows and columns of pixels per 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 multiclass classification.
#' For a binary classification or when only two classes are differentiated, use
#' the makeChips() function. If an irregular shaped raster grid is provided, only
#' chips and masks that contain no NA or NoDATA cells will be produced.
#'
#' Within the provided directory, image chips will be written to an "images" folder
#' and masks will be written to a "masks" 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 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. Default is FALSE.
#' @return Image and mask files written to disk in TIFF format. No R object is
#' returned.
#' @examples
#' \dontrun{
#' makeChipsMultiClass(image = "INPUT IMAGE NAME AND PATH",
#'                     mask = "INPUT RASTER MASK NAME AND PATH",
#'                     n_channels = 3,
#'                     size = 512,
#'                     stride_x = 512,
#'                     stride_y = 512,
#'                     outDir = "DIRECTORY IN WHICH TO SAVE CHIPS",
#'                     useExistingDir=FALSE)
#' }
#' @export
#' @importFrom utils stack
makeChipsMultiClass <- function(image,
                                mask,
                                n_channels = 3,
                                size = 256,
                                stride_x = 256,
                                stride_y = 256,
                                outDir,
                                useExistingDir = FALSE){

  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]
      }
      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"))
    }
  }

}