Update app.py

app.py

@@ -224,7 +224,7 @@ class BYTETracker:
 
 # ========================== # Optimized Configuration # ==========================
 CONFIG = {
-    "MODEL_NAME": "facebook/detr-resnet-50",  # Fine-tune with your dataset, e.g., "your-username/detr-resnet-50-finetuned-safety"
+    "MODEL_NAME": "facebook/detr-resnet-50",
     "VIOLATION_LABELS": {
         "no_helmet": "No Helmet",
         "no_harness": "No Harness",
@@ -263,25 +263,26 @@ CONFIG = {
     "MIN_VIOLATION_FRAMES": 2,
     "VIOLATION_COOLDOWN": 30.0,
     "WORKER_TRACKING_DURATION": 10.0,
-    "MAX_PROCESSING_TIME": 60,
-    "FRAME_SKIP": 2,
-    "BATCH_SIZE": 8,
+    "MAX_PROCESSING_TIME": 60,  # Reduced for early termination
+    "FRAME_SKIP": 4,  # Increased to reduce frames processed
+    "BATCH_SIZE": 4,  # Reduced for CPU efficiency
     "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": (384, 384),
+    "TARGET_RESOLUTION": (320, 320),  # Reduced for faster inference
     "HELMET_VALIDATION_FRAMES": 3
 }
 
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 logger.info(f"Using device: {device}")
+if device.type == "cpu":
+    logger.warning("Running on CPU, which may lead to slower processing. Consider using a GPU for better performance.")
 
 def load_model():
     try:
-        # Check for timm dependency
         import timm
         logger.info("timm library is available.")
     except ImportError as e:
@@ -312,7 +313,6 @@ def preprocess_frame(frame):
     return frame
 
 def is_unsafe_posture(box, frame_shape):
-    """Placeholder for unsafe posture detection. Replace with pose estimation (e.g., MediaPipe)."""
     x1, y1, x2, y2 = box
     height = y2 - y1
     width = x2 - x1
@@ -320,14 +320,12 @@ def is_unsafe_posture(box, frame_shape):
     return aspect_ratio > 2.0
 
 def is_improper_tool_use(person_box, tool_box):
-    """Placeholder for improper tool use. Fine-tune DETR for specific tools."""
     person_center = ((person_box[0] + person_box[2]) / 2, (person_box[1] + person_box[3]) / 2)
     tool_center = ((tool_box[0] + tool_box[2]) / 2, (tool_box[1] + tool_box[3]) / 2)
     dist = distance.euclidean(person_center, tool_center)
     return dist > 100
 
 def is_unsafe_zone(person_box, frame_shape):
-    """Check if person is in restricted area (e.g., top-left quadrant)."""
     px, py, pw, ph = person_box
     frame_h, frame_w = frame_shape
     person_center = (px + pw / 2, py + ph / 2)
@@ -508,7 +506,7 @@ def push_report_to_salesforce(violations, score, pdf_path, pdf_file):
                     sf.Safety_Video_Report__c.update(record_id, {"PDF_Report_URL": uploaded_url})
                     logger.info(f"Updated record {record_id} with PDF URL: {uploaded_url}")
                 except Exception as e:
-                    logger.error(f"Failed to update Safety_Video_Report__c: {e}")
+                    logger.error(f"Failed to update Safety_Violation_Report__c: {e}")
                     sf.Account.update(record_id, {"Description": uploaded_url})
                     logger.info(f"Updated account record {record_id} with PDF URL")
                 pdf_url = uploaded_url
@@ -617,6 +615,7 @@ def process_video(video_data, temp_dir):
         unique_violations = {}
         violation_frames = {}
         helmet_detections = {}
+        frame_detections = {}  # Store detections for snapshot reuse
         start_time = time.time()
         frame_skip = CONFIG["FRAME_SKIP"]
         processed_frames = 0
@@ -634,7 +633,7 @@ def process_video(video_data, temp_dir):
                     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}. Skipping.")
                     break
                 original_frame = frame.copy()
                 frame = preprocess_frame(frame)
@@ -644,12 +643,18 @@ def process_video(video_data, temp_dir):
                 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 + 1
+                processed_frames += frame_skip
 
             if not batch_frames:
                 logger.info("No more frames to process.")
                 break
 
+            # Check for timeout
+            elapsed_time = time.time() - start_time
+            if elapsed_time > CONFIG["MAX_PROCESSING_TIME"]:
+                logger.warning(f"Processing exceeded time limit of {CONFIG['MAX_PROCESSING_TIME']}s. Terminating early.")
+                break
+
             try:
                 inputs = processor(images=batch_frames, return_tensors="pt").to(device)
                 if device.type == "cuda":
@@ -671,7 +676,7 @@ def process_video(video_data, temp_dir):
                 progress = (processed_frames / total_frames) * 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)", "", "", "", "", ""
+                yield f"Processing video... {progress:.1f}% complete (Frame {processed_frames}/{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)):
@@ -680,6 +685,8 @@ def process_video(video_data, temp_dir):
                 person_boxes = []
                 tool_boxes = []
 
+                frame_detections[frame_idx] = []  # Store detections for this frame
+
                 for score, label, box in zip(result["scores"], result["labels"], result["boxes"]):
                     label_name = model.config.id2label[label.item()]
                     conf = float(score)
@@ -689,24 +696,27 @@ def process_video(video_data, temp_dir):
                     bbox_xywh = [x + w/2, y + h/2, w, h]
 
                     if label_name in ["no_helmet", "no_harness"] and conf >= CONFIG["CONFIDENCE_THRESHOLDS"].get(label_name, 0.25):
-                        if label_name == "no_helmet" and not validate_helmet(original_frame, bbox_xywh, conf):
-                            logger.info(f"Frame {frame_idx}: Height false positive violation filtered out at {conf:.2f} confidence")
+                        if label_name == "no_helmet" and not validate_helmet_detection(original_frame, bbox_xywh, conf):
+                            logger.info(f"Frame {frame_idx}: Helmet false positive filtered at {conf:.2f} confidence")
                             continue
                         track_inputs.append({"bbox": bbox_xywh, "conf": conf, "cls": label_name})
+                        frame_detections[frame_idx].append({"label": label_name, "conf": conf, "bbox": bbox_xywh})
                     elif label_name == "person":
                         person_boxes.append(bbox_xywh)
-                    elif label_name in ["hammer", "wrench"]:  # Example tools; update with your dataset
+                    elif label_name in ["hammer", "wrench"]:
                         tool_boxes.append(bbox_xywh)
 
-                # Handle Unsafe violations
                 for pbox in person_boxes:
                     if is_unsafe_posture(pbox, original_frame.shape[:2]):
                         track_inputs.append({"bbox": pbox, "conf": 0.9, "cls": "unsafe_posture"})
+                        frame_detections[frame_idx].append({"label": "unsafe_posture", "conf": 0.9, "bbox": pbox})
                     if is_unsafe_zone(pbox, original_frame.shape[:2]):
                         track_inputs.append({"bbox": pbox, "conf": 0.9, "cls": "unsafe_zone"})
+                        frame_detections[frame_idx].append({"label": "unsafe_zone", "conf": 0.9, "bbox": pbox})
                     for tbox in tool_boxes:
                         if is_improper_tool_use(pbox, tbox):
                             track_inputs.append({"bbox": pbox, "conf": 0.9, "cls": "improper_tool_use"})
+                            frame_detections[frame_idx].append({"label": "improper_tool_use", "conf": 0.9, "bbox": pbox})
 
                 if not track_inputs:
                     continue
@@ -722,6 +732,7 @@ def process_video(video_data, temp_dir):
                     tracker_id = obj['id']
                     label = obj['cls']
                     conf = obj['score']
+                    bbox = obj['bbox']
                     
                     if label not in CONFIG["VIOLATION_LABELS"]:
                         continue
@@ -761,100 +772,104 @@ def process_video(video_data, temp_dir):
 
         violations = []
         for (worker_id, label), detection_time in unique_violations.items():
+            frame_idx = violation_frames[(worker_id, label)]
+            conf = next((d["conf"] for d in frame_detections.get(frame_idx, []) if d["label"] == label), 0.0)
             violations.append({
                 "worker_id": worker_id,
                 "violation": label,
                 "timestamp": detection_time,
-                "confidence": 0.0,
+                "confidence": conf,
                 "frame_idx": violation_frames[(worker_id, label)]
             })
 
         if not violations:
             logger.info("No violations detected after processing")
-            yield "No violations detected in the video.", "Safety Score: 100%", "No snapshots captured.", "N/A", "N/A", ""
+            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))):
+            display_name = CONFIG["DISPLAY_NAMES"].get(v.get("violation", "Unknown"), "Unknown")
+            worker_id = v.get("worker_id", "Unknown")
+            timestamp = v.get("timestamp", 0.0)
+            confidence = v.get("confidence", 0.0)
+            violation_table += f"| {display_name} | {worker_id} | {timestamp:.2f} | {confidence:.2f} |\n"
+        yield violation_table, "", "", "", "", f"Violations detected in {processing_time:.1f}s"
+
         snapshots = []
         cap = cv2.VideoCapture(video_path)
         for violation in violations:
-            frame_idx = violation["frame_idx"]
-            cap.set(cv2.CAP_PROP_POS_FRAMES, frame_idx)
-            ret, frame = cap.read()
-            if not ret:
-                logger.warning(f"Failed to read frame {frame_idx} for snapshot.")
-                continue
+            try:
+                frame_idx = violation["frame_idx"]
+                cap.set(cv2.CAP_PROP_POS_FRAMES, frame_idx)
+                ret, frame = cap.read()
+                if not ret:
+                    logger.warning(f"Failed to read frame {frame_idx} for snapshot.")
+                    continue
 
-            frame = preprocess_frame(frame)
-            frame_pil = Image.fromarray(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))
-            inputs = processor(images=frame_pil, return_tensors="pt").to(device)
-            if device.type == "cuda":
-                inputs = {k: v.half() for k, v in inputs.items()}
-            with torch.no_grad():
-                outputs = model(**inputs)
-            target_sizes = torch.tensor([frame_pil.size[::-1]]).to(device)
-            result = processor.post_process_object_detection(outputs, target_sizes=target_sizes, threshold=0.1)[0]
-
-            for score, label, box in zip(result["scores"], result["labels"], result["boxes"]):
-                label_name = model.config.id2label[label.item()]
-                conf = float(score)
-                bbox = box.cpu().numpy()
-                x, y, x2, y2 = bbox
-                w, h = x2 - x, y2 - y
-                bbox_xywh = [x + w/2, y + h/2, w, h]
-                if label_name == violation["violation"]:
-                    violation["confidence"] = round(conf, 2)
-                    detection = {
-                        "worker_id": violation["worker_id"],
-                        "violation": label_name,
-                        "confidence": violation["confidence"],
-                        "bounding_box": bbox_xywh,
-                        "timestamp": violation["timestamp"]
-                    }
-                    snapshot_frame = frame.copy()
-                    snapshot_frame = draw_detections(snapshot_frame, [detection])
-                    cv2.putText(
-                        snapshot_frame,
-                        f"Time: {violation['timestamp']:.2f}s",
-                        (10, 30),
-                        cv2.FONT_HERSHEY_SIMPLEX,
-                        0.7,
-                        (255, 255, 255),
-                        2
-                    )
-                    snapshot_filename = f"violation_{label_name}_worker{violation['worker_id']}_{int(violation['timestamp']*100)}.jpg"
-                    snapshot_path = os.path.join(output_dir, snapshot_filename)
-                    cv2.imwrite(
-                        snapshot_path,
-                        snapshot_frame,
-                        [cv2.IMWRITE_JPEG_QUALITY, CONFIG["SNAPSHOT_QUALITY"]]
-                    )
-                    snapshots.append({
-                        "violation": label_name,
-                        "worker_id": violation["worker_id"],
-                        "timestamp": violation["timestamp"],
-                        "snapshot_path": snapshot_path,
-                        "snapshot_url": f"{CONFIG['PUBLIC_URL_BASE']}{snapshot_filename}",
-                        "confidence": violation["confidence"]
-                    })
-                    logger.info(f"Captured snapshot for {label_name} violation by worker {violation['worker_id']} at {violation['timestamp']:.2f}s")
-                    break
+                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"]:
+                        violation["confidence"] = round(det["conf"], 2)
+                        detection = {
+                            "worker_id": violation["worker_id"],
+                            "violation": det["label"],
+                            "confidence": violation["confidence"],
+                            "bounding_box": det["bbox"],
+                            "timestamp": violation["timestamp"]
+                        }
+                        snapshot_frame = frame.copy()
+                        snapshot_frame = draw_detections(snapshot_frame, [detection])
+                        cv2.putText(
+                            snapshot_frame,
+                            f"Time: {violation['timestamp']:.2f}s",
+                            (10, 30),
+                            cv2.FONT_HERSHEY_SIMPLEX,
+                            0.7,
+                            (255, 255, 255),
+                            2
+                        )
+                        snapshot_filename = f"violation_{det['label']}_worker{violation['worker_id']}_{int(violation['timestamp']*100)}.jpg"
+                        snapshot_path = os.path.join(output_dir, snapshot_filename)
+                        cv2.imwrite(
+                            snapshot_path,
+                            snapshot_frame,
+                            [cv2.IMWRITE_JPEG_QUALITY, CONFIG["SNAPSHOT_QUALITY"]]
+                        )
+                        snapshots.append({
+                            "violation": det["label"],
+                            "worker_id": violation["worker_id"],
+                            "timestamp": violation["timestamp"],
+                            "snapshot_path": snapshot_path,
+                            "snapshot_url": f"{CONFIG['PUBLIC_URL_BASE']}{snapshot_filename}",
+                            "confidence": violation["confidence"]
+                        })
+                        logger.info(f"Captured snapshot for {det['label']} violation by worker {violation['worker_id']} at {violation['timestamp']:.2f}s")
+                        break
+            except Exception as e:
+                logger.error(f"Error generating snapshot for violation: {e}")
+                continue
 
         cap.release()
 
         score = calculate_safety_score(violations)
-        pdf_path, pdf_url, pdf_file = generate_violation_pdf(violations, score, output_dir)
-        
-        record_id, final_pdf_url = push_report_to_salesforce(violations, score, pdf_path, pdf_file)
+        pdf_path, pdf_url, pdf_file = "", "", None
+        try:
+            pdf_path, pdf_url, pdf_file = generate_violation_pdf(violations, score, output_dir)
+        except Exception as e:
+            logger.error(f"PDF generation failed: {e}")
+            yield violation_table, f"Safety Score: {score}%", "Failed to generate snapshots due to PDF error.", "N/A", "N/A", f"Completed in {processing_time:.1f}s\nError: {str(e)}"
+            return
 
-        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))):
-            display_name = CONFIG["DISPLAY_NAMES"].get(v.get("violation", "Unknown"), "Unknown")
-            worker_id = v.get("worker_id", "Unknown")
-            timestamp = v.get("timestamp", 0.0)
-            confidence = v.get("confidence", 0.0)
-            violation_table += f"| {display_name} | {worker_id} | {timestamp:.2f} | {confidence:.2f} |\n"
+        record_id, final_pdf_url = "N/A", "Salesforce integration failed."
+        try:
+            record_id, final_pdf_url = push_report_to_salesforce(violations, score, pdf_path, pdf_file)
+        except Exception as e:
+            logger.error(f"Salesforce integration failed: {e}")
 
         snapshots_text = ""
         for s in snapshots:
@@ -873,12 +888,12 @@ def process_video(video_data, temp_dir):
             snapshots_text,
             f"Salesforce Record ID: {record_id}",
             final_pdf_url,
-            ""
+            f"Completed in {processing_time:.1f}s"
         )
 
     except Exception as e:
         logger.error(f"Error processing video: {str(e)}", exc_info=True)
-        yield f"Error processing video: {str(e)}", "", "", "", "", ""
+        yield f"Error processing video: {str(e)}", "", "", "", "", f"Failed after {time.time() - start_time:.1f}s"
     finally:
         if video_path and os.path.exists(video_path):
             try:
@@ -915,12 +930,12 @@ def gradio_interface(video_file):
         if not FFMPEG_AVAILABLE:
             return "FFmpeg is not available in the environment. Please install FFmpeg to process videos.", "", "", "", "", ""
 
-        for status, score, snapshots_text, record_id, details_url, _ in process_video(video_data, temp_dir):
-            yield status, score, snapshots_text, record_id, details_url, ""
+        for status, score, snapshots_text, record_id, details_url, log in process_video(video_data, temp_dir):
+            yield status, score, snapshots_text, record_id, details_url, log
             
     except Exception as e:
         logger.error(f"Error in Gradio interface: {e}", exc_info=True)
-        yield f"Error: {str(e)}", "", "Error in processing.", "", "", ""
+        yield f"Error: {str(e)}", "", "Error in processing.", "", "", str(e)
     finally:
         if local_video_path and os.path.exists(local_video_path):
             try:
@@ -945,7 +960,7 @@ interface = gr.Interface(
         gr.Markdown(label="Snapshots"),
         gr.Textbox(label="Salesforce Record ID"),
         gr.Textbox(label="Violation Details URL"),
-        gr.Textbox(label="Error Log", visible=False)
+        gr.Textbox(label="Processing Log")
     ],
     title="Worksite Safety Violation Analyzer",
     description="Upload site videos to detect safety violations (No Helmet, No Harness, Unsafe Posture, Unsafe Zone, Improper Tool Use). Each unique violation is detected only once per worker.",