Update app.py

app.py

@@ -263,16 +263,16 @@ CONFIG = {
     "MIN_VIOLATION_FRAMES": 2,
     "VIOLATION_COOLDOWN": 30.0,
     "WORKER_TRACKING_DURATION": 10.0,
-    "MAX_PROCESSING_TIME": 60,  # Reduced for early termination
-    "FRAME_SKIP": 4,  # Increased to reduce frames processed
-    "BATCH_SIZE": 4,  # Reduced for CPU efficiency
+    "MAX_PROCESSING_TIME": 120,  # Increased to allow more time for CPU processing
+    "FRAME_SKIP": 4,
+    "BATCH_SIZE": 2,  # Reduced for better CPU performance
     "PARALLEL_WORKERS": max(1, cpu_count() - 1),
     "TRACK_BUFFER": 150,
     "TRACK_THRESH": 0.3,
     "MATCH_THRESH": 0.5,
     "SNAPSHOT_QUALITY": 95,
     "MAX_WORKER_DISTANCE": 150,
-    "TARGET_RESOLUTION": (320, 320),  # Reduced for faster inference
+    "TARGET_RESOLUTION": (256, 256),  # Further reduced for faster inference
     "HELMET_VALIDATION_FRAMES": 3
 }
 
@@ -308,8 +308,7 @@ def preprocess_frame(frame):
     target_res = CONFIG["TARGET_RESOLUTION"]
     frame = cv2.resize(frame, target_res, interpolation=cv2.INTER_LINEAR)
     frame = cv2.convertScaleAbs(frame, alpha=1.3, beta=20)
-    kernel = np.array([[-1,-1,-1], [-1, 9,-1], [-1,-1,-1]])
-    frame = cv2.filter2D(frame, -1, kernel)
+    # Removed cv2.filter2D to reduce processing time
     return frame
 
 def is_unsafe_posture(box, frame_shape):
@@ -615,38 +614,45 @@ def process_video(video_data, temp_dir):
         unique_violations = {}
         violation_frames = {}
         helmet_detections = {}
-        frame_detections = {}  # Store detections for snapshot reuse
+        frame_detections = {}
         start_time = time.time()
         frame_skip = CONFIG["FRAME_SKIP"]
-        processed_frames = 0
+        processed_frames = 0  # Track actual frames processed
+        frames_read = 0  # Track frames read from video
         last_yield_time = start_time
         worker_counter = 1
 
-        while processed_frames < total_frames:
+        while True:
             batch_frames = []
             batch_indices = []
             batch_originals = []
             
             for _ in range(CONFIG["BATCH_SIZE"]):
                 frame_idx = int(cap.get(cv2.CAP_PROP_POS_FRAMES))
+                frames_read = frame_idx
                 if frame_idx >= total_frames:
+                    logger.info("Reached end of video.")
                     break
                 ret, frame = cap.read()
                 if not ret:
-                    logger.warning(f"Failed to read frame {frame_idx}. Skipping.")
+                    logger.warning(f"Failed to read frame {frame_idx}. Assuming end of video.")
+                    frames_read = total_frames  # Assume we've reached the end
                     break
                 original_frame = frame.copy()
                 frame = preprocess_frame(frame)
                 for _ in range(frame_skip - 1):
                     if not cap.grab():
+                        logger.warning(f"Failed to grab frame after {frame_idx}. Assuming end of video.")
+                        frames_read = total_frames
                         break
+                    frames_read += 1
                 batch_frames.append(Image.fromarray(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)))
                 batch_indices.append(frame_idx)
                 batch_originals.append(original_frame)
-                processed_frames += frame_skip
+                processed_frames += 1  # Increment for each frame actually processed
 
             if not batch_frames:
-                logger.info("No more frames to process.")
+                logger.info("No more frames to process in this batch.")
                 break
 
             # Check for timeout
@@ -673,10 +679,11 @@ def process_video(video_data, temp_dir):
 
             current_time = time.time()
             if current_time - last_yield_time > 0.1:
-                progress = (processed_frames / total_frames) * 100
+                progress = (frames_read / total_frames) * 100
+                progress = min(progress, 100.0)  # Ensure progress doesn't exceed 100%
                 elapsed_time = current_time - start_time
                 fps_processed = processed_frames / elapsed_time if elapsed_time > 0 else 0
-                yield f"Processing video... {progress:.1f}% complete (Frame {processed_frames}/{total_frames}, {fps_processed:.1f} FPS)", "", "", "", "", f"Elapsed: {elapsed_time:.1f}s"
+                yield f"Processing video... {progress:.1f}% complete (Frame {frames_read}/{total_frames}, {fps_processed:.1f} FPS)", "", "", "", "", f"Elapsed: {elapsed_time:.1f}s"
                 last_yield_time = current_time
 
             for i, (result, frame_idx, original_frame) in enumerate(zip(results, batch_indices, batch_originals)):
@@ -685,7 +692,7 @@ def process_video(video_data, temp_dir):
                 person_boxes = []
                 tool_boxes = []
 
-                frame_detections[frame_idx] = []  # Store detections for this frame
+                frame_detections[frame_idx] = []
 
                 for score, label, box in zip(result["scores"], result["labels"], result["boxes"]):
                     label_name = model.config.id2label[label.item()]
@@ -770,6 +777,11 @@ def process_video(video_data, temp_dir):
         logger.info(f"Processing complete in {processing_time:.2f}s")
         logger.info(f"Total unique workers detected: {len(set(worker_id_mapping.values()))}")
 
+        # Ensure final progress update
+        final_progress = (frames_read / total_frames) * 100
+        final_progress = min(final_progress, 100.0)
+        yield f"Processing video... {final_progress:.1f}% complete (Frame {frames_read}/{total_frames})", "", "", "", "", f"Elapsed: {processing_time:.1f}s"
+
         violations = []
         for (worker_id, label), detection_time in unique_violations.items():
             frame_idx = violation_frames[(worker_id, label)]
@@ -787,7 +799,6 @@ def process_video(video_data, temp_dir):
             yield "No violations detected in the video.", "Safety Score: 100%", "No snapshots captured.", "N/A", "N/A", f"Completed in {processing_time:.1f}s"
             return
 
-        # Generate violation table early for intermediate output
         violation_table = "| Violation | Worker ID | Time (s) | Confidence |\n"
         violation_table += "|-----------|-----------|----------|------------|\n"
         for v in sorted(violations, key=lambda x: (x.get("worker_id", "Unknown"), x.get("timestamp", 0.0))):
@@ -810,7 +821,6 @@ def process_video(video_data, temp_dir):
                     continue
 
                 frame = preprocess_frame(frame)
-                # Reuse detections instead of re-running inference
                 detections = frame_detections.get(frame_idx, [])
                 for det in detections:
                     if det["label"] == violation["violation"]: