Update app.py

app.py

@@ -1,255 +1,145 @@
+# Import required libraries
 import PIL
 import cv2
 import streamlit as st
 from ultralytics import YOLO
 import tempfile
 import time
-import requests
-import numpy as np
 import os
 
-# Page Config
-st.set_page_config(page_title="WildfireWatch", page_icon="🔥", layout="wide")
+# Replace the relative path to your weight file
+model_path = 'https://huggingface.co/spaces/tstone87/ccr-colorado/resolve/main/best.pt'  # Your correct model
 
-# CSS for layout stability and dark tab text
-st.markdown(
-    """
-    <style>
-    .stApp {
-        background-color: #f5f5f5;
-        color: #1a1a1a;
-    }
-    h1 {
-        color: #1a1a1a;
-    }
-    .stTabs > div > button {
-        background-color: #e0e0e0;
-        color: #333333;
-        font-weight: bold;
-    }
-    .stTabs > div > button:hover {
-        background-color: #d0d0d0;
-        color: #333333;
-    }
-    .stTabs > div > button[aria-selected="true"] {
-        background-color: #ffffff;
-        color: #333333;
-    }
-    .main .block-container {
-        max-height: 100vh;
-        overflow-y: auto;
-    }
-    .stImage > img {
-        max-height: 50vh;
-        object-fit: contain;
-    }
-    </style>
-    """,
-    unsafe_allow_html=True
+# Setting page layout
+st.set_page_config(
+    page_title="WildfireWatch",
+    page_icon="🔥",
+    layout="wide",
+    initial_sidebar_state="expanded"
 )
 
-# Load Model
-model_path = 'https://huggingface.co/spaces/tstone87/ccr-colorado/resolve/main/best.pt'
-try:
-    model = YOLO(model_path)
-except Exception as ex:
-    st.error(f"Unable to load model. Check the specified path: {model_path}")
-    st.error(ex)
-    st.stop()
-
-# Initialize Session State
-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
+# Creating sidebar
+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'))
+    confidence = float(st.slider("Select Model Confidence", 25, 100, 40)) / 100
+    sampling_options = {
+        "Every Frame": 0,
+        "1 FPS": 1,
+        "2 FPS": 2,
+        "5 FPS": 5,
+        "1 frame / 5s": 5,
+        "1 frame / 10s": 10,
+        "1 frame / 15s": 15
+    }
+    sampling_rate = st.selectbox("Analysis Rate", list(sampling_options.keys()), index=1)
 
-# Header
+# Creating main page heading
 st.title("WildfireWatch: Detecting Wildfire using AI")
+
+# Adding informative pictures and description about the motivation for the app
+col1, col2 = st.columns(2)
+with col1:
+    st.image("https://huggingface.co/spaces/ankitkupadhyay/fire_and_smoke/resolve/main/Fire_1.jpeg", use_column_width=True)
+with col2:
+    st.image("https://huggingface.co/spaces/ankitkupadhyay/fire_and_smoke/resolve/main/Fire_2.jpeg", use_column_width=True)
+
 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.
+    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, WildfireWatch, uses state-of-the-art YOLOv8 model for real-time wildfire and smoke detection in images and videos.
 """)
+
 st.markdown("---")
 
-# Tabs
-tabs = st.tabs(["Upload", "Webcam"])
+st.header("Let's Detect Wildfire")
 
-# Tab 1: Upload (Simplified to match original, with new features)
-with tabs[0]:
-    col1, col2 = st.columns(2)
-    with col1:
-        st.markdown("**Add Your File**")
-        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")
-        sampling_options = {
-            "Every Frame": 0,
-            "1 FPS": 1,
-            "2 FPS": 2,
-            "5 FPS": 5,
-            "1 frame / 5s": 5,
-            "1 frame / 10s": 10,
-            "1 frame / 15s": 15
-        }
-        sampling_rate = st.selectbox("Analysis Rate", list(sampling_options.keys()), index=1, key="sampling_rate")
-    
-    with col2:
-        frame_placeholder = st.empty()
-        status_placeholder = st.empty()
-        download_placeholder = st.empty()
-        
-        if source_file and st.button("Detect Wildfire", key="upload_detect"):
-            try:
-                st.write(f"File size: {source_file.size / 1024 / 1024:.2f} MB")  # Diagnostic
-                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':
-                    # Save input video
-                    tfile = tempfile.NamedTemporaryFile(delete=False, suffix='.mp4')
-                    tfile.write(source_file.read())
-                    tfile.close()
-                    
-                    # Open video
-                    vidcap = cv2.VideoCapture(tfile.name)
-                    if not vidcap.isOpened():
-                        status_placeholder.error("Failed to open video file.")
-                    else:
-                        total_frames = int(vidcap.get(cv2.CAP_PROP_FRAME_COUNT))
-                        fps = int(vidcap.get(cv2.CAP_PROP_FPS)) or 30
-                        frame_width = int(vidcap.get(cv2.CAP_PROP_FRAME_WIDTH))
-                        frame_height = int(vidcap.get(cv2.CAP_PROP_FRAME_HEIGHT))
-                        
-                        # Frame sampling
-                        target_rate = sampling_options[sampling_rate]
-                        frame_skip = 1 if target_rate == 0 else max(1, int(fps / target_rate) if target_rate <= 5 else int(fps * target_rate))
-                        
-                        # Output video (only analyzed frames)
-                        output_tfile = tempfile.NamedTemporaryFile(delete=False, suffix='_detected.mp4')
-                        fourcc = cv2.VideoWriter_fourcc(*'mp4v')
-                        output_fps = 1  # Fixed for short compilation
-                        out = cv2.VideoWriter(output_tfile.name, fourcc, output_fps, (frame_width, frame_height))
-                        
-                        success, frame = vidcap.read()
-                        frame_count = 0
-                        processed_count = 0
-                        
-                        while success:
-                            if frame_count % frame_skip == 0:
-                                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"Frame {frame_count}: Objects detected: {len(res[0].boxes)}")
-                                out.write(detected_frame[:, :, ::-1])
-                                processed_count += 1
-                            if total_frames > 0:
-                                progress = (frame_count + 1) / total_frames * 100
-                                st.write(f"Progress: {progress:.1f}% (Analyzed {processed_count} frames)")
-                            success, frame = vidcap.read()
-                            frame_count += 1
-                            time.sleep(0.05)
-                        
-                        vidcap.release()
-                        out.release()
-                        
-                        os.unlink(tfile.name)
-                        with open(output_tfile.name, 'rb') as f:
-                            download_placeholder.download_button(
-                                label="Download Analyzed Video",
-                                data=f,
-                                file_name="analyzed_video.mp4",
-                                mime="video/mp4"
-                            )
-                        status_placeholder.write(f"Video processing complete. Analyzed {processed_count} frames.")
-            except Exception as e:
-                status_placeholder.error(f"Error: {str(e)}")
+# Creating two columns on the main page
+col1, col2 = st.columns(2)
 
-# Tab 2: Webcam (Unchanged)
-with tabs[1]:
-    col1, col2 = st.columns([1, 1])
-    with col1:
-        st.markdown("**Webcam Feed**")
-        st.write("Provide a webcam URL (image or video stream) to monitor for hazards.")
-        webcam_url = st.text_input("Webcam URL", "http://<your_webcam_ip>/current.jpg", label_visibility="collapsed")
-        confidence = st.slider("Detection Threshold", 0.25, 1.0, 0.4, key="webcam_conf")
-        refresh_rate = st.slider("Refresh Rate (seconds)", 1, 60, 30, key="webcam_rate")
-        start = st.button("Begin Monitoring", key="webcam_start")
-        stop = st.button("Stop Monitoring", key="webcam_stop")
+# Adding image to the first column if image is uploaded
+with col1:
+    if source_file:
+        if source_file.type.split('/')[0] == 'image':
+            uploaded_image = PIL.Image.open(source_file)
+            st.image(source_file, caption="Uploaded Image", use_column_width=True)
+        else:
+            tfile = tempfile.NamedTemporaryFile(delete=False)
+            tfile.write(source_file.read())
+            vidcap = cv2.VideoCapture(tfile.name)
 
-        if start:
-            st.session_state.monitoring = True
-            st.session_state.current_webcam_url = webcam_url
-        if stop or (st.session_state.monitoring and webcam_url != st.session_state.current_webcam_url):
-            st.session_state.monitoring = False
-            st.session_state.current_webcam_url = None
+try:
+    model = YOLO(model_path)
+except Exception as ex:
+    st.error(f"Unable to load model. Check the specified path: {model_path}")
+    st.error(ex)
+    st.stop()
 
-    with col2:
-        frame_placeholder = st.empty()
-        status_placeholder = st.empty()
-        timer_placeholder = st.empty()
+if st.sidebar.button('Let\'s Detect Wildfire'):
+    if not source_file:
+        st.error("Please upload a file first!")
+    elif source_file.type.split('/')[0] == 'image':
+        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)
+            try:
+                with st.expander("Detection Results"):
+                    for box in boxes:
+                        st.write(box.xywh)
+            except Exception as ex:
+                st.write("No image is uploaded yet!")
+    else:
+        # Frame sampling setup
+        total_frames = int(vidcap.get(cv2.CAP_PROP_FRAME_COUNT))
+        fps = int(vidcap.get(cv2.CAP_PROP_FPS)) or 30
+        frame_width = int(vidcap.get(cv2.CAP_PROP_FRAME_WIDTH))
+        frame_height = int(vidcap.get(cv2.CAP_PROP_FRAME_HEIGHT))
+        target_rate = sampling_options[sampling_rate]
+        frame_skip = 1 if target_rate == 0 else max(1, int(fps / target_rate) if target_rate <= 5 else int(fps * target_rate))
         
-        if st.session_state.monitoring and st.session_state.current_webcam_url:
-            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 and cap.isOpened():
-                    try:
-                        ret, frame = cap.read()
-                        if not ret:
-                            status_placeholder.error("Video stream interrupted.")
-                            break
-                        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(res[0].boxes)}")
-                        time.sleep(0.1)
-                    except Exception as e:
-                        status_placeholder.error(f"Video error: {e}")
-                        st.session_state.monitoring = False
-                        break
-                cap.release()
-            else:
-                status_placeholder.write("Monitoring image-based webcam...")
-                while st.session_state.monitoring:
+        # Output video setup
+        output_tfile = tempfile.NamedTemporaryFile(delete=False, suffix='_detected.mp4')
+        fourcc = cv2.VideoWriter_fourcc(*'mp4v')
+        output_fps = 1  # Fixed for short compilation
+        out = cv2.VideoWriter(output_tfile.name, fourcc, output_fps, (frame_width, frame_height))
+        
+        success, image = vidcap.read()
+        frame_count = 0
+        processed_count = 0
+        
+        while success:
+            if frame_count % frame_skip == 0:
+                res = model.predict(image, conf=confidence)
+                boxes = res[0].boxes
+                res_plotted = res[0].plot()[:, :, ::-1]
+                with col2:
+                    st.image(res_plotted, caption=f'Detected Frame {frame_count}', use_column_width=True)
                     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}")
-                            break
-                        image_array = np.asarray(bytearray(response.content), dtype=np.uint8)
-                        frame = cv2.imdecode(image_array, cv2.IMREAD_COLOR)
-                        if frame is None:
-                            status_placeholder.error("Image decoding failed.")
-                            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(res[0].boxes)}")
-                        
-                        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 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.")
+                        with st.expander("Detection Results"):
+                            for box in boxes:
+                                st.write(box.xywh)
+                    except Exception as ex:
+                        st.write("No detection results available.")
+                out.write(res_plotted[:, :, ::-1])  # Write only analyzed frame
+                processed_count += 1
+            if total_frames > 0:
+                progress = (frame_count + 1) / total_frames * 100
+                st.write(f"Progress: {progress:.1f}% (Analyzed {processed_count} frames)")
+            success, image = vidcap.read()
+            frame_count += 1
+            time.sleep(0.05)
+        
+        vidcap.release()
+        out.release()
+        os.unlink(tfile.name)
+        
+        with col2:
+            with open(output_tfile.name, 'rb') as f:
+                st.download_button(
+                    label="Download Analyzed Video",
+                    data=f,
+                    file_name="analyzed_video.mp4",
+                    mime="video/mp4"
+                )
+            st.write(f"Video processing complete. Analyzed {processed_count} frames.")