app.py

import streamlit as st
import cv2
import PIL.Image
from ultralytics import YOLO
import tempfile
import time
import requests
import numpy as np
import streamlink

# Page Config
st.set_page_config(page_title="WildfireWatch", page_icon="🔥", layout="wide")

# Load Model
model_path = 'https://huggingface.co/spaces/ankitkupadhyay/fire_and_smoke/resolve/main/best.pt'
try:
    model = YOLO(model_path)
except Exception as ex:
    st.error(f"Model loading failed: {ex}")
    st.stop()

# Minimalist Header
st.title("WildfireWatch")
st.markdown("AI-powered detection of fire and smoke.")

# Tabs
tabs = st.tabs(["Upload", "Webcam", "YouTube"])

# Tab 1: File Upload
with tabs[0]:
    col1, col2 = st.columns([1, 1])
    with col1:
        st.markdown("**Upload an image or video**")
        uploaded_file = st.file_uploader("", type=["jpg", "jpeg", "png", "mp4"], label_visibility="collapsed")
        confidence = st.slider("Confidence", 0.25, 1.0, 0.4, key="upload_conf")
    with col2:
        if uploaded_file:
            file_type = uploaded_file.type.split('/')[0]
            if file_type == 'image':
                image = PIL.Image.open(uploaded_file)
                results = model.predict(image, conf=confidence)
                detected_image = results[0].plot()[:, :, ::-1]
                st.image(detected_image, use_column_width=True)
                st.write(f"Detections: {len(results[0].boxes)}")
            elif file_type == 'video':
                tfile = tempfile.NamedTemporaryFile(delete=False)
                tfile.write(uploaded_file.read())
                cap = cv2.VideoCapture(tfile.name)
                frame_placeholder = st.empty()
                while cap.isOpened():
                    ret, frame = cap.read()
                    if not ret:
                        break
                    results = model.predict(frame, conf=confidence)
                    detected_frame = results[0].plot()[:, :, ::-1]
                    frame_placeholder.image(detected_frame, use_column_width=True)
                    time.sleep(0.05)
                cap.release()

# Tab 2: Webcam / Image URL
with tabs[1]:
    col1, col2 = st.columns([1, 1])
    with col1:
        st.markdown("**Webcam snapshot or stream**")
        webcam_url = st.text_input("URL", "http://<your_webcam_ip>/current.jpg", label_visibility="collapsed")
        confidence = st.slider("Confidence", 0.25, 1.0, 0.4, key="webcam_conf")
        mode = st.radio("", ["Snapshot", "Stream"], label_visibility="collapsed")
        start = st.button("Start", key="webcam_start")
    with col2:
        if start:
            if mode == "Snapshot":
                placeholder = st.empty()
                timer_placeholder = st.empty()
                refresh_interval = 5  # seconds
                while True:
                    start_time = time.time()
                    try:
                        response = requests.get(webcam_url, timeout=5)
                        image_array = np.asarray(bytearray(response.content), dtype=np.uint8)
                        frame = cv2.imdecode(image_array, cv2.IMREAD_COLOR)
                        results = model.predict(frame, conf=confidence)
                        detected_frame = results[0].plot()[:, :, ::-1]
                        placeholder.image(detected_frame, use_column_width=True)
                        elapsed = time.time() - start_time
                        remaining = max(0, refresh_interval - elapsed)
                        timer_placeholder.write(f"Next refresh in: {int(remaining)}s")
                        time.sleep(1)
                        if remaining <= 0:
                            st.experimental_rerun()
                    except Exception as e:
                        st.error(f"Error: {e}")
                        break
            else:
                cap = cv2.VideoCapture(webcam_url)
                frame_placeholder = st.empty()
                while cap.isOpened():
                    ret, frame = cap.read()
                    if not ret:
                        st.error("Stream failed.")
                        break
                    results = model.predict(frame, conf=confidence)
                    detected_frame = results[0].plot()[:, :, ::-1]
                    frame_placeholder.image(detected_frame, use_column_width=True)
                    time.sleep(0.05)

# Tab 3: YouTube Live Stream
with tabs[2]:
    col1, col2 = st.columns([1, 1])
    with col1:
        st.markdown("**YouTube live stream**")
        youtube_url = st.text_input("URL", "https://www.youtube.com/watch?v=<id>", label_visibility="collapsed")
        confidence = st.slider("Confidence", 0.25, 1.0, 0.4, key="yt_conf")
        start = st.button("Start", key="yt_start")
    with col2:
        if start:
            streams = streamlink.streams(youtube_url)
            if streams:
                stream_url = streams["best"].to_url()
                cap = cv2.VideoCapture(stream_url)
                frame_placeholder = st.empty()
                while cap.isOpened():
                    ret, frame = cap.read()
                    if not ret:
                        st.error("Stream failed.")
                        break
                    results = model.predict(frame, conf=confidence)
                    detected_frame = results[0].plot()[:, :, ::-1]
                    frame_placeholder.image(detected_frame, use_column_width=True)
                    time.sleep(0.05)
            else:
                st.error("No stream found.")