Update app.py

app.py

@@ -5,7 +5,7 @@ import streamlit as st
 import PIL
 from ultralytics import YOLO
 
-# Ensure your model path points directly to the .pt file (not an HTML page)
+# Ensure your model path points directly to the .pt file
 model_path = 'https://huggingface.co/spaces/tstone87/ccr-colorado/resolve/main/best.pt'
 
 st.set_page_config(
@@ -29,7 +29,7 @@ with st.sidebar:
     progress_bar = st.progress(0)
 
 # --- MAIN PAGE TITLE AND IMAGES ---
-st.title("WildfireWatch: Detecting Wildfire using AI")
+st.title("Fire Watch: Detecting fire using AI vision models")
 col1, col2 = st.columns(2)
 with col1:
     st.image("https://huggingface.co/spaces/tstone87/ccr-colorado/resolve/main/Fire_1.jpeg", use_column_width=True)
@@ -66,19 +66,29 @@ except Exception as ex:
     st.error(f"Unable to load model. Check the specified path: {model_path}")
     st.error(ex)
 
-# --- SESSION STATE FOR PROCESSED FRAMES ---
+# --- SESSION STATE SETUP ---
 if "processed_frames" not in st.session_state:
     st.session_state["processed_frames"] = []
 
-# We'll keep the detection results for each frame (if you want them)
+# If you want bounding box data per frame:
 if "frame_detections" not in st.session_state:
     st.session_state["frame_detections"] = []
 
-# --- WHEN USER CLICKS DETECT ---
-if st.sidebar.button("Let's Detect Wildfire"):
+# We'll store the shortened video data so the download button remains visible
+if "shortened_video_data" not in st.session_state:
+    st.session_state["shortened_video_data"] = None
+if "shortened_video_ready" not in st.session_state:
+    st.session_state["shortened_video_ready"] = False
+
+# --- DETECT BUTTON ---
+if st.sidebar.button("Let's Detect fire"):
     if not source_file:
         st.warning("No file uploaded!")
     elif file_type == 'image':
+        # Reset previous video data
+        st.session_state["shortened_video_ready"] = False
+        st.session_state["shortened_video_data"] = None
+
         # IMAGE DETECTION
         res = model.predict(uploaded_image, conf=confidence)
         boxes = res[0].boxes
@@ -89,10 +99,11 @@ if st.sidebar.button("Let's Detect Wildfire"):
                 for box in boxes:
                     st.write(box.xywh)
     else:
-        # VIDEO DETECTION
-        # Clear previous frames from session_state
+        # Reset previous frames and video data
         st.session_state["processed_frames"] = []
         st.session_state["frame_detections"] = []
+        st.session_state["shortened_video_ready"] = False
+        st.session_state["shortened_video_data"] = None
 
         processed_frames = st.session_state["processed_frames"]
         frame_detections = st.session_state["frame_detections"]
@@ -131,8 +142,7 @@ if st.sidebar.button("Let's Detect Wildfire"):
                 res_plotted = res[0].plot()[:, :, ::-1]
 
                 processed_frames.append(res_plotted)
-                # If you want to store bounding boxes for each frame:
-                frame_detections.append(res[0].boxes)
+                frame_detections.append(res[0].boxes)  # optional
 
                 # Update progress
                 if total_frames > 0:
@@ -149,7 +159,7 @@ if st.sidebar.button("Let's Detect Wildfire"):
         progress_text.text("Video processing complete!")
         progress_bar.progress(100)
 
-        # Create shortened video from processed frames
+        # Create shortened video
         if processed_frames:
             temp_video_file = tempfile.NamedTemporaryFile(delete=False, suffix='.mp4')
             fourcc = cv2.VideoWriter_fourcc(*'mp4v')
@@ -158,33 +168,36 @@ if st.sidebar.button("Let's Detect Wildfire"):
                 out.write(frame)
             out.release()
 
-            st.success("Shortened video created successfully!")
+            # Store the video data in session_state
             with open(temp_video_file.name, 'rb') as video_file:
-                st.download_button(
-                    label="Download Shortened Video",
-                    data=video_file.read(),
-                    file_name="shortened_video.mp4",
-                    mime="video/mp4"
-                )
+                st.session_state["shortened_video_data"] = video_file.read()
+                st.session_state["shortened_video_ready"] = True
 
+            st.success("Shortened video created successfully!")
         else:
             st.error("No frames were processed from the video.")
 
-# --- DISPLAY THE PROCESSED FRAMES AFTER DETECTION ---
+# --- SHOW THE DOWNLOAD BUTTON IF READY ---
+if st.session_state["shortened_video_ready"] and st.session_state["shortened_video_data"]:
+    st.download_button(
+        label="Download Shortened Video",
+        data=st.session_state["shortened_video_data"],
+        file_name="shortened_video.mp4",
+        mime="video/mp4"
+    )
+
+# --- DISPLAY PROCESSED FRAMES IF ANY ---
 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,
@@ -195,7 +208,8 @@ if st.session_state["processed_frames"]:
         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:
+
+        # Optionally show bounding box data
         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]: