Update upload_view.py

upload_view.py

@@ -58,3 +58,19 @@ class Upload_View:
 
                 # Display tumor count and confidence information using Streamlit text element
                 st.subheader(f"Number of Tumors Detected: {num_tumors}")
+
+                # Convert annotated image to bytes for download
+                annotated_pil = Image.fromarray(annotated_image)
+                img_bytes = annotated_pil.convert("RGB")
+                from io import BytesIO
+                buf = BytesIO()
+                img_bytes.save(buf, format="PNG")
+                byte_im = buf.getvalue()
+
+                # Add download button
+                st.download_button(
+                    label="📥 Download Annotated Image",
+                    data=byte_im,
+                    file_name="tumor_prediction.png",
+                    mime="image/png"
+                )