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