tests/testthat/test-models-rfdetr.R

In torchvision: Models, Datasets and Transformations for Images

context("models-rfdetr")

test_that("test for non-pretrained model_rfdetr_nano", {
  model <- model_rfdetr_nano()
  input <- torch::torch_randn(1, 3, 384, 384)
  model$eval()
  out <- model(input)
  expect_true("detections" %in% names(out))
  expect_named(out$detections[[1]], c("boxes", "labels", "scores"), ignore.order = TRUE)
  expect_equal(out$detections[[1]]$boxes$shape[2], 4L)

  rm(model)
  gc()
})

test_that("test for pretrained model_rfdetr_nano", {

  skip_if(Sys.getenv("TEST_LARGE_MODELS", unset = 0) != 1,
      "Skipping test: set TEST_LARGE_MODELS=1 to enable tests requiring large downloads.")

  model <- model_rfdetr_nano(pretrained = TRUE)
  expect_coco_model_detects_cat(model, size = c(384, 384))

  rm(model)
  gc()
})

test_that("test for non-pretrained model_rfdetr_small", {
  model <- model_rfdetr_small()
  input <- torch::torch_randn(1, 3, 512, 512)
  model$eval()
  out <- model(input)
  expect_true("detections" %in% names(out))
  expect_named(out$detections[[1]], c("boxes", "labels", "scores"), ignore.order = TRUE)
  expect_equal(out$detections[[1]]$boxes$shape[2], 4L)

  rm(model)
  gc()
})

test_that("test for pretrained model_rfdetr_small", {

  skip_if(Sys.getenv("TEST_LARGE_MODELS", unset = 0) != 1,
      "Skipping test: set TEST_LARGE_MODELS=1 to enable tests requiring large downloads.")

  expect_warning(model <- model_rfdetr_small(pretrained = TRUE),
                 "model tensors not found in the checkpoint")
  # actually fails, detects a 'sink' with low confidence
  expect_coco_model_detects_cat(model, size = c(512, 512))

  input <- base_loader("assets/class/dog/dog.2.jpg") %>%
    transform_to_tensor() %>%
    transform_resize(c(512, 512)) %>%
    transform_normalize(mean = c(0.485, 0.456, 0.406), std = c(0.229, 0.224, 0.225)) %>%
    torch::torch_unsqueeze(1)
  model$eval()
  with_no_grad({
    out <- model(input)
  })
  expect_named(out, "detections")
  expect_named(out$detections[[1]], c("boxes", "labels", "scores"), ignore.order = TRUE)
  expect_bbox_is_xyxy(out$detections[[1]]$boxes, c(512, 512))


  rm(model)
  gc()
})

test_that("test for non-pretrained model_rfdetr_medium", {
  model <- model_rfdetr_medium()
  input <- torch::torch_randn(1, 3, 576, 576)
  model$eval()
  out <- model(input)
  expect_true("detections" %in% names(out))
  expect_named(out$detections[[1]], c("boxes", "labels", "scores"), ignore.order = TRUE)
  expect_equal(out$detections[[1]]$boxes$shape[2], 4L)

  rm(model)
  gc()
})

test_that("test for pretrained model_rfdetr_medium", {

  skip_if(Sys.getenv("TEST_LARGE_MODELS", unset = 0) != 1,
      "Skipping test: set TEST_LARGE_MODELS=1 to enable tests requiring large downloads.")

  expect_warning(model <- model_rfdetr_medium(pretrained = TRUE),
                 "model tensors not found in the checkpoint")

  # actually fails, detects a 'bed' with low confidence
  expect_coco_model_detects_cat(model, size = c(640, 640))

  input <- base_loader("assets/class/dog/dog.3.jpg") %>%
    transform_to_tensor() %>%
    transform_resize(c(640, 640)) %>%
    transform_normalize(mean = c(0.485, 0.456, 0.406), std = c(0.229, 0.224, 0.225)) %>%
    torch::torch_unsqueeze(1)
  model$eval()
  with_no_grad({
    out <- model(input)
  })
  expect_named(out, "detections")
  expect_named(out$detections[[1]], c("boxes", "labels", "scores"), ignore.order = TRUE)
  expect_bbox_is_xyxy(out$detections[[1]]$boxes, c(640, 640))

  rm(model)
  gc()
})

test_that("test for non-pretrained model_rfdetr_base", {
  model <- model_rfdetr_base()
  input <- torch::torch_randn(1, 3, 560, 560)
  model$eval()
  out <- model(input)
  expect_true("detections" %in% names(out))
  expect_named(out$detections[[1]], c("boxes", "labels", "scores"), ignore.order = TRUE)
  expect_equal(out$detections[[1]]$boxes$shape[2], 4L)

  rm(model)
  gc()
})

test_that("test for pretrained model_rfdetr_base", {

  skip_if(Sys.getenv("TEST_LARGE_MODELS", unset = 0) != 1,
      "Skipping test: set TEST_LARGE_MODELS=1 to enable tests requiring large downloads.")

  model <- model_rfdetr_base(pretrained = TRUE)
  expect_coco_model_detects_cat(model, size = c(640, 640))

  rm(model)
  gc()
})

test_that("test for non-pretrained model_rfdetr_base_2", {
  model <- model_rfdetr_base_2()
  input <- torch::torch_randn(1, 3, 560, 560)
  model$eval()
  out <- model(input)
  expect_true("detections" %in% names(out))
  expect_named(out$detections[[1]], c("boxes", "labels", "scores"), ignore.order = TRUE)
  expect_equal(out$detections[[1]]$boxes$shape[2], 4L)

  rm(model)
  gc()
})

test_that("test for pretrained model_rfdetr_base_2", {

  skip_if(Sys.getenv("TEST_LARGE_MODELS", unset = 0) != 1,
      "Skipping test: set TEST_LARGE_MODELS=1 to enable tests requiring large downloads.")

  model <- model_rfdetr_base_2(pretrained = TRUE)
  expect_coco_model_detects_cat(model, size = c(640, 640))

  rm(model)
  gc()
})

test_that("test for non-pretrained model_rfdetr_base_o365", {
  model <- model_rfdetr_base_o365()
  input <- torch::torch_randn(1, 3, 560, 560)
  model$eval()
  out <- model(input)
  expect_true("detections" %in% names(out))
  expect_named(out$detections[[1]], c("boxes", "labels", "scores"), ignore.order = TRUE)
  expect_equal(out$detections[[1]]$boxes$shape[2], 4L)

  rm(model)
  gc()
})

test_that("test for pretrained model_rfdetr_base_o365", {

  skip_if(Sys.getenv("TEST_LARGE_MODELS", unset = 0) != 1,
      "Skipping test: set TEST_LARGE_MODELS=1 to enable tests requiring large downloads.")

  model <- model_rfdetr_base_o365(pretrained = TRUE)

  input <- base_loader("assets/class/dog/dog.4.jpg") %>%
    transform_to_tensor() %>%
    transform_resize(c(640, 640)) %>%
    transform_normalize(mean = c(0.485, 0.456, 0.406), std = c(0.229, 0.224, 0.225)) %>%
    torch::torch_unsqueeze(1)
  model$eval()
  with_no_grad({
    out <- model(input)
  })
  expect_named(out, "detections")
  expect_bbox_is_xyxy(out$detections[[1]]$boxes, c(560, 560))

  # Expect model to detect a dog (class id 93 in object 365) among top predictions
  top_labels <- as.integer(out$detections[[1]]$labels$cpu())
  expect_true(93L %in% top_labels[1:5]) # Check that dog is in top 5
  rm(model)
  gc()
})

test_that("test for non-pretrained model_rfdetr_large", {
  model <- model_rfdetr_large()
  input <- torch::torch_randn(1, 3, 560, 560)
  model$eval()
  out <- model(input)
  expect_true("detections" %in% names(out))
  expect_named(out$detections[[1]], c("boxes", "labels", "scores"), ignore.order = TRUE)
  expect_equal(out$detections[[1]]$boxes$shape[2], 4L)

  rm(model)
  gc()
})

test_that("test for pretrained model_rfdetr_large", {

  skip_if(Sys.getenv("TEST_LARGE_MODELS", unset = 0) != 1,
      "Skipping test: set TEST_LARGE_MODELS=1 to enable tests requiring large downloads.")

  model <- model_rfdetr_large(pretrained = TRUE)
  expect_coco_model_detects_cat(model,  size = c(560, 560))

  rm(model)
  gc()
})

test_that("rfdetr_interpolate_bicubic_aa produces correct shapes", {
  x <- torch::torch_randn(1, 3, 14, 14)
  out <- rfdetr_interpolate_bicubic_aa(x, c(28, 28))
  expect_tensor_shape(out, c(1, 3, 28, 28))

  out_rect <- rfdetr_interpolate_bicubic_aa(x, c(16, 32))
  expect_tensor_shape(out_rect, c(1, 3, 16, 32))
})

test_that("gen_sineembed_for_position handles 2D and 4D tensors", {
  dim <- 128
  # 4D (cx, cy, w, h)
  pos_4d <- torch::torch_rand(2, 10, 4)
  out_4d <- gen_sineembed_for_position(pos_4d, dim)
  expect_tensor_shape(out_4d, c(2, 10, dim * 4))

  # 2D (cx, cy)
  pos_2d <- torch::torch_rand(2, 10, 2)
  out_2d <- gen_sineembed_for_position(pos_2d, dim)
  expect_tensor_shape(out_2d, c(2, 10, dim * 2))
})

test_that("rfdetr configs match the official rf-detr variants", {
  cfg <- rfdetr_configs
  expect_equal(cfg$small$dec_layers, 3)
  expect_equal(cfg$small$resolution, 512)
  expect_equal(cfg$medium$dec_layers, 4)
  expect_equal(cfg$medium$resolution, 576)
  for (v in c("base", "base_2", "base_o365", "large")) {
    # resolution must be a multiple of patch_size * num_windows
    expect_equal(cfg[[v]]$resolution, 560)
    expect_equal(cfg[[v]]$resolution %% (cfg[[v]]$patch_size * cfg[[v]]$num_windows), 0)
  }
})

test_that("rfdetr encoder proposals follow the row-major memory layout", {
  shapes <- matrix(c(2L, 3L), ncol = 2)
  memory <- torch::torch_zeros(2, 6, 4)
  props <- gen_encoder_output_proposals(memory, NULL, shapes, unsigmoid = FALSE)[[2]]
  expect_equal(props$shape, c(2, 6, 4))
  # flattened position k = row * w + col -> (x, y) = ((col + 0.5) / w, (row + 0.5) / h)
  expected_x <- rep((0:2 + 0.5) / 3, times = 2)
  expected_y <- rep((0:1 + 0.5) / 2, each = 3)
  expect_equal(as.numeric(props[2, , 1]), expected_x, tolerance = 1e-6)
  expect_equal(as.numeric(props[2, , 2]), expected_y, tolerance = 1e-6)

  # padded columns: the valid width (2 of 4) is used to normalise x
  shapes <- matrix(c(2L, 4L), ncol = 2)
  pad <- torch::torch_zeros(1, 2, 4, dtype = torch::torch_bool())
  pad[, , 3:4] <- TRUE
  props <- gen_encoder_output_proposals(torch::torch_zeros(1, 8, 4), pad$view(c(1, 8)), shapes,
                                        unsigmoid = FALSE)[[2]]
  expect_equal(as.numeric(props[1, 1:2, 1]), c(0.25, 0.75), tolerance = 1e-6)
})

test_that("model_rfdetr_nano works with a padding mask and non-square batches", {
  torch::torch_manual_seed(1)
  model <- model_rfdetr_nano()
  model$eval()
  x <- torch::torch_randn(2, 3, 240, 320)
  mask <- torch::torch_zeros(2, 240, 320, dtype = torch::torch_bool())
  mask[2, , 161:320] <- TRUE
  out <- torch::with_no_grad(model(x, mask = mask))
  expect_length(out$detections, 2)
  b1 <- out$detections[[1]]$boxes
  b2 <- out$detections[[2]]$boxes
  expect_true(b1[, c(1, 3)]$max()$item() <= 320)
  expect_true(b1[, c(2, 4)]$max()$item() <= 240)
  # image 2 only has 160 valid columns
  expect_true(b2[, c(1, 3)]$max()$item() <= 160)
  expect_true(b2[, c(2, 4)]$max()$item() <= 240)

  rm(model)
  gc()
})

test_that("model_rfdetr_nano respects custom target_sizes in eval mode", {
  model <- model_rfdetr_nano()
  model$eval()

  x <- torch::torch_randn(1, 3, 384, 384)
  # Force output boxes to be scaled to a 480x640 image space
  target_sizes <- torch::torch_tensor(matrix(c(480, 640), nrow = 1), dtype = torch::torch_int64())

  out <- torch::with_no_grad(model(x, target_sizes = target_sizes))
  boxes <- out$detections[[1]]$boxes

  expect_true(boxes[, 3]$max()$item() <= 640)
  expect_true(boxes[, 4]$max()$item() <= 480)

  rm(model)
  gc()
})



test_that("model_rfdetr_nano forward pass in training mode returns correct dict", {
  model <- model_rfdetr_nano()
  model$train()

  input <- torch::torch_randn(2, 3, 384, 384)
  out <- model(input)

  # Check required keys for DETR loss computation
  expect_true("pred_logits" %in% names(out))
  expect_true("pred_boxes" %in% names(out))
  expect_true("aux_outputs" %in% names(out))

  # In training mode, num_queries is multiplied by group_detr (13 for nano)
  # num_queries = 300, group_detr = 13 -> 3900 queries
  expect_tensor_shape(out$pred_logits, c(2, 300 * 13, 91))
  expect_tensor_shape(out$pred_boxes, c(2, 300 * 13, 4))

  # Check aux outputs length (dec_layers - 1 for nano which has 2 layers)
  expect_length(out$aux_outputs, 1)

  rm(model); gc()
})

test_that("model_rfdetr_base forward pass in training mode with padding mask", {
  model <- model_rfdetr_base()
  model$train()

  x <- torch::torch_randn(2, 3, 560, 560)
  mask <- torch::torch_zeros(2, 560, 560, dtype = torch::torch_bool())
  mask[2, , 281:560] <- TRUE # Half padded for second image

  out <- model(x, mask = mask)

  expect_true("pred_logits" %in% names(out))
  expect_tensor_shape(out$pred_logits, c(2, 300 * 13, 91))

  rm(model); gc()
})