-
Notifications
You must be signed in to change notification settings - Fork 34
/
prepare_video.py
executable file
·150 lines (129 loc) · 7.57 KB
/
prepare_video.py
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
149
150
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, use_full_clip=False): # Load the video file
video = VideoFileClip(input_video_path)
if use_full_clip:
cropped_video = video
else:
# 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
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, use_full_clip=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, use_full_clip)
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')
parser.add_argument('--use_full_clip', action='store_true', help='Use the full video clip without trimming')
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,
args.use_full_clip)
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,
args.use_full_clip)
if __name__ == "__main__":
main()