sample_youtube_video.py

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.")