import greengrasssdk
import io
from processing import *
from load_model import load_model
import json
from flask import Flask, request, send_file

client = greengrasssdk.client('iot-data')

app = Flask(__name__)

@app.route('/', methods=['POST'])
def transform():
	client.publish(topic='inference/brain_segmentation', payload='Image received.')
	try:
		img = decode_request(request)
		batch = prep_batch(img)
		net.forward(batch)
		raw_output = net.get_outputs()[0].asnumpy()
		postprocess(raw_output)
		client.publish(topic='inference/brain_segmentation', payload='Image processed.')
		with open('/tmp/mask.png', 'rb') as img:
			return send_file(io.BytesIO(img.read()),
							 attachment_filename='/tmp/mask.png',
							 mimetype='image/png')
	except Exception as e:
		client.publish(topic='inference/brain_segmentation', payload='Error: %s.' % (str(e)))

model_path = '/greengrass-machine-learning/mxnet/segmentation-net/'
net = load_model(model_path, 'model', 0)
client.publish(topic='inference/brain_segmentation', payload='Model loaded.')
client.publish(topic='inference/brain_segmentation', payload='App starting.')

app.run(debug=True, host='0.0.0.0')

def function_handler(event, context):
    return