Project import generated by Copybara.

GitOrigin-RevId: ff83882955f1a1e2a043ff4e71278be9d7217bbe

mediapipe/objc/BUILD

@@ -72,6 +72,7 @@ objc_library(
         ":util",
         "//mediapipe/framework:calculator_framework",
         "//mediapipe/framework:mediapipe_profiling",
+        "//mediapipe/framework/formats:image",
         "//mediapipe/framework/port:map_util",
         "//mediapipe/framework/port:ret_check",
         "//mediapipe/framework/port:source_location",

mediapipe/objc/MPPGraph.h

@@ -59,16 +59,21 @@ typedef NS_ENUM(int, MPPPacketType) {
   /// Calls mediapipeGraph:didOutputPacket:fromStream:
   MPPPacketTypeRaw,
 
-  /// CFHolder<CVPixelBufferRef>.
+  /// GpuBuffer packet.
   /// Calls mediapipeGraph:didOutputPixelBuffer:fromStream:
   /// Use this packet type to pass GPU frames to calculators.
   MPPPacketTypePixelBuffer,
 
-  /// ImageFrame.
+  /// Image packet.
+  /// Calls mediapipeGraph:didOutputPixelBuffer:fromStream:
+  /// Use this packet type to pass GPU frames to calculators.
+  MPPPacketTypeImage,
+
+  /// ImageFrame packet.
   /// Calls mediapipeGraph:didOutputPixelBuffer:fromStream:
   MPPPacketTypeImageFrame,
 
-  /// RGBA ImageFrame, but do not swap the channels if the input pixel buffer
+  /// RGBA ImageFrame packet, but do not swap the channels if the input pixel buffer
   /// is BGRA. This is useful when the graph needs RGBA ImageFrames, but the
   /// calculators do not care about the order of the channels, so BGRA data can
   /// be used as-is.
@@ -164,6 +169,9 @@ typedef NS_ENUM(int, MPPPacketType) {
 - (mediapipe::Packet)packetWithPixelBuffer:(CVPixelBufferRef)pixelBuffer
                                 packetType:(MPPPacketType)packetType;
 
+/// Creates a MediaPipe packet of type Image, wrapping the given CVPixelBufferRef.
+- (mediapipe::Packet)imagePacketWithPixelBuffer:(CVPixelBufferRef)pixelBuffer;
+
 /// Sends a pixel buffer into a graph input stream, using the specified packet
 /// type. The graph must have been started before calling this. Drops frames and
 /// returns NO if maxFramesInFlight is exceeded. If allowOverwrite is set to YES,

mediapipe/objc/MPPGraph.mm

@@ -21,6 +21,7 @@
 
 #include "absl/memory/memory.h"
 #include "mediapipe/framework/calculator_framework.h"
+#include "mediapipe/framework/formats/image.h"
 #include "mediapipe/framework/formats/image_frame.h"
 #include "mediapipe/framework/graph_service.h"
 #include "mediapipe/gpu/MPPGraphGPUData.h"
@@ -163,10 +164,15 @@ void CallFrameDelegate(void* wrapperVoid, const std::string& streamName,
         _GTMDevLog(@"unsupported ImageFormat: %d", format);
       }
 #if MEDIAPIPE_GPU_BUFFER_USE_CV_PIXEL_BUFFER
-    } else if (packetType == MPPPacketTypePixelBuffer) {
+    } else if (packetType == MPPPacketTypePixelBuffer ||
+               packetType == MPPPacketTypeImage) {
       wrapper->_framesInFlight--;
-      CVPixelBufferRef pixelBuffer = packet.Get<mediapipe::GpuBuffer>().GetCVPixelBufferRef();
-      if ([wrapper.delegate
+      CVPixelBufferRef pixelBuffer;
+      if (packetType == MPPPacketTypePixelBuffer)
+        pixelBuffer = packet.Get<mediapipe::GpuBuffer>().GetCVPixelBufferRef();
+      else
+        pixelBuffer = packet.Get<mediapipe::Image>().GetCVPixelBufferRef();
+if ([wrapper.delegate
               respondsToSelector:@selector
               (mediapipeGraph:didOutputPixelBuffer:fromStream:timestamp:)]) {
         [wrapper.delegate mediapipeGraph:wrapper
@@ -315,6 +321,16 @@ void CallFrameDelegate(void* wrapperVoid, const std::string& streamName,
 #if MEDIAPIPE_GPU_BUFFER_USE_CV_PIXEL_BUFFER
   } else if (packetType == MPPPacketTypePixelBuffer) {
     packet = mediapipe::MakePacket<mediapipe::GpuBuffer>(imageBuffer);
+#endif  // MEDIAPIPE_GPU_BUFFER_USE_CV_PIXEL_BUFFER
+  } else if (packetType == MPPPacketTypeImage) {
+#if MEDIAPIPE_GPU_BUFFER_USE_CV_PIXEL_BUFFER
+    // GPU
+    packet = mediapipe::MakePacket<mediapipe::Image>(imageBuffer);
+#else
+    // CPU
+    auto frame = CreateImageFrameForCVPixelBuffer(imageBuffer, /* canOverwrite = */ false,
+                                                  /* bgrAsRgb = */ false);
+    packet = mediapipe::MakePacket<mediapipe::Image>(std::move(frame));
 #endif  // MEDIAPIPE_GPU_BUFFER_USE_CV_PIXEL_BUFFER
   } else {
     _GTMDevLog(@"unsupported packet type: %d", packetType);
@@ -322,6 +338,10 @@ void CallFrameDelegate(void* wrapperVoid, const std::string& streamName,
   return packet;
 }
 
+- (mediapipe::Packet)imagePacketWithPixelBuffer:(CVPixelBufferRef)pixelBuffer {
+  return [self packetWithPixelBuffer:(pixelBuffer) packetType:(MPPPacketTypeImage)];
+}
+
 - (BOOL)sendPixelBuffer:(CVPixelBufferRef)imageBuffer
              intoStream:(const std::string&)inputName
              packetType:(MPPPacketType)packetType

mediapipe/objc/MPPGraphTests.mm

@@ -16,6 +16,7 @@
 #import <XCTest/XCTest.h>
 
 #include "absl/memory/memory.h"
+#include "mediapipe/framework/formats/image.h"
 #import "mediapipe/objc/MPPGraph.h"
 #import "mediapipe/objc/MPPGraphTestBase.h"
 #import "mediapipe/objc/NSError+util_status.h"
@@ -333,4 +334,21 @@ REGISTER_CALCULATOR(ErrorCalculator);
   [self waitForExpectationsWithTimeout:3.0 handler:NULL];
 }
 
+- (void)testPixelBufferToImage {
+  CFHolder<CVPixelBufferRef> pixelBufferIn;
+  absl::Status status = CreateCVPixelBufferFromCGImage(_sourceImage.CGImage, &pixelBufferIn);
+  XCTAssert(status.ok());
+
+  mediapipe::CalculatorGraphConfig config;
+  _graph = [[MPPGraph alloc] initWithGraphConfig:config];
+
+  mediapipe::Packet packet = [_graph imagePacketWithPixelBuffer:*pixelBufferIn];
+  CVPixelBufferRef pixelBufferOut = packet.Get<mediapipe::Image>().GetCVPixelBufferRef();
+
+  XCTAssertTrue([self pixelBuffer:*pixelBufferIn
+                        isCloseTo:pixelBufferOut
+               maxLocalDifference:0
+             maxAverageDifference:0]);
+}
+
 @end