library(imageseg)
library(magick)
library(tidyverse)
library(keras)
library(tensorflow)

wd_images <- system.file("images/canopy/resized",
                         package = "imageseg")
images <- loadImages(wd_images)

wd_masks <- system.file("images/canopy/masks",
                        package = "imageseg")

masks <- loadImages(wd_masks)

split_data <- function(n, frac_test = NULL, frac_val = NULL, seed) {
  set.seed(seed) # for reproducible splitting
  if(is.null(frac_test)) frac_test <- 0
  if(is.null(frac_val)) frac_val <- 0
  split_tmp <- sample(rep(1:3,
                          diff(floor(n * c(0, frac_test, (frac_test + frac_val), 1)))))
  test <- which(split_tmp == 1)
  validation <- which(split_tmp == 2)
  train <- which(split_tmp == 3)
  return(list(index_test = test,
              index_val = validation,
              index_train = train))
}

test_split <- 0.2           # 20% of images for testing
validation_split <- 0.1     # 10% of images for validation

split_tmp <- split_data(n = nrow(images$info), frac_test = test_split, frac_val = validation_split, seed = 100)

# get indices of values 1 (test), 2 (validation), 3 (training images)
test_index   <- split_tmp$index_test
val_index    <- split_tmp$index_val
train_index  <- split_tmp$index_train

images_aug <- dataAugmentation(images = images, 
                               subset = train_index,
                               rotation_angles = c(0, 90, 180, 270),
                               flip = TRUE, 
                               flop = TRUE)

masks_aug <- dataAugmentation(images = masks, 
                              subset = train_index,
                              rotation_angles = c(0, 90, 180, 270),
                              flip = TRUE, 
                              flop = TRUE)

x_train <- imagesToKerasInput(images_aug, type = "image")
y_train <- imagesToKerasInput(masks_aug, type = "mask")

x_test <- imagesToKerasInput(images_aug, type = "image", subset = test_index)
y_test <- imagesToKerasInput(masks_aug, type = "mask", subset = test_index)

x_val <- imagesToKerasInput(images_aug, type = "image", subset = val_index)
y_val <- imagesToKerasInput(masks_aug, type = "mask", subset = val_index)

model <- u_net(net_h = 256, 
               net_w = 256,  
               filters = 16)

bce_dice_loss <- imageseg:::bce_dice_loss
dice_coef     <- imageseg:::dice_coef
jaccard_index <- imageseg:::jaccard_index

model %>% compile(
  optimizer = optimizer_adam(),
  loss = bce_dice_loss,
  metrics = list(custom_metric("dice_coef", dice_coef),
                 custom_metric("jaccard_index", jaccard_index))
)

epochs <- 2
batch_size <- 8 

history <- model %>% fit(
  x_train, 
  y_train, 
  epochs = epochs, 
  batch_size = batch_size,   
  # validation_split = 0.15
  validation_data = list(x_val, y_val),
)

plot_history <- plot(history)
plot_history

scores <- model %>% evaluate(
  x_test, y_test, verbose = 0
)
print(scores)

# works with purrr 0.3.5, not with 1.0.2
out <- imageSegmentation(model, x = x_test)

√  checking for file 'C:\Users\Fiorentini\AppData\Local\Temp\RtmpeuGFE1\remotes1b64158335a6\EcoDynIZW-imageseg-388d34b/DESCRIPTION' ... 
-  preparing 'imageseg':
√  checking DESCRIPTION meta-information ... 
-  installing the package to process help pages
-  saving partial Rd database (6.7s)
   Errore in loadVignetteBuilder(pkgdir, TRUE) : 
     vignette builder 'R.rsp' not found
   Esecuzione interrotta
Error: Failed to install 'imageseg' from GitHub:
  System command 'Rcmd.exe' failed, exit status: 1, stdout & stderr were printed