Project import generated by Copybara.

GitOrigin-RevId: 6e5aa035cd1f6a9333962df5d3ab97a05bd5744e

mediapipe/python/solution_base.py

@@ -28,6 +28,7 @@ from typing import Any, Iterable, List, Mapping, NamedTuple, Optional, Union
 
 import numpy as np
 
+from google.protobuf.internal import containers
 from google.protobuf import descriptor
 from google.protobuf import message
 # resources dependency
@@ -216,6 +217,7 @@ class SolutionBase:
       binary_graph_path: Optional[str] = None,
       graph_config: Optional[calculator_pb2.CalculatorGraphConfig] = None,
       calculator_params: Optional[Mapping[str, Any]] = None,
+      graph_options: Optional[message.Message] = None,
       side_inputs: Optional[Mapping[str, Any]] = None,
       outputs: Optional[List[str]] = None,
       stream_type_hints: Optional[Mapping[str, PacketDataType]] = None):
@@ -227,6 +229,7 @@ class SolutionBase:
         format.
       calculator_params: A mapping from the
         {calculator_name}.{options_field_name} str to the field value.
+      graph_options: The graph options protobuf for the mediapipe graph.
       side_inputs: A mapping from the side packet name to the packet raw data.
       outputs: A list of the graph output stream names to observe. If the list
         is empty, all the output streams listed in the graph config will be
@@ -267,6 +270,10 @@ class SolutionBase:
     if calculator_params:
       self._modify_calculator_options(canonical_graph_config_proto,
                                       calculator_params)
+    if graph_options:
+      self._set_extension(canonical_graph_config_proto.graph_options,
+                          graph_options)
+
     self._graph = calculator_graph.CalculatorGraph(
         graph_config=canonical_graph_config_proto)
     self._simulated_timestamp = 0
@@ -530,6 +537,50 @@ class SolutionBase:
     if num_modified < len(nested_calculator_params):
       raise ValueError('Not all calculator params are valid.')
 
+  def create_graph_options(self, options_message: message.Message,
+                           values: Mapping[str, Any]) -> message.Message:
+    """Sets protobuf field values.
+
+    Args:
+      options_message: the options protobuf message.
+      values: field value pairs, where each field may be a "." separated path.
+
+    Returns:
+      the options protobuf message.
+    """
+
+    if hasattr(values, 'items'):
+      values = values.items()
+    for pair in values:
+      (field, value) = pair
+      fields = field.split('.')
+      m = options_message
+      while len(fields) > 1:
+        m = getattr(m, fields[0])
+        del fields[0]
+      v = getattr(m, fields[0])
+      if hasattr(v, 'append'):
+        del v[:]
+        v.extend(value)
+      elif hasattr(v, 'CopyFrom'):
+        v.CopyFrom(value)
+      else:
+        setattr(m, fields[0], value)
+    return options_message
+
+  def _set_extension(self,
+                     extension_list: containers.RepeatedCompositeFieldContainer,
+                     extension_value: message.Message) -> None:
+    """Sets one value in a repeated protobuf.Any extension field."""
+    for extension_any in extension_list:
+      if extension_any.Is(extension_value.DESCRIPTOR):
+        v = type(extension_value)()
+        extension_any.Unpack(v)
+        v.MergeFrom(extension_value)
+        extension_any.Pack(v)
+        return
+    extension_list.add().Pack(extension_value)
+
   def _make_packet(self, packet_data_type: PacketDataType,
                    data: Any) -> packet.Packet:
     if (packet_data_type == PacketDataType.IMAGE_FRAME or