Update app.py

app.py

@@ -1,22 +1,21 @@
-import streamlit as st
+import PIL
 import cv2
-import PIL.Image
+import streamlit as st
 from ultralytics import YOLO
 import tempfile
 import time
 import requests
 import numpy as np
-import streamlink
 
 # Page Config
-st.set_page_config(page_title="AI Fire Watch", page_icon="🌍", layout="wide")
+st.set_page_config(page_title="WildfireWatch", page_icon="🔥", layout="wide")
 
-# Lighter Background CSS with Darker Text
+# CSS for layout stability and dark tab text
 st.markdown(
     """
     <style>
     .stApp {
-        background-color: #cdc0b0;
+        background-color: #f5f5f5;
         color: #1a1a1a;
     }
     h1 {
@@ -24,7 +23,7 @@ st.markdown(
     }
     .stTabs > div > button {
         background-color: #e0e0e0;
-        color: #333333;  /* Darker tab text */
+        color: #333333;
         font-weight: bold;
     }
     .stTabs > div > button:hover {
@@ -35,22 +34,12 @@ st.markdown(
         background-color: #ffffff;
         color: #333333;
     }
-    .stButton > button {
-        background-color: #e0e0e0;
-        color: #1a1a1a;
-        font-weight: bold;
-    }
-    .stButton > button:hover {
-        background-color: #d0d0d0;
-        color: #1a1a1a;
-    }
-    /* Fix container height to prevent scrolling */
     .main .block-container {
         max-height: 100vh;
         overflow-y: auto;
     }
     .stImage > img {
-        max-height: 50vh;  /* Limit image height */
+        max-height: 50vh;
         object-fit: contain;
     }
     </style>
@@ -59,11 +48,12 @@ st.markdown(
 )
 
 # Load Model
-model_path = 'https://huggingface.co/spaces/tstone87/ccr-colorado/resolve/main/best.pt'
+model_path = 'https://huggingface.co/spaces/tstone87/ccr-colorado/resolve/main/best.pt'  # Your updated model
 try:
     model = YOLO(model_path)
 except Exception as ex:
-    st.error(f"Model loading failed: {ex}")
+    st.error(f"Unable to load model. Check the specified path: {model_path}")
+    st.error(ex)
     st.stop()
 
 # Initialize Session State
@@ -71,77 +61,61 @@ if 'monitoring' not in st.session_state:
     st.session_state.monitoring = False
 if 'current_webcam_url' not in st.session_state:
     st.session_state.current_webcam_url = None
-if 'yt_monitoring' not in st.session_state:
-    st.session_state.yt_monitoring = False
 
 # Header
-st.title("AI Fire Watch")
-st.markdown("Monitor fire and smoke in real-time with AI vision models.")
+st.title("WildfireWatch: Detecting Wildfire using AI")
+st.markdown("""
+    Wildfires are a major environmental issue, causing substantial losses to ecosystems, human livelihoods, and potentially leading to loss of life. Early detection of wildfires can prevent these losses. Our application uses state-of-the-art YOLOv8 model for real-time wildfire and smoke detection.
+""")
+st.markdown("---")
 
 # Tabs
-tabs = st.tabs(["Upload", "Webcam", "YouTube"])
+tabs = st.tabs(["Upload", "Webcam"])
 
-# Tab 1: Upload
+# Tab 1: Upload (Your original working version)
 with tabs[0]:
-    col1, col2 = st.columns([1, 1])
+    col1, col2 = st.columns(2)
     with col1:
         st.markdown("**Add Your File**")
-        st.write("Upload an image or video (max 200MB) to scan for fire or smoke.")
-        uploaded_file = st.file_uploader("", type=["jpg", "jpeg", "png", "mp4"], label_visibility="collapsed")
+        st.write("Upload an image or video to scan for fire or smoke.")
+        source_file = st.file_uploader("", type=["jpg", "jpeg", "png", "mp4"], label_visibility="collapsed")
         confidence = st.slider("Detection Threshold", 0.25, 1.0, 0.4, key="upload_conf")
+    
     with col2:
         frame_placeholder = st.empty()
         status_placeholder = st.empty()
-        if uploaded_file:
-            try:
-                # Log file details
-                file_size_mb = uploaded_file.size / 1024 / 1024
-                status_placeholder.write(f"File size: {file_size_mb:.2f} MB, Type: {uploaded_file.type}")
-                
-                if file_size_mb > 200:
-                    status_placeholder.error("File exceeds 200MB limit. Please use a smaller file.")
+        if source_file and st.button("Detect Wildfire", key="upload_detect"):
+            file_type = source_file.type.split('/')[0]
+            if file_type == 'image':
+                uploaded_image = PIL.Image.open(source_file)
+                res = model.predict(uploaded_image, conf=confidence)
+                detected_image = res[0].plot()[:, :, ::-1]
+                frame_placeholder.image(detected_image, use_column_width=True)
+                status_placeholder.write(f"Objects detected: {len(res[0].boxes)}")
+            elif file_type == 'video':
+                tfile = tempfile.NamedTemporaryFile(delete=False, suffix='.mp4')
+                tfile.write(source_file.read())
+                tfile.close()
+                vidcap = cv2.VideoCapture(tfile.name)
+                if not vidcap.isOpened():
+                    status_placeholder.error("Failed to open video file.")
                 else:
-                    file_type = uploaded_file.type.split('/')[0]
-                    status_placeholder.write(f"Processing {file_type} file...")
-                    
-                    if file_type == 'image':
-                        image = PIL.Image.open(uploaded_file)
-                        results = model.predict(image, conf=confidence)
-                        detected_frame = results[0].plot()[:, :, ::-1]
+                    success, frame = vidcap.read()
+                    frame_count = 0
+                    while success:
+                        res = model.predict(frame, conf=confidence)
+                        detected_frame = res[0].plot()[:, :, ::-1]
                         frame_placeholder.image(detected_frame, use_column_width=True)
-                        status_placeholder.write(f"Objects detected: {len(results[0].boxes)}")
-                    
-                    elif file_type == 'video':
-                        # Save to temporary file
-                        tfile = tempfile.NamedTemporaryFile(delete=False, suffix='.mp4')
-                        tfile.write(uploaded_file.read())
-                        tfile.close()
-                        
-                        # Process video
-                        cap = cv2.VideoCapture(tfile.name)
-                        if not cap.isOpened():
-                            status_placeholder.error("Failed to open video file. Check format or codec.")
-                        else:
-                            frame_count = 0
-                            while cap.isOpened():
-                                ret, frame = cap.read()
-                                if not ret:
-                                    status_placeholder.write(f"Finished processing video. Processed {frame_count} frames.")
-                                    break
-                                results = model.predict(frame, conf=confidence)
-                                detected_frame = results[0].plot()[:, :, ::-1]
-                                frame_placeholder.image(detected_frame, use_column_width=True)
-                                status_placeholder.write(f"Frame {frame_count}: Objects detected: {len(results[0].boxes)}")
-                                frame_count += 1
-                                time.sleep(0.05)
-                            cap.release()
-                        # Clean up
-                        import os
-                        os.unlink(tfile.name)
-            except Exception as e:
-                status_placeholder.error(f"Error processing file: {str(e)}")
+                        status_placeholder.write(f"Frame {frame_count}: Objects detected: {len(res[0].boxes)}")
+                        success, frame = vidcap.read()
+                        frame_count += 1
+                        time.sleep(0.05)
+                    vidcap.release()
+                    import os
+                    os.unlink(tfile.name)
+                    status_placeholder.write(f"Video processing complete. Processed {frame_count} frames.")
 
-# Tab 2: Webcam
+# Tab 2: Webcam (Enhanced with video and image support)
 with tabs[1]:
     col1, col2 = st.columns([1, 1])
     with col1:
@@ -166,22 +140,24 @@ with tabs[1]:
         timer_placeholder = st.empty()
         
         if st.session_state.monitoring and st.session_state.current_webcam_url:
-            # Try as video stream first
             cap = cv2.VideoCapture(webcam_url)
             is_video_stream = cap.isOpened()
             
             if is_video_stream:
                 status_placeholder.write("Connected to video stream...")
-                while st.session_state.monitoring:
+                while st.session_state.monitoring and cap.isOpened():
                     try:
                         ret, frame = cap.read()
                         if not ret:
                             status_placeholder.error("Video stream interrupted.")
                             break
-                        results = model.predict(frame, conf=confidence)
-                        detected_frame = results[0].plot()[:, :, ::-1]
+                        if webcam_url != st.session_state.current_webcam_url:
+                            status_placeholder.write("URL changed. Stopping video monitoring.")
+                            break
+                        res = model.predict(frame, conf=confidence)
+                        detected_frame = res[0].plot()[:, :, ::-1]
                         frame_placeholder.image(detected_frame, use_column_width=True)
-                        status_placeholder.write(f"Objects detected: {len(results[0].boxes)}")
+                        status_placeholder.write(f"Objects detected: {len(res[0].boxes)}")
                         time.sleep(0.1)  # Fast update for video
                     except Exception as e:
                         status_placeholder.error(f"Video error: {e}")
@@ -189,11 +165,13 @@ with tabs[1]:
                         break
                 cap.release()
             else:
-                # Image-based webcam
                 status_placeholder.write("Monitoring image-based webcam...")
                 while st.session_state.monitoring:
                     try:
                         start_time = time.time()
+                        if webcam_url != st.session_state.current_webcam_url:
+                            status_placeholder.write("URL changed. Stopping image monitoring.")
+                            break
                         response = requests.get(webcam_url, timeout=5)
                         if response.status_code != 200:
                             status_placeholder.error(f"Fetch failed: HTTP {response.status_code}")
@@ -204,71 +182,22 @@ with tabs[1]:
                             status_placeholder.error("Image decoding failed.")
                             break
 
-                        results = model.predict(frame, conf=confidence)
-                        detected_frame = results[0].plot()[:, :, ::-1]
+                        res = model.predict(frame, conf=confidence)
+                        detected_frame = res[0].plot()[:, :, ::-1]
                         frame_placeholder.image(detected_frame, use_column_width=True)
-                        status_placeholder.write(f"Objects detected: {len(results[0].boxes)}")
+                        status_placeholder.write(f"Objects detected: {len(res[0].boxes)}")
                         
-                        # Proper refresh timing for images
                         elapsed = time.time() - start_time
                         remaining = max(0, refresh_rate - elapsed)
                         for i in range(int(remaining), -1, -1):
-                            if not st.session_state.monitoring:
+                            if not st.session_state.monitoring or webcam_url != st.session_state.current_webcam_url:
+                                status_placeholder.write("Monitoring interrupted or URL changed.")
                                 break
                             timer_placeholder.write(f"Next scan: {i}s")
                             time.sleep(1)
-                        
                     except Exception as e:
                         status_placeholder.error(f"Image fetch error: {e}")
                         st.session_state.monitoring = False
                         break
-                if not st.session_state.monitoring:
-                    timer_placeholder.write("Monitoring stopped.")
-
-# Tab 3: YouTube
-with tabs[2]:
-    col1, col2 = st.columns([1, 1])
-    with col1:
-        st.markdown("**YouTube Live**")
-        st.write("Enter a live YouTube URL to auto-analyze the stream.")
-        youtube_url = st.text_input("YouTube URL", "https://www.youtube.com/watch?v=<id>", label_visibility="collapsed")
-        confidence = st.slider("Detection Threshold", 0.25, 1.0, 0.4, key="yt_conf")
-        start_yt = st.button("Start Analysis", key="yt_start")
-        stop_yt = st.button("Stop Analysis", key="yt_stop")
-
-        if start_yt:
-            st.session_state.yt_monitoring = True
-        if stop_yt:
-            st.session_state.yt_monitoring = False
-
-    with col2:
-        frame_placeholder = st.empty()
-        status_placeholder = st.empty()
-        
-        if st.session_state.yt_monitoring and youtube_url and youtube_url != "https://www.youtube.com/watch?v=<id>":
-            try:
-                status_placeholder.write("Initializing stream...")
-                streams = streamlink.streams(youtube_url)
-                if not streams:
-                    status_placeholder.error("No streams found. Check if the URL is a live stream.")
-                else:
-                    stream_url = streams["best"].url
-                    cap = cv2.VideoCapture(stream_url)
-                    if not cap.isOpened():
-                        status_placeholder.error("Unable to open stream.")
-                    else:
-                        status_placeholder.write("Analyzing live stream...")
-                        while st.session_state.yt_monitoring and cap.isOpened():
-                            ret, frame = cap.read()
-                            if not ret:
-                                status_placeholder.error("Stream interrupted.")
-                                break
-                            results = model.predict(frame, conf=confidence)
-                            detected_frame = results[0].plot()[:, :, ::-1]
-                            frame_placeholder.image(detected_frame, use_column_width=True)
-                            status_placeholder.write(f"Objects detected: {len(results[0].boxes)}")
-                            time.sleep(0.1)
-                        cap.release()
-            except Exception as e:
-                status_placeholder.error(f"Error: {e}")
-            st.session_state.yt_monitoring = False
+            if not st.session_state.monitoring:
+                timer_placeholder.write("Monitoring stopped.")