support arbitrary channel names in exr

pgrit · pgrit · commit e53ac264bcf0 · 2026-03-07T11:26:14.000+01:00
diff --git a/Core/imageio.cpp b/Core/imageio.cpp
@@ -59,15 +59,41 @@ struct ExrChannelLayout {
     int idxG = -1;
     int idxB = -1;
     int idxA = -1;
-    int idxY = -1;
+
+    std::map<std::string, int> customChannels = {};
 
     int CountChannels() const {
         return
             (idxR >= 0 ? 1 : 0) +
             (idxG >= 0 ? 1 : 0) +
             (idxB >= 0 ? 1 : 0) +
             (idxA >= 0 ? 1 : 0) +
-            (idxY >= 0 ? 1 : 0);
+            customChannels.size();
+    }
+
+    std::string NameOf(int idx) const {
+        int rgbaOffset =
+            (idxR >= 0 ? 1 : 0) +
+            (idxG >= 0 ? 1 : 0) +
+            (idxB >= 0 ? 1 : 0) +
+            (idxA >= 0 ? 1 : 0);
+
+        if (idx < rgbaOffset) {
+            char* fixed = (char*)alloca(4);
+            int offset = 0;
+            if (idxR >= 0) fixed[offset++] = 'R';
+            if (idxG >= 0) fixed[offset++] = 'G';
+            if (idxB >= 0) fixed[offset++] = 'B';
+            if (idxA >= 0) fixed[offset++] = 'A';
+
+            return std::string { fixed[idx] };
+        }
+
+        int i = idx - rgbaOffset;
+        auto it = customChannels.cbegin();
+        for (; it != customChannels.cend(); ++it)
+            if (i-- == 0) break;
+        return it->first;
     }
 };
 
@@ -156,48 +182,43 @@ int CacheExrImage(const char* filename) {
     }
 
     // Read the number of channels in each layer
-    int freeChannels = 0;
     for (int chan = 0; chan < result.header.num_channels; ++chan) {
-        // Extract the layer name by assuming a channel name of the form "layername.R", "layername.B" and so on
-        size_t len = strlen(result.header.channels[chan].name);
-
-        std::string layerName;
-        if (len <= 2) // Plain channels without a layer will be merged into an unnamed default layer (e.g., "R", "G", and "B")
+        // Extract the layer name by assuming a channel name of the form "layername.channelname"
+        std::string name = result.header.channels[chan].name;
+        size_t len = name.size();
+
+        size_t idxSep = name.rfind('.');
+        std::string layerName, chanName;
+        if (idxSep == name.npos) {
+            // Plain channels without a layer will be merged into an unnamed default layer (e.g., "R", "G", and "B")
             layerName = "";
-        else if (result.header.channels[chan].name[len - 2] != '.') // Keep the full name as the layer name
-            layerName = std::string(result.header.channels[chan].name, len);
-        else // Remove channel name from the layer name (e.g., .R / .G / .B)
-            layerName = std::string(result.header.channels[chan].name, len - 2);
-
-        char chanName = result.header.channels[chan].name[len - 1];
+            chanName = name;
+        } else {
+            // Remove channel name from the layer name
+            layerName = name.substr(0, idxSep);
+            chanName = name.substr(idxSep + 1);
+        }
 
         // Update the channel layout info
         auto iter = result.channelsPerLayer.find(layerName);
         if (iter == result.channelsPerLayer.end()) {
             result.channelsPerLayer[layerName] = ExrChannelLayout();
             result.layerNames.emplace_back(layerName);
         }
-
         auto& layout = result.channelsPerLayer[layerName];
 
-        switch (chanName) {
-        case 'R':
+        // Handle known channel names R G B and A so we can reorder them
+        // All other channel names will be alphabethically sorted
+        if (chanName == "R")
             layout.idxR = chan;
-            break;
-        case 'G':
+        else if (chanName == "G")
             layout.idxG = chan;
-            break;
-        case 'B':
+        else if (chanName == "B")
             layout.idxB = chan;
-            break;
-        case 'A':
+        else if (chanName == "A")
             layout.idxA = chan;
-            break;
-        case 'Y':
-        default:
-            layout.idxY = chan;
-            break;
-        }
+        else if (!layout.customChannels.emplace(chanName, chan).second)
+            std::cerr << "Duplicate channel '" << chanName << "' in layer '" << layerName << "' ignored." << std::endl;
     }
 
     cacheMutex.lock();
@@ -234,44 +255,24 @@ bool CopyCachedExrLayer(int id, std::string layerName, float* out) {
     int numChannels = layerInfo.CountChannels();
 
     auto swizzle = [numChannels, &layerInfo, &layerName](unsigned char** images, int srcIdx, int dstIdx, float* out) {
-        if (numChannels == 1) { // Y
-            assert(layerInfo.idxY >= 0);
-            auto chanImg = (float*)images[layerInfo.idxY];
-            out[dstIdx + 0] = chanImg[srcIdx];
-        } else if (numChannels == 3) { // RGB
-            assert(layerInfo.idxR >= 0);
-            auto chanImg = (float*)images[layerInfo.idxR];
-            out[dstIdx + 0] = chanImg[srcIdx];
-
-            assert(layerInfo.idxG >= 0);
-            chanImg = (float*)images[layerInfo.idxG];
-            out[dstIdx + 1] = chanImg[srcIdx];
-
-            assert(layerInfo.idxB >= 0);
-            chanImg = (float*)images[layerInfo.idxB];
-            out[dstIdx + 2] = chanImg[srcIdx];
-        } else if (numChannels == 4) { // RGBA
-            assert(layerInfo.idxR >= 0);
-            auto chanImg = (float*)images[layerInfo.idxR];
-            out[dstIdx + 0] = chanImg[srcIdx];
-
-            assert(layerInfo.idxG >= 0);
-            chanImg = (float*)images[layerInfo.idxG];
-            out[dstIdx + 1] = chanImg[srcIdx];
-
-            assert(layerInfo.idxB >= 0);
-            chanImg = (float*)images[layerInfo.idxB];
-            out[dstIdx + 2] = chanImg[srcIdx];
-
-            assert(layerInfo.idxA >= 0);
-            chanImg = (float*)images[layerInfo.idxA];
-            out[dstIdx + 3] = chanImg[srcIdx];
-        } else {
-            std::cerr << "ERROR while reading .exr layer " << layerName << ": Images with "
-                    << numChannels << " channels are currently not supported. " << std::endl;
-            cacheMutex.unlock();
-            return false;
-        }
+        int offset = 0;
+        auto add = [&](int idx) {
+            auto chanImg = (float*)images[idx];
+            out[dstIdx + offset] = chanImg[srcIdx];
+            offset++;
+        };
+
+        if (layerInfo.idxR >= 0)
+            add(layerInfo.idxR);
+        if (layerInfo.idxG >= 0)
+            add(layerInfo.idxG);
+        if (layerInfo.idxB >= 0)
+            add(layerInfo.idxB);
+        if (layerInfo.idxA >= 0)
+            add(layerInfo.idxA);
+        for (auto it = layerInfo.customChannels.begin(); it != layerInfo.customChannels.end(); ++it)
+            add(it->second);
+
         return true;
     };
 
@@ -896,6 +897,20 @@ SIIO_API void GetExrLayerName(int id, int layerIdx, char* out) {
     cacheMutex.unlock();
 }
 
+SIIO_API int GetExrChannelNameLen(int id, const char* layerName, int channelIdx) {
+    cacheMutex.lock();
+    const auto& chanName = exrImages[id].channelsPerLayer[layerName].NameOf(channelIdx);
+    cacheMutex.unlock();
+    return chanName.size();
+}
+
+SIIO_API void GetExrChannelName(int id, const char* layerName, int channelIdx, char* out) {
+    cacheMutex.lock();
+    const auto& chanName = exrImages[id].channelsPerLayer[layerName].NameOf(channelIdx);
+    strcpy(out, chanName.c_str());
+    cacheMutex.unlock();
+}
+
 SIIO_API bool CopyCachedLayer(int id, const char* name, float* out) {
     return CopyCachedExrLayer(id, name, out);
 }
@@ -960,7 +975,6 @@ SIIO_API int GetExrLayerNames(const char* filename, char*** names) {
     }
 
     // Read the number of channels in each layer
-    int freeChannels = 0;
     std::unordered_set<std::string> layerNames;
     for (int chan = 0; chan < header.num_channels; ++chan) {
         // Extract the layer name by assuming a channel name of the form "layername.R", "layername.B" and so on
diff --git a/PyTest/test_io.py b/PyTest/test_io.py
@@ -204,11 +204,21 @@ def test_write_layers_read_all(self):
         ], dtype=np.float32)
 
         sio.write_layered_exr("layered.exr", {"": img, "albedo": other})
-        layers = sio.read_layered_exr("layered.exr")
+        channel_names = {}
+        layers = sio.read_layered_exr("layered.exr", channel_names)
 
         self.assertTrue("" in layers)
         self.assertTrue("albedo" in layers)
 
+        self.assertEqual(len(channel_names[""]), 3)
+        self.assertEqual(len(channel_names["albedo"]), 3)
+        self.assertEqual(channel_names["albedo"][0], "R")
+        self.assertEqual(channel_names["albedo"][1], "G")
+        self.assertEqual(channel_names["albedo"][2], "B")
+        self.assertEqual(channel_names[""][0], "R")
+        self.assertEqual(channel_names[""][1], "G")
+        self.assertEqual(channel_names[""][2], "B")
+
         self.assertEqual(other[0,0,0], layers["albedo"][0,0,0])
         self.assertEqual(other[0,0,1], layers["albedo"][0,0,1])
         self.assertEqual(other[0,0,2], layers["albedo"][0,0,2])
@@ -248,4 +258,4 @@ def test_write_layers_read_all(self):
         os.remove("layered.exr")
 
 if __name__ == "__main__":
-    unittest.main()
+    unittest.main()
diff --git a/PyWrapper/simpleimageio/image.py b/PyWrapper/simpleimageio/image.py
@@ -45,6 +45,14 @@
 _get_layer_name.argtypes = [c_int, c_int, POINTER(c_char)]
 _get_layer_name.restype = None
 
+_get_channel_name_len = corelib.core.GetExrChannelNameLen
+_get_channel_name_len.argtypes = [c_int, c_char_p, c_int]
+_get_channel_name_len.restype = c_int
+
+_get_channel_name = corelib.core.GetExrChannelName
+_get_channel_name.argtypes = [c_int, c_char_p, c_int, POINTER(c_char)]
+_get_channel_name.restype = None
+
 _copy_layer = corelib.core.CopyCachedLayer
 _copy_layer.argtypes = [c_int, c_char_p, POINTER(c_float)]
 _copy_layer.restype = None
@@ -69,7 +77,14 @@ def read(filename: str):
 
     return buffer
 
-def read_layered_exr(filename: str):
+def read_layered_exr(filename: str, layout = None):
+    '''
+    Reads all layers from an .exr file
+
+    Arguments:
+    filename -- the file to load
+    layout -- optional dictionary that will be filled with the per-layer channel name configurations
+    '''
     w = c_int()
     h = c_int()
     c = c_int()
@@ -92,6 +107,16 @@ def read_layered_exr(filename: str):
 
         _copy_layer(idx, name, buffer.ctypes.data_as(POINTER(c_float)))
 
+        # Read the channel names
+        if layout is not None:
+            channel_names = []
+            for c in range(num_chans):
+                nc = _get_channel_name_len(idx, name, c)
+                cstr_buf = create_string_buffer(nc + 1)
+                _get_channel_name(idx, name, c, cstr_buf);
+                channel_names.append(cstr_buf.value.decode('utf-8'))
+            layout[name.decode('utf-8')] = channel_names
+
         layers[name.decode('utf-8')] = buffer
 
     _delete_image(idx)
@@ -142,4 +167,4 @@ def base64_jpg(img, quality = 80):
     mem = corelib.invoke(_write_to_mem, img, ".jpg".encode('utf-8'), quality, byref(numbytes))
     b64 = base64.b64encode(bytearray(mem[:numbytes.value]))
     _free_mem(mem)
-    return b64
+    return b64
diff --git a/SimpleImageIO.Integration/Program.cs b/SimpleImageIO.Integration/Program.cs
@@ -26,4 +26,4 @@
 
 if (args.Contains("falsecolor") || args.Length == 0) {
     TestFalseColor.TestGradient();
-}
+}
diff --git a/SimpleImageIO/Image.cs b/SimpleImageIO/Image.cs
@@ -23,6 +23,12 @@ public unsafe class Image : IDisposable {
     /// </summary>
     public int NumChannels { get; protected set; }
 
+    /// <summary>
+    /// Names of each channel in this image (e.g., "R" for the red color channel)
+    /// If null or empty, then a default interpretation (RGBA order) is assumed
+    /// </summary>
+    public string[] ChannelNames { get; protected set; }
+
     /// <summary>
     /// Pointer to the native memory containing the image data
     /// </summary>
@@ -77,6 +83,7 @@ public Image(int w, int h, int numChannels) {
         Width = w;
         Height = h;
         this.NumChannels = numChannels;
+        this.ChannelNames = new string[numChannels];
         Alloc();
 
         // Zero out the values to avoid undefined contents
@@ -140,6 +147,7 @@ public static void Move(Image src, Image dest) {
         dest.Width = src.Width;
         dest.Height = src.Height;
         dest.NumChannels = src.NumChannels;
+        dest.ChannelNames = src.ChannelNames;
     }
 
     [MethodImpl(MethodImplOptions.AggressiveInlining)]
@@ -366,6 +374,7 @@ protected void LoadFromFile(string filename) {
         Width = w;
         Height = h;
         NumChannels = n;
+        ChannelNames = new string[n];
         if (id < 0 || Width <= 0 || Height <= 0)
             throw new IOException($"ERROR: Could not load image file '{filename}'");
 
@@ -441,6 +450,7 @@ protected void Zoom(Image other, int scale) {
         Width = other.Width * scale;
         Height = other.Height * scale;
         NumChannels = other.NumChannels;
+        ChannelNames = other.ChannelNames;
         Alloc();
 
         SimpleImageIOCore.ZoomWithNearestInterp(other.DataPointer, NumChannels * other.Width, DataPointer,
@@ -468,6 +478,7 @@ public static bool Matches(Image a, Image b)
         result.Width = a.Width;
         result.Height = a.Height;
         result.NumChannels = a.NumChannels;
+        result.ChannelNames = a.ChannelNames;
         result.Alloc();
         Parallel.For(0, a.Height, row => {
             for (int col = 0; col < a.Width; ++col) {
@@ -492,6 +503,7 @@ public static bool Matches(Image a, Image b)
         result.Width = a.Width;
         result.Height = a.Height;
         result.NumChannels = a.NumChannels;
+        result.ChannelNames = a.ChannelNames;
         result.Alloc();
         Parallel.For(0, a.Height, row => {
             for (int col = 0; col < a.Width; ++col) {
diff --git a/SimpleImageIO/Layers.cs b/SimpleImageIO/Layers.cs
@@ -41,6 +41,14 @@ public static Dictionary<string, Image> LoadFromFile(string filename) {
             else
                 layers[name] = new(width, height, numChans);
 
+            // Read the channel names
+            for (int c = 0; c < numChans; ++c) {
+                int clen = SimpleImageIOCore.GetExrChannelNameLen(id, name, c);
+                StringBuilder chanNameBuilder = new(clen);
+                SimpleImageIOCore.GetExrChannelName(id, name, c, chanNameBuilder);
+                layers[name].ChannelNames[c] = chanNameBuilder.ToString();
+            }
+
             SimpleImageIOCore.CopyCachedLayer(id, name, layers[name].DataPointer);
         }
 
diff --git a/SimpleImageIO/SimpleImageIOCore.cs b/SimpleImageIO/SimpleImageIOCore.cs
@@ -16,6 +16,12 @@ public static extern int CacheImage(out int width, out int height, out int numCh
     [DllImport("SimpleImageIOCore", CallingConvention = CallingConvention.Cdecl)]
     public static extern int GetExrLayerChannelCount(int id, string name);
 
+    [DllImport("SimpleImageIOCore", CallingConvention = CallingConvention.Cdecl)]
+    public static extern int GetExrChannelNameLen(int id, string layerName, int channelIdx);
+
+    [DllImport("SimpleImageIOCore", CallingConvention = CallingConvention.Cdecl)]
+    public static extern int GetExrChannelName(int id, string layerName, int channelIdx, StringBuilder name);
+
     [DllImport("SimpleImageIOCore", CallingConvention = CallingConvention.Cdecl)]
     public static extern int GetExrLayerNameLen(int id, int layerIdx);
 
