fffiloni commited on
Commit
4264aae
1 Parent(s): 8ab8669

add extract_frames_with_labels function

Browse files
Files changed (1) hide show
  1. app.py +36 -0
app.py CHANGED
@@ -14,6 +14,42 @@ hf_hub_download(
14
  local_dir="checkpoint"
15
  )
16
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
17
  # Define a function to run your script with selected inputs
18
  def run_xportrait(
19
  model_config,
 
14
  local_dir="checkpoint"
15
  )
16
 
17
+ def extract_frames_with_labels(video_path, output_dir="frames"):
18
+ # Ensure output directory exists
19
+ os.makedirs(output_dir, exist_ok=True)
20
+
21
+ # Open the video file
22
+ video_capture = cv2.VideoCapture(video_path)
23
+ if not video_capture.isOpened():
24
+ raise ValueError(f"Cannot open video file: {video_path}")
25
+
26
+ frame_data = []
27
+ frame_index = 0
28
+
29
+ # Loop through the video frames
30
+ while True:
31
+ ret, frame = video_capture.read()
32
+ if not ret:
33
+ break # Exit the loop if there are no frames left to read
34
+
35
+ # Zero-padded frame index for filename and label
36
+ frame_label = f"{frame_index:04}"
37
+ frame_filename = os.path.join(output_dir, f"frame_{frame_label}.jpg")
38
+
39
+ # Save the frame as a .jpg file
40
+ cv2.imwrite(frame_filename, frame)
41
+
42
+ # Append the tuple (filename, label) to the list
43
+ frame_data.append((frame_filename, frame_label))
44
+
45
+ # Increment frame index
46
+ frame_index += 1
47
+
48
+ # Release the video capture object
49
+ video_capture.release()
50
+
51
+ return frame_data
52
+
53
  # Define a function to run your script with selected inputs
54
  def run_xportrait(
55
  model_config,