Update app.py

app.py

@@ -129,6 +129,73 @@ 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['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,
@@ -242,68 +309,3 @@ def process_video(video_data):
             "violation_details_url": "",
             "message": f"Error processing video: {e}"
         }
-
-def gradio_interface(video_file):
-    if not video_file:
-        return "No file uploaded.", "", "No file uploaded.", "", ""
-    try:
-        yield "Processing video... please wait.", "", "", "", ""
-
-        with open(video_file, "rb") as f:
-            video_data = f.read()
-
-        result = process_video(video_data)
-
-        if result.get("message"):
-            yield result["message"], "", "", "", ""
-            return
-
-        violation_table = "No violations detected."
-        if result["violations"]:
-            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}      | {v['bounding_box']} |"
-                rows.append(row)
-            violation_table = header + separator + "\n".join(rows)
-
-        snapshots_text = "No snapshots captured."
-        if result["snapshots"]:
-            violation_name_map = CONFIG["DISPLAY_NAMES"]
-            snapshots_text = "\n".join(
-                f"- Snapshot for {violation_name_map.get(s['violation'], s['violation'])} at frame {s['frame']}: ![]({s['snapshot_base64']})"
-                for s in result["snapshots"]
-            )
-
-        yield (
-            violation_table,
-            f"Safety Score: {result['score']}%",
-            snapshots_text,
-            f"Salesforce Record ID: {result['salesforce_record_id'] or 'N/A'}",
-            result["violation_details_url"] or "N/A"
-        )
-    except Exception as e:
-        logger.error(f"Error in Gradio interface: {e}", exc_info=True)
-        yield f"Error: {str(e)}", "", "Error in processing.", "", ""
-
-interface = gr.Interface(
-    fn=gradio_interface,
-    inputs=gr.Video(label="Upload Site Video"),
-    outputs=[
-        gr.Markdown(label="Detected Safety Violations"),
-        gr.Textbox(label="Compliance Score"),
-        gr.Markdown(label="Snapshots"),
-        gr.Textbox(label="Salesforce Record ID"),
-        gr.Textbox(label="Violation Details URL")
-    ],
-    title="Worksite Safety Violation Analyzer",
-    description="Upload site videos to detect safety violations (No Helmet Violation, No Harness Violation, Unsafe Posture Violation). Non-violations are ignored.",
-    allow_flagging="never"
-)
-
-if __name__ == "__main__":
-    logger.info("Launching Safety Analyzer App...")
-    interface.launch()