Update pages/Data Collection.py

pages/Data Collection.py

@@ -1,6 +1,9 @@
 import streamlit as st
 import pandas as pd
 import matplotlib.pyplot as plt
+import io
+import zipfile
+import numpy as np
 
 
 # Define functions for individual pages
@@ -1656,6 +1659,8 @@ def Image_Augmentation_page():
         cv2.destroyAllWindows()
         """, language="python")
 
+
+
         # Function to apply affine transformations
         def apply_affine_transformation(image, transformation_type):
             transformed_images = []
@@ -1678,8 +1683,9 @@ def Image_Augmentation_page():
                     # Simple cropping: reduce the size incrementally
                     x1, y1 = i * 5, i * 5
                     x2, y2 = cols - i * 5, rows - i * 5
-                    transformed_image = image[y1:y2, x1:x2]
-                    transformed_images.append(transformed_image)
+                    if x1 < x2 and y1 < y2:  # Ensure cropping dimensions are valid
+                        transformed_image = image[y1:y2, x1:x2]
+                        transformed_images.append(transformed_image)
                     continue  # Skip warpAffine for cropping
                 else:
                     st.error("Invalid transformation type!")
@@ -1688,8 +1694,8 @@ def Image_Augmentation_page():
                 transformed_image = cv2.warpAffine(image, M, (cols, rows))
                 transformed_images.append(transformed_image)
         
-        return transformed_images
-            
+            return transformed_images
+        
         # Streamlit App
         st.title("Dynamic Affine Transformation Tool")
         st.write("Select a transformation type to proceed and learn how it works before uploading an image.")
@@ -1729,7 +1735,7 @@ def Image_Augmentation_page():
                 transformed_images = apply_affine_transformation(image, transformation)
         
                 if transformed_images:
-                    st.write(f"Generated 10 images using {transformation}:")
+                    st.write(f"Generated {len(transformed_images)} images using {transformation}:")
         
                     # Display all transformed images
                     for i, img in enumerate(transformed_images):
@@ -1750,6 +1756,10 @@ def Image_Augmentation_page():
                         file_name=f"{transformation}_transformed_images.zip",
                         mime="application/zip"
                     )
+                else:
+                    st.warning("No transformed images generated. Please check your transformation type.")
+            else:
+                st.warning("Please upload an image to proceed.")
         else:
             st.warning("Please select a valid transformation type to proceed.")