Update app.py

app.py

@@ -6,6 +6,11 @@ import numpy as np
 from ultralytics import YOLO
 import time
 from simple_salesforce import Salesforce
+from reportlab.lib.pagesizes import letter
+from reportlab.pdfgen import canvas
+from reportlab.lib.units import inch
+from io import BytesIO
+import base64
 
 # ==========================
 # Configuration
@@ -66,58 +71,131 @@ def violations_to_text(violations):
         lines.append(line)
     return "\n".join(lines)
 
-def generate_violation_text_file(violations):
+def generate_violation_pdf(violations, score):
     try:
         # Ensure the directory is writable
         if not os.access(STATIC_OUTPUT_DIR, os.W_OK):
             raise PermissionError(f"Cannot write to directory: {STATIC_OUTPUT_DIR}")
 
-        text_filename = f"violations_{int(time.time())}.txt"
-        text_path = os.path.join(STATIC_OUTPUT_DIR, text_filename)
-        
-        # Generate the content
-        content = "Worksite Safety Violation Details\n" + "=" * 40 + "\n"
-        content += violations_to_text(violations)
-        
-        # Write the file
-        with open(text_path, "w") as f:
-            f.write(content)
-        
-        print(f"✅ Text file generated at {text_path}")
-        if not os.path.exists(text_path):
-            raise FileNotFoundError(f"Text file was not created 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, content
+        pdf_filename = f"violations_{int(time.time())}.pdf"
+        pdf_path = os.path.join(STATIC_OUTPUT_DIR, pdf_filename)
+
+        # Generate PDF
+        pdf_file = BytesIO()
+        c = canvas.Canvas(pdf_file, pagesize=letter)
+        c.setFont("Helvetica", 12)
+        c.drawString(1 * inch, 10 * inch, "Worksite Safety Violation Report")
+        c.setFont("Helvetica", 10)
+        y_position = 9.5 * inch
+
+        # Add report details
+        report_data = {
+            "Compliance Score": f"{score}%",
+            "Violations Found": len(violations),
+            "Timestamp": time.strftime("%Y-%m-%d %H:%M:%S")
+        }
+        for key, value in report_data.items():
+            c.drawString(1 * inch, y_position, f"{key}: {value}")
+            y_position -= 0.3 * inch
+
+        # Add violation details
+        y_position -= 0.3 * inch
+        c.drawString(1 * inch, y_position, "Violation Details:")
+        y_position -= 0.3 * inch
+        for v in violations:
+            violation_text = f"{v['violation']} at {v['timestamp']:.2f}s (Confidence: {v['confidence']})"
+            c.drawString(1 * inch, y_position, violation_text)
+            y_position -= 0.3 * inch
+            if y_position < 1 * inch:  # Handle page overflow
+                c.showPage()
+                c.setFont("Helvetica", 10)
+                y_position = 10 * inch
+
+        c.showPage()
+        c.save()
+        pdf_file.seek(0)
+
+        # Save PDF to disk
+        with open(pdf_path, "wb") as f:
+            f.write(pdf_file.getvalue())
+
+        print(f"✅ PDF generated at {pdf_path}")
+        if not os.path.exists(pdf_path):
+            raise FileNotFoundError(f"PDF was not created at {pdf_path}")
+
+        public_url = f"https://huggingface.co/spaces/PrashanthB461/AI_Safety_Demo1/resolve/main/static/output/{pdf_filename}"
+        print(f"✅ PDF URL: {public_url}")
+        return pdf_path, public_url, pdf_file
+    except Exception as e:
+        print(f"❌ Error generating PDF: {e}")
+        return "", "", None
+
+def upload_pdf_to_salesforce(pdf_file: BytesIO, report_id: str):
+    try:
+        sf = connect_to_salesforce()
+        if not pdf_file:
+            print("❌ No PDF file provided for upload")
+            return ""
+
+        encoded_pdf_data = 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_data,
+            "FirstPublishLocationId": report_id
+        }
+
+        content_version = sf.ContentVersion.create(content_version_data)
+        content_version_id = content_version["id"]
+
+        result = sf.query(f"SELECT Id, ContentDocumentId FROM ContentVersion WHERE Id = '{content_version_id}'")
+        if not result['records']:
+            print("❌ Failed to retrieve ContentVersion")
+            return ""
+
+        file_url = f"https://{sf.sf_instance}/sfc/servlet.shepherd/version/download/{content_version_id}"
+        print(f"✅ PDF uploaded to Salesforce, URL: {file_url}")
+        return file_url
     except Exception as e:
-        print(f"❌ Error generating text file: {e}")
-        return "", "", "Error: Could not generate violation details."
+        print(f"❌ Error uploading PDF to Salesforce: {e}")
+        return ""
 
-def push_report_to_salesforce(score, violations, text_filename):
+def push_report_to_salesforce(score, violations, pdf_path, pdf_file):
     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/"
-        text_url = f"{PUBLIC_BASE_URL}{os.path.basename(text_filename)}"
+        pdf_url = f"{PUBLIC_BASE_URL}{os.path.basename(pdf_path)}"
 
         record_data = {
             "Compliance_Score__c": score,
             "Violations_Found__c": violations_count,
             "Violations_Details__c": violations_details,
             "Status__c": "Pending",
-            "PDF_Report_URL__c": text_url  # Store text file URL in PDF_Report_URL__c
+            "PDF_Report_URL__c": pdf_url
         }
 
         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')
+        record_id = record.get('id')
+        print(f"✅ Salesforce record created with Id: {record_id}")
+
+        # Upload PDF to Salesforce
+        if pdf_file:
+            uploaded_url = upload_pdf_to_salesforce(pdf_file, record_id)
+            if uploaded_url:
+                sf.Safety_Video_Report__c.update(
+                    record_id,
+                    {"PDF_Report_URL__c": uploaded_url}
+                )
+                pdf_url = uploaded_url
+
+        return record_id, pdf_url
     except Exception as e:
         print(f"❌ Salesforce record creation failed: {e}")
-        return None
+        return None, ""
 
 # ==========================
 # Safety Score Calculation
@@ -204,22 +282,22 @@ def process_video(video_data, frame_skip=5, max_frames=100):
         os.remove(video_path)
 
         score = calculate_safety_score(violations)
-        text_path, text_url, text_content = generate_violation_text_file(violations)
+        pdf_path, pdf_url, pdf_file = generate_violation_pdf(violations, score)
 
-        if text_path:
-            report_id = push_report_to_salesforce(score, violations, text_path)
+        if pdf_path:
+            report_id, final_pdf_url = push_report_to_salesforce(score, violations, pdf_path, pdf_file)
             print(f"✅ Salesforce record created with Id: {report_id}")
         else:
-            report_id = None
-            print("❌ No text file generated, skipping Salesforce record creation")
+            report_id, final_pdf_url = None, ""
+            print("❌ No PDF generated, skipping Salesforce record creation")
 
         return {
             "violations": violations,
             "snapshots": snapshots,
             "score": score,
             "salesforce_record_id": report_id,
-            "violation_details_url": text_url,
-            "violation_details_content": text_content
+            "violation_details_url": final_pdf_url,
+            "violation_details_content": violations_to_text(violations) if violations else "No violations detected."
         }
 
     except Exception as e: