annotate-maps.py

import base64
import os
import re
import sys

from googleapiclient import discovery
from googleapiclient import errors
import nltk
from nltk.stem.snowball import EnglishStemmer
from oauth2client.client import GoogleCredentials
import pickle

DISCOVERY_URL = 'https://{api}.googleapis.com/$discovery/rest?version={apiVersion}'  # noqa
BATCH_SIZE = 10


class VisionApi:
    """Construct and use the Google Vision API service."""

    def __init__(self, api_discovery_file='vision_api.json'):
        self.credentials = GoogleCredentials.get_application_default()
        self.service = discovery.build(
            'vision', 'v1', credentials=self.credentials,
            discoveryServiceUrl=DISCOVERY_URL)

    def safe_search(self, input_filenames, num_retries=3, max_results=6):
        """Uses the Vision API to detect unsafe content in a given image.
        """
        images = {}
        for filename in input_filenames:
            with open(filename, 'rb') as image_file:
                images[filename] = image_file.read()

        batch_request = []
        for filename in images:
            batch_request.append({
                'image': {
                    'content': base64.b64encode(
                            images[filename]).decode('UTF-8')
                },
                'features': [{
                    'type': 'SAFE_SEARCH_DETECTION',
                    'maxResults': max_results,
                }]
            })
        request = self.service.images().annotate(
            body={'requests': batch_request})

        try:
            responses = request.execute(num_retries=num_retries)
            if 'responses' not in responses:
                return {}
            text_response = {}
            for filename, response in zip(images, responses['responses']):
                if 'error' in response:
                    print("API Error for %s: %s" % (
                            filename,
                            response['error']['message']
                            if 'message' in response['error']
                            else ''))
                    continue
                if 'safeSearchAnnotation' in response:
                    text_response[filename] = response['safeSearchAnnotation']
                else:
                    text_response[filename] = []
            return text_response
        except errors.HttpError as e:
            print("Http Error for %s: %s" % (filename, e))
        except KeyError as e2:
            print("Key error: %s" % e2)

print("OK, here we go.")

api = VisionApi()

# load any previous responses
responses = pickle.load(open( "responses.p", "r"))

# process 100 blocks of 10 images each
for j in range(0,100):
    files = []

    # make array of 10 scraped images
    for i in range(j*10,(j+1)*10):
        files.append(str(i) + ".png")

    # apply the Google Vision API
    r = api.safe_search(files)

    responses.update(r)

# save new responses
pickle.dump( responses, open( "responses.p", "wb" ) )