reab5555 commited on
Commit
8be35d2
·
verified ·
1 Parent(s): af91dd4

Update video_processing.py

Browse files
Files changed (1) hide show
  1. 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.5)
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"