Update src/streamlit_app.py

src/streamlit_app.py

@@ -89,19 +89,22 @@ def fetch_random_image():
 def preprocess_image(image, target_size=(64, 64)):
 
     """Resizes, normalizes, and ensures 3 channels for the image."""
-
-    image = image.resize(target_size)
-
-    if image.mode == 'RGBA':
-        image = image.convert('RGB')
-    elif image.mode == 'P':
-        image = image.convert('RGB')
-
-    image = np.array(image) / 255.0  # Normalize to [0, 1]
-
-    if len(image.shape) == 2:  # Convert grayscale to RGB
-
-        image = np.stack([image] * 3, axis=-1)
+    try:
+        image = image.resize(target_size)
+    
+        if image.mode == 'RGBA':
+            image = image.convert('RGB')
+        elif image.mode == 'P':
+            image = image.convert('RGB')
+    
+        image = np.array(image) / 255.0  # Normalize to [0, 1]
+    
+        if len(image.shape) == 2:  # Convert grayscale to RGB
+    
+            image = np.stack([image] * 3, axis=-1)
+    except Exception as e:
+        print(f"Preprocessing error: {e}")
+        return None
 
     return image
 
@@ -146,21 +149,24 @@ def create_model(input_shape=(64, 64, 3)):
 def fetch_image():
     print("entering fetch_image")
 
-    image = fetch_random_image()
-
-    if image is not None:
-
-        st.session_state.unprocessed_image = image
-        st.session_state.current_image = preprocess_image(image)
-
-        st.session_state.current_label = None
-
-        st.session_state.current_prediction = None
-        st.info("image ready")
-
-    else:
-
-        st.error("Failed to fetch an image. Please try again.")
+    try:
+        image = fetch_random_image()
+    
+        if image is not None:
+    
+            st.session_state.unprocessed_image = image
+            st.session_state.current_image = preprocess_image(image)
+    
+            st.session_state.current_label = None
+    
+            st.session_state.current_prediction = None
+            st.info("image ready")
+    
+        else:
+    
+            st.error("Failed to fetch an image. Please try again.")
+    except Exception as e:
+        st.error(f"An error occurred while fetching the image: {e}")
 
 
 # ---------------------------