# R Script for the Fashion NIST Example using convnet

# Deep Learning with R Cookbook

library(keras3)

# 1. We import the Fashion-MNIST dataset in our environment:

fashion <- dataset_fashion_mnist()
x_train <- fashion$train$x
y_train <- fashion$train$y
x_test <- fashion$test$x
y_test <- fashion$test$y

dim(x_train)
dim(x_test)

x_train[1,,]

paste("label of first image is:  " ,y_train[1])

label_names = c('T-shirt/top', 'Trouser', 'Pullover', 'Dress', 'Coat',  'Sandal',
                'Shirt', 'Sneaker', 'Bag', 'Ankle boot')

label_names[10]  # The y_train labels 0 to 9, so 9 is the 10th name

# Visualize images
par(mfcol=c(3,3))
par(mar=c(2,2,2,2),xaxs = "i",yaxs = "i")
for (idx in 1:9) { 
  img <- x_train[idx,,]
  img <- t(apply(img, 2, rev)) 
  image(1:28,1:28,img, main=paste(label_names[y_train[idx]+1]),xaxt = 'n',yaxt = 'n',col= gray((0:255)/255))
}

# 2. Next, we reshape the data, normalize it, and convert the target label to a binary class matrix:

# Resize the shape of inputs, this add the color channel
x_train <- array_reshape(x_train, c(nrow(x_train), 28, 28, 1))
x_test <- array_reshape(x_test, c(nrow(x_test), 28, 28, 1))

# Transform RGB values into [0,1] range
x_train <- x_train / 255
x_test <- x_test / 255

# Convert class vectors to binary class matrices, one-hot encoding of the 10 outputs
y_train <- to_categorical(y_train, 10)
y_test <- to_categorical(y_test, 10)

# 3, Once we are done with data preparation, we build, compile, and train our CNN model:

# Define model
cnn_model <- keras_model_sequential() |>
  layer_conv_2d(filters = 8, kernel_size = c(4,4), activation = 'relu',
                input_shape = c(28,28,1)) |> 
  layer_conv_2d(filters = 16, kernel_size = c(3,3), activation = 'relu') |> 
  layer_flatten() |> 
  layer_dense(units = 16, activation = 'relu') |> 
  layer_dense(units = 10, activation = 'softmax')

cnn_model |> summary()

loss_entropy <- function(y_pred, y_true) {
  loss_categorical_crossentropy(y_pred, y_true)
}

# Compile model
cnn_model |> compile(
  loss = loss_entropy,
  optimizer = optimizer_sgd(),
  metrics = c('accuracy')
)

# train the model
cnn_model |> fit(
  x_train, y_train,
  batch_size = 128,
  epochs = 5,
  validation_split = 0.2
)

# 4. Finally, we evaluate the performance of the trained model and print the evaluation metrics:

scores <- cnn_model |> evaluate(x_test, y_test, verbose = 1)

# Output metrics
paste('Test loss:', scores[[1]])
paste('Test accuracy:', scores[[2]])


# prediction
predicted_label <- cnn_model |> predict(x_test) 

str(predicted_label)

predicted_label[1, ]

sum(predicted_label[1, ])

which.max(predicted_label[1, ])

y_test[1:5,]

# Using the new keras3 operation op_argmax()

predicted_label |> 
  op_argmax(axis = -1) |>
  as.numeric()