Update app.py

app.py

@@ -35,7 +35,7 @@ logger = logging.getLogger(__name__)
 # ByteTrack Implementation
 # ==========================
 class BYTETracker:
-    """Custom implementation of ByteTrack to avoid installation issues"""
+    """Robust ByteTrack implementation with fallback"""
     def __init__(self, track_thresh=0.5, track_buffer=30, match_thresh=0.8, frame_rate=30):
         self.track_thresh = track_thresh
         self.track_buffer = track_buffer
@@ -112,12 +112,7 @@ CONFIG = {
     "MATCH_THRESH": 0.8
 }
 
-# Setup logging
-logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s")
-logger = logging.getLogger(__name__)
-
-os.makedirs(CONFIG["OUTPUT_DIR"], exist_ok=True)
-
+# Initialize device and model
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 logger.info(f"Using device: {device}")
 
@@ -140,12 +135,8 @@ 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
+# Helper Functions
 # ==========================
 def draw_detections(frame, detections):
     for det in detections:
@@ -226,9 +217,6 @@ def generate_violation_pdf(violations, score):
         logger.error(f"Error generating PDF: {e}")
         return "", "", None
 
-# ==========================
-# Salesforce Integration
-# ==========================
 @retry(stop_max_attempt_number=3, wait_fixed=2000)
 def connect_to_salesforce():
     try:
@@ -307,9 +295,6 @@ def push_report_to_salesforce(violations, score, pdf_path, pdf_file):
         logger.error(f"Salesforce record creation failed: {e}", exc_info=True)
         return None, ""
 
-# ==========================
-# Fast Video Processing
-# ==========================
 def process_video(video_data):
     try:
         # Create temp video file
@@ -412,8 +397,8 @@ def process_video(video_data):
 
                 # Process tracked objects
                 for obj, track_input in zip(tracked_objects, track_inputs):
-                    worker_id = obj.id
-                    label = CONFIG["VIOLATION_LABELS"].get(int(obj.cls), None)
+                    worker_id = obj['id']
+                    label = CONFIG["VIOLATION_LABELS"].get(int(obj['cls']), None)
                     bbox = track_input["bbox"]
                     conf = track_input["conf"]
                     
@@ -499,55 +484,6 @@ def process_video(video_data):
         logger.error(f"Error processing video: {e}", exc_info=True)
         yield f"Error processing video: {e}", "", "", "", ""
 
-
-# Initialize device and model
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-logger.info(f"Using device: {device}")
-
-def load_model():
-    try:
-        if os.path.isfile(CONFIG["MODEL_PATH"]):
-            model_path = CONFIG["MODEL_PATH"]
-            logger.info(f"Model loaded: {model_path}")
-        else:
-            model_path = CONFIG["FALLBACK_MODEL"]
-            logger.warning("Using fallback model. Train yolov8_safety.pt for best results.")
-            if not os.path.isfile(model_path):
-                logger.info(f"Downloading fallback model: {model_path}")
-                torch.hub.download_url_to_file('https://github.com/ultralytics/assets/releases/download/v8.3.0/yolov8n.pt', model_path)
-        model = YOLO(model_path).to(device)
-        return model
-    except Exception as e:
-        logger.error(f"Failed to load model: {e}")
-        raise
-
-model = load_model()
-
-# ==========================
-# Helper Functions
-# ==========================
-def draw_detections(frame, detections):
-    # ... [your existing implementation] ...
-
-def calculate_safety_score(violations):
-    # ... [your existing implementation] ...
-
-def generate_violation_pdf(violations, score):
-    # ... [your existing implementation] ...
-
-@retry(stop_max_attempt_number=3, wait_fixed=2000)
-def connect_to_salesforce():
-    # ... [your existing implementation] ...
-
-def upload_pdf_to_salesforce(sf, pdf_file, report_id):
-    # ... [your existing implementation] ...
-
-def push_report_to_salesforce(violations, score, pdf_path, pdf_file):
-    # ... [your existing implementation] ...
-
-def process_video(video_data):
-    # ... [your existing implementation] ...
-
 def gradio_interface(video_file):
     if not video_file:
         return "No file uploaded.", "", "No file uploaded.", "", ""