Support 3-channel RGB images for Mac Python

PiperOrigin-RevId: 577240413
This commit is contained in:
Sebastian Schmidt 2023-10-27 10:32:11 -07:00 committed by Copybara-Service
parent eaf0807849
commit d73ef24406

View File

@ -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,