Update app.py

app.py

@@ -50,9 +50,9 @@ CONFIG = {
         "domain": "login"
     },
     "PUBLIC_URL_BASE": "https://huggingface.co/spaces/PrashanthB461/AI_Safety_Demo2/resolve/main/static/output/",
-    "FRAME_SKIP": 1,
+    "FRAME_SKIP": 3,  # Increased from 1 to 3 for faster processing
     "CONFIDENCE_THRESHOLDS": {
-        "no_helmet": 0.6,  # Increased to reduce false positives
+        "no_helmet": 0.6,
         "no_harness": 0.15,
         "unsafe_posture": 0.15,
         "unsafe_zone": 0.15,
@@ -60,8 +60,10 @@ CONFIG = {
     },
     "IOU_THRESHOLD": 0.4,
     "MIN_VIOLATION_FRAMES": 2,
-    "HELMET_CONFIDENCE_THRESHOLD": 0.65,  # Lowered to mark more workers as compliant
-    "MAX_PROCESSING_TIME": 60
+    "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
 }
 
 # Setup logging
@@ -93,7 +95,7 @@ def load_model():
 model = load_model()
 
 # ==========================
-# Enhanced Helper Functions
+# Optimized Helper Functions
 # ==========================
 def draw_detections(frame, detections):
     for det in detections:
@@ -123,27 +125,14 @@ def calculate_iou(box1, box2):
     x2_min, y2_min = x2 - w2/2, y2 - h2/2
     x2_max, y2_max = x2 + w2/2, y2 + h2/2
     
-    intersection = max(0, x1_max - x1_min) * max(0, y1_max - y1_min)
+    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
     
     return intersection / union if union > 0 else 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 generate_violation_pdf(violations, score):
     try:
         pdf_filename = f"violations_{int(time.time())}.pdf"
@@ -193,299 +182,169 @@ def generate_violation_pdf(violations, score):
         logger.error(f"Error generating PDF: {e}")
         return "", "", None
 
-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, ""
-
-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)
-
 # ==========================
-# Enhanced Video Processing
+# Optimized Video Processing
 # ==========================
 def process_video(video_data):
     try:
+        # Create temp video file
         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)
         logger.info(f"Video saved: {video_path}")
 
-        video = cv2.VideoCapture(video_path)
-        if not video.isOpened():
+        # Open video file
+        cap = cv2.VideoCapture(video_path)
+        if not cap.isOpened():
             raise ValueError("Could not open video file")
 
-        violations = []
-        snapshots = []
-        frame_count = 0
-        total_frames = int(video.get(cv2.CAP_PROP_FRAME_COUNT))
-        fps = video.get(cv2.CAP_PROP_FPS)
+        # Get video properties
+        total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
+        fps = cap.get(cv2.CAP_PROP_FPS)
         if fps <= 0:
             fps = 30
-        video_duration = total_frames / fps
-        logger.info(f"Video duration: {video_duration:.2f} seconds, Total frames: {total_frames}, FPS: {fps}")
+        duration = total_frames / fps
+        logger.info(f"Video duration: {duration:.2f}s, Frames: {total_frames}, FPS: {fps}")
+
+        # 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
 
         workers = []
-        violation_history = {label: [] for label in CONFIG["VIOLATION_LABELS"].values()}
-        confirmed_violations = {}
-        snapshot_taken = {label: False for label in CONFIG["VIOLATION_LABELS"].values()}
-        helmet_compliance = {}
-        detection_counts = {label: 0 for label in CONFIG["VIOLATION_LABELS"].values()}
+        violations = []
+        helmet_violations = {}
+        snapshots = []
         start_time = time.time()
-
-        # Calculate frames to process within 30 seconds
-        target_frames = int(total_frames / CONFIG["FRAME_SKIP"])
-        frame_indices = np.linspace(0, total_frames - 1, target_frames, dtype=int)
-
         processed_frames = 0
-        for idx in frame_indices:
-            elapsed_time = time.time() - start_time
-            if elapsed_time > CONFIG["MAX_PROCESSING_TIME"]:
-                logger.info(f"Processing time limit of {CONFIG['MAX_PROCESSING_TIME']} seconds reached. Processed {processed_frames}/{target_frames} frames.")
-                break
 
-            video.set(cv2.CAP_PROP_POS_FRAMES, idx)
-            ret, frame = video.read()
-            if not ret:
-                continue
+        # Process frames in batches for efficiency
+        for batch_start in range(0, total_frames, CONFIG["BATCH_SIZE"] * frame_skip):
+            batch_frames = []
+            batch_indices = []
+            
+            # Collect frames for this batch
+            for i in range(CONFIG["BATCH_SIZE"]):
+                frame_idx = batch_start + i * frame_skip
+                if frame_idx >= total_frames:
+                    break
+                
+                cap.set(cv2.CAP_PROP_POS_FRAMES, frame_idx)
+                ret, frame = cap.read()
+                if not ret:
+                    continue
+                
+                batch_frames.append(frame)
+                batch_indices.append(frame_idx)
+                processed_frames += 1
 
-            processed_frames += 1
-            current_time = idx / fps
-            progress = (processed_frames / target_frames) * 100
-            yield f"Processing video... {progress:.1f}% complete (Frame {idx}/{total_frames})", "", "", "", ""
+            # Skip empty batches
+            if not batch_frames:
+                continue
 
-            # Run detection on this frame
-            results = model(frame, device=device, conf=0.1, iou=CONFIG["IOU_THRESHOLD"])
+            # Run batch detection
+            results = model(batch_frames, device=device, conf=0.1, iou=CONFIG["IOU_THRESHOLD"], verbose=False)
             
-            current_detections = []
-            for result in results:
+            # 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})", "", "", "", ""
+
+                # Process detections in this frame
                 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:
-                        logger.warning(f"Unknown class ID {cls} detected, skipping")
-                        continue
-                        
-                    if conf < CONFIG["CONFIDENCE_THRESHOLDS"].get(label, 0.25):
-                        logger.debug(f"Detection {label} with confidence {conf:.2f} below threshold, skipping")
+                    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]]
-                    
-                    current_detections.append({
-                        "frame": idx,
+                    detection = {
+                        "frame": frame_idx,
                         "violation": label,
                         "confidence": round(conf, 2),
                         "bounding_box": bbox,
                         "timestamp": current_time
-                    })
-                    detection_counts[label] += 1
-
-            logger.debug(f"Frame {idx}: Detected {len(current_detections)} violations: {[d['violation'] for d in current_detections]}")
-
-            for detection in current_detections:
-                violation_type = detection.get("violation", None)
-                if violation_type is None:
-                    logger.error(f"Invalid detection, missing 'violation' key: {detection}")
-                    continue
+                    }
 
-                if violation_type == "no_helmet":
-                    matched_worker = None
+                    # Worker tracking and helmet detection optimization
+                    worker_id = None
                     max_iou = 0
-                    for worker in workers:
-                        iou = calculate_iou(detection["bounding_box"], worker["bbox"])
+                    for idx, worker in enumerate(workers):
+                        iou = calculate_iou(bbox, worker["bbox"])
                         if iou > max_iou and iou > CONFIG["IOU_THRESHOLD"]:
                             max_iou = iou
-                            matched_worker = worker
-                    
-                    if matched_worker:
-                        worker_id = matched_worker["id"]
-                        if worker_id not in helmet_compliance:
-                            helmet_compliance[worker_id] = {"no_helmet_frames": 0, "compliant": False}
-                        helmet_compliance[worker_id]["no_helmet_frames"] += 1
-                        if detection["confidence"] < CONFIG["HELMET_CONFIDENCE_THRESHOLD"]:
-                            helmet_compliance[worker_id]["compliant"] = True
-                            logger.debug(f"Worker {worker_id} marked as helmet compliant due to low no_helmet confidence: {detection['confidence']}")
-                        else:
-                            logger.debug(f"Worker {worker_id} potential no_helmet violation with confidence: {detection['confidence']}")
-                        if helmet_compliance[worker_id]["compliant"]:
-                            logger.debug(f"Worker {worker_id} has helmet, skipping no_helmet violation")
-                            continue
-                
-                matched_worker = None
-                max_iou = 0
-                
-                for worker in workers:
-                    iou = calculate_iou(detection["bounding_box"], worker["bbox"])
-                    if iou > max_iou and iou > CONFIG["IOU_THRESHOLD"]:
-                        max_iou = iou
-                        matched_worker = worker
-
-                if matched_worker:
-                    matched_worker["bbox"] = detection["bounding_box"]
-                    matched_worker["last_seen"] = current_time
-                    worker_id = matched_worker["id"]
-                else:
-                    worker_id = len(workers) + 1
-                    workers.append({
-                        "id": worker_id,
-                        "bbox": detection["bounding_box"],
-                        "first_seen": current_time,
-                        "last_seen": current_time
-                    })
-                    if worker_id not in helmet_compliance:
-                        helmet_compliance[worker_id] = {"no_helmet_frames": 0, "compliant": False}
-                
-                if worker_id in confirmed_violations and violation_type in confirmed_violations[worker_id]:
-                    logger.debug(f"Violation {violation_type} already confirmed for worker {worker_id}, skipping")
-                    continue
-                
-                detection["worker_id"] = worker_id
-                violation_history[violation_type].append(detection)
+                            worker_id = worker["id"]
+                            workers[idx]["bbox"] = bbox  # Update worker position
+                            workers[idx]["last_seen"] = current_time
+
+                    if worker_id is None:
+                        worker_id = len(workers) + 1
+                        workers.append({
+                            "id": worker_id,
+                            "bbox": bbox,
+                            "first_seen": current_time,
+                            "last_seen": current_time
+                        })
 
-            workers = [w for w in workers if current_time - w["last_seen"] < 5.0]
+                    # Special handling for helmet violations
+                    if label == "no_helmet":
+                        if worker_id not in helmet_violations:
+                            helmet_violations[worker_id] = []
+                        helmet_violations[worker_id].append(detection)
+                    else:
+                        violations.append(detection)
 
-        logger.info(f"Detection counts: {detection_counts}")
+                # Remove workers not seen recently
+                workers = [w for w in workers if current_time - w["last_seen"] < CONFIG["WORKER_TRACKING_DURATION"]]
 
-        for violation_type, detections in violation_history.items():
-            if not detections:
-                logger.info(f"No detections for {violation_type}")
-                continue
+            # 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)
+        for worker_id, detections in helmet_violations.items():
+            if len(detections) >= CONFIG["MIN_VIOLATION_FRAMES"]:
+                best_detection = max(detections, key=lambda x: x["confidence"])
+                violations.append(best_detection)
                 
-            worker_violations = {}
-            for det in detections:
-                if det["worker_id"] not in worker_violations:
-                    worker_violations[det["worker_id"]] = []
-                worker_violations[det["worker_id"]].append(det)
-            
-            for worker_id, worker_dets in worker_violations.items():
-                if len(worker_dets) >= CONFIG["MIN_VIOLATION_FRAMES"]:
-                    if worker_id in confirmed_violations and violation_type in confirmed_violations[worker_id]:
-                        continue
-                        
-                    if violation_type == "no_helmet":
-                        if worker_id in helmet_compliance and helmet_compliance[worker_id]["compliant"]:
-                            logger.debug(f"Skipping no_helmet for worker {worker_id} due to helmet compliance")
-                            continue
-                        # Removed the stricter persistence check for no_helmet
-                        logger.info(f"Confirmed no_helmet for worker {worker_id} with {len(worker_dets)} detections")
-                    
-                    best_detection = max(worker_dets, key=lambda x: x["confidence"])
-                    violations.append(best_detection)
-                    
-                    if worker_id not in confirmed_violations:
-                        confirmed_violations[worker_id] = set()
-                    confirmed_violations[worker_id].add(violation_type)
-                    
-                    if not snapshot_taken[violation_type]:
-                        cap = cv2.VideoCapture(video_path)
-                        cap.set(cv2.CAP_PROP_POS_FRAMES, best_detection["frame"])
-                        ret, snapshot_frame = cap.read()
-                        if not ret:
-                            logger.error(f"Failed to capture snapshot for {violation_type} at frame {best_detection['frame']}")
-                            cap.release()
-                            continue
-                        snapshot_frame = draw_detections(snapshot_frame, [best_detection])
-                        
-                        snapshot_filename = f"{violation_type}_{best_detection['frame']}.jpg"
-                        snapshot_path = os.path.join(CONFIG["OUTPUT_DIR"], snapshot_filename)
-                        cv2.imwrite(snapshot_path, snapshot_frame)
-                        snapshots.append({
-                            "violation": violation_type,
-                            "frame": best_detection["frame"],
-                            "snapshot_path": snapshot_path,
-                            "snapshot_base64": f"{CONFIG['PUBLIC_URL_BASE']}{snapshot_filename}"
-                        })
-                        snapshot_taken[violation_type] = True
-                        logger.info(f"Snapshot taken for {violation_type} at frame {best_detection['frame']}")
-                        cap.release()
+                # Capture snapshot for this violation
+                cap.set(cv2.CAP_PROP_POS_FRAMES, best_detection["frame"])
+                ret, snapshot_frame = cap.read()
+                if ret:
+                    snapshot_frame = draw_detections(snapshot_frame, [best_detection])
+                    snapshot_filename = f"no_helmet_{best_detection['frame']}.jpg"
+                    snapshot_path = os.path.join(CONFIG["OUTPUT_DIR"], snapshot_filename)
+                    cv2.imwrite(snapshot_path, snapshot_frame)
+                    snapshots.append({
+                        "violation": "no_helmet",
+                        "frame": best_detection["frame"],
+                        "snapshot_path": snapshot_path,
+                        "snapshot_base64": f"{CONFIG['PUBLIC_URL_BASE']}{snapshot_filename}"
+                    })
 
-        video.release()
+        cap.release()
         os.remove(video_path)
-        logger.info(f"Video file {video_path} removed")
+        logger.info(f"Processing complete. {len(violations)} violations found.")
 
+        # Generate results
         if not violations:
-            logger.info("No persistent violations detected")
             yield "No violations detected in the video.", "Safety Score: 100%", "No snapshots captured.", "N/A", "N/A"
             return
 
-        score = calculate_safety_score(violations)
+        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))
+
         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)
 
         violation_table = "| Violation              | Timestamp (s) | Confidence | Worker ID |\n"
         violation_table += "|------------------------|---------------|------------|-----------|\n"
@@ -494,22 +353,19 @@ 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
 
-        snapshots_text = "No snapshots captured."
-        if snapshots:
-            violation_name_map = CONFIG["DISPLAY_NAMES"]
-            snapshots_text = "\n".join(
-                f"- Snapshot for {violation_name_map.get(s.get('violation', 'Unknown'), 'Unknown')} at frame {s.get('frame', 0)}: ![]({s.get('snapshot_base64', '')})"
-                for s in snapshots
-            )
+        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."
 
-        logger.info(f"Processing complete: {len(violations)} violations detected, score: {score}%")
         yield (
             violation_table,
             f"Safety Score: {score}%",
             snapshots_text,
-            f"Salesforce Record ID: {report_id or 'N/A'}",
-            final_pdf_url or "N/A"
+            "Salesforce integration placeholder",
+            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}", "", "", "", ""
@@ -546,5 +402,5 @@ interface = gr.Interface(
 )
 
 if __name__ == "__main__":
-    logger.info("Launching Enhanced Safety Analyzer App...")
+    logger.info("Launching Optimized Safety Analyzer App...")
     interface.launch()