File size: 7,383 Bytes
26853cd |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 |
import argparse
from moviepy.editor import VideoFileClip
import os
import glob
import random
import numpy as np
from PIL import Image
def extract_frames(video_path, frame_count=16):
clip = VideoFileClip(video_path)
duration = clip.duration
frames = []
# Calculate the time interval at which to extract frames
times = np.linspace(0, duration, frame_count, endpoint=False)
for t in times:
# Extract the frame at the specific timestamp
frame = clip.get_frame(t)
# Convert the frame (numpy array) to a PIL Image
pil_img = Image.fromarray(frame)
frames.append(pil_img)
return frames
def crop_and_resize_video(input_video_path, output_folder, clip_duration, width=None, height=None, start_time=None, end_time=None, n_frames=16, center_crop=False, x_offset=0, y_offset=0, longest_to_width=False): # Load the video file
video = VideoFileClip(input_video_path)
# Calculate start and end times for cropping
if start_time is not None:
start_time = float(start_time)
end_time = start_time + clip_duration
elif end_time is not None:
end_time = float(end_time)
start_time = end_time - clip_duration
else:
# Default to random cropping if neither start nor end time is specified
video_duration = video.duration
if video_duration <= clip_duration:
print(f"Skipping {input_video_path}: duration is less than or equal to the clip duration.")
return
max_start_time = video_duration - clip_duration
start_time = random.uniform(0, max_start_time)
end_time = start_time + clip_duration
# Crop the video
cropped_video = video.subclip(start_time, end_time)
if center_crop:
# Calculate scale to ensure the desired crop size fits within the video
video_width, video_height = cropped_video.size
scale_width = video_width / width
scale_height = video_height / height
if longest_to_width:
scale = max(scale_width, scale_height)
else:
scale = min(scale_width, scale_height)
# Resize video to ensure the crop area fits within the frame
# This step ensures that the smallest dimension matches or exceeds 512 pixels
new_width = int(video_width / scale)
new_height = int(video_height / scale)
resized_video = cropped_video.resize(newsize=(new_width, new_height))
print(f"Resized video to ({new_width}, {new_height})")
# Calculate crop position with offset, ensuring the crop does not go out of bounds
# The offset calculation needs to ensure that the cropping area remains within the video frame
offset_x = int(((x_offset + 1) / 2) * (new_width - width)) # Adjusted for [-1, 1] scale
offset_y = int(((y_offset + 1) / 2) * (new_height - height)) # Adjusted for [-1, 1] scale
# Ensure offsets do not push the crop area out of the video frame
offset_x = max(0, min(new_width - width, offset_x))
offset_y = max(0, min(new_height - height, offset_y))
# Apply center crop with offsets
cropped_video = resized_video.crop(x1=offset_x, y1=offset_y, width=width, height=height)
elif width and height:
# Directly resize the video to specified width and height if no center crop is specified
cropped_video = cropped_video.resize(newsize=(width, height))
# After resizing and cropping, set the frame rate to fps
fps = n_frames // clip_duration
final_video = cropped_video.set_fps(fps)
# Prepare the output video path
if not os.path.exists(output_folder):
os.makedirs(output_folder)
filename = os.path.basename(input_video_path)
output_video_path = os.path.join(output_folder, filename)
# Write the result to the output file
final_video.write_videofile(output_video_path, codec='libx264', audio_codec='aac', fps=fps)
print(f"Processed {input_video_path}, saved to {output_video_path}")
return output_video_path
def process_videos(input_folder, output_base_folder, clip_duration, width=None, height=None, start_time=None, end_time=None, n_frames=16, center_crop=False, x_offset=0, y_offset=0, longest_to_width=False):
video_files = glob.glob(os.path.join(input_folder, '*.mp4')) # Adjust the pattern if needed
if video_files == []:
print(f"No video files found in {input_folder}")
return
for video_file in video_files:
crop_and_resize_video(video_file, output_base_folder, clip_duration, width, height, start_time, end_time, n_frames, center_crop, x_offset, y_offset, longest_to_width)
return
def main():
parser = argparse.ArgumentParser(description='Crop and resize video segments.')
parser.add_argument('--input_folder', type=str, help='Path to the input folder containing video files')
parser.add_argument('--video_path', type=str, default=None, required=False, help='Path to the input video file')
parser.add_argument('--output_folder', type=str, default="processed_video_data", help='Path to the folder for the output videos')
parser.add_argument('--clip_duration', type=int, default=2, required=False, help='Duration of the video clips in seconds')
parser.add_argument('--width', type=int, default=512, help='Width of the output video (optional)')
parser.add_argument('--height', type=int, default=512, help='Height of the output video (optional)')
parser.add_argument('--start_time', type=float, help='Start time for cropping (optional)')
parser.add_argument('--end_time', type=float, help='End time for cropping (optional)')
parser.add_argument('--n_frames', type=int, default=16, help='Number of frames to extract from each video')
parser.add_argument('--center_crop', action='store_true', help='Center crop the video')
parser.add_argument('--x_offset', type=float, default=0, required=False, help='Horizontal offset for center cropping, range -1 to 1 (optional)')
parser.add_argument('--y_offset', type=float, default=0, required=False, help='Vertical offset for center cropping, range -1 to 1 (optional)')
parser.add_argument('--longest_to_width', action='store_true', help='Resize the longest dimension to the specified width')
args = parser.parse_args()
if args.start_time and args.end_time:
print("Please specify only one of start_time or end_time, not both.")
return
if args.video_path:
crop_and_resize_video(args.video_path,
args.output_folder,
args.clip_duration,
args.width, args.height,
args.start_time, args.end_time,
args.n_frames,
args.center_crop, args.x_offset, args.y_offset, args.longest_to_width)
else:
process_videos(args.input_folder,
args.output_folder,
args.clip_duration,
args.width, args.height,
args.start_time, args.end_time,
args.n_frames,
args.center_crop, args.x_offset, args.y_offset, args.longest_to_width)
if __name__ == "__main__":
main()
|