Update app.py

app.py

@@ -10,6 +10,7 @@ from reportlab.pdfgen import canvas
 from reportlab.lib.utils import ImageReader
 import base64
 from PIL import Image
+from simple_salesforce import Salesforce
 
 # ==========================
 # Configuration
@@ -44,6 +45,112 @@ except Exception as e:
     print(f"❌ Failed to load model: {e}")
     raise
 
+# ==========================
+# Salesforce connection
+# ==========================
+def connect_to_salesforce():
+    sf = Salesforce(
+        username=os.getenv("prashanth1ai@safety.com"),
+        password=os.getenv("SaiPrash461"),
+        security_token=os.getenv("AP4AQnPoidIKPvSvNEfAHyoK"),
+        domain=os.getenv("SF_DOMAIN", "login")
+    )
+    return sf
+
+def violations_to_text(violations):
+    lines = []
+    for v in violations:
+        line = f"{v['violation']} at {v['timestamp']:.2f}s (Confidence: {v['confidence']})"
+        lines.append(line)
+    return "\n".join(lines)
+
+def push_report_to_salesforce(score, violations, pdf_filename):
+    sf = connect_to_salesforce()
+
+    violations_count = len(violations)
+    violations_details = violations_to_text(violations)
+
+    # Public URL base for your PDF reports on HuggingFace space or other hosting
+    PUBLIC_PDF_BASE_URL = "https://huggingface.co/spaces/PrashanthB461/AI_Safety_Demo1/resolve/main/output/"
+
+    pdf_url = f"{PUBLIC_PDF_BASE_URL}{os.path.basename(pdf_filename)}"
+
+    record = sf.Safety_Video_Report__c.create({
+        "Compliance_Score__c": score,
+        "Violations_Found__c": violations_count,
+        "Violations_Details__c": violations_details,
+        "Status__c": "Pending",
+        "PDF_Report_URL__c": pdf_url
+    })
+
+    return record.get("id")
+
+# ==========================
+# Safety Score Calculation
+# ==========================
+def calculate_safety_score(violations):
+    base_score = 100
+    penalties = {
+        "no_helmet": 25,
+        "no_harness": 30,
+        "unsafe_posture": 20,
+        "unsafe_zone": 25
+    }
+    for v in violations:
+        base_score -= penalties.get(v["violation"], 0)
+    return max(base_score, 0)
+
+# ==========================
+# PDF Report Generation
+# ==========================
+def generate_pdf_report(violations, snapshots, score):
+    try:
+        pdf_path = os.path.join(OUTPUT_DIR, f"report_{int(time.time())}.pdf")
+        c = canvas.Canvas(pdf_path, pagesize=letter)
+        width, height = letter
+
+        # Title
+        c.setFont("Helvetica-Bold", 16)
+        c.drawString(50, height - 50, "Worksite Safety Compliance Report")
+
+        # Compliance Score
+        c.setFont("Helvetica", 12)
+        c.drawString(50, height - 80, f"Compliance Score: {score}%")
+
+        # Violations Table
+        y = height - 120
+        c.setFont("Helvetica-Bold", 12)
+        c.drawString(50, y, "Detected Violations:")
+        y -= 20
+
+        for v in violations:
+            c.setFont("Helvetica", 10)
+            text = f"Violation: {v['violation']}, Timestamp: {v['timestamp']:.2f}s, Confidence: {v['confidence']}"
+            c.drawString(50, y, text)
+            y -= 20
+
+            # Add snapshot if available and exists locally
+            snapshot = next((s for s in snapshots if s["frame"] == v["frame"] and s["violation"] == v["violation"]), None)
+            if snapshot:
+                local_img_path = os.path.join(OUTPUT_DIR, os.path.basename(snapshot["snapshot_url"]))
+                if os.path.exists(local_img_path):
+                    img = ImageReader(local_img_path)
+                    c.drawImage(img, 50, y - 100, width=200, height=150)
+                    y -= 170
+
+            if y < 50:
+                c.showPage()
+                y = height - 50
+
+        c.save()
+        print(f"PDF generated at {pdf_path}")
+
+        return pdf_path
+
+    except Exception as e:
+        print(f"❌ Error generating PDF: {e}")
+        return ""
+
 # ==========================
 # Video Processing
 # ==========================
@@ -100,7 +207,7 @@ def process_video(video_data, frame_skip=5, max_frames=100):
                     snapshots.append({
                         "violation": label,
                         "frame": frame_count,
-                        "snapshot_url": f"https://huggingface.co/spaces/PrashanthB461/AI_Safety_Demo1/output/{os.path.basename(snapshot_path)}"
+                        "snapshot_url": f"https://huggingface.co/spaces/PrashanthB461/AI_Safety_Demo1/resolve/main/output/{os.path.basename(snapshot_path)}"
                     })
 
             frame_count += 1
@@ -117,13 +224,25 @@ def process_video(video_data, frame_skip=5, max_frames=100):
         os.remove(video_path)
 
         score = calculate_safety_score(violations)
-        pdf_base64 = generate_pdf_report(violations, snapshots, score)
+        pdf_path = generate_pdf_report(violations, snapshots, score)
+
+        # Push report to Salesforce with public PDF URL
+        if pdf_path:
+            report_id = push_report_to_salesforce(score, violations, pdf_path)
+            print(f"Salesforce record created with Id: {report_id}")
+        else:
+            report_id = None
+
+        # Convert PDF to base64 for UI display (optional)
+        with open(pdf_path, "rb") as f:
+            pdf_base64 = base64.b64encode(f.read()).decode('utf-8')
 
         return {
             "violations": violations,
             "snapshots": snapshots,
             "score": score,
-            "pdf_base64": pdf_base64
+            "pdf_base64": pdf_base64,
+            "salesforce_record_id": report_id
         }
 
     except Exception as e:
@@ -133,100 +252,32 @@ def process_video(video_data, frame_skip=5, max_frames=100):
             "snapshots": [],
             "score": 0,
             "pdf_base64": "",
+            "salesforce_record_id": None,
             "error": str(e)
         }
 
-# ==========================
-# Safety Score Calculation
-# ==========================
-def calculate_safety_score(violations):
-    base_score = 100
-    penalties = {
-        "no_helmet": 25,
-        "no_harness": 30,
-        "unsafe_posture": 20,
-        "unsafe_zone": 25
-    }
-    for v in violations:
-        base_score -= penalties.get(v["violation"], 0)
-    return max(base_score, 0)
-
-# ==========================
-# PDF Report Generation
-# ==========================
-def generate_pdf_report(violations, snapshots, score):
-    try:
-        pdf_path = os.path.join(OUTPUT_DIR, f"report_{int(time.time())}.pdf")
-        c = canvas.Canvas(pdf_path, pagesize=letter)
-        width, height = letter
-
-        # Title
-        c.setFont("Helvetica-Bold", 16)
-        c.drawString(50, height - 50, "Worksite Safety Compliance Report")
-
-        # Compliance Score
-        c.setFont("Helvetica", 12)
-        c.drawString(50, height - 80, f"Compliance Score: {score}%")
-
-        # Violations Table
-        y = height - 120
-        c.setFont("Helvetica-Bold", 12)
-        c.drawString(50, y, "Detected Violations:")
-        y -= 20
-
-        for v in violations:
-            c.setFont("Helvetica", 10)
-            text = f"Violation: {v['violation']}, Timestamp: {v['timestamp']:.2f}s, Confidence: {v['confidence']}"
-            c.drawString(50, y, text)
-            y -= 20
-
-            # Add snapshot if available
-            snapshot = next((s for s in snapshots if s["frame"] == v["frame"] and s["violation"] == v["violation"]), None)
-            if snapshot and os.path.exists(snapshot["snapshot_url"].split('/')[-1]):
-                img = ImageReader(snapshot["snapshot_url"].split('/')[-1])
-                c.drawImage(img, 50, y - 100, width=200, height=150)
-                y -= 170
-
-            if y < 50:
-                c.showPage()
-                y = height - 50
-
-        c.save()
-        print(f"PDF generated at {pdf_path}")
-
-        # Convert PDF to base64
-        with open(pdf_path, "rb") as f:
-            pdf_base64 = base64.b64encode(f.read()).decode('utf-8')
-        
-        # Clean up
-        os.remove(pdf_path)
-        print("PDF converted to base64 and file removed")
-        return pdf_base64
-    except Exception as e:
-        print(f"❌ Error generating PDF: {e}")
-        return ""
-
 # ==========================
 # Gradio Interface
 # ==========================
 def gradio_interface(video_file):
     try:
         if not video_file:
-            return {"error": "Please upload a video file."}, "", "", []
+            return {"error": "Please upload a video file."}, "", "", [], ""
 
         with open(video_file, "rb") as f:
             video_data = f.read()
 
         result = process_video(video_data)
         return (
-            result["violations"],
-            f"Safety Score: {result['score']}%",
-            result["pdf_base64"],
-            result["snapshots"]
+            result.get("violations", []),
+            f"Safety Score: {result.get('score', 0)}%",
+            result.get("pdf_base64", ""),
+            result.get("snapshots", []),
+            f"Salesforce Record ID: {result.get('salesforce_record_id', '')}"
         )
     except Exception as e:
         print(f"❌ Error in gradio_interface: {e}")
-        return {"error": str(e)}, "", "", []
+        return {"error": str(e)}, "", "", [], ""
 
 interface = gr.Interface(
     fn=gradio_interface,
@@ -235,7 +286,8 @@ interface = gr.Interface(
         gr.JSON(label="Detected Safety Violations"),
         gr.Textbox(label="Compliance Score"),
         gr.Textbox(label="PDF Base64 (for API use)"),
-        gr.JSON(label="Snapshots")
+        gr.JSON(label="Snapshots"),
+        gr.Textbox(label="Salesforce Record ID")
     ],
     title="Worksite Safety Violation Analyzer",
     description="Upload short site videos to detect safety violations (e.g., no helmet, no harness, unsafe posture)."
@@ -243,4 +295,4 @@ interface = gr.Interface(
 
 if __name__ == "__main__":
     print("🚀 Launching Safety Analyzer App...")
-    interface.launch()
+    interface.launch()