Update app.py

app.py

@@ -1,6 +1,5 @@
 import os
 import sys
-import subprocess
 import logging
 import warnings
 import cv2
@@ -16,10 +15,8 @@ from reportlab.lib.units import inch
 from io import BytesIO
 import base64
 from retrying import retry
-import uuid
-from multiprocessing import Pool, cpu_count
-from functools import partial
 from collections import defaultdict
+from multiprocessing import cpu_count
 
 # ========================== # Configuration and Setup # ==========================
 os.environ['YOLO_CONFIG_DIR'] = '/tmp/Ultralytics'
@@ -29,7 +26,7 @@ logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(
 logger = logging.getLogger(__name__)
 warnings.filterwarnings("ignore")
 
-# ========================== # Position-Based Tracker Implementation # ==========================
+# ========================== # Optimized Tracker Implementation (No Face Recognition) # ==========================
 class SafetyTracker:
     def __init__(self, track_thresh=0.3, track_buffer=30, match_thresh=0.7, frame_rate=30):
         self.track_thresh = track_thresh
@@ -37,13 +34,10 @@ class SafetyTracker:
         self.match_thresh = match_thresh
         self.frame_rate = frame_rate
         self.next_id = 1
+        self.worker_tracks = {}
+        self.violation_history = defaultdict(dict)
+        self.position_history = defaultdict(list)
         
-        # Tracking stores
-        self.worker_tracks = {}  # Active tracks
-        self.violation_history = defaultdict(dict)  # {worker_id: {violation_type: last_detection_time}}
-        self.position_history = defaultdict(list)  # {worker_id: [positions]}
-        
-        # Violation cooldowns (seconds)
         self.VIOLATION_COOLDOWNS = {
             "no_helmet": 30.0,
             "no_harness": 20.0,
@@ -52,8 +46,7 @@ class SafetyTracker:
             "improper_tool_use": 15.0
         }
 
-    def update(self, detections):
-        """Update tracks with new detections using position-based matching"""
+    def update(self, detections, frame):
         current_time = time.time()
         new_violations = []
         
@@ -62,14 +55,11 @@ class SafetyTracker:
             label = det['violation']
             confidence = det['confidence']
             
-            # Match by position
             worker_id = self._match_by_position(bbox, label)
-            
             if worker_id is None:
                 worker_id = self.next_id
                 self.next_id += 1
                 
-            # Check if new violation
             if self._is_new_violation(worker_id, label, current_time):
                 violation = {
                     'worker_id': worker_id,
@@ -81,62 +71,41 @@ class SafetyTracker:
                 new_violations.append(violation)
                 self.violation_history[worker_id][label] = current_time
             
-            # Update position history
-            x, y, w, h = bbox
-            self.position_history[worker_id].append((x, y))
-            
-            # Update active tracks
             self.worker_tracks[worker_id] = {
                 'bbox': bbox,
                 'last_seen': current_time,
                 'label': label
             }
+            self.position_history[worker_id].append((bbox[0], bbox[1]))
         
-        # Cleanup old tracks
         self._cleanup_tracks(current_time)
-        
         return new_violations
 
     def _match_by_position(self, bbox, label):
-        """Match detection to existing tracks using position"""
-        x, y, w, h = bbox
-        current_pos = (x, y)
-        
+        x, y = bbox[0], bbox[1]
         for worker_id, positions in self.position_history.items():
-            # Only match if worker has had this violation type before
-            if label not in self.violation_history[worker_id]:
-                continue
-                
-            # Check distance to historical positions
             for pos in positions[-5:]:  # Check last 5 positions
-                distance = np.sqrt((current_pos[0]-pos[0])**2 + (current_pos[1]-pos[1])**2)
-                if distance < 100:  # Within 100 pixels
+                if np.sqrt((x-pos[0])**2 + (y-pos[1])**2) < 100:
                     return worker_id
         return None
 
     def _is_new_violation(self, worker_id, label, current_time):
-        """Check if violation is new based on cooldown"""
         if label not in self.violation_history[worker_id]:
             return True
-            
-        last_time = self.violation_history[worker_id][label]
-        cooldown = self.VIOLATION_COOLDOWNS.get(label, 10.0)
-        return (current_time - last_time) > cooldown
+        return (current_time - self.violation_history[worker_id][label]) > self.VIOLATION_COOLDOWNS.get(label, 10.0)
 
     def _cleanup_tracks(self, current_time):
-        """Remove inactive tracks"""
         inactive_ids = [
-            id for id, track in self.worker_tracks.items()
+            wid for wid, track in self.worker_tracks.items()
             if (current_time - track['last_seen']) > (self.track_buffer / self.frame_rate)
         ]
-        for id in inactive_ids:
-            self.worker_tracks.pop(id, None)
-            self.position_history.pop(id, None)
-            # Keep violation history for longer
-            if (current_time - max(self.violation_history[id].values(), default=0)) > 300:
-                self.violation_history.pop(id, None)
+        for wid in inactive_ids:
+            self.worker_tracks.pop(wid, None)
+            self.position_history.pop(wid, None)
+            if (current_time - max(self.violation_history[wid].values(), default=0)) > 300:
+                self.violation_history.pop(wid, None)
 
-# ========================== # App Configuration # ==========================
+# ========================== # Configuration # ==========================
 CONFIG = {
     "MODEL_PATH": "yolov8_safety.pt",
     "FALLBACK_MODEL": "yolov8n.pt",
@@ -156,10 +125,10 @@ CONFIG = {
         "improper_tool_use": (255, 255, 0)
     },
     "DISPLAY_NAMES": {
-        "no_helmet": "No Helmet Violation",
-        "no_harness": "No Harness Violation",
+        "no_helmet": "No Helmet",
+        "no_harness": "No Harness",
         "unsafe_posture": "Unsafe Posture",
-        "unsafe_zone": "Unsafe Zone Entry",
+        "unsafe_zone": "Unsafe Zone",
         "improper_tool_use": "Improper Tool Use"
     },
     "SF_CREDENTIALS": {
@@ -177,7 +146,7 @@ CONFIG = {
         "improper_tool_use": 0.3
     },
     "FRAME_SKIP": 2,
-    "BATCH_SIZE": 8,  # Reduced for stability
+    "BATCH_SIZE": 8,
     "SNAPSHOT_QUALITY": 90
 }
 
@@ -188,9 +157,7 @@ logger.info(f"Using device: {device}")
 def load_model():
     try:
         model_path = CONFIG["MODEL_PATH"] if os.path.exists(CONFIG["MODEL_PATH"]) else CONFIG["FALLBACK_MODEL"]
-        logger.info(f"Loading model: {model_path}")
         if not os.path.exists(model_path):
-            logger.info("Downloading fallback model...")
             torch.hub.download_url_to_file('https://github.com/ultralytics/assets/releases/download/v8.3.0/yolov8n.pt', model_path)
         return YOLO(model_path).to(device)
     except Exception as e:
@@ -199,35 +166,23 @@ def load_model():
 
 model = load_model()
 
-def preprocess_frame(frame):
-    """Basic image enhancement"""
-    return cv2.convertScaleAbs(frame, alpha=1.2, beta=20)
-
 def draw_detections(frame, detections):
-    """Draw bounding boxes with labels"""
-    result = frame.copy()
+    annotated = frame.copy()
     for det in detections:
         x, y, w, h = det['bbox']
         x1, y1 = int(x-w/2), int(y-h/2)
         x2, y2 = int(x+w/2), int(y+h/2)
-        label = CONFIG["DISPLAY_NAMES"].get(det['violation'], det['violation'])
-        color = CONFIG["CLASS_COLORS"].get(det['violation'], (0,0,255))
-        
-        cv2.rectangle(result, (x1, y1), (x2, y2), color, 3)
-        cv2.putText(result, f"{label} (Worker {det['worker_id']})", 
-                   (x1, y1-10), cv2.FONT_HERSHEY_SIMPLEX, 0.6, (255,255,255), 2)
-    return result
+        color = CONFIG["CLASS_COLORS"][det['violation']]
+        cv2.rectangle(annotated, (x1, y1), (x2, y2), color, 2)
+        label = f"{CONFIG['DISPLAY_NAMES'][det['violation']]} (Worker {det['worker_id']})"
+        cv2.putText(annotated, label, (x1, y1-10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255,255,255), 2)
+    return annotated
 
 def calculate_safety_score(violations):
-    penalty_map = {
-        "no_helmet": 25, "no_harness": 30, 
-        "unsafe_posture": 20, "unsafe_zone": 35, 
-        "improper_tool_use": 25
-    }
+    penalties = {"no_helmet":25, "no_harness":30, "unsafe_posture":20, "unsafe_zone":35, "improper_tool_use":25}
     unique_violations = {v['violation'] for v in violations}
-    return max(0, 100 - sum(penalty_map.get(v,0) for v in unique_violations))
+    return max(0, 100 - sum(penalties.get(v,0) for v in unique_violations))
 
-# ========================== # Reporting Functions # ==========================
 def generate_violation_pdf(violations, score):
     try:
         pdf_buffer = BytesIO()
@@ -237,16 +192,16 @@ def generate_violation_pdf(violations, score):
         c.setFont("Helvetica-Bold", 16)
         c.drawString(1*inch, 10*inch, "Safety Violation Report")
         c.setFont("Helvetica", 12)
-        c.drawString(1*inch, 9.5*inch, f"Generated: {time.strftime('%Y-%m-%d %H:%M:%S')}")
+        c.drawString(1*inch, 9.5*inch, f"Date: {time.strftime('%Y-%m-%d %H:%M:%S')}")
         c.drawString(1*inch, 9*inch, f"Safety Score: {score}%")
         
-        # Violations list
+        # Violations List
         y = 8.5*inch
         c.setFont("Helvetica-Bold", 14)
         c.drawString(1*inch, y, "Violations Detected:")
         y -= 0.3*inch
-        
         c.setFont("Helvetica", 10)
+        
         for v in violations:
             text = f"Worker {v['worker_id']}: {CONFIG['DISPLAY_NAMES'][v['violation']]} at {v['timestamp']:.1f}s"
             c.drawString(1.2*inch, y, text)
@@ -254,6 +209,7 @@ def generate_violation_pdf(violations, score):
             if y < 1*inch:
                 c.showPage()
                 y = 10*inch
+                c.setFont("Helvetica", 10)
         
         c.save()
         pdf_buffer.seek(0)
@@ -272,85 +228,63 @@ def generate_violation_pdf(violations, score):
 def connect_to_salesforce():
     try:
         sf = Salesforce(**CONFIG["SF_CREDENTIALS"])
-        logger.info("Salesforce connection established")
+        sf.describe()
         return sf
     except Exception as e:
         logger.error(f"Salesforce connection failed: {e}")
         raise
 
-def push_report_to_salesforce(violations, score, pdf_path, pdf_buffer):
+def upload_to_salesforce(sf, pdf_file, violations, score):
     try:
-        sf = connect_to_salesforce()
-        
         # Create record
-        record_data = {
+        violations_text = "\n".join(
+            f"Worker {v['worker_id']}: {CONFIG['DISPLAY_NAMES'][v['violation']]} at {v['timestamp']:.1f}s"
+            for v in violations
+        )
+        
+        record = sf.Safety_Video_Report__c.create({
             "Compliance_Score__c": score,
             "Violations_Found__c": len(violations),
-            "Violations_Details__c": "\n".join(
-                f"Worker {v['worker_id']}: {CONFIG['DISPLAY_NAMES'][v['violation']]}"
-                for v in violations
-            ),
-            "Status__c": "New"
-        }
+            "Violations_Details__c": violations_text
+        })
         
-        record = sf.Safety_Video_Report__c.create(record_data)
-        record_id = record['id']
-        logger.info(f"Created Salesforce record: {record_id}")
+        # Upload PDF
+        encoded = base64.b64encode(pdf_file.getvalue()).decode()
+        content = sf.ContentVersion.create({
+            "Title": f"Safety_Report_{int(time.time())}",
+            "PathOnClient": "report.pdf",
+            "VersionData": encoded,
+            "FirstPublishLocationId": record['id']
+        })
         
-        # Upload PDF if available
-        pdf_url = ""
-        if pdf_buffer:
-            encoded = base64.b64encode(pdf_buffer.getvalue()).decode()
-            content_version = sf.ContentVersion.create({
-                "Title": f"Safety_Report_{record_id}",
-                "PathOnClient": "report.pdf",
-                "VersionData": encoded,
-                "FirstPublishLocationId": record_id
-            })
-            pdf_url = f"https://{sf.sf_instance}/sfc/servlet.shepherd/version/download/{content_version['id']}"
-            logger.info(f"PDF uploaded: {pdf_url}")
-        
-        return record_id, pdf_url
+        return record['id'], f"https://{sf.sf_instance}/sfc/servlet.shepherd/version/download/{content['id']}"
     except Exception as e:
         logger.error(f"Salesforce upload failed: {e}")
-        return None, ""
+        return None, None
 
 # ========================== # Video Processing # ==========================
 def process_video(video_data):
     try:
         os.makedirs(CONFIG["OUTPUT_DIR"], exist_ok=True)
-        
-        # Save video
-        video_path = os.path.join(CONFIG["OUTPUT_DIR"], f"input_{int(time.time())}.mp4")
+        video_path = os.path.join(CONFIG["OUTPUT_DIR"], f"temp_{int(time.time())}.mp4")
         with open(video_path, "wb") as f:
             f.write(video_data)
             
         cap = cv2.VideoCapture(video_path)
-        if not cap.isOpened():
-            raise ValueError("Failed to open video")
-            
         fps = cap.get(cv2.CAP_PROP_FPS) or 30
-        total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
         tracker = SafetyTracker(frame_rate=fps)
         snapshots = []
         
-        frame_count = 0
-        while True:
+        while cap.isOpened():
             ret, frame = cap.read()
             if not ret:
                 break
                 
-            if frame_count % CONFIG["FRAME_SKIP"] != 0:
-                frame_count += 1
-                continue
-                
-            # Process frame
-            frame = preprocess_frame(frame)
-            results = model(frame, verbose=False)[0]
+            frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
+            results = model(frame, verbose=False)
             
-            # Get detections
             detections = []
-            for box in results.boxes:
+            for box in results[0].boxes:
                 cls = int(box.cls)
                 label = CONFIG["VIOLATION_LABELS"].get(cls)
                 if label and box.conf > CONFIG["CONFIDENCE_THRESHOLDS"].get(label, 0.3):
@@ -360,72 +294,49 @@ def process_video(video_data):
                         'confidence': float(box.conf)
                     })
             
-            # Update tracker
-            new_violations = tracker.update(detections)
+            new_violations = tracker.update(detections, frame)
             
-            # Capture snapshots for new violations
             for violation in new_violations:
                 snapshot = draw_detections(frame.copy(), [violation])
                 timestamp = time.strftime("%Y%m%d_%H%M%S")
-                img_path = os.path.join(
-                    CONFIG["OUTPUT_DIR"],
-                    f"violation_{violation['worker_id']}_{violation['violation']}_{timestamp}.jpg"
-                )
-                cv2.imwrite(img_path, snapshot, [cv2.IMWRITE_JPEG_QUALITY, CONFIG["SNAPSHOT_QUALITY"]])
+                img_path = os.path.join(CONFIG["OUTPUT_DIR"], f"violation_{violation['worker_id']}_{timestamp}.jpg")
+                cv2.imwrite(img_path, cv2.cvtColor(snapshot, cv2.COLOR_RGB2BGR), 
+                           [cv2.IMWRITE_JPEG_QUALITY, CONFIG["SNAPSHOT_QUALITY"]])
                 snapshots.append({
                     'path': img_path,
                     'url': f"{CONFIG['PUBLIC_URL_BASE']}{os.path.basename(img_path)}",
                     'violation': violation
                 })
-            
-            # Update progress
-            if frame_count % 10 == 0:
-                progress = min(100, frame_count / total_frames * 100)
-                yield f"Processing... {progress:.1f}%", "", "", "", ""
                 
-            frame_count += 1
-            
+            yield f"Processing frame {int(cap.get(cv2.CAP_PROP_POS_FRAMES))}...", "", "", "", ""
+        
         cap.release()
         os.remove(video_path)
         
-        # Generate report
-        violations = [
-            {
-                'worker_id': worker_id,
-                'violation': violation_type,
-                'timestamp': detection_time
-            }
-            for worker_id, violations in tracker.violation_history.items()
-            for violation_type, detection_time in violations.items()
-        ]
-        
-        if not violations:
-            yield "No violations found", "Safety Score: 100%", "No snapshots", "N/A", "N/A"
+        if not snapshots:
+            yield "No violations detected", "Safety Score: 100%", "No snapshots", "N/A", "N/A"
             return
             
-        score = calculate_safety_score(violations)
-        pdf_path, pdf_url, pdf_buffer = generate_violation_pdf(violations, score)
-        record_id, salesforce_url = push_report_to_salesforce(violations, score, pdf_path, pdf_buffer)
-        
-        # Format output
-        violations_table = "| Violation | Worker ID | Time |\n|-----------|-----------|------|\n"
-        violations_table += "\n".join(
-            f"| {CONFIG['DISPLAY_NAMES'][v['violation']]} | {v['worker_id']} | {v['timestamp']:.1f}s |"
-            for v in violations
-        )
+        score = calculate_safety_score([v['violation'] for v in snapshots])
+        pdf_path, pdf_url, pdf_file = generate_violation_pdf([v['violation'] for v in snapshots], score)
         
-        snapshots_md = "\n\n".join(
-            f"### {CONFIG['DISPLAY_NAMES'][s['violation']['violation']]} (Worker {s['violation']['worker_id']})\n"
-            f"![Snapshot]({s['url']})"
-            for s in snapshots
+        if pdf_file:
+            record_id, sf_url = upload_to_salesforce(connect_to_salesforce(), pdf_file, 
+                                                   [v['violation'] for v in snapshots], score)
+        else:
+            record_id, sf_url = None, None
+            
+        snapshots_md = "\n".join(
+            f"![{v['violation']['violation']}]({v['url']})" 
+            for v in snapshots
         )
         
         yield (
-            violations_table,
+            "\n".join(f"- {v['violation']['violation']} (Worker {v['violation']['worker_id']})" for v in snapshots),
             f"Safety Score: {score}%",
-            snapshots_md or "No snapshots",
+            snapshots_md,
             f"Salesforce ID: {record_id or 'N/A'}",
-            salesforce_url or pdf_url or "N/A"
+            sf_url or pdf_url or "N/A"
         )
         
     except Exception as e:
@@ -435,30 +346,22 @@ def process_video(video_data):
         yield f"Error: {str(e)}", "", "", "", ""
 
 # ========================== # Gradio Interface # ==========================
-def gradio_interface(video_file):
-    if not video_file:
+def gradio_interface(video):
+    if not video:
         return "Upload a video file", "", "", "", ""
-        
-    try:
-        with open(video_file, "rb") as f:
-            video_data = f.read()
-            
-        for output in process_video(video_data):
-            yield output
-            
-    except Exception as e:
-        logger.error(f"Interface error: {e}")
-        yield f"Error: {str(e)}", "", "", "", ""
+    
+    for update in process_video(open(video, "rb").read()):
+        yield update
 
 interface = gr.Interface(
     fn=gradio_interface,
-    inputs=gr.Video(label="Upload Safety Video"),
+    inputs=gr.Video(),
     outputs=[
-        gr.Markdown("## Detected Violations"),
-        gr.Textbox(label="Safety Score"),
-        gr.Markdown("## Evidence Snapshots"),
-        gr.Textbox(label="Salesforce Record"),
-        gr.Textbox(label="Report URL")
+        gr.Markdown("Detected Violations"),
+        gr.Textbox("Safety Score"),
+        gr.Markdown("Evidence Snapshots"),
+        gr.Textbox("Salesforce Record"),
+        gr.Textbox("Report URL")
     ],
     title="AI Safety Compliance Analyzer",
     description="Detects PPE and safety violations in worksite videos"