Download Button Issue: Downloads the file and not the Image

Using Streamlit
2

Hi @Shenin_Francies :wave:

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",
        )

download-pencil-sketch

Hope this helps! :balloon:

2 Likes