Create utils.py

utils.py

@@ -0,0 +1,157 @@
+import os
+import cv2
+import numpy as np
+from reportlab.lib.pagesizes import letter
+from reportlab.pdfgen import canvas
+from reportlab.lib.units import inch
+from io import BytesIO
+import base64
+import asyncio
+import logging
+from simple_salesforce import Salesforce
+import tenacity
+from config import CONFIG
+import dlib
+
+logger = logging.getLogger(__name__)
+
+async def preprocess_frame(frame):
+    frame = cv2.resize(frame, CONFIG["TARGET_RESOLUTION"], interpolation=cv2.INTER_LINEAR)
+    frame = cv2.convertScaleAbs(frame, alpha=1.2, beta=20)
+    return frame
+
+def draw_detections(frame, detections):
+    result_frame = frame.copy()
+    for det in detections:
+        label = det.get("violation", "Unknown")
+        confidence = det.get("confidence", 0.0)
+        x, y, w, h = det.get("bounding_box", [0, 0, 0, 0])
+        worker_id = det.get("worker_id", "Unknown")
+        x1, y1 = int(x - w/2), int(y - h/2)
+        x2, y2 = int(x + w/2), int(y + h/2)
+        color = CONFIG["CLASS_COLORS"].get(label, (0, 0, 255))
+        cv2.rectangle(result_frame, (x1, y1), (x2, y2), color, 3)
+        display_text = f"{CONFIG['DISPLAY_NAMES'].get(label, label)} (Worker {worker_id})"
+        text_size = cv2.getTextSize(display_text, cv2.FONT_HERSHEY_SIMPLEX, 0.6, 2)[0]
+        cv2.rectangle(result_frame, (x1, y1-text_size[1]-10), (x1+text_size[0]+10, y1), (0, 0, 0), -1)
+        cv2.putText(result_frame, display_text, (x1+5, y1-5), cv2.FONT_HERSHEY_SIMPLEX, 0.6, (255, 255, 255), 2)
+        cv2.putText(result_frame, f"Conf: {confidence:.2f}", (x1+5, y2+20), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 255, 255), 2)
+    return result_frame
+
+def calculate_safety_score(violations):
+    penalties = {
+        "no_helmet": 25, "no_harness": 30, "unsafe_posture": 20,
+        "unsafe_zone": 35, "improper_tool_use": 25
+    }
+    worker_violations = {}
+    for v in violations:
+        worker_id = v.get("worker_id", "Unknown")
+        worker_violations.setdefault(worker_id, set()).add(v.get("violation", "Unknown"))
+    total_penalty = sum(sum(penalties.get(v, 0) for v in worker_violations[wid]) for wid in worker_violations)
+    return max(0, 100 - total_penalty)
+
+async def generate_violation_pdf(violations, score, output_dir):
+    try:
+        pdf_filename = f"violations_{int(asyncio.get_event_loop().time())}.pdf"
+        pdf_path = os.path.join(output_dir, pdf_filename)
+        pdf_file = BytesIO()
+        c = canvas.Canvas(pdf_file, pagesize=letter)
+        c.setFont("Helvetica-Bold", 16)
+        c.drawString(1 * inch, 10 * inch, "Worksite Safety Violation Report")
+        c.setFont("Helvetica", 12)
+        c.drawString(1 * inch, 9.5 * inch, f"Date: {asyncio.get_event_loop().time():.0f}")
+        c.setFont("Helvetica-Bold", 14)
+        c.drawString(1 * inch, 8.7 * inch, f"Safety Compliance Score: {score}%")
+        c.setFont("Helvetica", 10)
+        y_position = 8.2 * inch
+        worker_violations = {}
+        for v in violations:
+            worker_id = v.get("worker_id", "Unknown")
+            worker_violations.setdefault(worker_id, []).append(v)
+        for worker_id, worker_vios in worker_violations.items():
+            c.drawString(1 * inch, y_position, f"Worker {worker_id}:")
+            y_position -= 0.2 * inch
+            for v in worker_vios:
+                display_name = CONFIG["DISPLAY_NAMES"].get(v.get("violation", "Unknown"), "Unknown")
+                c.drawString(1.2 * inch, y_position, f"  - {display_name} at {v.get('timestamp', 0.0):.2f}s (Conf: {v.get('confidence', 0.0):.2f})")
+                y_position -= 0.2 * inch
+                if y_position < 1 * inch:
+                    c.showPage()
+                    y_position = 10 * inch
+        c.save()
+        pdf_file.seek(0)
+        with open(pdf_path, "wb") as f:
+            f.write(pdf_file.getvalue())
+        return pdf_path, f"{CONFIG['PUBLIC_URL_BASE']}{pdf_filename}", pdf_file
+    except Exception as e:
+        logger.error(f"Error generating PDF: {e}")
+        return "", "", None
+
+@tenacity.retry(stop=tenacity.stop_after_attempt(3), wait=tenacity.wait_fixed(2))
+async def connect_to_salesforce():
+    sf = Salesforce(**CONFIG["SF_CREDENTIALS"])
+    logger.info("Connected to Salesforce")
+    return sf
+
+async def upload_pdf_to_salesforce(sf, pdf_file, record_id):
+    try:
+        encoded_pdf = base64.b64encode(pdf_file.getvalue()).decode('utf-8')
+        content_version = sf.ContentVersion.create({
+            "Title": f"Safety_Violation_Report_{int(asyncio.get_event_loop().time())}",
+            "PathOnClient": f"safety_violation_{int(asyncio.get_event_loop().time())}.pdf",
+            "VersionData": encoded_pdf,
+            "FirstPublishLocationId": record_id
+        })
+        result = sf.query(f"SELECT Id, ContentDocumentId FROM ContentVersion WHERE Id = '{content_version['id']}'")
+        return f"https://{sf.sf_instance}/sfc/servlet.shepherd/version/download/{content_version['id']}"
+    except Exception as e:
+        logger.error(f"Error uploading PDF to Salesforce: {e}")
+        return ""
+
+async def push_report_to_salesforce(violations, score, pdf_path, pdf_file):
+    try:
+        sf = await connect_to_salesforce()
+        violations_text = "".join(
+            f"Worker {v.get('worker_id', 'Unknown')}: {CONFIG['DISPLAY_NAMES'].get(v.get('violation', 'Unknown'), 'Unknown')} at {v.get('timestamp', 0.0):.2f}s (Conf: {v.get('confidence', 0.0):.2f})\n"
+            for v in violations
+        ) or "No violations detected."
+        record_data = {
+            "Compliance_Score__c": score,
+            "Violations_Found__c": len(violations),
+            "Violations_Details__c": violations_text,
+            "Status__c": "Pending",
+            "PDF_Report_URL__c": f"{CONFIG['PUBLIC_URL_BASE']}{os.path.basename(pdf_path)}" if pdf_path else ""
+        }
+        record = sf.Safety_Video_Report__c.create(record_data)
+        if pdf_file:
+            uploaded_url = await upload_pdf_to_salesforce(sf, pdf_file, record["id"])
+            if uploaded_url:
+                sf.Safety_Video_Report__c.update(record["id"], {"PDF_Report_URL__c": uploaded_url})
+        return record["id"], uploaded_url or record_data["PDF_Report_URL__c"]
+    except Exception as e:
+        logger.error(f"Salesforce record creation failed: {e}")
+        return "N/A", "Salesforce integration failed."
+
+@tenacity.retry(
+    stop=tenacity.stop_after_attempt(3),
+    wait=tenacity.wait_fixed(1),
+    retry=tenacity.retry_if_exception_type((IOError, OSError))
+)
+def verify_and_open_video(video_path):
+    if not os.path.exists(video_path) or os.path.getsize(video_path) == 0:
+        raise ValueError(f"Invalid video file: {video_path}")
+    cap = cv2.VideoCapture(video_path)
+    if not cap.isOpened():
+        raise ValueError("Could not open video file.")
+    return cap
+
+def blur_faces(frame):
+    detector = dlib.get_frontal_face_detector()
+    gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
+    faces = detector(gray)
+    for face in faces:
+        x, y, w, h = face.left(), face.top(), face.width(), face.height()
+        roi = frame[y:y+h, x:x+w]
+        roi = cv2.GaussianBlur(roi, (23, 23), 30)
+        frame[y:y+h, x:x+w] = roi
+    return frame