import io
import boto3
import json
import numpy as np
from PIL import Image, ImageOps
import cv2
import random

runtime= boto3.client('runtime.sagemaker')
s3_client = boto3.client('s3', aws_access_key_id='xxx',aws_secret_access_key='xxx')

def handler(event,context):
  print("Received event")
  jsondata={}
  filename = str(random.randint(1,1001))+".jpg"
  print(filename)
  obj = s3_client.get_object(Bucket='coco-valid-images', Key=filename)
  object_content = obj['Body'].read()
  img=Image.open(io.BytesIO(object_content))
  res = cv2.resize(np.asarray(img), dsize=(224, 224))
  data = np.array(res)
  data = data.reshape([1] + list(data.shape))
  print(data.shape)
  jsondata['instances']=data.tolist()
  payload = json.dumps(jsondata)

  response = runtime.invoke_endpoint(EndpointName='tf-mobilenet-v2-coco-2018-03-29',
                            ContentType='application/json', Body=payload);
  result = json.loads(response['Body'].read().decode())

  return {
    'body': result,
    'headers': {
      'Content-Type': 'text/plain'
    },
    'statusCode': 200
  }