Fix grayscale channel dimension (1) lost in DataLoader

NAMESPACE

@@ -6,6 +6,7 @@ S3method(draw_bounding_boxes,torch_tensor)
 S3method(draw_segmentation_masks,default)
 S3method(draw_segmentation_masks,image_with_segmentation_mask)
 S3method(draw_segmentation_masks,torch_tensor)
+S3method(get_image_size,"magick-image")
 S3method(transform_adjust_brightness,default)
 S3method(transform_adjust_brightness,torch_tensor)
 S3method(transform_adjust_contrast,default)

R/dataset-mnist.R

@@ -118,7 +118,7 @@ mnist_dataset <- dataset(
       archive <- download_and_cache(r[1], prefix = class(self)[1])
       fs::file_copy(archive, destpath)
 
-      if (!tools::md5sum(destpath) == r[2])
+      if (!(tools::md5sum(destpath) == r[2]))
         runtime_error("Corrupt file! Delete the file in {archive} and try again.")
 
     }
@@ -149,18 +149,45 @@ mnist_dataset <- dataset(
   },
 
   .getitem = function(index) {
-    x <- self$data[index, ,]
+    index <- as.integer(index)
+    x <- self$data[index, , , drop = FALSE]
     y <- self$targets[index]
 
-    if (!is.null(self$transform))
-      x <- self$transform(x)
+    if (!is.null(self$transform)) {
+      if (length(index) > 1) {
+        x_list <- lapply(seq_along(index), function(i) {
+          self$transform(x[i, , , drop = FALSE])
+        })
+        if (length(x_list) > 0 && inherits(x_list[[1]], "torch_tensor")) {
+          x <- torch::torch_stack(x_list)
+        } else {
+          x <- x_list
+        }
+      } else {
+        x <- self$transform(x)
+      }
+    }
 
-    if (!is.null(self$target_transform))
-      y <- self$target_transform(y)
+    if (!is.null(self$target_transform)) {
+      if (length(index) > 1) {
+        y_list <- lapply(as.list(y), self$target_transform)
+        if (length(y_list) > 0 && inherits(y_list[[1]], "torch_tensor")) {
+          y <- torch::torch_stack(y_list)
+        } else {
+          y <- unlist(y_list, use.names = FALSE)
+        }
+      } else {
+        y <- self$target_transform(y)
+      }
+    }
 
     list(x = x, y = y)
   },
 
+  .getbatch = function(index) {
+    self$.getitem(index)
+  },
+
   .length = function() {
     dim(self$data)[1]
   },
@@ -179,7 +206,7 @@ mnist_dataset <- dataset(
 #' @describeIn mnist_dataset Kuzushiji-MNIST cursive Japanese character dataset.
 #' @export
 kmnist_dataset <- dataset(
-  name = "kminst_dataset",
+  name = "kmnist_dataset",
   inherit = mnist_dataset,
   archive_size = "21 MB",
   resources = list(
@@ -252,55 +279,63 @@ qmnist_dataset <- dataset(
   },
 
   download = function() {
-    if (self$check_exists())
-      return(NULL)
-
     fs::dir_create(self$raw_folder)
     fs::dir_create(self$processed_folder)
 
-    cli_inform("Downloading {.cls {class(self)[[1]]}} ...")
-    for (r in self$resources[[self$split]]) {
-      filename <- basename(r[1])
-      destpath <- file.path(self$raw_folder, filename)
-
-      archive <- download_and_cache(r[1], prefix = glue::glue("qmnist-{self$split}"))
-      fs::file_copy(archive, destpath, overwrite = TRUE)
-
-      if (!tools::md5sum(destpath) == r[2])
-        runtime_error("Corrupt file! Delete the file in {archive} and try again.")
+    split_exists <- function(split_name) {
+      fs::file_exists(file.path(self$processed_folder, self$files[[split_name]]))
     }
 
-    cli_inform("Processing {.cls {class(self)[[1]]}} ...")
-
-
-    if (self$split == "train") {
-      saveRDS(
-        list(
-          read_sn3_pascalvincent(file.path(self$raw_folder, "qmnist-train-images-idx3-ubyte.gz")),
-          read_sn3_pascalvincent(file.path(self$raw_folder, "qmnist-train-labels-idx2-int.gz"))
-        ),
-        file.path(self$processed_folder, self$files$train)
-      )
-    }
-
-    if (self$split == "test") {
-      saveRDS(
-        list(
-          read_sn3_pascalvincent(file.path(self$raw_folder, "qmnist-test-images-idx3-ubyte.gz")),
-          read_sn3_pascalvincent(file.path(self$raw_folder, "qmnist-test-labels-idx2-int.gz"))
-        ),
-        file.path(self$processed_folder, self$files$test)
-      )
-    }
+    if (all(vapply(names(self$resources), split_exists, logical(1))))
+      return(NULL)
 
-    if (self$split == "nist") {
-      saveRDS(
-        list(
-          read_sn3_pascalvincent(file.path(self$raw_folder, "xnist-images-idx3-ubyte.xz")),
-          read_sn3_pascalvincent(file.path(self$raw_folder, "xnist-labels-idx2-int.xz"))
-        ),
-        file.path(self$processed_folder, self$files$nist)
-      )
+    cli_inform("Downloading {.cls {class(self)[[1]]}} ...")
+    for (split_name in names(self$resources)) {
+      if (split_exists(split_name))
+        next
+
+      for (r in self$resources[[split_name]]) {
+        filename <- basename(r[1])
+        destpath <- file.path(self$raw_folder, filename)
+
+        archive <- download_and_cache(r[1], prefix = glue::glue("qmnist-{split_name}"))
+        fs::file_copy(archive, destpath, overwrite = TRUE)
+
+        if (!(tools::md5sum(destpath) == r[2]))
+          runtime_error("Corrupt file! Delete the file in {archive} and try again.")
+      }
+
+      cli_inform("Processing {.val {split_name}} split for {.cls {class(self)[[1]]}} ...")
+
+      if (split_name == "train") {
+        saveRDS(
+          list(
+            read_sn3_pascalvincent(file.path(self$raw_folder, "qmnist-train-images-idx3-ubyte.gz")),
+            read_sn3_pascalvincent(file.path(self$raw_folder, "qmnist-train-labels-idx2-int.gz"))
+          ),
+          file.path(self$processed_folder, self$files$train)
+        )
+      }
+
+      if (split_name == "test") {
+        saveRDS(
+          list(
+            read_sn3_pascalvincent(file.path(self$raw_folder, "qmnist-test-images-idx3-ubyte.gz")),
+            read_sn3_pascalvincent(file.path(self$raw_folder, "qmnist-test-labels-idx2-int.gz"))
+          ),
+          file.path(self$processed_folder, self$files$test)
+        )
+      }
+
+      if (split_name == "nist") {
+        saveRDS(
+          list(
+            read_sn3_pascalvincent(file.path(self$raw_folder, "xnist-images-idx3-ubyte.xz")),
+            read_sn3_pascalvincent(file.path(self$raw_folder, "xnist-labels-idx2-int.xz"))
+          ),
+          file.path(self$processed_folder, self$files$nist)
+        )
+      }
     }
 
     cli_inform("Dataset {.cls {class(self)[[1]]}} downloaded and extracted successfully.")
@@ -388,7 +423,8 @@ emnist_collection <- dataset(
     self$processed_folder <- file.path(root, class(self)[1], "processed")
     self$transform <- transform
     self$target_transform <- target_transform
-    self$class <- self$classes_all_dataset[[self$dataset]]
+    self$classes <- self$classes_all_dataset[[self$dataset]]
+    self$class <- self$classes
 
     if (download) {
       cli_inform("Dataset {.val {self$dataset}} split {.val {self$split}} of {.cls {class(self)[[1]]}} (~{.emph {self$archive_size}}) will be downloaded and processed if not already available.")
@@ -414,7 +450,7 @@ emnist_collection <- dataset(
     url <- self$resources[[1]][1]
     archive <- download_and_cache(url, prefix = class(self)[1])
 
-    if (!tools::md5sum(archive) == self$resources[[1]][2])
+    if (!(tools::md5sum(archive) == self$resources[[1]][2]))
       runtime_error("Corrupt file! Delete the file in {archive} and try again.")
 
     unzip_dir <- file.path(self$raw_folder, "unzipped")
@@ -437,27 +473,57 @@ emnist_collection <- dataset(
   },
 
   .getitem = function(index) {
-
-    x <- self$data[index, , ]
+    index <- as.integer(index)
+    x <- self$data[index, , , drop = FALSE]
     y <- self$targets[index]
 
-    if (!is.null(self$transform))
-      x <- self$transform(x)
+    if (!is.null(self$transform)) {
+      if (length(index) > 1) {
+        x_list <- lapply(seq_along(index), function(i) {
+          self$transform(x[i, , , drop = FALSE])
+        })
+        if (length(x_list) > 0 && inherits(x_list[[1]], "torch_tensor")) {
+          x <- torch::torch_stack(x_list)
+        } else {
+          x <- x_list
+        }
+      } else {
+        x <- self$transform(x)
+      }
+    }
 
-    if (!is.null(self$target_transform))
-      y <- self$target_transform(y)
+    if (!is.null(self$target_transform)) {
+      if (length(index) > 1) {
+        y_list <- lapply(as.list(y), self$target_transform)
+        if (length(y_list) > 0 && inherits(y_list[[1]], "torch_tensor")) {
+          y <- torch::torch_stack(y_list)
+        } else {
+          y <- unlist(y_list, use.names = FALSE)
+        }
+      } else {
+        y <- self$target_transform(y)
+      }
+    }
 
     list(x = x, y = y)
   },
 
+  .getbatch = function(index) {
+    self$.getitem(index)
+  },
+
   .length = function() {
     dim(self$data)[1]
   }
 
 )
 
 read_sn3_pascalvincent <- function(path) {
-  x <- gzfile(path, open = "rb")
+  x <- if (grepl("\\\\.xz$", path, ignore.case = TRUE)) {
+    xzfile(path, open = "rb")
+  } else {
+    gzfile(path, open = "rb")
+  }
   on.exit({close(x)})
 
   magic <- readBin(x, endian = "big", what = integer(), n = 1)

R/transforms-array.R

@@ -5,7 +5,16 @@ transform_to_tensor.array <- function(img) {
   if (length(dim(img)) == 2)
     dim(img) <- c(dim(img), 1)
 
-  res <- torch::torch_tensor(img)$permute(c(3, 1, 2))
+  dims <- dim(img)
+  if (length(dims) != 3)
+    value_error("Expected a 2D or 3D array.")
+
+  # Support both HWC (default image arrays) and CHW (channel-first arrays).
+  if (dims[1] <= 4 && dims[3] > 4) {
+    res <- torch::torch_tensor(img)
+  } else {
+    res <- torch::torch_tensor(img)$permute(c(3, 1, 2))
+  }
 
   if (res$dtype == torch::torch_long())
     res <- res/255

R/transforms-magick.R

@@ -53,6 +53,8 @@
 
 # Utils -------------------------------------------------------------------
 
+#' @method get_image_size magick-image
+#' @export
 `get_image_size.magick-image` <- function(img) {
   info <- magick::image_info(img)
   c(info$width, info$height)

tests/testthat/test-dataset-mnist.R

@@ -11,11 +11,23 @@ test_that("tests for the mnist dataset", {
   ds <- mnist_dataset(dir, download = TRUE)
 
   i <- ds[1]
-  expect_equal(dim(i[[1]]), c(28, 28))
+  expect_equal(dim(i[[1]]), c(1, 28, 28))
   expect_equal(i[[2]], 6)
   expect_length(ds, 60000)
 
+  raw_items <- ds$.getitem(c(1, 2))
+  expect_length(raw_items, 2)
+  expect_named(raw_items, c("x", "y"))
+  expect_equal(dim(raw_items$x), c(2, 28, 28))
+  expect_equal(length(raw_items$y), 2)
+
   ds <- mnist_dataset(dir, transform = transform_to_tensor)
+  items <- ds$.getitem(c(1, 2))
+  expect_length(items, 2)
+  expect_named(items, c("x", "y"))
+  expect_tensor_shape(items$x, c(2, 1, 28, 28))
+  expect_equal(length(items$y), 2)
+
   dl <- torch::dataloader(ds, batch_size = 32)
   expect_length(dl, 1875)
   iter <- dataloader_make_iter(dl)
@@ -25,6 +37,12 @@ test_that("tests for the mnist dataset", {
   expect_true((torch_max(i[[1]]) <= 1)$item())
   expect_named(i, c("x", "y"))
 
+  # Regression: grayscale channel must be preserved for larger batches.
+  dl_128 <- torch::dataloader(ds, batch_size = 128)
+  i_128 <- dataloader_next(dataloader_make_iter(dl_128))
+  expect_tensor_shape(i_128$x, c(128, 1, 28, 28))
+  expect_tensor_shape(i_128$y, 128)
+
 })
 
 test_that("tests for the kmnist dataset", {
@@ -36,7 +54,7 @@ test_that("tests for the kmnist dataset", {
   ds <- kmnist_dataset(dir, download = TRUE)
 
   i <- ds[1]
-  expect_equal(dim(i[[1]]), c(28, 28))
+  expect_equal(dim(i[[1]]), c(1, 28, 28))
   expect_equal(i[[2]], 6)
   expect_length(ds, 60000)
 
@@ -62,7 +80,7 @@ test_that("fashion_mnist_dataset loads correctly", {
   expect_s3_class(ds, "fashion_mnist_dataset")
   expect_type(ds$.getitem(1), "list")
   expect_named(ds$.getitem(1), c("x", "y"))
-  expect_equal(dim(as.array(ds$.getitem(1)$x)), c(28, 28))
+  expect_equal(dim(as.array(ds$.getitem(1)$x)), c(1, 28, 28))
   expect_true(ds$.getitem(1)$y >= 1 && ds$.getitem(1)$y <= 10)
 
   ds2 <- fashion_mnist_dataset(dir, transform = transform_to_tensor)
@@ -96,7 +114,7 @@ test_that("tests for the emnist dataset", {
   first_item <- emnist[1]
   expect_named(first_item, c("x", "y"))
   expect_s3_class(first_item$x, "array")
-  expect_equal(dim(first_item$x), c(28,28))
+  expect_equal(dim(first_item$x), c(1, 28, 28))
   expect_equal((first_item[[2]]), 19)
 
   emnist <- emnist_collection(dir, dataset = "bymerge", download = TRUE)
@@ -157,7 +175,7 @@ test_that("tests for the qmnist dataset", {
     ds <- qmnist_dataset(dir, split = split, download = TRUE)
 
     i <- ds[1]
-    expect_equal(dim(i[[1]]), c(28, 28))
+    expect_equal(dim(i[[1]]), c(1, 28, 28))
     expect_true(i[[2]] %in% 1:10)
 
     expect_gt(length(ds), 0)