From d73ef2440631d36e90bc3964129066e096b63127 Mon Sep 17 00:00:00 2001
From: Sebastian Schmidt <mrschmidt@google.com>
Date: Fri, 27 Oct 2023 10:32:11 -0700
Subject: [PATCH] Support 3-channel RGB images for Mac Python

PiperOrigin-RevId: 577240413
---
 mediapipe/python/pybind/image.cc | 22 +++++++++++++++++++---
 1 file changed, 19 insertions(+), 3 deletions(-)

diff --git a/mediapipe/python/pybind/image.cc b/mediapipe/python/pybind/image.cc
index 62437e439..f3d89bada 100644
--- a/mediapipe/python/pybind/image.cc
+++ b/mediapipe/python/pybind/image.cc
@@ -244,12 +244,26 @@ void ImageSubmodule(pybind11::module* module) {
   image.def_static(
       "create_from_file",
       [](const std::string& file_name) {
+        unsigned char* image_data = nullptr;
         int width;
         int height;
         int channels;
-        auto* image_data =
-            stbi_load(file_name.c_str(), &width, &height, &channels,
-                      /*desired_channels=*/0);
+
+#if TARGET_OS_OSX && !MEDIAPIPE_DISABLE_GPU
+        // Our ObjC layer does not support 3-channel images, so we read the
+        // number of channels first and request RGBA if needed.
+        if (stbi_info(file_name.c_str(), &width, &height, &channels)) {
+          if (channels == 3) {
+            channels = 4;
+          }
+          int unused;
+          image_data =
+              stbi_load(file_name.c_str(), &width, &height, &unused, channels);
+        }
+#else
+        image_data = stbi_load(file_name.c_str(), &width, &height, &channels,
+                               /*desired_channels=*/0);
+#endif  // TARGET_OS_OSX && !MEDIAPIPE_DISABLE_GPU
         if (image_data == nullptr) {
           throw RaisePyError(PyExc_RuntimeError,
                              absl::StrFormat("Image decoding failed (%s): %s",
@@ -263,11 +277,13 @@ void ImageSubmodule(pybind11::module* module) {
                 ImageFormat::GRAY8, width, height, width, image_data,
                 stbi_image_free);
             break;
+#if !TARGET_OS_OSX || MEDIAPIPE_DISABLE_GPU
           case 3:
             image_frame = std::make_shared<ImageFrame>(
                 ImageFormat::SRGB, width, height, 3 * width, image_data,
                 stbi_image_free);
             break;
+#endif  // !TARGET_OS_OSX || MEDIAPIPE_DISABLE_GPU
           case 4:
             image_frame = std::make_shared<ImageFrame>(
                 ImageFormat::SRGBA, width, height, 4 * width, image_data,