Update app.py

app.py

@@ -5,7 +5,11 @@ import streamlit as st
 import PIL
 from ultralytics import YOLO
 
-# Required libraries: streamlit, opencv-python-headless, ultralytics, Pillow
+# Required libraries:
+#   streamlit
+#   opencv-python-headless
+#   ultralytics
+#   Pillow
 
 model_path = 'https://huggingface.co/spaces/tstone87/ccr-colorado/resolve/main/best.pt'
 
@@ -18,14 +22,14 @@ st.set_page_config(
 
 with st.sidebar:
     st.header("IMAGE/VIDEO UPLOAD")
-    source_file = st.file_uploader("Choose an image or video...", 
-                                   type=("jpg", "jpeg", "png", "bmp", "webp", "mp4"))
+    source_file = st.file_uploader("Choose an image or video...", type=("jpg", "jpeg", "png", "bmp", "webp", "mp4"))
     confidence = float(st.slider("Select Model Confidence", 25, 100, 40)) / 100
     video_option = st.selectbox("Select Video Shortening Option", 
                                 ["Original FPS", "1 fps", "1 frame per 5 seconds", "1 frame per 10 seconds", "1 frame per 15 seconds"])
     progress_text = st.empty()
     progress_bar = st.progress(0)
-    slider_container = st.empty()  # For dynamic slider widget
+    # A container where our dynamic slider (frame viewer) will be placed.
+    slider_container = st.empty()
 
 st.title("WildfireWatch: Detecting Wildfire using AI")
 col1, col2 = st.columns(2)
@@ -35,11 +39,13 @@ with col2:
     st.image("https://huggingface.co/spaces/tstone87/ccr-colorado/resolve/main/Fire_3.png", use_column_width=True)
 
 st.markdown("""
-Fires in Colorado present a serious challenge, threatening urban communities, highways, and even remote areas. Early detection is critical. WildfireWatch leverages YOLOv8 for real‐time fire and smoke detection in images and videos.
+Fires in Colorado present a serious challenge, threatening urban communities, highways, and even remote areas. Early detection is critical. 
+WildfireWatch leverages YOLOv8 for real‐time fire and smoke detection in images and videos.
 """)
 st.markdown("---")
 st.header("Fire Detection:")
 
+# Left column for the uploaded file, right for detection results.
 col1, col2 = st.columns(2)
 if source_file:
     if source_file.type.split('/')[0] == 'image':
@@ -52,22 +58,30 @@ if source_file:
 else:
     st.info("Please upload an image or video file to begin.")
 
+# Attempt to load the model.
 try:
     model = YOLO(model_path)
 except Exception as ex:
     st.error(f"Unable to load model. Check the specified path: {model_path}")
     st.error(ex)
 
-viewer_slot = st.empty()  # Container for currently viewed frame
+# We'll store processed frames persistently in session_state.
+if "processed_frames" not in st.session_state:
+    st.session_state["processed_frames"] = []
 
-# Initialize session state for slider value if not already set.
+# Also store the last slider value (if the user manually changes it).
 if "slider_value" not in st.session_state:
-    st.session_state.slider_value = 0
+    st.session_state["slider_value"] = 0
 
+# Container to display the currently viewed frame.
+viewer_slot = st.empty()
+
+# --- Processing and Viewer Update ---
 if st.sidebar.button("Let's Detect Wildfire"):
     if not source_file:
         st.warning("No file uploaded!")
     elif source_file.type.split('/')[0] == 'image':
+        # Process image input.
         res = model.predict(uploaded_image, conf=confidence)
         boxes = res[0].boxes
         res_plotted = res[0].plot()[:, :, ::-1]
@@ -77,8 +91,11 @@ if st.sidebar.button("Let's Detect Wildfire"):
                 for box in boxes:
                     st.write(box.xywh)
     else:
-        processed_frames = []
+        # For video input, process frames.
+        processed_frames = st.session_state["processed_frames"]
         frame_count = 0
+
+        # Video properties.
         orig_fps = vidcap.get(cv2.CAP_PROP_FPS)
         total_frames = int(vidcap.get(cv2.CAP_PROP_FRAME_COUNT))
         width  = int(vidcap.get(cv2.CAP_PROP_FRAME_WIDTH))
@@ -109,7 +126,9 @@ if st.sidebar.button("Let's Detect Wildfire"):
                 res = model.predict(image, conf=confidence)
                 res_plotted = res[0].plot()[:, :, ::-1]
                 processed_frames.append(res_plotted)
+                st.session_state["processed_frames"] = processed_frames
 
+                # Update progress info.
                 if total_frames > 0:
                     progress_pct = int((frame_count / total_frames) * 100)
                     progress_text.text(f"Processing frame {frame_count} / {total_frames} ({progress_pct}%)")
@@ -117,30 +136,44 @@ if st.sidebar.button("Let's Detect Wildfire"):
                 else:
                     progress_text.text(f"Processing frame {frame_count}")
 
-                # Only update the slider if we have frames.
-                if processed_frames:
-                    slider_container.empty()  # Clear previous slider
-                    # Use stored slider_value if within bounds, otherwise default to last frame.
-                    default_val = st.session_state.slider_value if st.session_state.slider_value < len(processed_frames) else len(processed_frames)-1
-                    # Create slider with a unique key based on number of processed frames.
-                    slider_val = slider_container.slider(
-                        "Frame Viewer", 
-                        min_value=0, 
-                        max_value=len(processed_frames)-1, 
-                        value=default_val, 
-                        step=1, 
-                        key=f"frame_slider_{len(processed_frames)}"
-                    )
-                    st.session_state.slider_value = slider_val
-
-                    if slider_val == len(processed_frames)-1:
-                        viewer_slot.image(processed_frames[-1], caption=f"Frame {len(processed_frames)-1}", use_column_width=True)
+                # --- Update the frame viewer slider dynamically ---
+                # Retrieve user's last slider selection.
+                last_slider = st.session_state.slider_value
+                # If the user was at the end, default to the new end.
+                if last_slider >= len(processed_frames):
+                    default_val = len(processed_frames) - 1
+                else:
+                    default_val = last_slider
+
+                # Clear the slider container and recreate the slider.
+                slider_container.empty()
+                # Use a dynamic key to avoid duplicate key errors.
+                slider_key = f"frame_slider_{len(processed_frames)}"
+                slider_val = slider_container.slider("Frame Viewer", 
+                                                       min_value=0, 
+                                                       max_value=len(processed_frames) - 1, 
+                                                       value=default_val, 
+                                                       step=1, 
+                                                       key=slider_key)
+                st.session_state.slider_value = slider_val
+
+                # If the slider is at the most recent frame, update the viewer.
+                if slider_val == len(processed_frames) - 1:
+                    viewer_slot.image(processed_frames[-1], 
+                                      caption=f"Frame {len(processed_frames) - 1}", 
+                                      use_column_width=True)
+                else:
+                    # Otherwise, show the frame corresponding to the slider.
+                    viewer_slot.image(processed_frames[slider_val], 
+                                      caption=f"Frame {slider_val}", 
+                                      use_column_width=True)
             frame_count += 1
             success, image = vidcap.read()
 
         progress_text.text("Video processing complete!")
         progress_bar.progress(100)
 
+        # --- Video Download Section ---
         if processed_frames:
             temp_video_file = tempfile.NamedTemporaryFile(delete=False, suffix='.mp4')
             fourcc = cv2.VideoWriter_fourcc(*'mp4v')