-
Notifications
You must be signed in to change notification settings - Fork 67
/
google.py
123 lines (87 loc) · 4.11 KB
/
google.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
''' Google API-based feature extraction classes. '''
from pliers.extractors.image import ImageExtractor
from pliers.transformers import GoogleVisionAPITransformer
from pliers.extractors.base import ExtractorResult
import numpy as np
import pandas as pd
class GoogleVisionAPIExtractor(GoogleVisionAPITransformer, ImageExtractor):
''' Base class for all Extractors that use the Google Vision API. '''
VERSION = '1.0'
def _extract(self, stims):
request = self._build_request(stims)
responses = self._query_api(request)
results = []
for i, response in enumerate(responses):
if response and self.response_object in response:
raw = response[self.response_object]
results.append(ExtractorResult(None, stims[i], self, raw=raw))
elif 'error' in response:
raise Exception(response['error']['message'])
else:
results.append(ExtractorResult(None, stims[i], self, raw=[{}]))
return results
class GoogleVisionAPIFaceExtractor(GoogleVisionAPIExtractor):
''' Identifies faces in images using the Google Cloud Vision API. '''
request_type = 'FACE_DETECTION'
response_object = 'faceAnnotations'
def _to_df(self, result):
annotations = result.raw
if self.handle_annotations == 'first':
annotations = [annotations[0]]
face_results = []
for i, annotation in enumerate(annotations):
data_dict = {}
for field, val in annotation.items():
if 'Confidence' in field:
data_dict['face_' + field] = val
elif 'oundingPoly' in field:
for j, vertex in enumerate(val['vertices']):
for dim in ['x', 'y']:
name = '%s_vertex%d_%s' % (field, j+1, dim)
val = vertex[dim] if dim in vertex else np.nan
data_dict[name] = val
elif field == 'landmarks':
for lm in val:
name = 'landmark_' + lm['type'] + '_%s'
lm_pos = {name %
k: v for (k, v) in lm['position'].items()}
data_dict.update(lm_pos)
else:
data_dict[field] = val
face_results.append(data_dict)
return pd.DataFrame(face_results)
class GoogleVisionAPILabelExtractor(GoogleVisionAPIExtractor):
''' Labels objects in images using the Google Cloud Vision API. '''
request_type = 'LABEL_DETECTION'
response_object = 'labelAnnotations'
def _to_df(self, result):
res = {label['description']: label['score'] for label in result.raw}
return pd.DataFrame([res])
class GoogleVisionAPIPropertyExtractor(GoogleVisionAPIExtractor):
''' Extracts image properties using the Google Cloud Vision API. '''
request_type = 'IMAGE_PROPERTIES'
response_object = 'imagePropertiesAnnotation'
def _to_df(self, result):
colors = result.raw['dominantColors']['colors']
data_dict = {}
for color in colors:
rgb = color['color']
data_dict[(rgb['red'], rgb['green'], rgb['blue'])] = color['score']
return pd.DataFrame([data_dict])
class GoogleVisionAPISafeSearchExtractor(GoogleVisionAPIExtractor):
''' Extracts safe search detection using the Google Cloud Vision API. '''
request_type = 'SAFE_SEARCH_DETECTION'
response_object = 'safeSearchAnnotation'
def _to_df(self, result):
return pd.DataFrame([result.raw])
class GoogleVisionAPIWebEntitiesExtractor(GoogleVisionAPIExtractor):
''' Extracts web entities using the Google Cloud Vision API. '''
request_type = 'WEB_DETECTION'
response_object = 'webDetection'
def _to_df(self, result):
data_dict = {}
if 'webEntities' in result.raw:
for entity in result.raw['webEntities']:
if 'description' in entity and 'score' in entity:
data_dict[entity['description']] = entity['score']
return pd.DataFrame([data_dict])