diff --git a/vision/google/cloud/vision/annotations.py b/vision/google/cloud/vision/annotations.py index 861ef7e4bd730..0b521378404e3 100644 --- a/vision/google/cloud/vision/annotations.py +++ b/vision/google/cloud/vision/annotations.py @@ -119,6 +119,7 @@ def _process_image_annotations(image): :returns: Dictionary populated with entities from response. """ return { + 'faces': _make_faces_from_pb(image.face_annotations), 'labels': _make_entity_from_pb(image.label_annotations), 'landmarks': _make_entity_from_pb(image.landmark_annotations), 'logos': _make_entity_from_pb(image.logo_annotations), diff --git a/vision/google/cloud/vision/face.py b/vision/google/cloud/vision/face.py index 0dc525a3bbe19..7fdda10a75ebd 100644 --- a/vision/google/cloud/vision/face.py +++ b/vision/google/cloud/vision/face.py @@ -17,21 +17,13 @@ from enum import Enum +from google.cloud.grpc.vision.v1 import image_annotator_pb2 + from google.cloud.vision.geometry import BoundsBase from google.cloud.vision.likelihood import Likelihood from google.cloud.vision.geometry import Position -LIKELIHOOD_PB = { - 0: 'UNKNOWN', - 1: 'VERY_UNLIKELY', - 2: 'UNLIKELY', - 3: 'POSSIBLE', - 4: 'LIKELY', - 5: 'VERY_LIKELY', -} - - def _get_pb_likelihood(likelihood): """Convert gRPC Likelihood integer value to Likelihood instance. @@ -41,7 +33,7 @@ def _get_pb_likelihood(likelihood): :rtype: :class:`~google.cloud.vision.likelihood.Likelihood` :returns: Instance of ``Likelihood`` converted from gRPC value. """ - likelihood_pb = LIKELIHOOD_PB.get(likelihood) + likelihood_pb = image_annotator_pb2.Likelihood.Name(likelihood) return Likelihood[likelihood_pb] @@ -256,8 +248,7 @@ def from_pb(cls, response): detection_confidence = response.detection_confidence emotions = Emotions.from_pb(response) fd_bounds = FDBounds.from_pb(response.fd_bounding_poly) - headwear_likelihood = Likelihood[ - LIKELIHOOD_PB[response.headwear_likelihood]] + headwear_likelihood = _get_pb_likelihood(response.headwear_likelihood) image_properties = FaceImageProperties.from_pb(response) landmarks = Landmarks.from_pb(response.landmarks) landmarking_confidence = response.landmarking_confidence