prepare_video.py

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()