Update app.py

app.py

@@ -129,75 +129,6 @@ def generate_violation_pdf(violations, score):
         logger.error(f"Error generating PDF: {e}")
         return "", "", None
 
-@retry(stop_max_attempt_number=3, wait_fixed=2000)
-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 ""
-
-@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()
-        violations_text = "\n".join(
-            f"{CONFIG['DISPLAY_NAMES'].get(v['violation'], v['violation'])} at {v['timestamp']:.2f}s (Confidence: {v['confidence']})"
-            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,
@@ -224,7 +155,7 @@ def process_video(video_data):
         fps = video.get(cv2.CAP_PROP_FPS)
 
         snapshot_taken = {label: False for label in CONFIG["VIOLATION_LABELS"].values()}
-        seen_violations = set()  # to track violations and avoid repeating
+        detected_violations = {}
 
         while True:
             ret, frame = video.read()
@@ -244,21 +175,21 @@ 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 CONFIG["VIOLATION_LABELS"].values():
-                        continue
-                    if conf < CONFIG["CONFIDENCE_THRESHOLD"]:
+                    if label not in CONFIG["VIOLATION_LABELS"].values() or conf < CONFIG["CONFIDENCE_THRESHOLD"]:
                         continue
-                    if label in seen_violations:
+
+                    violation_key = f"{label}_{frame_count}"
+                    if violation_key in detected_violations:
                         continue
-                    seen_violations.add(label)
 
-                    violation = {
+                    detected_violations[violation_key] = {
                         "frame": frame_count,
                         "violation": label,
                         "confidence": round(conf, 2),
+                        "bounding_box": [round(x, 2) for x in box.xywh.cpu().numpy()[0]],
                         "timestamp": frame_count / fps
                     }
-                    violations.append(violation)
+                    violations.append(detected_violations[violation_key])
 
                     if not snapshot_taken[label]:
                         snapshot_path = os.path.join(CONFIG["OUTPUT_DIR"], f"snapshot_{frame_count}_{label}.jpg")
@@ -329,13 +260,13 @@ def gradio_interface(video_file):
 
         violation_table = "No violations detected."
         if result["violations"]:
-            header = "| Violation              | Timestamp (s) | Confidence | \n"
-            separator = "|------------------------|---------------|------------|\n"
+            header = "| Violation              | Timestamp (s) | Confidence | Bounding Box             |\n"
+            separator = "|------------------------|---------------|------------|--------------------------|\n"
             rows = []
             violation_name_map = CONFIG["DISPLAY_NAMES"]
             for v in result["violations"]:
                 display_name = violation_name_map.get(v["violation"], v["violation"])
-                row = f"| {display_name:<22} | {v['timestamp']:.2f}        | {v['confidence']:.2f}      |"
+                row = f"| {display_name:<22} | {v['timestamp']:.2f}        | {v['confidence']:.2f}      | {v['bounding_box']} |"
                 rows.append(row)
             violation_table = header + separator + "\n".join(rows)