Downloading 3D numpy arrays

[MLbyML]

Hello everyone!

I am trying to build a streamlit app that produces a 3D numpy array (stored in a variable called predicted_map) having size 10 x 100 x 100 as a result of a calculation. I wanted to add a st.download_button that provides the user of the app, access to this calculated array.

Initially, I tried

st.download_button(label='Download model prediction', data=predicted_map, file_name='prediction')

This generates an error because a bytes object (and not a numpy array) is expected in data

I think I don’t want to do something like:

st.download_button(label='Download model prediction', data=predicted_map.tobytes(), file_name='prediction')

because then the user of the app would have to do the decoding process back to numpy array on their side. The ideal scenario would be that they are directly able to receive a numpy array.

Wrapping the 3D numpy array as an Image object using Image.fromarray also seems to not work for such 3D arrays.

How could I download such an array directly? Any suggestions would be very appreciated. Thank you!

[snehankekre]

Hi @MLbyML, welcome to the Streamlit community!!

Thank you for walking us through what you’ve tried and what you don’t want to do.

Here’s how you can save predicted_map, a 3D numpy arrray, as a .npy file with st.download_button. The trick is to create an in-memory buffer with io.BytesIO, write to it using np.save(), and use st.download_button() to download this in-memory buffer as as an .npy file:

import streamlit as st
import numpy as np
import io

@st.experimental_memo
def load_data():
    return np.random.rand(10, 100, 100)

# Create a 3D numpy array
predicted_map = load_data()

# Create an in-memory buffer
with io.BytesIO() as buffer:
    # Write array to buffer
    np.save(buffer, predicted_map)
    btn = st.download_button(
        label="Download numpy array (.npy)",
        data = buffer, # Download buffer
        file_name = 'predicted_map.npy'
    ) 

Your users will be “directly able to receive a numpy array”, which they can load using np.load(predicted_map.npy).

Hope this helps!

Happy Streamlit-ing,
Snehan

[MLbyML]

Hello Snehan,

This is such an elegant solution. Works like a charm. Thanks very much!

[n3urovirtual]

I am trying to do the same with an image which is stored as an array. I use the code you provided but when I download the image, it says that I “file format is not supported” and I cannot see the image. You can see my code below:

 edges = cv2.Canny(array, l_threshold, h_threshold)
 st.image(edges,caption='Edges', use_column_width=True)
                  
# Create an in-memory buffer
with io.BytesIO() as buffer:
    # Write array to buffer
    np.save(buffer, edges)
    btn = st.download_button(
        label="Download Image",
        data = buffer, # Download buffer
        file_name = 'Image.npy'
    ) 

Any ideas?

[MLbyML]

Converting edges back to a numpy array, by enclosing it with np.asarray() will do the trick:

img = cv2.imread('logo.png', 0)
edges = cv2.Canny(img, 100, 200)
st.image(edges, caption='Edges', use_column_width=True)
with io.BytesIO() as buffer:
    np.save(buffer, np.asarray(edges))
    btn = st.download_button(
        label="Download Image",
        data=buffer,  # Download buffer
        file_name='Image.npy'
    )

[n3urovirtual]

Nope. It doesn’t. I am still getting the same error.

[snehankekre]

Perhaps if you share a minimal reproducible code example, we can be of more help.

Best,
Snehan

[n3urovirtual]

I thought what I posted was an MRC. There you go:

array = cv2.imread('original_image.png',0)
edges = cv2.Canny(array, l_threshold, h_threshold)
st.image(edges,caption='Edges', use_column_width=True)
                  
# Create an in-memory buffer
with io.BytesIO() as buffer:
    # Write array to buffer
    np.save(buffer, edges)
    btn = st.download_button(
        label="Download Image",
        data = buffer, # Download buffer
        file_name = 'Image.npy'
    ) 

It downloads the image properly but you can’t view it because it says format not supported. Ideally, would like to get a .png or sth similar.

[n3urovirtual]

I found the solution. Posting it here for others:

array = cv2.imread('original_image.png',0)
edges = cv2.Canny(array, l_threshold, h_threshold)
st.image(edges,caption='Edged', use_column_width=True)

# Download button
im = Image.fromarray(edge)
im.save("Edge.png")
with open("Edge.png", "rb") as file:
        btn = st.download_button(
                 label="Download Image",
                 data=file,
                 file_name="Edge_image.png",
                 mime="image/png"
               )