update app.py

app.py

@@ -59,7 +59,7 @@ def get_transform():
 
 # Define function to load model
 @st.cache_resource
-def load_model(model_path="models/model.pth"):
+def load_model(model_path="model_comparison/best_model/model.pth"):
     # Try to find model_info.json to determine which model architecture to use
     model_dir = os.path.dirname(model_path)
     model_info_path = os.path.join(model_dir, "model_info.json")
@@ -290,41 +290,36 @@ def main():
         with upload_container:
             uploaded_file = st.file_uploader("Choose a leaf image to analyze", type=["jpg", "jpeg", "png"])
             
-            # Add example images
-            st.markdown("<h3>Or try an example:</h3>", unsafe_allow_html=True)
-            example_cols = st.columns(4)
+            # Add test images from test_images folder
+            st.markdown("<h3>Or select a test image:</h3>", unsafe_allow_html=True)
+            test_cols = st.columns(4)
             
-            # Example images from URLs
-            example_images = [
-                {"name": "Healthy Leaf", "path": "https://www.gardendesign.com/pictures/images/900x705Max/site_3/helianthus-yellow-flower-pixabay_12708.jpg"},
-                {"name": "Mild Infestation", "path": "https://extension.umn.edu/sites/extension.umn.edu/files/styles/caption_medium/public/two-spotted-spider-mite-feeding-damage-JWeiland-iStock.jpg?itok=VLzrPzLc"},
-                {"name": "Severe Infestation", "path": "https://www.planetnatural.com/wp-content/uploads/2012/12/spider-mite-damage-1.jpg"},
-                {"name": "Spider Mite Closeup", "path": "https://bugguide.net/images/raw/ZR0/ZXR/ZR0ZXRCZMLJLUL1LJLQZ5RRR9RHH5RQZIRFZIRCZKRQZ5RNLKRFZ5RULARCZIRGZIZ9RTLSZPR.jpg"},
-            ]
+            # Get test images from test_images folder
+            test_images = []
+            for root, dirs, files in os.walk("test_images"):
+                for file in files:
+                    if file.lower().endswith((".jpg", ".jpeg", ".png")):
+                        test_images.append({"name": file, "path": os.path.join(root, file)})
             
-            use_example = None
-            for i, example in enumerate(example_images):
-                with example_cols[i]:
-                    st.image(example["path"], caption=example["name"], width=150)
-                    if st.button(f"Use Example {i+1}", key=f"example_{i}"):
-                        use_example = example["path"]
+            # Display test images
+            use_test = None
+            for i, test_img in enumerate(test_images):
+                with test_cols[i % 4]:
+                    folder_type = "Healthy" if "healthy" in test_img["path"].lower() else "Spider Mite Infested"
+                    st.image(test_img["path"], caption=folder_type, width=150)
+                    if st.button(f"Use {folder_type}", key=f"test_button_{i}"):
+                        use_test = test_img["path"]
+                        uploaded_file = use_test
         
-        # Process image
-        if uploaded_file is not None or use_example is not None:
-            # Handle example image
-            if use_example is not None:
+        # Process image if uploaded
+        if uploaded_file is not None:
+            # Handle test image path
+            if isinstance(uploaded_file, str):
                 try:
-                    import requests
-                    from io import BytesIO
-                    response = requests.get(use_example)
-                    image = Image.open(BytesIO(response.content)).convert('RGB')
-                except Exception as e:
-                    st.error(f"Error loading example image: {e}")
+                    image = Image.open(uploaded_file).convert('RGB')
+                except:
+                    st.error(f"Could not open test image: {uploaded_file}")
                     return
-            else:
-                # User uploaded image
-                image = Image.open(uploaded_file).convert('RGB')
-            
             st.markdown("---")
             
             analysis_container = st.container()
@@ -333,6 +328,7 @@ def main():
                 
                 with col1:
                     st.markdown('<h3 class="sub-header">Uploaded Image</h3>', unsafe_allow_html=True)
+                    image = Image.open(uploaded_file).convert('RGB')
                     st.image(image, caption="", use_container_width=True)
                 
                 with col2: