Update app.py

app.py

@@ -1,145 +1,143 @@
-# Import required libraries
-import PIL
+import os
+import tempfile
 import cv2
 import streamlit as st
+import PIL
 from ultralytics import YOLO
-import tempfile
-import time
-import os
 
-# Replace the relative path to your weight file
-model_path = 'https://huggingface.co/spaces/tstone87/ccr-colorado/resolve/main/best.pt'  # Your correct model
+# Required libraries (ensure these are in your requirements.txt):
+# streamlit
+# opencv-python-headless
+# ultralytics
+# Pillow
+
+# Replace with your model's URL or local path
+model_path = 'https://huggingface.co/spaces/tstone87/ccr-colorado/blob/main/best.pt'
 
-# Setting page layout
+# Configure the page for Hugging Face Spaces
 st.set_page_config(
-    page_title="WildfireWatch",
+    page_title="Fire Watch using AI vision models",
     page_icon="🔥",
     layout="wide",
     initial_sidebar_state="expanded"
 )
 
-# Creating sidebar
+# Sidebar for file upload and settings
 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'))
+        "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)
+    video_option = st.selectbox(
+        "Select Video Shortening Option", 
+        ["Original FPS", "1 fps", "1 frame per 5 seconds", "1 frame per 10 seconds", "1 frame per 15 seconds"]
+    )
 
-# Creating main page heading
+# Main page header and introduction images
 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)
+    st.image("https://huggingface.co/spaces/tstone87/ccr-colorado/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.image("https://huggingface.co/spaces/tstone87/ccr-colorado/resolve/main/Fire_3.png", 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, WildfireWatch, uses state-of-the-art YOLOv8 model for real-time wildfire and smoke detection in images and videos.
+Fires in Colorado present a serious challenge, threatening urban communities, highways, and even remote areas. Early detection is critical to mitigating risks. WildfireWatch leverages the YOLOv8 model for real-time fire and smoke detection in images and videos.
 """)
-
 st.markdown("---")
+st.header("Fire Detection:")
 
-st.header("Let's Detect Wildfire")
-
-# Creating two columns on the main page
 col1, col2 = st.columns(2)
+if source_file:
+    if source_file.type.split('/')[0] == 'image':
+        uploaded_image = PIL.Image.open(source_file)
+        st.image(uploaded_image, caption="Uploaded Image", use_column_width=True)
+    else:
+        tfile = tempfile.NamedTemporaryFile(delete=False)
+        tfile.write(source_file.read())
+        vidcap = cv2.VideoCapture(tfile.name)
+else:
+    st.info("Please upload an image or video file to begin.")
 
-# 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)
-
+# Load the YOLO model
 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()
 
-if st.sidebar.button('Let\'s Detect Wildfire'):
+if st.sidebar.button("Let's Detect Wildfire"):
     if not source_file:
-        st.error("Please upload a file first!")
+        st.warning("No file uploaded!")
     elif source_file.type.split('/')[0] == 'image':
+        # Process image input
         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!")
+            with st.expander("Detection Results"):
+                for box in boxes:
+                    st.write(box.xywh)
     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))
-        
-        # 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()
+        # Process video input and shorten video based on sampling option
+        processed_frames = []
         frame_count = 0
-        processed_count = 0
+
+        # Video properties
+        orig_fps = vidcap.get(cv2.CAP_PROP_FPS)
+        width  = int(vidcap.get(cv2.CAP_PROP_FRAME_WIDTH))
+        height = int(vidcap.get(cv2.CAP_PROP_FRAME_HEIGHT))
         
+        # Determine sampling interval and output fps
+        if video_option == "Original FPS":
+            sample_interval = 1  # process every frame
+            output_fps = orig_fps
+        elif video_option == "1 fps":
+            sample_interval = int(orig_fps) if orig_fps > 0 else 1
+            output_fps = 1
+        elif video_option == "1 frame per 5 seconds":
+            sample_interval = int(orig_fps * 5) if orig_fps > 0 else 5
+            output_fps = 1
+        elif video_option == "1 frame per 10 seconds":
+            sample_interval = int(orig_fps * 10) if orig_fps > 0 else 10
+            output_fps = 1
+        elif video_option == "1 frame per 15 seconds":
+            sample_interval = int(orig_fps * 15) if orig_fps > 0 else 15
+            output_fps = 1
+        else:
+            sample_interval = 1
+            output_fps = orig_fps
+
+        success, image = vidcap.read()
         while success:
-            if frame_count % frame_skip == 0:
+            if frame_count % sample_interval == 0:
                 res = model.predict(image, conf=confidence)
-                boxes = res[0].boxes
                 res_plotted = res[0].plot()[:, :, ::-1]
+                processed_frames.append(res_plotted)
                 with col2:
                     st.image(res_plotted, caption=f'Detected Frame {frame_count}', 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 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()
+                    with st.expander("Detection Results"):
+                        for box in res[0].boxes:
+                            st.write(box.xywh)
             frame_count += 1
-            time.sleep(0.05)
-        
-        vidcap.release()
-        out.release()
-        os.unlink(tfile.name)
+            success, image = vidcap.read()
         
-        with col2:
-            with open(output_tfile.name, 'rb') as f:
+        if processed_frames:
+            temp_video_file = tempfile.NamedTemporaryFile(delete=False, suffix='.mp4')
+            fourcc = cv2.VideoWriter_fourcc(*'mp4v')
+            out = cv2.VideoWriter(temp_video_file.name, fourcc, output_fps, (width, height))
+            for frame in processed_frames:
+                out.write(frame)
+            out.release()
+            
+            st.success("Shortened video created successfully!")
+            with open(temp_video_file.name, 'rb') as video_file:
                 st.download_button(
-                    label="Download Analyzed Video",
-                    data=f,
-                    file_name="analyzed_video.mp4",
+                    label="Download Shortened Video",
+                    data=video_file.read(),
+                    file_name="shortened_video.mp4",
                     mime="video/mp4"
                 )
-            st.write(f"Video processing complete. Analyzed {processed_count} frames.")
+        else:
+            st.error("No frames were processed from the video.")