Code for the lambda function

We have already written the code for the lambda function; we just need to combine it into a single script named hw.py.

from io import BytesIO
from urllib import request

from PIL import Image

import numpy as np
import os

import tflite_runtime.interpreter as tflite

MODEL_NAME = os.getenv('MODEL_NAME', 'model_2024_hairstyle_v2.tflite')

interpreter = tflite.Interpreter(model_path=MODEL_NAME)
interpreter.allocate_tensors()
		
input_details = interpreter.get_input_details()
input_index = input_details[0]['index']
		
output_details = interpreter.get_output_details()
output_index = output_details[0]['index']

def download_image(url):
    with request.urlopen(url) as resp:
        buffer = resp.read()
    stream = BytesIO(buffer)
    img = Image.open(stream)
    return img
    
def prepare_image(img, target_size):
    if img.mode != 'RGB':
        img = img.convert('RGB')
    img = img.resize(target_size, Image.NEAREST)
    return img
    
def prepare_input(x):
    return x / 255.0    
    
def predict(X):
    interpreter.set_tensor(input_index, X)
    interpreter.invoke()
    preds = interpreter.get_tensor(output_index)
    return float(preds[0,0])
    
def decode_predictions(pred):
	return ["curly", "straight"][pred > 0.5]
		
def lambda_handler(event, context):
    url = event['url']
    img = download_image(url)
    img = prepare_image(img, target_size=(200, 200))
    x = np.array(img, dtype='float32')
    X = np.array([x])
    X = prepare_input(X)
    preds = predict(X)
    results = decode_predictions(preds)
    return {"prediction" : results}

lambda_handler() is the function invoked by the AWS Lambda environment. The event parameter contains all the information we pass to the lambda function in our request. The context parameter is typically not used.

Preparing the Docker Image

First, create a file named Dockerfile:

FROM agrigorev/model-2024-hairstyle:v3

RUN pip install Pillow
RUN pip install <https://github.com/alexeygrigorev/tflite-aws-lambda/raw/main/tflite/tflite_runtime-2.14.0-cp310-cp310-linux_x86_64.whl>
RUN pip install numpy==1.23.1

COPY hw.py .
ENV MODEL_NAME=model_2024_hairstyle_v2.tflite

CMD [ "hw.lambda_handler" ]

Let’s build this image:

docker build -t hair-type-prediction .

Next, we need to check that the lambda function works. Let’s run the image:

docker run -it --rm -p 8080:8080 hair-type-prediction:latest

image.png

It’s running! We can test it now using the following script named test.py :

import requests

url = '<http://localhost:8080/2015-03-31/functions/function/invocations>'

data = {'url':'<https://habrastorage.org/webt/yf/_d/ok/yf_dokzqy3vcritme8ggnzqlvwa.jpeg>'}
print(requests.post(url, json=data).json())

image.png