Receive image from Cloud PubSub

recv.py

def receive_image(project_id, sub_name, prefix, extension, duration, binary):
    """Receieve images transmitted to a PubSub subscription."""
    import base64
    import binascii
    import io
    import time

    from google.cloud import pubsub
    from PIL import Image

    subscriber = pubsub.SubscriberClient()
    subscription_path = subscriber.subscription_path(project_id, sub_name)

    global count
    count = 0
    file_pattern = '{}-{}.{}'

    def callback(message):
        global count
        try:
            count = count + 1
            print('Received image {}:'.format(count))
            image_data = message.data
            if not binary:
                image_data = base64.b64decode(image_data)

            filename = file_pattern.format(prefix, count, extension)
            with io.open(filename, 'wb') as f:
                f.write(image_data)
                message.ack()
                with Image.open(filename) as img:
                    img.show()

        except binascii.Error:
            message.ack()  # To move forward if a message can't be processed

    subscriber.subscribe(subscription_path, callback=callback)

    sleep_count = 0
    print('Listening for messages on {}'.format(subscription_path))
    while sleep_count < duration:
        time.sleep(1)
        sleep_count = sleep_count + 1