def main(host, port):
ctx = zmq.Context()
sock2 = ctx.socket(zmq.REP)
sockstr = "tcp://%s:%d" % (host, port)
sock2.bind(sockstr)
print ("Bound to %s" % sockstr)
while True:
msg = sock2.recv_json()
print "Message received! ",
# print(msg)
initialize(msg["path"])
keys = msg["keys"]
# print keys
if "size" in msg: # If client asks for size, return only the size
send_array(sock2, np.array(rget(file, keys[0]).shape[0]))
print "Sending data size..."
continue
if "batch_size" in msg:
N = rget(file, keys[0]).shape[0]
bn = msg["batch_num"]
batch_size = msg["batch_size"]
start = (bn * batch_size) % N
end = ((bn + 1) * batch_size - 1) % N + 1
sl = slice(start, end)
else:
sl = slice(None)
print ("Sending batch %d" % bn)
for ind, k in enumerate(keys):
data = rget(file, k)[sl]
if ind < len(keys) - 1:
send_array(sock2, data, flags=zmq.SNDMORE)
else:
send_array(sock2, data)
valid[start: end] = True
file.flush()
while True:
msg = sock2.recv_json()
print(msg)
if msg.get('command') == 'get_valid':
va = np.asarray(valid)
undone = (va == 0).nonzero()[0]
if len(undone) > 0:
bn = undone[0] / BATCH_SIZE
make_new_batch(bn)
va = np.asarray(valid).nonzero()[0]
batches = np.unique(np.floor(va / BATCH_SIZE)).astype(np.int)
send_array(sock2, batches)
elif msg.get('command') == 'get_info':
bn = 0
make_new_batch(bn)
infopath = os.path.join(infodir, str(bn) + '.json')
infolist = json.loads(open(infopath).read())
msg = {'info': infolist[1],
'image_size': 256,
'num_channels': 3}
sock2.send_json(msg)
else:
bn = msg['batch_num']
bsize = BATCH_SIZE
start = (bn * bsize) % N
end = ((bn + 1) * bsize - 1) % N + 1
make_new_batch(bn)
