Update app.py

app.py

@@ -73,7 +73,7 @@ st.markdown("---")
 # Tabs
 tabs = st.tabs(["Upload", "Webcam"])
 
-# Tab 1: Upload (Output only analyzed frames)
+# Tab 1: Upload (Simplified to match original, with new features)
 with tabs[0]:
     col1, col2 = st.columns(2)
     with col1:
@@ -95,12 +95,11 @@ with tabs[0]:
     with col2:
         frame_placeholder = st.empty()
         status_placeholder = st.empty()
-        progress_placeholder = st.empty()
         download_placeholder = st.empty()
         
-        if source_file:
-            st.write(f"File size: {source_file.size / 1024 / 1024:.2f} MB")
-            if st.button("Detect Wildfire", key="upload_detect"):
+        if source_file and st.button("Detect Wildfire", key="upload_detect"):
+            try:
+                st.write(f"File size: {source_file.size / 1024 / 1024:.2f} MB")  # Diagnostic
                 file_type = source_file.type.split('/')[0]
                 if file_type == 'image':
                     uploaded_image = PIL.Image.open(source_file)
@@ -109,94 +108,64 @@ 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':
-                    try:
-                        # Save input video
-                        input_tfile = tempfile.NamedTemporaryFile(delete=False, suffix='.mp4')
-                        input_tfile.write(source_file.read())
-                        input_tfile.close()
-
-                        # Open video
-                        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))
-                            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))
-                            
-                            # Calculate frame skip
-                            target_rate = sampling_options[sampling_rate]
-                            if target_rate == 0:
-                                frame_skip = 1
-                            elif target_rate <= 5:
-                                frame_skip = max(1, int(fps / target_rate))
-                            else:
-                                frame_skip = max(1, int(fps * target_rate))
-                            
-                            # Batch processing setup
-                            batch_size = 10
-                            frames_to_process = []
-                            frame_indices = []
-                            
-                            # Output video (only analyzed frames)
-                            output_tfile = tempfile.NamedTemporaryFile(delete=False, suffix='_detected.mp4')
-                            fourcc = cv2.VideoWriter_fourcc(*'mp4v')
-                            # Use 1 FPS for output to make it viewable, adjust as needed
-                            output_fps = 1
-                            out = cv2.VideoWriter(output_tfile.name, fourcc, output_fps, (frame_width, frame_height))
-                            
+                    # Save input video
+                    tfile = tempfile.NamedTemporaryFile(delete=False, suffix='.mp4')
+                    tfile.write(source_file.read())
+                    tfile.close()
+                    
+                    # Open video
+                    vidcap = cv2.VideoCapture(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))
+                        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))
+                        
+                        # Frame sampling
+                        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 (only analyzed frames)
+                        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, frame = vidcap.read()
+                        frame_count = 0
+                        processed_count = 0
+                        
+                        while success:
+                            if frame_count % frame_skip == 0:
+                                res = model.predict(frame, conf=confidence)
+                                detected_frame = res[0].plot()[:, :, ::-1]
+                                frame_placeholder.image(detected_frame, use_column_width=True)
+                                status_placeholder.write(f"Frame {frame_count}: Objects detected: {len(res[0].boxes)}")
+                                out.write(detected_frame[:, :, ::-1])
+                                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, frame = vidcap.read()
-                            frame_count = 0
-                            processed_count = 0
-                            last_ui_update = time.time()
-                            
-                            while success:
-                                if frame_count % frame_skip == 0:
-                                    frames_to_process.append(frame)
-                                    frame_indices.append(frame_count)
-                                    processed_count += 1
-                                
-                                # Process batch when full or at end
-                                if len(frames_to_process) >= batch_size or (not success and frames_to_process):
-                                    res = model.predict(frames_to_process, conf=confidence)
-                                    for i, (result, idx) in enumerate(zip(res, frame_indices)):
-                                        detected_frame = result.plot()[:, :, ::-1]
-                                        # Update UI sparingly
-                                        if time.time() - last_ui_update >= 1.0:
-                                            frame_placeholder.image(detected_frame, use_column_width=True)
-                                            status_placeholder.write(f"Frame {idx}: Objects detected: {len(result.boxes)}")
-                                            last_ui_update = time.time()
-                                        # Write only analyzed frames
-                                        out.write(detected_frame[:, :, ::-1])
-                                
-                                # Progress
-                                if total_frames > 0:
-                                    progress_percent = (frame_count + 1) / total_frames * 100
-                                    progress_placeholder.write(f"Progress: {progress_percent:.1f}% (Analyzed {processed_count} frames)")
-                                else:
-                                    progress_placeholder.write(f"Progress: {frame_count} frames processed")
-                                
-                                success, frame = vidcap.read()
-                                frame_count += 1
-                                if len(frames_to_process) >= batch_size:
-                                    frames_to_process = []
-                                    frame_indices = []
-                            
-                            vidcap.release()
-                            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. Analyzed {processed_count} frames.")
-                    except Exception as e:
-                        status_placeholder.error(f"Error processing video: {str(e)}")
+                            frame_count += 1
+                            time.sleep(0.05)
+                        
+                        vidcap.release()
+                        out.release()
+                        
+                        os.unlink(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. Analyzed {processed_count} frames.")
+            except Exception as e:
+                status_placeholder.error(f"Error: {str(e)}")
 
 # Tab 2: Webcam (Unchanged)
 with tabs[1]: