Update app.py

app.py

@@ -216,7 +216,7 @@ CONFIG = {
     "MATCH_THRESH": 0.5,
     "SNAPSHOT_QUALITY": 95,
     "MAX_WORKER_DISTANCE": 300,
-    "MODEL_INPUT_SIZE": (640, 640)  # Updated to match YOLO input requirements
+    "TARGET_RESOLUTION": (384, 384)  # Changed to 384x384 (divisible by 32)
 }
 
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
@@ -246,56 +246,11 @@ def load_model():
 model = load_model()
 
 # ========================== # Helper Functions # ==========================
-def preprocess_frame(frame, original_shape):
-    # Resize while preserving aspect ratio, then pad to MODEL_INPUT_SIZE (640x640)
-    target_size = CONFIG["MODEL_INPUT_SIZE"]  # (640, 640)
-    h, w = frame.shape[:2]
-    scale = min(target_size[0] / w, target_size[1] / h)
-    new_w, new_h = int(w * scale), int(h * scale)
-    
-    # Resize the frame
-    frame_resized = cv2.resize(frame, (new_w, new_h), interpolation=cv2.INTER_LINEAR)
-    
-    # Create a new 640x640 image with padding
-    padded_frame = np.zeros((target_size[1], target_size[0], 3), dtype=np.uint8)
-    top = (target_size[1] - new_h) // 2
-    left = (target_size[0] - new_w) // 2
-    padded_frame[top:top+new_h, left:left+new_w] = frame_resized
-    
-    # Apply contrast adjustment
-    padded_frame = cv2.convertScaleAbs(padded_frame, alpha=1.2, beta=20)
-    
-    # Store padding info to adjust bounding boxes later
-    padding_info = {
-        "scale": scale,
-        "top": top,
-        "left": left,
-        "original_shape": original_shape
-    }
-    
-    return padded_frame, padding_info
-
-def adjust_bbox(bbox, padding_info):
-    # Adjust bounding box coordinates from padded 640x640 space back to original frame space
-    scale = padding_info["scale"]
-    top = padding_info["top"]
-    left = padding_info["left"]
-    
-    x, y, w, h = bbox
-    # Remove padding offset and scale back
-    x = (x - left) / scale
-    y = (y - top) / scale
-    w = w / scale
-    h = h / scale
-    
-    # Ensure coordinates are within original frame bounds
-    orig_h, orig_w = padding_info["original_shape"][:2]
-    x = max(0, min(x, orig_w))
-    y = max(0, min(y, orig_h))
-    w = max(0, min(w, orig_w - x))
-    h = max(0, min(h, orig_h - y))
-    
-    return [x, y, w, h]
+def preprocess_frame(frame):
+    target_res = CONFIG["TARGET_RESOLUTION"]
+    frame = cv2.resize(frame, target_res, interpolation=cv2.INTER_LINEAR)
+    frame = cv2.convertScaleAbs(frame, alpha=1.2, beta=20)
+    return frame
 
 def draw_detections(frame, detections):
     result_frame = frame.copy()
@@ -496,7 +451,7 @@ def push_report_to_salesforce(violations, score, pdf_path, pdf_file):
         
         try:
             record = sf.Safety_Video_Report__c.create(record_data)
-            logger.info(f"Created Safety_Violation_Report__c record: {record['id']}")
+            logger.info(f"Created Safety_Video_Report__c record: {record['id']}")
         except Exception as e:
             logger.error(f"Failed to create Safety_Video_Report__c: {e}")
             record = sf.Account.create({"Name": f"Safety_Report_{int(time.time())}"})
@@ -593,16 +548,15 @@ def process_video(video_data, temp_dir):
 
         worker_id_mapping = {}
         unique_violations = {}
-        snapshots = []
+        violation_frames = {}  # Store frame indices for violations
         start_time = time.time()
         frame_skip = CONFIG["FRAME_SKIP"]
         processed_frames = 0
+        last_yield_time = start_time
 
         while processed_frames < total_frames:
             batch_frames = []
             batch_indices = []
-            batch_padding_info = []
-            batch_original_frames = []
             
             for _ in range(CONFIG["BATCH_SIZE"]):
                 frame_idx = int(cap.get(cv2.CAP_PROP_POS_FRAMES))
@@ -614,9 +568,7 @@ def process_video(video_data, temp_dir):
                     logger.warning(f"Failed to read frame {frame_idx}. Skipping.")
                     break
                 
-                # Keep a copy of the original frame for drawing detections
-                original_frame = frame.copy()
-                frame, padding_info = preprocess_frame(frame, original_shape=frame.shape)
+                frame = preprocess_frame(frame)
                 
                 for _ in range(frame_skip - 1):
                     if not cap.grab():
@@ -624,8 +576,6 @@ def process_video(video_data, temp_dir):
                 
                 batch_frames.append(frame)
                 batch_indices.append(frame_idx)
-                batch_padding_info.append(padding_info)
-                batch_original_frames.append(original_frame)
                 processed_frames += 1
 
             if not batch_frames:
@@ -633,9 +583,9 @@ def process_video(video_data, temp_dir):
                 break
 
             try:
-                # Convert frames to tensor and move to device
-                batch_frames_tensor = [torch.from_numpy(frame).permute(2, 0, 1).float() / 255.0 for frame in batch_frames]
-                batch_frames_tensor = torch.stack(batch_frames_tensor).to(device)
+                batch_frames_np = np.array(batch_frames)  # Shape: (batch, height, width, channels)
+                batch_frames_tensor = torch.from_numpy(batch_frames_np).permute(0, 3, 1, 2).float() / 255.0
+                batch_frames_tensor = batch_frames_tensor.to(device)
                 if device.type == "cuda":
                     batch_frames_tensor = batch_frames_tensor.half()
 
@@ -648,16 +598,17 @@ def process_video(video_data, temp_dir):
                 if device.type == "cuda":
                     torch.cuda.empty_cache()
 
-            for i, (result, frame_idx, padding_info, original_frame) in enumerate(zip(results, batch_indices, batch_padding_info, batch_original_frames)):
+            current_time = time.time()
+            if current_time - last_yield_time > 0.1:
+                progress = (processed_frames / total_frames) * 100
+                elapsed_time = current_time - start_time
+                fps_processed = processed_frames / elapsed_time if elapsed_time > 0 else 0
+                yield f"Processing video... {progress:.1f}% complete (Frame {processed_frames}/{total_frames}, {fps_processed:.1f} FPS)", "", "", "", ""
+                last_yield_time = current_time
+
+            for i, (result, frame_idx) in enumerate(zip(results, batch_indices)):
                 current_time = frame_idx / fps
                 
-                if time.time() - start_time > 0.5:
-                    progress = (processed_frames / total_frames) * 100
-                    elapsed_time = time.time() - start_time
-                    fps_processed = processed_frames / elapsed_time if elapsed_time > 0 else 0
-                    yield f"Processing video... {progress:.1f}% complete (Frame {processed_frames}/{total_frames}, {fps_processed:.1f} FPS)", "", "", "", ""
-                    start_time = time.time()
-
                 boxes = result.boxes
                 track_inputs = []
                 
@@ -673,8 +624,6 @@ def process_video(video_data, temp_dir):
                         continue
 
                     bbox = box.xywh.cpu().numpy()[0]
-                    # Adjust bounding box coordinates to original frame space
-                    bbox = adjust_bbox(bbox, padding_info)
                     track_inputs.append({
                         "bbox": bbox,
                         "conf": conf,
@@ -710,79 +659,92 @@ def process_video(video_data, temp_dir):
                     
                     if violation_key not in unique_violations:
                         unique_violations[violation_key] = current_time
-                        
-                        detection = {
-                            "worker_id": worker_id,
-                            "violation": label,
-                            "confidence": round(float(conf), 2),
-                            "bounding_box": bbox,
-                            "timestamp": current_time
-                        }
-                        
-                        # Use the original frame for drawing detections
-                        snapshot_frame = original_frame.copy()
-                        snapshot_frame = draw_detections(snapshot_frame, [detection])
-                        
-                        cv2.putText(
-                            snapshot_frame,
-                            f"Time: {current_time:.2f}s",
-                            (10, 30),
-                            cv2.FONT_HERSHEY_SIMPLEX,
-                            0.7,
-                            (255, 255, 255),
-                            2
-                        )
-                        
-                        snapshot_filename = f"violation_{label}_worker{worker_id}_{int(current_time*100)}.jpg"
-                        snapshot_path = os.path.join(output_dir, snapshot_filename)
-                        
-                        cv2.imwrite(
-                            snapshot_path,
-                            snapshot_frame,
-                            [cv2.IMWRITE_JPEG_QUALITY, CONFIG["SNAPSHOT_QUALITY"]]
-                        )
-                        
-                        snapshots.append({
-                            "violation": label,
-                            "worker_id": worker_id,
-                            "timestamp": current_time,
-                            "snapshot_path": snapshot_path,
-                            "snapshot_url": f"{CONFIG['PUBLIC_URL_BASE']}{snapshot_filename}",
-                            "confidence": round(float(conf), 2)
-                        })
-                        
-                        logger.info(f"Captured snapshot for {label} violation by worker {worker_id} at {current_time:.2f}s")
-
-                if len(snapshots) > 100:
-                    snapshots = snapshots[-10:]
+                        violation_frames[violation_key] = frame_idx
 
         cap.release()
         processing_time = time.time() - start_time
         logger.info(f"Processing complete in {processing_time:.2f}s")
 
-        logger.info(f"Snapshots: {snapshots}")
-
         violations = []
         for (worker_id, label), detection_time in unique_violations.items():
-            confidence = next(
-                (float(s["confidence"]) for s in snapshots if s["worker_id"] == worker_id and s["violation"] == label),
-                0.0
-            )
-            violation = {
+            violations.append({
                 "worker_id": worker_id,
                 "violation": label,
                 "timestamp": detection_time,
-                "confidence": confidence
-            }
-            violations.append(violation)
-
-        logger.info(f"Violations: {violations}")
+                "confidence": 0.0,  # Will be updated after reprocessing frames
+                "frame_idx": violation_frames[(worker_id, label)]
+            })
 
         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
 
+        # Reopen video to capture snapshots for violations
+        snapshots = []
+        cap = cv2.VideoCapture(video_path)
+        for violation in violations:
+            frame_idx = violation["frame_idx"]
+            cap.set(cv2.CAP_PROP_POS_FRAMES, frame_idx)
+            ret, frame = cap.read()
+            if not ret:
+                logger.warning(f"Failed to read frame {frame_idx} for snapshot.")
+                continue
+
+            frame = preprocess_frame(frame)
+            frame_tensor = torch.from_numpy(frame).permute(2, 0, 1).float() / 255.0
+            frame_tensor = frame_tensor.unsqueeze(0).to(device)
+            if device.type == "cuda":
+                frame_tensor = frame_tensor.half()
+
+            result = model(frame_tensor, device=device, conf=0.1, verbose=False)[0]
+            boxes = result.boxes
+
+            for box in boxes:
+                cls = int(box.cls)
+                conf = float(box.conf)
+                label = CONFIG["VIOLATION_LABELS"].get(cls, None)
+                if label == violation["violation"]:
+                    violation["confidence"] = round(conf, 2)
+                    bbox = box.xywh.cpu().numpy()[0]
+                    detection = {
+                        "worker_id": violation["worker_id"],
+                        "violation": label,
+                        "confidence": violation["confidence"],
+                        "bounding_box": bbox,
+                        "timestamp": violation["timestamp"]
+                    }
+                    snapshot_frame = frame.copy()
+                    snapshot_frame = draw_detections(snapshot_frame, [detection])
+                    cv2.putText(
+                        snapshot_frame,
+                        f"Time: {violation['timestamp']:.2f}s",
+                        (10, 30),
+                        cv2.FONT_HERSHEY_SIMPLEX,
+                        0.7,
+                        (255, 255, 255),
+                        2
+                    )
+                    snapshot_filename = f"violation_{label}_worker{violation['worker_id']}_{int(violation['timestamp']*100)}.jpg"
+                    snapshot_path = os.path.join(output_dir, snapshot_filename)
+                    cv2.imwrite(
+                        snapshot_path,
+                        snapshot_frame,
+                        [cv2.IMWRITE_JPEG_QUALITY, CONFIG["SNAPSHOT_QUALITY"]]
+                    )
+                    snapshots.append({
+                        "violation": label,
+                        "worker_id": violation["worker_id"],
+                        "timestamp": violation["timestamp"],
+                        "snapshot_path": snapshot_path,
+                        "snapshot_url": f"{CONFIG['PUBLIC_URL_BASE']}{snapshot_filename}",
+                        "confidence": violation["confidence"]
+                    })
+                    logger.info(f"Captured snapshot for {label} violation by worker {violation['worker_id']} at {violation['timestamp']:.2f}s")
+                    break
+
+        cap.release()
+
         score = calculate_safety_score(violations)
         pdf_path, pdf_url, pdf_file = generate_violation_pdf(violations, score, output_dir)
         
@@ -795,12 +757,7 @@ def process_video(video_data, temp_dir):
             display_name = CONFIG["DISPLAY_NAMES"].get(v.get("violation", "Unknown"), "Unknown")
             worker_id = v.get("worker_id", "Unknown")
             timestamp = v.get("timestamp", 0.0)
-            try:
-                confidence = float(v.get("confidence", 0.0))
-            except (ValueError, TypeError) as e:
-                logger.error(f"Invalid confidence value in violation {v}: {e}")
-                confidence = 0.0
-            
+            confidence = v.get("confidence", 0.0)
             violation_table += f"| {display_name} | {worker_id} | {timestamp:.2f} | {confidence:.2f} |\n"
 
         snapshots_text = ""
@@ -808,7 +765,6 @@ def process_video(video_data, temp_dir):
             display_name = CONFIG["DISPLAY_NAMES"].get(s["violation"], "Unknown")
             worker_id = s.get("worker_id", "Unknown")
             timestamp = s.get("timestamp", 0.0)
-            
             snapshots_text += f"### {display_name} - Worker {worker_id} at {timestamp:.2f}s\n\n"
             snapshots_text += f"![Violation]({s['snapshot_url']})\n\n"