🚀 ニフティ’s Notion

import json
import base64
import boto3

# 入力されたテキストデータから機械学習エンドポイントに対して、リクエストを送信し、結果を返す
def query_endpoint(text):
    client = boto3.client('runtime.sagemaker')
    endpoint_name = '<エンドポイント名>'

    encoded_text = text.encode("utf-8")
    response = client.invoke_endpoint(EndpointName=endpoint_name, ContentType='application/x-text', Body=encoded_text, Accept='application/json;jpeg')

    return response

# 機械学習エンドポイントからの出力結果のJSONを解析し、画像を返す
def parse_response(query_response):
    response_dict = json.loads(query_response['Body'].read())
    return response_dict['generated_image']

# ここが実行される
def lambda_handler(event, context):
    body = json.loads(event['body'])
    text = body['text']
    response = query_endpoint(text)
    img = parse_response(response)

    return {
        'statusCode': 200,
        'headers': {
            'Content-Type': "'image/jpeg'",
        },
        'body': img,
        'isBase64Encoded': True
    }