Update app.py

app.py

@@ -5,12 +5,7 @@ import streamlit as st
 import PIL
 from ultralytics import YOLO
 
-# Required libraries:
-#   streamlit
-#   opencv-python-headless
-#   ultralytics
-#   Pillow
-
+# Ensure your model path points directly to the .pt file (not an HTML page)
 model_path = 'https://huggingface.co/spaces/tstone87/ccr-colorado/resolve/main/best.pt'
 
 st.set_page_config(
@@ -20,17 +15,20 @@ st.set_page_config(
     initial_sidebar_state="expanded"
 )
 
+# --- SIDEBAR ---
 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"])
+    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)
-    # A container where our dynamic slider (frame viewer) will be placed.
-    slider_container = st.empty()
 
+# --- MAIN PAGE TITLE AND IMAGES ---
 st.title("WildfireWatch: Detecting Wildfire using AI")
 col1, col2 = st.columns(2)
 with col1:
@@ -39,49 +37,49 @@ 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.
+# --- DISPLAY UPLOADED FILE ---
 col1, col2 = st.columns(2)
 if source_file:
-    if source_file.type.split('/')[0] == 'image':
+    file_type = source_file.type.split('/')[0]
+    if file_type == 'image':
         uploaded_image = PIL.Image.open(source_file)
         st.image(uploaded_image, caption="Uploaded Image", use_column_width=True)
     else:
+        # Temporarily store the uploaded video
         tfile = tempfile.NamedTemporaryFile(delete=False)
         tfile.write(source_file.read())
         vidcap = cv2.VideoCapture(tfile.name)
 else:
     st.info("Please upload an image or video file to begin.")
 
-# Attempt to load the model.
+# --- LOAD YOLO 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)
 
-# We'll store processed frames persistently in session_state.
+# --- SESSION STATE FOR PROCESSED FRAMES ---
 if "processed_frames" not in st.session_state:
     st.session_state["processed_frames"] = []
 
-# 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
-
-# Container to display the currently viewed frame.
-viewer_slot = st.empty()
+# We'll keep the detection results for each frame (if you want them)
+if "frame_detections" not in st.session_state:
+    st.session_state["frame_detections"] = []
 
-# --- Processing and Viewer Update ---
+# --- WHEN USER CLICKS DETECT ---
 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.
+    elif file_type == 'image':
+        # IMAGE DETECTION
         res = model.predict(uploaded_image, conf=confidence)
         boxes = res[0].boxes
         res_plotted = res[0].plot()[:, :, ::-1]
@@ -91,16 +89,21 @@ if st.sidebar.button("Let's Detect Wildfire"):
                 for box in boxes:
                     st.write(box.xywh)
     else:
-        # For video input, process frames.
+        # VIDEO DETECTION
+        # Clear previous frames from session_state
+        st.session_state["processed_frames"] = []
+        st.session_state["frame_detections"] = []
+
         processed_frames = st.session_state["processed_frames"]
-        frame_count = 0
+        frame_detections = st.session_state["frame_detections"]
 
-        # Video properties.
+        frame_count = 0
         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))
         height = int(vidcap.get(cv2.CAP_PROP_FRAME_HEIGHT))
 
+        # Determine sampling interval
         if video_option == "Original FPS":
             sample_interval = 1
             output_fps = orig_fps
@@ -123,12 +126,15 @@ if st.sidebar.button("Let's Detect Wildfire"):
         success, image = vidcap.read()
         while success:
             if frame_count % sample_interval == 0:
+                # Run detection
                 res = model.predict(image, conf=confidence)
                 res_plotted = res[0].plot()[:, :, ::-1]
+
                 processed_frames.append(res_plotted)
-                st.session_state["processed_frames"] = processed_frames
+                # If you want to store bounding boxes for each frame:
+                frame_detections.append(res[0].boxes)
 
-                # Update progress info.
+                # Update progress
                 if total_frames > 0:
                     progress_pct = int((frame_count / total_frames) * 100)
                     progress_text.text(f"Processing frame {frame_count} / {total_frames} ({progress_pct}%)")
@@ -136,44 +142,14 @@ if st.sidebar.button("Let's Detect Wildfire"):
                 else:
                     progress_text.text(f"Processing frame {frame_count}")
 
-                # --- 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()
 
+        # Processing complete
         progress_text.text("Video processing complete!")
         progress_bar.progress(100)
 
-        # --- Video Download Section ---
+        # Create shortened video from processed frames
         if processed_frames:
             temp_video_file = tempfile.NamedTemporaryFile(delete=False, suffix='.mp4')
             fourcc = cv2.VideoWriter_fourcc(*'mp4v')
@@ -190,5 +166,37 @@ if st.sidebar.button("Let's Detect Wildfire"):
                     file_name="shortened_video.mp4",
                     mime="video/mp4"
                 )
+
         else:
             st.error("No frames were processed from the video.")
+
+# --- DISPLAY THE PROCESSED FRAMES AFTER DETECTION ---
+if st.session_state["processed_frames"]:
+    st.markdown("### Browse Detected Frames")
+    num_frames = len(st.session_state["processed_frames"])
+
+    if num_frames == 1:
+        # Only one frame was processed
+        st.image(st.session_state["processed_frames"][0], caption="Frame 0", use_column_width=True)
+        # If you want to show bounding boxes:
+        if st.session_state["frame_detections"]:
+            with st.expander("Detection Results for Frame 0"):
+                for box in st.session_state["frame_detections"][0]:
+                    st.write(box.xywh)
+    else:
+        # Multiple frames
+        frame_idx = st.slider(
+            "Select Frame",
+            min_value=0,
+            max_value=num_frames - 1,
+            value=0,
+            step=1
+        )
+        st.image(st.session_state["processed_frames"][frame_idx],
+                 caption=f"Frame {frame_idx}",
+                 use_column_width=True)
+        # If you want to show bounding boxes:
+        if st.session_state["frame_detections"]:
+            with st.expander(f"Detection Results for Frame {frame_idx}"):
+                for box in st.session_state["frame_detections"][frame_idx]:
+                    st.write(box.xywh)