Update app.py

app.py

@@ -23,20 +23,17 @@ from functools import partial
 # ==========================
 # Configuration and Setup
 # ==========================
-# Handle Ultralytics config directory
 os.environ['YOLO_CONFIG_DIR'] = '/tmp/Ultralytics'
 os.makedirs('/tmp/Ultralytics', exist_ok=True)
 
-# Setup logging
-logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s")
+logging.basicConfig(level=logging.DEBUG, format="%(asctime)s - %(levelname)s - %(message)s")
 logger = logging.getLogger(__name__)
 
 # ==========================
 # ByteTrack Implementation
 # ==========================
 class BYTETracker:
-    """Robust ByteTrack implementation with fallback"""
-    def __init__(self, track_thresh=0.5, track_buffer=30, match_thresh=0.8, frame_rate=30):
+    def __init__(self, track_thresh=0.3, track_buffer=30, match_thresh=0.7, frame_rate=30):
         self.track_thresh = track_thresh
         self.track_buffer = track_buffer
         self.match_thresh = match_thresh
@@ -47,6 +44,7 @@ class BYTETracker:
         tracks = []
         for i, (det, score, cl) in enumerate(zip(dets, scores, cls)):
             if score < self.track_thresh:
+                logger.debug(f"Skipping detection with score {score} below threshold {self.track_thresh}")
                 continue
                 
             x, y, w, h = det
@@ -74,11 +72,11 @@ CONFIG = {
         4: "improper_tool_use"
     },
     "CLASS_COLORS": {
-        "no_helmet": (0, 0, 255),  # Red
-        "no_harness": (0, 165, 255),  # Orange
-        "unsafe_posture": (0, 255, 0),  # Green
-        "unsafe_zone": (255, 0, 0),  # Blue
-        "improper_tool_use": (255, 255, 0)  # Yellow
+        "no_helmet": (0, 0, 255),
+        "no_harness": (0, 165, 255),
+        "unsafe_posture": (0, 255, 0),
+        "unsafe_zone": (255, 0, 0),
+        "improper_tool_use": (255, 255, 0)
     },
     "DISPLAY_NAMES": {
         "no_helmet": "No Helmet Violation",
@@ -93,26 +91,25 @@ CONFIG = {
         "security_token": "AP4AQnPoidIKPvSvNEfAHyoK",
         "domain": "login"
     },
-    "PUBLIC_URL_BASE": "https://huggingface.co/spaces/PrashanthB461/AI_Safety_Demo2/resolve/main/static/output/",
+    "PUBLIC_URL_BASE": "https://huggingface.co/spaces/PrashanthB461/AI_Sadio2/resolve/main/static/output/",
     "CONFIDENCE_THRESHOLDS": {
-        "no_helmet": 0.75,
-        "no_harness": 0.4,
-        "unsafe_posture": 0.4,
-        "unsafe_zone": 0.4,
-        "improper_tool_use": 0.4
+        "no_helmet": 0.5,  # Lowered from 0.75
+        "no_harness": 0.3,  # Lowered from 0.4
+        "unsafe_posture": 0.3,  # Lowered from 0.4
+        "unsafe_zone": 0.3,  # Lowered from 0.4
+        "improper_tool_use": 0.3  # Lowered from 0.4
     },
-    "MIN_VIOLATION_FRAMES": 3,
+    "MIN_VIOLATION_FRAMES": 1,  # Lowered from 3
     "WORKER_TRACKING_DURATION": 5.0,
     "MAX_PROCESSING_TIME": 60,
     "FRAME_SKIP": 1,
-    "BATCH_SIZE": 32,
+    "BATCH_SIZE": 16,  # Reduced from 32 to prevent memory issues
     "PARALLEL_WORKERS": max(1, cpu_count() - 1),
     "TRACK_BUFFER": 30,
-    "TRACK_THRESH": 0.4,
-    "MATCH_THRESH": 0.8
+    "TRACK_THRESH": 0.3,  # Lowered from 0.4
+    "MATCH_THRESH": 0.7  # Lowered from 0.8
 }
 
-# Initialize device and model
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 logger.info(f"Using device: {device}")
 
@@ -128,6 +125,7 @@ def load_model():
                 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)
+        logger.info(f"Model classes: {model.names}")
         return model
     except Exception as e:
         logger.error(f"Failed to load model: {e}")
@@ -138,6 +136,11 @@ model = load_model()
 # ==========================
 # Helper Functions
 # ==========================
+def preprocess_frame(frame):
+    """Apply basic preprocessing to enhance detection"""
+    frame = cv2.convertScaleAbs(frame, alpha=1.2, beta=20)  # Increase contrast
+    return frame
+
 def draw_detections(frame, detections):
     for det in detections:
         label = det.get("violation", "Unknown")
@@ -297,23 +300,19 @@ def push_report_to_salesforce(violations, score, pdf_path, pdf_file):
 
 def process_video(video_data):
     try:
-        # Ensure output directory exists
         os.makedirs(CONFIG["OUTPUT_DIR"], exist_ok=True)
         logger.info(f"Output directory ensured: {CONFIG['OUTPUT_DIR']}")
 
-        # Create temp video file
         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)
         logger.info(f"Video saved: {video_path}")
 
-        # Open video file
         cap = cv2.VideoCapture(video_path)
         if not cap.isOpened():
             os.remove(video_path)
             raise ValueError("Could not open video file")
 
-        # Get video properties
         total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
         fps = cap.get(cv2.CAP_PROP_FPS) or 30
         duration = total_frames / fps
@@ -321,7 +320,6 @@ def process_video(video_data):
         height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
         logger.info(f"Video properties: {duration:.2f}s, {total_frames} frames, {fps:.1f} FPS, {width}x{height}")
 
-        # Initialize ByteTrack
         tracker = BYTETracker(
             track_thresh=CONFIG["TRACK_THRESH"],
             track_buffer=CONFIG["TRACK_BUFFER"],
@@ -329,18 +327,15 @@ def process_video(video_data):
             frame_rate=fps
         )
 
-        # Track violations by worker ID and type
-        violation_tracker = {}  # {worker_id: {violation_type: [detections]}}
+        violation_tracker = {}
         snapshots = []
         start_time = time.time()
         frame_skip = CONFIG["FRAME_SKIP"]
 
-        # Process frames in batches
         while True:
             batch_frames = []
             batch_indices = []
             
-            # Collect frames for this batch
             for _ in range(CONFIG["BATCH_SIZE"]):
                 frame_idx = int(cap.get(cv2.CAP_PROP_POS_FRAMES))
                 if frame_idx >= total_frames:
@@ -350,7 +345,8 @@ def process_video(video_data):
                 if not ret:
                     break
                 
-                # Skip frames if needed
+                frame = preprocess_frame(frame)
+                
                 for _ in range(frame_skip - 1):
                     if not cap.grab():
                         break
@@ -358,24 +354,19 @@ def process_video(video_data):
                 batch_frames.append(frame)
                 batch_indices.append(frame_idx)
 
-            # Break if no more frames
             if not batch_frames:
                 break
 
-            # Run batch detection
             results = model(batch_frames, device=device, conf=0.1, verbose=False)
             
-            # Process results for each frame in batch
             for i, (result, frame_idx) in enumerate(zip(results, batch_indices)):
                 current_time = frame_idx / fps
                 
-                # Update progress
                 if time.time() - start_time > 1.0:
                     progress = (frame_idx / total_frames) * 100
                     yield f"Processing video... {progress:.1f}% complete (Frame {frame_idx}/{total_frames})", "", "", "", ""
                     start_time = time.time()
 
-                # Prepare detections for ByteTrack
                 boxes = result.boxes
                 track_inputs = []
                 for box in boxes:
@@ -383,24 +374,27 @@ def process_video(video_data):
                     conf = float(box.conf)
                     label = CONFIG["VIOLATION_LABELS"].get(cls, None)
                     
-                    if label is None or conf < CONFIG["CONFIDENCE_THRESHOLDS"].get(label, 0.25):
+                    if label is None:
+                        logger.debug(f"Unknown class ID {cls} detected, skipping")
+                        continue
+                    if conf < CONFIG["CONFIDENCE_THRESHOLDS"].get(label, 0.25):
+                        logger.debug(f"Detection for {label} with confidence {conf} below threshold {CONFIG['CONFIDENCE_THRESHOLDS'].get(label, 0.25)}")
                         continue
 
                     bbox = box.xywh.cpu().numpy()[0]
                     track_inputs.append({
-                        "bbox": bbox,  # [x, y, w, h]
+                        "bbox": bbox,
                         "conf": conf,
                         "cls": cls
                     })
 
-                # Update tracker
                 tracked_objects = tracker.update(
                     np.array([t["bbox"] for t in track_inputs]),
                     np.array([t["conf"] for t in track_inputs]),
                     np.array([t["cls"] for t in track_inputs])
                 )
+                logger.debug(f"Frame {frame_idx}: {len(tracked_objects)} objects tracked")
 
-                # 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)
@@ -415,8 +409,8 @@ def process_video(video_data):
                         "timestamp": current_time,
                         "worker_id": worker_id
                     }
+                    logger.debug(f"Detection: {detection}")
 
-                    # Track violations by worker_id and type
                     if worker_id not in violation_tracker:
                         violation_tracker[worker_id] = {}
                     if label not in violation_tracker[worker_id]:
@@ -429,19 +423,19 @@ def process_video(video_data):
         processing_time = time.time() - start_time
         logger.info(f"Processing complete in {processing_time:.2f}s")
 
-        # Consolidate violations
         violations = []
         for worker_id, worker_violations in violation_tracker.items():
             for label, detections in worker_violations.items():
                 if len(detections) >= CONFIG["MIN_VIOLATION_FRAMES"]:
-                    # Select highest-confidence detection
                     best_detection = max(detections, key=lambda x: x["confidence"])
                     best_detection["start_timestamp"] = min(d["timestamp"] for d in detections)
                     best_detection["end_timestamp"] = max(d["timestamp"] for d in detections)
                     violations.append(best_detection)
 
-                    # Capture snapshot for confirmed violation
                     cap = cv2.VideoCapture(video_path)
+                    if not cap.isOpened():
+                        logger.warning(f"Could not reopen video for snapshot at frame {best_detection['frame']}")
+                        continue
                     cap.set(cv2.CAP_PROP_POS_FRAMES, best_detection["frame"])
                     ret, snapshot_frame = cap.read()
                     if ret:
@@ -457,8 +451,8 @@ def process_video(video_data):
                         })
                     cap.release()
 
-        # Generate results
         if not violations:
+            logger.info("No violations detected after processing")
             yield "No violations detected in the video.", "Safety Score: 100%", "No snapshots captured.", "N/A", "N/A"
             return