Update app.py

app.py

@@ -1,114 +1,112 @@
-# Import required libraries
-import PIL
-import cv2
 import streamlit as st
+import cv2
+import PIL
 from ultralytics import YOLO
 import tempfile
+import time
 
+# ----------------------------------------------------------------
+# Load the model (using a URL to your weight file)
+model_path = 'https://huggingface.co/spaces/tstone87/ccr-colorado/blob/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)
 
-# Replace the relative path to your weight file
-model_path = 'https://huggingface.co/spaces/ankitkupadhyay/fire_and_smoke/blob/main/best.pt'
-
-# Setting page layout
+# ----------------------------------------------------------------
+# Set page configuration
 st.set_page_config(
-    page_title="WildfireWatch",  # Setting page title
-    page_icon="🔥",     # Setting page icon
-    layout="wide",      # Setting layout to wide
-    initial_sidebar_state="expanded"    # Expanding sidebar by default
+    page_title="WildfireWatch",
+    page_icon="🔥",
+    layout="wide",
+    initial_sidebar_state="expanded"
 )
 
-# Creating sidebar
-with st.sidebar:
-    st.header("IMAGE/VIDEO UPLOAD")     # Adding header to sidebar
-    # Adding file uploader to sidebar for selecting images and videos
-    source_file = st.file_uploader(
-        "Choose an image or video...", type=("jpg", "jpeg", "png", 'bmp', 'webp', 'mp4'))
-
-    # Model Options
-    confidence = float(st.slider(
-        "Select Model Confidence", 25, 100, 40)) / 100
-
-# Creating main page heading
+# ----------------------------------------------------------------
+# App Title and Description
 st.title("WildfireWatch: Detecting Wildfire using AI")
+st.markdown(
+    """
+    **Wildfires are a critical threat to ecosystems and communities.**  
+    Early detection can save lives and reduce environmental damage.  
+    Use this app to analyze images, videos, or live webcam streams for signs of fire.
+    """
+)
 
-# 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)
+# ----------------------------------------------------------------
+# Create two tabs: one for file uploads and one for live webcam stream detection
+tab_upload, tab_live = st.tabs(["Upload Image/Video", "Live Webcam Stream"])
 
-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, WildfireWatch, uses state-of-the-art YOLOv8 model for real-time wildfire and smoke detection in images and videos.
-""")
+# =========================
+# Tab 1: File Upload for Image/Video Detection
+with tab_upload:
+    st.header("Upload an Image or Video")
+    uploaded_file = st.file_uploader(
+        "Choose an image or video...",
+        type=["jpg", "jpeg", "png", "bmp", "webp", "mp4"]
+    )
+    confidence = st.slider("Select Model Confidence", 25, 100, 40) / 100
 
-st.markdown("---")  # Adding a horizontal line
+    if uploaded_file is not None:
+        if uploaded_file.type.split('/')[0] == 'image':
+            # Process uploaded image
+            image = PIL.Image.open(uploaded_file)
+            st.image(image, caption="Uploaded Image", use_column_width=True)
+            if st.button("Detect Wildfire in Image"):
+                results = model.predict(image, conf=confidence)
+                annotated_image = results[0].plot()[:, :, ::-1]
+                st.image(annotated_image, caption="Detection Result", use_column_width=True)
+                with st.expander("Detection Details"):
+                    for box in results[0].boxes:
+                        st.write("Box coordinates (xywh):", box.xywh)
+        elif uploaded_file.type.split('/')[0] == 'video':
+            # Process uploaded video
+            tfile = tempfile.NamedTemporaryFile(delete=False)
+            tfile.write(uploaded_file.read())
+            cap = cv2.VideoCapture(tfile.name)
+            if st.button("Detect Wildfire in Video"):
+                frame_placeholder = st.empty()
+                while cap.isOpened():
+                    ret, frame = cap.read()
+                    if not ret:
+                        break
+                    results = model.predict(frame, conf=confidence)
+                    annotated_frame = results[0].plot()[:, :, ::-1]
+                    frame_placeholder.image(annotated_frame, channels="BGR", use_column_width=True)
+                    time.sleep(0.05)  # Adjust delay for processing speed
+                cap.release()
 
-st.header("Let's Detect Wildfire")
+# =========================
+# Tab 2: Live Webcam Stream Detection
+with tab_live:
+    st.header("Live Webcam Stream Detection")
+    st.markdown("Enter the URL for your online hosted webcam stream (e.g., an IP camera stream).")
+    webcam_url = st.text_input("Webcam Stream URL", value="http://<your_webcam_stream_url>")
+    live_confidence = st.slider("Select Live Detection Confidence", 25, 100, 40) / 100
 
-# Creating two columns on the main page
-col1, col2 = st.columns(2)
+    # Initialize a session state flag for stopping the live loop
+    if "stop_live" not in st.session_state:
+        st.session_state.stop_live = False
 
-# Adding image to the first column if image is uploaded
-with col1:
-    if source_file:
-        # Check if the file is an image
-        if source_file.type.split('/')[0] == 'image':
-            # Opening the uploaded image
-            uploaded_image = PIL.Image.open(source_file)
-            # Adding the uploaded image to the page with a caption
-            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)
+    def stop_live_detection():
+        st.session_state.stop_live = True
 
-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)
-
-if st.sidebar.button('Let\'s Detect Wildfire'):
-    if 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:
-        # Open the video file
-        success, image = vidcap.read()
-        while success:
-            res = model.predict(image,
-                                conf=confidence
-                                )
-            boxes = res[0].boxes
-            res_plotted = res[0].plot()[:, :, ::-1]
-            with col2:
-                st.image(res_plotted,
-                         caption='Detected Frame',
-                         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 video is uploaded yet!")
-            success, image = vidcap.read()
+    if st.button("Start Live Detection"):
+        cap = cv2.VideoCapture(webcam_url)
+        if not cap.isOpened():
+            st.error("Unable to open webcam stream. Please check the URL.")
+        else:
+            live_frame_placeholder = st.empty()
+            st.button("Stop Live Detection", on_click=stop_live_detection)
+            while cap.isOpened() and not st.session_state.stop_live:
+                ret, frame = cap.read()
+                if not ret:
+                    st.error("Failed to retrieve frame from stream.")
+                    break
+                results = model.predict(frame, conf=live_confidence)
+                annotated_frame = results[0].plot()[:, :, ::-1]
+                live_frame_placeholder.image(annotated_frame, channels="BGR", use_column_width=True)
+                time.sleep(0.05)
+            cap.release()
+            st.session_state.stop_live = False