Skip to content

Commit

Permalink
fix: add numpy json encoder
Browse files Browse the repository at this point in the history
  • Loading branch information
MrBlenny committed Jun 9, 2022
1 parent 594564c commit 26d8ac4
Showing 1 changed file with 17 additions and 1 deletion.
18 changes: 17 additions & 1 deletion nuclio_lighting_flash/main.py
Original file line number Diff line number Diff line change
Expand Up @@ -3,12 +3,24 @@
import base64
from PIL import Image
import io
import numpy as np
import yaml

from flash.image import ObjectDetector
from flash_model_handler import FlashModelHandler


class NpEncoder(json.JSONEncoder):
def default(self, obj):
if isinstance(obj, np.integer):
return int(obj)
if isinstance(obj, np.floating):
return float(obj)
if isinstance(obj, np.ndarray):
return obj.tolist()
return json.JSONEncoder.default(self, obj)


def init_context(context):
context.logger.info("Init context... 0%")

Expand All @@ -20,6 +32,10 @@ def init_context(context):
labels_spec = annotations["spec"]
labels = {item["id"]: item["name"] for item in json.loads(labels_spec)}

print(f"Model head: {annotations['head']}")
print(f"Model backbone: {annotations['backbone']}")
print(f"Num classes: {len(labels)}")

# Read the DL model
# Either "checkpoint_path" or "head" and "backbone" should be specified
model = (
Expand Down Expand Up @@ -48,7 +64,7 @@ def handler(context, event):
results = context.user_data.model.infer(image, threshold)

return context.Response(
body=json.dumps(results),
body=json.dumps(results, cls=NpEncoder),
headers={},
content_type="application/json",
status_code=200,
Expand Down

0 comments on commit 26d8ac4

Please sign in to comment.