Update app.py

app.py

@@ -11,84 +11,84 @@ from reportlab.pdfgen import canvas
 from reportlab.lib.units import inch
 from io import BytesIO
 import base64
+import logging
 
 # ==========================
 # Configuration
 # ==========================
-DEFAULT_MODEL_PATH = "models/yolov8_safety.pt"
-FALLBACK_MODEL = "yolov8n.pt"
-MODEL_PATH = os.getenv("SAFETY_MODEL_PATH", DEFAULT_MODEL_PATH)
+CONFIG = {
+    "MODEL_PATH": os.getenv("SAFETY_MODEL_PATH", "models/yolov8_safety.pt"),
+    "FALLBACK_MODEL": "yolov8n.pt",
+    "OUTPUT_DIR": "static/output",
+    "VIOLATION_LABELS": {
+        0: "no_helmet",
+        1: "no_harness",
+        2: "unsafe_posture",
+        3: "unsafe_zone"
+    },
+    "SF_CREDENTIALS": {
+        "username": "prashanth1ai@safety.com",
+        "password": "SaiPrash461",
+        "security_token": "AP4AQnPoidIKPvSvNEfAHyoK",
+        "domain": "login"
+    },
+    "PUBLIC_URL_BASE": "https://huggingface.co/spaces/PrashanthB461/AI_Safety_Demo1/resolve/main/static/output/",
+    "FRAME_SKIP": 5,
+    "MAX_FRAMES": 100,
+    "MAX_PROCESSING_TIME": 30
+}
 
-# Use static/output for publicly accessible files
-STATIC_OUTPUT_DIR = "static/output"
-os.makedirs(STATIC_OUTPUT_DIR, exist_ok=True)
+# Setup logging
+logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s")
+logger = logging.getLogger(__name__)
 
-VIOLATION_LABELS = {
-    0: "no_helmet",
-    1: "no_harness",
-    2: "unsafe_posture",
-    3: "unsafe_zone"
-}
+# Ensure output directory exists
+os.makedirs(CONFIG["OUTPUT_DIR"], exist_ok=True)
 
 # ==========================
 # Device Setup
 # ==========================
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-print(f"✅ Using device: {device}")
+logger.info(f"Using device: {device}")
 
 # ==========================
-# Load Model
+# Model Loading
 # ==========================
-try:
-    selected_model = MODEL_PATH if os.path.isfile(MODEL_PATH) else FALLBACK_MODEL
-    model = YOLO(selected_model)
-    print(f"✅ Model loaded: {selected_model}")
-except Exception as e:
-    print(f"❌ Failed to load model: {e}")
-    raise
+def load_model():
+    model_path = CONFIG["MODEL_PATH"] if os.path.isfile(CONFIG["MODEL_PATH"]) else CONFIG["FALLBACK_MODEL"]
+    try:
+        model = YOLO(model_path).to(device)
+        logger.info(f"Model loaded: {model_path}")
+        return model
+    except Exception as e:
+        logger.error(f"Failed to load model: {e}")
+        raise
+
+model = load_model()
 
 # ==========================
-# Salesforce connection
+# Salesforce Integration
 # ==========================
 def connect_to_salesforce():
     try:
-        sf = Salesforce(
-            username="prashanth1ai@safety.com",
-            password="SaiPrash461",
-            security_token="AP4AQnPoidIKPvSvNEfAHyoK",
-            domain="login"
-        )
-        print("✅ Connected to Salesforce")
+        sf = Salesforce(**CONFIG["SF_CREDENTIALS"])
+        logger.info("Connected to Salesforce")
         return sf
     except Exception as e:
-        print(f"❌ Salesforce connection failed: {e}")
+        logger.error(f"Salesforce connection failed: {e}")
         raise
 
-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 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}")
-
         pdf_filename = f"violations_{int(time.time())}.pdf"
-        pdf_path = os.path.join(STATIC_OUTPUT_DIR, pdf_filename)
-
-        # Generate PDF
+        pdf_path = os.path.join(CONFIG["OUTPUT_DIR"], pdf_filename)
         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
+        y_position = 9.5 * inch
         report_data = {
             "Compliance Score": f"{score}%",
             "Violations Found": len(violations),
@@ -98,15 +98,14 @@ def generate_violation_pdf(violations, score):
             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)
+            text = f"{v['violation']} at {v['timestamp']:.2f}s (Confidence: {v['confidence']})"
+            c.drawString(1 * inch, y_position, text)
             y_position -= 0.3 * inch
-            if y_position < 1 * inch:  # Handle page overflow
+            if y_position < 1 * inch:
                 c.showPage()
                 c.setFont("Helvetica", 10)
                 y_position = 10 * inch
@@ -115,142 +114,115 @@ def generate_violation_pdf(violations, score):
         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}")
+        public_url = f"{CONFIG['PUBLIC_URL_BASE']}{pdf_filename}"
+        logger.info(f"PDF generated: {public_url}")
         return pdf_path, public_url, pdf_file
     except Exception as e:
-        print(f"❌ Error generating PDF: {e}")
+        logger.error(f"Error generating PDF: {e}")
         return "", "", None
 
-def upload_pdf_to_salesforce(pdf_file: BytesIO, report_id: str):
+def upload_pdf_to_salesforce(sf, pdf_file, report_id):
     try:
-        sf = connect_to_salesforce()
         if not pdf_file:
-            print("❌ No PDF file provided for upload")
+            logger.error("No PDF file provided for upload")
             return ""
-
-        encoded_pdf_data = base64.b64encode(pdf_file.getvalue()).decode('utf-8')
+        encoded_pdf = 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,
+            "VersionData": encoded_pdf,
             "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}'")
+        result = sf.query(f"SELECT Id, ContentDocumentId FROM ContentVersion WHERE Id = '{content_version['id']}'")
         if not result['records']:
-            print("❌ Failed to retrieve ContentVersion")
+            logger.error("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}")
+        file_url = f"https://{sf.sf_instance}/sfc/servlet.shepherd/version/download/{content_version['id']}"
+        logger.info(f"PDF uploaded to Salesforce: {file_url}")
         return file_url
     except Exception as e:
-        print(f"❌ Error uploading PDF to Salesforce: {e}")
+        logger.error(f"Error uploading PDF to Salesforce: {e}")
         return ""
 
-def push_report_to_salesforce(score, violations, pdf_path, pdf_file):
+def push_report_to_salesforce(violations, score, 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/"
-        pdf_url = f"{PUBLIC_BASE_URL}{os.path.basename(pdf_path)}"
+        violations_text = "\n".join(
+            f"{v['violation']} at {v['timestamp']:.2f}s (Confidence: {v['confidence']})"
+            for v in violations
+        ) or "No violations detected."
+        pdf_url = f"{CONFIG['PUBLIC_URL_BASE']}{os.path.basename(pdf_path)}" if pdf_path else ""
 
         record_data = {
             "Compliance_Score__c": score,
-            "Violations_Found__c": violations_count,
-            "Violations_Details__c": violations_details,
+            "Violations_Found__c": len(violations),
+            "Violations_Details__c": violations_text,
             "Status__c": "Pending",
             "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)
-        record_id = record.get('id')
-        print(f"✅ Salesforce record created with Id: {record_id}")
+        record_id = record["id"]
+        logger.info(f"Salesforce record created: {record_id}")
 
-        # Upload PDF to Salesforce
         if pdf_file:
-            uploaded_url = upload_pdf_to_salesforce(pdf_file, record_id)
+            uploaded_url = upload_pdf_to_salesforce(sf, pdf_file, record_id)
             if uploaded_url:
-                sf.Safety_Video_Report__c.update(
-                    record_id,
-                    {"PDF_Report_URL__c": 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}")
+        logger.error(f"Salesforce record creation failed: {e}")
         return None, ""
 
 # ==========================
 # 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)
+    score = 100 - sum(penalties.get(v["violation"], 0) for v in violations)
+    return max(score, 0)
 
 # ==========================
 # Video Processing
 # ==========================
-def process_video(video_data, frame_skip=5, max_frames=100):
+def process_video(video_data):
     try:
-        print("Processing video data...")
-        video_path = os.path.join(STATIC_OUTPUT_DIR, f"temp_{int(time.time())}.mp4")
+        video_path = os.path.join(CONFIG["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}")
+        logger.info(f"Video saved: {video_path}")
 
         video = cv2.VideoCapture(video_path)
         if not video.isOpened():
-            raise ValueError("Could not open video file.")
+            raise ValueError("Could not open video file")
 
-        frame_count = 0
-        violations = []
-        snapshots = []
-        processed_frame_count = 0
+        violations, snapshots = [], []
+        frame_count, processed_frames = 0, 0
         start_time = time.time()
 
-        while True:
+        while processed_frames < CONFIG["MAX_FRAMES"]:
             ret, frame = video.read()
             if not ret:
                 break
-
-            if frame_count % frame_skip != 0:
+            if frame_count % CONFIG["FRAME_SKIP"] != 0:
                 frame_count += 1
                 continue
 
             results = model(frame, device=device)
-
             for result in results:
                 for box in result.boxes:
-                    cls = int(box.cls)
-                    conf = float(box.conf)
+                    cls, conf = int(box.cls), float(box.conf)
                     xywh = box.xywh.cpu().numpy()[0]
-
-                    label = VIOLATION_LABELS.get(cls, f"class_{cls}")
+                    label = CONFIG["VIOLATION_LABELS"].get(cls, f"class_{cls}")
                     violation = {
                         "frame": frame_count,
                         "violation": label,
@@ -260,36 +232,25 @@ 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_path = os.path.join(CONFIG["OUTPUT_DIR"], f"snapshot_{frame_count}_{label}.jpg")
                     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"{CONFIG['PUBLIC_URL_BASE']}{os.path.basename(snapshot_path)}"
                     })
 
             frame_count += 1
-            processed_frame_count += 1
-
-            if processed_frame_count >= max_frames:
-                break
-
-            if time.time() - start_time > 30:
-                print("⏰ Exceeded 30 seconds of processing time.")
+            processed_frames += 1
+            if time.time() - start_time > CONFIG["MAX_PROCESSING_TIME"]:
+                logger.warning("Processing time limit exceeded")
                 break
 
         video.release()
         os.remove(video_path)
-
         score = calculate_safety_score(violations)
         pdf_path, pdf_url, pdf_file = generate_violation_pdf(violations, score)
-
-        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, final_pdf_url = None, ""
-            print("❌ No PDF generated, skipping Salesforce record creation")
+        report_id, final_pdf_url = push_report_to_salesforce(violations, score, pdf_path, pdf_file)
 
         return {
             "violations": violations,
@@ -297,43 +258,42 @@ def process_video(video_data, frame_skip=5, max_frames=100):
             "score": score,
             "salesforce_record_id": report_id,
             "violation_details_url": final_pdf_url,
-            "violation_details_content": violations_to_text(violations) if violations else "No violations detected."
+            "violation_details_content": "\n".join(
+                f"{v['violation']} at {v['timestamp']:.2f}s (Confidence: {v['confidence']})"
+                for v in violations
+            ) or "No violations detected."
         }
-
     except Exception as e:
-        print(f"❌ Error processing video: {e}")
+        logger.error(f"Error processing video: {e}")
         return {
             "violations": [],
             "snapshots": [],
             "score": 0,
             "salesforce_record_id": None,
             "violation_details_url": "",
-            "violation_details_content": "Error: Could not process video.",
-            "error": str(e)
+            "violation_details_content": f"Error: {str(e)}"
         }
 
 # ==========================
 # Gradio Interface
 # ==========================
 def gradio_interface(video_file):
+    if not video_file:
+        return {"error": "No video file uploaded"}, "", [], "", "", "No file uploaded."
     try:
-        if not video_file:
-            return {"error": "Please upload a video file."}, "", [], "", "", "No file uploaded."
-
         with open(video_file, "rb") as f:
             video_data = f.read()
-
         result = process_video(video_data)
         return (
-            result.get("violations", []),
-            f"Safety Score: {result.get('score', 0)}%",
-            result.get("snapshots", []),
-            f"Salesforce Record ID: {result.get('salesforce_record_id', '')}",
-            f"Violation Details URL: {result.get('violation_details_url', '')}",
-            result.get("violation_details_content", "No content available.")
+            result["violations"],
+            f"Safety Score: {result['score']}%",
+            result["snapshots"],
+            f"Salesforce Record ID: {result['salesforce_record_id'] or 'N/A'}",
+            f"Violation Details URL: {result['violation_details_url'] or 'N/A'}",
+            result["violation_details_content"]
         )
     except Exception as e:
-        print(f"❌ Error in gradio_interface: {e}")
+        logger.error(f"Error in Gradio interface: {e}")
         return {"error": str(e)}, "", [], "", "", "Error in processing."
 
 interface = gr.Interface(
@@ -345,12 +305,12 @@ interface = gr.Interface(
         gr.JSON(label="Snapshots"),
         gr.Textbox(label="Salesforce Record ID"),
         gr.Textbox(label="Violation Details URL"),
-        gr.Textbox(label="Violation Details Content (Preview)")
+        gr.Textbox(label="Violation Details Content")
     ],
     title="Worksite Safety Violation Analyzer",
     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...")
+    logger.info("Launching Safety Analyzer App...")
     interface.launch()