Add PNG support and error handling to image read (#19)

* Add PNG support and error handling to image read

* Placate flake8

* Placate flake8 take 2

* Placate flake8 take 3

* Placate flake8 take 4

* Add Pillow dep for travis tests

* Uncomment image conversion code

Author: MLnick

.travis.yml

@@ -11,7 +11,7 @@ install:
   - sleep 45
 
 before_script:
-  - pip install pytest requests flake8
+  - pip install pytest requests flake8 Pillow
   - flake8 . --max-line-length=127
 
 script:

core/model.py

@@ -6,6 +6,7 @@
 import io
 import numpy as np
 import logging
+from flask import abort
 from config import DEFAULT_MODEL_PATH, MODEL_INPUT_IMG_SIZE, MODEL_META_DATA as model_meta
 from maxfw.model import MAXModelWrapper
 
@@ -26,8 +27,12 @@ def __init__(self, path=DEFAULT_MODEL_PATH):
         logger.info('Loaded model: {}'.format(self.model.name))
 
     def _read_image(self, image_data):
-        image = Image.open(io.BytesIO(image_data))
-        return image
+        try:
+            image = Image.open(io.BytesIO(image_data)).convert('RGB')
+            return image
+        except IOError as e:
+            logger.error(str(e))
+            abort(400, "The provided input is not a valid image (PNG or JPG required).")
 
     def _pre_process(self, image, target, mode='tf'):
         image = image.resize(target)

tests/test.py

@@ -1,5 +1,7 @@
 import pytest
 import requests
+from PIL import Image
+import tempfile
 
 
 def test_swagger():
@@ -29,20 +31,47 @@ def test_metadata():
     assert metadata['license'] == 'Apache v2'
 
 
+def _check_predict(r):
+
+    assert r.status_code == 200
+    response = r.json()
+    assert response['status'] == 'ok'
+    assert response['predictions'][0]['label_id'] == 'n02123045'
+    assert response['predictions'][0]['label'] == 'tabby'
+    assert response['predictions'][0]['probability'] > 0.6
+
+
 def test_predict():
+
+    formats = ['JPEG', 'PNG']
     model_endpoint = 'http://localhost:5000/model/predict'
     file_path = 'assets/cat.jpg'
+    jpg = Image.open(file_path)
+
+    for f in formats:
+        temp = tempfile.TemporaryFile()
+        if f == 'PNG':
+            jpg.convert('RGBA').save(temp, f)
+        else:
+            jpg.save(temp, f)
+        temp.seek(0)
+        file_form = {'image': (file_path, temp, 'image/{}'.format(f.lower()))}
+        r = requests.post(url=model_endpoint, files=file_form)
+        _check_predict(r)
+
+
+def test_invalid_input():
+
+    model_endpoint = 'http://localhost:5000/model/predict'
+    file_path = 'assets/README.md'
 
     with open(file_path, 'rb') as file:
         file_form = {'image': (file_path, file, 'image/jpeg')}
         r = requests.post(url=model_endpoint, files=file_form)
 
-    assert r.status_code == 200
+    assert r.status_code == 400
     response = r.json()
-    assert response['status'] == 'ok'
-    assert response['predictions'][0]['label_id'] == 'n02123045'
-    assert response['predictions'][0]['label'] == 'tabby'
-    assert response['predictions'][0]['probability'] > 0.6
+    assert 'input is not a valid image' in response['message']
 
 
 if __name__ == '__main__':