Update app.py

app.py

@@ -6,6 +6,7 @@ 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")
@@ -48,7 +49,7 @@ st.markdown(
 )
 
 # Load Model
-model_path = 'https://huggingface.co/spaces/tstone87/ccr-colorado/resolve/main/best.pt'  # Your updated model
+model_path = 'https://huggingface.co/spaces/tstone87/ccr-colorado/resolve/main/best.pt'
 try:
     model = YOLO(model_path)
 except Exception as ex:
@@ -72,7 +73,7 @@ st.markdown("---")
 # Tabs
 tabs = st.tabs(["Upload", "Webcam"])
 
-# Tab 1: Upload (Your original working version)
+# Tab 1: Upload (With preview and progress)
 with tabs[0]:
     col1, col2 = st.columns(2)
     with col1:
@@ -84,6 +85,9 @@ with tabs[0]:
     with col2:
         frame_placeholder = st.empty()
         status_placeholder = st.empty()
+        progress_placeholder = st.empty()  # For progress percentage
+        download_placeholder = st.empty()
+        
         if source_file and st.button("Detect Wildfire", key="upload_detect"):
             file_type = source_file.type.split('/')[0]
             if file_type == 'image':
@@ -93,29 +97,65 @@ with tabs[0]:
                 frame_placeholder.image(detected_image, use_column_width=True)
                 status_placeholder.write(f"Objects detected: {len(res[0].boxes)}")
             elif file_type == 'video':
-                tfile = tempfile.NamedTemporaryFile(delete=False, suffix='.mp4')
-                tfile.write(source_file.read())
-                tfile.close()
-                vidcap = cv2.VideoCapture(tfile.name)
+                # Save input video to temporary file
+                input_tfile = tempfile.NamedTemporaryFile(delete=False, suffix='.mp4')
+                input_tfile.write(source_file.read())
+                input_tfile.close()
+
+                # Open video to get total frame count
+                vidcap = cv2.VideoCapture(input_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))
+                    if total_frames <= 0:
+                        total_frames = None  # Fallback if frame count unavailable
+                    
+                    # Prepare output video
+                    output_tfile = tempfile.NamedTemporaryFile(delete=False, suffix='_detected.mp4')
+                    frame_width = int(vidcap.get(cv2.CAP_PROP_FRAME_WIDTH))
+                    frame_height = int(vidcap.get(cv2.CAP_PROP_FRAME_HEIGHT))
+                    fps = int(vidcap.get(cv2.CAP_PROP_FPS)) or 30
+                    fourcc = cv2.VideoWriter_fourcc(*'mp4v')
+                    out = cv2.VideoWriter(output_tfile.name, fourcc, fps, (frame_width, frame_height))
+                    
                     success, frame = vidcap.read()
                     frame_count = 0
                     while success:
                         res = model.predict(frame, conf=confidence)
                         detected_frame = res[0].plot()[:, :, ::-1]
-                        frame_placeholder.image(detected_frame, use_column_width=True)
+                        frame_placeholder.image(detected_frame, use_column_width=True)  # Preview
                         status_placeholder.write(f"Frame {frame_count}: Objects detected: {len(res[0].boxes)}")
+                        
+                        # Write to output video
+                        detected_frame_bgr = detected_frame[:, :, ::-1]
+                        out.write(detected_frame_bgr)
+                        
+                        # Update progress
+                        if total_frames:
+                            progress_percent = (frame_count + 1) / total_frames * 100
+                            progress_placeholder.write(f"Progress: {progress_percent:.1f}%")
+                        else:
+                            progress_placeholder.write(f"Progress: {frame_count} frames processed (total unknown)")
+                        
                         success, frame = vidcap.read()
                         frame_count += 1
                         time.sleep(0.05)
+                    
                     vidcap.release()
-                    import os
-                    os.unlink(tfile.name)
+                    out.release()
+                    
+                    os.unlink(input_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. Processed {frame_count} frames.")
 
-# Tab 2: Webcam (Enhanced with video and image support)
+# Tab 2: Webcam (Unchanged)
 with tabs[1]:
     col1, col2 = st.columns([1, 1])
     with col1:
@@ -158,7 +198,7 @@ with tabs[1]:
                         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)  # Fast update for video
+                        time.sleep(0.1)
                     except Exception as e:
                         status_placeholder.error(f"Video error: {e}")
                         st.session_state.monitoring = False