Update app.py

app.py

@@ -19,10 +19,12 @@ 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 (adjust based on your hosting platform)
+BASE_PUBLIC_URL = "https://huggingface.co/spaces/PrashanthB461/AI_Safety_Demo1/resolve/main/static/output/"
+
 VIOLATION_LABELS = {
     0: "no_helmet",
     1: "no_harness",
@@ -48,7 +50,7 @@ except Exception as e:
     raise
 
 # ==========================
-# Salesforce connection (hardcoded creds for testing)
+# Salesforce Connection
 # ==========================
 def connect_to_salesforce():
     try:
@@ -74,43 +76,27 @@ def violations_to_text(violations):
         lines.append(line)
     return "\n".join(lines)
 
-def push_report_to_salesforce(score, violations, pdf_filename):
+# ==========================
+# 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)
 
-        # Upload the PDF to Salesforce Files
-        with open(pdf_filename, "rb") as f:
-            pdf_data = f.read()
-        file_name = os.path.basename(pdf_filename)
-
-        # Upload to ContentVersion (Salesforce Files)
-        print(f"Uploading file to Salesforce Files: {file_name}")
-        file_upload = sf.ContentVersion.create({
-            'Title': file_name,
-            'PathOnClient': file_name,
-            'VersionData': pdf_data
-        })
-
-        file_id = file_upload.get('id')
-        print(f"File uploaded with ID: {file_id}")
-        
-        # Get the public URL for the uploaded file (direct download link)
-        public_url = f"https://{sf.sf_instance}.salesforce.com/sfc/servlet.shepherd/version/download/{file_id}"
-
-        # Create the record in Salesforce
+        # 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}")
         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": public_url
+            "PDF_Report_URL__c": pdf_url
         })
-        print(f"✅ Salesforce record created with Id: {record.get('id')}")
-        return record.get("id")
+        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}")
@@ -140,7 +126,8 @@ def calculate_safety_score(violations):
 # ==========================
 def generate_pdf_report(violations, snapshots, score):
     try:
-        pdf_path = os.path.join(STATIC_OUTPUT_DIR, f"report_{int(time.time())}.pdf")
+        pdf_filename = f"report_{int(time.time())}.pdf"
+        pdf_path = os.path.join(STATIC_OUTPUT_DIR, pdf_filename)
         c = canvas.Canvas(pdf_path, pagesize=letter)
         width, height = letter
 
@@ -176,11 +163,15 @@ def generate_pdf_report(violations, snapshots, score):
         c.save()
         print(f"PDF generated at {pdf_path}")
 
-        return 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
 
     except Exception as e:
         print(f"❌ Error generating PDF: {e}")
-        return ""
+        return "", ""
 
 # ==========================
 # Video Processing
@@ -230,12 +221,13 @@ def process_video(video_data, frame_skip=5, max_frames=100):
                     }
                     violations.append(violation)
 
-                    snapshot_path = os.path.join(STATIC_OUTPUT_DIR, f"snapshot_{frame_count}_{label}.jpg")
+                    snapshot_filename = f"snapshot_{frame_count}_{label}.jpg"
+                    snapshot_path = os.path.join(STATIC_OUTPUT_DIR, snapshot_filename)
                     cv2.imwrite(snapshot_path, frame)
                     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)}"
+                        "snapshot_url": f"{BASE_PUBLIC_URL}{snapshot_filename}"
                     })
 
             frame_count += 1
@@ -252,13 +244,14 @@ 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)
+        pdf_path, pdf_url = generate_pdf_report(violations, snapshots, score)
+
+        if not pdf_path:
+            raise Exception("Failed to generate PDF report")
 
-        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
+        # Create Salesforce record with the PDF URL
+        record_id = push_report_to_salesforce(score, violations, pdf_url)
+        print(f"Salesforce record created with Id: {record_id}")
 
         with open(pdf_path, "rb") as f:
             pdf_base64 = base64.b64encode(f.read()).decode('utf-8')
@@ -268,7 +261,8 @@ def process_video(video_data, frame_skip=5, max_frames=100):
             "snapshots": snapshots,
             "score": score,
             "pdf_base64": pdf_base64,
-            "salesforce_record_id": report_id
+            "pdf_url": pdf_url,
+            "salesforce_record_id": record_id
         }
 
     except Exception as e:
@@ -278,6 +272,7 @@ 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)
         }
@@ -288,7 +283,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()
@@ -296,15 +291,16 @@ def gradio_interface(video_file):
         result = process_video(video_data)
         violation_text = violations_to_text(result.get("violations", []))
         return (
-            violation_text,  # Returning violation details as plain text
+            violation_text,
             f"Safety Score: {result.get('score', 0)}%",
             result.get("pdf_base64", ""),
             result.get("snapshots", []),
+            f"PDF URL: {result.get('pdf_url', '')}",
             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 f"Error: {e}", "", "", [], "", ""
 
 interface = gr.Interface(
     fn=gradio_interface,
@@ -314,12 +310,13 @@ interface = gr.Interface(
         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)."
+    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."
 )
 
 if __name__ == "__main__":
     print("🚀 Launching Safety Analyzer App...")
-    interface.launch()
+    interface.launch()