Update app.py

app.py

@@ -1,63 +1,34 @@
 import os
 import tempfile
-import base64
-import time
 import cv2
 import streamlit as st
+import PIL
 import requests
 from ultralytics import YOLO
-from huggingface_hub import hf_hub_download
-import imageio
+import time
 import numpy as np
 
-# Page config must be first
+# Page config first
 st.set_page_config(
-    page_title="Wildfire Detection Demo",
+    page_title="WildfireWatch: AI Detection",
     page_icon="🔥",
     layout="wide",
     initial_sidebar_state="expanded"
 )
 
-# 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 {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)
+# Model path
+model_path = 'https://huggingface.co/spaces/tstone87/ccr-colorado/resolve/main/best.pt'
 
-# Initialize session state
-for key in ["processed_video", "processing_complete", "start_time", "progress"]:
+# Session state initialization
+for key in ["processed_frames", "slider_value", "processed_video", "start_time"]:
     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()
+        st.session_state[key] = [] if key == "processed_frames" else 0 if key == "slider_value" else None
 
 # Sidebar
 with st.sidebar:
-    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)
+    st.header("Upload & Settings")
+    source_file = st.file_uploader("Upload image/video", type=["jpg", "jpeg", "png", "bmp", "webp", "mp4"])
+    confidence = float(st.slider("Confidence Threshold", 25, 100, 40)) / 100
     fps_options = {
         "Original FPS": None,
         "3 FPS": 3,
@@ -67,113 +38,128 @@ with st.sidebar:
         "1 frame/15s": 0.0667,
         "1 frame/30s": 0.0333
     }
-    selected_fps = st.selectbox("Output FPS", list(fps_options.keys()), index=0)
-    process_button = st.button("Process Video")
+    video_option = st.selectbox("Output Frame Rate", list(fps_options.keys()))
+    process_button = st.button("Detect Wildfire")
     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!")
+# Main page
+st.title("WildfireWatch: AI-Powered Detection")
+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)
+with col2:
+    st.image("https://huggingface.co/spaces/tstone87/ccr-colorado/resolve/main/Fire_3.png", use_column_width=True)
 
-# Example videos
-example_videos = {
-    "T Example": ("T1.mp4", "T2.mpg"),
-    "LA Example": ("LA1.mp4", "LA2.mp4")
-}
+st.markdown("""
+Early wildfire detection using YOLOv8 AI vision model. See examples below or upload your own content!
+""")
 
-for example_name in example_videos:
+# Example videos
+st.header("Example Results")
+for example in [("T1.mp4", "T2.mpg"), ("LA1.mp4", "LA2.mp4")]:
     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)}")
+    orig_url = f"https://huggingface.co/spaces/tstone87/ccr-colorado/resolve/main/{example[0]}"
+    proc_url = f"https://huggingface.co/spaces/tstone87/ccr-colorado/resolve/main/{example[1]}"
+    orig_data = requests.get(orig_url).content
+    proc_data = requests.get(proc_url).content
+    with col1:
+        st.video(orig_data)
+    with col2:
+        st.video(proc_data)
+
+st.header("Your Results")
+result_cols = st.columns(2)
+viewer_slot = st.empty()
+
+# Load model
+try:
+    model = YOLO(model_path)
+except Exception as ex:
+    st.error(f"Model loading failed: {str(ex)}")
+    model = None
 
-# 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
+# Processing
+if process_button and source_file and model:
+    st.session_state.processed_frames = []
+    if source_file.type.split('/')[0] == 'image':
+        image = PIL.Image.open(source_file)
+        res = model.predict(image, conf=confidence)
+        result = res[0].plot()[:, :, ::-1]
+        with result_cols[0]:
+            st.image(image, caption="Original", use_column_width=True)
+        with result_cols[1]:
+            st.image(result, caption="Detected", use_column_width=True)
+    else:
+        # Video processing
+        with tempfile.NamedTemporaryFile(suffix=".mp4", delete=False) as tmp:
+            tmp.write(source_file.read())
+            vidcap = cv2.VideoCapture(tmp.name)
         
-        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
+        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))
         
-        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))
+        output_fps = fps_options[video_option] if fps_options[video_option] else orig_fps
+        sample_interval = max(1, int(orig_fps / output_fps)) if output_fps else 1
         
         st.session_state.start_time = time.time()
+        frame_count = 0
         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)
+        success, frame = vidcap.read()
+        while success:
+            if frame_count % sample_interval == 0:
+                res = model.predict(frame, conf=confidence)
+                processed_frame = res[0].plot()[:, :, ::-1]
+                st.session_state.processed_frames.append(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)
+                progress = frame_count / total_frames
                 
-                if elapsed > 0:
-                    frames_left = total_frames - i - 1
+                if elapsed > 0 and processed_count > 0:
                     time_per_frame = elapsed / processed_count
-                    eta = frames_left * time_per_frame / frame_interval
+                    frames_left = (total_frames - frame_count) / sample_interval
+                    eta = frames_left * time_per_frame
                     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}")
+                progress_bar.progress(min(progress, 1.0))
+                progress_text.text(f"Progress: {progress:.1%} | ETA: {eta_str}")
+            
+            frame_count += 1
+            success, frame = vidcap.read()
         
-        writer.release()
-        reader.close()
+        vidcap.release()
+        os.unlink(tmp.name)
         
-        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):
+        if st.session_state.processed_frames:
+            out_path = tempfile.NamedTemporaryFile(suffix='.mp4', delete=False).name
+            writer = cv2.VideoWriter(out_path, cv2.VideoWriter_fourcc(*'mp4v'), output_fps or orig_fps, (width, height))
+            for frame in st.session_state.processed_frames:
+                writer.write(frame)
+            writer.release()
+            
+            with open(out_path, 'rb') as f:
+                st.session_state.processed_video = f.read()
             os.unlink(out_path)
+            
+            progress_bar.progress(1.0)
+            progress_text.text("Processing complete!")
+            with result_cols[0]:
+                st.video(source_file)
+            with result_cols[1]:
+                st.video(st.session_state.processed_video)
+            download_slot.download_button(
+                label="Download Processed Video",
+                data=st.session_state.processed_video,
+                file_name="processed_wildfire.mp4",
+                mime="video/mp4"
+            )
 
-# 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.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.")
+if not source_file:
+    st.info("Please upload a file to begin.")