Update app.py

app.py

@@ -5,11 +5,7 @@ import torch
 import numpy as np
 from ultralytics import YOLO
 import time
-from reportlab.lib.pagesizes import letter
-from reportlab.pdfgen import canvas
-from reportlab.lib.utils import ImageReader
 import base64
-from PIL import Image
 from simple_salesforce import Salesforce
 
 # ==========================
@@ -51,13 +47,18 @@ except Exception as e:
 # Salesforce connection
 # ==========================
 def connect_to_salesforce():
-    sf = Salesforce(
-        username="prashanth1ai@safety.com",
-        password="SaiPrash461",
-        security_token="AP4AQnPoidIKPvSvNEfAHyoK",
-        domain="login"
-    )
-    return sf
+    try:
+        sf = Salesforce(
+            username="prashanth1ai@safety.com",
+            password="SaiPrash461",
+            security_token="AP4AQnPoidIKPvSvNEfAHyoK",
+            domain="login"
+        )
+        print("✅ Connected to Salesforce")
+        return sf
+    except Exception as e:
+        print(f"❌ Salesforce connection failed: {e}")
+        raise
 
 def violations_to_text(violations):
     lines = []
@@ -74,31 +75,33 @@ def generate_violation_text_file(violations):
             f.write("Worksite Safety Violation Details\n")
             f.write("=" * 40 + "\n")
             f.write(violations_to_text(violations))
-        print(f"Text file generated at {text_path}")
-        return text_path
+        print(f"✅ Text file generated at {text_path}")
+        public_url = f"https://huggingface.co/spaces/PrashanthB461/AI_Safety_Demo1/resolve/main/static/output/{text_filename}"
+        print(f"✅ Text file URL: {public_url}")
+        return text_path, public_url
     except Exception as e:
         print(f"❌ Error generating text file: {e}")
-        return ""
-
-def push_report_to_salesforce(score, violations, pdf_filename, text_filename):
-    sf = connect_to_salesforce()
+        return "", ""
 
-    violations_count = len(violations)
-    violations_details = violations_to_text(violations)
+def push_report_to_salesforce(score, violations, text_filename):
+    try:
+        sf = connect_to_salesforce()
+        violations_count = len(violations)
+        violations_details = violations_to_text(violations)
 
-    PUBLIC_BASE_URL = "https://huggingface.co/spaces/PrashanthB461/AI_Safety_Demo1/resolve/main/static/output/"
-    pdf_url = f"{PUBLIC_BASE_URL}{os.path.basename(pdf_filename)}"
-    text_url = f"{PUBLIC_BASE_URL}{os.path.basename(text_filename)}" if text_filename else ""
+        PUBLIC_BASE_URL = "https://huggingface.co/spaces/PrashanthB461/AI_Safety_Demo1/resolve/main/static/output/"
+        text_url = f"{PUBLIC_BASE_URL}{os.path.basename(text_filename)}"
 
-    try:
-        record = sf.Safety_Video_Report__c.create({
+        record_data = {
             "Compliance_Score__c": score,
             "Violations_Found__c": violations_count,
             "Violations_Details__c": violations_details,
             "Status__c": "Pending",
-            "PDF_Report_URL__c": pdf_url,
-            "Violation_Details_URL__c": text_url  # New field for text file URL
-        })
+            "PDF_Report_URL__c": text_url  # Store text file URL in PDF_Report_URL__c
+        }
+
+        print(f"Attempting to create Salesforce record with data: {record_data}")
+        record = sf.Safety_Video_Report__c.create(record_data)
         print(f"✅ Salesforce record created with Id: {record.get('id')}")
         return record.get("id")
     except Exception as e:
@@ -120,51 +123,6 @@ def calculate_safety_score(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(STATIC_OUTPUT_DIR, f"report_{int(time.time())}.pdf")
-        c = canvas.Canvas(pdf_path, pagesize=letter)
-        width, height = letter
-
-        c.setFont("Helvetica-Bold", 16)
-        c.drawString(50, height - 50, "Worksite Safety Compliance Report")
-
-        c.setFont("Helvetica", 12)
-        c.drawString(50, height - 80, f"Compliance Score: {score}%")
-
-        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
-
-            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(STATIC_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
 # ==========================
@@ -174,7 +132,7 @@ def process_video(video_data, frame_skip=5, max_frames=100):
         video_path = os.path.join(STATIC_OUTPUT_DIR, f"temp_{int(time.time())}.mp4")
         with open(video_path, "wb") as f:
             f.write(video_data)
-        print(f"Video saved to {video_path}")
+        print(f"✅ Video saved to {video_path}")
 
         video = cv2.VideoCapture(video_path)
         if not video.isOpened():
@@ -235,24 +193,21 @@ def process_video(video_data, frame_skip=5, max_frames=100):
         os.remove(video_path)
 
         score = calculate_safety_score(violations)
-        pdf_path = generate_pdf_report(violations, snapshots, score)
-        text_path = generate_violation_text_file(violations)
+        text_path, text_url = generate_violation_text_file(violations)
 
-        if pdf_path and text_path:
-            report_id = push_report_to_salesforce(score, violations, pdf_path, text_path)
-            print(f"Salesforce record created with Id: {report_id}")
+        if text_path:
+            report_id = push_report_to_salesforce(score, violations, text_path)
+            print(f"✅ Salesforce record created with Id: {report_id}")
         else:
             report_id = None
-
-        with open(pdf_path, "rb") as f:
-            pdf_base64 = base64.b64encode(f.read()).decode('utf-8')
+            print("❌ No text file generated, skipping Salesforce record creation")
 
         return {
             "violations": violations,
             "snapshots": snapshots,
             "score": score,
-            "pdf_base64": pdf_base64,
-            "salesforce_record_id": report_id
+            "salesforce_record_id": report_id,
+            "violation_details_url": text_url
         }
 
     except Exception as e:
@@ -261,8 +216,8 @@ def process_video(video_data, frame_skip=5, max_frames=100):
             "violations": [],
             "snapshots": [],
             "score": 0,
-            "pdf_base64": "",
             "salesforce_record_id": None,
+            "violation_details_url": "",
             "error": str(e)
         }
 
@@ -272,7 +227,7 @@ def process_video(video_data, frame_skip=5, max_frames=100):
 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()
@@ -281,13 +236,13 @@ def gradio_interface(video_file):
         return (
             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', '')}"
+            f"Salesforce Record ID: {result.get('salesforce_record_id', '')}",
+            f"Violation Details URL: {result.get('violation_details_url', '')}"
         )
     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,
@@ -295,9 +250,9 @@ interface = gr.Interface(
     outputs=[
         gr.JSON(label="Detected Safety Violations"),
         gr.Textbox(label="Compliance Score"),
-        gr.Textbox(label="PDF Base64 (for API use)"),
         gr.JSON(label="Snapshots"),
-        gr.Textbox(label="Salesforce Record ID")
+        gr.Textbox(label="Salesforce Record ID"),
+        gr.Textbox(label="Violation Details URL")
     ],
     title="Worksite Safety Violation Analyzer",
     description="Upload short site videos to detect safety violations (e.g., no helmet, no harness, unsafe posture)."