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 'cookiefile': 'youtube_cookies.txt' } # 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.")