PyTorch张量转FastAPI返回图像失败，问题出在哪？

你的第一种写法核心错误在于：torch.save() 是用来**序列化PyTorch对象（比如张量、模型）**的工具，它输出的是PyTorch专属的.pt格式字节数据，并非浏览器能识别的标准JPEG图像字节流。浏览器拿到这种序列化数据后，无法解析成图像，自然无法正常显示。

而第二种写法里，image.save(return_image, "JPEG") 是将PIL图像编码成了标准的JPEG格式字节流，符合image/jpeg媒体类型的要求，所以能被正确渲染。

正确的转换流程 #

要把PyTorch张量转换成可返回的JPEG字节流，需要先将张量转换成PIL图像，再编码为JPEG格式，步骤如下：

• 调整张量维度：PyTorch图像张量通常是[C, H, W]（通道、高度、宽度）的格式，而PIL图像要求是[H, W, C]，需要用permute调整。
• 修正数值范围：如果张量是归一化后的[0,1]或[-1,1]，要转回[0,255]的uint8类型。
• 转换为PIL图像，再保存为JPEG字节流。

修正后的代码 #

import io
from PIL import Image
import torch
from starlette.responses import StreamingResponse

def some_unimportant_function(params):
    # 假设some_img是PyTorch图像张量，格式为[C, H, W]
    some_img = ...  # 你的图像张量
    
    # 处理张量：调整维度+转回0-255的uint8
    # 适配不同数值范围的张量
    if some_img.min() < 0:
        img_tensor = (some_img + 1) / 2  # 从[-1,1]转到[0,1]
    else:
        img_tensor = some_img.clamp(0, 1)  # 确保在[0,1]范围
    img_tensor = (img_tensor * 255).byte().permute(1, 2, 0)
    
    # 转成PIL图像
    pil_img = Image.fromarray(img_tensor.numpy())
    
    # 保存为JPEG字节流
    return_image = io.BytesIO()
    pil_img.save(return_image, "JPEG")
    return_image.seek(0)
    
    return StreamingResponse(content=return_image, media_type="image/jpeg")

内容的提问来源于stack exchange，提问作者Just Another Programmer