github.com/apache/beam/sdks/v2@v2.48.2/python/apache_beam/ml/gcp/videointelligenceml.py (about)

     1  #
     2  # Licensed to the Apache Software Foundation (ASF) under one or more
     3  # contributor license agreements.  See the NOTICE file distributed with
     4  # this work for additional information regarding copyright ownership.
     5  # The ASF licenses this file to You under the Apache License, Version 2.0
     6  # (the "License"); you may not use this file except in compliance with
     7  # the License.  You may obtain a copy of the License at
     8  #
     9  #    http://www.apache.org/licenses/LICENSE-2.0
    10  #
    11  # Unless required by applicable law or agreed to in writing, software
    12  # distributed under the License is distributed on an "AS IS" BASIS,
    13  # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    14  # See the License for the specific language governing permissions and
    15  # limitations under the License.
    16  #
    17  
    18  """A connector for sending API requests to the GCP Video Intelligence API."""
    19  
    20  from typing import Optional
    21  from typing import Tuple
    22  from typing import Union
    23  
    24  from apache_beam import typehints
    25  from apache_beam.metrics import Metrics
    26  from apache_beam.transforms import DoFn
    27  from apache_beam.transforms import ParDo
    28  from apache_beam.transforms import PTransform
    29  from cachetools.func import ttl_cache
    30  
    31  try:
    32    from google.cloud import videointelligence
    33  except ImportError:
    34    raise ImportError(
    35        'Google Cloud Video Intelligence not supported for this execution '
    36        'environment (could not import google.cloud.videointelligence).')
    37  
    38  __all__ = ['AnnotateVideo', 'AnnotateVideoWithContext']
    39  
    40  
    41  @ttl_cache(maxsize=128, ttl=3600)
    42  def get_videointelligence_client():
    43    """Returns a Cloud Video Intelligence client."""
    44    _client = videointelligence.VideoIntelligenceServiceClient()
    45    return _client
    46  
    47  
    48  class AnnotateVideo(PTransform):
    49    """A ``PTransform`` for annotating video using the GCP Video Intelligence API
    50    ref: https://cloud.google.com/video-intelligence/docs
    51  
    52    Sends each element to the GCP Video Intelligence API. Element is a
    53    Union[str, bytes] of either an URI (e.g. a GCS URI) or
    54    bytes base64-encoded video data.
    55    Accepts an `AsDict` side input that maps each video to a video context.
    56    """
    57    def __init__(
    58        self,
    59        features,
    60        location_id=None,
    61        metadata=None,
    62        timeout=120,
    63        context_side_input=None):
    64      """
    65      Args:
    66        features: (List[``videointelligence_v1.Feature``]) Required.
    67          The Video Intelligence API features to detect
    68        location_id: (str) Optional.
    69          Cloud region where annotation should take place.
    70          If no region is specified, a region will be determined
    71          based on video file location.
    72        metadata: (Sequence[Tuple[str, str]]) Optional.
    73          Additional metadata that is provided to the method.
    74        timeout: (int) Optional.
    75          The time in seconds to wait for the response from the
    76          Video Intelligence API
    77        context_side_input: (beam.pvalue.AsDict) Optional.
    78          An ``AsDict`` of a PCollection to be passed to the
    79          _VideoAnnotateFn as the video context mapping containing additional
    80          video context and/or feature-specific parameters.
    81          Example usage::
    82  
    83            video_contexts =
    84              [('gs://cloud-samples-data/video/cat.mp4', Union[dict,
    85              ``videointelligence_v1.VideoContext``]),
    86              ('gs://some-other-video/sample.mp4', Union[dict,
    87              ``videointelligence_v1.VideoContext``]),]
    88  
    89            context_side_input =
    90              (
    91                p
    92                | "Video contexts" >> beam.Create(video_contexts)
    93              )
    94  
    95            videointelligenceml.AnnotateVideo(features,
    96              context_side_input=beam.pvalue.AsDict(context_side_input)))
    97      """
    98      super().__init__()
    99      self.features = features
   100      self.location_id = location_id
   101      self.metadata = metadata
   102      self.timeout = timeout
   103      self.context_side_input = context_side_input
   104  
   105    def expand(self, pvalue):
   106      return pvalue | ParDo(
   107          _VideoAnnotateFn(
   108              features=self.features,
   109              location_id=self.location_id,
   110              metadata=self.metadata,
   111              timeout=self.timeout),
   112          context_side_input=self.context_side_input)
   113  
   114  
   115  @typehints.with_input_types(
   116      Union[str, bytes], Optional[videointelligence.VideoContext])
   117  class _VideoAnnotateFn(DoFn):
   118    """A DoFn that sends each input element to the GCP Video Intelligence API
   119    service and outputs an element with the return result of the API
   120    (``google.cloud.videointelligence_v1.AnnotateVideoResponse``).
   121    """
   122    def __init__(self, features, location_id, metadata, timeout):
   123      super().__init__()
   124      self._client = None
   125      self.features = features
   126      self.location_id = location_id
   127      self.metadata = metadata
   128      self.timeout = timeout
   129      self.counter = Metrics.counter(self.__class__, "API Calls")
   130  
   131    def start_bundle(self):
   132      self._client = get_videointelligence_client()
   133  
   134    def _annotate_video(self, element, video_context):
   135      if isinstance(element, str):  # Is element an URI to a GCS bucket
   136        response = self._client.annotate_video(
   137            input_uri=element,
   138            features=self.features,
   139            video_context=video_context,
   140            location_id=self.location_id,
   141            metadata=self.metadata)
   142      else:  # Is element raw bytes
   143        response = self._client.annotate_video(
   144            input_content=element,
   145            features=self.features,
   146            video_context=video_context,
   147            location_id=self.location_id,
   148            metadata=self.metadata)
   149      return response
   150  
   151    def process(self, element, context_side_input=None, *args, **kwargs):
   152      if context_side_input:  # If we have a side input video context, use that
   153        video_context = context_side_input.get(element)
   154      else:
   155        video_context = None
   156      response = self._annotate_video(element, video_context)
   157      self.counter.inc()
   158      yield response.result(timeout=self.timeout)
   159  
   160  
   161  class AnnotateVideoWithContext(AnnotateVideo):
   162    """A ``PTransform`` for annotating video using the GCP Video Intelligence API
   163    ref: https://cloud.google.com/video-intelligence/docs
   164  
   165    Sends each element to the GCP Video Intelligence API.
   166    Element is a tuple of
   167  
   168      (Union[str, bytes],
   169      Optional[videointelligence.VideoContext])
   170  
   171    where the former is either an URI (e.g. a GCS URI) or
   172    bytes base64-encoded video data
   173    """
   174    def __init__(self, features, location_id=None, metadata=None, timeout=120):
   175      """
   176        Args:
   177          features: (List[``videointelligence_v1.Feature``]) Required.
   178            the Video Intelligence API features to detect
   179          location_id: (str) Optional.
   180            Cloud region where annotation should take place.
   181            If no region is specified, a region will be determined
   182            based on video file location.
   183          metadata: (Sequence[Tuple[str, str]]) Optional.
   184            Additional metadata that is provided to the method.
   185          timeout: (int) Optional.
   186            The time in seconds to wait for the response from the
   187            Video Intelligence API
   188      """
   189      super().__init__(
   190          features=features,
   191          location_id=location_id,
   192          metadata=metadata,
   193          timeout=timeout)
   194  
   195    def expand(self, pvalue):
   196      return pvalue | ParDo(
   197          _VideoAnnotateFnWithContext(
   198              features=self.features,
   199              location_id=self.location_id,
   200              metadata=self.metadata,
   201              timeout=self.timeout))
   202  
   203  
   204  @typehints.with_input_types(
   205      Tuple[Union[str, bytes], Optional[videointelligence.VideoContext]])
   206  class _VideoAnnotateFnWithContext(_VideoAnnotateFn):
   207    """A DoFn that unpacks each input tuple to element, video_context variables
   208    and sends these to the GCP Video Intelligence API service and outputs
   209    an element with the return result of the API
   210    (``google.cloud.videointelligence_v1.AnnotateVideoResponse``).
   211    """
   212    def __init__(self, features, location_id, metadata, timeout):
   213      super().__init__(
   214          features=features,
   215          location_id=location_id,
   216          metadata=metadata,
   217          timeout=timeout)
   218  
   219    def process(self, element, *args, **kwargs):
   220      element, video_context = element  # Unpack (video, video_context) tuple
   221      response = self._annotate_video(element, video_context)
   222      self.counter.inc()
   223      yield response.result(timeout=self.timeout)