forked from cvat-ai/cvat
-
Notifications
You must be signed in to change notification settings - Fork 0
/
Copy pathmain.py
27 lines (21 loc) · 766 Bytes
/
main.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
import json
import base64
from PIL import Image
import io
from model_handler import ModelHandler
def init_context(context):
context.logger.info("Init context... 0%")
# Read the DL model
model = ModelHandler()
setattr(context.user_data, 'model', model)
context.logger.info("Init context...100%")
def handler(context, event):
context.logger.info("Run SiamMask model")
data = event.body
buf = io.BytesIO(base64.b64decode(data["image"].encode('utf-8')))
shape = data.get("shape")
state = data.get("state")
image = Image.open(buf)
results = context.user_data.model.infer(image, shape, state)
return context.Response(body=json.dumps(results), headers={},
content_type='application/json', status_code=200)