Update app.py

app.py

@@ -50,7 +50,7 @@ CONFIG = {
         "domain": "login"
     },
     "PUBLIC_URL_BASE": "https://huggingface.co/spaces/PrashanthB461/AI_Safety_Demo2/resolve/main/static/output/",
-    "FRAME_SKIP": 3,  # Increased from 1 to 3 for faster processing
+    "FRAME_SKIP": 2,  # Process every 2nd frame for balance of speed/accuracy
     "CONFIDENCE_THRESHOLDS": {
         "no_helmet": 0.6,
         "no_harness": 0.15,
@@ -59,11 +59,12 @@ CONFIG = {
         "improper_tool_use": 0.15
     },
     "IOU_THRESHOLD": 0.4,
-    "MIN_VIOLATION_FRAMES": 2,
+    "MIN_VIOLATION_FRAMES": 3,  # Require more consistent detections
     "HELMET_CONFIDENCE_THRESHOLD": 0.65,
-    "MAX_PROCESSING_TIME": 30,  # Reduced from 60 to 30 seconds
-    "BATCH_SIZE": 10,  # Process frames in batches for efficiency
-    "WORKER_TRACKING_DURATION": 3.0  # Seconds to track a worker without updates
+    "WORKER_TRACKING_DURATION": 3.0,  # Seconds to track a worker
+    "MIN_FRAME_RATE": 5,  # Minimum frames per second to process
+    "MAX_FRAME_RATE": 15,  # Maximum frames per second to process
+    "BATCH_SIZE": 8  # Number of frames to process at once
 }
 
 # Setup logging
@@ -120,18 +121,21 @@ def calculate_iou(box1, box2):
     x1, y1, w1, h1 = box1
     x2, y2, w2, h2 = box2
     
-    x1_min, y1_min = x1 - w1/2, y1 - h1/2
-    x1_max, y1_max = x1 + w1/2, y1 + h1/2
-    x2_min, y2_min = x2 - w2/2, y2 - h2/2
-    x2_max, y2_max = x2 + w2/2, y2 + h2/2
+    # Calculate coordinates of the intersection rectangle
+    x_left = max(x1 - w1/2, x2 - w2/2)
+    y_top = max(y1 - h1/2, y2 - h2/2)
+    x_right = min(x1 + w1/2, x2 + w2/2)
+    y_bottom = min(y1 + h1/2, y2 + h2/2)
     
-    intersection = max(0, min(x1_max, x2_max) - max(x1_min, x2_min)) * \
-                   max(0, min(y1_max, y2_max) - max(y1_min, y2_min))
-    area1 = w1 * h1
-    area2 = w2 * h2
-    union = area1 + area2 - intersection
+    if x_right < x_left or y_bottom < y_top:
+        return 0.0
     
-    return intersection / union if union > 0 else 0
+    intersection_area = (x_right - x_left) * (y_bottom - y_top)
+    box1_area = w1 * h1
+    box2_area = w2 * h2
+    union_area = box1_area + box2_area - intersection_area
+    
+    return intersection_area / union_area
 
 def generate_violation_pdf(violations, score):
     try:
@@ -182,6 +186,18 @@ def generate_violation_pdf(violations, score):
         logger.error(f"Error generating PDF: {e}")
         return "", "", None
 
+def calculate_safety_score(violations):
+    penalties = {
+        "no_helmet": 25,
+        "no_harness": 30,
+        "unsafe_posture": 20,
+        "unsafe_zone": 35,
+        "improper_tool_use": 25
+    }
+    total_penalty = sum(penalties.get(v.get("violation", "Unknown"), 0) for v in violations)
+    score = 100 - total_penalty
+    return max(score, 0)
+
 # ==========================
 # Optimized Video Processing
 # ==========================
@@ -202,16 +218,19 @@ def process_video(video_data):
         total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
         fps = cap.get(cv2.CAP_PROP_FPS)
         if fps <= 0:
-            fps = 30
+            fps = 30  # Default assumption if FPS not available
         duration = total_frames / fps
-        logger.info(f"Video duration: {duration:.2f}s, Frames: {total_frames}, FPS: {fps}")
+        width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
+        height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
+
+        logger.info(f"Video properties: {duration:.2f}s, {total_frames} frames, {fps:.1f} FPS, {width}x{height}")
 
-        # Calculate frames to process
-        frame_skip = CONFIG["FRAME_SKIP"]
-        frames_to_process = total_frames // frame_skip
-        if frames_to_process < 10:  # Ensure we process at least 10 frames
-            frame_skip = max(1, total_frames // 10)
-            frames_to_process = total_frames // frame_skip
+        # Calculate optimal frame skipping
+        original_frame_skip = CONFIG["FRAME_SKIP"]
+        target_fps = min(max(fps / original_frame_skip, CONFIG["MIN_FRAME_RATE"]), CONFIG["MAX_FRAME_RATE"])
+        actual_frame_skip = max(1, int(fps / target_fps))
+        frames_to_process = total_frames // actual_frame_skip
+        logger.info(f"Processing strategy: Frame skip={actual_frame_skip}, Target FPS={target_fps:.1f}, Frames to process={frames_to_process}")
 
         workers = []
         violations = []
@@ -219,30 +238,35 @@ def process_video(video_data):
         snapshots = []
         start_time = time.time()
         processed_frames = 0
+        last_progress_update = 0
 
-        # Process frames in batches for efficiency
-        for batch_start in range(0, total_frames, CONFIG["BATCH_SIZE"] * frame_skip):
+        # Process frames in batches
+        while True:
             batch_frames = []
             batch_indices = []
             
             # Collect frames for this batch
-            for i in range(CONFIG["BATCH_SIZE"]):
-                frame_idx = batch_start + i * frame_skip
+            for _ in range(CONFIG["BATCH_SIZE"]):
+                frame_idx = int(cap.get(cv2.CAP_PROP_POS_FRAMES))
                 if frame_idx >= total_frames:
                     break
                 
-                cap.set(cv2.CAP_PROP_POS_FRAMES, frame_idx)
                 ret, frame = cap.read()
                 if not ret:
-                    continue
+                    break
                 
                 batch_frames.append(frame)
                 batch_indices.append(frame_idx)
                 processed_frames += 1
+                
+                # Skip frames according to our strategy
+                for _ in range(actual_frame_skip - 1):
+                    if not cap.grab():
+                        break
 
-            # Skip empty batches
+            # Break if no more frames
             if not batch_frames:
-                continue
+                break
 
             # Run batch detection
             results = model(batch_frames, device=device, conf=0.1, iou=CONFIG["IOU_THRESHOLD"], verbose=False)
@@ -250,8 +274,12 @@ def process_video(video_data):
             # Process results for each frame in batch
             for i, (result, frame_idx) in enumerate(zip(results, batch_indices)):
                 current_time = frame_idx / fps
-                progress = (processed_frames / frames_to_process) * 100
-                yield f"Processing video... {progress:.1f}% complete (Frame {frame_idx}/{total_frames})", "", "", "", ""
+                
+                # Update progress periodically
+                if time.time() - last_progress_update > 1.0:  # Update every second
+                    progress = (frame_idx / total_frames) * 100
+                    yield f"Processing video... {progress:.1f}% complete (Frame {frame_idx}/{total_frames})", "", "", "", ""
+                    last_progress_update = time.time()
 
                 # Process detections in this frame
                 boxes = result.boxes
@@ -272,7 +300,7 @@ def process_video(video_data):
                         "timestamp": current_time
                     }
 
-                    # Worker tracking and helmet detection optimization
+                    # Worker tracking
                     worker_id = None
                     max_iou = 0
                     for idx, worker in enumerate(workers):
@@ -292,6 +320,8 @@ def process_video(video_data):
                             "last_seen": current_time
                         })
 
+                    detection["worker_id"] = worker_id
+
                     # Special handling for helmet violations
                     if label == "no_helmet":
                         if worker_id not in helmet_violations:
@@ -303,14 +333,10 @@ def process_video(video_data):
                 # Remove workers not seen recently
                 workers = [w for w in workers if current_time - w["last_seen"] < CONFIG["WORKER_TRACKING_DURATION"]]
 
-            # Check processing time limit
-            if time.time() - start_time > CONFIG["MAX_PROCESSING_TIME"]:
-                logger.info(f"Processing time limit reached at frame {frame_idx}")
-                break
-
-        # Process helmet violations (more strict criteria)
+        # Process helmet violations (require consistent detections)
         for worker_id, detections in helmet_violations.items():
             if len(detections) >= CONFIG["MIN_VIOLATION_FRAMES"]:
+                # Find the detection with highest confidence
                 best_detection = max(detections, key=lambda x: x["confidence"])
                 violations.append(best_detection)
                 
@@ -331,38 +357,54 @@ def process_video(video_data):
 
         cap.release()
         os.remove(video_path)
-        logger.info(f"Processing complete. {len(violations)} violations found.")
+        processing_time = time.time() - start_time
+        logger.info(f"Processing complete in {processing_time:.2f}s. {len(violations)} violations found.")
 
         # Generate results
         if not violations:
             yield "No violations detected in the video.", "Safety Score: 100%", "No snapshots captured.", "N/A", "N/A"
             return
 
-        score = max(0, 100 - sum(25 if v["violation"] == "no_helmet" else 
-                                30 if v["violation"] == "no_harness" else 
-                                20 if v["violation"] == "unsafe_posture" else 
-                                35 if v["violation"] == "unsafe_zone" else 
-                                25 for v in violations))
-
+        score = calculate_safety_score(violations)
         pdf_path, pdf_url, pdf_file = generate_violation_pdf(violations, score)
 
+        # Generate violation table
         violation_table = "| Violation              | Timestamp (s) | Confidence | Worker ID |\n"
         violation_table += "|------------------------|---------------|------------|-----------|\n"
-        for v in violations:
+        for v in sorted(violations, key=lambda x: x["timestamp"]):
             display_name = CONFIG["DISPLAY_NAMES"].get(v.get("violation", "Unknown"), "Unknown")
             row = f"| {display_name:<22} | {v.get('timestamp', 0.0):.2f}        | {v.get('confidence', 0.0):.2f}      | {v.get('worker_id', 'N/A')} |\n"
             violation_table += row
 
+        # Generate snapshots text
         snapshots_text = "\n".join(
             f"- Snapshot for {CONFIG['DISPLAY_NAMES'].get(s['violation'], 'Unknown')} at frame {s['frame']}: ![]({s['snapshot_base64']})"
             for s in snapshots
         ) if snapshots else "No snapshots captured."
 
+        # Push to Salesforce
+        try:
+            sf = connect_to_salesforce()
+            record_data = {
+                "Compliance_Score__c": score,
+                "Violations_Found__c": len(violations),
+                "Status__c": "Completed",
+                "Processing_Time__c": f"{processing_time:.2f}s"
+            }
+            record = sf.Safety_Video_Report__c.create(record_data)
+            record_id = record["id"]
+            
+            if pdf_file:
+                pdf_url = upload_pdf_to_salesforce(sf, pdf_file, record_id)
+        except Exception as e:
+            logger.error(f"Salesforce integration failed: {e}")
+            record_id = "N/A (Salesforce error)"
+
         yield (
             violation_table,
             f"Safety Score: {score}%",
             snapshots_text,
-            "Salesforce integration placeholder",
+            f"Salesforce Record ID: {record_id}",
             pdf_url or "N/A"
         )
 
@@ -370,6 +412,40 @@ def process_video(video_data):
         logger.error(f"Error processing video: {e}", exc_info=True)
         yield f"Error processing video: {e}", "", "", "", ""
 
+# ==========================
+# Salesforce Integration
+# ==========================
+@retry(stop_max_attempt_number=3, wait_fixed=2000)
+def connect_to_salesforce():
+    try:
+        sf = Salesforce(**CONFIG["SF_CREDENTIALS"])
+        logger.info("Connected to Salesforce")
+        return sf
+    except Exception as e:
+        logger.error(f"Salesforce connection failed: {e}")
+        raise
+
+def upload_pdf_to_salesforce(sf, pdf_file, report_id):
+    try:
+        encoded_pdf = base64.b64encode(pdf_file.getvalue()).decode('utf-8')
+        content_version_data = {
+            "Title": f"Safety_Violation_Report_{int(time.time())}",
+            "PathOnClient": f"safety_violation_{int(time.time())}.pdf",
+            "VersionData": encoded_pdf,
+            "FirstPublishLocationId": report_id
+        }
+        content_version = sf.ContentVersion.create(content_version_data)
+        result = sf.query(f"SELECT Id, ContentDocumentId FROM ContentVersion WHERE Id = '{content_version['id']}'")
+        if not result['records']:
+            logger.error("Failed to retrieve ContentVersion")
+            return ""
+        file_url = f"https://{sf.sf_instance}/sfc/servlet.shepherd/version/download/{content_version['id']}"
+        logger.info(f"PDF uploaded to Salesforce: {file_url}")
+        return file_url
+    except Exception as e:
+        logger.error(f"Error uploading PDF to Salesforce: {e}")
+        return ""
+
 # ==========================
 # Gradio Interface
 # ==========================
@@ -402,5 +478,5 @@ interface = gr.Interface(
 )
 
 if __name__ == "__main__":
-    logger.info("Launching Optimized Safety Analyzer App...")
+    logger.info("Launching Enhanced Safety Analyzer App...")
     interface.launch()