Update app.py

app.py

@@ -10,7 +10,7 @@ from reportlab.pdfgen import canvas
 from reportlab.lib.utils import ImageReader
 import base64
 from PIL import Image
-from simple_salesforce import Salesforce, SalesforceMalformedRequest
+from simple_salesforce import Salesforce
 
 # ==========================
 # Configuration
@@ -19,12 +19,10 @@ DEFAULT_MODEL_PATH = "models/yolov8_safety.pt"
 FALLBACK_MODEL = "yolov8n.pt"
 MODEL_PATH = os.getenv("SAFETY_MODEL_PATH", DEFAULT_MODEL_PATH)
 
+# Use static/output for publicly accessible files
 STATIC_OUTPUT_DIR = "static/output"
 os.makedirs(STATIC_OUTPUT_DIR, exist_ok=True)
 
-# Base URL for publicly accessible files (Hugging Face Spaces)
-BASE_PUBLIC_URL = "https://huggingface.co/spaces/PrashanthB461/AI_Safety_Demo1/raw/main/static/output/"
-
 VIOLATION_LABELS = {
     0: "no_helmet",
     1: "no_harness",
@@ -50,59 +48,44 @@ except Exception as e:
     raise
 
 # ==========================
-# Salesforce Connection
+# Salesforce connection (hardcoded creds for testing)
 # ==========================
 def connect_to_salesforce():
-    try:
-        sf = Salesforce(
-            username="prashanth1ai@safety.com",
-            password="SaiPrash461",
-            security_token="AP4AQnPoidIKPvSvNEfAHyoK",
-            domain="login"
-        )
-        print("✅ Salesforce connection successful.")
-        return sf
-    except Exception as e:
-        print(f"❌ Salesforce connection failed: {e}")
-        raise
+    sf = Salesforce(
+        username="prashanth1ai@safety.com",
+        password="SaiPrash461",
+        security_token="AP4AQnPoidIKPvSvNEfAHyoK",
+        domain="login"
+    )
+    return sf
 
-# ==========================
-# Violation Format Update
-# ==========================
 def violations_to_text(violations):
     lines = []
     for v in violations:
-        line = f"Violation: {v['violation']} | Timestamp: {v['timestamp']:.2f}s | Confidence: {v['confidence']}"
+        line = f"{v['violation']} at {v['timestamp']:.2f}s (Confidence: {v['confidence']})"
         lines.append(line)
     return "\n".join(lines)
 
-# ==========================
-# Salesforce Record Creation
-# ==========================
-def push_report_to_salesforce(score, violations, pdf_url):
-    try:
-        sf = connect_to_salesforce()
-        violations_count = len(violations)
-        violations_details = violations_to_text(violations)
+def push_report_to_salesforce(score, violations, pdf_filename):
+    sf = connect_to_salesforce()
+
+    violations_count = len(violations)
+    violations_details = violations_to_text(violations)
+
+    # Updated public URL base with static/output path
+    PUBLIC_PDF_BASE_URL = "https://huggingface.co/spaces/PrashanthB461/AI_Safety_Demo1/resolve/main/static/output/"
+    pdf_url = f"{PUBLIC_PDF_BASE_URL}{os.path.basename(pdf_filename)}"
 
-        # Create the record in the custom object Safety_Video_Report__c
-        print(f"Creating Salesforce record with compliance score: {score} and violations found: {violations_count}")
+    try:
         record = sf.Safety_Video_Report__c.create({
             "Compliance_Score__c": score,
             "Violations_Found__c": violations_count,
-            "Violfilepath: app.py
             "Violations_Details__c": violations_details,
             "Status__c": "Pending",
             "PDF_Report_URL__c": pdf_url
         })
-        record_id = record.get('id')
-        print(f"✅ Salesforce record created with Id: {record_id}")
-        return record_id
-
-    except SalesforceMalformedRequest as e:
-        print(f"❌ Salesforce malformed request: {e}")
-        print(f"Error details: {e.content}")
-        return None
+        print(f"✅ Salesforce record created with Id: {record.get('id')}")
+        return record.get("id")
     except Exception as e:
         print(f"❌ Salesforce record creation failed: {e}")
         return None
@@ -127,8 +110,7 @@ def calculate_safety_score(violations):
 # ==========================
 def generate_pdf_report(violations, snapshots, score):
     try:
-        pdf_filename = f"report_{int(time.time())}.pdf"
-        pdf_path = os.path.join(STATIC_OUTPUT_DIR, pdf_filename)
+        pdf_path = os.path.join(STATIC_OUTPUT_DIR, f"report_{int(time.time())}.pdf")
         c = canvas.Canvas(pdf_path, pagesize=letter)
         width, height = letter
 
@@ -145,7 +127,7 @@ def generate_pdf_report(violations, snapshots, score):
 
         for v in violations:
             c.setFont("Helvetica", 10)
-            text = f"Violation: {v['violation']} | Timestamp: {v['timestamp']:.2f}s | Confidence: {v['confidence']}"
+            text = f"Violation: {v['violation']}, Timestamp: {v['timestamp']:.2f}s, Confidence: {v['confidence']}"
             c.drawString(50, y, text)
             y -= 20
 
@@ -162,36 +144,24 @@ def generate_pdf_report(violations, snapshots, score):
                 y = height - 50
 
         c.save()
-        print(f"✅ PDF generated at {pdf_path}")
-
-        # Verify the file exists with retry
-        for _ in range(3):
-            if os.path.exists(pdf_path):
-                break
-            time.sleep(1)
-        else:
-            raise Exception(f"PDF file not found at {pdf_path} after retries")
+        print(f"PDF generated at {pdf_path}")
 
-        # Generate public URL for the PDF
-        pdf_url = f"{BASE_PUBLIC_URL}{pdf_filename}"
-        print(f"📎 Public PDF URL: {pdf_url}")
-
-        return pdf_path, pdf_url
+        return pdf_path
 
     except Exception as e:
         print(f"❌ Error generating PDF: {e}")
-        return "", ""
+        return ""
 
 # ==========================
 # Video Processing
 # ==========================
 def process_video(video_data, frame_skip=5, max_frames=100):
     try:
-        print("📹 Processing video data...")
+        print("Processing video data...")
         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():
@@ -230,17 +200,13 @@ def process_video(video_data, frame_skip=5, max_frames=100):
                     }
                     violations.append(violation)
 
-                    snapshot_filename = f"snapshot_{frame_count}_{label}.jpg"
-                    snapshot_path = os.path.join(STATIC_OUTPUT_DIR, snapshot_filename)
+                    snapshot_path = os.path.join(STATIC_OUTPUT_DIR, f"snapshot_{frame_count}_{label}.jpg")
                     cv2.imwrite(snapshot_path, frame)
-                    if os.path.exists(snapshot_path):
-                        snapshots.append({
-                            "violation": label,
-                            "frame": frame_count,
-                            "snapshot_url": f"{BASE_PUBLIC_URL}{snapshot_filename}"
-                        })
-                    else:
-                        print(f"⚠️ Snapshot not saved at {snapshot_path}")
+                    snapshots.append({
+                        "violation": label,
+                        "frame": frame_count,
+                        "snapshot_url": f"https://huggingface.co/spaces/PrashanthB461/AI_Safety_Demo1/resolve/main/static/output/{os.path.basename(snapshot_path)}"
+                    })
 
             frame_count += 1
             processed_frame_count += 1
@@ -254,18 +220,15 @@ def process_video(video_data, frame_skip=5, max_frames=100):
 
         video.release()
         os.remove(video_path)
-        print(f"🗑️ Temporary video file removed: {video_path}")
 
         score = calculate_safety_score(violations)
-        pdf_path, pdf_url = generate_pdf_report(violations, snapshots, score)
+        pdf_path = generate_pdf_report(violations, snapshots, score)
 
-        if not pdf_path:
-            raise Exception("Failed to generate PDF report")
-
-        # Create Salesforce record with the PDF URL
-        record_id = push_report_to_salesforce(score, violations, pdf_url)
-        if not record_id:
-            print("⚠️ Salesforce record creation failed, but PDF URL is available")
+        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
 
         with open(pdf_path, "rb") as f:
             pdf_base64 = base64.b64encode(f.read()).decode('utf-8')
@@ -275,8 +238,7 @@ def process_video(video_data, frame_skip=5, max_frames=100):
             "snapshots": snapshots,
             "score": score,
             "pdf_base64": pdf_base64,
-            "pdf_url": pdf_url,
-            "salesforce_record_id": record_id
+            "salesforce_record_id": report_id
         }
 
     except Exception as e:
@@ -286,7 +248,6 @@ def process_video(video_data, frame_skip=5, max_frames=100):
             "snapshots": [],
             "score": 0,
             "pdf_base64": "",
-            "pdf_url": "",
             "salesforce_record_id": None,
             "error": str(e)
         }
@@ -297,40 +258,37 @@ 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()
 
         result = process_video(video_data)
-        violation_text = violations_to_text(result.get("violations", []))
         return (
-            violation_text,
+            result.get("violations", []),
             f"Safety Score: {result.get('score', 0)}%",
-              f"PDF URL: {result.get('pdf_url', '')}",
             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 f"Error: {e}", "", "", [], "", ""
+        return {"error": str(e)}, "", "", [], ""
 
 interface = gr.Interface(
     fn=gradio_interface,
     inputs=gr.Video(label="Upload Site Video"),
     outputs=[
-        gr.Textbox(label="Detected Safety Violations"),
+        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="PDF URL"),
         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). The PDF URL is saved in a Salesforce custom object."
+    description="Upload short site videos to detect safety violations (e.g., no helmet, no harness, unsafe posture)."
 )
 
 if __name__ == "__main__":
     print("🚀 Launching Safety Analyzer App...")
-    interface.launch()
+    interface.launch()