Update Image Augmentation.py

Image Augmentation.py

@@ -5,8 +5,9 @@ from io import BytesIO
 from PIL import Image
 import base64
 from fpdf import FPDF
+import os
 
-
+# Define augmentation functions
 def rotate_image(image, angle):
     (h, w) = image.shape[:2]
     center = (w // 2, h // 2)
@@ -14,7 +15,6 @@ def rotate_image(image, angle):
     rotated_image = cv2.warpAffine(image, rotation_matrix, (w, h))
     return rotated_image
 
-
 def add_noise(image):
     s_vs_p = 0.5
     amount = 0.02
@@ -27,45 +27,48 @@ def add_noise(image):
     noisy_image[coords[0], coords[1], :] = 0
     return noisy_image
 
-
 def blur_image(image, kernel_size):
     blurred_image = cv2.GaussianBlur(image, (kernel_size, kernel_size), 0)
     return blurred_image
 
-
 def translate_image(image, tx, ty):
     translation_matrix = np.float32([[1, 0, tx], [0, 1, ty]])
     height, width = image.shape[:2]
     translated_image = cv2.warpAffine(image, translation_matrix, (width, height))
     return translated_image
 
-
+# Download functions
 def download_image(image, file_format):
+    if file_format.lower() == "jpg":
+        file_format = "jpeg"
     pil_image = Image.fromarray(cv2.cvtColor(image, cv2.COLOR_BGR2RGB))
     buffer = BytesIO()
     pil_image.save(buffer, format=file_format.upper())
     buffer.seek(0)
     b64 = base64.b64encode(buffer.read()).decode()
-    href = f'<a href="data:file/{file_format};base64,{b64}" download="augmented_image.{file_format}">Download {file_format.upper()} File</a>'
+    href = f'<a href="data:file/{file_format};base64,{b64}" download="augmented_image.{file_format.lower()}">Download {file_format.upper()} File</a>'
     return href
 
-
 def download_pdf(image):
     pdf = FPDF()
     pdf.add_page()
+
+    temp_image_path = "temp_image.jpg"
     pil_image = Image.fromarray(cv2.cvtColor(image, cv2.COLOR_BGR2RGB))
-    buffer = BytesIO()
-    pil_image.save(buffer, format="JPEG")
-    buffer.seek(0)
-    pdf.image(buffer, x=10, y=10, w=190)
+    pil_image.save(temp_image_path, format="JPEG")
+
+    pdf.image(temp_image_path, x=10, y=10, w=190)
+    os.remove(temp_image_path)
+
     pdf_output = BytesIO()
     pdf.output(pdf_output, "F")
     pdf_output.seek(0)
+
     b64 = base64.b64encode(pdf_output.read()).decode()
     href = f'<a href="data:application/octet-stream;base64,{b64}" download="augmented_image.pdf">Download PDF</a>'
     return href
 
-
+# Streamlit app
 st.title("Image Augmentation Tool")
 
 # File upload
@@ -96,7 +99,9 @@ if uploaded_file:
         st.image(cv2.cvtColor(transformed_image, cv2.COLOR_BGR2RGB), caption="Transformed Image", use_column_width=True)
 
         st.markdown("### Download Transformed Image:")
-        st.markdown(download_image(transformed_image, "jpeg"), unsafe_allow_html=True)
-        st.markdown(download_image(transformed_image, "png"), unsafe_allow_html=True)
-        st.markdown(download_pdf(transformed_image), unsafe_allow_html=True)
+        download_format = st.radio("Select Download Format:", ["JPG", "PNG", "PDF"])
 
+        if download_format == "JPG" or download_format == "PNG":
+            st.markdown(download_image(transformed_image, download_format), unsafe_allow_html=True)
+        elif download_format == "PDF":
+            st.markdown(download_pdf(transformed_image), unsafe_allow_html=True)