Update app.py

app.py

@@ -37,50 +37,28 @@ logger = logging.getLogger(__name__)
 class BYTETracker:
     """Custom implementation of ByteTrack to avoid installation issues"""
     def __init__(self, track_thresh=0.5, track_buffer=30, match_thresh=0.8, frame_rate=30):
-        try:
-            from yolox.tracker.byte_tracker import BYTETracker as OriginalBYTETracker
-            self.tracker = OriginalBYTETracker(
-                track_thresh=track_thresh,
-                track_buffer=track_buffer,
-                match_thresh=match_thresh,
-                frame_rate=frame_rate
-            )
-            self._original = True
-        except ImportError:
-            logger.warning("Using simplified ByteTrack implementation")
-            self._original = False
-            self.track_thresh = track_thresh
-            self.track_buffer = track_buffer
-            self.match_thresh = match_thresh
-            self.frame_rate = frame_rate
-            self.tracked_objects = {}
-            self.next_id = 1
-
-    def update(self, dets, scores, cls):
-        if self._original:
-            return self.tracker.update(dets, scores, cls)
-        
-        # Simplified tracking logic for fallback
-        if len(dets) == 0:
-            return []
+        self.track_thresh = track_thresh
+        self.track_buffer = track_buffer
+        self.match_thresh = match_thresh
+        self.frame_rate = frame_rate
+        self.next_id = 1
         
+    def update(self, dets, scores, cls):
         tracks = []
         for i, (det, score, cl) in enumerate(zip(dets, scores, cls)):
             if score < self.track_thresh:
                 continue
                 
             x, y, w, h = det
-            track_id = self.next_id
-            self.next_id += 1
             tracks.append({
-                'id': track_id,
+                'id': self.next_id,
                 'bbox': [x, y, w, h],
                 'score': score,
                 'cls': cl
             })
+            self.next_id += 1
         return tracks
 
-
 # ==========================
 # Optimized Configuration
 # ==========================
@@ -522,6 +500,67 @@ def process_video(video_data):
         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.", "", ""
+    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.", "", ""
+
 # ==========================
 # Gradio Interface
 # ==========================