From 8909c9445f50b6c9cda4cede980e04b42ea1e8ad Mon Sep 17 00:00:00 2001 From: John Pope Date: Fri, 11 Oct 2024 07:30:43 +1100 Subject: [PATCH] ok --- extract_frames.py | 39 ++++++++++++++++++++++++++++++++------- 1 file changed, 32 insertions(+), 7 deletions(-) diff --git a/extract_frames.py b/extract_frames.py index c87db97..5be28d3 100644 --- a/extract_frames.py +++ b/extract_frames.py @@ -3,13 +3,30 @@ from tqdm import tqdm import random -def extract_frames(video_path, output_folder, frame_skip=0, max_frames=None): +def extract_quadrants(image): + height, width = image.shape[:2] + mid_x, mid_y = width // 2, height // 2 + + top_left = image[0:mid_y, 0:mid_x] + top_right = image[0:mid_y, mid_x:width] + bottom_left = image[mid_y:height, 0:mid_x] + bottom_right = image[mid_y:height, mid_x:width] + + return [top_left, top_right, bottom_left, bottom_right] + +def extract_frames(video_path, output_folder, frame_skip=0, max_frames=None, save_quadrants=False): video = cv2.VideoCapture(video_path) frame_count = int(video.get(cv2.CAP_PROP_FRAME_COUNT)) if not os.path.exists(output_folder): os.makedirs(output_folder) + if save_quadrants: + for i in range(4): + quadrant_folder = os.path.join(output_folder + f"_{i}") + if not os.path.exists(quadrant_folder): + os.makedirs(quadrant_folder) + success, image = video.read() count = 0 saved_count = 0 @@ -17,8 +34,15 @@ def extract_frames(video_path, output_folder, frame_skip=0, max_frames=None): for _ in tqdm(range(frame_count), desc=f"Processing {os.path.basename(video_path)}"): if success: if count % (frame_skip + 1) == 0: - frame_name = f"frame_{saved_count:06d}.jpg" + frame_name = f"frame_{saved_count:06d}.png" cv2.imwrite(os.path.join(output_folder, frame_name), image) + + if save_quadrants: + quadrants = extract_quadrants(image) + for i, quadrant in enumerate(quadrants): + quadrant_folder = os.path.join(output_folder + f"_{i}") + cv2.imwrite(os.path.join(quadrant_folder, frame_name), quadrant) + saved_count += 1 if max_frames and saved_count >= max_frames: break @@ -28,7 +52,7 @@ def extract_frames(video_path, output_folder, frame_skip=0, max_frames=None): video.release() print(f"Extracted {saved_count} frames from {os.path.basename(video_path)}") -def process_videos(input_folder, output_base_folder, frame_skip=0, max_videos=1000, max_frames_per_video=100): +def process_videos(input_folder, output_base_folder, frame_skip=0, max_videos=1000, max_frames_per_video=100, save_quadrants=False): video_files = [] for root, _, files in os.walk(input_folder): for file in files: @@ -43,13 +67,14 @@ def process_videos(input_folder, output_base_folder, frame_skip=0, max_videos=10 relative_path = os.path.relpath(video_path, input_folder) video_name = os.path.splitext(relative_path)[0] output_folder = os.path.join(output_base_folder, video_name) - extract_frames(video_path, output_folder, frame_skip, max_frames_per_video) + extract_frames(video_path, output_folder, frame_skip, max_frames_per_video, save_quadrants) # Usage -input_folder = '/media/oem/12TB/Downloads/CelebV-HQ/celebvhq/35666' -output_base_folder = '/media/oem/12TB/Downloads/CelebV-HQ/celebvhq/35666/images' +input_folder = './' +output_base_folder = './data' frame_skip = 0 # Adjust this if you want to skip frames max_videos = 100 # Limit the number of videos to process max_frames_per_video = 1000 # Limit the number of frames per video +save_quadrants = True # Set to True to save quadrants as separate PNG files -process_videos(input_folder, output_base_folder, frame_skip, max_videos, max_frames_per_video) \ No newline at end of file +process_videos(input_folder, output_base_folder, frame_skip, max_videos, max_frames_per_video, save_quadrants) \ No newline at end of file