Update app.py

app.py

@@ -19,7 +19,7 @@ from retrying import retry
 # ==========================
 CONFIG = {
     "MODEL_PATH": "yolov8_safety.pt",
-    "FALLBACK_MODEL_PATH": "yolov8n.pt",
+    "FALLBACK_MODEL": "yolov8n.pt",  # updated key to match first code
     "OUTPUT_DIR": "static/output",
     "VIOLATION_LABELS": {
         0: "no_helmet",
@@ -39,7 +39,7 @@ CONFIG = {
     },
     "PUBLIC_URL_BASE": "https://huggingface.co/spaces/PrashanthB461/AI_Safety_Demo2/resolve/main/static/output/",
     "FRAME_SKIP": 15,
-    "MAX_PROCESSING_TIME": 30,  # seconds
+    "MAX_PROCESSING_TIME": 30,
     "CONFIDENCE_THRESHOLD": 0.5
 }
 
@@ -54,14 +54,11 @@ logger.info(f"Using device: {device}")
 
 def load_model():
     try:
-        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_path = CONFIG["MODEL_PATH"] if os.path.isfile(CONFIG["MODEL_PATH"]) else CONFIG["FALLBACK_MODEL"]
+        if model_path == CONFIG["FALLBACK_MODEL"]:
+            logger.warning("Using fallback model. Detection accuracy may be poor. Train yolov8_safety.pt for best results.")
         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}")
@@ -69,12 +66,15 @@ def load_model():
 
 model = load_model()
 
-@retry(stop_max_attempt_number=2, wait_fixed=1000)
+# ==========================
+# 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()
+        sf.describe()  # verify connection and metadata fetch
         return sf
     except Exception as e:
         logger.error(f"Salesforce connection failed: {e}")
@@ -129,7 +129,7 @@ def generate_violation_pdf(violations, score):
         logger.error(f"Error generating PDF: {e}")
         return "", "", None
 
-@retry(stop_max_attempt_number=2, wait_fixed=1000)
+@retry(stop_max_attempt_number=3, wait_fixed=2000)
 def upload_pdf_to_salesforce(sf, pdf_file, report_id):
     try:
         if not pdf_file:
@@ -143,7 +143,7 @@ def upload_pdf_to_salesforce(sf, pdf_file, report_id):
             "FirstPublishLocationId": report_id
         }
         content_version = sf.ContentVersion.create(content_version_data)
-        result = sf.query(f"SELECT Id FROM ContentVersion WHERE Id = '{content_version['id']}'")
+        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 ""
@@ -154,7 +154,7 @@ def upload_pdf_to_salesforce(sf, pdf_file, report_id):
         logger.error(f"Error uploading PDF to Salesforce: {e}")
         return ""
 
-@retry(stop_max_attempt_number=2, wait_fixed=1000)
+@retry(stop_max_attempt_number=3, wait_fixed=2000)
 def push_report_to_salesforce(violations, score, pdf_path, pdf_file):
     try:
         sf = connect_to_salesforce()
@@ -173,10 +173,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_Violation_Report__c.create(record_data)
-            logger.info(f"Created Safety_Violation_Report__c record: {record['id']}")
+            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_Violation_Report__c: {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"]
@@ -185,17 +185,17 @@ 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_Violation_Report__c.update(record_id, {"PDF_Report_URL__c": uploaded_url})
+                    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_Violation_Report__c: {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}")
+        logger.error(f"Salesforce record creation failed: {e}", exc_info=True)
         return None, ""
 
 def calculate_safety_score(violations):
@@ -204,10 +204,7 @@ def calculate_safety_score(violations):
         "no_harness": 30,
         "unsafe_posture": 20
     }
-    score = 100
-    for v in violations:
-        if v["violation"] in penalties:
-            score -= penalties[v["violation"]]
+    score = 100 - sum(penalties.get(v["violation"], 0) for v in violations)
     return max(score, 0)
 
 def process_video(video_data):
@@ -226,19 +223,19 @@ def process_video(video_data):
         start_time = time.time()
         fps = video.get(cv2.CAP_PROP_FPS)
 
-        snapshot_taken = {"no_helmet": False, "no_harness": False, "unsafe_posture": False}
+        snapshot_taken = {label: False for label in CONFIG["VIOLATION_LABELS"].values()}
 
         while True:
             ret, frame = video.read()
             if not ret:
-                break  # End of video
+                break
 
             if frame_count % CONFIG["FRAME_SKIP"] != 0:
                 frame_count += 1
                 continue
 
             if time.time() - start_time > CONFIG["MAX_PROCESSING_TIME"]:
-                logger.info("Processing time limit of 30 seconds reached")
+                logger.info("Processing time limit reached")
                 break
 
             results = model(frame, device=device)
@@ -247,7 +244,7 @@ def process_video(video_data):
                 for box in result.boxes:
                     cls, conf = int(box.cls), float(box.conf)
                     label = CONFIG["VIOLATION_LABELS"].get(cls, None)
-                    if label not in ["no_helmet", "no_harness", "unsafe_posture"]:
+                    if label not in CONFIG["VIOLATION_LABELS"].values():
                         continue
                     if conf < CONFIG["CONFIDENCE_THRESHOLD"]:
                         continue
@@ -306,7 +303,7 @@ def process_video(video_data):
             "message": ""
         }
     except Exception as e:
-        logger.error(f"Error processing video: {e}")
+        logger.error(f"Error processing video: {e}", exc_info=True)
         return {
             "violations": [],
             "snapshots": [],
@@ -320,7 +317,6 @@ def gradio_interface(video_file):
     if not video_file:
         return "No file uploaded.", "", "No file uploaded.", "", ""
     try:
-        # Show processing message early in the UI
         yield "Processing video... please wait.", "", "", "", ""
 
         with open(video_file, "rb") as f:
@@ -329,7 +325,6 @@ def gradio_interface(video_file):
         result = process_video(video_data)
 
         if result.get("message"):
-            # If message present (either no violations or error), show it plainly
             yield result["message"], "", "", "", ""
             return
 
@@ -361,7 +356,7 @@ def gradio_interface(video_file):
             result["violation_details_url"] or "N/A"
         )
     except Exception as e:
-        logger.error(f"Error in Gradio interface: {e}")
+        logger.error(f"Error in Gradio interface: {e}", exc_info=True)
         yield f"Error: {str(e)}", "", "Error in processing.", "", ""
 
 interface = gr.Interface(