app.py

import os
import tempfile
import cv2
import streamlit as st
import PIL
import requests
from ultralytics import YOLO
import time
import numpy as np
import imageio_ffmpeg as ffmpeg
import base64

# Page config first
st.set_page_config(
    page_title="Fire Watch: Fire and Smoke Detection with an AI Vision Model",
    page_icon="🔥",
    layout="wide",
    initial_sidebar_state="expanded"
)

# Model path
model_path = 'https://huggingface.co/spaces/tstone87/ccr-colorado/resolve/main/best.pt'

# 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] = [] if key == "processed_frames" else 0 if key == "slider_value" else None

# Sidebar
with st.sidebar:
    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", 10, 100, 20)) / 100
    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
    }
    video_option = st.selectbox("Output Frame Rate", list(fps_options.keys()))
    process_button = st.button("Detect fire")
    progress_bar = st.progress(0)
    progress_text = st.empty()
    download_slot = st.empty()

# Main page
st.title("Fire Watch: AI-Powered Fire and Smoke 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)

st.markdown("""
Early wildfire detection using YOLOv8 AI vision model. See examples below or upload your own content!
Hover over example videos to see detection results!
""")

# Function to create synchronized video pair HTML with hover images
def create_synced_video_pair(orig_url, proc_url, pair_id, hover_image_url):
    try:
        orig_bytes = requests.get(orig_url).content
        proc_bytes = requests.get(proc_url).content
        hover_bytes = requests.get(hover_image_url).content
        orig_b64 = base64.b64encode(orig_bytes).decode('utf-8')
        proc_b64 = base64.b64encode(proc_bytes).decode('utf-8')
        hover_b64 = base64.b64encode(hover_bytes).decode('utf-8')
        
        html = f"""
        <div style="display: flex; gap: 10px; margin-bottom: 20px; position: relative;">
            <div style="flex: 1; position: relative;">
                <h4>Original</h4>
                <video id="orig_{pair_id}" width="100%" controls preload="metadata">
                    <source src="data:video/mp4;base64,{orig_b64}" type="video/mp4">
                    Your browser does not support the video tag.
                </video>
            </div>
            <div style="flex: 1; position: relative;">
                <h4>Processed</h4>
                <video id="proc_{pair_id}" width="100%" controls preload="metadata">
                    <source src="data:video/mp4;base64,{proc_b64}" type="video/mp4">
                    Your browser does not support the video tag.
                </video>
            </div>
            <img id="hover_{pair_id}" src="data:image/jpeg;base64,{hover_b64}" 
                 style="position: absolute; top: 50%; left: 50%; transform: translate(-50%, -50%); 
                        max-width: 80%; max-height: 80%; display: none; z-index: 10;">
        </div>
        <script>
            document.addEventListener('DOMContentLoaded', function() {{
                const origVideo = document.getElementById('orig_{pair_id}');
                const procVideo = document.getElementById('proc_{pair_id}');
                const hoverImg = document.getElementById('hover_{pair_id}');
                const container = origVideo.parentElement.parentElement;
                
                if (origVideo && procVideo && hoverImg) {{
                    origVideo.addEventListener('play', function() {{
                        procVideo.currentTime = origVideo.currentTime;
                        procVideo.play().catch(e => console.log('Play error:', e));
                    }});
                    procVideo.addEventListener('play', function() {{
                        origVideo.currentTime = procVideo.currentTime;
                        origVideo.play().catch(e => console.log('Play error:', e));
                    }});
                    
                    origVideo.addEventListener('pause', function() {{
                        procVideo.pause();
                    }});
                    procVideo.addEventListener('pause', function() {{
                        origVideo.pause();
                    }});
                    
                    origVideo.addEventListener('seeked', function() {{
                        procVideo.currentTime = origVideo.currentTime;
                    }});
                    procVideo.addEventListener('seeked', function() {{
                        origVideo.currentTime = procVideo.currentTime;
                    }});

                    // Hover events
                    container.addEventListener('mouseover', function() {{
                        hoverImg.style.display = 'block';
                    }});
                    container.addEventListener('mouseout', function() {{
                        hoverImg.style.display = 'none';
                    }});
                }} else {{
                    console.log('Elements not found for {pair_id}');
                }}
            }});
        </script>
        """
        return html
    except Exception as e:
        return f"<p>Error loading videos or hover image: {str(e)}</p>"

if not source_file:
    st.info("Please upload a file to begin.")

st.header("Your Results")
result_cols = st.columns(2)
viewer_slot = st.empty()

# Example videos with synchronization and hover images
st.header("Example Results")
examples = [
    ("T Example", "T1.mp4", "T2.mp4", "File_3a.jpg"),
    ("LA Example", "LA1.mp4", "LA2.mp4", "File_4a.jpg")
]
for title, orig_file, proc_file, hover_file in examples:
    st.subheader(title)
    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}"
    hover_url = f"https://huggingface.co/spaces/tstone87/ccr-colorado/resolve/main/{hover_file}"
    pair_id = title.replace(" ", "").lower()
    video_html = create_synced_video_pair(orig_url, proc_url, pair_id, hover_url)
    st.markdown(video_html, unsafe_allow_html=True)

# Load model
try:
    model = YOLO(model_path)
except Exception as ex:
    st.error(f"Model loading failed: {str(ex)}")
    model = None

# 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)
        
        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))
        
        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
        
        success, frame = vidcap.read()
        while success:
            if frame_count % sample_interval == 0:
                res = model.predict(frame, conf=confidence)
                processed_frame = res[0].plot()[:, :, ::-1]
                if not processed_frame.flags['C_CONTIGUOUS']:
                    processed_frame = np.ascontiguousarray(processed_frame)
                st.session_state.processed_frames.append(processed_frame)
                
                processed_count += 1
                elapsed = time.time() - st.session_state.start_time
                progress = frame_count / total_frames
                
                if elapsed > 0 and processed_count > 0:
                    time_per_frame = elapsed / processed_count
                    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(min(progress, 1.0))
                progress_text.text(f"Progress: {progress:.1%} | ETA: {eta_str}")
            
            frame_count += 1
            success, frame = vidcap.read()
        
        vidcap.release()
        os.unlink(tmp.name)
        
        if st.session_state.processed_frames:
            out_path = tempfile.NamedTemporaryFile(suffix='.mp4', delete=False).name
            writer = ffmpeg.write_frames(
                out_path,
                (width, height),
                fps=output_fps or orig_fps,
                codec='libx264',
                pix_fmt_in='bgr24',
                pix_fmt_out='yuv420p'
            )
            writer.send(None)  # Initialize writer
            
            for frame in st.session_state.processed_frames:
                writer.send(frame)
            writer.close()
            
            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_fire_analysis.mp4",
                mime="video/mp4"
            )

if not source_file:
    st.info("Please upload a file to begin.")