Update app.py

app.py

@@ -71,7 +71,7 @@ if uploaded_files:
 
         col1, col2 = st.columns([1, 1.5])
         with col1:
-            # Controls
+            # Reset button
             if st.button("🔁 Reset", key=f"reset_{key}"):
                 st.session_state[f"angle_{key}"] = 0
                 st.session_state[f"enhance_{key}"] = False
@@ -81,29 +81,30 @@ if uploaded_files:
                 st.session_state[f"{key}_right"] = w
                 st.session_state[f"{key}_bottom"] = h
 
+            # Controls
             st.slider("🌀 Rotate", -180, 180, st.session_state[f"angle_{key}"], key=f"angle_{key}")
             st.checkbox("✨ Enhance", value=st.session_state[f"enhance_{key}"], key=f"enhance_{key}")
             left, top, right, bottom = crop_sliders(original, key)
             st.text_input("📁 Rename this file", value=st.session_state[f"filename_{key}"], key=f"filename_{key}_input")
 
         with col2:
-            # Real-time preview
-            img = st.session_state[f"original_{key}"].crop((left, top, right, bottom))
-            img = resize_image(img)
+            # Live preview logic here (always runs)
+            preview = st.session_state[f"original_{key}"].crop((left, top, right, bottom))
+            preview = resize_image(preview)
 
             if st.session_state[f"enhance_{key}"]:
-                img = ImageEnhance.Contrast(img).enhance(1.5)
-                img = ImageEnhance.Sharpness(img).enhance(2.0)
+                preview = ImageEnhance.Contrast(preview).enhance(1.5)
+                preview = ImageEnhance.Sharpness(preview).enhance(2.0)
 
             angle = st.session_state[f"angle_{key}"]
             if angle != 0:
-                img = img.rotate(angle, expand=True)
+                preview = preview.rotate(angle, expand=True)
 
-            st.session_state[f"preview_{key}"] = img
-            st.image(img, caption="🎯 Live Preview", use_container_width=True)
+            st.session_state[f"preview_{key}"] = preview
+            st.image(preview, caption="🎯 Live Preview", use_container_width=True)
 
             if st.button("✅ Save This Image", key=f"save_{key}"):
-                st.session_state[f"final_{key}"] = img.copy()
+                st.session_state[f"final_{key}"] = preview.copy()
                 st.success("Saved for export!")
 
         final_img = st.session_state.get(f"final_{key}", original)
@@ -111,34 +112,45 @@ if uploaded_files:
         processed_images.append(final_img)
         file_names.append(filename)
 
-    # --- Export Section ---
+    # --- EXPORT SECTION ---
     st.markdown("### 🧾 Export Options")
-    with st.container():
-        col1, col2 = st.columns([1, 1])
-        with col1:
-            if st.button("📦 Generate ZIP", key="zip_btn"):
-                zip_buffer = BytesIO()
-                with zipfile.ZipFile(zip_buffer, "w") as zipf:
-                    for name, img in zip(file_names, processed_images):
-                        buffer = BytesIO()
-                        img.save(buffer, format="PNG")
-                        zipf.writestr(f"{name}.png", buffer.getvalue())
-                zip_buffer.seek(0)
-                st.download_button("⬇️ Download ZIP", zip_buffer, "images.zip", mime="application/zip", key="zip_dl")
+    col1, col2 = st.columns([1, 1])
 
-        with col2:
-            if st.button("📄 Generate PDF", key="pdf_btn"):
-                pdf = FPDF()
-                tmp_paths = []
+    show_zip = st.session_state.get("show_zip", False)
+    show_pdf = st.session_state.get("show_pdf", False)
+
+    with col1:
+        if st.button("📦 Generate ZIP"):
+            zip_buffer = BytesIO()
+            with zipfile.ZipFile(zip_buffer, "w") as zipf:
                 for name, img in zip(file_names, processed_images):
-                    path = f"/tmp/{name}.jpg"
-                    img.save(path, "JPEG")
-                    pdf.add_page()
-                    pdf.image(path, x=10, y=10, w=190)
-                    tmp_paths.append(path)
-                out_path = "/tmp/final.pdf"
-                pdf.output(out_path)
-                with open(out_path, "rb") as f:
-                    st.download_button("⬇️ Download PDF", f, "scanned.pdf", mime="application/pdf", key="pdf_dl")
-                for p in tmp_paths + [out_path]:
-                    os.remove(p)
+                    buffer = BytesIO()
+                    img.save(buffer, format="PNG")
+                    zipf.writestr(f"{name}.png", buffer.getvalue())
+            zip_buffer.seek(0)
+            st.session_state["zip_file"] = zip_buffer
+            st.session_state["show_zip"] = True
+
+        if st.session_state.get("show_zip", False):
+            st.download_button("⬇️ Download ZIP", st.session_state["zip_file"], "images.zip", mime="application/zip")
+
+    with col2:
+        if st.button("📄 Generate PDF"):
+            pdf = FPDF()
+            tmp_paths = []
+            for name, img in zip(file_names, processed_images):
+                path = f"/tmp/{name}.jpg"
+                img.save(path, "JPEG")
+                pdf.add_page()
+                pdf.image(path, x=10, y=10, w=190)
+                tmp_paths.append(path)
+            out_path = "/tmp/final.pdf"
+            pdf.output(out_path)
+            with open(out_path, "rb") as f:
+                st.session_state["pdf_file"] = f.read()
+            for p in tmp_paths + [out_path]:
+                os.remove(p)
+            st.session_state["show_pdf"] = True
+
+        if st.session_state.get("show_pdf", False):
+            st.download_button("⬇️ Download PDF", st.session_state["pdf_file"], "scanned.pdf", mime="application/pdf")