import os
import json
from base64 import b64encode, b64decode
from matplotlib import pyplot as plt

import mxnet as mx
import sagemaker
from gluoncv.utils import download, viz

# 이미지를 로컬에 다운로드
download('https://sportshub.cbsistatic.com/i/r/2019/11/15/10869f78-1378-4aa5-b36b-085607ae3387/thumbnail/770x433/f3276ac966a56b7cb45987869098cddb/lionel-messi-argentina-brazil.jpg', path='messi.jpg')

# 이미지를 바이트스트링으로 읽는다
bimage = None
with open('messi.jpg', 'rb') as fp:
    bimage = fp.read()

    
# 보내는 데이터를 serialization
def serializer(data): 
    return json.dumps(data).encode('utf-8')

# 받은 데이터를 deserialization
def deserializer(body, content_type):
    return json.loads(body.read().decode('utf-8'))


predictor = sagemaker.predictor.RealTimePredictor(
    endpoint='sagemaker-yolo3-3', 
    content_type='application/json', 
    accept='application/json',
    serializer=serializer,
    deserializer=deserializer,
)

# 보낼 이미지를 b64로 인코딩한다
s = b64encode(bimage).decode('utf-8')

# 파라미터는 short, image 두개만 사용한다. short 는 이미지의 짧은 부분 길이
res = predictor.predict({
    'short': 320,
    'image': s
})
print(res['shape'])

# 화면 출력용 코드
ax = viz.plot_bbox(mx.image.imresize(mx.image.imdecode(bimage), res['shape'][3], res['shape'][2]), mx.nd.array(res['bbox']), mx.nd.array(res['score']), mx.nd.array(res['cid']), class_names=['person'])
plt.show()