Project import generated by Copybara.

GitOrigin-RevId: 612e50bb8db2ec3dc1c30049372d87a80c3848db

mediapipe/python/packet_creator.py

@@ -16,6 +16,7 @@
 """The public facing packet creator APIs."""
 
 from typing import List, Union
+import warnings
 
 import numpy as np
 
@@ -48,34 +49,60 @@ create_string_to_packet_map = _packet_creator.create_string_to_packet_map
 create_matrix = _packet_creator.create_matrix
 
 
-def create_image_frame(
-    data: Union[image_frame.ImageFrame, np.ndarray],
-    *,
-    image_format: image_frame.ImageFormat = None) -> packet.Packet:
+def create_image_frame(data: Union[image_frame.ImageFrame, np.ndarray],
+                       *,
+                       image_format: image_frame.ImageFormat = None,
+                       copy: bool = None) -> packet.Packet:
   """Create a MediaPipe ImageFrame packet.
 
-  A MediaPipe ImageFrame packet can be created from either the raw pixel data
+  A MediaPipe ImageFrame packet can be created from an existing MediaPipe
+  ImageFrame object and the data will be realigned and copied into a new
+  ImageFrame object inside of the packet.
+
+  A MediaPipe ImageFrame packet can also be created from the raw pixel data
   represented as a numpy array with one of the uint8, uint16, and float data
-  types or an existing MediaPipe ImageFrame object. The data will be realigned
-  and copied into an ImageFrame object inside of the packet.
+  types. There are three data ownership modes depending on how the 'copy' arg
+  is set.
+
+  i) Default mode
+  If copy is not set, mutable data is always copied while the immutable data
+  is by reference.
+
+  ii) Copy mode (safe)
+  If copy is set to True, the data will be realigned and copied into an
+  ImageFrame object inside of the packet regardless the immutablity of the
+  original data.
+
+  iii) Reference mode (dangerous)
+  If copy is set to False, the data will be forced to be shared. If the data is
+  mutable (data.flags.writeable is True), a warning will be raised.
 
   Args:
     data: A MediaPipe ImageFrame object or the raw pixel data that is
       represnted as a numpy ndarray.
     image_format: One of the image_frame.ImageFormat enum types.
+    copy: Indicate if the packet should copy the data from the numpy nparray.
 
   Returns:
     A MediaPipe ImageFrame Packet.
 
   Raises:
     ValueError:
-      i) When "data" is a numpy ndarray, "image_format" is not provided.
+      i) When "data" is a numpy ndarray, "image_format" is not provided or
+        the "data" array is not c_contiguous in the reference mode.
       ii) When "data" is an ImageFrame object, the "image_format" arg doesn't
-        match the image format of the "data" ImageFrame object.
+        match the image format of the "data" ImageFrame object or "copy" is
+        explicitly set to False.
     TypeError: If "image format" doesn't match "data" array's data type.
 
   Examples:
     np_array = np.random.randint(255, size=(321, 123, 3), dtype=np.uint8)
+    # Copy mode by default if the data array is writable.
+    image_frame_packet = mp.packet_creator.create_image_frame(
+        image_format=mp.ImageFormat.SRGB, data=np_array)
+
+    # Make the array unwriteable to trigger the reference mode.
+    np_array.flags.writeable = False
     image_frame_packet = mp.packet_creator.create_image_frame(
         image_format=mp.ImageFormat.SRGB, data=np_array)
 
@@ -87,14 +114,33 @@ def create_image_frame(
     if image_format is not None and data.image_format != image_format:
       raise ValueError(
           'The provided image_format doesn\'t match the one from the data arg.')
+    if copy is not None and not copy:
+      raise ValueError(
+          'Creating image frame packet by taking a reference of another image frame object is not supported yet.'
+      )
     # pylint:disable=protected-access
-    return _packet_creator._create_image_frame_with_copy(data)
+    return _packet_creator._create_image_frame_from_image_frame(data)
     # pylint:enable=protected-access
   else:
     if image_format is None:
       raise ValueError('Please provide \'image_format\' with \'data\'.')
+    # If copy arg is not set, copying the data if it's immutable. Otherwise,
+    # take a reference of the immutable data to avoid data copy.
+    if copy is None:
+      copy = True if data.flags.writeable else False
+    if not copy:
+      # TODO: Investigate why the first 2 bytes of the data has data
+      # corruption when "data" is not c_contiguous.
+      if not data.flags.c_contiguous:
+        raise ValueError(
+            'Reference mode is unavailable if \'data\' is not c_contiguous.')
+      if data.flags.writeable:
+        warnings.warn(
+            '\'data\' is still writeable. Taking a reference of the data to create ImageFrame packet is dangerous.',
+            RuntimeWarning, 2)
     # pylint:disable=protected-access
-    return _packet_creator._create_image_frame_with_copy(image_format, data)
+    return _packet_creator._create_image_frame_from_pixel_data(
+        image_format, data, copy)
     # pylint:enable=protected-access