Source code for apache_beam.ml.gcp.videointelligenceml

#
# Licensed to the Apache Software Foundation (ASF) under one or more
# contributor license agreements.  See the NOTICE file distributed with
# this work for additional information regarding copyright ownership.
# The ASF licenses this file to You under the Apache License, Version 2.0
# (the "License"); you may not use this file except in compliance with
# the License.  You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#

"""A connector for sending API requests to the GCP Video Intelligence API."""

from typing import Optional
from typing import Tuple
from typing import Union

from apache_beam import typehints
from apache_beam.metrics import Metrics
from apache_beam.transforms import DoFn
from apache_beam.transforms import ParDo
from apache_beam.transforms import PTransform
from cachetools.func import ttl_cache

try:
  from google.cloud import videointelligence
except ImportError:
  raise ImportError(
      'Google Cloud Video Intelligence not supported for this execution '
      'environment (could not import google.cloud.videointelligence).')

__all__ = ['AnnotateVideo', 'AnnotateVideoWithContext']


@ttl_cache(maxsize=128, ttl=3600)
def get_videointelligence_client():
  """Returns a Cloud Video Intelligence client."""
  _client = videointelligence.VideoIntelligenceServiceClient()
  return _client


[docs]class AnnotateVideo(PTransform): """A ``PTransform`` for annotating video using the GCP Video Intelligence API ref: https://cloud.google.com/video-intelligence/docs Sends each element to the GCP Video Intelligence API. Element is a Union[str, bytes] of either an URI (e.g. a GCS URI) or bytes base64-encoded video data. Accepts an `AsDict` side input that maps each video to a video context. """ def __init__( self, features, location_id=None, metadata=None, timeout=120, context_side_input=None): """ Args: features: (List[``videointelligence_v1.Feature``]) Required. The Video Intelligence API features to detect location_id: (str) Optional. Cloud region where annotation should take place. If no region is specified, a region will be determined based on video file location. metadata: (Sequence[Tuple[str, str]]) Optional. Additional metadata that is provided to the method. timeout: (int) Optional. The time in seconds to wait for the response from the Video Intelligence API context_side_input: (beam.pvalue.AsDict) Optional. An ``AsDict`` of a PCollection to be passed to the _VideoAnnotateFn as the video context mapping containing additional video context and/or feature-specific parameters. Example usage:: video_contexts = [('gs://cloud-samples-data/video/cat.mp4', Union[dict, ``videointelligence_v1.VideoContext``]), ('gs://some-other-video/sample.mp4', Union[dict, ``videointelligence_v1.VideoContext``]),] context_side_input = ( p | "Video contexts" >> beam.Create(video_contexts) ) videointelligenceml.AnnotateVideo(features, context_side_input=beam.pvalue.AsDict(context_side_input))) """ super().__init__() self.features = features self.location_id = location_id self.metadata = metadata self.timeout = timeout self.context_side_input = context_side_input
[docs] def expand(self, pvalue): return pvalue | ParDo( _VideoAnnotateFn( features=self.features, location_id=self.location_id, metadata=self.metadata, timeout=self.timeout), context_side_input=self.context_side_input)
@typehints.with_input_types( Union[str, bytes], Optional[videointelligence.VideoContext]) class _VideoAnnotateFn(DoFn): """A DoFn that sends each input element to the GCP Video Intelligence API service and outputs an element with the return result of the API (``google.cloud.videointelligence_v1.AnnotateVideoResponse``). """ def __init__(self, features, location_id, metadata, timeout): super().__init__() self._client = None self.features = features self.location_id = location_id self.metadata = metadata self.timeout = timeout self.counter = Metrics.counter(self.__class__, "API Calls") def start_bundle(self): self._client = get_videointelligence_client() def _annotate_video(self, element, video_context): if isinstance(element, str): # Is element an URI to a GCS bucket response = self._client.annotate_video( input_uri=element, features=self.features, video_context=video_context, location_id=self.location_id, metadata=self.metadata) else: # Is element raw bytes response = self._client.annotate_video( input_content=element, features=self.features, video_context=video_context, location_id=self.location_id, metadata=self.metadata) return response def process(self, element, context_side_input=None, *args, **kwargs): if context_side_input: # If we have a side input video context, use that video_context = context_side_input.get(element) else: video_context = None response = self._annotate_video(element, video_context) self.counter.inc() yield response.result(timeout=self.timeout)
[docs]class AnnotateVideoWithContext(AnnotateVideo): """A ``PTransform`` for annotating video using the GCP Video Intelligence API ref: https://cloud.google.com/video-intelligence/docs Sends each element to the GCP Video Intelligence API. Element is a tuple of (Union[str, bytes], Optional[videointelligence.VideoContext]) where the former is either an URI (e.g. a GCS URI) or bytes base64-encoded video data """ def __init__(self, features, location_id=None, metadata=None, timeout=120): """ Args: features: (List[``videointelligence_v1.Feature``]) Required. the Video Intelligence API features to detect location_id: (str) Optional. Cloud region where annotation should take place. If no region is specified, a region will be determined based on video file location. metadata: (Sequence[Tuple[str, str]]) Optional. Additional metadata that is provided to the method. timeout: (int) Optional. The time in seconds to wait for the response from the Video Intelligence API """ super().__init__( features=features, location_id=location_id, metadata=metadata, timeout=timeout)
[docs] def expand(self, pvalue): return pvalue | ParDo( _VideoAnnotateFnWithContext( features=self.features, location_id=self.location_id, metadata=self.metadata, timeout=self.timeout))
@typehints.with_input_types( Tuple[Union[str, bytes], Optional[videointelligence.VideoContext]]) class _VideoAnnotateFnWithContext(_VideoAnnotateFn): """A DoFn that unpacks each input tuple to element, video_context variables and sends these to the GCP Video Intelligence API service and outputs an element with the return result of the API (``google.cloud.videointelligence_v1.AnnotateVideoResponse``). """ def __init__(self, features, location_id, metadata, timeout): super().__init__( features=features, location_id=location_id, metadata=metadata, timeout=timeout) def process(self, element, *args, **kwargs): element, video_context = element # Unpack (video, video_context) tuple response = self._annotate_video(element, video_context) self.counter.inc() yield response.result(timeout=self.timeout)