Fix bug where getting image data in python fails if the mask has a di…

…fferent image size * add mask channel to consideration for number of channels
EmilDohne · Dec 18, 2024 · da2c180 · da2c180
1 parent ab6e15f
commit da2c180
Show file tree

Hide file tree

Showing 3 changed files with 64 additions and 4 deletions.
diff --git a/python/psapi-test/documents/MismatchedMaskChannel.psb b/python/psapi-test/documents/MismatchedMaskChannel.psb
diff --git a/python/psapi-test/test_imagelayer.py b/python/psapi-test/test_imagelayer.py
@@ -11,6 +11,7 @@
 
 class TestImageLayer(unittest.TestCase):
     path = os.path.join(os.path.dirname(__file__), "documents", "BaseFile.psb")
+    mask_path = os.path.join(os.path.dirname(__file__), "documents", "MismatchedMaskChannel.psb")
     bin_data_path = os.path.join(os.path.dirname(__file__), "bin_data", "monza_npy.bin")
 
     def test_get_image_data(self):
@@ -26,6 +27,24 @@ def test_get_image_data(self):
         self.assertTrue(len(layer.channels) == 4)
         self.assertTrue(layer.num_channels == 4)
 
+    def test_get_image_data_mismatched_mask(self):
+        """
+        Test that the mask channel gets read properly even if it has a different size than 
+        """
+        file = psapi.LayeredFile.read(self.mask_path)
+        layer: psapi.ImageLayer_16bit = file["MonzaSP1_DawnShot_V1_v002_ED.BaseAOV"]
+
+        image_data = layer.image_data
+        image_data_2 = layer.get_image_data()
+
+        self.assertTrue(-2 in image_data and -2 in image_data_2)
+        self.assertTrue(-1 in image_data and -1 in image_data_2)
+        self.assertTrue(0 in image_data and 0 in image_data_2)
+        self.assertTrue(1 in image_data and 1 in image_data_2)
+        self.assertTrue(2 in image_data and 2 in image_data_2)
+        self.assertTrue(len(layer.channels) == 5)
+        self.assertTrue(layer.num_channels == 5)
+
     def test_channel_setting_roundtrip(self):
         file = psapi.LayeredFile.read(self.path)
         layer: psapi.ImageLayer_16bit = file["Render"]["AOV"]["Beauties"]["MonzaSP1_DawnShot_V1_v002_ED.BaseAOV"]

diff --git a/python/src/DeclareImageLayer.h b/python/src/DeclareImageLayer.h
@@ -488,14 +488,33 @@ void declareImageLayer(py::module& m, const std::string& extension) {
         {
             auto data = self.getImageData(do_copy);
             std::unordered_map<int, py::array_t<T>> outData;
+
+            constexpr auto mask_id = Enum::ChannelIDInfo{ Enum::ChannelID::UserSuppliedLayerMask, -2 };
             for (auto& [key, value] : data)
             {
-                outData[key.index] = to_py_array(std::move(value), self.m_Width, self.m_Height);
+                // Mask channels may have a different resolution compared to the actual layer so we must account for this while parsing.
+                if (key == mask_id)
+                {
+                    if (!self.m_LayerMask)
+                    {
+                        throw py::value_error("Internal error: Encountered mask channel but layer does not have mask");
+                    }
+                    auto& mask = self.m_LayerMask.value();
+                    auto width = mask.maskData->getWidth();
+                    auto height = mask.maskData->getHeight();
+
+                    outData[key.index] = to_py_array(std::move(value), width, height);
+                }
+                else
+                {
+                    outData[key.index] = to_py_array(std::move(value), self.m_Width, self.m_Height);
+                }
             }
             return outData;
         }, py::arg("do_copy") = true, R"pbdoc(
 
-        Extract all the channels of the ImageLayer into an unordered_map.
+        Extract all the channels of the ImageLayer into an unordered_map. The channels may have differing sizes 
+        as photoshop optimizes mask channels differently than the pixel data
                 
         :param do_copy: Defaults to true, Whether to copy the data
         :type do_copy: bool
@@ -616,9 +635,27 @@ void declareImageLayer(py::module& m, const std::string& extension) {
         {
 			auto data = self.getImageData(true);
 			std::unordered_map<int, py::array_t<T>> outData;
+
+			constexpr auto mask_id = Enum::ChannelIDInfo{ Enum::ChannelID::UserSuppliedLayerMask, -2 };
 			for (auto& [key, value] : data)
 			{
-				outData[key.index] = to_py_array(std::move(value), self.m_Width, self.m_Height);
+				// Mask channels may have a different resolution compared to the actual layer so we must account for this while parsing.
+				if (key == mask_id)
+				{
+					if (!self.m_LayerMask)
+					{
+						throw py::value_error("Internal error: Encountered mask channel but layer does not have mask");
+					}
+					auto& mask = self.m_LayerMask.value();
+					auto width = mask.maskData->getWidth();
+					auto height = mask.maskData->getHeight();
+
+					outData[key.index] = to_py_array(std::move(value), width, height);
+				}
+				else
+				{
+					outData[key.index] = to_py_array(std::move(value), self.m_Width, self.m_Height);
+				}
 			}
 			return outData;
         }, R"pbdoc(
@@ -627,7 +664,7 @@ void declareImageLayer(py::module& m, const std::string& extension) {
 
     imageLayer.def_property_readonly("num_channels", [](Class& self)
         {
-            return self.m_ImageData.size();
+            return self.m_ImageData.size() + self.m_LayerMask.has_value();
 		});
 
     imageLayer.def_property_readonly("channels", [](Class& self)
@@ -637,6 +674,10 @@ void declareImageLayer(py::module& m, const std::string& extension) {
             {
                 indices.push_back(key.index);
             }
+            if (self.m_LayerMask)
+            {
+                indices.push_back(-2);
+            }
             return indices;
         });
 }