Update app.py

app.py

@@ -6,160 +6,126 @@ import tempfile
 import time
 import requests
 import numpy as np
-import streamlink  # pip install streamlink
+import streamlink
 
-# -------------------------------
-# Page & Model Setup
-st.set_page_config(
-    page_title="WildfireWatch",
-    page_icon="🔥",
-    layout="wide",
-    initial_sidebar_state="expanded"
-)
+# Page Config
+st.set_page_config(page_title="WildfireWatch", page_icon="🔥", layout="wide")
 
-# IMPORTANT: Ensure the model URL returns the raw .pt file. Otherwise, use a local path.
+# 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"Unable to load model from: {model_path}")
-    st.error(ex)
-    st.stop()  # Stop the app if the model cannot load
+    st.error(f"Model loading failed: {ex}")
+    st.stop()
 
-st.title("WildfireWatch: Detecting Fire and Smoke using AI")
-st.markdown("""
-Wildfires and smoke are a major threat. This app uses a YOLOv8 model to detect fire/smoke in:
-- Uploaded images or videos,
-- Webcam snapshots or live streams,
-- YouTube live streams.
-""")
+# Minimalist Header
+st.title("WildfireWatch")
+st.markdown("AI-powered detection of fire and smoke.")
 
-# -------------------------------
-# Create three tabs for different modes
-tabs = st.tabs(["File Upload", "Webcam / Image URL", "YouTube Live Stream"])
+# Tabs
+tabs = st.tabs(["Upload", "Webcam", "YouTube"])
 
-# ===============================
-# Tab 1: File Upload Mode
+# Tab 1: File Upload
 with tabs[0]:
-    st.header("File Upload Mode")
-    col_left, col_right = st.columns(2)
-    with col_left:
-        uploaded_file = st.file_uploader("Choose an image or video...", type=["jpg", "jpeg", "png", "bmp", "webp", "mp4"])
-        file_confidence = st.slider("Select Detection Confidence", 25, 100, 40) / 100
-        display_mode = st.radio("Display Mode", options=["Detection", "Original"], index=0)
-    if uploaded_file:
-        file_type = uploaded_file.type.split('/')[0]
-        if file_type == 'image':
-            image = PIL.Image.open(uploaded_file)
-            col_left.image(image, caption="Uploaded Image", use_column_width=True)
-            results = model.predict(image, conf=file_confidence)
-            detected_image = results[0].plot()[:, :, ::-1]
-            if display_mode == "Detection":
-                col_right.image(detected_image, caption="Detection Result", use_column_width=True)
-            else:
-                col_right.image(image, caption="Original Image", use_column_width=True)
-            with col_right.expander("Detection Details"):
-                for box in results[0].boxes:
-                    st.write("Box (xywh):", box.xywh)
-        elif file_type == 'video':
-            tfile = tempfile.NamedTemporaryFile(delete=False)
-            tfile.write(uploaded_file.read())
-            cap = cv2.VideoCapture(tfile.name)
-            if st.button("Start Video Detection"):
+    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=file_confidence)
+                    results = model.predict(frame, conf=confidence)
                     detected_frame = results[0].plot()[:, :, ::-1]
-                    col_right.image(detected_frame, caption="Detection Frame", channels="BGR", use_column_width=True)
+                    frame_placeholder.image(detected_frame, use_column_width=True)
                     time.sleep(0.05)
                 cap.release()
 
-# ===============================
-# Tab 2: Webcam / Image URL Mode
+# Tab 2: Webcam / Image URL
 with tabs[1]:
-    st.header("Webcam / Image URL Mode")
-    col_left, col_right = st.columns(2)
-    with col_left:
-        # Enter a webcam snapshot URL (e.g. an updating JPEG) or a live stream URL
-        webcam_url = st.text_input("Enter Webcam URL", value="http://<your_webcam_ip>/current.jpg")
-        webcam_confidence = st.slider("Select Detection Confidence", 25, 100, 40, key="webcam_conf") / 100
-        input_mode = st.radio("Input Mode", options=["Snapshot (Updating Image)", "Live Stream"], index=0)
-        start_webcam = st.button("Start Detection", key="start_webcam")
-    if start_webcam:
-        if input_mode == "Snapshot (Updating Image)":
-            placeholder = col_right.empty()
-            # Loop to periodically fetch the snapshot
-            while True:
-                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)
-                    if frame is None:
-                        col_right.error("Failed to decode image from URL.")
+    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
-                    results = model.predict(frame, conf=webcam_confidence)
-                    detected_frame = results[0].plot()[:, :, ::-1]
-                    placeholder.image(detected_frame, channels="BGR", use_column_width=True)
-                    time.sleep(1)  # Adjust refresh rate as needed
-                    st.experimental_rerun()
-                except Exception as e:
-                    col_right.error(f"Error fetching image: {e}")
-                    break
-        else:
-            cap = cv2.VideoCapture(webcam_url)
-            if not cap.isOpened():
-                col_right.error("Unable to open webcam stream. Check the URL and network connectivity.")
             else:
-                stop_button = st.button("Stop Live Detection", key="stop_webcam")
-                while cap.isOpened() and not stop_button:
+                cap = cv2.VideoCapture(webcam_url)
+                frame_placeholder = st.empty()
+                while cap.isOpened():
                     ret, frame = cap.read()
                     if not ret:
-                        col_right.error("Failed to retrieve frame from stream.")
+                        st.error("Stream failed.")
                         break
-                    results = model.predict(frame, conf=webcam_confidence)
+                    results = model.predict(frame, conf=confidence)
                     detected_frame = results[0].plot()[:, :, ::-1]
-                    col_right.image(detected_frame, channels="BGR", use_column_width=True)
+                    frame_placeholder.image(detected_frame, use_column_width=True)
                     time.sleep(0.05)
-                    stop_button = st.button("Stop Live Detection", key="stop_webcam_loop")
-                cap.release()
 
-# ===============================
-# Tab 3: YouTube Live Stream Mode
+# Tab 3: YouTube Live Stream
 with tabs[2]:
-    st.header("YouTube Live Stream Mode")
-    col_left, col_right = st.columns(2)
-    with col_left:
-        youtube_url = st.text_input("Enter YouTube Live URL", value="https://www.youtube.com/watch?v=<live_stream_id>")
-        yt_confidence = st.slider("Select Detection Confidence", 25, 100, 40, key="yt_conf") / 100
-        start_yt = st.button("Start Live Detection", key="start_yt")
-    if start_yt:
-        try:
+    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 not streams:
-                st.error("No streams found. Please check the URL or ensure the stream is live.")
+            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:
-                best_stream = streams.get("best")
-                if best_stream is None:
-                    st.error("No suitable stream found.")
-                else:
-                    stream_url = best_stream.to_url()
-                    cap = cv2.VideoCapture(stream_url)
-                    if not cap.isOpened():
-                        st.error("Unable to open YouTube live stream.")
-                    else:
-                        stop_yt = st.button("Stop Live Detection", key="stop_yt")
-                        while cap.isOpened() and not stop_yt:
-                            ret, frame = cap.read()
-                            if not ret:
-                                col_right.error("Failed to retrieve frame from YouTube stream.")
-                                break
-                            results = model.predict(frame, conf=yt_confidence)
-                            detected_frame = results[0].plot()[:, :, ::-1]
-                            col_right.image(detected_frame, channels="BGR", use_column_width=True)
-                            time.sleep(0.05)
-                            stop_yt = st.button("Stop Live Detection", key="stop_yt_loop")
-                        cap.release()
-        except Exception as e:
-            st.error(f"Error processing YouTube stream: {e}")
+                st.error("No stream found.")