import base64
import cv2
import numpy as np
import mxnet as mx
import mxnet.ndarray as F
import png

def decode_request(request):
	data = request.files['Body'].read()
	b64_bytes = base64.b64encode(data)
	b64_string = b64_bytes.decode()
	return F.array(cv2.imdecode(np.fromstring(base64.b64decode(b64_string)), 1).astype(np.float32))

def prep_batch(img):
	img = F.expand_dims(F.transpose(img, (2, 0, 1)), 0)
	img = F.sum_axis(F.array([[[[0.3]], [[0.59]],[[0.11]]]]) * img, 1, keepdims=True)
	return mx.io.DataBatch([img])

def postprocess(raw_output):
	mask = np.argmax(raw_output, axis=(1))[0]*85
	png.from_array(mask.astype(np.uint8), 'L').save('/tmp/mask.png')
	return