Spaces:
Sleeping
Sleeping
| 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="AI Fire Watch", page_icon="🌍", layout="wide") | |
| # Lighter Background CSS with Darker Text | |
| st.markdown( | |
| """ | |
| <style> | |
| .stApp { | |
| background-color: #cdc0b0; | |
| color: #1a1a1a; | |
| } | |
| h1 { | |
| color: #1a1a1a; | |
| } | |
| .stTabs > div > button { | |
| background-color: #e0e0e0; | |
| color: #333333; /* Darker tab text */ | |
| font-weight: bold; | |
| } | |
| .stTabs > div > button:hover { | |
| background-color: #d0d0d0; | |
| color: #333333; | |
| } | |
| .stTabs > div > button[aria-selected="true"] { | |
| 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 */ | |
| object-fit: contain; | |
| } | |
| </style> | |
| """, | |
| unsafe_allow_html=True | |
| ) | |
| # 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"Model loading failed: {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 | |
| 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.") | |
| # Tabs | |
| tabs = st.tabs(["Upload", "Webcam", "YouTube"]) | |
| # Tab 1: Upload | |
| with tabs[0]: | |
| col1, col2 = st.columns([1, 1]) | |
| with col1: | |
| st.markdown("**Add Your File**") | |
| st.write("Upload an image or video to scan for fire or smoke.") | |
| uploaded_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: | |
| 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_image = results[0].plot()[:, :, ::-1] | |
| frame_placeholder.image(detected_image, use_column_width=True) | |
| status_placeholder.write(f"Objects detected: {len(results[0].boxes)}") | |
| elif file_type == 'video': | |
| # Save uploaded file to temporary location | |
| tfile = tempfile.NamedTemporaryFile(delete=False, suffix='.mp4') | |
| tfile.write(uploaded_file.read()) | |
| tfile.close() | |
| # Open video with OpenCV | |
| cap = cv2.VideoCapture(tfile.name) | |
| if not cap.isOpened(): | |
| status_placeholder.error("Failed to open video file.") | |
| 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) # Control playback speed | |
| cap.release() | |
| # Clean up temporary file | |
| import os | |
| os.unlink(tfile.name) | |
| except Exception as e: | |
| status_placeholder.error(f"Error processing file: {str(e)}") | |
| # Tab 2: Webcam | |
| 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") | |
| 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 | |
| with col2: | |
| frame_placeholder = st.empty() | |
| status_placeholder = st.empty() | |
| 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: | |
| 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] | |
| frame_placeholder.image(detected_frame, use_column_width=True) | |
| status_placeholder.write(f"Objects detected: {len(results[0].boxes)}") | |
| time.sleep(0.1) # Fast update for video | |
| except Exception as e: | |
| status_placeholder.error(f"Video error: {e}") | |
| st.session_state.monitoring = False | |
| break | |
| cap.release() | |
| else: | |
| # Image-based webcam | |
| status_placeholder.write("Monitoring image-based webcam...") | |
| while st.session_state.monitoring: | |
| try: | |
| start_time = time.time() | |
| 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 | |
| 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)}") | |
| # 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: | |
| 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 |