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