Spaces:
Sleeping
Sleeping
| import cv2 | |
| import numpy as np | |
| import matplotlib.pyplot as plt | |
| import mediapipe as mp | |
| from app.face_utils import get_box | |
| mp_face_mesh = mp.solutions.face_mesh | |
| def preprocess_video_and_predict_sleep_quality(video): | |
| cap = cv2.VideoCapture(video) | |
| w = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)) | |
| h = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)) | |
| fps = np.round(cap.get(cv2.CAP_PROP_FPS)) | |
| path_save_video_original = 'result_original.mp4' | |
| path_save_video_face = 'result_face.mp4' | |
| path_save_video_sleep = 'result_sleep.mp4' | |
| vid_writer_original = cv2.VideoWriter(path_save_video_original, cv2.VideoWriter_fourcc(*'mp4v'), fps, (w, h)) | |
| vid_writer_face = cv2.VideoWriter(path_save_video_face, cv2.VideoWriter_fourcc(*'mp4v'), fps, (224, 224)) | |
| vid_writer_sleep = cv2.VideoWriter(path_save_video_sleep, cv2.VideoWriter_fourcc(*'mp4v'), fps, (224, 224)) | |
| frames = [] | |
| sleep_quality_scores = [] | |
| eye_bags_images = [] | |
| with mp_face_mesh.FaceMesh( | |
| max_num_faces=1, | |
| refine_landmarks=False, | |
| min_detection_confidence=0.5, | |
| min_tracking_confidence=0.5) as face_mesh: | |
| while cap.isOpened(): | |
| ret, frame = cap.read() | |
| if not ret: | |
| break | |
| frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB) | |
| results = face_mesh.process(frame_rgb) | |
| if results.multi_face_landmarks: | |
| for fl in results.multi_face_landmarks: | |
| startX, startY, endX, endY = get_box(fl, w, h) | |
| cur_face = frame_rgb[startY:endY, startX:endX] | |
| sleep_quality_score, eye_bags_image = analyze_sleep_quality(cur_face) | |
| sleep_quality_scores.append(sleep_quality_score) | |
| eye_bags_images.append(cv2.resize(eye_bags_image, (224, 224))) | |
| sleep_quality_viz = create_sleep_quality_visualization(cur_face, sleep_quality_score) | |
| cur_face = cv2.resize(cur_face, (224, 224)) | |
| vid_writer_face.write(cv2.cvtColor(cur_face, cv2.COLOR_RGB2BGR)) | |
| vid_writer_sleep.write(sleep_quality_viz) | |
| vid_writer_original.write(frame) | |
| frames.append(len(frames) + 1) | |
| cap.release() | |
| vid_writer_original.release() | |
| vid_writer_face.release() | |
| vid_writer_sleep.release() | |
| sleep_stat = sleep_quality_statistics_plot(frames, sleep_quality_scores) | |
| if eye_bags_images: | |
| average_eye_bags_image = np.mean(np.array(eye_bags_images), axis=0).astype(np.uint8) | |
| else: | |
| average_eye_bags_image = np.zeros((224, 224, 3), dtype=np.uint8) | |
| return (path_save_video_original, path_save_video_face, path_save_video_sleep, | |
| average_eye_bags_image, sleep_stat) | |
| def analyze_sleep_quality(face_image): | |
| # Placeholder function - implement your sleep quality analysis here | |
| sleep_quality_score = np.random.random() | |
| eye_bags_image = cv2.resize(face_image, (224, 224)) | |
| return sleep_quality_score, eye_bags_image | |
| def create_sleep_quality_visualization(face_image, sleep_quality_score): | |
| viz = face_image.copy() | |
| cv2.putText(viz, f"Sleep Quality: {sleep_quality_score:.2f}", (10, 30), | |
| cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 2) | |
| return cv2.cvtColor(viz, cv2.COLOR_RGB2BGR) | |
| def sleep_quality_statistics_plot(frames, sleep_quality_scores): | |
| fig, ax = plt.subplots() | |
| ax.plot(frames, sleep_quality_scores) | |
| ax.set_xlabel('Frame') | |
| ax.set_ylabel('Sleep Quality Score') | |
| ax.set_title('Sleep Quality Over Time') | |
| plt.tight_layout() | |
| return fig |