github.com/apache/beam/sdks/v2@v2.48.2/python/apache_beam/ml/gcp/visionml_test_it.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  # pytype: skip-file
    18  
    19  import unittest
    20  
    21  import pytest
    22  
    23  import apache_beam as beam
    24  from apache_beam.testing.test_pipeline import TestPipeline
    25  from apache_beam.testing.util import assert_that
    26  from apache_beam.testing.util import equal_to
    27  
    28  # Protect against environments where Google Cloud Vision client is not
    29  # available.
    30  try:
    31    from apache_beam.ml.gcp.visionml import AnnotateImage
    32    from google.cloud import vision
    33  except ImportError:
    34    vision = None
    35  
    36  
    37  def extract(response):
    38    for r in response.responses:
    39      for text_annotation in r.text_annotations:
    40        yield text_annotation.description
    41  
    42  
    43  @pytest.mark.it_postcommit
    44  @unittest.skipIf(vision is None, 'GCP dependencies are not installed')
    45  class VisionMlTestIT(unittest.TestCase):
    46    def test_text_detection_with_language_hint(self):
    47      IMAGES_TO_ANNOTATE = [
    48          'gs://apache-beam-samples/advanced_analytics/vision/sign.jpg'
    49      ]
    50  
    51      IMAGE_CONTEXT = [vision.ImageContext({'language_hints': ['en']})]
    52  
    53      with TestPipeline(is_integration_test=True) as p:
    54        contexts = p | 'Create context' >> beam.Create(
    55            dict(zip(IMAGES_TO_ANNOTATE, IMAGE_CONTEXT)))
    56  
    57        output = (
    58            p
    59            | beam.Create(IMAGES_TO_ANNOTATE)
    60            | AnnotateImage(
    61                features=[
    62                    vision.Feature({'type_': vision.Feature.Type.TEXT_DETECTION})
    63                ],
    64                context_side_input=beam.pvalue.AsDict(contexts))
    65            | beam.ParDo(extract))
    66  
    67      assert_that(
    68          output,
    69          equal_to([
    70              'WAITING?\nPLEASE\nTURN OFF\nYOUR\nENGINE',
    71              'WAITING?',
    72              'PLEASE',
    73              'TURN',
    74              'OFF',
    75              'YOUR',
    76              'ENGINE'
    77          ]))
    78  
    79  
    80  if __name__ == '__main__':
    81    unittest.main()