build_and_upload.py

import os
import json
import shutil
from pathlib import Path
import subprocess
import argparse
import tempfile
from typing import Optional, Dict, List, Tuple
import sys


def read_and_validate_data(metadata_dir: str, video_dir: str) -> List[Tuple[str, dict, str]]:
    """
    Read metadata and video files, matching them by filename.
    Returns list of tuples: (filename_base, metadata_content, video_path)
    """
    paired_data = []
    missing_pairs = []
    
    # Get all json files
    json_files = list(Path(metadata_dir).glob('*.json'))
    
    print(f"Found {len(json_files)} JSON files in metadata directory")
    
    for json_path in json_files:
        base_name = json_path.stem
        video_path = Path(video_dir) / f"{base_name}.mp4"
        
        # Check if corresponding video exists
        if not video_path.exists():
            missing_pairs.append(base_name)
            continue
            
        try:
            with open(json_path, 'r') as f:
                metadata = json.load(f)
        except Exception as e:
            print(f"Error reading metadata file {json_path}: {str(e)}")
            continue
            
        paired_data.append((base_name, metadata, str(video_path)))
    
    # Report statistics
    print(f"\nDataset Statistics:")
    print(f"Total paired files: {len(paired_data)}")
    if missing_pairs:
        print(f"Missing video files for: {', '.join(missing_pairs)}")
    
    return paired_data

def create_video_dataset(
    source_video_dir: str,
    source_metadata_dir: str,
    output_base_dir: str,
    examples_per_folder: int = 9500,
    max_total_examples: Optional[int] = None,
) -> int:
    """
    Create a video dataset organized in folders with accompanying metadata.
    
    Args:
        source_video_dir: Directory containing source video files
        source_metadata_dir: Directory containing metadata JSON files
        output_base_dir: Base directory for the organized dataset
        examples_per_folder: Maximum number of examples per folder
        max_total_examples: Maximum total examples to process (None for all)
    
    Returns:
        int: Total processed examples
    """
    
    # Create base output directory
    Path(output_base_dir).mkdir(parents=True, exist_ok=True)
    
    # Read and validate data
    paired_data = read_and_validate_data(source_metadata_dir, source_video_dir)
    if not paired_data:
        raise ValueError("No valid paired files found")
    
    current_folder = 0
    processed_examples = 0
    current_folder_examples = 0

    # Create the first folder
    current_folder_path = os.path.join(output_base_dir, f"{current_folder:04d}")
    Path(current_folder_path).mkdir(exist_ok=True)
    
    # Initialize metadata file for the current folder
    metadata_file = open(os.path.join(current_folder_path, "metadata.jsonl"), "w")
    
    for base_name, metadata, video_path in paired_data:
        # Break if we've reached the maximum examples
        if max_total_examples and processed_examples >= max_total_examples:
            break
            
        # If we've reached the examples per folder limit, create a new folder
        if current_folder_examples >= examples_per_folder:
            metadata_file.close()
            current_folder += 1
            current_folder_examples = 0
            current_folder_path = os.path.join(output_base_dir, f"{current_folder:04d}")
            Path(current_folder_path).mkdir(exist_ok=True)
            metadata_file = open(os.path.join(current_folder_path, "metadata.jsonl"), "w")
        
        # Copy video to new location
        video_filename = f"{base_name}.mp4"
        destination_path = os.path.join(current_folder_path, video_filename)
        shutil.copy2(video_path, destination_path)
        
        # Create metadata entry (including original filename and all metadata)
        metadata_entry = {
            "file_name": video_filename,
            **metadata  # Include all fields from original metadata
        }
        
        # Write metadata entry
        metadata_file.write(json.dumps(metadata_entry) + "\n")
        
        current_folder_examples += 1
        processed_examples += 1
        
        if processed_examples % 100 == 0:  # Progress update every 100 examples
            print(f"Processed {processed_examples} examples")
    
    # Close the last metadata file
    metadata_file.close()
    
    print(f"Dataset creation complete. Processed {processed_examples} examples across {current_folder + 1} folders")
    return processed_examples

def upload_to_huggingface(dataset_path: str, hf_dataset_name: str) -> None:
    """Upload the dataset to Hugging Face"""
    try:
        cmd = [
            "huggingface-cli", 
            "upload-large-folder", 
            hf_dataset_name, 
            dataset_path,
            "--repo-type=dataset"
        ]
        subprocess.run(cmd, check=True)
        print(f"Successfully uploaded dataset to {hf_dataset_name}")
    except subprocess.CalledProcessError as e:
        print(f"Error uploading to Hugging Face: {str(e)}")
        sys.exit(1)

def main():
    parser = argparse.ArgumentParser(description="Create and upload a video dataset to Hugging Face")
    parser.add_argument("--video-dir", required=True, help="Directory containing source video files")
    parser.add_argument("--metadata-dir", required=True, help="Directory containing metadata JSON files")
    parser.add_argument("--hf-dataset-name", required=True, help="Hugging Face dataset name (e.g., 'username/dataset-name')")
    parser.add_argument("--examples-per-folder", type=int, default=9500, 
                       help="Maximum examples per folder (max 10000)")
    parser.add_argument("--max-examples", type=int, help="Maximum total examples to process")
    parser.add_argument("--temp-dir", help="Temporary directory for dataset creation (default: system temp directory)")
    
    args = parser.parse_args()
    
    # Validate examples_per_folder
    if args.examples_per_folder > 10000:
        print("Error: examples-per-folder cannot exceed 10000")
        sys.exit(1)
    elif args.examples_per_folder <= 0:
        print("Error: examples-per-folder must be greater than 0")
        sys.exit(1)
    
    # Use provided temp directory or create one
    temp_base_dir = args.temp_dir or tempfile.mkdtemp()
    print(f"Using temporary directory: {temp_base_dir}")
    
    try:
        # Create dataset
        processed_examples = create_video_dataset(
            source_video_dir=args.video_dir,
            source_metadata_dir=args.metadata_dir,
            output_base_dir=temp_base_dir,
            examples_per_folder=args.examples_per_folder,
            max_total_examples=args.max_examples
        )
        
        if processed_examples > 0:
            # Upload to Hugging Face
            upload_to_huggingface(temp_base_dir, args.hf_dataset_name)
        else:
            print("No examples were processed. Aborting upload.")
            sys.exit(1)
            
    finally:
        if not args.temp_dir:  # Only remove if we created the temp directory
            shutil.rmtree(temp_base_dir, ignore_errors=True)

if __name__ == "__main__":
    main()