Update app.py

app.py

@@ -64,8 +64,10 @@ CONFIG = {
     "HELMET_CONFIDENCE_THRESHOLD": 0.6,
     "WORKER_TRACKING_DURATION": 2.5,
     "MAX_PROCESSING_TIME": 30,
-    "PARALLEL_WORKERS": max(1, cpu_count() - 1),
-    "CHUNK_SIZE": 15  # Increased chunk size for faster processing
+    "PARALLEL_WORKERS": 2,  # Reduced for Hugging Face Spaces
+    "CHUNK_SIZE": 20,  # Increased for faster batch processing
+    "FRAME_SAMPLE_RATE": 2,  # Process every 2nd frame
+    "MAX_FRAME_WIDTH": 640  # Resize frames to this width
 }
 
 # Setup logging
@@ -99,6 +101,15 @@ model = load_model()
 # ==========================
 # Optimized Helper Functions
 # ==========================
+def resize_frame(frame, max_width):
+    height, width = frame.shape[:2]
+    if width > max_width:
+        scale = max_width / width
+        new_width = int(width * scale)
+        new_height = int(height * scale)
+        frame = cv2.resize(frame, (new_width, new_height), interpolation=cv2.INTER_AREA)
+    return frame
+
 def draw_detections(frame, detections):
     for det in detections:
         label = det.get("violation", "Unknown")
@@ -139,7 +150,9 @@ def calculate_iou(box1, box2):
 
 def process_frame_batch(frame_batch, frame_indices, fps):
     batch_results = []
-    results = model(frame_batch, device=device, conf=0.05, iou=CONFIG["IOU_THRESHOLD"], verbose=False)
+    start_inference = time.time()
+    results = model(frame_batch, device=device, conf=0.1, iou=CONFIG["IOU_THRESHOLD"], verbose=False)
+    logger.info(f"Inference time for batch of {len(frame_batch)} frames: {time.time() - start_inference:.2f}s")
     
     for idx, (result, frame_idx) in enumerate(zip(results, frame_indices)):
         current_time = frame_idx / fps
@@ -169,6 +182,7 @@ def process_frame_batch(frame_batch, frame_indices, fps):
 
 def generate_violation_pdf(violations, score):
     try:
+        start_pdf = time.time()
         pdf_filename = f"violations_{int(time.time())}.pdf"
         pdf_path = os.path.join(CONFIG["OUTPUT_DIR"], pdf_filename)
         pdf_file = BytesIO()
@@ -210,6 +224,7 @@ def generate_violation_pdf(violations, score):
         with open(pdf_path, "wb") as f:
             f.write(pdf_file.getvalue())
         public_url = f"{CONFIG['PUBLIC_URL_BASE']}{pdf_filename}"
+        logger.info(f"PDF generation time: {time.time() - start_pdf:.2f}s")
         logger.info(f"PDF generated: {public_url}")
         return pdf_path, public_url, pdf_file
     except Exception as e:
@@ -241,6 +256,7 @@ def process_video(video_data):
         logger.info(f"Video saved: {video_path}")
 
         # Open video file
+        start_read = time.time()
         cap = cv2.VideoCapture(video_path)
         if not cap.isOpened():
             raise ValueError("Could not open video file")
@@ -264,19 +280,32 @@ def process_video(video_data):
             yield "Video duration too long. Please upload a shorter video.", "", "", "", ""
             return
 
-        # Read all frames upfront
+        # Estimate processing feasibility
+        estimated_frames = total_frames // CONFIG["FRAME_SAMPLE_RATE"]
+        if estimated_frames * 0.1 > CONFIG["MAX_PROCESSING_TIME"]:  # Rough estimate: 0.1s per frame
+            logger.warning(f"Too many frames ({estimated_frames}) to process within {CONFIG['MAX_PROCESSING_TIME']}s")
+            cap.release()
+            os.remove(video_path)
+            yield "Video has too many frames to process within 30 seconds.", "", "", "", ""
+            return
+
+        # Read frames with sampling
         frame_batches = []
         frame_indices_batches = []
         current_batch = []
         current_indices = []
         frame_count = 0
+        sampled_frame_count = 0
 
         while True:
             ret, frame = cap.read()
             if not ret:
                 break
-            current_batch.append(frame)
-            current_indices.append(frame_count)
+            if frame_count % CONFIG["FRAME_SAMPLE_RATE"] == 0:
+                frame = resize_frame(frame, CONFIG["MAX_FRAME_WIDTH"])
+                current_batch.append(frame)
+                current_indices.append(frame_count)
+                sampled_frame_count += 1
             frame_count += 1
 
             if len(current_batch) >= CONFIG["CHUNK_SIZE"]:
@@ -290,6 +319,8 @@ def process_video(video_data):
             frame_indices_batches.append(current_indices)
 
         cap.release()
+        logger.info(f"Frame reading time: {time.time() - start_read:.2f}s")
+        logger.info(f"Total frames: {frame_count}, Sampled frames: {sampled_frame_count}")
 
         # Process frames in parallel
         violations = []
@@ -297,6 +328,7 @@ def process_video(video_data):
         snapshots = []
         last_progress_time = start_time
 
+        start_parallel = time.time()
         with Pool(processes=CONFIG["PARALLEL_WORKERS"]) as pool:
             process_func = partial(process_frame_batch, fps=fps)
             results = pool.starmap(process_func, zip(frame_batches, frame_indices_batches))
@@ -307,15 +339,18 @@ def process_video(video_data):
                 all_detections.extend(batch_result)
             all_detections.sort(key=lambda x: x[0])
 
+        logger.info(f"Parallel processing time: {time.time() - start_parallel:.2f}s")
+
         # Worker tracking
+        start_tracking = time.time()
         workers = []
         for frame_idx, detections in all_detections:
             current_time = frame_idx / fps
 
             # Update progress every second
             if time.time() - last_progress_time > 1.0:
-                progress = (frame_idx / total_frames) * 100
-                yield f"Processing video... {progress:.1f}% complete (Frame {frame_idx}/{total_frames})", "", "", "", ""
+                progress = (frame_idx / frame_count) * 100
+                yield f"Processing video... {progress:.1f}% complete (Frame {frame_idx}/{frame_count})", "", "", "", ""
                 last_progress_time = time.time()
 
             # Early termination if time limit approached
@@ -354,7 +389,10 @@ def process_video(video_data):
 
             workers = [w for w in workers if current_time - w["last_seen"] < CONFIG["WORKER_TRACKING_DURATION"]]
 
+        logger.info(f"Worker tracking time: {time.time() - start_tracking:.2f}s")
+
         # Process helmet violations
+        start_snapshot = time.time()
         for worker_id, detections in helmet_violations.items():
             if len(detections) >= CONFIG["MIN_VIOLATION_FRAMES"]:
                 best_detection = max(detections, key=lambda x: x["confidence"])
@@ -378,6 +416,8 @@ def process_video(video_data):
                         })
                     cap.release()
 
+        logger.info(f"Snapshot generation time: {time.time() - start_snapshot:.2f}s")
+
         os.remove(video_path)
         processing_time = time.time() - start_time
         logger.info(f"Processing complete in {processing_time:.2f}s. {len(violations)} violations found.")
@@ -401,6 +441,7 @@ def process_video(video_data):
             for s in snapshots
         ) if snapshots else "No snapshots captured."
 
+        start_salesforce = time.time()
         try:
             sf = connect_to_salesforce()
             record_data = {
@@ -417,6 +458,8 @@ def process_video(video_data):
         except Exception as e:
             logger.error(f"Salesforce integration failed: {e}")
             record_id = "N/A (Salesforce error)"
+        
+        logger.info(f"Salesforce integration time: {time.time() - start_salesforce:.2f}s")
 
         yield (
             violation_table,