Add download button

app/app.py

@@ -17,6 +17,11 @@ def describe_service():
     st.subheader("Just upload your low-light image and get the processed one!")
 
 
+@st.experimental_memo
+def call_model(uploaded_file):
+    return model.run_model(uploaded_file, st.session_state["model"])
+
+
 def process_image():
     uploaded_file: io.BytesIO = st.file_uploader(
         label="Choose a file (you can upload new files without refreshing the page)",
@@ -25,9 +30,11 @@ def process_image():
     if uploaded_file:
         placeholder: st.delta_generator.DeltaGenerator = st.empty()
         placeholder.info("The image is being processed. It may take some time. Wait, please...")
-        image: Image.Image = model.run_model(uploaded_file, st.session_state["model"])
+        image: Image.Image = call_model(uploaded_file)
         placeholder.empty()
         placeholder.image(image)
+        with open("user_data/output.jpg", "rb") as file:
+            st.download_button(label="Download lightened image", data=file, file_name="lightened.jpg", mime="image/jpg")
 
 
 def main():

app/model.py

@@ -3,6 +3,7 @@ from tensorflow import keras
 from PIL import Image
 import numpy as np
 import io
+import os
 
 
 def create_model() -> keras.Model:
@@ -27,6 +28,7 @@ def run_model(image_bytes: io.BytesIO, model: keras.Model) -> Image.Image:
     output_image_array5: np.ndarray = output_image_array3.astype(np.uint8)
     output_image1: Image.Image = Image.fromarray(output_image_array5)
     output_image2: Image.Image = output_image1.resize((width, height))
-    # Uncomment if necessary:
-    # output_image.save("user_data/output.jpg")
+    if not os.path.exists("user_data"):
+        os.makedirs("user_data")
+    output_image2.save("user_data/output.jpg")
     return output_image2