Add format_callback option to TextureLoader for texture conversion

Author: wlewNV

slangpy/tests/utils/test_texture_loader.py

@@ -7,7 +7,7 @@
 from dataclasses import dataclass
 
 import slangpy as spy
-from slangpy import TextureLoader, Bitmap, Format, DataStruct, FormatSupport
+from slangpy import TextureLoader, Bitmap, Format, DataStruct, FormatSupport, FormatOverride
 from slangpy.testing import helpers
 
 
@@ -292,6 +292,106 @@ def test_load_textures(device_type: spy.DeviceType):
     assert len(textures) == 2
 
 
+@dataclass
+class FormatCallbackTestCase:
+    component_type: ComponentType
+    rg_format: Format
+    y_value: float
+    a_value: float
+    dtype: npt.DTypeLike  # type: ignore
+
+
+FORMAT_CALLBACK_TEST_CASES = [
+    FormatCallbackTestCase(ComponentType.uint8, Format.rg8_unorm, 128, 200, np.uint8),
+    FormatCallbackTestCase(ComponentType.uint16, Format.rg16_unorm, 32768, 50000, np.uint16),
+    FormatCallbackTestCase(ComponentType.float32, Format.rg32_float, 0.5, 0.8, np.float32),
+]
+
+
+@pytest.mark.parametrize("device_type", helpers.DEFAULT_DEVICE_TYPES)
+@pytest.mark.parametrize("test_case", FORMAT_CALLBACK_TEST_CASES)
+def test_format_callback_ya_to_rg(device_type: spy.DeviceType, test_case: FormatCallbackTestCase):
+    """Test format_callback to load YA bitmap as RG texture instead of RGBA."""
+    device = helpers.get_device(type=device_type)
+
+    # Check if the RG format is supported
+    format_support = device.get_format_support(test_case.rg_format)
+    if not FormatSupport.shader_load in format_support:
+        pytest.skip(f"Format {test_case.rg_format} not supported as shader resource")
+
+    # Create YA bitmap
+    bitmap = Bitmap(
+        pixel_format=PixelFormat.ya,
+        component_type=test_case.component_type,
+        width=100,
+        height=50,
+    )
+
+    # Fill with test data
+    a = np.array(bitmap, copy=False)
+    a[:, :, 0] = test_case.dtype(test_case.y_value)  # Y (luminance)
+    a[:, :, 1] = test_case.dtype(test_case.a_value)  # A (alpha)
+
+    # Define callback to keep YA as RG (Y->R, A->G)
+    # Since YA and RG have the same memory layout (2 channels), we don't need
+    # to convert - just upload the YA data directly as RG texture
+    rg_format = test_case.rg_format
+
+    def ya_to_rg_callback(device, bmp):
+        if bmp.pixel_format == PixelFormat.ya:
+            # Keep as RG without conversion: Y maps to R, A maps to G
+            # convert_to=None means use bitmap data as-is
+            return FormatOverride(format=rg_format, convert_to=None)
+        return None
+
+    # Load with callback
+    loader = TextureLoader(device)
+    options = TextureLoader.Options()
+    options.load_as_normalized = True
+    options.format_callback = ya_to_rg_callback
+    texture = loader.load_texture(bitmap=bitmap, options=options)
+
+    # Verify the texture is RG format, not RGBA
+    assert texture.format == test_case.rg_format
+    assert texture.width == bitmap.width
+    assert texture.height == bitmap.height
+
+    # Verify data: R should be Y, G should be A
+    data = texture.to_numpy()
+    assert data.shape == (50, 100, 2)
+    assert np.allclose(data[:, :, 0], test_case.y_value, atol=1)  # R = Y
+    assert np.allclose(data[:, :, 1], test_case.a_value, atol=1)  # G = A
+
+
+@pytest.mark.parametrize("device_type", helpers.DEFAULT_DEVICE_TYPES)
+def test_format_callback_default_behavior(device_type: spy.DeviceType):
+    """Test that returning None from callback uses default behavior."""
+    device = helpers.get_device(type=device_type)
+
+    # Create YA bitmap
+    bitmap = Bitmap(
+        pixel_format=PixelFormat.ya,
+        component_type=ComponentType.uint8,
+        width=100,
+        height=50,
+    )
+
+    # Callback that returns None (use default)
+    def passthrough_callback(device, bmp):
+        return None
+
+    # Load with callback that returns None
+    loader = TextureLoader(device)
+    options = TextureLoader.Options()
+    options.load_as_normalized = True
+    options.load_as_srgb = False  # Disable sRGB to get rgba8_unorm
+    options.format_callback = passthrough_callback
+    texture = loader.load_texture(bitmap=bitmap, options=options)
+
+    # Should use default behavior: YA -> RGBA
+    assert texture.format == Format.rgba8_unorm
+
+
 @pytest.mark.parametrize("device_type", helpers.DEFAULT_DEVICE_TYPES)
 def test_load_texture_array(device_type: spy.DeviceType):
     device = helpers.get_device(type=device_type)

src/sgl/utils/texture_loader.cpp

@@ -38,12 +38,14 @@ struct SourceImage {
  *   8-bit RGBA bitmap with sRGB gamma will be determined as \c Format::rgba8_unorm_srgb.
  * - \c Options::load_as_normalized
  *   8/16-bit integer bitmap will be determined as normalized resource format.
+ * - \c Options::format_callback
+ *   Custom callback for handling formats not in the standard format table.
  *
  * \param bitmap Bitmap to determine format for.
  * \param options Texture loading options.
- * \return A pair containing the determined format and flag if the bitmap needs to be converted to RGBA to match the format.
+ * \return A pair containing the determined format and optional pixel format to convert bitmap to.
  */
-inline std::pair<Format, bool>
+inline std::pair<Format, std::optional<Bitmap::PixelFormat>>
 determine_texture_format(Device* device, const Bitmap* bitmap, const TextureLoader::Options& options)
 {
     SGL_ASSERT(bitmap != nullptr);
@@ -126,7 +128,7 @@ determine_texture_format(Device* device, const Bitmap* bitmap, const TextureLoad
         format_flags = FormatFlags::normalized;
 
     // Check if bitmap is RGB and we can convert to RGBA.
-    bool convert_to_rgba = false;
+    std::optional<PixelFormat> convert_to_format;
     if (options.extend_alpha && pixel_format == PixelFormat::rgb) {
         // Find if the RGB format exists, if it does check if the device supports it
         bool rgb_format_supported = false;
@@ -137,15 +139,23 @@ determine_texture_format(Device* device, const Bitmap* bitmap, const TextureLoad
         bool rgba_format_supported
             = FORMAT_TABLE.find(make_key(PixelFormat::rgba, component_type, format_flags)) != FORMAT_TABLE.end();
         if (!rgb_format_supported && rgba_format_supported) {
-            convert_to_rgba = true;
+            convert_to_format = PixelFormat::rgba;
             pixel_format = PixelFormat::rgba;
         }
     }
 
-    // If format is greyscale with alpha, we must convert to rgba
+    // Handle formats not in FORMAT_TABLE (e.g., YA/greyscale+alpha).
+    // First check if a callback is provided, otherwise use default behavior.
     if (pixel_format == PixelFormat::ya) {
+        if (options.format_callback) {
+            auto override_result = options.format_callback(device, bitmap);
+            if (override_result) {
+                return {override_result->format, override_result->convert_to};
+            }
+        }
+        // Default behavior: convert YA to RGBA
         pixel_format = PixelFormat::rgba;
-        convert_to_rgba = true;
+        convert_to_format = PixelFormat::rgba;
     }
 
     // Use sRGB format if requested and supported.
@@ -155,10 +165,18 @@ determine_texture_format(Device* device, const Bitmap* bitmap, const TextureLoad
 
     // Find texture format.
     auto it = FORMAT_TABLE.find(make_key(pixel_format, component_type, format_flags));
-    if (it == FORMAT_TABLE.end())
+    if (it == FORMAT_TABLE.end()) {
+        // Format not found - try callback if provided
+        if (options.format_callback) {
+            auto override_result = options.format_callback(device, bitmap);
+            if (override_result) {
+                return {override_result->format, override_result->convert_to};
+            }
+        }
         SGL_THROW("Unsupported bitmap format: {} {}", pixel_format, component_type);
+    }
 
-    return {it->second, convert_to_rgba};
+    return {it->second, convert_to_format};
 }
 
 inline std::pair<TextureType, uint32_t> get_texture_type_and_layer_count(DDSFile::TextureType type, uint32_t array_size)
@@ -185,10 +203,10 @@ inline std::pair<TextureType, uint32_t> get_texture_type_and_layer_count(DDSFile
 
 inline SourceImage convert_bitmap(Device* device, ref<Bitmap> bitmap, const TextureLoader::Options& options)
 {
-    auto [format, convert_to_rgba] = determine_texture_format(device, bitmap, options);
+    auto [format, convert_to_format] = determine_texture_format(device, bitmap, options);
     return SourceImage{
-        .bitmap = convert_to_rgba
-            ? bitmap->convert(Bitmap::PixelFormat::rgba, bitmap->component_type(), bitmap->srgb_gamma())
+        .bitmap = convert_to_format
+            ? bitmap->convert(*convert_to_format, bitmap->component_type(), bitmap->srgb_gamma())
             : bitmap,
         .format = format,
     };

src/sgl/utils/texture_loader.h

@@ -7,11 +7,25 @@
 
 #include "sgl/core/fwd.h"
 #include "sgl/core/object.h"
+#include "sgl/core/bitmap.h"
 
 #include <filesystem>
+#include <functional>
 
 namespace sgl {
 
+struct SGL_API FormatOverride {
+    /// The texture format to use.
+    Format format;
+    /// If set, convert the bitmap to this pixel format before uploading.
+    std::optional<Bitmap::PixelFormat> convert_to;
+};
+
+/// Callback type for custom format conversion.
+/// Invoked when the standard format table doesn't have an acceptable mapping
+/// for the bitmap's pixel format (e.g., for YA/greyscale+alpha bitmaps).
+using FormatCallback = std::function<std::optional<FormatOverride>(Device* device, const Bitmap* bitmap)>;
+
 /**
  * \brief Utility class for loading textures from bitmaps and image files.
  */
@@ -35,6 +49,8 @@ class SGL_API TextureLoader : public sgl::Object {
         /// Resource usage flags for the texture.
         /// \c TextureUsage::render_target will be added automatically if \c generate_mips is true.
         TextureUsage usage{TextureUsage::shader_resource};
+        /// Optional callback for custom format conversion.
+        FormatCallback format_callback;
     };
 
     /**

src/slangpy_ext/utils/texture_loader.cpp

@@ -24,6 +24,21 @@ SGL_PY_EXPORT(utils_texture_loader)
 {
     using namespace sgl;
 
+    nb::class_<FormatOverride>(
+        m,
+        "FormatOverride",
+        "Result of a format conversion callback. Specifies the texture format to use "
+        "and optionally what pixel format to convert the bitmap to before uploading."
+    )
+        .def(nb::init<>())
+        .def(nb::init<Format, std::optional<Bitmap::PixelFormat>>(), "format"_a, "convert_to"_a = nb::none())
+        .def_rw("format", &FormatOverride::format, "The texture format to use.")
+        .def_rw(
+            "convert_to",
+            &FormatOverride::convert_to,
+            "If set, convert the bitmap to this pixel format before uploading."
+        );
+
     nb::class_<TextureLoader, Object> texture_loader(m, "TextureLoader", D(TextureLoader));
 
     nb::class_<TextureLoader::Options>(texture_loader, "Options", D(TextureLoader, Options))
@@ -44,7 +59,14 @@ SGL_PY_EXPORT(utils_texture_loader)
         .def_rw("extend_alpha", &TextureLoader::Options::extend_alpha, D(TextureLoader, Options, extend_alpha))
         .def_rw("allocate_mips", &TextureLoader::Options::allocate_mips, D(TextureLoader, Options, allocate_mips))
         .def_rw("generate_mips", &TextureLoader::Options::generate_mips, D(TextureLoader, Options, generate_mips))
-        .def_rw("usage", &TextureLoader::Options::usage);
+        .def_rw("usage", &TextureLoader::Options::usage)
+        .def_rw(
+            "format_callback",
+            &TextureLoader::Options::format_callback,
+            "Optional callback for custom format conversion. Invoked when the standard "
+            "format table doesn't have an acceptable mapping (e.g., for YA bitmaps). "
+            "Return None to use default behavior."
+        );
 
     nb::implicitly_convertible<nb::dict, TextureLoader::Options>();