Update app.py

app.py

@@ -1,3 +1,4 @@
+```python
 import os
 import cv2
 import gradio as gr
@@ -13,19 +14,18 @@ from io import BytesIO
 import base64
 import logging
 from retrying import retry
-import uuid
 
 # ==========================
 # Configuration
 # ==========================
 CONFIG = {
-    "MODEL_PATH": "yolov8n.pt",  # Lightweight model, must be trained for violations only
+    "MODEL_PATH": "yolov8_safety.pt",  # Custom-trained model for specific violations
+    "FALLBACK_MODEL_PATH": "yolov8n.pt",  # Fallback if custom model is missing
     "OUTPUT_DIR": "static/output",
     "VIOLATION_LABELS": {
         0: "no_helmet",
         1: "no_harness",
-        2: "unsafe_posture",
-        3: "unsafe_zone"  # Ignored in processing
+        2: "unsafe_posture"
     },
     "DISPLAY_NAMES": {  # Mapping for user-friendly violation names
         "no_helmet": "Missing Helmet",
@@ -33,10 +33,10 @@ CONFIG = {
         "unsafe_posture": "Unsafe Posture"
     },
     "SF_CREDENTIALS": {
-        "username": "prashanth1ai@safety.com",
-        "password": "SaiPrash461",
-        "security_token": "AP4AQnPoidIKPvSvNEfAHyoK",
-        "domain": "login"
+        "username": "your_username@safety.com",
+        "password": "your_password",
+        "security_token": "your_security_token",
+        "domain": "login"  # Use "test" for sandbox
     },
     "PUBLIC_URL_BASE": "https://huggingface.co/spaces/PrashanthB461/AI_Safety_Demo1/resolve/main/static/output/",
     "FRAME_SKIP": 15,  # Process every 15th frame
@@ -62,9 +62,14 @@ logger.info(f"Using device: {device}")
 # ==========================
 def load_model():
     try:
-        model = YOLO(CONFIG["MODEL_PATH"]).to(device)
-        logger.info(f"Model loaded: {CONFIG['MODEL_PATH']}")
-        logger.warning("Ensure yolov8n.pt is trained to detect ONLY 'no_helmet', 'no_harness', 'unsafe_posture'. Replace with custom-trained yolov8_safety.pt if unexpected classes are detected.")
+        model_path = CONFIG["MODEL_PATH"]
+        if not os.path.exists(model_path):
+            logger.warning(f"Custom model {model_path} not found. Falling back to {CONFIG['FALLBACK_MODEL_PATH']}")
+            model_path = CONFIG["FALLBACK_MODEL_PATH"]
+        model = YOLO(model_path).to(device)
+        logger.info(f"Model loaded: {model_path}")
+        if model_path == CONFIG["FALLBACK_MODEL_PATH"]:
+            logger.warning("Using fallback model. Detection accuracy may be poor. Train yolov8_safety.pt for best results.")
         return model
     except Exception as e:
         logger.error(f"Failed to load model: {e}")
@@ -179,10 +184,10 @@ def push_report_to_salesforce(violations, score, pdf_path, pdf_file):
         }
         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']}")
+            record = sf.Safety_Violation_Report__c.create(record_data)
+            logger.info(f"Created Safety_Violation_Report__c record: {record['id']}")
         except Exception as e:
-            logger.error(f"Failed to create Safety_Video_Report__c: {e}")
+            logger.error(f"Failed to create Safety_Violation_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"]
@@ -191,10 +196,10 @@ def push_report_to_salesforce(violations, score, pdf_path, 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})
+                    sf.Safety_Violation_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}")
+                    logger.error(f"Failed to update Safety_Violation_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
@@ -264,7 +269,7 @@ def process_video(video_data):
                     label = CONFIG["VIOLATION_LABELS"].get(cls, f"unknown_class_{cls}")
                     # Only process specified violations
                     if label not in ["no_helmet", "no_harness", "unsafe_posture"]:
-                        logger.info(f"Ignoring detection: {label} (cls: {cls}, conf: {conf}) - not a target violation")
+                        logger.warning(f"Unexpected detection: {label} (cls: {cls}, conf: {conf}) - ignored")
                         continue
                     # Apply confidence threshold
                     if conf < CONFIG["CONFIDENCE_THRESHOLD"]:
@@ -389,4 +394,5 @@ interface = gr.Interface(
 
 if __name__ == "__main__":
     logger.info("Launching Safety Analyzer App...")
-    interface.launch()
+    interface.launch()
+```