Spaces:
Sleeping
Sleeping
File size: 2,468 Bytes
4c80c04 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 |
import cv2
import os
import logging
import argparse
# Setup logging
logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s")
logger = logging.getLogger(__name__)
def extract_frames(video_path, output_dir, frame_interval=30):
try:
os.makedirs(output_dir, exist_ok=True)
logger.info(f"Output directory: {output_dir}")
video = cv2.VideoCapture(video_path)
if not video.isOpened():
raise ValueError(f"Could not open video: {video_path}")
fps = video.get(cv2.CAP_PROP_FPS)
total_frames = int(video.get(cv2.CAP_PROP_FRAME_COUNT))
logger.info(f"Video: {video_path}, FPS: {fps}, Total frames: {total_frames}")
frame_count = 0
saved_count = 0
video_name = os.path.splitext(os.path.basename(video_path))[0]
while True:
ret, frame = video.read()
if not ret:
break
if frame_count % frame_interval == 0:
frame_path = os.path.join(output_dir, f"{video_name}_frame_{frame_count}.jpg")
cv2.imwrite(frame_path, frame)
saved_count += 1
logger.info(f"Saved frame: {frame_path}")
frame_count += 1
video.release()
logger.info(f"Extracted {saved_count} frames from {video_path}")
except Exception as e:
logger.error(f"Error extracting frames from {video_path}: {e}")
raise
def main():
parser = argparse.ArgumentParser(description="Extract frames from videos for training dataset.")
parser.add_argument("--video-dir", default="videos", help="Directory containing input videos")
parser.add_argument("--output-dir", default="dataset/images/train", help="Directory to save extracted frames")
parser.add_argument("--frame-interval", type=int, default=30, help="Extract every nth frame")
args = parser.parse_args()
video_dir = args.video_dir
output_dir = args.output_dir
frame_interval = args.frame_interval
if not os.path.exists(video_dir):
logger.error(f"Video directory does not exist: {video_dir}")
return
for video_file in os.listdir(video_dir):
if video_file.endswith((".mp4", ".avi", ".mov")):
video_path = os.path.join(video_dir, video_file)
extract_frames(video_path, output_dir, frame_interval)
if __name__ == "__main__":
logger.info("Starting frame extraction...")
main()
|