Update app.py

app.py

@@ -4,8 +4,10 @@ import base64
 import cv2
 import streamlit as st
 import PIL
-from ultralytics import YOLO
 import requests
+import imageio
+from ultralytics import YOLO
+from huggingface_hub import hf_hub_download
 
 ###############################################################################
 # Helper: Embed an HTML5 video that autoplays (muted) with controls.
@@ -35,9 +37,21 @@ if "shortened_video_ready" not in st.session_state:
     st.session_state["shortened_video_ready"] = False
 
 ###############################################################################
-# Configure YOLO model path and page layout
+# Download YOLO Model from Hugging Face
+###############################################################################
+repo_id = "tstone87/ccr-colorado"
+model_filename = "best.pt"
+
+try:
+    local_model_path = hf_hub_download(repo_id=repo_id, filename=model_filename)
+    model = YOLO(local_model_path)
+except Exception as ex:
+    st.error(f"Unable to load model. Check model path: {model_filename}")
+    st.error(ex)
+
+###############################################################################
+# Configure Streamlit Page Layout
 ###############################################################################
-model_path = 'https://huggingface.co/spaces/tstone87/ccr-colorado/resolve/main/best.pt'
 st.set_page_config(
     page_title="Fire Detection: Original vs. Processed Video",
     page_icon="🔥",
@@ -46,7 +60,7 @@ st.set_page_config(
 )
 
 ###############################################################################
-# SIDEBAR: Video input options, confidence, sampling options, and example selection
+# SIDEBAR: Video input options
 ###############################################################################
 with st.sidebar:
     st.header("Video Input Options")
@@ -65,7 +79,7 @@ with st.sidebar:
     )
     progress_text = st.empty()
     progress_bar = st.progress(0)
-    download_placeholder = st.empty()  # This placeholder will hold the download button
+    download_placeholder = st.empty()  # Placeholder for download button
 
 ###############################################################################
 # MAIN TITLE
@@ -73,33 +87,27 @@ with st.sidebar:
 st.title("Fire Detection: Original vs. Processed Video")
 
 ###############################################################################
-# Load YOLO model
-###############################################################################
-try:
-    model = YOLO(model_path)
-except Exception as ex:
-    st.error(f"Unable to load model. Check model path: {model_path}")
-    st.error(ex)
-
-###############################################################################
-# Determine source video(s): Example pair or uploaded file.
+# Load Example Video Data
 ###############################################################################
 original_video_data = None
-processed_video_data = None  # For example pairs
+processed_video_data = None  
 
 if example_option != "None":
     # Use example videos from remote 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"
-    elif example_option == "LA Example":
-        orig_url = "https://huggingface.co/spaces/tstone87/ccr-colorado/resolve/main/LA1.mp4"
-        proc_url = "https://huggingface.co/spaces/tstone87/ccr-colorado/resolve/main/LA2.mp4"
-    try:
-        original_video_data = requests.get(orig_url).content
-        processed_video_data = requests.get(proc_url).content
-    except Exception as ex:
-        st.error("Error loading example videos. Check your URLs.")
+    example_videos = {
+        "T Example": ("T1.mp4", "T2.mpg"),
+        "LA Example": ("LA1.mp4", "LA2.mp4")
+    }
+    orig_filename, proc_filename = example_videos.get(example_option, (None, None))
+
+    if orig_filename and proc_filename:
+        try:
+            orig_url = f"https://huggingface.co/spaces/tstone87/ccr-colorado/resolve/main/{orig_filename}"
+            proc_url = f"https://huggingface.co/spaces/tstone87/ccr-colorado/resolve/main/{proc_filename}"
+            original_video_data = requests.get(orig_url).content
+            processed_video_data = requests.get(proc_url).content
+        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:
@@ -111,18 +119,12 @@ else:
             with open(buf.name, "rb") as f:
                 original_video_data = f.read()
         else:
-            tfile = tempfile.NamedTemporaryFile(delete=False, suffix=".mp4")
-            tfile.write(source_file.read())
-            tfile.flush()
-            with open(tfile.name, "rb") as vf:
-                original_video_data = vf.read()
-            # Open with OpenCV for processing.
-            vidcap = cv2.VideoCapture(tfile.name)
-    else:
-        st.info("Please select an example pair or upload a video file.")
+            with tempfile.NamedTemporaryFile(delete=False, suffix=".mp4") as tfile:
+                tfile.write(source_file.read())
+                original_video_data = tfile.name
 
 ###############################################################################
-# Layout: Two columns for Original and Processed videos.
+# Layout: Two Columns for Original and Processed Videos
 ###############################################################################
 col1, col2 = st.columns(2)
 
@@ -135,7 +137,6 @@ with col1:
 
 with col2:
     st.subheader("Result File")
-    # Create a dedicated placeholder for the processed video.
     viewer_slot = st.empty()
     if example_option != "None":
         if processed_video_data:
@@ -146,73 +147,40 @@ with col2:
         viewer_slot.info("Processed video will appear here once detection is run.")
 
 ###############################################################################
-# DETECTION: Process the uploaded video if no example is selected.
+# Process 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.
         st.session_state["processed_frames"] = []
-        st.session_state["shortened_video_data"] = None
-        st.session_state["shortened_video_ready"] = False
-
         processed_frames = st.session_state["processed_frames"]
 
+        vid_reader = imageio.get_reader(original_video_data)
+        fps = vid_reader.get_meta_data()['fps']
+        width, height = vid_reader.get_meta_data()['size']
+
         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
-        elif video_option == "1 fps":
-            sample_interval = int(orig_fps) if orig_fps > 0 else 1
-            output_fps = 1
-        elif video_option == "1 frame per 5 seconds":
-            sample_interval = int(orig_fps * 5) if orig_fps > 0 else 5
-            output_fps = 1
-        elif video_option == "1 frame per 10 seconds":
-            sample_interval = int(orig_fps * 10) if orig_fps > 0 else 10
-            output_fps = 1
-        elif video_option == "1 frame per 15 seconds":
-            sample_interval = int(orig_fps * 15) if orig_fps > 0 else 15
-            output_fps = 1
-        else:
-            sample_interval = 1
-            output_fps = orig_fps
-
-        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)
+        total_frames = vid_reader.get_length()
+
+        for frame in vid_reader:
+            res = model.predict(frame, conf=confidence)
+            res_plotted = res[0].plot()[:, :, ::-1]
+            processed_frames.append(res_plotted)
+
+            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))
             frame_count += 1
-            success, image = vidcap.read()
 
         progress_text.text("Video processing complete!")
         progress_bar.progress(100)
 
-        # Create shortened video from processed frames.
         if processed_frames:
             temp_video_file = tempfile.NamedTemporaryFile(delete=False, suffix='.mp4')
-            # 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))
+            fourcc = cv2.VideoWriter_fourcc(*'mp4v')  # Use widely supported codec
+            out = cv2.VideoWriter(temp_video_file.name, fourcc, fps, (width, height))
+
             for frame in processed_frames:
-                frame_out = cv2.convertScaleAbs(frame)
-                out.write(frame_out)
+                out.write(frame)
             out.release()
 
             with open(temp_video_file.name, 'rb') as video_file:
@@ -220,16 +188,12 @@ 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.
+# Show Download Button if Ready
 ###############################################################################
-if st.session_state["shortened_video_ready"] and st.session_state["shortened_video_data"]:
+if st.session_state["shortened_video_ready"]:
     download_placeholder.download_button(
         label="Download Processed Video",
         data=st.session_state["shortened_video_data"],