panda70m_downloader.py

import os
import shutil
import pandas as pd
from vidfetch import compress_folder, pull_from_hf


def download_video_links(hf_token: str, filename: str, save_dir: str):
 # check save dir
 if not os.path.exists(save_dir):
 os.makedirs(save_dir)
 
 # download
 pull_from_hf(
 hf_token=hf_token,
 hf_repo_id="OpenVideo/Panda-70M-Original-Links",
 filename=filename,
 save_dir=save_dir
 )
 

def download_videos_by_csv(
 csv_file_path: str,
 save_dir: str, 
 targz_filename: str,
):
 try:
 import youtube_dl
 except:
 raise ModuleNotFoundError(
 "youtube_dl missed, please install it by ``vidfetch.package.youtube.youtube_dl_install_helper``"
 )
 # path/dir
 folder_name = targz_filename.replace(".tar.gz", "")
 download_videos_dir = os.path.join(save_dir, folder_name, "download_raw") 
 log_path = os.path.join(download_videos_dir, "log.txt")
 targz_path = os.path.join(save_dir, targz_filename)
 
 # make dirs
 if not os.path.exists(download_videos_dir):
 os.makedirs(download_videos_dir)
 
 # read from csv
 csv_filename = os.path.basename(csv_file_path)
 shutil.copy(src=csv_file_path, dst=os.path.join(download_videos_dir, csv_filename))
 data = pd.read_csv(csv_file_path)
 links = data["url"].tolist()
 videos_id = data["videoID"].to_list() 

 failed_links = [] # record failed links
 for link, video_id in zip(links, videos_id):
 # check if downloaded
 video_save_path = os.path.join(download_videos_dir, video_id[1:]+".mp4")
 if os.path.exists(video_save_path):
 continue
 
 # download
 ydl_opts = {
 'format': 'best',
 'quiet': False,
 'outtmpl': os.path.join(download_videos_dir, video_id[1:]+".mp4"),
 }
 with youtube_dl.YoutubeDL(ydl_opts) as ydl:
 try:
 ydl.download([link])
 except:
 failed_links.append(link)

 # delete videos larger than 100MB
 video_files = os.listdir(download_videos_dir)
 delete_videos = []
 for file in video_files:
 file_path = os.path.join(download_videos_dir, file)
 file_size_mb = os.path.getsize(file_path) / (1024 * 1024) # Convert to megabytes
 if file_size_mb > 500:
 delete_videos.append(file_path)
 os.remove(file_path)
 
 # Write to log file
 with open(log_path, 'w') as file:
 file.write('Fail to download\n')
 file.write('\n'.join(failed_links))
 file.write('Delete videos larger than 500MB\n')
 file.write('\n'.join(failed_links))
 
 compress_folder(download_videos_dir, targz_path)