How to output an .nii.gz nifiti image with FastAPI whilst keeping everything on memory?

Question

I would like to find a way to send a .nii.gz file from my FastAPI 'server' without saving the image, but instead keeping it in the memory. I managed to upload the image like that, but I struggle to do the inverse. For the context, the idea is to input an 3D CT scan, perform an organ segmentation and retrieve the segmented mask.

Here is the input function (simplified):

import io
import gzip
import nibabel as nib
from fastapi import FastAPI,HTTPException

@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}

Here's what I tried in order to output an .nii.gz image for which I can have download link. I defined a function to compress my image first then the one where I tried to send it.

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

The get_mask() function here return a nifti image of the target organ.

Test 1 -> error: "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 -> error: "File "/app/server/server.py", line 148, in aorta_mask io.BytesIO(self.get_mask()), 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 -> error: "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 -> I don't get an error, but I don't get a download link

@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

Thanks in advance for your help !

Answer 1

After couple of days looking for a way to do this, I managed to find the answer. I post it here if it can help someone one day:

@app.get("/image_4")
async def upload_image_4():
    """Function used to download a 3D CT scan"""
    # Loading the image just for testing purpose
    img_in = nib.load("test.nii.gz")

    # Compress data and write them in a file
    nifti_data = nib.Nifti1Image.to_bytes(img_in)
    compressed_data = gzip.compress(nifti_data)
    output_bytesio = io.BytesIO()
    output_bytesio.write(compressed_data)
    output_bytesio.seek(0)

    return StreamingResponse(output_bytesio, media_type="application/octet-stream")

To download the image, either use the docs of the local sever and save the image as an .nii.gz file, or use this Curl cmd:

curl -X 'GET'   'http://0.0.0.0:8000/image_4'   -H 'accept: application/json' --output myfile.nii.gz

The last option is to write a python script like this one:

import requests
import shutil
import uuid


def download_file(url):
    local_filename = "download.nii.gz"
    with requests.get(url, stream=True) as r:
        with open(local_filename, "wb") as f:
            shutil.copyfileobj(r.raw, f)
    return local_filename


url = "http://localhost:8000/image_4"
download_file(url)