import yt_dlp
import os
import imageio

# Set the YouTube video URL and directories
video_url = 'https://www.youtube.com/watch?v=rfscVS0vtbw'  # Replace with another URL
data_root = 'C:/Users/Vishal R/OneDrive/Desktop/esophageal/data/'  # Replace with the correct path
video_dir = 'surgical_video_src'
output_dir = 'surgical_video_frames'

# Check if the root path exists
if not os.path.exists(data_root):
    print(f"The path {data_root} does not exist. Please check the path.")
    exit()

# Create directories for video and output frames if they don't exist
video_path = os.path.join(data_root, video_dir)
output_path = os.path.join(data_root, output_dir)

if not os.path.exists(video_path):
    os.makedirs(video_path)

if not os.path.exists(output_path):
    os.makedirs(output_path)

# Download the video using yt-dlp
ydl_opts = {
    'outtmpl': os.path.join(video_path, '%(id)s.%(ext)s'),  # Set the output file name
}

with yt_dlp.YoutubeDL(ydl_opts) as ydl:
    info_dict = ydl.extract_info(video_url, download=True)
    filename = ydl.prepare_filename(info_dict)

# Read the video using imageio
vid = imageio.get_reader(filename, 'ffmpeg')

# Define the start and end frames for extraction
start_frame = 340
end_frame = 500

# Loop through the frames and save them as images
for frame in range(start_frame, end_frame):
    image = vid.get_data(frame)
    output_filename = os.path.join(output_path, f"frame_{frame - start_frame:04d}.png")
    imageio.imwrite(output_filename, image)
    print(f"Saved: {output_filename}")

print("Video frames extracted successfully.")