Update app.py

app.py

@@ -1,207 +1,179 @@
 import os
 import tempfile
 import base64
+import time
 import cv2
 import streamlit as st
-import PIL
 import requests
-import imageio
 from ultralytics import YOLO
 from huggingface_hub import hf_hub_download
+import imageio
+import numpy as np
 
-# Move set_page_config to the very top as it must be the first Streamlit command
+# Page config must be first
 st.set_page_config(
-    page_title="Fire Detection: Original vs. Processed Video",
+    page_title="Wildfire Detection Demo",
     page_icon="🔥",
     layout="wide",
     initial_sidebar_state="expanded"
 )
 
-###############################################################################
-# Helper: Embed an HTML5 video that autoplays (muted) with controls.
-###############################################################################
-def show_autoplay_video(video_bytes: bytes, title: str = "Video"):
+# Helper function to display videos
+def show_video(video_bytes: bytes, title: str, loop=True):
     if not video_bytes:
         st.warning(f"No {title} video available.")
         return
     video_base64 = base64.b64encode(video_bytes).decode()
+    loop_attr = "loop" if loop else ""
     video_html = f"""
     <h4>{title}</h4>
-    <video width="100%" controls autoplay muted>
+    <video width="100%" controls autoplay muted {loop_attr}>
       <source src="data:video/mp4;base64,{video_base64}" type="video/mp4">
       Your browser does not support the video tag.
     </video>
     """
     st.markdown(video_html, unsafe_allow_html=True)
 
-###############################################################################
-# Session state initialization (for processed results)
-###############################################################################
-if "processed_frames" not in st.session_state:
-    st.session_state["processed_frames"] = []
-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
-
-###############################################################################
-# Download YOLO Model from Hugging Face
-###############################################################################
-repo_id = "tstone87/ccr-colorado"  # Verify this is correct
-model_filename = "best.pt"
-
-try:
-    # Add error handling for repository access
-    local_model_path = hf_hub_download(
-        repo_id=repo_id, 
-        filename=model_filename,
-        repo_type="model"  # Explicitly specify repo_type
-    )
-    model = YOLO(local_model_path)
-except Exception as ex:
-    st.error(f"Unable to load model. Error details: {str(ex)}")
-    st.error("Please verify:")
-    st.error("1. Repository 'tstone87/ccr-colorado' exists and is accessible")
-    st.error("2. 'best.pt' file is uploaded to the repository")
-    st.error("3. You have proper access permissions")
-    model = None  # Set model to None if loading fails
-
-###############################################################################
-# SIDEBAR: Video input options
-###############################################################################
+# Initialize session state
+for key in ["processed_video", "processing_complete", "start_time", "progress"]:
+    if key not in st.session_state:
+        st.session_state[key] = None if key in ["processed_video", "start_time"] else False if key == "processing_complete" else 0
+
+# Load model
+@st.cache_resource
+def load_model():
+    repo_id = "tstone87/ccr-colorado"
+    filename = "best.pt"
+    try:
+        model_path = hf_hub_download(repo_id=repo_id, filename=filename, repo_type="model")
+        return YOLO(model_path)
+    except Exception as e:
+        st.error(f"Failed to load model: {str(e)}")
+        return None
+
+model = load_model()
+
+# Sidebar
 with st.sidebar:
-    st.header("Video Input Options")
-    example_option = st.selectbox(
-        "Select Example Pair (optional)",
-        ["None", "T Example", "LA Example"]
-    )
-    source_file = st.file_uploader(
-        "Or upload your own file...",
-        type=("mp4", "jpg", "jpeg", "png", "bmp", "webp")
-    )
-    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)
-    download_placeholder = st.empty()
-
-###############################################################################
-# MAIN TITLE
-###############################################################################
-st.title("Fire Detection: Original vs. Processed Video")
-
-###############################################################################
-# Load Example Video Data
-###############################################################################
-original_video_data = None
-processed_video_data = None  
-
-if example_option != "None":
-    example_videos = {
-        "T Example": ("T1.mp4", "T2.mpg"),
-        "LA Example": ("LA1.mp4", "LA2.mp4")
+    st.header("Process Your Own Video")
+    uploaded_file = st.file_uploader("Upload a video", type=["mp4"])
+    confidence = st.slider("Detection Confidence", 0.25, 1.0, 0.4)
+    fps_options = {
+        "Original FPS": None,
+        "3 FPS": 3,
+        "1 FPS": 1,
+        "1 frame/4s": 0.25,
+        "1 frame/10s": 0.1,
+        "1 frame/15s": 0.0667,
+        "1 frame/30s": 0.0333
     }
-    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(f"Error loading example videos: {str(ex)}")
-else:
-    if source_file:
-        file_type = source_file.type.split('/')[0]
-        if file_type == 'image':
-            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:
-            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
-###############################################################################
-col1, col2 = st.columns(2)
-
-with col1:
-    st.subheader("Original File")
-    if original_video_data:
-        show_autoplay_video(original_video_data, title="Original Video")
-    else:
-        st.info("No original video available.")
-
-with col2:
-    st.subheader("Result File")
-    viewer_slot = st.empty()
-    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.")
-    else:
-        viewer_slot.info("Processed video will appear here once detection is run.")
-
-###############################################################################
-# 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") and model is not None:
-        st.session_state["processed_frames"] = []
-        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
-        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
-
-        progress_text.text("Video processing complete!")
-        progress_bar.progress(100)
-
-        if processed_frames:
-            temp_video_file = tempfile.NamedTemporaryFile(delete=False, suffix='.mp4')
-            fourcc = cv2.VideoWriter_fourcc(*'mp4v')
-            out = cv2.VideoWriter(temp_video_file.name, fourcc, fps, (width, height))
-
-            for frame in processed_frames:
-                out.write(frame)
-            out.release()
-
-            with open(temp_video_file.name, 'rb') as video_file:
-                st.session_state["shortened_video_data"] = video_file.read()
-                st.session_state["shortened_video_ready"] = True
-
-            st.success("Processed video created successfully!")
-            show_autoplay_video(st.session_state["shortened_video_data"], title="Processed Video")
-
-###############################################################################
-# Show Download Button if Ready
-###############################################################################
-if st.session_state["shortened_video_ready"]:
-    download_placeholder.download_button(
+    selected_fps = st.selectbox("Output FPS", list(fps_options.keys()), index=0)
+    process_button = st.button("Process Video")
+    progress_bar = st.progress(0)
+    progress_text = st.empty()
+    download_slot = st.empty()
+
+# Main content
+st.title("Wildfire Detection Demo")
+st.markdown("Watch our example videos below or upload your own in the sidebar!")
+
+# Example videos
+example_videos = {
+    "T Example": ("T1.mp4", "T2.mpg"),
+    "LA Example": ("LA1.mp4", "LA2.mp4")
+}
+
+for example_name in example_videos:
+    col1, col2 = st.columns(2)
+    orig_file, proc_file = example_videos[example_name]
+    try:
+        orig_url = f"https://huggingface.co/spaces/tstone87/ccr-colorado/resolve/main/{orig_file}"
+        proc_url = f"https://huggingface.co/spaces/tstone87/ccr-colorado/resolve/main/{proc_file}"
+        orig_data = requests.get(orig_url).content
+        proc_data = requests.get(proc_url).content
+        
+        with col1:
+            show_video(orig_data, f"{example_name} - Original", loop=True)
+        with col2:
+            show_video(proc_data, f"{example_name} - Processed", loop=True)
+    except Exception as e:
+        st.error(f"Failed to load {example_name}: {str(e)}")
+
+# Video processing
+def process_video(video_file, target_fps, confidence):
+    with tempfile.NamedTemporaryFile(suffix=".mp4", delete=False) as tmp:
+        tmp.write(video_file.read())
+        tmp_path = tmp.name
+    
+    try:
+        reader = imageio.get_reader(tmp_path)
+        meta = reader.get_meta_data()
+        original_fps = meta['fps']
+        width, height = meta['size']
+        total_frames = meta['nframes'] if meta['nframes'] != float('inf') else 1000  # Fallback for unknown length
+        
+        output_fps = fps_options[target_fps] if fps_options[target_fps] else original_fps
+        frame_interval = max(1, int(original_fps / output_fps)) if output_fps else 1
+        
+        out_path = tempfile.NamedTemporaryFile(suffix='.mp4', delete=False).name
+        writer = cv2.VideoWriter(out_path, cv2.VideoWriter_fourcc(*'mp4v'), output_fps or original_fps, (width, height))
+        
+        st.session_state.start_time = time.time()
+        processed_count = 0
+        
+        for i, frame in enumerate(reader):
+            if i % frame_interval == 0:
+                frame_rgb = np.array(frame)
+                results = model.predict(frame_rgb, conf=confidence)
+                processed_frame = results[0].plot()[:, :, ::-1]
+                writer.write(processed_frame)
+                
+                processed_count += 1
+                elapsed = time.time() - st.session_state.start_time
+                progress = (i + 1) / total_frames
+                st.session_state.progress = min(progress, 1.0)
+                
+                if elapsed > 0:
+                    frames_left = total_frames - i - 1
+                    time_per_frame = elapsed / processed_count
+                    eta = frames_left * time_per_frame / frame_interval
+                    eta_str = f"{int(eta // 60)}m {int(eta % 60)}s"
+                else:
+                    eta_str = "Calculating..."
+                
+                progress_bar.progress(st.session_state.progress)
+                progress_text.text(f"Progress: {st.session_state.progress:.1%} | ETA: {eta_str}")
+        
+        writer.release()
+        reader.close()
+        
+        with open(out_path, 'rb') as f:
+            return f.read()
+    
+    finally:
+        if os.path.exists(tmp_path):
+            os.unlink(tmp_path)
+        if os.path.exists(out_path):
+            os.unlink(out_path)
+
+# Process uploaded video
+if process_button and uploaded_file and model:
+    with st.spinner("Processing video..."):
+        st.session_state.processed_video = process_video(uploaded_file, selected_fps, confidence)
+        st.session_state.processing_complete = True
+        progress_bar.progress(1.0)
+        progress_text.text("Processing complete!")
+
+# Show processed video and download button
+if st.session_state.processing_complete and st.session_state.processed_video:
+    st.subheader("Your Processed Video")
+    show_video(st.session_state.processed_video, "Processed Result", loop=False)
+    download_slot.download_button(
         label="Download Processed Video",
-        data=st.session_state["shortened_video_data"],
-        file_name="processed_video.mp4",
+        data=st.session_state.processed_video,
+        file_name="processed_wildfire.mp4",
         mime="video/mp4"
-    )
+    )
+
+if not model:
+    st.error("Model loading failed. Please check the repository and model file availability.")