fsblog-20170302-zmq_utils.py

import numpy as np
import json
import zmq


def send_image(socket, image, channel = b"image"):
    # Serialize a Numpy array
    dtype = str(image.dtype).encode('ascii')
    shape = json.dumps(image.shape).encode('ascii')
    data = image.tostring('C')

    socket.send_multipart([channel, dtype, shape, data])


def recv_image(socket, flags = 0):
    channel, dtype, shape, data = socket.recv_multipart(flags)

    # Deserialize a numpy array
    image = np.frombuffer(data, dtype=dtype.decode('ascii'))
    image.shape = json.loads(shape.decode('ascii'))
    return image