SOCR ≫ DSPA ≫ DSPA2 Topics ≫

This DSPA2 module represents Part 3 of the DSPA2 Chapter 14 (Deep Learning, Neural Networks). Learners are encouraged to first complete Part 1 of the DSPA2 Chapter 14 (Deep Learning, Neural Networks) and Part 2 of the DSPA2 Chapter 14 (Deep Learning, Neural Networks) prior to continuing with this Part 3 of the DSPA2 Chapter 14 (Deep Learning, Neural Networks) and the next Part 4 of the DSPA2 Chapter 14 (Deep Learning, Neural Networks), which cover the Torch and Tensorflow Image Pre-processing Image Classification Pipelines.

1 Torch and Tensorflow Image Pre-processing Image Classification Pipelines

We use keras to build and train a ML model and also use the tfhub package (TensorFlow Hub), which provides reusable machine learning libraries.

venv_name <- "r-tensorflow"
reticulate::use_virtualenv(virtualenv = venv_name, required = TRUE)
library(reticulate)

# load the necessary libraries
# May need some installations first, e.g., 
#    in conda%> pip install tensorflow_datasets
#    install.packages("remotes")
#    remotes::install_github("rstudio/tfds")
#    tfds::install_tfds()

library(keras)
# library(reticulate)

# Install package TFhub: https://github.com/rstudio/tfhub
# devtools::install_github("rstudio/tfhub")
library(tfhub)
# library(tfds)

# Install TFdatasets: https://cran.r-project.org/web/packages/tfdatasets/vignettes/introduction.html
# devtools::install_github("rstudio/tfdatasets")
library(tfdatasets)
library(utf8)

# specify r-reticulate or r-tensorflow python Anaconda environment
# use_condaenv("r-tensorflow")
# use_condaenv("r-reticulate", required = TRUE)
# there are many ways to "finding" your conda environments, and using the reticulate package to set them

# conda_list()[[1]][1] %>%  use_condaenv(required = TRUE)

# Check tensorflow install configuration
tensorflow::tf_config()
## TensorFlow v2.15.1 (C:\Users\IvoD\DOCUME~1\VIRTUA~1\R-TENS~1\lib\site-packages\tensorflow_hub\__init__.p)
## Python v3.10 (C:/Users/IvoD/Documents/.virtualenvs/r-tensorflow/Scripts/python.exe)
# py_module_available("tensorflow_hub")

# py_install("tensorflow_hub", pip = TRUE) # py_install("tensorflow_hub")
# py_install("tfds", pip = TRUE) # py_install("tfds")
# py_install("tensorflow_datasets", pip = TRUE)

# py_module_available("tensorflow_datasets")
# py_module_available("tfds")  # tensorflow_datasets

Below, we define a function brain_dataset() required for iterative retrieval of pytorch datasets. All such datasets need to have a method called initialize() instantiating the inventory of imaging and mask file names that will be used by the second method, .getitem(), to ingest the imaging data from these files, return input-image plus mask-target pairs, and perform data-augmentation. The parameter random_sampling=TRUE, .getitem() controls the weighted sample loading of image-mask pairs with larger in size tumors. This option is used with the training set to counter any class-label imbalances.

The training sets, but not the validation sets, use of data augmentation to ensure DCNN-invariance to specific spatiotemporal and intensity transformations. During imaging-data augmentation, the training images/masks may be flipped, resized, and rotated with specifiable probabilities for each type of augmentation transformation.

# install.packages(torch)
library(torch)
library(torchvision)

# data wrangling
library(tidyverse)
library(zeallot)   # needed for the piping function "%<-%" in "brain_dataset()"

# image processing and visualization
library(magick)
#library(cowplot)

# dataset loading
library(pins)
library(zip)

torch_manual_seed(1234)
set.seed(1234)



train_dir <- file.path(getwd(),"data","data")
valid_dir <- file.path(getwd(),"data","mri_valid")


brain_dataset <- dataset(
  name = "brain_dataset",
  # 1. Initialize
  initialize = function(img_dir, augmentation_params = NULL, random_sampling = FALSE) {
    self$images <- tibble(img = grep(list.files(img_dir, full.names = TRUE,
                                                pattern = "tif", recursive = TRUE),
                                     pattern = 'mask', invert = TRUE, value = TRUE),
                          mask = grep(list.files(img_dir, full.names = TRUE,
                                                 pattern = "tif", recursive = TRUE),
                                      pattern = 'mask',value = TRUE)
                          )
    self$slice_weights <- self$calc_slice_weights(self$images$mask)
    self$augmentation_params <- augmentation_params
    self$random_sampling <- random_sampling
  },

  # 2. Load and transform images from files into TF object elements (tensors)
  .getitem = function(i) {index <- if (self$random_sampling == TRUE) sample(1:self$.length(), 1, prob = self$slice_weights)
    else i
    img <- self$images$img[index] %>% image_read() %>% transform_to_tensor()
    mask <- self$images$mask[index] %>% image_read() %>% transform_to_tensor() %>% transform_rgb_to_grayscale() %>%
      torch_unsqueeze(1)
    img <- self$min_max_scale(img)

    if (!is.null(self$augmentation_params)) {
      scale_param <- self$augmentation_params[1]
      c(img, mask) %<-% self$resize(img, mask, scale_param)

      rot_param <- self$augmentation_params[2]
      c(img, mask) %<-% self$rotate(img, mask, rot_param)

      flip_param <- self$augmentation_params[3]
      c(img, mask) %<-% self$flip(img, mask, flip_param)
    }
    list(img = img, mask = mask)
  },

  # 3. Save the total number of imaging files
  .length = function() { nrow(self$images)  },

  # 4. Estimate 2D image weights: Bigger tumor-masks correspond to higher weights
  calc_slice_weights = function(masks) {
    weights <- map_dbl(masks, function(m) {
      img <- as.integer(magick::image_data(image_read(m), channels = "gray"))
      sum(img / 255)
    })

    sum_weights <- sum(weights)
    num_weights <- length(weights)

    weights <- weights %>% map_dbl(function(w) {
      w <- (w + sum_weights * 0.1 / num_weights) / (sum_weights * 1.1)
    })
    weights
  },

  # 5. Estimate the image intensity range (min, max)
  min_max_scale = function(x) {
    min = x$min()$item()
    max = x$max()$item()
    x$clamp_(min = min, max = max)
    x$add_(-min)$div_(max - min + 1e-5)
    x
  },

  # 6. Image tensor shape resizing (when necessary)
  resize = function(img, mask, scale_param) {
    img_size <- dim(img)[2]
    rnd_scale <- runif(1, 1 - scale_param, 1 + scale_param)
    img <- transform_resize(img, size = rnd_scale * img_size)
    mask <- transform_resize(mask, size = rnd_scale * img_size)
    diff <- dim(img)[2] - img_size
    if (diff > 0) {
      top <- ceiling(diff / 2)
      left <- ceiling(diff / 2)
      img <- transform_crop(img, top, left, img_size, img_size)
      mask <- transform_crop(mask, top, left, img_size, img_size)
    } else {
      img <- transform_pad(img,padding = -c(ceiling(diff/2),floor(diff/2),ceiling(diff/2),floor(diff/2)))
      mask <- transform_pad(mask, padding = -c(ceiling(diff/2), floor(diff/2),ceiling(diff/2),floor(diff/2)))
    }
    list(img, mask)
  },

  # 7. Rotation (if/when augmentation is requested)
  rotate = function(img, mask, rot_param) {
    rnd_rot <- runif(1, 1 - rot_param, 1 + rot_param)
    img <- transform_rotate(img, angle = rnd_rot)
    mask <- transform_rotate(mask, angle = rnd_rot)
    list(img, mask)
  },

  # 8. Flipping (if/when augmentation is requested)
  flip = function(img, mask, flip_param) {
    rnd_flip <- runif(1)
    if (rnd_flip > flip_param) {
      img <- transform_hflip(img)
      mask <- transform_hflip(mask)
    }
    list(img, mask)
  }
)

Next, using the brain_dataset() method, we actually generate the (training and validation) imaging datasets as computable objects using the raw filenames in the training and validation data-frames defined above.

train_ds <- brain_dataset(
  train_dir,
  augmentation_params = c(0.05, 15, 0.5),
  random_sampling = TRUE
)

length(train_ds)
## [1] 3317
# ~3K

valid_ds <- brain_dataset(
  valid_dir,
  augmentation_params = NULL,
  random_sampling = FALSE
)
length(valid_ds)
## [1] 612
# ~700

Let’s visualize one testing and one validation image-mask pairs using plot_ly().

# if (!require("BiocManager", quietly = TRUE))
#     install.packages("BiocManager")
#
# BiocManager::install("EBImage")

library (plotly)
rasterPlotly <- function (image, name="", hovermode = NULL) {
  myPlot <- plot_ly(type="image", z=image, name=name, hoverlabel=name, text=name,
            hovertext=name, hoverinfo="name+x+y") %>%
    layout(hovermode = hovermode, xaxis = list(hoverformat = '.1f'), yaxis = list(hoverformat = '.1f'))
  return(myPlot)
}

# Training Case 20
img_and_mask <- train_ds[20]
img <- img_and_mask[[1]]  # 3-channel image
img <- img_and_mask[[1]]$permute(c(2, 3, 1)) %>% as.array()  # 3-channel image
mask <- img_and_mask[[2]]$squeeze() %>% as.array() # tumor mask
mask <- EBImage::rgbImage(255*mask, 255*mask, 255*mask)  # manually generate a gray scale mask image
# plot_ly(z=255*bird_images[i+ (j-1)*10,,,], type="image", showscale=FALSE)

p1 <- rasterPlotly(image=255*img, name = "RGB Image", hovermode = "y unified")
p2 <- rasterPlotly(mask, name = "Tumor Mask", hovermode = "y unified")
subplot(p1,p2, shareY = TRUE) %>% layout(title="Training Case 20: 3-Channel Image (Left) & Tumor Mask (Right)")
# Validation Case 18
img_and_mask <- valid_ds[18]
img <- img_and_mask[[1]]
mask <- img_and_mask[[2]]
img <- img_and_mask[[1]]$permute(c(2, 3, 1)) %>% as.array()  # 3-channel image
mask <- img_and_mask[[2]]$squeeze() %>% as.array() # tumor mask
mask <- EBImage::rgbImage(255*mask, 255*mask, 255*mask)  # manually generate a gray scale mask image

p1 <- rasterPlotly(image=255*img, name = "RGB Image", hovermode = "y unified")
p2 <- rasterPlotly(mask, name = "Tumor Mask", hovermode = "y unified")
subplot(p1,p2, shareY = TRUE) %>% layout(title="Validation Case 18: 3-Channel Image (Left) & Tumor Mask (Right)")

Demonstrate training-data image-augmentation by using one image to generate 7 rows of 4 augmented images.

# # Use plot_ly()  & subplot graphics in interactive runnig mode
# # avoid for Rmd to HTML knit compilation!!!
# N <- 7*4
# img_and_mask <- valid_ds[77]
# img <- img_and_mask[[1]]
# mask <- img_and_mask[[2]]
# 
# imgs <- map (1:N, function(i) {
#   # spatial-scale factor
#   c(img, mask) %<-% train_ds$resize(img, mask, 0.25)
#   c(img, mask) %<-% train_ds$flip(img, mask, 0.3)
#   c(img, mask) %<-% train_ds$rotate(img, mask, 45)
# 
#   img %>%
#     transform_rgb_to_grayscale() %>%
#     as.array() %>%
#     plot_ly(z=., type="heatmap", showscale = FALSE, name=paste0("AugmImg=", i)) %>%
#     layout(showlegend=FALSE,  hovermode = "y unified") #,
#           # yaxis = list(scaleratio = 1, scaleanchor = 'x'))
# })
# 
# # imgs[[2]]
# imgs %>%
#   subplot(nrows = 7, shareX = TRUE, shareY = TRUE, which_layout=1) %>%
#   layout(title="Validation Case 77: Random Rotation/Flop/Scale Image Augmentation")
library(ggplot2)
library(dplyr)      # for bind_rows
library(tidyr)      # for pivot_longer (alternative approach)
library(purrr)      # for map_dfr

N <- 7 * 4
img_and_mask <- valid_ds[77]
img <- img_and_mask[[1]]
mask <- img_and_mask[[2]]

# Collect all augmented image data frames
all_imgs_df <- map_dfr(1:N, function(i) {
  # Apply augmentations (same as before)
  c(img, mask) %<-% train_ds$resize(img, mask, 0.25)
  c(img, mask) %<-% train_ds$flip(img, mask, 0.3)
  c(img, mask) %<-% train_ds$rotate(img, mask, 45)
  
  # Convert to grayscale matrix
  mat <- img %>%
    transform_rgb_to_grayscale() %>%
    as.array()
  
  # Build a data frame with x, y, and intensity value
  # (rows = y, columns = x; we'll reverse y later for image orientation)
  df <- expand.grid(
    x = 1:ncol(mat),
    y = 1:nrow(mat)
  )
  df$value <- as.vector(mat)
  df$aug_id <- factor(paste0("AugmImg=", i))
  
  df
})

# Static ggplot heatmap grid
ggplot(all_imgs_df, aes(x = x, y = y, fill = value)) +
  geom_raster() +
  facet_wrap(~ aug_id, nrow = 7) +
  scale_y_reverse() +                           # row 1 at top (image convention)
  scale_fill_gradient(low = "black", high = "white") +  # grayscale
  coord_fixed() +                               # square pixels
  theme_void() +                                # remove axes, grid, background
  theme(
    legend.position = "none",
    plot.title = element_text(hjust = 0.5)
  ) +
  labs(title = "Validation Case 77: Random Rotation/Flip/Scale Image Augmentation")

Instantiate training/validation data loaders using torch::dataloader(), which combines a data-set and a sampler-iterator into a single/multi-process iterator over the entire dataset.

batch_size <- 25
# train_dl <- dataloader(train_ds, batch_size)
# valid_dl <- dataloader(valid_ds, batch_size)

train_dl <- train_ds %>% dataloader(batch_size = batch_size, shuffle = TRUE)
valid_dl <- valid_ds %>% dataloader(batch_size = batch_size, shuffle = FALSE)

Next, we specify the Unet model indicating the number of “down” or encoding (analysis) depth for shrinking the input images and incrementing the number of filters, as well as specify how do we go “up” again during the decoding (synthesis) phase.

# forward(), keeps track of layer outputs seen going “down,” to be added back in going “up.”
unet <- nn_module(name="unet",
  initialize = function(channels_in = 3, n_classes = 1, depth = 5, n_filters = 6) {
    self$down_path <- nn_module_list()
    prev_channels <- channels_in
    for (i in 1:depth) {
      self$down_path$append(down_block(prev_channels, 2^(n_filters+i-1)))
      prev_channels <- 2^(n_filters+i-1)
    }

    self$up_path <- nn_module_list()

    for (i in ((depth - 1):1)) {
      self$up_path$append(up_block(prev_channels, 2^(n_filters+i-1)))
      prev_channels <- 2^(n_filters+i-1)
    }

    self$last = nn_conv2d(prev_channels, n_classes, kernel_size = 1)
  },

  forward = function(x) {
    blocks <- list()
    for (i in 1:length(self$down_path)) {
      x <- self$down_path[[i]](x)
      if (i != length(self$down_path)) {
        blocks <- c(blocks, x)
        x <- nnf_max_pool2d(x, 2)
      }
    }

    for (i in 1:length(self$up_path)) {
      x <- self$up_path[[i]](x, blocks[[length(blocks)-i+1]]$to(device=device))
    }

    torch_sigmoid(self$last(x))
  }
)

# unet utilizes `down_block` and `up_block`
# down_block delegates to its own workhorse, conv_block, and up_block “bridges” the UNET

down_block <- nn_module(
  classname="down_block",
  initialize = function(in_size, out_size) {
    self$conv_block <- conv_block(in_size, out_size)
  },

  forward = function(x) {   self$conv_block(x)  }
)

up_block <- nn_module(
  classname="up_block",
  initialize = function(in_size, out_size) {
    self$up = nn_conv_transpose2d(in_size, out_size, kernel_size=2, stride=2)
    self$conv_block = conv_block(in_size, out_size)
  },

  forward = function(x, bridge) {
    up <- self$up(x)
    torch_cat(list(up, bridge), 2) %>% self$conv_block()
  }
)

conv_block <- nn_module(
  classname="conv_block",
  initialize = function(in_size, out_size) {
    self$conv_block <- nn_sequential(
      nn_conv2d(in_size, out_size, kernel_size = 3, padding = 1),
      nn_relu(),
      nn_dropout(0.6),
      nn_conv2d(out_size, out_size, kernel_size = 3, padding = 1),
      nn_relu()
    )
  },

  forward = function(x){
    self$conv_block(x)
  }
)

Before we start the unet model training, we need to specify CPU/GPU device and optimization scheme.

# Initialize the model with appropriate CPU/GPU
device <- torch_device(if(cuda_is_available()) "cuda" else "cpu")
device <-"cpu"
model <- unet(depth = 5)$to(device = device)

# OPTIMIZATION

# DCNN model training using cross_entropy and dice_loss.
calc_dice_loss <- function(y_pred, y_true) {
  smooth <- 1
  y_pred <- y_pred$view(-1)
  y_true <- y_true$view(-1)
  intersection <- (y_pred * y_true)$sum()
  1 - ((2*intersection+smooth)/(y_pred$sum()+y_true$sum()+smooth))
}

dice_weight <- 0.3
#     learning_rate = 0.1
optimizer <- optim_sgd(model$parameters, lr = 0.1, momentum = 0.9)

On a multi-core machine, training the UNET model will require ~2-hrs per epoch. This step can be skipped if you load in a previously pre-trained model (model) that is saved (torch_save(model, "/path/model.pt")) and available for import (torch_load("/path/model.ptt")), see several *.pt model here.

Using num_epochs <- 1 # for knitting, otherwise increase to 5+

num_epochs <- 1  # for knitting, otherwise increase to 5+

scheduler <- lr_one_cycle(optimizer, max_lr = 0.1, steps_per_epoch = length(train_dl), epochs = num_epochs)

# TRAINING
train_batch <- function(b) {
  optimizer$zero_grad()
  output <- model(b[[1]]$to(device = device))
  target <- b[[2]]$to(device = device)

  bce_loss <- nnf_binary_cross_entropy(output, target)
  dice_loss <- calc_dice_loss(output, target)
  loss <-  dice_weight*dice_loss + (1-dice_weight)*bce_loss

  loss$backward()
  optimizer$step()
  scheduler$step()

  list(bce_loss$item(), dice_loss$item(), loss$item())
}

valid_batch <- function(b) {
  output <- model(b[[1]]$to(device = device))
  target <- b[[2]]$to(device = device)

  bce_loss <- nnf_binary_cross_entropy(output, target)
  dice_loss <- calc_dice_loss(output, target)
  loss <-  dice_weight * dice_loss + (1-dice_weight)*bce_loss

  list(bce_loss$item(), dice_loss$item(), loss$item())
}

Start the Unet training process (this is suspended here to enable real-time knitting of the HTML doc).

for (epoch in 1:num_epochs) {
  model$train()
  train_bce <- c()
  train_dice <- c()
  train_loss <- c()

  coro::loop(for (b in train_dl) {
    c(bce_loss, dice_loss, loss) %<-% train_batch(b)
    train_bce <- c(train_bce, bce_loss)
    train_dice <- c(train_dice, dice_loss)
    train_loss <- c(train_loss, loss)
  })

  torch_save(model, paste0(getwd(),"/model_", epoch, ".pt"))

  cat(sprintf("\nEpoch %d, training: loss:%3f, bce: %3f, dice: %3f\n",
              epoch, mean(train_loss), mean(train_bce), mean(train_dice)))

  model$eval()
  valid_bce <- c()
  valid_dice <- c()
  valid_loss <- c()

  i <- 0
  coro::loop(for (b in valid_dl) {

    i <<- i + 1
    c(bce_loss, dice_loss, loss) %<-% valid_batch(b)
    valid_bce <- c(valid_bce, bce_loss)
    valid_dice <- c(valid_dice, dice_loss)
    valid_loss <- c(valid_loss, loss)

  })

  cat(sprintf("\n Epoch %d, validation: loss: %3f, bce: %3f, dice: %3f\n",
              epoch, mean(valid_loss), mean(valid_bce), mean(valid_dice)))
}

# note that DCNN model mask prediction will be based on the model complexity!

# Epoch 1, training: loss:0.495716, bce: 0.307674, dice: 0.934480
# Epoch 1, validation: loss: 0.336556, bce: 0.070743, dice: 0.956785
#
# Epoch 2, training: loss:0.290850, bce: 0.114878, dice: 0.701452
# Epoch 2, validation: loss: 0.213417, bce: 0.046690, dice: 0.602447
#
# Epoch 3, training: loss:0.204054, bce: 0.102973, dice: 0.439909
# Epoch 3, validation: loss: 0.218209, bce: 0.051383, dice: 0.607468
#
# Epoch 4, training: loss:0.199831, bce: 0.101564, dice: 0.429123
# Epoch 4, validation: loss: 0.243201, bce: 0.068757, dice: 0.650236
#
# Epoch 5, training: loss:0.201025, bce: 0.101525, dice: 0.433192
# Epoch 5, validation: loss: 0.258278, bce: 0.077821, dice: 0.679345
#
# Epoch 6, training: loss:0.188058, bce: 0.094279, dice: 0.406877
# Epoch 6, validation: loss: 0.206198, bce: 0.038769, dice: 0.596865
#
# Epoch 7, training: loss:0.187285, bce: 0.094894, dice: 0.402866
# Epoch 7, validation: loss: 0.216684, bce: 0.050536, dice: 0.604361
#
# Epoch 8, training: loss:0.184097, bce: 0.093609, dice: 0.395237
# Epoch 8, validation: loss: 0.235675, bce: 0.060247, dice: 0.645006
#
# Epoch 9, training: loss:0.179406, bce: 0.090458, dice: 0.386951
# Epoch 9, validation: loss: 0.238004, bce: 0.055283, dice: 0.664353
#
# Epoch 10, training: loss:0.175083, bce: 0.089144, dice: 0.375608
# Epoch 10, validation: loss: 0.258206, bce: 0.080033, dice: 0.673942
#
# Epoch 11, training: loss:0.175821, bce: 0.089795, dice: 0.376549
# Epoch 11, validation: loss: 0.213895, bce: 0.037936, dice: 0.624464
#
# Epoch 12, training: loss:0.174915, bce: 0.087026, dice: 0.379987
# Epoch 12, validation: loss: 0.264673, bce: 0.083884, dice: 0.686514
#
# Epoch 13, training: loss:0.171973, bce: 0.087210, dice: 0.369754
# Epoch 13, validation: loss: 0.236396, bce: 0.065432, dice: 0.635311
#
# Epoch 14, training: loss:0.167253, bce: 0.084660, dice: 0.359970
# Epoch 14, validation: loss: 0.210641, bce: 0.047418, dice: 0.591495
#
# Epoch 15, training: loss:0.170835, bce: 0.083572, dice: 0.374449
# Epoch 15, validation: loss: 0.246905, bce: 0.067056, dice: 0.666552
#
# Epoch 16, training: loss:0.167938, bce: 0.083285, dice: 0.365460
# Epoch 16, validation: loss: 0.228558, bce: 0.057135, dice: 0.628546
#
# Epoch 17, training: loss:0.161383, bce: 0.080367, dice: 0.350421
# Epoch 17, validation: loss: 0.221012, bce: 0.055709, dice: 0.606719
#
# Epoch 18, training: loss:0.158193, bce: 0.078068, dice: 0.345150
# Epoch 18, validation: loss: 0.222881, bce: 0.052606, dice: 0.620188
#
# Epoch 19, training: loss:0.158190, bce: 0.078865, dice: 0.343280
# Epoch 19, validation: loss: 0.223159, bce: 0.054643, dice: 0.616361
#
# Epoch 20, training: loss:0.161496, bce: 0.079807, dice: 0.352102
# Epoch 20, validation: loss: 0.221487, bce: 0.051793, dice: 0.617438

# > proc.time()  # In Seconds! About 2,154,523 sec ~ 598 hours
#      user    system   elapsed
# 2154523.7  776791.3  141279.2

Evaluate the trained DCNN model using model_6.pt and a batch of \(n=10\) validation datasets (2D brain imaging scans).

# EVALUATION
# without random sampling, we'd mainly see lesion-free patches
N <- 10 # number of cases to predict the masks for

eval_ds <- brain_dataset(valid_dir, augmentation_params = NULL, random_sampling = TRUE)
eval_dl <- dataloader(eval_ds, batch_size = N)

batch <- eval_dl %>% dataloader_make_iter() %>% dataloader_next()

# Load a previously save torch model
# epoch = 1
# torch_load(model, paste0("model_", epoch, ".pt"))

# torch_save(model, "C:/Users/Dinov/Desktop/model_epoch_1.pt")

##### Requirements for loading a pre-trained model .....
# load torch packages above
# The pre-computed model files are available on Canvas:
# https://umich.instructure.com/courses/38100/files/folder/Case_Studies/36_TCGA_MRI_Segmentation_Data_Phenotypes
#  localModelsFolder <- "C:/Users/IvoD/Desktop/Ivo.dir/Research/UMichigan/Publications_Books/2023/DSPA_Springer_2nd_Edition_2023/Rmd_HTML/appendix/"

localModelsFolder <-getwd()

model <- torch_load(paste0(localModelsFolder, "/appendix/model_6.pt"))
# train_dir <- "/data/data"
# valid_dir <- "/data/mri_valid"
# brain_dataset ...
# train_ds; valid_ds
# eval_ds <- brain_dataset(valid_dir,augmentation_params=NULL,random_sampling=TRUE)
# eval_dl <- dataloader(eval_ds, batch_size = 8)
# batch <- eval_dl %>% dataloader_make_iter() %>% dataloader_next()
# device <- torch_device(if(cuda_is_available()) "cuda" else "cpu")
# calc_dice_loss
library(plotly)

imgsTruePred <- map (1:N, function(i) {
  # Get the 3 images
  img <- batch[[1]][i, .., drop = FALSE]                            # Image to predict the tumor mask for
  inferred_mask <- model(img$to(device = device))                   # Predicted MASK
  true_mask <- batch[[2]][i, .., drop = FALSE]$to(device = device)  # True manual mask delineation

  # compute/report the BCE/Dice performance metrics
  bce <- nnf_binary_cross_entropy(inferred_mask, true_mask)$to(device = "cpu") %>% as.numeric()
  dc <- calc_dice_loss(inferred_mask, true_mask)$to(device = "cpu") %>% as.numeric()
  cat(sprintf("\nSample %d, bce: %3f, dice: %3f\n", i, bce, dc))

  # extract the inferred predicted mask as a 2D image/array of probability values per voxel!
  inferred_mask <- inferred_mask$to(device = "cpu") %>% as.array() %>% .[1, 1, , ]
  # Binarize the probability tumor prediction to binary mask
  inferred_mask <- ifelse(inferred_mask > 0.48, 1, 0)  # In a real run, use "inferred_mask > 0.5, 1, 0"
  # hist(as.matrix(inferred_mask[1,1,,]))

  imgs <- img[1, 1, ,] %>% as.array() %>% as.array() %>%
    plot_ly(z=., type="heatmap", showscale = FALSE, name="Image")  %>%
      layout(showlegend=FALSE,  yaxis = list(scaleanchor = "x", scaleratio = 1))    # hovermode = "y unified"
  masks <- true_mask$to(device = "cpu")[1, 1, ,] %>% as.array() %>% as.array() %>%
    plot_ly(z=., type="heatmap", showscale = FALSE, name="True Mask") %>%
      layout(showlegend=FALSE,  yaxis = list(scaleanchor = "x", scaleratio = 1))    # hovermode = "y unified",
  predMasks <- inferred_mask %>% as.array() %>%
    plot_ly(z=., type="heatmap", showscale = FALSE, name="Unet-Derived Mask") %>%
        layout(showlegend=FALSE,  yaxis = list(scaleanchor = "x", scaleratio = 1))  # hovermode = "y unified",
  rowSubPlots <- subplot(imgs, masks, predMasks, nrows = 1, shareY = TRUE) %>%
            # , widths = c(0.3, 0.3, 0.3))
    layout(hovermode = "y unified") #, yaxis = list(scaleanchor = "x", scaleratio = 1))
})
## 
## Sample 1, bce: 0.316758, dice: 0.771990
## 
## Sample 2, bce: 0.057257, dice: 0.983731
## 
## Sample 3, bce: 0.033557, dice: 0.126262
## 
## Sample 4, bce: 0.031052, dice: 0.566501
## 
## Sample 5, bce: 0.085775, dice: 0.217008
## 
## Sample 6, bce: 0.086270, dice: 0.215835
## 
## Sample 7, bce: 0.197902, dice: 0.480702
## 
## Sample 8, bce: 0.032555, dice: 0.177114
## 
## Sample 9, bce: 0.157686, dice: 0.461385
## 
## Sample 10, bce: 0.189908, dice: 0.350955

Finally, we can assess the unet model performance on the independent validation images, report the DCNN-derived tumor masks (as images) and some quantitative measures (e.g., Binary Cross-Entropy, Dice coefficient). The figure displays the results - brain image, true mask, and DCNN-estimated mask.

# imgsTruePred[[2]]
# p1 <- imgsTruePred[c(1:2)] %>% subplot(nrows = 2)
# p2 <- imgsTruePred[c(3:4)] %>% subplot(nrows = 2)
# p3 <- imgsTruePred[c(5:6)] %>% subplot(nrows = 2)
# p4 <- imgsTruePred[c(7:8)] %>% subplot(nrows = 2)
# p5 <- imgsTruePred[c(9:10)] %>% subplot(nrows = 2)
# subplot(p1, p2, p3, p4, p5, nrows = 5) %>%  layout(title="Model Validation using N=10 Cases")

imgsTruePred %>% subplot(nrows = N) %>%  layout(title="Model Validation using N=10 Cases")

1.1 Torch-based DNN-based Image Reconstruction and Synthetic Image Generation

1.1.1 Data Reparametrization

This image-normalization (standardization) is necessary to avoid extreme differences in image intensities (e.g., sMRI, fMRI, PET, SPECT, dMDI/DTI, MRS, and other imaging modalities tend to have vastly different scales).

reparametrize <- function(mu, logvar){
    std = torch_exp(0.5*logvar)
    eps = torch_randn_like(std)
    z = mu + std * eps
    return (z)
}

1.1.2 Define the Architecture of the UNet Encoder Branch

h_dim <- c(32, 64, 128, 256, 512)

net_encoder <- nn_module(
        initialize = function(latent_size=1000){

        self$latent_size <- latent_size  # Z

        # h_dim <- c(32, 64, 128, 256, 512)
        in_channels = 3

        # Encoder

        self$encoder <- nn_sequential(
                nn_conv2d(in_channels, out_channels=h_dim[1],
                          kernel_size= 3, stride= 2, padding = 1),
                nn_batch_norm2d(h_dim[1]),
                nn_leaky_relu(),

                nn_conv2d(h_dim[1], out_channels=h_dim[2],
                          kernel_size= 3, stride= 2, padding = 1),
                nn_batch_norm2d(h_dim[2]),
                nn_leaky_relu(),

                nn_conv2d(h_dim[2], out_channels=h_dim[3],
                          kernel_size= 3, stride= 2, padding = 1),
                nn_batch_norm2d(h_dim[3]),
                nn_leaky_relu(),

                nn_conv2d(h_dim[3], out_channels=h_dim[4],
                          kernel_size= 3, stride= 2, padding = 1),
                nn_batch_norm2d(h_dim[4]),
                nn_leaky_relu(),

                nn_conv2d(h_dim[4], out_channels=h_dim[5],
                          kernel_size= 3, stride= 2, padding = 1),
                nn_batch_norm2d(h_dim[5]),
                nn_leaky_relu()
                )

        self$mu_layer <- nn_linear(h_dim[5]*64, latent_size)
        self$logvar_layer <- nn_linear(h_dim[5]*64, latent_size)
        },

        forward = function(x){

            list_output = c()
            hidden = self$encoder(x)
            hidden = torch_flatten(hidden, start_dim=2)
            mu = self$mu_layer(hidden)
            logvar = self$logvar_layer(hidden)
            z = reparametrize(mu, logvar)
            list_output = append(list_output, z)
            list_output = append(list_output, mu)
            list_output = append(list_output, logvar)

            return (list_output)
        }
)

1.1.3 Define the Architecture of the UNet Decoder Branch

net_decoder <- nn_module(
        initialize = function(latent_size=1000){

        self$latent_size <- latent_size  # Z

        self$decoder_input <- nn_linear(latent_size, h_dim[5]*64)

        hidden_dims = c(32, 64, 128, 256, 512)

        # decoder
        self$decoder <- nn_sequential(
                nn_conv_transpose2d(hidden_dims[5],
                                   hidden_dims[4],
                                   kernel_size=3,
                                   stride = 2,
                                   padding=1,
                                   output_padding=1),
                nn_batch_norm2d(hidden_dims[4]),
                nn_leaky_relu(),

                nn_conv_transpose2d(hidden_dims[4],
                                   hidden_dims[3],
                                   kernel_size=3,
                                   stride = 2,
                                   padding=1,
                                   output_padding=1),
                nn_batch_norm2d(hidden_dims[3]),
                nn_leaky_relu(),

                nn_conv_transpose2d(hidden_dims[3],
                                   hidden_dims[2],
                                   kernel_size=3,
                                   stride = 2,
                                   padding=1,
                                   output_padding=1),
                nn_batch_norm2d(hidden_dims[2]),
                nn_leaky_relu(),

                nn_conv_transpose2d(hidden_dims[2],
                                   hidden_dims[1],
                                   kernel_size=3,
                                   stride = 2,
                                   padding=1,
                                   output_padding=1),
                nn_batch_norm2d(hidden_dims[1]),
                nn_leaky_relu()
                )

        self$final_layer <- nn_sequential(
                            nn_conv_transpose2d(hidden_dims[1],
                                               hidden_dims[1],
                                               kernel_size=3,
                                               stride=2,
                                               padding=1,
                                               output_padding=1),
                            nn_batch_norm2d(hidden_dims[1]),
                            nn_leaky_relu(),
                            nn_conv2d(hidden_dims[1], out_channels= 3,
                                      kernel_size= 3, padding= 1),
                            nn_sigmoid())
        },

        forward = function(z){

            z = self$decoder_input(z)
            z = z$view(c(-1, 512, 8, 8))
            x_hat = self$final_layer(self$decoder(z))

            return (x_hat)
        }
)

1.1.4 Define the Loss Function & Optimizer

loss_function <- function(x_hat, x, mu, logvar){
    N = mu$shape[1]
    Z = mu$shape[2]
    loss = nnf_binary_cross_entropy(x_hat, x, reduction='sum') -0.5*torch_sum(logvar-mu^2-torch_exp(logvar))-0.5*Z*N
    return (loss/N)
}
1.1.4.0.1 Train DNN UNet (Encoder & Decoder Branches)

For GPU optimization of this long calculation see this GPUComputingWithR module.

model_encoder = net_encoder()
model_decoder = net_decoder()
# model = model$to(device='cuda')
optimizer <- optim_adam(model_decoder$parameters)
# optimizer <- optim_adam(model$parameters)
# train_data <- train_data / 255 # torch tensor with shape (N, C, H, W)

# torch_model$cuda()

# Epochs
capture.output(     # Capture/Suppress Pytorch output, which overwhelms the knitted HTML report
  for (epoch in 1:5) {   # could increase the number of epochs for better results
    l <- c()
    coro::loop(for (b in train_dl) {
      optimizer$zero_grad()
      gnd = b[[1]]
      list_output <- model_encoder(gnd)
      z = list_output[[1]]
      mu = list_output[[2]]
      logvar = list_output[[3]]
      output <- model_decoder(z)

      loss <- loss_function(output, gnd, mu, logvar)
      print(loss)
      loss$backward()
      optimizer$step()
      l <- c(l, loss$item())
    })

    cat(sprintf("Loss at epoch %d: %3f\n", epoch, mean(l)))
  })

# capture.output(     # Capture/Suppress Pytorch output, which overwhelms the knitted HTML report
#   for (epoch in 1:5) {   # could increase the number of epochs for better results
#     l <- c()
#     coro::loop(for (b in train_dl) {
#       # optimizer$zero_grad()
#       encoderOptimizer$zero_grad()
#       decoderOptimizer$zero_grad()
#       gnd = b[[1]]
#       list_output <- model_encoder(gnd)
#       z = list_output[[1]]
#       mu = list_output[[2]]
#       logvar = list_output[[3]]
#       output <- model_decoder(z)
#
#       loss <- loss_function(output, gnd, mu, logvar)
#       print(loss)
#       loss$backward()
#       # optimizer$step()
#       encoderOptimizer$step()
#       decoderOptimizer$step()
#       l <- c(l, loss$item())
#     })
#
#     cat(sprintf("Loss at epoch %d: %3f\n", epoch, mean(l)))
#   })

1.1.5 Obtain the DNN-derived Image Reconstructions and Random Synthetic Images

The synthetic images are derived by randomly sampling the latent space of the DNN (bottom of network) and using solely the decoding branch of the UNet.

img <- (train_ds[1][1]$img)$view(c(1, 3, 256, 256))
z = (model_encoder(img))[[1]]
reconstructed <-  model_decoder(z)
reconstructed <- (reconstructed$view(c(3, 256, 256)))$permute(c(2, 3, 1)) %>% as.array()

# Generate a random tensor of numbers in the Latent space using a uniform distribution on the interval [0,1)
# these will be used to seed the VAE decoder branch of the Unet and synthetically generate 2D brain images
z <- torch_rand(c(1, 1000))
generated_images <- model_decoder(z)
generated_images <- (generated_images$view(c(3, 256, 256)))$permute(c(2, 3, 1)) %>% as.array()

1.1.6 Result Visualization

Finally, plot examples of the DNN-derived image reconstruction and synthetic image generation.

library (plotly)
rasterPlotly <- function (image, name="", hovermode = NULL) {
  myPlot <- plot_ly(type="image", z=image, name=name, hoverlabel=name, text=name,
            hovertext=name, hoverinfo="name+x+y+z") %>%
    layout(title=name, hovermode = hovermode, xaxis = list(hoverformat = '.1f'), yaxis = list(hoverformat = '.1f'))
  return(myPlot)
}

p1 <- rasterPlotly(image=255*reconstructed, name = "DNN-reconstructed RGB Image")
# rasterPlotly(image=255*generated_images, name = "Synthetic DNN-generated RGB Image")
p2 <- rasterPlotly(image=20/(0.1+generated_images), name = "Synthetic DNN-generated RGB Image")
# subplot(p1, p2) %>% layout(title="DNN (Unet) Image Reconstruction and Synthetic Generation")
p1
p2
# # Ensure you have the 'kaleido' python package installed for this to work seamlessly
# # reticulate::py_install('python-kaleido')
#
# p1_static <- export(p1, file = "p1.png")
# p2_static <- export(p2, file = "p2.png")
#
# # To display them in the Rmd without crashing Pandoc:
# knitr::include_graphics(c("p1.png", "p2.png"))

PNG image conversion and other image processing tasks –>

PNG conversion worked, inspect one case –>

2 References

This DSPA2 module represents Part 3 of the DSPA2 Chapter 14 (Deep Learning, Neural Networks). Learners are encouraged to first complete Part 1 of the DSPA2 Chapter 14 (Deep Learning, Neural Networks) and Part 2 of the DSPA2 Chapter 14 (Deep Learning, Neural Networks) and after completing this Part 3 of the DSPA2 Chapter 14 (Deep Learning, Neural Networks), prior to continuing with the next Part 4 of the DSPA2 Chapter 14 (Deep Learning, Neural Networks), which cover the Torch and Tensorflow Image Pre-processing Image Classification Pipelines.

SOCR Resource Visitor number Web Analytics SOCR Email