Spaces:
Sleeping
Sleeping
File size: 2,954 Bytes
f66d8b7 |
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 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 |
import cv2
import yt_dlp
import os
from typing import List
from utils import encode_frame_to_base64
# Nota: Precisa ter a biblioteca ffmpeg instalada na maquina!
def sample_youtube_video(youtube_url: str, sample_rate: int) -> List[str]:
"""
Downloads a YouTube video and samples frames from it at a specified interval.
Args:
youtube_url (str): The URL of the YouTube video.
sample_rate (int): The interval (in seconds) at which to sample frames.
Returns:
List[cv2.Mat]: A list of sampled frames as cv2.Mat objects. Returns an empty list on error.
"""
frames: List[cv2.Mat] = []
# Construct the yt_dlp options
ydl_opts = {
'format': 'bestvideo[ext=mp4]+bestaudio[ext=m4a]/mp4', # Get the best quality video and audio in mp4 format
'outtmpl': '%(title)s.%(ext)s', # Save with the video title
'quiet': True, # Suppress console output
}
# Download the video
try:
with yt_dlp.YoutubeDL(ydl_opts) as ydl:
info_dict = ydl.extract_info(youtube_url, download=True)
video_file = ydl.prepare_filename(info_dict) # Get the downloaded filename
except Exception as e:
print(f"Error downloading video: {e}")
return []
# Open the video file
try:
cap = cv2.VideoCapture(video_file)
except Exception as e:
print(f"Error opening video file: {e}")
return []
# Get the video's frame rate
fps = cap.get(cv2.CAP_PROP_FPS)
if fps <= 0:
print("Error: Could not determine video frame rate. OpenCV may not be able to read this file.")
cap.release()
os.remove(video_file)
return []
# Calculate the frame interval
frame_interval = int(sample_rate * fps)
frame_count = 0
success = True
duration_seconds = float(info_dict.get('duration', 0))
while success:
success, frame = cap.read()
if not success:
break
if frame_count % frame_interval == 0:
# Get the current time stamp
current_time = cap.get(cv2.CAP_PROP_POS_MSEC) / 1000.0 # in seconds
if current_time > duration_seconds and duration_seconds != 0:
break
frames.append(frame)
frame_count += 1
# Release the video capture object and remove the downloaded video file
cap.release()
try:
os.remove(video_file)
print(f"Deleted downloaded video file: {video_file}")
except Exception as e:
print(f"Error deleting video file: {e}")
print(f"Finished sampling. {len(frames)} frames sampled.")
return [encode_frame_to_base64(frame) for frame in frames]
if __name__ == "__main__":
sampled_frames = sample_youtube_video("https://www.youtube.com/watch?v=L1vXCYZAYYM", 5)
if sampled_frames:
print(f"Number of frames returned: {len(sampled_frames)}")
else:
print("No frames were returned.")
|