ITの隊長のブログ

ITの隊長のブログです。Rubyを使って仕事しています。最近も色々やっているお(^ω^ = ^ω^)

LambdaからSageMakerを叩いたときのスクリプトメモ

スポンサードリンク

import boto3
import base64
import uuid
import json


def convert_b64_string_to_bynary(s):
    """base64をデコードする"""
    return base64.b64decode(s.split(',')[1].encode("UTF-8"))
    
    
def query_endpoint(base_64ed_image):
    img = convert_b64_string_to_bynary(base_64ed_image)
    print(type(img))
    endpoint_name = 'jumpstart-ftc-tf-ic-resnet-50-classification-1-1'
    client = boto3.client('runtime.sagemaker')
    response = client.invoke_endpoint(EndpointName=endpoint_name, ContentType='application/x-image', Body=img)
    model_predictions = json.loads(response['Body'].read())['predictions'][0]
    return model_predictions


def lambda_handler(event, context):
    # requestbodeyの中のjsonはeventに辞書型に変化されて保存されている
    # なので、eventの中には {"mypng": "base64でエンコードされた文字列"}が入力される想定。
    base_64ed_image = event['mypng']
    model_predictions = query_endpoint(base_64ed_image)
    # convert_b64_string_to_bynary(base_64ed_image)
    # とりあえずOKを返す
    
    return {
        'statusCode': 200,
        'body': json.dumps(model_predictions)
    }

もう記憶がないけど削除するのでメモ