Add MAX Framework package (#17)

Author: bdwyer2

.travis.yml

@@ -1,17 +1,14 @@
 language: python
 python:
-- 3.6
-
+  - 3.6
 services:
-- docker
-
+  - docker
 install:
-- docker build -t max-resnet-50 .
-- docker run -it -d -p 5000:5000 max-resnet-50
-- sleep 30 # container needs a few seconds before it will accept requests
-
+  - docker build -t max-resnet-50 .
+  - docker run -it -d -p 5000:5000 max-resnet-50
+  - pip install pytest requests flake8
 before_script:
-- pip install pytest requests
-
+  - flake8 . --max-line-length=127
+  - sleep 30 # container needs a few seconds before it will accept requests
 script:
-- pytest tests/test.py
+  - pytest tests/test.py

Dockerfile

@@ -1,4 +1,4 @@
-FROM codait/max-base:v1.0.0
+FROM codait/max-base:v1.1.0
 
 ARG model_bucket=http://max-assets.s3-api.us-geo.objectstorage.softlayer.net/keras
 ARG model_file=resnet50.h5
@@ -10,7 +10,9 @@ COPY requirements.txt /workspace
 RUN pip install -r requirements.txt
 
 COPY . /workspace
-RUN md5sum -c md5sums.txt # check file integrity
+
+# check file integrity
+RUN md5sum -c md5sums.txt
 
 EXPOSE 5000
 

api/__init__.py

@@ -1,12 +1,2 @@
-from flask_restplus import Api
-
-from config import API_TITLE, API_VERSION, API_DESC
-from .model import api as model_ns
-
-api = Api(
-	title=API_TITLE,
-	version=API_VERSION,
-	description=API_DESC)
-
-api.namespaces.clear()
-api.add_namespace(model_ns)
+from .metadata import ModelMetadataAPI  # noqa
+from .predict import ModelPredictAPI  # noqa

api/metadata.py

@@ -0,0 +1,10 @@
+from core.model import ModelWrapper
+from maxfw.core import MAX_API, MetadataAPI, METADATA_SCHEMA
+
+
+class ModelMetadataAPI(MetadataAPI):
+
+    @MAX_API.marshal_with(METADATA_SCHEMA)
+    def get(self):
+        """Return the metadata associated with the model"""
+        return ModelWrapper.MODEL_META_DATA

api/model.py

@@ -0,0 +1,41 @@
+from flask_restplus import fields
+from werkzeug.datastructures import FileStorage
+from maxfw.core import MAX_API, PredictAPI
+from core.model import ModelWrapper
+
+input_parser = MAX_API.parser()
+input_parser.add_argument('image', type=FileStorage, location='files', required=True, help="An image file")
+
+label_prediction = MAX_API.model('LabelPrediction', {
+    'label_id': fields.String(required=False, description='Class label identifier'),
+    'label': fields.String(required=True, description='Class label'),
+    'probability': fields.Float(required=True)
+})
+
+predict_response = MAX_API.model('ModelPredictResponse', {
+    'status': fields.String(required=True, description='Response status message'),
+    'predictions': fields.List(fields.Nested(label_prediction), description='Predicted labels and probabilities')
+})
+
+
+class ModelPredictAPI(PredictAPI):
+
+    model_wrapper = ModelWrapper()
+
+    @MAX_API.doc('predict')
+    @MAX_API.expect(input_parser)
+    @MAX_API.marshal_with(predict_response)
+    def post(self):
+        """Make a prediction given input data"""
+        result = {'status': 'error'}
+
+        args = input_parser.parse_args()
+        image_data = args['image'].read()
+        image = self.model_wrapper.read_image(image_data)
+        preds = self.model_wrapper.predict(image)
+
+        label_preds = [{'label_id': p[0], 'label': p[1], 'probability': p[2]} for p in [x for x in preds]]
+        result['predictions'] = label_preds
+        result['status'] = 'ok'
+
+        return result

api/predict.py

@@ -1,14 +1,8 @@
-import os
-from flask import Flask
-from api import api
+from maxfw.core import MAXApp
+from api import ModelMetadataAPI, ModelPredictAPI
+from config import API_TITLE, API_DESC, API_VERSION
 
-app = Flask(__name__)
-# load default config
-app.config.from_object('config')
-# load override config if exists
-if 'APP_CONFIG' in os.environ:
-	app.config.from_envvar('APP_CONFIG')
-api.init_app(app)
-
-if __name__ == '__main__':
-	app.run(host='0.0.0.0')
+max_app = MAXApp(API_TITLE, API_DESC, API_VERSION)
+max_app.add_api(ModelMetadataAPI, '/metadata')
+max_app.add_api(ModelPredictAPI, '/predict')
+max_app.run()

app.py

@@ -1,33 +1,15 @@
 # Application settings
 
-# Flask settings 
+# Flask settings
 DEBUG = False
 
 # Flask-restplus settings
 RESTPLUS_MASK_SWAGGER = False
 SWAGGER_UI_DOC_EXPANSION = 'none'
 
 # API metadata
-API_TITLE = 'Model Asset Exchange Server'
-API_DESC = 'An API for serving models'
+API_TITLE = 'MAX ResNet 50'
+API_DESC = 'Identify objects in images using a first-generation deep residual network.'
 API_VERSION = '0.1'
 
-# Model settings
-models = {
-	'resnet50': {'size': (224, 224), 'license': 'MIT'}
-}
-
-# default model
-MODEL_NAME = 'resnet50'
-DEFAULT_MODEL_PATH = 'assets/{}.h5'.format(MODEL_NAME)
-# for image models, may not be required
-MODEL_INPUT_IMG_SIZE = models[MODEL_NAME]['size']
-MODEL_LICENSE = models[MODEL_NAME]['license']
-
-MODEL_META_DATA = {
-    'id': '{}-keras-imagenet'.format(MODEL_NAME.lower()),
-    'name': '{} Keras Model'.format(MODEL_NAME),
-    'description': '{} Keras model trained on ImageNet'.format(MODEL_NAME),
-    'type': 'image_classification',
-    'license': '{}'.format(MODEL_LICENSE)
-}
+DEFAULT_MODEL_PATH = 'assets/resnet50.h5'

config.py

@@ -0,0 +1,53 @@
+import io
+import logging
+from PIL import Image
+from keras.backend import clear_session
+from keras import models
+from keras.preprocessing.image import img_to_array
+from keras.applications import imagenet_utils
+import numpy as np
+from maxfw.model import MAXModelWrapper
+from config import DEFAULT_MODEL_PATH
+
+
+logger = logging.getLogger()
+
+
+class ModelWrapper(MAXModelWrapper):
+    """Model wrapper for Keras models"""
+
+    MODEL_NAME = 'resnet50'
+    MODEL_INPUT_IMG_SIZE = (224, 224)
+    MODEL_LICENSE = 'MIT'
+    MODEL_MODE = 'caffe'
+    MODEL_META_DATA = {
+        'id': '{}-keras-imagenet'.format(MODEL_NAME.lower()),
+        'name': '{} Keras Model'.format(MODEL_NAME),
+        'description': '{} Keras model trained on ImageNet'.format(MODEL_NAME),
+        'type': 'image_classification',
+        'license': '{}'.format(MODEL_LICENSE),
+        'source': 'https://developer.ibm.com/exchanges/models/all/max-resnet-50/'
+    }
+
+    def __init__(self, path=DEFAULT_MODEL_PATH):
+        logger.info('Loading model from: {}...'.format(path))
+        clear_session()
+        self.model = models.load_model(path)
+        # this seems to be required to make Keras models play nicely with threads
+        self.model._make_predict_function()
+        logger.info('Loaded model: {}'.format(self.model.name))
+
+    def read_image(self, image_data):
+        return Image.open(io.BytesIO(image_data))
+
+    def _pre_process(self, image):
+        image = image.resize(self.MODEL_INPUT_IMG_SIZE)
+        image = img_to_array(image)
+        image = np.expand_dims(image, axis=0)
+        return imagenet_utils.preprocess_input(image, mode=self.MODEL_MODE)
+
+    def _post_process(self, preds):
+        return imagenet_utils.decode_predictions(preds)[0]
+
+    def _predict(self, x):
+        return self.model.predict(x)