Update app.py

app.py

@@ -8,7 +8,7 @@ from ultralytics import YOLO
 import requests
 
 ###############################################################################
-# Helper function to embed an HTML5 video that autoplays (muted) with controls.
+# Helper: Embed an HTML5 video for autoplay (muted) display.
 ###############################################################################
 def show_autoplay_video(video_bytes: bytes, title: str = "Video"):
     if not video_bytes:
@@ -25,7 +25,7 @@ def show_autoplay_video(video_bytes: bytes, title: str = "Video"):
     st.markdown(video_html, unsafe_allow_html=True)
 
 ###############################################################################
-# Session state initialization (for processed results)
+# Session state initialization for processed results.
 ###############################################################################
 if "processed_frames" not in st.session_state:
     st.session_state["processed_frames"] = []
@@ -35,7 +35,7 @@ if "shortened_video_ready" not in st.session_state:
     st.session_state["shortened_video_ready"] = False
 
 ###############################################################################
-# Configure YOLO model path and page layout
+# Configure YOLO model path and page layout.
 ###############################################################################
 model_path = 'https://huggingface.co/spaces/tstone87/ccr-colorado/resolve/main/best.pt'
 st.set_page_config(
@@ -46,11 +46,10 @@ st.set_page_config(
 )
 
 ###############################################################################
-# SIDEBAR: Video input options, confidence, sampling options, and example selection
+# SIDEBAR: Video input options, confidence, sampling options, and example selection.
 ###############################################################################
 with st.sidebar:
     st.header("Video Input Options")
-    # Option to select an example pair; "None" means use an uploaded file.
     example_option = st.selectbox(
         "Select Example Pair (optional)",
         ["None", "T Example", "LA Example"]
@@ -66,8 +65,6 @@ with st.sidebar:
     )
     progress_text = st.empty()
     progress_bar = st.progress(0)
-    # Placeholder for the download button so it persists.
-    download_placeholder = st.empty()
 
 ###############################################################################
 # MAIN TITLE
@@ -75,7 +72,7 @@ with st.sidebar:
 st.title("Fire Detection: Original vs. Processed Video")
 
 ###############################################################################
-# Load YOLO model
+# Load YOLO model.
 ###############################################################################
 try:
     model = YOLO(model_path)
@@ -90,7 +87,7 @@ original_video_data = None
 processed_video_data = None  # For example pairs
 
 if example_option != "None":
-    # Use example videos from remote URLs.
+    # Load example videos from URLs.
     if example_option == "T Example":
         orig_url = "https://huggingface.co/spaces/tstone87/ccr-colorado/resolve/main/T1.mp4"
         proc_url = "https://huggingface.co/spaces/tstone87/ccr-colorado/resolve/main/T2.mpg"
@@ -103,16 +100,17 @@ if example_option != "None":
     except Exception as ex:
         st.error("Error loading example videos. Check your URLs.")
 else:
-    # No example selected. If a file is uploaded, use it.
     if source_file:
         file_type = source_file.type.split('/')[0]
         if file_type == 'image':
+            # For images, show the image preview.
             original_image = PIL.Image.open(source_file)
             buf = tempfile.NamedTemporaryFile(suffix=".png", delete=False)
             original_image.save(buf.name, format="PNG")
             with open(buf.name, "rb") as f:
                 original_video_data = f.read()
         else:
+            # For video uploads, save to a temporary file.
             tfile = tempfile.NamedTemporaryFile(delete=False, suffix=".mp4")
             tfile.write(source_file.read())
             tfile.flush()
@@ -124,7 +122,7 @@ else:
         st.info("Please select an example pair or upload a video file.")
 
 ###############################################################################
-# Layout: Two columns for Original and Processed videos
+# Layout: Two columns for Original and Processed videos.
 ###############################################################################
 col1, col2 = st.columns(2)
 
@@ -135,25 +133,27 @@ with col1:
     else:
         st.info("No original video available.")
 
-with col2:
-    st.subheader("Result File")
-    if example_option != "None":
-        if processed_video_data:
-            show_autoplay_video(processed_video_data, title="Processed Video")
-        else:
-            st.info("No processed video available in example.")
+# Create a placeholder for the processed video display.
+viewer_slot = col2.empty()
+col2.subheader("Result File")
+
+if example_option != "None":
+    if processed_video_data:
+        show_autoplay_video(processed_video_data, title="Processed Video")
     else:
-        if st.session_state["shortened_video_ready"] and st.session_state["shortened_video_data"]:
-            show_autoplay_video(st.session_state["shortened_video_data"], title="Processed Video")
-        else:
-            st.info("Processed video will appear here once detection is run.")
+        col2.info("No processed video available in example.")
+else:
+    if st.session_state["shortened_video_ready"] and st.session_state["shortened_video_data"]:
+        show_autoplay_video(st.session_state["shortened_video_data"], title="Processed Video")
+    else:
+        col2.info("Processed video will appear here once detection is run.")
 
 ###############################################################################
 # DETECTION: Process the uploaded video if no example is selected.
 ###############################################################################
 if example_option == "None" and source_file and source_file.type.split('/')[0] != 'image':
     if st.sidebar.button("Let's Detect Wildfire"):
-        # Reset previous processed results.
+        # Reset previous results.
         st.session_state["processed_frames"] = []
         st.session_state["shortened_video_data"] = None
         st.session_state["shortened_video_ready"] = False
@@ -186,18 +186,22 @@ if example_option == "None" and source_file and source_file.type.split('/')[0] !
             sample_interval = 1
             output_fps = orig_fps
 
+        # Process frames.
         success, image = vidcap.read()
         while success:
             if frame_count % sample_interval == 0:
                 res = model.predict(image, conf=confidence)
                 res_plotted = res[0].plot()[:, :, ::-1]
                 processed_frames.append(res_plotted)
+                # 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}%)")
                     progress_bar.progress(min(100, progress_pct))
                 else:
                     progress_text.text(f"Processing frame {frame_count}")
+                # Update the viewer with the most recent processed frame.
+                viewer_slot.image(res_plotted, caption=f"Frame {frame_count}", use_column_width=True)
             frame_count += 1
             success, image = vidcap.read()
 
@@ -207,7 +211,7 @@ if example_option == "None" and source_file and source_file.type.split('/')[0] !
         # Create shortened video from processed frames.
         if processed_frames:
             temp_video_file = tempfile.NamedTemporaryFile(delete=False, suffix='.mp4')
-            # Use the 'avc1' codec (H.264) for better compatibility
+            # Use 'avc1' codec (H.264) for better compatibility.
             fourcc = cv2.VideoWriter_fourcc(*'avc1')
             out = cv2.VideoWriter(temp_video_file.name, fourcc, output_fps, (width, height))
             for frame in processed_frames:
@@ -220,14 +224,17 @@ if example_option == "None" and source_file and source_file.type.split('/')[0] !
                 st.session_state["shortened_video_ready"] = True
 
             st.success("Processed video created successfully!")
+            # Update the viewer with the final processed video.
+            viewer_slot.empty()
+            show_autoplay_video(st.session_state["shortened_video_data"], title="Processed Video")
         else:
             st.error("No frames were processed from the video.")
 
 ###############################################################################
-# Always display the download button if a processed video is ready.
+# ALWAYS display the download button if a processed video is ready.
 ###############################################################################
 if st.session_state["shortened_video_ready"] and st.session_state["shortened_video_data"]:
-    download_placeholder.download_button(
+    st.download_button(
         label="Download Processed Video",
         data=st.session_state["shortened_video_data"],
         file_name="processed_video.mp4",