Update app.py

app.py

@@ -42,6 +42,12 @@ def preprocess_image(image_path, target_size=(224, 224)):
 
 def load_image_dataset(zip_path, target_size=(224, 224), problem_type="Classification"):
     """Load and preprocess an image dataset from a zip file."""
+    # Check file size (5GB = 5 * 1024 * 1024 * 1024 bytes)
+    file_size = os.path.getsize(zip_path) if isinstance(zip_path, str) else zip_path.size
+    max_size = 5 * 1024 * 1024 * 1024  # 5GB in bytes
+    if file_size > max_size:
+        raise ValueError(f"Uploaded file size ({file_size / (1024 * 1024):.2f} MB) exceeds the 5GB limit.")
+    
     # Extract zip file to a temporary directory
     with zipfile.ZipFile(zip_path, 'r') as zip_ref:
         zip_ref.extractall('temp_images')
@@ -365,17 +371,25 @@ if app_mode == "Data Upload":
                 with col2: st.metric("Columns", df.shape[1])
                 with col3: st.metric("Missing Values", df.isna().sum().sum())
         else:  # Image
-            uploaded_file = st.file_uploader("Upload Zip File with Images", type=["zip"])
+            uploaded_file = st.file_uploader("Upload Zip File with Images (Max 5GB)", type=["zip"])
             if uploaded_file:
-                problem_type = st.selectbox("Problem Type for Image Data", ["Image Classification", "Compression", "Clustering"])
-                images, labels, class_names = load_image_dataset(uploaded_file, problem_type=problem_type)
-                st.session_state.images = images
-                st.session_state.labels = labels
-                st.session_state.class_names = class_names if problem_type == "Image Classification" else None
-                st.write(f"Loaded {len(images)} images.")
-                if problem_type == "Image Classification":
-                    st.write(f"Classes: {class_names}")
-                st.image(images[:5], caption=["Sample " + str(i+1) for i in range(min(5, len(images)))], width=100)
+                # Save uploaded file temporarily to check size
+                with open("temp_upload.zip", "wb") as f:
+                    f.write(uploaded_file.getbuffer())
+                try:
+                    problem_type = st.selectbox("Problem Type for Image Data", ["Image Classification", "Compression", "Clustering"])
+                    images, labels, class_names = load_image_dataset("temp_upload.zip", problem_type=problem_type)
+                    st.session_state.images = images
+                    st.session_state.labels = labels
+                    st.session_state.class_names = class_names if problem_type == "Image Classification" else None
+                    st.write(f"Loaded {len(images)} images.")
+                    if problem_type == "Image Classification":
+                        st.write(f"Classes: {class_names}")
+                    st.image(images[:5], caption=["Sample " + str(i+1) for i in range(min(5, len(images)))], width=100)
+                except ValueError as e:
+                    st.error(str(e))
+                finally:
+                    os.remove("temp_upload.zip")
 
 elif app_mode == "Model Training":
     st.title("🧠 Model Training")