Update app.py

app.py

@@ -180,8 +180,7 @@ if uploaded_file is not None:
     with col5: decompose = st.button('Decompose')
 
     with col6:
-        shutil.make_archive("predictions", 'zip', "predictions")
-        with open('predictions.zip', 'rb') as f:
+        with open(uploaded_file.name, 'rb') as f:
             res = f.read()
         st.markdown("")
         st.markdown("")
@@ -207,4 +206,9 @@ if uploaded_file is not None:
             
             image_decomposed = decompose_cavity(y_pred_th)
     
-            plot_decomposed(image_decomposed)
+            plot_decomposed(image_decomposed)
+
+            shutil.make_archive("predictions", 'zip', "predictions")
+            with open('predictions.zip', 'rb') as f:
+                res = f.read()
+            download.data=res