Update app.py

app.py

@@ -4,7 +4,6 @@ import gradio as gr
 import torch
 import numpy as np
 from ultralytics import YOLO
-from bytetrack import BYTETracker  # Added ByteTrack
 import time
 from simple_salesforce import Salesforce
 from reportlab.lib.pagesizes import letter
@@ -17,6 +16,20 @@ from retrying import retry
 import uuid
 from multiprocessing import Pool, cpu_count
 from functools import partial
+import subprocess
+import sys
+
+# ==========================
+# ByteTrack Installation & Import
+# ==========================
+try:
+    from bytetrack import BYTETracker
+except ImportError:
+    logger = logging.getLogger(__name__)
+    logger.warning("ByteTrack not found, installing...")
+    subprocess.check_call([sys.executable, "-m", "pip", "install", "cython", "lap"])
+    subprocess.check_call([sys.executable, "-m", "pip", "install", "git+https://github.com/ifzhang/ByteTrack.git@main"])
+    from bytetrack import BYTETracker
 
 # ==========================
 # Optimized Configuration
@@ -54,21 +67,21 @@ CONFIG = {
     },
     "PUBLIC_URL_BASE": "https://huggingface.co/spaces/PrashanthB461/AI_Safety_Demo2/resolve/main/static/output/",
     "CONFIDENCE_THRESHOLDS": {
-        "no_helmet": 0.75,  # Increased for stricter helmet detection
+        "no_helmet": 0.75,
         "no_harness": 0.4,
         "unsafe_posture": 0.4,
         "unsafe_zone": 0.4,
         "improper_tool_use": 0.4
     },
     "MIN_VIOLATION_FRAMES": 3,
-    "WORKER_TRACKING_DURATION": 5.0,  # Increased to 5s for better continuity
-    "MAX_PROCESSING_TIME": 60,  # 1 minute limit
-    "FRAME_SKIP": 1,  # Process all frames to avoid missing violations
-    "BATCH_SIZE": 32,  # Increased for performance
-    "PARALLEL_WORKERS": max(1, cpu_count() - 1),  # Use all CPU cores except one
-    "TRACK_BUFFER": 30,  # Frames to keep a track alive
-    "TRACK_THRESH": 0.4,  # Tracking confidence threshold
-    "MATCH_THRESH": 0.8  # IOU threshold for matching
+    "WORKER_TRACKING_DURATION": 5.0,
+    "MAX_PROCESSING_TIME": 60,
+    "FRAME_SKIP": 1,
+    "BATCH_SIZE": 32,
+    "PARALLEL_WORKERS": max(1, cpu_count() - 1),
+    "TRACK_BUFFER": 30,
+    "TRACK_THRESH": 0.4,
+    "MATCH_THRESH": 0.8
 }
 
 # Setup logging
@@ -99,6 +112,10 @@ def load_model():
 
 model = load_model()
 
+# [Rest of your existing functions remain exactly the same...]
+# draw_detections(), calculate_safety_score(), generate_violation_pdf(),
+# connect_to_salesforce(), upload_pdf_to_salesforce(), push_report_to_salesforce(),
+# process_video(), and gradio_interface() functions should be kept exactly as they were
 # ==========================
 # Optimized Helper Functions
 # ==========================
@@ -454,22 +471,10 @@ def process_video(video_data):
         logger.error(f"Error processing video: {e}", exc_info=True)
         yield f"Error processing video: {e}", "", "", "", ""
 
+
 # ==========================
 # Gradio Interface
 # ==========================
-def gradio_interface(video_file):
-    if not video_file:
-        return "No file uploaded.", "", "No file uploaded.", "", ""
-    try:
-        with open(video_file, "rb") as f:
-            video_data = f.read()
-
-        for status, score, snapshots_text, record_id, details_url in process_video(video_data):
-            yield status, score, snapshots_text, record_id, details_url
-    except Exception as e:
-        logger.error(f"Error in Gradio interface: {e}", exc_info=True)
-        yield f"Error: {str(e)}", "", "Error in processing.", "", ""
-
 interface = gr.Interface(
     fn=gradio_interface,
     inputs=gr.Video(label="Upload Site Video"),