✨ support npy format response

signcl · Mar 20, 2020 · 5ffeeac · 5ffeeac
1 parent 2a95ef9
commit 5ffeeac
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 2 deletions.
diff --git a/example.py b/example.py
@@ -61,7 +61,7 @@
 resp = requests.post(
     'http://localhost:8001/v1/predict',
     headers={
-        'Accept': 'application/plain+json',
+        'Accept': 'application/npy+json',
         'Content-Type': 'application/npy+msgpack',
     },
     data=msgpack.packb({

diff --git a/src/serve.py b/src/serve.py
@@ -143,11 +143,12 @@ def parse_type(tp):
     COMPATIBLE  = (
         ('base64', 'json'),
         ('plain',  'json'),
+        ('npy',  'json'),
 
         ('bytes',  'msgpack'),
         ('base64', 'msgpack'),
         ('plain',  'msgpack'),
-        ('npy',  'msgpack'),
+        ('npy',  'msgpack')
     )
 
     if (fmt, enc) not in COMPATIBLE:
@@ -218,6 +219,10 @@ def make_output(results):
     for k, v in results.items():
         if fmt == 'bytes':
             rst[k] = v.tobytes()
+        elif fmt == 'npy':
+            bytesdata = io.BytesIO()
+            np.save(bytesdata, v, allow_pickle=False)
+            rst[k] = bytesdata.getvalue().decode('latin-1')
         elif fmt == 'base64':
             rst[k] = base64.b64encode(v.tobytes()).decode('utf-8')
         elif fmt == 'plain':