import video_handler as vh
import frame_handler_yolo as fh

model_path = "yolov8n.pt"  # YOLOv8 model path
original_fps = 30  # Original FPS of the input videos

video_path = "../input_data/football.mp4"  # Replace with your video path

output_folder = "../output_data/"  # Folder to save extracted frames
all_frames_folder = output_folder + "all_frames"

# =============  Step 1: Extract frames from the input videos  ========

vh.extract_all_frames(video_path, all_frames_folder)

#frame_rate = 2  # Extract 2 frames per second
#vh.extract_frames_by_rate(video_path, all_frames_folder, frame_rate)


#============== Step 2: Extract key frames from the extracted frames  ========
# key frames = frames contains a ball
# if a previous frame of a key frame is a non-key frame - major movement detected
# -> reclassify up to 30 previous frames (~ 1 second) as key frames to add context the major movement

key_frames_folder = output_folder + "key_frames"  # Save key frames here
nonkey_frames_folder = output_folder + "nonkey_frames"  # Save non-key frames here

fh.extract_key_frames(all_frames_folder, key_frames_folder, original_fps, model_path)

#============== Step 3: Crop the key frames to align with 9:16 ratio aspect while keeping the key object - football ball
key_frames_9_16_folder = output_folder + "key_frames_9_16"  # Save processed frames here

target_resolution = (360, 640)  # Output resolution (9:16)

fh.crop_preserve_key_objects(key_frames_folder, key_frames_9_16_folder, model_path, target_resolution)

#============== Step 4: Create a video from the processed frames  ========

output_video_path_9_16 = output_folder + "output_video_9_16.mp4"  # Output video path
target_frame_rate = 30  # Frames per second of the output videos
vh.create_video_from_frames(key_frames_9_16_folder, output_video_path_9_16, target_frame_rate, target_resolution)