Download Button Issue: Downloads the file and not the Image

[Shenin_Francies]

So I am trying to use a download button to download an image after a filter has been applied to it.

The idea is to take an uploaded image, make a Pencil Sketch of the uploaded image and then download the sketched image. I am able to display the pencil sketched image well enough as you can see below:

I cannot however download the pencil sketched image. The output is a file and not an image.
Please ignore the label on the button, in case it is misleading.

Here is how I have tried so far:

# Convert Image
def save_image(byteImage):
    bytesImg = io.BytesIO(byteImage)
    imgFile = Image.open(bytesImg)
    return imgFile


# Check if the file has been uploaded
if uploaded_file is not None:
    file = uploaded_file.read()
    filename = uploaded_file.name.split('.')[0]
    path = save_image(file)
    st.title(option)
    # st.markdown('## Original Image')
    col1, col2, col3 = st.columns([1, 1, 1])
    col2.image(path, width=500, caption='Uploaded Image')
    file_bytes = np.asarray(bytearray(file), dtype=np.uint8)
    opencv_image = cv2.imdecode(file_bytes, 1)
    if option == 'Pencil Sketch':
        sketch = PencilSketch(opencv_image)
        col1, col2, col3 = st.columns([1, 1, 1])
        col2.image(sketch, width=500, caption='Pencil Sketch')
        col1, col2, col3 = st.columns([1, 1, 1])
        btn = col2.download_button(
            label="Download {}".format(filename),
            data=file,
            file_name=filename,
            mime="image/png",
        )

Also, is there a way that I can center the download button further, so that it is straight at the center below the image? I am aware that Streamlit does not support nested columns, so I am not sure about the other ways to do this.

Thank you!

[snehankekre]

Hi @Shenin_Francies

It looks like you were close to getting the download button working correctly. There were a few issues in your original code that were causing the download to not work as expected.

In the save_image function, you were converting the image to a PIL Image object, but then not doing anything with it. Instead, you should be passing the image data directly to the download_button function.

Another issue I see in your code is that you’re trying to download the original file instead of the sketch image.

In the fixed code below:

• I’ve added a new buffer variable that is used to hold the image data in memory with io.BytesIO
• Since PencilSketch was not defined in your snippet, I’ve used cv2.pencilSketch instead to generate the sketched image sketch.
• cv2.pencilSketch returns a tuple of two image arrays sketch[0], and sketch[1].
• I convert sketch[0] to a PIL image and save it to the in-memory buffer
• I create two columns col1: to display original image; col2: to display the sketch
• Lastly, I create 3 columns and display the download button widget in the middle column, and pass the in-memory buffer to st.download_button’s data parameter

import io

import cv2
import numpy as np
from PIL import Image

import streamlit as st

uploaded_file = st.file_uploader("Choose an image...")
# Create an in-memory buffer to hold the image
buffer = io.BytesIO()

option = st.selectbox(
    "Select the type of sketch you want to generate",
    ("Pencil Sketch", "Color Sketch"),
)

# Check if the file has been uploaded
if uploaded_file is not None:
    file = uploaded_file.read()
    filename = uploaded_file.name.split(".")[0]
    st.title(option)
    # Create two columns
    # col1: display original image; col2: display sketch
    col1, col2 = st.columns(2)
    # Display the original image
    col1.image(uploaded_file, caption="Uploaded Image")
    file_bytes = np.asarray(bytearray(file), dtype=np.uint8)
    opencv_image = cv2.imdecode(file_bytes, 1)
    if option == "Pencil Sketch":
        # PencilSketch in your code was not defined
        # Using cv2.pencilSketch instead
        sketch = cv2.pencilSketch(
            opencv_image, sigma_s=60, sigma_r=0.07, shade_factor=0.05
        )  # sketch is a tuple of two image arrays
        col2.image(sketch[0], caption="Pencil Sketch")
        # Convert the image to a PIL image
        # and save it to the in-memory buffer
        im = Image.fromarray(sketch[0])
        im.save(buffer, format="PNG")
        # Create 3 columns and display the download button
        # in the middle column
        col1, col2, col3 = st.columns(3)
        btn = col2.download_button(
            label=f"Download **{filename}.png**",
            data=buffer,  # download image from the in-memory buffer
            file_name=f"{filename}.png",
            mime="image/png",
        )

Hope this helps!