Implement Image.create_from_file and update the object_detector_test.py file accordingly.
PiperOrigin-RevId: 477682930
This commit is contained in:
committed by
Copybara-Service
parent
80fd47b820
commit
554e2a9d69
@@ -48,6 +48,8 @@ pybind_extension(
|
||||
"//mediapipe/python/pybind:timestamp",
|
||||
"//mediapipe/python/pybind:validated_graph_config",
|
||||
"//mediapipe/tasks/python/core/pybind:task_runner",
|
||||
"@com_google_absl//absl/strings:str_format",
|
||||
"@stblib//:stb_image",
|
||||
# Type registration.
|
||||
"//mediapipe/framework:basic_types_registration",
|
||||
"//mediapipe/framework/formats:classification_registration",
|
||||
|
||||
@@ -15,6 +15,7 @@
|
||||
"""Tests for mediapipe.python._framework_bindings.image."""
|
||||
|
||||
import gc
|
||||
import os
|
||||
import random
|
||||
import sys
|
||||
|
||||
@@ -23,6 +24,7 @@ import cv2
|
||||
import numpy as np
|
||||
import PIL.Image
|
||||
|
||||
# resources dependency
|
||||
from mediapipe.python._framework_bindings import image
|
||||
from mediapipe.python._framework_bindings import image_frame
|
||||
|
||||
@@ -185,6 +187,16 @@ class ImageTest(absltest.TestCase):
|
||||
gc.collect()
|
||||
self.assertEqual(sys.getrefcount(rgb_image), initial_ref_count)
|
||||
|
||||
def test_image_create_from_file(self):
|
||||
image_path = os.path.join(
|
||||
resources.GetRunfilesDir(),
|
||||
'mediapipe/tasks/testdata/vision/cat.jpg')
|
||||
loaded_image = Image.create_from_file(image_path)
|
||||
self.assertEqual(loaded_image.width, 600)
|
||||
self.assertEqual(loaded_image.height, 400)
|
||||
self.assertEqual(loaded_image.channels, 3)
|
||||
self.assertEqual(loaded_image.image_format, ImageFormat.SRGB)
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
absltest.main()
|
||||
|
||||
@@ -45,6 +45,8 @@ pybind_library(
|
||||
":util",
|
||||
"//mediapipe/framework:type_map",
|
||||
"//mediapipe/framework/formats:image",
|
||||
"@com_google_absl//absl/strings:str_format",
|
||||
"@stblib//:stb_image",
|
||||
],
|
||||
)
|
||||
|
||||
|
||||
@@ -16,9 +16,11 @@
|
||||
|
||||
#include <memory>
|
||||
|
||||
#include "absl/strings/str_format.h"
|
||||
#include "mediapipe/python/pybind/image_frame_util.h"
|
||||
#include "mediapipe/python/pybind/util.h"
|
||||
#include "pybind11/stl.h"
|
||||
#include "stb_image.h"
|
||||
|
||||
namespace mediapipe {
|
||||
namespace python {
|
||||
@@ -225,6 +227,62 @@ void ImageSubmodule(pybind11::module* module) {
|
||||
image.is_aligned(16)
|
||||
)doc");
|
||||
|
||||
image.def_static(
|
||||
"create_from_file",
|
||||
[](const std::string& file_name) {
|
||||
int width;
|
||||
int height;
|
||||
int channels;
|
||||
auto* image_data =
|
||||
stbi_load(file_name.c_str(), &width, &height, &channels,
|
||||
/*desired_channels=*/0);
|
||||
if (image_data == nullptr) {
|
||||
throw RaisePyError(PyExc_RuntimeError,
|
||||
absl::StrFormat("Image decoding failed (%s): %s",
|
||||
stbi_failure_reason(), file_name)
|
||||
.c_str());
|
||||
}
|
||||
ImageFrameSharedPtr image_frame;
|
||||
switch (channels) {
|
||||
case 1:
|
||||
image_frame = std::make_shared<ImageFrame>(
|
||||
ImageFormat::GRAY8, width, height, width, image_data,
|
||||
stbi_image_free);
|
||||
break;
|
||||
case 3:
|
||||
image_frame = std::make_shared<ImageFrame>(
|
||||
ImageFormat::SRGB, width, height, 3 * width, image_data,
|
||||
stbi_image_free);
|
||||
break;
|
||||
case 4:
|
||||
image_frame = std::make_shared<ImageFrame>(
|
||||
ImageFormat::SRGBA, width, height, 4 * width, image_data,
|
||||
stbi_image_free);
|
||||
break;
|
||||
default:
|
||||
throw RaisePyError(
|
||||
PyExc_RuntimeError,
|
||||
absl::StrFormat(
|
||||
"Expected image with 1 (grayscale), 3 (RGB) or 4 "
|
||||
"(RGBA) channels, found %d channels.",
|
||||
channels)
|
||||
.c_str());
|
||||
}
|
||||
return Image(std::move(image_frame));
|
||||
},
|
||||
R"doc(Creates `Image` object from the image file.
|
||||
|
||||
Args:
|
||||
file_name: Image file name.
|
||||
|
||||
Returns:
|
||||
`Image` object.
|
||||
|
||||
Raises:
|
||||
RuntimeError if the image file can't be decoded.
|
||||
)doc",
|
||||
py::arg("file_name"));
|
||||
|
||||
image.def_property_readonly("width", &Image::width)
|
||||
.def_property_readonly("height", &Image::height)
|
||||
.def_property_readonly("channels", &Image::channels)
|
||||
|
||||
Reference in New Issue
Block a user