vidaio / scripts /capture_video_frames.py
DavidHstar's picture
Update model files
f72237c verified
import cv2
import os
import argparse
import numpy as np
from pathlib import Path
def capture_frames(video_path, output_dir, num_frames=3):
"""
Captures the first num_frames frames from a video file and resizes them to 128x128 pixels.
Args:
video_path (str): Path to the video file
output_dir (str): Directory to save the captured frames
num_frames (int): Number of frames to capture from the start
Returns:
list: Paths to the saved frames
"""
# Check if video exists
if not os.path.isfile(video_path):
raise FileNotFoundError(f"Video file not found: {video_path}")
# Create output directory if it doesn't exist
os.makedirs(output_dir, exist_ok=True)
# Get video filename without extension
video_name = Path(video_path).stem
# Open the video file
cap = cv2.VideoCapture(video_path)
if not cap.isOpened():
raise Exception(f"Error opening video file: {video_path}")
# Get video properties
width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
fps = cap.get(cv2.CAP_PROP_FPS)
total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
print(f"Video Info:")
print(f" Resolution: {width}x{height}")
print(f" FPS: {fps:.2f}")
print(f" Total Frames: {total_frames}")
# Adjust num_frames if video is shorter
if total_frames < num_frames:
print(f"Warning: Video has only {total_frames} frames, adjusting to capture {total_frames} frames.")
num_frames = total_frames
saved_frames = []
for i in range(num_frames):
ret, frame = cap.read()
if not ret:
print(f"Error reading frame {i}")
break
resized_frame = cv2.resize(frame, (10, 10), interpolation=cv2.INTER_AREA)
# Save resized frame
frame_path = os.path.join(output_dir, f"{video_name}_frame_{i:03d}.png")
cv2.imwrite(frame_path, resized_frame)
saved_frames.append(frame_path)
print(f"Saved frame {i+1}/{num_frames}: {frame_path}")
cap.release()
return saved_frames
def main():
parser = argparse.ArgumentParser(description="Capture and resize frames from a video to 128x128 pixels")
parser.add_argument("--video", required=True, help="Path to the video file")
parser.add_argument("--output", default="frames", help="Directory to save the frames")
parser.add_argument("--frames", type=int, default=3, help="Number of frames to capture from the start")
args = parser.parse_args()
try:
saved_frames = capture_frames(args.video, args.output, args.frames)
print(f"Successfully captured {len(saved_frames)} frames from the video")
except Exception as e:
print(f"Error: {e}")
if __name__ == "__main__":
main()