| |
|
|
| import cv2 |
| import os |
| from tqdm import tqdm |
|
|
| def extract_frames(video_path, output_dir='frames', prefix='frame'): |
| """ |
| Extracts frames from a video file. |
| |
| Args: |
| video_path (str): Path to the input video file. |
| output_dir (str): Directory to save extracted frames. |
| prefix (str): Prefix for the frame filenames. |
| """ |
| if not os.path.exists(output_dir): |
| os.makedirs(output_dir) |
| |
| cap = cv2.VideoCapture(video_path) |
| |
| if not cap.isOpened(): |
| print(f"Error: Could not open video {video_path}.") |
| return |
| |
| frame_count = 0 |
| frame_total = int(cap.get(cv2.CAP_PROP_FRAME_COUNT)) |
| |
| for _ in tqdm(range(frame_total), desc=f"Extracting frames from {os.path.basename(video_path)}"): |
| ret, frame = cap.read() |
| if not ret: |
| break |
| frame_filename = os.path.join(output_dir, f"{prefix}_frame_{frame_count}.jpg") |
| cv2.imwrite(frame_filename, frame) |
| frame_count += 1 |
| |
| cap.release() |
| print(f"Extracted {frame_count} frames from {video_path}.") |
|
|
| def process_all_videos(videos_dir='videos', frames_dir='frames'): |
| """ |
| Processes all videos in the specified directory and extracts frames. |
| |
| Args: |
| videos_dir (str): Directory containing video files. |
| frames_dir (str): Directory to save extracted frames. |
| """ |
| if not os.path.exists(frames_dir): |
| os.makedirs(frames_dir) |
| |
| video_files = [f for f in os.listdir(videos_dir) if f.endswith('.avi') or f.endswith('.mp4')] |
| |
| for video_file in video_files: |
| label = video_file.split('_')[0] |
| video_path = os.path.join(videos_dir, video_file) |
| output_subdir = os.path.join(frames_dir, label) |
| if not os.path.exists(output_subdir): |
| os.makedirs(output_subdir) |
| extract_frames(video_path, output_dir=output_subdir, prefix=video_file.split('.')[0]) |
|
|
| if __name__ == "__main__": |
| process_all_videos(videos_dir='videos', frames_dir='frames') |
|
|
