[tab] data_download/download_scripts/download_videos.py (python)
import json
import requests
import os
import tyro
import logging
from urllib.parse import urljoin
from dataclasses import dataclass
from pathlib import Path
from tqdm import tqdm
from multiprocessing import Pool, cpu_count
import time

@dataclass
class DownloadVideos:
    index_file_path: str ='index_json/all_6xx_Jun_29.json'
    output_dir: str ='/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/data/open_ai_minecraft'

def download_single_file(args):
    """Download a single file - designed to be used with multiprocessing"""
    url, output_path, relpath = args
    
    if os.path.exists(output_path):
        return f"Skipped {relpath} (already exists)"
    
    # Create parent directories if they don't exist
    os.makedirs(os.path.dirname(output_path), exist_ok=True)
    
    try:
        response = requests.get(url, stream=True, timeout=30)
        if response.status_code == 200:
            file_size = 0
            with open(output_path, 'wb') as f:
                for chunk in response.iter_content(chunk_size=8192):
                    if chunk:
                        f.write(chunk)
                        file_size += len(chunk)
            
            # Convert to MB for logging
            file_size_mb = file_size / (1024 * 1024)
            return f"Downloaded {relpath} ({file_size_mb:.2f} MB)"
        else:
            return f"Failed to download {relpath}: HTTP {response.status_code}"
    except requests.exceptions.RequestException as e:
        return f"Request failed for {relpath}: {e}"
    except Exception as e:
        return f"Unexpected error downloading {relpath}: {e}"

def download_dataset(index_file_path, output_dir, num_workers=64):
    # Load the index file
    with open(index_file_path, 'r') as f:
        index_data = json.load(f)
    
    basedir = index_data['basedir']
    relpaths = index_data['relpaths']
    
    # Filter for mp4 files only
    mp4_files = [(relpath, urljoin(basedir, relpath), os.path.join(output_dir, relpath)) 
                 for relpath in relpaths if relpath.endswith('.mp4')]
    
    print(f"Found {len(mp4_files)} MP4 files to download")
    print(f"Using {num_workers} workers for parallel downloads")
    
    # Use multiprocessing pool to download files in parallel
    start_time = time.time()
    
    if num_workers > len(mp4_files):
        num_workers = len(mp4_files)
    
    # Create a progress bar for overall progress
    with tqdm(total=len(mp4_files), desc="Overall Download Progress", unit="files") as pbar:
        with Pool(processes=num_workers) as pool:
            # Map the download function to all files (without passing pbar)
            results = []
            for result in pool.imap_unordered(download_single_file, 
                                             [(url, output_path, relpath) for relpath, url, output_path in mp4_files]):
                results.append(result)
                pbar.update(1)
                # Print every 100th result to avoid overwhelming output
                if len(results) % 100 == 0:
                    print(f"Completed {len(results)} downloads...")
    
    # Print final results summary
    successful_downloads = sum(1 for r in results if "Downloaded" in r)
    skipped_files = sum(1 for r in results if "Skipped" in r)
    failed_downloads = len(results) - successful_downloads - skipped_files
    
    print(f"
Download Summary:")
    print(f"  Successful downloads: {successful_downloads}")
    print(f"  Skipped files: {skipped_files}")
    print(f"  Failed downloads: {failed_downloads}")
    
    end_time = time.time()
    total_time = end_time - start_time
    print(f"Download completed in {total_time:.2f} seconds")

if __name__ == "__main__":
    args = tyro.cli(DownloadVideos)

    output_dir = os.path.join(args.output_dir, os.path.basename(args.index_file_path).split('.')[0])
    os.makedirs(output_dir, exist_ok=True)

    num_workers = os.cpu_count()

    # print all args
    print(f"Index file path: {args.index_file_path}")
    print(f"Output directory: {output_dir}")
    print(f"Number of workers: {num_workers}")

    # get cpu count
    download_dataset(args.index_file_path, output_dir, num_workers)