Serverless Deployment: Deploying a Hair Type Classifier

This document outlines the steps taken to prepare and deploy a Straight vs. Curly Hair Type classifier model using ONNX, including local inference and an attempt at Dockerization for a Lambda function.

1. Model Download and Inspection

We started by downloading the ONNX model files provided for the homework.

Commands:

PREFIX="https://github.com/alexeygrigorev/large-datasets/releases/download/hairstyle"
DATA_URL = f"{PREFIX}/hair_classifier_v1.onnx.data"
MODEL_URL = f"{PREFIX}/hair_classifier_v1.onnx"
!wget {DATA_URL}
!wget {MODEL_URL}

Output of Model Download:

... 'hair_classifier_v1.onnx.data' saved ...
... 'hair_classifier_v1.onnx' saved ...

Next, we needed to identify the input and output node names of the ONNX model to correctly interact with it. This was done using the onnx library.

Identifying Output Node:

!pip install onnx
import onnx

model_path = 'hair_classifier_v1.onnx'
onnx_model = onnx.load(model_path)
output_name = onnx_model.graph.output[0].name
print(f"The name of the output node is: {output_name}")

Output:

The name of the output node is: output

Identifying Input Node:

import onnx

model_path = 'hair_classifier_v1.onnx'
onnx_model = onnx.load(model_path)
input_name = onnx_model.graph.input[0].name
print(f"The name of the input node is: {input_name}")

Output:

The name of the input node is: input

2. Image Preparation and Preprocessing

To prepare an image for inference, we used utility functions to download and resize it, followed by preprocessing steps (converting to a tensor and normalizing) as established in the previous homework.

Utility Functions:

!pip install Pillow
from io import BytesIO
from urllib import request
from PIL import Image

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

Downloading and Resizing an Image:

We used the image https://habrastorage.org/webt/yf/_d/ok/yf_dokzqy3vcritme8ggnzqlvwa.jpeg with a target size of (200, 200).

image_url = 'https://habrastorage.org/webt/yf/_d/ok/yf_dokzqy3vcritme8ggnzqlvwa.jpeg'
target_size = (200, 200)

img = download_image(image_url)
img_prepared = prepare_image(img, target_size)

# Optional: Display the prepared image
import matplotlib.pyplot as plt
plt.imshow(img_prepared)
plt.title(f"Prepared Image (Size: {img_prepared.size})")
plt.axis('off')
plt.show()

Preprocessing the Image:

The preprocessing involved converting the image to a PyTorch tensor and normalizing it using ImageNet's mean and standard deviation.

from torchvision import transforms
import numpy as np

mean = [0.485, 0.456, 0.406]
std = [0.229, 0.224, 0.225]

preprocess_transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize(mean=mean, std=std)
])

img_tensor = preprocess_transform(img_prepared)
img_array = img_tensor.numpy()

first_pixel_r_channel = img_array[0, 0, 0]
print(f"Value in the first pixel, R channel after preprocessing: {first_pixel_r_channel}")

Output:

Value in the first pixel, R channel after preprocessing: -1.0732940435409546

3. Local Model Inference

With the image prepared, we performed local inference using onnxruntime.

Performing Inference:

!pip install onnxruntime
import onnxruntime

sess = onnxruntime.InferenceSession('hair_classifier_v1.onnx')

input_name = sess.get_inputs()[0].name
output_name = sess.get_outputs()[0].name

input_data = np.expand_dims(img_tensor.numpy(), axis=0).astype(np.float32)

outputs = sess.run([output_name], {input_name: input_data})
model_output = outputs[0]

print(f"The output of the model is: {model_output}")

Output:

The output of the model is: [[0.09156641]]

4. Preparing Lambda Function Code

All the necessary logic was consolidated into a lambda_function.py file, ready for potential deployment or local testing.

lambda_function.py Content:

%%writefile lambda_function.py

import onnxruntime
import numpy as np
from io import BytesIO
from urllib import request
from PIL import Image
from torchvision import transforms

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

mean = [0.485, 0.456, 0.406]
std = [0.229, 0.224, 0.225]

preprocess_transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize(mean=mean, std=std)
])

MODEL_PATH = 'hair_classifier_v1.onnx'

def predict_image(image_url):
    img = download_image(image_url)
    target_size = (200, 200)
    img_prepared = prepare_image(img, target_size)

    img_tensor = preprocess_transform(img_prepared)
    input_data = np.expand_dims(img_tensor.numpy(), axis=0).astype(np.float32)

    sess = onnxruntime.InferenceSession(MODEL_PATH)

    input_name = sess.get_inputs()[0].name
    output_name = sess.get_outputs()[0].name

    outputs = sess.run([output_name], {input_name: input_data})

    return outputs[0]

def lambda_handler(event, context):
    if 'url' not in event:
        return {
            'statusCode': 400,
            'body': 'Error: Missing image URL in event.'
        }

    image_url = event['url']
    try:
        prediction = predict_image(image_url)
        return {
            'statusCode': 200,
            'body': {'prediction': prediction.tolist()}}
    except Exception as e:
        return {
            'statusCode': 500,
            'body': f'Error during prediction: {str(e)}'
        }

Local Test of Lambda Function:

from lambda_function import predict_image

image_url = 'https://habrastorage.org/webt/yf/_d/ok/yf_dokzqy3vcritme8ggnzqlvwa.jpeg'

prediction_output = predict_image(image_url)

print(f"Local prediction output: {prediction_output}")

Output:

Local prediction output: [[0.09156641]]

5. Dockerization (Question 6)

The goal was to extend a provided Docker base image, install dependencies, add the lambda code, and run it locally to score an image. The base image agrigorev/model-2025-hairstyle:v1 uses Python 3.13 and already contains a model (hair_classifier_empty.onnx).

Proposed Dockerfile.extended:

FROM agrigorev/model-2025-hairstyle:v1

RUN pip install Pillow==11.3.0 torchvision==0.18.0 onnx==1.20.0 onnxruntime==1.23.2

COPY lambda_function.py ${LAMBDA_TASK_ROOT}

ENV MODEL_PATH=hair_classifier_empty.onnx