Update app.py

app.py

@@ -1,25 +1,62 @@
 import os
 import tempfile
+import base64
 import cv2
 import streamlit as st
 import PIL
 from ultralytics import YOLO
 
-# Ensure your model path points directly to the .pt file
-model_path = 'https://huggingface.co/spaces/tstone87/ccr-colorado/resolve/main/best.pt'
+###############################################################################
+# Helper function: Display an HTML5 video with autoplay, controls, and muted
+###############################################################################
+def show_autoplay_video(video_data: bytes, title: str = "Video"):
+    if not video_data:
+        st.warning(f"No {title} video available.")
+        return
+    video_base64 = base64.b64encode(video_data).decode()
+    video_html = f"""
+    <h4>{title}</h4>
+    <video width="100%" height="auto" controls autoplay muted>
+        <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 (for uploaded files)
+###############################################################################
+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
 
+###############################################################################
+# Configure YOLO model path and Streamlit page
+###############################################################################
+model_path = 'https://huggingface.co/spaces/tstone87/ccr-colorado/resolve/main/best.pt'
 st.set_page_config(
-    page_title="Fire Watch using AI vision models",
+    page_title="Fire Detection: Original vs. Processed Video",
     page_icon="🔥",
     layout="wide",
     initial_sidebar_state="expanded"
 )
 
-# --- SIDEBAR ---
+###############################################################################
+# SIDEBAR: Upload file, set confidence, video option, and select an example pair
+###############################################################################
 with st.sidebar:
-    st.header("IMAGE/VIDEO UPLOAD")
-    source_file = st.file_uploader("Choose an image or video...",
-                                   type=("jpg", "jpeg", "png", "bmp", "webp", "mp4"))
+    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=("jpg", "jpeg", "png", "bmp", "webp", "mp4")
+    )
     confidence = float(st.slider("Select Model Confidence", 25, 100, 40)) / 100
     video_option = st.selectbox(
         "Select Video Shortening Option",
@@ -28,189 +65,172 @@ with st.sidebar:
     progress_text = st.empty()
     progress_bar = st.progress(0)
 
-# --- MAIN PAGE TITLE AND IMAGES ---
-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)
-with col2:
-    st.image("https://huggingface.co/spaces/tstone87/ccr-colorado/resolve/main/Fire_3.png", use_column_width=True)
-
-st.markdown("""
-Fires in Colorado present a serious challenge, threatening urban communities, highways, and even remote areas. 
-Early detection is critical. WildfireWatch leverages YOLOv8 for real-time fire and smoke detection 
-in images and videos.
-""")
-st.markdown("---")
-st.header("Fire Detection:")
-
-# --- DISPLAY UPLOADED FILE ---
-col1, col2 = st.columns(2)
-if source_file:
-    file_type = source_file.type.split('/')[0]
-    if file_type == 'image':
-        uploaded_image = PIL.Image.open(source_file)
-        st.image(uploaded_image, caption="Uploaded Image", use_column_width=True)
-    else:
-        # Temporarily store the uploaded video
-        tfile = tempfile.NamedTemporaryFile(delete=False)
-        tfile.write(source_file.read())
-        vidcap = cv2.VideoCapture(tfile.name)
-else:
-    st.info("Please upload an image or video file to begin.")
+###############################################################################
+# MAIN PAGE TITLE
+###############################################################################
+st.title("Fire Detection: Original vs. Processed Video")
 
-# --- LOAD YOLO MODEL ---
+###############################################################################
+# Load YOLO model
+###############################################################################
 try:
     model = YOLO(model_path)
 except Exception as ex:
-    st.error(f"Unable to load model. Check the specified path: {model_path}")
+    st.error(f"Unable to load model. Check model path: {model_path}")
     st.error(ex)
 
-# --- SESSION STATE SETUP ---
-if "processed_frames" not in st.session_state:
+###############################################################################
+# Determine source: Example or Uploaded File
+###############################################################################
+original_video_data = None
+processed_video_data = None  # For example pairs, these are loaded directly
+
+if example_option != "None":
+    # An example pair was chosen. Load the videos from disk.
+    if example_option == "T Example":
+        # T1.mp4: original, T2.mpg: processed (analysis completed video)
+        try:
+            with open("T1.mp4", "rb") as f:
+                original_video_data = f.read()
+            with open("T2.mpg", "rb") as f:
+                processed_video_data = f.read()
+        except Exception as ex:
+            st.error("Error loading T Example videos. Ensure T1.mp4 and T2.mpg are in your repo.")
+    elif example_option == "LA Example":
+        # LA1.mp4: original, LA2.mp4: processed
+        try:
+            with open("LA1.mp4", "rb") as f:
+                original_video_data = f.read()
+            with open("LA2.mp4", "rb") as f:
+                processed_video_data = f.read()
+        except Exception as ex:
+            st.error("Error loading LA Example videos. Ensure LA1.mp4 and LA2.mp4 are in your repo.")
+else:
+    # No example selected. Use uploaded file if available.
+    if source_file:
+        file_type = source_file.type.split('/')[0]
+        if file_type == 'image':
+            # For images, simply show the uploaded image (and detection result below)
+            original_image = PIL.Image.open(source_file)
+            # Convert image to bytes for display if needed
+            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()  # Actually, this is just an image preview.
+        else:
+            # For video, save to a temporary file and load its bytes.
+            tfile = tempfile.NamedTemporaryFile(delete=False)
+            tfile.write(source_file.read())
+            tfile.flush()
+            with open(tfile.name, "rb") as vf:
+                original_video_data = vf.read()
+            # Also open video with OpenCV for processing below.
+            vidcap = cv2.VideoCapture(tfile.name)
+    else:
+        st.info("Please select an example pair or upload a file.")
+
+###############################################################################
+# Display the Original and Result columns side-by-side
+###############################################################################
+col1, col2 = st.columns(2)
+
+# Left column: Original video
+with col1:
+    st.subheader("Original File")
+    if original_video_data:
+        show_autoplay_video(original_video_data, title="Original")
+    else:
+        st.info("No original video available.")
+
+###############################################################################
+# DETECTION: For uploaded video files (not example pairs) run YOLO analysis
+###############################################################################
+# We only run detection if no example pair is selected and if an upload is provided.
+if example_option == "None" and source_file and source_file.type.split('/')[0] != 'image':
+    # Reset processed frames for a new analysis
     st.session_state["processed_frames"] = []
+    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]
+            st.session_state["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}")
+        frame_count += 1
+        success, image = vidcap.read()
 
-# If you want bounding box data per frame:
-if "frame_detections" not in st.session_state:
-    st.session_state["frame_detections"] = []
+    progress_text.text("Video processing complete!")
+    progress_bar.progress(100)
 
-# 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
+    # Create shortened video from processed frames
+    processed_frames = st.session_state["processed_frames"]
+    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, output_fps, (width, height))
+        for frame in processed_frames:
+            out.write(frame)
+        out.release()
 
-# --- 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
-        res_plotted = res[0].plot()[:, :, ::-1]
-        with col2:
-            st.image(res_plotted, caption='Detected Image', use_column_width=True)
-            with st.expander("Detection Results"):
-                for box in boxes:
-                    st.write(box.xywh)
-    else:
-        # 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"]
-
-        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
+        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
 
-        success, image = vidcap.read()
-        while success:
-            if frame_count % sample_interval == 0:
-                # Run detection
-                res = model.predict(image, conf=confidence)
-                res_plotted = res[0].plot()[:, :, ::-1]
-
-                processed_frames.append(res_plotted)
-                frame_detections.append(res[0].boxes)  # optional
-
-                # 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}")
-
-            frame_count += 1
-            success, image = vidcap.read()
-
-        # Processing complete
-        progress_text.text("Video processing complete!")
-        progress_bar.progress(100)
-
-        # Create shortened video
-        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, output_fps, (width, height))
-            for frame in processed_frames:
-                out.write(frame)
-            out.release()
-
-            # Store the video data in session_state
-            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("Shortened video created successfully!")
-        else:
-            st.error("No frames were processed from the video.")
+        st.success("Processed video created successfully!")
 
-# --- SHOW THE DOWNLOAD BUTTON IF READY ---
+###############################################################################
+# Right column: Display the Processed (Result) video
+###############################################################################
+with col2:
+    st.subheader("Result File")
+    # For example pairs, use the preloaded processed_video_data
+    if processed_video_data:
+        show_autoplay_video(processed_video_data, title="Processed")
+    # Otherwise, if a processed video has been generated from an upload, show it
+    elif st.session_state["shortened_video_ready"] and st.session_state["shortened_video_data"]:
+        show_autoplay_video(st.session_state["shortened_video_data"], title="Processed")
+    else:
+        st.info("No processed video available yet. Run detection if you uploaded a file.")
+
+###############################################################################
+# 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"]:
     st.download_button(
-        label="Download Shortened Video",
+        label="Download Processed Video",
         data=st.session_state["shortened_video_data"],
-        file_name="shortened_video.mp4",
+        file_name="processed_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:
-        st.image(st.session_state["processed_frames"][0], caption="Frame 0", use_column_width=True)
-        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:
-        frame_idx = st.slider(
-            "Select Frame",
-            min_value=0,
-            max_value=num_frames - 1,
-            value=0,
-            step=1
-        )
-        st.image(st.session_state["processed_frames"][frame_idx],
-                 caption=f"Frame {frame_idx}",
-                 use_column_width=True)
-
-        # 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]:
-                    st.write(box.xywh)