app.py

import cv2
import streamlit as st
from ultralytics import YOLO
import time
import numpy as np
from datetime import datetime
import pytz

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

# --- Session State Initialization ---
if "streams" not in st.session_state:
    st.session_state.streams = []
if "num_streams" not in st.session_state:
    st.session_state.num_streams = 1
if "confidence" not in st.session_state:
    st.session_state.confidence = 0.30  # default 30%
if "target_fps" not in st.session_state:
    st.session_state.target_fps = 1.0   # default 1 FPS

# --- Default URLs and Names for 10 Streams ---
default_m3u8_urls = [
    "https://publicstreamer2.cotrip.org/rtplive/070E27555CAM1RP1/playlist.m3u8",  # EB @ York St Denver
    "https://publicstreamer1.cotrip.org/rtplive/225N00535CAM1RP1/playlist.m3u8",  # NB at Iliff Denver
    "https://publicstreamer2.cotrip.org/rtplive/070W28220CAM1RHS/playlist.m3u8",  # WB Half Mile West of I225 Denver
    "https://publicstreamer1.cotrip.org/rtplive/070W26805CAM1RHS/playlist.m3u8",  # 1 mile E of Kipling Denver
    "https://publicstreamer4.cotrip.org/rtplive/076W03150CAM1RP1/playlist.m3u8",  # Main St Hudson
    "https://publicstreamer2.cotrip.org/rtplive/070E27660CAM1NEC/playlist.m3u8",   # EB Colorado Blvd i70 Denver
    "https://publicstreamer2.cotrip.org/rtplive/070W27475CAM1RHS/playlist.m3u8",   # E of Washington St Denver
    "https://publicstreamer3.cotrip.org/rtplive/070W28155CAM1RHS/playlist.m3u8",   # WB Peroia St Underpass Denver
    "https://publicstreamer3.cotrip.org/rtplive/070E11660CAM1RHS/playlist.m3u8",    # Grand Ave Glenwood
    "https://publicstreamer4.cotrip.org/rtplive/070E27890CAM1RHS/playlist.m3u8"   # EB at i270
]
default_names = [
    "EB @ York St Denver",
    "NB at Iliff Denver",
    "WB Half Mile West of I225 Denver",
    "1 mile E of Kipling Denver",
    "Main St Hudson",
    "EB Colorado Blvd i70 Denver",
    "E of Washington St Denver",
    "WB Peroia St Underpass Denver",
    "Grand Ave Glenwood",
    "EB at i270"
]

# --- Sidebar Settings ---
with st.sidebar:
    st.header("Stream Settings")
    # Custom configuration for stream 1 only.
    custom_m3u8 = st.text_input("Custom M3U8 URL for Stream 1 (optional)", value="", key="custom_m3u8")
    custom_name = st.text_input("Custom Webcam Name for Stream 1 (optional)", value="", key="custom_name")
    
    # Choose number of streams (1 to 10)
    num_streams = st.selectbox("Number of Streams", list(range(1, 11)), index=0)
    st.session_state.num_streams = num_streams

    # Global settings for confidence and processing rate.
    confidence = float(st.slider("Confidence Threshold", 5, 100, 30)) / 100
    st.session_state.confidence = confidence
    fps_options = {
        "1 FPS": 1,
        "1 frame/2s": 0.5,
        "1 frame/3s": 0.3333,
        "1 frame/5s": 0.2,
        "1 frame/15s": 0.0667,
        "1 frame/30s": 0.0333
    }
    video_option = st.selectbox("Processing Rate", list(fps_options.keys()), index=3)
    st.session_state.target_fps = fps_options[video_option]

# Update or initialize the streams using defaults (with custom override for stream 1).
if len(st.session_state.streams) != st.session_state.num_streams:
    st.session_state.streams = []
    for i in range(st.session_state.num_streams):
        if i == 0:
            url = custom_m3u8.strip() if custom_m3u8.strip() else default_m3u8_urls[0]
            display_name = custom_name.strip() if custom_name.strip() else default_names[0]
        else:
            url = default_m3u8_urls[i] if i < len(default_m3u8_urls) else ""
            display_name = default_names[i] if i < len(default_names) else f"Stream {i+1}"
        st.session_state.streams.append({
            "current_m3u8_url": url,
            "processed_frame": np.zeros((480, 640, 3), dtype=np.uint8),
            "start_time": time.time(),
            "processed_count": 0,
            "detected_frames": [],
            "last_processed_time": 0,
            "stats_text": "Processing FPS: 0.00\nFrame Delay: 0.00 sec\nTensor Results: No detections",
            "highest_match": 0.0,
            "display_name": display_name
        })
else:
    if st.session_state.num_streams > 0:
        url = custom_m3u8.strip() if custom_m3u8.strip() else default_m3u8_urls[0]
        display_name = custom_name.strip() if custom_name.strip() else default_names[0]
        st.session_state.streams[0]["current_m3u8_url"] = url
        st.session_state.streams[0]["display_name"] = display_name

confidence = st.session_state.confidence
target_fps = st.session_state.target_fps

# --- Load Model ---
model_path = 'https://huggingface.co/spaces/tstone87/ccr-colorado/resolve/main/best.pt'
@st.cache_resource
def load_model():
    return YOLO(model_path)
try:
    model = load_model()
except Exception as ex:
    st.error(f"Model loading failed: {str(ex)}")
    st.stop()

# --- Create Placeholders for Streams in a 2-Column Grid ---
num_streams = st.session_state.num_streams
feed_placeholders = []
stats_placeholders = []
cols = st.columns(2)
for i in range(num_streams):
    col_index = i % 2
    if i >= 2 and col_index == 0:
        cols = st.columns(2)
    feed_placeholders.append(cols[col_index].empty())
    stats_placeholders.append(cols[col_index].empty())
if num_streams == 1:
    _ = st.columns(2)

def update_stream(i):
    current_time = time.time()
    sleep_time = 1.0 / target_fps
    stream_state = st.session_state.streams[i]
    if current_time - stream_state["last_processed_time"] >= sleep_time:
        url = stream_state["current_m3u8_url"]
        cap = cv2.VideoCapture(url)
        if not cap.isOpened():
            stats_placeholders[i].text("Failed to open M3U8 stream.")
            return
        ret, frame = cap.read()
        cap.release()
        if not ret:
            stats_placeholders[i].text("Stream interrupted or ended.")
            return
        res = model.predict(frame, conf=confidence)
        processed_frame = res[0].plot()[:, :, ::-1]
        # Extract detection results.
        tensor_info = "No detections"
        max_conf = 0.0
        try:
            boxes = res[0].boxes
            if boxes is not None and len(boxes) > 0:
                max_conf = float(boxes.conf.max())
                tensor_info = f"Detections: {len(boxes)} | Max Confidence: {max_conf:.2f}"
        except Exception as ex:
            tensor_info = f"Error extracting detections: {ex}"
        # Only update if new detection's confidence is >= current highest.
        if max_conf >= stream_state["highest_match"]:
            stream_state["highest_match"] = max_conf
            stream_state["detected_frames"].append(processed_frame)
        stream_state["processed_count"] += 1
        stream_state["last_processed_time"] = current_time
        mt_time = datetime.now(pytz.timezone('America/Denver')).strftime('%Y-%m-%d %H:%M:%S MT')
        stream_state["processed_frame"] = processed_frame
        stream_state["stats_text"] = (
            f"Processing FPS: {stream_state['processed_count'] / (current_time - stream_state['start_time']):.2f}\n"
            f"{tensor_info}\n"
            f"Highest Match: {stream_state['highest_match']:.2f}"
        )
        feed_placeholders[i].image(
            processed_frame,
            caption=f"Stream {i+1} - {stream_state['display_name']} - {mt_time}",
            use_container_width=True
        )
        stats_placeholders[i].text(stream_state["stats_text"])

# --- Continuous Processing Loop ---
while True:
    for i in range(num_streams):
        update_stream(i)
   # time.sleep(1.0 / target_fps)