| | 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.") |