Update app.py

app.py

@@ -188,7 +188,7 @@ if uploaded_file is not None:
             ccd.write(f"predicted_{i+1}.fits", overwrite=True)
             image_decomposed += (i+1) * np.where(cav > 0, 1, 0)
 
-        shutil.make_archive("predictions", 'zip', "predictions")
+        shutil.make_archive("predictions.zip", 'zip', "predictions")
         np.save("decomposed.npy", image_decomposed)
 
     try: image_decomposed = np.load("decomposed.npy")
@@ -196,8 +196,14 @@ if uploaded_file is not None:
     plot_decomposed(image_decomposed)
 
     with col6:
-        with open('predicted.zip', 'rb') as f:
-            res = f.read()
+        try:
+            name = 'predictions.zip'
+            with open(name, 'rb') as f:
+                res = f.read()
+        except:
+            name = uploaded_file
+            with open(name, 'rb') as f:
+                res = f.read()
         # st.markdown("""<style>[data-baseweb="select"] {margin-top: 16px;}</style>""", unsafe_allow_html=True)
     #     # download = st.button('Download')
-        download = st.download_button(label="Download", data=res, file_name="predicted.zip", mime="application/octet-stream")
+        download = st.download_button(label="Download", data=res, file_name=name, mime="application/octet-stream")