| import cv2 |
| from ultralytics import YOLO |
| import json |
| import os |
|
|
| def analyze_video_for_ppe(video_path, model_path='yolov8n.pt', frames_per_sec=1.0): |
| """ |
| Analyzes a video for PPE compliance using a YOLOv8 model. |
| """ |
| |
| |
| try: |
| model = YOLO(model_path) |
| except Exception as e: |
| print(f"Error loading model: {e}. Ensure you have a valid YOLOv8 model path.") |
| return [] |
|
|
| |
| cap = cv2.VideoCapture(video_path) |
| if not cap.isOpened(): |
| print(f"Error: Could not open video file {video_path}") |
| return [] |
|
|
| |
| fps = cap.get(cv2.CAP_PROP_FPS) |
| frame_interval = int(fps / frames_per_sec) |
| frame_count = 0 |
| analysis_results = [] |
| |
| |
|
|
| print(f"Video FPS: {fps}, Analyzing every {frame_interval} frames...") |
|
|
| while cap.isOpened(): |
| |
| ret, frame = cap.read() |
|
|
| if not ret: |
| break |
|
|
| |
| if frame_count % frame_interval == 0: |
| timestamp_sec = frame_count / fps |
| |
| |
| results = model(frame, verbose=False) |
| |
| |
| detections = [] |
| for r in results: |
| |
| for box in r.boxes.data.tolist(): |
| x1, y1, x2, y2, conf, cls = box |
| label = model.names[int(cls)] |
| |
| |
| detections.append({ |
| 'label': label, |
| 'confidence': round(conf, 2), |
| 'bbox': [int(x1), int(y1), int(x2), int(y2)] |
| }) |
| |
| |
| analysis_results.append({ |
| 'video_id': os.path.basename(video_path), |
| 'frame_id': frame_count, |
| 'timestamp_sec': round(timestamp_sec, 2), |
| 'detections': detections |
| }) |
|
|
| frame_count += 1 |
|
|
| |
| cap.release() |
| print(f"Analysis complete. Total frames analyzed: {len(analysis_results)}") |
| return analysis_results |
|
|
| |
| if __name__ == '__main__': |
| |
| |
| |
| if not os.path.exists('construction.mp4'): |
| print("Please place a video file named 'construction.mp4' in the current directory.") |
| else: |
| results = analyze_video_for_ppe('construction.mp4', frames_per_sec=0.5) |
| |
| |
| with open('raw_analysis.json', 'w') as f: |
| json.dump(results, f, indent=4) |
| |
| print(f"Raw analysis saved to raw_analysis.json. {len(results)} records created.") |