Add a PadToSquare image transform

shared/api/image_processor.cc

@@ -38,6 +38,7 @@ Operation::KernelRegistry ImageProcessor::kernel_registry_ = {
     {"Rescale", []() { return CreateKernelInstance(&Rescale::Compute); }},
     {"Normalize", []() { return CreateKernelInstance(&Normalize::Compute); }},
     {"CenterCrop", []() { return CreateKernelInstance(&CenterCrop::Compute); }},
+    {"PadToSquare", []() { return CreateKernelInstance(&PadToSquare::Compute); }},
     {"ConvertRGB", []() { return CreateKernelInstance(convert_to_rgb); }},
     {"Permute3D", []() { return CreateKernelInstance(&Permute3D::Compute); }},
     {"Phi3ImageTransform", []() { return CreateKernelInstance(phi3_hd_transform); }},

shared/api/image_transforms.hpp

@@ -453,6 +453,77 @@ struct CenterCrop {
   int64_t target_w_{224};
 };
 
+// Pads an (H, W, C) image to a (S, S, C) square with S = max(H, W). The image is pasted at row
+// (S - H) / 2 and column (S - W) / 2, rounded down, and the border is filled with `fill`, one value
+// per channel or a single value for all of them. This is the "pad to square" (LLaVA expand2square)
+// step some vision encoders are trained with before a fixed-size Resize; square images pass through.
+struct PadToSquare {
+  template <typename DictT>
+  OrtxStatus Init(const DictT& attrs) {
+    for (const auto& [key, value] : attrs) {
+      if (key == "fill") {
+        if (const auto* values = std::get_if<std::vector<int64_t>>(&value)) {
+          fill_ = *values;
+        } else if (const auto* single = std::get_if<int64_t>(&value)) {
+          fill_ = {*single};
+        } else {
+          return {kOrtxErrorInvalidArgument, "[PadToSquare]: fill must be an integer or a list of integers"};
+        }
+      } else {
+        return {kOrtxErrorInvalidArgument, "[PadToSquare]: Invalid attribute " + key};
+      }
+    }
+
+    if (fill_.empty()) {
+      return {kOrtxErrorInvalidArgument, "[PadToSquare]: fill must have at least one value"};
+    }
+    for (auto v : fill_) {
+      if (v < 0 || v > 255) {
+        return {kOrtxErrorInvalidArgument, "[PadToSquare]: fill values must be in [0, 255]"};
+      }
+    }
+
+    return {};
+  }
+
+  OrtxStatus Compute(const ortc::Tensor<uint8_t>& input, ortc::Tensor<uint8_t>& output) {
+    auto& dimensions = input.Shape();
+    if (dimensions.size() != 3ULL) {
+      return {kOrtxErrorInvalidArgument, "[PadToSquare]: Only raw image formats"};
+    }
+
+    const int64_t h = dimensions[0];
+    const int64_t w = dimensions[1];
+    const int64_t c = dimensions[2];
+    if (fill_.size() != 1 && static_cast<int64_t>(fill_.size()) != c) {
+      return {kOrtxErrorInvalidArgument, "[PadToSquare]: fill has " + std::to_string(fill_.size()) +
+                                             " values but the image has " + std::to_string(c) + " channels"};
+    }
+
+    const int64_t side = (std::max)(h, w);
+    const int64_t top = (side - h) / 2;
+    const int64_t left = (side - w) / 2;
+    auto* input_data = input.Data();
+    auto* p_output_image = output.Allocate({side, side, c});
+
+    for (int64_t j = 0; j < side; ++j) {
+      for (int64_t k = 0; k < side; ++k) {
+        for (int64_t l = 0; l < c; ++l) {
+          p_output_image[(j * side + k) * c + l] = static_cast<uint8_t>(fill_[fill_.size() == 1 ? 0 : l]);
+        }
+      }
+    }
+    for (int64_t j = 0; j < h; ++j) {
+      std::memcpy(p_output_image + ((j + top) * side + left) * c, input_data + j * w * c, static_cast<size_t>(w * c));
+    }
+
+    return {};
+  }
+
+ private:
+  std::vector<int64_t> fill_{0, 0, 0};
+};
+
 struct Permute3D {
 
   OrtxStatus Compute(const ortc::Tensor<float>& input, ortc::Tensor<float>& output) {

test/data/processor/decode_only.json

@@ -0,0 +1,25 @@
+{
+  "processor": {
+    "name": "decode_only",
+    "transforms": [
+      {
+        "operation": {
+          "name": "decode_image",
+          "type": "DecodeImage",
+          "attrs": {
+            "color_space": "RGB"
+          }
+        }
+      },
+      {
+        "operation": {
+          "name": "rescale",
+          "type": "Rescale",
+          "attrs": {
+            "rescale_factor": 1.0
+          }
+        }
+      }
+    ]
+  }
+}

test/data/processor/pad_to_square.json

@@ -0,0 +1,38 @@
+{
+  "processor": {
+    "name": "pad_to_square",
+    "transforms": [
+      {
+        "operation": {
+          "name": "decode_image",
+          "type": "DecodeImage",
+          "attrs": {
+            "color_space": "RGB"
+          }
+        }
+      },
+      {
+        "operation": {
+          "name": "pad_to_square",
+          "type": "PadToSquare",
+          "attrs": {
+            "fill": [
+              255,
+              0,
+              128
+            ]
+          }
+        }
+      },
+      {
+        "operation": {
+          "name": "rescale",
+          "type": "Rescale",
+          "attrs": {
+            "rescale_factor": 1.0
+          }
+        }
+      }
+    ]
+  }
+}

test/data/processor/pad_to_square_invalid_fill.json

@@ -0,0 +1,38 @@
+{
+  "processor": {
+    "name": "pad_to_square_invalid_fill",
+    "transforms": [
+      {
+        "operation": {
+          "name": "decode_image",
+          "type": "DecodeImage",
+          "attrs": {
+            "color_space": "RGB"
+          }
+        }
+      },
+      {
+        "operation": {
+          "name": "pad_to_square",
+          "type": "PadToSquare",
+          "attrs": {
+            "fill": [
+              0,
+              0,
+              256
+            ]
+          }
+        }
+      },
+      {
+        "operation": {
+          "name": "rescale",
+          "type": "Rescale",
+          "attrs": {
+            "rescale_factor": 1.0
+          }
+        }
+      }
+    ]
+  }
+}

test/data/processor/pad_to_square_scalar_fill.json

@@ -0,0 +1,34 @@
+{
+  "processor": {
+    "name": "pad_to_square_scalar_fill",
+    "transforms": [
+      {
+        "operation": {
+          "name": "decode_image",
+          "type": "DecodeImage",
+          "attrs": {
+            "color_space": "RGB"
+          }
+        }
+      },
+      {
+        "operation": {
+          "name": "pad_to_square",
+          "type": "PadToSquare",
+          "attrs": {
+            "fill": 64
+          }
+        }
+      },
+      {
+        "operation": {
+          "name": "rescale",
+          "type": "Rescale",
+          "attrs": {
+            "rescale_factor": 1.0
+          }
+        }
+      }
+    ]
+  }
+}

test/pp_api_test/test_processor.cc

@@ -992,3 +992,98 @@ TEST(ProcessorTest, TestPixtralImageProcessingMultiImage) {
     }
   }
 }
+
+TEST(ProcessorTest, TestPadToSquare) {
+  // Runs a processor config on one image and holds its first output, [1, H, W, C] float.
+  auto run = [](const char* config, OrtxRawImages* raw_images, OrtxObjectPtr<OrtxTensorResult>& result,
+                OrtxObjectPtr<OrtxTensor>& tensor) {
+    OrtxObjectPtr<OrtxProcessor> processor;
+    ASSERT_EQ(OrtxCreateProcessor(processor.ToBeAssigned(), config), kOrtxOK) << OrtxGetLastErrorMessage();
+    ASSERT_EQ(OrtxImagePreProcess(processor.get(), raw_images, result.ToBeAssigned()), kOrtxOK)
+        << OrtxGetLastErrorMessage();
+    ASSERT_EQ(OrtxTensorResultGetAt(result.get(), 0, tensor.ToBeAssigned()), kOrtxOK);
+  };
+
+  struct Case {
+    const char* image_path;
+    const char* config;
+    std::vector<float> fill;
+    int64_t top, bottom, left, right;  // expected border widths
+  };
+  // exceltable.png is 487x206 (wide), passport.png is 600x805 (tall), standard_s.jpg is 1344x1344 (square).
+  // Both odd differences put the extra border row or column at the bottom or on the right.
+  const Case cases[] = {
+      {"data/processor/exceltable.png", "data/processor/pad_to_square.json", {255, 0, 128}, 140, 141, 0, 0},
+      {"data/processor/passport.png", "data/processor/pad_to_square.json", {255, 0, 128}, 0, 0, 102, 103},
+      {"data/processor/standard_s.jpg", "data/processor/pad_to_square.json", {255, 0, 128}, 0, 0, 0, 0},
+      {"data/processor/passport.png", "data/processor/pad_to_square_scalar_fill.json", {64, 64, 64}, 0, 0, 102, 103},
+  };
+  for (const auto& test_case : cases) {
+    const char* paths[] = {test_case.image_path};
+    OrtxObjectPtr<OrtxRawImages> raw_images{};
+    ASSERT_EQ(OrtxLoadImages(raw_images.ToBeAssigned(), paths, 1, nullptr), kOrtxOK);
+
+    OrtxObjectPtr<OrtxTensorResult> decoded_result, padded_result;
+    OrtxObjectPtr<OrtxTensor> decoded_tensor, padded_tensor;
+    run("data/processor/decode_only.json", raw_images.get(), decoded_result, decoded_tensor);
+    run(test_case.config, raw_images.get(), padded_result, padded_tensor);
+
+    const float *decoded{}, *padded{};
+    const int64_t *decoded_shape{}, *padded_shape{};
+    size_t decoded_dims{}, padded_dims{};
+    ASSERT_EQ(OrtxGetTensorData(decoded_tensor.get(), reinterpret_cast<const void**>(&decoded), &decoded_shape,
+                                &decoded_dims),
+              kOrtxOK);
+    ASSERT_EQ(
+        OrtxGetTensorData(padded_tensor.get(), reinterpret_cast<const void**>(&padded), &padded_shape, &padded_dims),
+        kOrtxOK);
+    ASSERT_EQ(decoded_dims, 4);  // [1, H, W, C]
+    ASSERT_EQ(padded_dims, 4);
+
+    const int64_t h = decoded_shape[1], w = decoded_shape[2], c = decoded_shape[3];
+    const int64_t side = std::max(h, w);
+    ASSERT_EQ(c, static_cast<int64_t>(test_case.fill.size()));
+    ASSERT_EQ(padded_shape[1], side);
+    ASSERT_EQ(padded_shape[2], side);
+    ASSERT_EQ(padded_shape[3], c);
+
+    const int64_t top = (side - h) / 2, left = (side - w) / 2;
+    ASSERT_EQ(top, test_case.top);
+    ASSERT_EQ(side - h - top, test_case.bottom);
+    ASSERT_EQ(left, test_case.left);
+    ASSERT_EQ(side - w - left, test_case.right);
+
+    auto pixel = [&](const float* data, int64_t width, int64_t y, int64_t x, int64_t k) {
+      return data[(y * width + x) * c + k];
+    };
+    if (left > 0) {
+      // The last left border column, both image edges, the first right border column and the extra one.
+      for (int64_t y = 0; y < side; ++y) {
+        for (int64_t k = 0; k < c; ++k) {
+          ASSERT_EQ(pixel(padded, side, y, 0, k), test_case.fill[k]) << "y=" << y;
+          ASSERT_EQ(pixel(padded, side, y, left - 1, k), test_case.fill[k]) << "y=" << y;
+          ASSERT_EQ(pixel(padded, side, y, left, k), pixel(decoded, w, y, 0, k)) << "y=" << y;
+          ASSERT_EQ(pixel(padded, side, y, left + w - 1, k), pixel(decoded, w, y, w - 1, k)) << "y=" << y;
+          ASSERT_EQ(pixel(padded, side, y, left + w, k), test_case.fill[k]) << "y=" << y;
+          ASSERT_EQ(pixel(padded, side, y, side - 1, k), test_case.fill[k]) << "y=" << y;
+        }
+      }
+    }
+
+    for (int64_t y = 0; y < side; ++y) {
+      for (int64_t x = 0; x < side; ++x) {
+        const bool inside = y >= top && y < top + h && x >= left && x < left + w;
+        for (int64_t k = 0; k < c; ++k) {
+          const float expected = inside ? decoded[((y - top) * w + (x - left)) * c + k] : test_case.fill[k];
+          ASSERT_EQ(padded[(y * side + x) * c + k], expected)
+              << test_case.image_path << " " << test_case.config << " y=" << y << " x=" << x << " c=" << k;
+        }
+      }
+    }
+  }
+
+  OrtxObjectPtr<OrtxProcessor> processor;
+  ASSERT_NE(OrtxCreateProcessor(processor.ToBeAssigned(), "data/processor/pad_to_square_invalid_fill.json"), kOrtxOK);
+  EXPECT_NE(std::string(OrtxGetLastErrorMessage()).find("fill values must be in [0, 255]"), std::string::npos)
+      << OrtxGetLastErrorMessage();
+}