Spaces:
Runtime error
Runtime error
File size: 3,652 Bytes
fc286f6 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 |
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 |