Update app.py

app.py

@@ -16,7 +16,6 @@ from io import BytesIO
 import base64
 from retrying import retry
 from collections import defaultdict
-from multiprocessing import cpu_count
 
 # ========================== # Configuration and Setup # ==========================
 os.environ['YOLO_CONFIG_DIR'] = '/tmp/Ultralytics'
@@ -26,7 +25,7 @@ logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(
 logger = logging.getLogger(__name__)
 warnings.filterwarnings("ignore")
 
-# ========================== # Optimized Tracker Implementation (No Face Recognition) # ==========================
+# ========================== # Position-Based Tracker (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
@@ -34,9 +33,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)
+        
+        self.worker_tracks = {}  # Active worker tracks
+        self.violation_history = defaultdict(dict)  # Track violations per worker
+        self.position_history = defaultdict(list)  # Track positions for all violations
         
         self.VIOLATION_COOLDOWNS = {
             "no_helmet": 30.0,
@@ -46,7 +46,7 @@ class SafetyTracker:
             "improper_tool_use": 15.0
         }
 
-    def update(self, detections, frame):
+    def update(self, detections):
         current_time = time.time()
         new_violations = []
         
@@ -56,6 +56,7 @@ class SafetyTracker:
             confidence = det['confidence']
             
             worker_id = self._match_by_position(bbox, label)
+            
             if worker_id is None:
                 worker_id = self.next_id
                 self.next_id += 1
@@ -82,30 +83,40 @@ class SafetyTracker:
         return new_violations
 
     def _match_by_position(self, bbox, label):
-        x, y = bbox[0], bbox[1]
+        x, y, w, h = bbox
+        current_pos = (x, y)
+        
         for worker_id, positions in self.position_history.items():
-            for pos in positions[-5:]:  # Check last 5 positions
-                if np.sqrt((x-pos[0])**2 + (y-pos[1])**2) < 100:
-                    return worker_id
+            if not positions:
+                continue
+                
+            last_pos = positions[-1]
+            distance = np.sqrt((current_pos[0]-last_pos[0])**2 + (current_pos[1]-last_pos[1])**2)
+            if distance < 100:  # Within 100 pixels
+                return worker_id
         return None
 
     def _is_new_violation(self, worker_id, label, current_time):
         if label not in self.violation_history[worker_id]:
             return True
-        return (current_time - self.violation_history[worker_id][label]) > self.VIOLATION_COOLDOWNS.get(label, 10.0)
+            
+        last_detection = self.violation_history[worker_id][label]
+        cooldown = self.VIOLATION_COOLDOWNS.get(label, 10.0)
+        return (current_time - last_detection) > cooldown
 
     def _cleanup_tracks(self, current_time):
         inactive_ids = [
-            wid for wid, track in self.worker_tracks.items()
+            worker_id for worker_id, track in self.worker_tracks.items()
             if (current_time - track['last_seen']) > (self.track_buffer / self.frame_rate)
         ]
-        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)
+        
+        for worker_id in inactive_ids:
+            self.worker_tracks.pop(worker_id, None)
+            self.position_history.pop(worker_id, None)
+            if (current_time - max(self.violation_history[worker_id].values(), default=0)) > 300:
+                self.violation_history.pop(worker_id, None)
 
-# ========================== # Configuration # ==========================
+# ========================== # Optimized Configuration # ==========================
 CONFIG = {
     "MODEL_PATH": "yolov8_safety.pt",
     "FALLBACK_MODEL": "yolov8n.pt",
@@ -125,10 +136,10 @@ CONFIG = {
         "improper_tool_use": (255, 255, 0)
     },
     "DISPLAY_NAMES": {
-        "no_helmet": "No Helmet",
-        "no_harness": "No Harness",
+        "no_helmet": "No Helmet Violation",
+        "no_harness": "No Harness Violation",
         "unsafe_posture": "Unsafe Posture",
-        "unsafe_zone": "Unsafe Zone",
+        "unsafe_zone": "Unsafe Zone Entry",
         "improper_tool_use": "Improper Tool Use"
     },
     "SF_CREDENTIALS": {
@@ -153,35 +164,53 @@ CONFIG = {
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 logger.info(f"Using device: {device}")
 
-# ========================== # Core Functions # ==========================
 def load_model():
     try:
-        model_path = CONFIG["MODEL_PATH"] if os.path.exists(CONFIG["MODEL_PATH"]) else CONFIG["FALLBACK_MODEL"]
-        if not os.path.exists(model_path):
-            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)
+        if os.path.isfile(CONFIG["MODEL_PATH"]):
+            model = YOLO(CONFIG["MODEL_PATH"]).to(device)
+            logger.info(f"Loaded custom model: {CONFIG['MODEL_PATH']}")
+        else:
+            model = YOLO(CONFIG["FALLBACK_MODEL"]).to(device)
+            logger.warning("Using fallback YOLOv8n model")
+        return model
     except Exception as e:
         logger.error(f"Model loading failed: {e}")
         raise
 
 model = load_model()
 
+# ========================== # Core Functions # ==========================
+def preprocess_frame(frame):
+    frame = cv2.convertScaleAbs(frame, alpha=1.2, beta=20)
+    return frame
+
 def draw_detections(frame, detections):
-    annotated = frame.copy()
+    result_frame = 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)
-        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
+        label = det["violation"]
+        confidence = det["confidence"]
+        x, y, w, h = det["bbox"]
+        worker_id = det["worker_id"]
+
+        x1, y1 = int(x - w/2), int(y - h/2)
+        x2, y2 = int(x + w/2), int(y + h/2)
+        color = CONFIG["CLASS_COLORS"][label]
+        
+        cv2.rectangle(result_frame, (x1, y1), (x2, y2), color, 3)
+        text = f"{CONFIG['DISPLAY_NAMES'][label]} (Worker {worker_id})"
+        (tw, th), _ = cv2.getTextSize(text, cv2.FONT_HERSHEY_SIMPLEX, 0.6, 2)
+        cv2.rectangle(result_frame, (x1, y1-th-10), (x1+tw+10, y1), (0,0,0), -1)
+        cv2.putText(result_frame, 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}
-    unique_violations = {v['violation'] for v in violations}
-    return max(0, 100 - sum(penalties.get(v,0) for v in unique_violations))
+    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(penalties.get(v, 0) for v in unique_violations))
 
 def generate_violation_pdf(violations, score):
     try:
@@ -190,7 +219,7 @@ def generate_violation_pdf(violations, score):
         
         # Header
         c.setFont("Helvetica-Bold", 16)
-        c.drawString(1*inch, 10*inch, "Safety Violation Report")
+        c.drawString(1*inch, 10*inch, "Worksite Safety Violation Report")
         c.setFont("Helvetica", 12)
         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}%")
@@ -198,28 +227,29 @@ def generate_violation_pdf(violations, score):
         # Violations List
         y = 8.5*inch
         c.setFont("Helvetica-Bold", 14)
-        c.drawString(1*inch, y, "Violations Detected:")
+        c.drawString(1*inch, y, "Detected Violations:")
         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)
-            y -= 0.2*inch
+            text = (f"Worker {v['worker_id']}: {CONFIG['DISPLAY_NAMES'][v['violation']} "
+                   f"at {v['timestamp']:.2f}s (Confidence: {v['confidence']:.2f})")
             if y < 1*inch:
                 c.showPage()
                 y = 10*inch
-                c.setFont("Helvetica", 10)
+            c.drawString(1.2*inch, y, text)
+            y -= 0.2*inch
         
         c.save()
         pdf_buffer.seek(0)
         
         # Save to file
-        pdf_path = os.path.join(CONFIG["OUTPUT_DIR"], f"report_{int(time.time())}.pdf")
+        pdf_filename = f"violation_report_{int(time.time())}.pdf"
+        pdf_path = os.path.join(CONFIG["OUTPUT_DIR"], pdf_filename)
         with open(pdf_path, "wb") as f:
             f.write(pdf_buffer.getvalue())
             
-        return pdf_path, f"{CONFIG['PUBLIC_URL_BASE']}{os.path.basename(pdf_path)}", pdf_buffer
+        return pdf_path, f"{CONFIG['PUBLIC_URL_BASE']}{pdf_filename}", pdf_buffer
     except Exception as e:
         logger.error(f"PDF generation failed: {e}")
         return None, None, None
@@ -228,119 +258,173 @@ def generate_violation_pdf(violations, score):
 def connect_to_salesforce():
     try:
         sf = Salesforce(**CONFIG["SF_CREDENTIALS"])
-        sf.describe()
+        logger.info("Salesforce connection established")
         return sf
     except Exception as e:
         logger.error(f"Salesforce connection failed: {e}")
         raise
 
-def upload_to_salesforce(sf, pdf_file, violations, score):
+def upload_to_salesforce(sf, pdf_file, record_id):
+    try:
+        encoded = base64.b64encode(pdf_file.getvalue()).decode('utf-8')
+        file_data = {
+            "Title": f"Safety_Report_{int(time.time())}",
+            "PathOnClient": "safety_report.pdf",
+            "VersionData": encoded,
+            "FirstPublishLocationId": record_id
+        }
+        result = sf.ContentVersion.create(file_data)
+        return f"https://{sf.sf_instance}/sfc/servlet.shepherd/version/download/{result['id']}"
+    except Exception as e:
+        logger.error(f"Salesforce upload failed: {e}")
+        return None
+
+def create_salesforce_record(violations, score, pdf_url=None):
     try:
-        # Create record
+        sf = connect_to_salesforce()
         violations_text = "\n".join(
-            f"Worker {v['worker_id']}: {CONFIG['DISPLAY_NAMES'][v['violation']]} at {v['timestamp']:.1f}s"
+            f"Worker {v['worker_id']}: {CONFIG['DISPLAY_NAMES'][v['violation']]} "
+            f"at {v['timestamp']:.2f}s (Confidence: {v['confidence']:.2f})"
             for v in violations
         )
         
-        record = sf.Safety_Video_Report__c.create({
+        record_data = {
             "Compliance_Score__c": score,
             "Violations_Found__c": len(violations),
-            "Violations_Details__c": violations_text
-        })
-        
-        # 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']
-        })
+            "Violations_Details__c": violations_text,
+            "Status__c": "Pending",
+            "PDF_Report_URL__c": pdf_url or ""
+        }
         
-        return record['id'], f"https://{sf.sf_instance}/sfc/servlet.shepherd/version/download/{content['id']}"
+        try:
+            record = sf.Safety_Video_Report__c.create(record_data)
+            return record["id"], None
+        except:
+            record = sf.Account.create({"Name": f"Safety_Report_{int(time.time())}"})
+            return record["id"], "Used Account as fallback"
     except Exception as e:
-        logger.error(f"Salesforce upload failed: {e}")
-        return None, None
+        logger.error(f"Salesforce record creation failed: {e}")
+        return None, str(e)
 
-# ========================== # Video Processing # ==========================
 def process_video(video_data):
     try:
         os.makedirs(CONFIG["OUTPUT_DIR"], exist_ok=True)
         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 = []
-        
-        while cap.isOpened():
-            ret, frame = cap.read()
-            if not ret:
+        processed_frames = 0
+        last_update = time.time()
+
+        while processed_frames < total_frames:
+            batch_frames = []
+            for _ in range(CONFIG["BATCH_SIZE"]):
+                ret, frame = cap.read()
+                if not ret:
+                    break
+                batch_frames.append(preprocess_frame(frame))
+                processed_frames += 1
+                if CONFIG["FRAME_SKIP"] > 1:
+                    for _ in range(CONFIG["FRAME_SKIP"]-1):
+                        cap.grab()
+                        processed_frames += 1
+
+            if not batch_frames:
                 break
-                
-            frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
-            results = model(frame, verbose=False)
-            
-            detections = []
-            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):
-                    detections.append({
-                        'bbox': box.xywh[0].cpu().numpy(),
-                        'violation': label,
-                        'confidence': float(box.conf)
-                    })
-            
-            new_violations = tracker.update(detections, frame)
+
+            results = model(batch_frames, device=device, conf=0.1, verbose=False)
             
-            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']}_{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
-                })
+            for i, result in enumerate(results):
+                detections = []
+                for box in result.boxes:
+                    cls = int(box.cls)
+                    conf = float(box.conf)
+                    label = CONFIG["VIOLATION_LABELS"].get(cls)
+                    if label and conf >= CONFIG["CONFIDENCE_THRESHOLDS"].get(label, 0.3):
+                        detections.append({
+                            "bbox": box.xywh.cpu().numpy()[0],
+                            "violation": label,
+                            "confidence": conf
+                        })
+
+                new_violations = tracker.update(detections)
                 
-            yield f"Processing frame {int(cap.get(cv2.CAP_PROP_POS_FRAMES))}...", "", "", "", ""
-        
+                for violation in new_violations:
+                    frame_with_det = draw_detections(batch_frames[i].copy(), [violation])
+                    timestamp = f"Time: {violation['timestamp']:.2f}s"
+                    cv2.putText(frame_with_det, timestamp, (10, 30), 
+                               cv2.FONT_HERSHEY_SIMPLEX, 0.7, (255,255,255), 2)
+                    
+                    snap_name = f"violation_{violation['violation']}_worker{violation['worker_id']}_{int(violation['timestamp']*100)}.jpg"
+                    snap_path = os.path.join(CONFIG["OUTPUT_DIR"], snap_name)
+                    cv2.imwrite(snap_path, frame_with_det, 
+                               [cv2.IMWRITE_JPEG_QUALITY, CONFIG["SNAPSHOT_QUALITY"]])
+                    
+                    snapshots.append({
+                        "violation": violation['violation'],
+                        "worker_id": violation['worker_id'],
+                        "timestamp": violation['timestamp'],
+                        "snapshot_url": f"{CONFIG['PUBLIC_URL_BASE']}{snap_name}"
+                    })
+
+            if time.time() - last_update > 1:
+                progress = (processed_frames / total_frames) * 100
+                yield f"Processing... {progress:.1f}%", "", "", "", ""
+                last_update = time.time()
+
         cap.release()
-        os.remove(video_path)
-        
-        if not snapshots:
-            yield "No violations detected", "Safety Score: 100%", "No snapshots", "N/A", "N/A"
+        if os.path.exists(video_path):
+            os.remove(video_path)
+
+        violations = [
+            {"worker_id": wid, "violation": v, "timestamp": t, "confidence": 0}  # Confidence placeholder
+            for wid, violations in tracker.violation_history.items()
+            for v, t in violations.items()
+        ]
+
+        if not violations:
+            yield "No violations found", "Safety Score: 100%", "No snapshots", "N/A", "N/A"
             return
-            
-        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)
+
+        score = calculate_safety_score(violations)
+        pdf_path, pdf_url, pdf_file = generate_violation_pdf(violations, score)
+        record_id, sf_error = create_salesforce_record(violations, score, pdf_url)
         
-        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
+        if pdf_file and record_id:
+            uploaded_url = upload_to_salesforce(connect_to_salesforce(), pdf_file, record_id)
+            if uploaded_url:
+                pdf_url = uploaded_url
+
+        violation_table = "| Violation | Worker ID | Time (s) |\n|-----------|-----------|----------|\n"
+        violation_table += "\n".join(
+            f"| {CONFIG['DISPLAY_NAMES'][v['violation']]} | {v['worker_id']} | {v['timestamp']:.2f} |"
+            for v in sorted(violations, key=lambda x: x['timestamp'])
         )
-        
+
+        snapshots_md = "\n\n".join(
+            f"### {CONFIG['DISPLAY_NAMES'][s['violation']]} - Worker {s['worker_id']} at {s['timestamp']:.2f}s\n\n"
+            f"![Snapshot]({s['snapshot_url']})"
+            for s in snapshots
+        ) if snapshots else "No snapshots captured"
+
         yield (
-            "\n".join(f"- {v['violation']['violation']} (Worker {v['violation']['worker_id']})" for v in snapshots),
+            violation_table,
             f"Safety Score: {score}%",
             snapshots_md,
-            f"Salesforce ID: {record_id or 'N/A'}",
-            sf_url or pdf_url or "N/A"
+            f"Salesforce ID: {record_id or 'N/A'} {sf_error or ''}",
+            pdf_url or "N/A"
         )
-        
+
     except Exception as e:
-        logger.error(f"Processing failed: {e}")
+        logger.error(f"Video processing failed: {e}")
         if 'video_path' in locals() and os.path.exists(video_path):
             os.remove(video_path)
         yield f"Error: {str(e)}", "", "", "", ""
@@ -350,22 +434,34 @@ def gradio_interface(video):
     if not video:
         return "Upload a video file", "", "", "", ""
     
-    for update in process_video(open(video, "rb").read()):
-        yield update
-
-interface = gr.Interface(
-    fn=gradio_interface,
-    inputs=gr.Video(),
-    outputs=[
-        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"
-)
+    try:
+        with open(video, "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)}", "", "", "", ""
+
+with gr.Blocks(title="Safety Compliance Analyzer") as app:
+    gr.Markdown("# Worksite Safety Violation Analyzer")
+    gr.Markdown("Upload site videos to detect safety violations (No Helmet, No Harness, etc.)")
+    
+    with gr.Row():
+        video_input = gr.Video(label="Site Video", sources=["upload"])
+        with gr.Column():
+            violations_out = gr.Markdown(label="Detected Violations")
+            score_out = gr.Textbox(label="Safety Score")
+            snapshots_out = gr.Markdown(label="Violation Snapshots")
+            salesforce_out = gr.Textbox(label="Salesforce Record")
+            pdf_out = gr.Textbox(label="Report PDF URL")
+    
+    video_input.change(
+        gradio_interface,
+        inputs=video_input,
+        outputs=[violations_out, score_out, snapshots_out, salesforce_out, pdf_out]
+    )
 
 if __name__ == "__main__":
-    interface.launch()
+    app.launch(server_port=7860, server_name="0.0.0.0")