如何在FastAPI中内存直出.nii.gz格式NIfTI图像

我需要在FastAPI服务端不落地存储，直接在内存中生成并发送.nii.gz格式的NIfTI图像。目前已实现内存上传3D CT扫描图像的功能，但返回分割掩码（.nii.gz格式）时遇到问题。业务流程：上传3D CT图像→完成器官分割→返回分割掩码的.nii.gz文件。

已实现的上传函数 #

import io
import gzip
import nibabel as nib
from fastapi import FastAPI, HTTPException, File, UploadFile
import time

@self.app.post("/input/aorta/image")
async def upload_image(file: UploadFile = File(...)):
    """Function used to upload a 3D CT scan"""
    start = time.time()

    # Load the file into a nifti image
    compressed_data = await file.read()
    compressed_stream = io.BytesIO(compressed_data)
    with gzip.GzipFile(
        fileobj=compressed_stream, mode="rb"
    ) as decompressed_stream:
        nifti_data = io.BytesIO(decompressed_stream.read())
    compressed_stream.close()
    nifti_image = nib.Nifti1Image.from_bytes(nifti_data.read())

    # Saving the nifti image
    self.filename = file.filename
    self.nii_image = nifti_image

    return {"message": "OK", "upload in (s)": time.time() - start}

自定义的NIfTI压缩函数 #

def nifti_zip(self, nifti_image):
    """ Function to compress an nifti image"""
    # Serialize the NIfTI image to bytes
    nifti_bytes = nifti_image.to_bytes()

    # Create a BytesIO object to store the compressed data
    compressed_nifti = io.BytesIO()

    # Use gzip to compress the NIfTI data in memory
    with gzip.GzipFile(fileobj=compressed_nifti, mode="wb") as compressed_stream:
        # Write the NIfTI data into the compressed stream
        compressed_stream.write(nifti_bytes)
    return compressed_nifti

四次测试情况 #

Test 1 → 错误：ValueError: stat: embedded null character in path #

@self.app.get("/output/aorta/mask")
async def aorta_mask():
    return FileResponse(
        nib.Nifti1Image.to_bytes(self.get_mask()),
        media_type="application/gzip",
        filename=self.filename,
    )

Test 2 → 错误：TypeError: a bytes-like object is required, not 'Nifti1Image' #

@self.app.get("/output/aorta/mask")
async def aorta_mask():
    return FileResponse(
        io.BytesIO(self.get_mask()),
        media_type="application/gzip",
        filename=self.filename,
    )

Test 3 → 错误：ValueError: stat: embedded null character in path #

@self.app.get("/output/aorta/mask")
async def aorta_mask():
    return FileResponse(
        self.nifti_zip(self.get_mask()).getvalue(),
        media_type="application/octet-stream",
        filename=self.filename,
    )

Test 4 → 无错误，但未触发下载 #

@self.app.get("/output/aorta/mask")
async def aorta_mask():
    response = StreamingResponse(
        self.nifti_zip(self.get_mask()), media_type="application/gzip"
    )
    response.headers["Content-Disposition"] = "attachment; filename={}".format(
        self.filename
    )
    return response

解决方案 #

1. 修复压缩函数的指针位置 #

nifti_zip返回的BytesIO对象在写入数据后，指针停在流的末尾，导致StreamingResponse无法读取内容，需要手动将指针移到开头：

def nifti_zip(self, nifti_image):
    """ Function to compress an nifti image"""
    nifti_bytes = nifti_image.to_bytes()
    compressed_nifti = io.BytesIO()

    with gzip.GzipFile(fileobj=compressed_nifti, mode="wb") as compressed_stream:
        compressed_stream.write(nifti_bytes)
    
    # 将文件指针移到流的起始位置
    compressed_nifti.seek(0)
    return compressed_nifti

2. 修正返回接口的文件名与响应头 #

确保文件名以.nii.gz结尾，同时规范Content-Disposition的格式（用引号包裹文件名，避免特殊字符问题）：

@self.app.get("/output/aorta/mask")
async def aorta_mask():
    mask_image = self.get_mask()
    compressed_stream = self.nifti_zip(mask_image)
    
    # 确保下载文件名是.nii.gz格式
    if not self.filename.endswith(".nii.gz"):
        base_name = self.filename.rsplit('.', 1)[0] if '.' in self.filename else self.filename
        download_filename = f"{base_name}.nii.gz"
    else:
        download_filename = self.filename
    
    response = StreamingResponse(compressed_stream, media_type="application/gzip")
    response.headers["Content-Disposition"] = f'attachment; filename="{download_filename}"'
    return response

错误原因说明 #

• Test1/Test3：FileResponse要求第一个参数是文件路径字符串，传入字节数据会被当作路径解析，导致非法路径错误。
• Test2：io.BytesIO只能接收字节数据，直接传入Nifti1Image对象会触发类型不匹配错误。
• Test4：BytesIO流指针在末尾，StreamingResponse读取不到内容；同时文件名格式不规范可能导致浏览器无法识别文件类型，无法触发下载。

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