| import os |
| import cv2 |
| import numpy as np |
| import shutil |
| import argparse |
| import sys |
|
|
| def process_video(input_path: str, output_folder: str) -> str: |
| """Process a single video file to 14 frames""" |
| os.makedirs(output_folder, exist_ok=True) |
|
|
| video = cv2.VideoCapture(input_path) |
| frame_count = int(video.get(cv2.CAP_PROP_FRAME_COUNT)) |
| fps = video.get(cv2.CAP_PROP_FPS) |
| width = int(video.get(cv2.CAP_PROP_FRAME_WIDTH)) |
| height = int(video.get(cv2.CAP_PROP_FRAME_HEIGHT)) |
|
|
| frame_indices = np.linspace(0, frame_count - 1, num=14, dtype=int) |
| sampled_frames = [] |
| current_frame = 0 |
| frames_read = 0 |
|
|
| while True: |
| ret, frame = video.read() |
| if not ret: |
| break |
| if current_frame in frame_indices: |
| sampled_frames.append(frame) |
| frames_read += 1 |
| if frames_read == len(frame_indices): |
| break |
| current_frame += 1 |
|
|
| video.release() |
|
|
| output_video_path = os.path.join(output_folder, os.path.basename(input_path)) |
| fourcc = cv2.VideoWriter_fourcc(*'mp4v') |
| out_video = cv2.VideoWriter(output_video_path, fourcc, fps, (width, height)) |
|
|
| for frame in sampled_frames: |
| out_video.write(frame) |
|
|
| out_video.release() |
|
|
| |
| frame_folder = os.path.join(output_folder, os.path.splitext(os.path.basename(input_path))[0]) |
| os.makedirs(frame_folder, exist_ok=True) |
|
|
| for idx, frame in enumerate(sampled_frames): |
| frame_image_path = os.path.join(frame_folder, f'frame_{idx:03d}.png') |
| cv2.imwrite(frame_image_path, frame) |
|
|
| return output_video_path |
|
|
| def main(): |
| parser = argparse.ArgumentParser(description='Process video to 14 frames') |
| parser.add_argument('--input', type=str, required=True, help='Input video file') |
| parser.add_argument('--output', type=str, required=True, help='Output folder path') |
| |
| args = parser.parse_args() |
| |
| try: |
| output_path = process_video(args.input, args.output) |
| print(f"Successfully processed video to: {output_path}") |
| except Exception as e: |
| print(f"Error processing video: {str(e)}", file=sys.stderr) |
| sys.exit(1) |
|
|
| if __name__ == '__main__': |
| main() |
|
|
