AWS Lambda调用语义分割模型返回RecordIO-Protobuf而非图片的问题

问题原因 #

SageMaker内置或通过SDK训练的语义分割模型，在Notebook中调用时，SageMaker SDK会自动处理application/x-recordio-protobuf格式的解析与图片转换；但直接用boto3的invoke_endpoint调用时，模型端点默认返回原始推理结果的Protobuf格式，而非可视化图片，因为推理脚本未配置为返回图片类型的响应。

解决方案1：修改模型推理脚本，让端点直接返回图片 #

调整模型的inference.py脚本，将推理后的分割掩码转换为JPEG图片，并设置响应格式为image/jpeg。重新部署模型端点后，Lambda调用即可直接获取图片响应。

示例推理脚本（inference.py） #

import numpy as np
from PIL import Image
import io

def model_fn(model_dir):
    # 加载你的语义分割模型（根据实际训练方式调整）
    # 示例：如果是TensorFlow模型，用tf.keras.models.load_model(model_dir)
    pass

def input_fn(request_body, request_content_type):
    # 解析输入的JPEG图片
    if request_content_type == 'image/jpeg':
        img = Image.open(io.BytesIO(request_body))
        return np.array(img)
    raise ValueError(f"不支持的输入格式: {request_content_type}")

def predict_fn(input_data, model):
    # 执行模型推理，得到分割掩码（shape通常为(H,W,num_classes)或(num_classes,H,W)）
    mask = model.predict(input_data)
    # 转换为单通道类别索引掩码（取概率最高的类别）
    if len(mask.shape) == 3 and mask.shape[-1] > 1:
        mask = np.argmax(mask, axis=-1)
    elif len(mask.shape) == 3 and mask.shape[0] > 1:
        mask = np.argmax(mask, axis=0)
    # 定义类别颜色映射（根据你的模型类别调整）
    color_map = {
        0: (0, 0, 0),    # 背景
        1: (255, 0, 0),  # 类别1
        2: (0, 255, 0)   # 类别2
    }
    # 转换为彩色图片
    mask_color = np.zeros((mask.shape[0], mask.shape[1], 3), dtype=np.uint8)
    for cls_idx, color in color_map.items():
        mask_color[mask == cls_idx] = color
    return mask_color

def output_fn(prediction, accept):
    # 根据Accept头返回对应格式
    if accept == 'image/jpeg':
        img = Image.fromarray(prediction)
        buf = io.BytesIO()
        img.save(buf, format='JPEG')
        return buf.getvalue(), 'image/jpeg'
    raise ValueError(f"不支持的输出格式: {accept}")

部署更新后的模型 #

将修改后的inference.py与模型文件打包成.tar.gz，上传至S3，然后更新SageMaker模型端点的模型数据源，重新部署端点。之后Lambda原代码即可直接获取JPEG图片响应。

解决方案2：在Lambda内解析Protobuf并转换为图片 #

如果无法修改模型端点，可在Lambda中解析返回的Protobuf数据，将分割掩码转换为图片。需创建包含依赖库的Lambda层，因为Lambda默认不包含解析Protobuf所需的库。

步骤1：创建Lambda层 #

1. 在本地创建Python环境，安装依赖：

   mkdir python && cd python
   pip install mxnet protobuf -t . --platform manylinux2014_x86_64 --only-binary=:all:
   cd .. && zip -r mxnet-proto-layer.zip python/
   

2. 登录AWS控制台，进入Lambda的「层」页面，上传mxnet-proto-layer.zip创建层，关联到你的Lambda函数。

步骤2：Lambda处理代码示例 #

import boto3
import io
import base64
import numpy as np
from PIL import Image
import mxnet as mx
from mxnet.recordio import MXRecordIO

def lambda_handler(event, context):
    runtime = boto3.client('runtime.sagemaker')
    ENDPOINT_NAME = '你的端点名称'
    
    # 解析输入的图片（假设请求body是base64编码的JPEG）
    image_body = base64.b64decode(event['body'])
    image_buf = io.BytesIO(image_body)
    
    # 调用模型端点
    response = runtime.invoke_endpoint(
        EndpointName=ENDPOINT_NAME,
        ContentType='image/jpeg',
        Body=image_buf,
        Accept='application/x-recordio-protobuf'
    )
    
    # 解析RecordIO Protobuf
    payload = response['Body'].read()
    rec_buf = io.BytesIO(payload)
    rec = MXRecordIO(rec_buf, 'r')
    item = rec.read()
    _, mask_data = mx.recordio.unpack(item)
    
    # 转换为numpy数组（根据模型输出形状调整reshape参数）
    # 示例：模型输出为(1, num_classes, H, W)，需转换为(H,W,num_classes)
    mask = mx.ndarray.frombuffer(mask_data, dtype=np.float32).asnumpy()
    mask = mask.reshape((3, 512, 512)).transpose((1, 2, 0))  # 替换为你的模型输出维度
    mask = np.argmax(mask, axis=2)  # 取类别索引
    
    # 转换为彩色图片
    color_map = {
        0: (0, 0, 0),
        1: (255, 0, 0),
        2: (0, 255, 0)
    }
    mask_color = np.zeros((mask.shape[0], mask.shape[1], 3), dtype=np.uint8)
    for cls_idx, color in color_map.items():
        mask_color[mask == cls_idx] = color
    
    # 转换为JPEG并返回
    img = Image.fromarray(mask_color)
    output_buf = io.BytesIO()
    img.save(output_buf, format='JPEG')
    jpeg_bytes = output_buf.getvalue()
    
    return {
        'statusCode': 200,
        'headers': {'Content-Type': 'image/jpeg'},
        'body': base64.b64encode(jpeg_bytes).decode('utf-8'),
        'isBase64Encoded': True
    }

注意事项 #

• 需根据你的模型实际输出形状调整reshape和transpose的参数。
• 若模型不是基于MXNet，可能需要使用对应框架的Protobuf解析工具（如TensorFlow的tf.io.parse_single_example）。

内容的提问来源于stack exchange，提问作者user3466555