Update video_processing.py
Browse files- video_processing.py +7 -4
video_processing.py
CHANGED
@@ -21,7 +21,7 @@ device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
|
21 |
mtcnn = MTCNN(keep_all=False, device=device, thresholds=[0.9, 0.9, 0.9], min_face_size=50)
|
22 |
|
23 |
mp_face_mesh = mp.solutions.face_mesh
|
24 |
-
face_mesh = mp_face_mesh.FaceMesh(static_image_mode=False, max_num_faces=1, min_detection_confidence=0.
|
25 |
|
26 |
mp_pose = mp.solutions.pose
|
27 |
pose = mp_pose.Pose(static_image_mode=False, min_detection_confidence=0.6, min_tracking_confidence=0.6)
|
@@ -52,6 +52,7 @@ def process_frames(frames_folder, aligned_faces_folder, frame_count, progress):
|
|
52 |
embeddings_by_frame = {}
|
53 |
posture_scores_by_frame = {}
|
54 |
posture_landmarks_by_frame = {}
|
|
|
55 |
aligned_face_paths = []
|
56 |
frame_files = sorted([f for f in os.listdir(frames_folder) if f.endswith('.jpg')])
|
57 |
|
@@ -64,7 +65,9 @@ def process_frames(frames_folder, aligned_faces_folder, frame_count, progress):
|
|
64 |
posture_score, posture_landmarks = calculate_posture_score(frame)
|
65 |
posture_scores_by_frame[frame_num] = posture_score
|
66 |
posture_landmarks_by_frame[frame_num] = posture_landmarks
|
67 |
-
|
|
|
|
|
68 |
boxes, probs = mtcnn.detect(frame)
|
69 |
|
70 |
if boxes is not None and len(boxes) > 0 and probs[0] >= 0.99:
|
@@ -85,8 +88,8 @@ def process_frames(frames_folder, aligned_faces_folder, frame_count, progress):
|
|
85 |
|
86 |
progress((i + 1) / len(frame_files), f"Processing frame {i + 1} of {len(frame_files)}")
|
87 |
|
88 |
-
return embeddings_by_frame, posture_scores_by_frame, posture_landmarks_by_frame, aligned_face_paths
|
89 |
-
|
90 |
def process_video(video_path, anomaly_threshold, desired_fps, progress=None):
|
91 |
start_time = time.time()
|
92 |
output_folder = "output"
|
|
|
21 |
mtcnn = MTCNN(keep_all=False, device=device, thresholds=[0.9, 0.9, 0.9], min_face_size=50)
|
22 |
|
23 |
mp_face_mesh = mp.solutions.face_mesh
|
24 |
+
face_mesh = mp_face_mesh.FaceMesh(static_image_mode=False, max_num_faces=1, min_detection_confidence=0.6)
|
25 |
|
26 |
mp_pose = mp.solutions.pose
|
27 |
pose = mp_pose.Pose(static_image_mode=False, min_detection_confidence=0.6, min_tracking_confidence=0.6)
|
|
|
52 |
embeddings_by_frame = {}
|
53 |
posture_scores_by_frame = {}
|
54 |
posture_landmarks_by_frame = {}
|
55 |
+
facial_landmarks_by_frame = {}
|
56 |
aligned_face_paths = []
|
57 |
frame_files = sorted([f for f in os.listdir(frames_folder) if f.endswith('.jpg')])
|
58 |
|
|
|
65 |
posture_score, posture_landmarks = calculate_posture_score(frame)
|
66 |
posture_scores_by_frame[frame_num] = posture_score
|
67 |
posture_landmarks_by_frame[frame_num] = posture_landmarks
|
68 |
+
if results.multi_face_landmarks:
|
69 |
+
facial_landmarks_by_frame[frame_num] = results.multi_face_landmarks[0]
|
70 |
+
|
71 |
boxes, probs = mtcnn.detect(frame)
|
72 |
|
73 |
if boxes is not None and len(boxes) > 0 and probs[0] >= 0.99:
|
|
|
88 |
|
89 |
progress((i + 1) / len(frame_files), f"Processing frame {i + 1} of {len(frame_files)}")
|
90 |
|
91 |
+
return embeddings_by_frame, posture_scores_by_frame, posture_landmarks_by_frame, aligned_face_paths, facial_landmarks_by_frame
|
92 |
+
|
93 |
def process_video(video_path, anomaly_threshold, desired_fps, progress=None):
|
94 |
start_time = time.time()
|
95 |
output_folder = "output"
|