Update app.py

app.py

@@ -14,6 +14,8 @@ import base64
 import logging
 from retrying import retry
 import uuid
+from multiprocessing import Pool, cpu_count
+from functools import partial
 
 # ==========================
 # Optimized Configuration
@@ -61,7 +63,8 @@ CONFIG = {
     "WORKER_TRACKING_DURATION": 3.0,
     "MAX_PROCESSING_TIME": 60,  # 1 minute limit
     "FRAME_SKIP": 2,  # Process every 2nd frame for speed
-    "BATCH_SIZE": 16  # Frames per batch
+    "BATCH_SIZE": 16,  # Frames per batch
+    "PARALLEL_WORKERS": max(1, cpu_count() - 1)  # Use all CPU cores except one
 }
 
 # Setup logging
@@ -132,6 +135,36 @@ def calculate_iou(box1, box2):
     
     return intersection_area / union_area
 
+def process_frame_batch(frame_batch, frame_indices, fps):
+    batch_results = []
+    results = model(frame_batch, device=device, conf=0.1, verbose=False)
+    
+    for idx, (result, frame_idx) in enumerate(zip(results, frame_indices)):
+        current_time = frame_idx / fps
+        detections = []
+        
+        boxes = result.boxes
+        for box in boxes:
+            cls = int(box.cls)
+            conf = float(box.conf)
+            label = CONFIG["VIOLATION_LABELS"].get(cls, None)
+            
+            if label is None or conf < CONFIG["CONFIDENCE_THRESHOLDS"].get(label, 0.25):
+                continue
+
+            bbox = [round(x, 2) for x in box.xywh.cpu().numpy()[0]]
+            detections.append({
+                "frame": frame_idx,
+                "violation": label,
+                "confidence": round(conf, 2),
+                "bounding_box": bbox,
+                "timestamp": current_time
+            })
+        
+        batch_results.append((frame_idx, detections))
+    
+    return batch_results
+
 def generate_violation_pdf(violations, score):
     try:
         pdf_filename = f"violations_{int(time.time())}.pdf"
@@ -193,6 +226,87 @@ def calculate_safety_score(violations):
     score = 100 - total_penalty
     return max(score, 0)
 
+# ==========================
+# 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")
+        sf.describe()
+        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:
+        if not pdf_file:
+            logger.error("No PDF file provided for upload")
+            return ""
+        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 ""
+
+def push_report_to_salesforce(violations, score, pdf_path, pdf_file):
+    try:
+        sf = connect_to_salesforce()
+        violations_text = "\n".join(
+            f"{CONFIG['DISPLAY_NAMES'].get(v.get('violation', 'Unknown'), 'Unknown')} at {v.get('timestamp', 0.0):.2f}s (Confidence: {v.get('confidence', 0.0):.2f})"
+            for v in violations
+        ) or "No violations detected."
+        pdf_url = f"{CONFIG['PUBLIC_URL_BASE']}{os.path.basename(pdf_path)}" if pdf_path else ""
+
+        record_data = {
+            "Compliance_Score__c": score,
+            "Violations_Found__c": len(violations),
+            "Violations_Details__c": violations_text,
+            "Status__c": "Pending",
+            "PDF_Report_URL__c": pdf_url
+        }
+        logger.info(f"Creating Salesforce record with data: {record_data}")
+        try:
+            record = sf.Safety_Video_Report__c.create(record_data)
+            logger.info(f"Created Safety_Video_Report__c record: {record['id']}")
+        except Exception as e:
+            logger.error(f"Failed to create Safety_Video_Report__c: {e}")
+            record = sf.Account.create({"Name": f"Safety_Report_{int(time.time())}"})
+            logger.warning(f"Fell back to Account record: {record['id']}")
+        record_id = record["id"]
+
+        if pdf_file:
+            uploaded_url = upload_pdf_to_salesforce(sf, pdf_file, record_id)
+            if uploaded_url:
+                try:
+                    sf.Safety_Video_Report__c.update(record_id, {"PDF_Report_URL__c": 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}")
+                    sf.Account.update(record_id, {"Description": uploaded_url})
+                    logger.info(f"Updated Account record {record_id} with PDF URL")
+                pdf_url = uploaded_url
+
+        return record_id, pdf_url
+    except Exception as e:
+        logger.error(f"Salesforce record creation failed: {e}", exc_info=True)
+        return None, ""
+
 # ==========================
 # Fast Video Processing
 # ==========================
@@ -225,7 +339,6 @@ def process_video(video_data):
         helmet_violations = {}
         snapshots = []
         start_time = time.time()
-        processed_frames = 0
         frame_skip = CONFIG["FRAME_SKIP"]
 
         # Process frames in batches
@@ -250,7 +363,6 @@ def process_video(video_data):
                 
                 batch_frames.append(frame)
                 batch_indices.append(frame_idx)
-                processed_frames += 1
 
             # Break if no more frames
             if not batch_frames:
@@ -356,8 +468,8 @@ def process_video(video_data):
 
         score = calculate_safety_score(violations)
         pdf_path, pdf_url, pdf_file = generate_violation_pdf(violations, score)
+        report_id, final_pdf_url = push_report_to_salesforce(violations, score, pdf_path, pdf_file)
 
-        # Generate violation table
         violation_table = "| Violation              | Timestamp (s) | Confidence | Worker ID |\n"
         violation_table += "|------------------------|---------------|------------|-----------|\n"
         for v in sorted(violations, key=lambda x: x["timestamp"]):
@@ -365,76 +477,23 @@ def process_video(video_data):
             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,
-            f"Salesforce Record ID: {record_id}",
-            pdf_url or "N/A"
+            f"Salesforce Record ID: {report_id or 'N/A'}",
+            final_pdf_url or "N/A"
         )
 
     except Exception as e:
         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
 # ==========================
@@ -462,10 +521,10 @@ interface = gr.Interface(
         gr.Textbox(label="Violation Details URL")
     ],
     title="Worksite Safety Violation Analyzer",
-    description="Upload site videos to detect safety violations (No Helmet, No Harness, Unsafe Posture, Unsafe Zone, Improper Tool Use).",
+    description="Upload site videos to detect safety violations (No Helmet, No Harness, Unsafe Posture, Unsafe Zone, Improper Tool Use). Non-violations are ignored.",
     allow_flagging="never"
 )
 
 if __name__ == "__main__":
-    logger.info("Launching Safety Analyzer App...")
+    logger.info("Launching Enhanced Safety Analyzer App...")
     interface.launch()