basketball_generation

add_basketball_court.py

from spatialstudio import splv
from utils import vv_to_splv
from time import perf_counter
import numpy as np
import os
import math
from tqdm import tqdm
import multiprocessing
import tempfile
import shutil
import glob
import librosa

class AudioProcessor:
    def __init__(self, path, sampleRate=44100):
        self.sampleRate = sampleRate
        self.rawAudio, _ = librosa.load(path, sr=sampleRate)
        self.numFrames = len(self.rawAudio)
        
    def get_audio_frame(self, startSample, samplesPerFrame):
        endSample = min(startSample + samplesPerFrame, len(self.rawAudio))
        frameData = np.clip(self.rawAudio[startSample:endSample], -1.0, 1.0)
        
        if len(frameData) < samplesPerFrame:
            frameData = np.pad(frameData, (0, samplesPerFrame - len(frameData)), mode='constant')
        
        audioInt16 = (frameData * 32767).astype(np.int16)
        return audioInt16.tobytes()

def create_rectangular_hardwood_floor(resolution):
    """Create a rectangular hardwood floor extending 1.5x back in the direction the camera is facing"""
    floor_voxels = {}
    
    # Uniform hardwood color - consistent basketball court color
    base_wood_color = (158, 117, 76)  # Saddle brown
    
    # Court dimensions
    min_x = 0
    max_x = 512  # Original width
    min_z = 0
    max_z = 768  # 1.5x depth
    
    # Corner radius for all corners
    corner_radius = max(8, int(512 * 0.08))  # 8% of original resolution
    
    print(f"Creating rectangular hardwood floor with 1.5x back extension...")
    print(f"  Court dimensions: {max_x}x{max_z}")
    print(f"  X bounds: {min_x} to {max_x}")
    print(f"  Z bounds: {min_z} to {max_z}")
    print(f"  Corner radius: {corner_radius}")
    
    # Create the floor
    for x in range(min_x, max_x - 1):
        for z in range(min_z, max_z - 1):
            # Determine if this point should be included
            is_inside = True
            
            # Check if we're near any edge
            near_x_edge = x < corner_radius or x > (max_x - 1 - corner_radius)
            near_z_edge = z < corner_radius or z > (max_z - 1 - corner_radius)
            
            # If near both edges, we're in a corner - check the radius
            if near_x_edge and near_z_edge:
                # Calculate distance from nearest corner
                corner_x = corner_radius if x < corner_radius else (max_x - 1 - corner_radius)
                corner_z = corner_radius if z < corner_radius else (max_z - 1 - corner_radius)
                
                dx = x - corner_x
                dz = z - corner_z
                corner_distance = math.sqrt(dx * dx + dz * dz)
                is_inside = corner_distance <= corner_radius
            
            if is_inside:
                # Add subtle wood grain variation based on position
                position_hash = (x * 73 + z * 131) % 1000
                variation = (position_hash % 17) - 8  # -8 to +8 variation
                
                # Apply variation to base color
                r = max(0, min(255, base_wood_color[0] + variation))
                g = max(0, min(255, base_wood_color[1] + variation))
                b = max(0, min(255, base_wood_color[2] + variation))
                
                floor_voxels[(x, 0, z)] = (r, g, b)
    
    print(f"Created rectangular floor: {len(floor_voxels)} voxels")
    return floor_voxels

def apply_shadow_blur(floor_voxels, shadow_positions, resolution):
    """Apply blur kernel to shadow areas for realistic soft shadows"""
    if not shadow_positions:
        return floor_voxels
    
    # Define blur kernel - 3x3 Gaussian-like kernel
    blur_kernel = [
        [0.0625, 0.125, 0.0625],  # 1/16, 1/8, 1/16
        [0.125,  0.25,  0.125],   # 1/8, 1/4, 1/8  
        [0.0625, 0.125, 0.0625]   # 1/16, 1/8, 1/16
    ]
    
    # Create a copy of floor_voxels to work with
    blurred_floor = dict(floor_voxels)
    
    # Apply blur to each shadow position
    for shadow_pos in shadow_positions:
        sx, sy, sz = shadow_pos
        
        # Only blur shadows on the floor surface (y=0)
        if sy == 0:
            # Apply blur kernel in X-Z plane (floor plane)
            total_weight = 0.0
            blurred_r = 0.0
            blurred_g = 0.0
            blurred_b = 0.0
            
            for kernel_x in range(3):
                for kernel_z in range(3):
                    # Calculate neighbor position
                    neighbor_x = sx + (kernel_x - 1)  # -1, 0, 1
                    neighbor_z = sz + (kernel_z - 1)  # -1, 0, 1
                    
                    # Check bounds
                    if (0 <= neighbor_x < resolution and 
                        0 <= neighbor_z < resolution and 
                        (neighbor_x, 0, neighbor_z) in floor_voxels):
                        
                        weight = blur_kernel[kernel_x][kernel_z]
                        neighbor_color = floor_voxels[(neighbor_x, 0, neighbor_z)]
                        
                        blurred_r += neighbor_color[0] * weight
                        blurred_g += neighbor_color[1] * weight
                        blurred_b += neighbor_color[2] * weight
                        total_weight += weight
            
            # Apply blurred color if we have valid weights
            if total_weight > 0:
                final_r = int(blurred_r / total_weight)
                final_g = int(blurred_g / total_weight)
                final_b = int(blurred_b / total_weight)
                
                blurred_floor[(sx, 0, sz)] = (final_r, final_g, final_b)
    
    return blurred_floor

def apply_soft_edge_blur(floor_voxels, shadow_positions, resolution):
    """Apply soft edge blur around shadow perimeter for gradual falloff"""
    if not shadow_positions:
        return floor_voxels
    
    # Create expanded shadow area with gradient falloff
    edge_positions = set()
    
    # Find edge positions around shadows
    for shadow_pos in shadow_positions:
        sx, sy, sz = shadow_pos
        
        # Check 5x5 area around each shadow voxel (only on floor y=0)
        if sy == 0:
            for dx in range(-2, 3):
                for dz in range(-2, 3):
                    edge_x = sx + dx
                    edge_z = sz + dz
                    
                    if (0 <= edge_x < resolution and 
                        0 <= edge_z < resolution and 
                        (edge_x, 0, edge_z) in floor_voxels and
                        (edge_x, 0, edge_z) not in shadow_positions):
                        
                        # Calculate distance from shadow center
                        distance = max(abs(dx), abs(dz))  # Manhattan distance
                        
                        if distance <= 2:  # Within 2 voxels of shadow
                            edge_positions.add((edge_x, 0, edge_z, distance))
    
    # Apply gradient falloff to edge positions
    for edge_pos in edge_positions:
        ex, ey, ez, distance = edge_pos
        
        # Calculate falloff strength based on distance
        if distance == 1:
            falloff_strength = 0.9  # Slight darkening
        elif distance == 2:
            falloff_strength = 0.95  # Very slight darkening
        else:
            continue
        
        # Apply falloff
        if (ex, ey, ez) in floor_voxels:
            existing_color = floor_voxels[(ex, ey, ez)]
            
            new_r = int(existing_color[0] * falloff_strength)
            new_g = int(existing_color[1] * falloff_strength)
            new_b = int(existing_color[2] * falloff_strength)
            
            floor_voxels[(ex, ey, ez)] = (new_r, new_g, new_b)
    
    return floor_voxels

def cast_shadows_from_human(frame, floor_voxels, resolution):
    """Cast shadows on the floor based on human body position (lighting from top front)"""
    # Find all human body voxels - only iterate through the player frame's actual bounds
    # The player frame is 512x512x512, so we need to limit our search to that space
    player_frame_size = 512  # Original player frame size
    
    # human_voxels = []
    # for x in range(player_frame_size):
    #     for y in range(player_frame_size):
    #         for z in range(player_frame_size):
    #             voxel = frame.get_voxel(x, y, z)
    #             if voxel is not None:
    #                 human_voxels.append((x, y, z))

    
    # Track which floor positions already have shadows to prevent darkening
    shadow_positions = set()
    
    # Cast shadows - light comes from top front (high Y, low Z)
    # Shadow direction: toward back (negative Z for behind player), stretched out
    
    for (hx, hy, hz), color in frame:
        # Cast shadow onto floor - more stretched out and behind character
        shadow_stretch = 1.0  # How much to stretch the shadow
        shadow_offset_z = -0.8  # Shadow cast toward back (NEGATIVE for behind player)
        
        # Calculate shadow position - stretch based on height
        shadow_x = hx
        shadow_z = int(hz + (hy * shadow_offset_z * shadow_stretch + 256))
        
        # Cast shadow on floor (y=0 only) - single layer
        shadow_y = 0
        
        # Check if shadow position is within the floor bounds (can be larger than 512)
        if (shadow_x, shadow_y, shadow_z) in floor_voxels:
            # Check if we've already cast a shadow here
            shadow_key = (shadow_x, shadow_y, shadow_z)
            if shadow_key not in shadow_positions:
                # First shadow at this position - darken it
                existing_color = floor_voxels[(shadow_x, shadow_y, shadow_z)]
                
                # Darker shadow for visibility
                shadow_strength = 0.75  # More visible shadow
                new_r = int(existing_color[0] * shadow_strength)
                new_g = int(existing_color[1] * shadow_strength)
                new_b = int(existing_color[2] * shadow_strength)
                
                floor_voxels[(shadow_x, shadow_y, shadow_z)] = (new_r, new_g, new_b)
                shadow_positions.add(shadow_key)
            # If shadow already exists at this position, don't darken it further
    
    # Apply blur to shadow areas for realism
    if shadow_positions:
        print(f"Applying blur to {len(shadow_positions)} shadow voxels...")
        floor_voxels = apply_shadow_blur(floor_voxels, shadow_positions, resolution)
        
        # Apply soft edge blur for gradual falloff
        print(f"Applying soft edge blur for gradual falloff...")
        floor_voxels = apply_soft_edge_blur(floor_voxels, shadow_positions, resolution)
    
    return floor_voxels

def add_floor_to_frame_with_shadows(frame, resolution):
    """Add rectangular floor with shadows cast by the human body"""
    # Create rectangular floor
    floor_voxels = create_rectangular_hardwood_floor(resolution)
    
    # Cast shadows based on human body position
    floor_voxels = cast_shadows_from_human(frame, floor_voxels, resolution)
    
    # Add floor voxels to frame
    for (x, y, z), color in floor_voxels.items():
        if 0 <= x < resolution and 0 <= y < resolution and 0 <= z < resolution:
            frame.set_voxel(x, y, z, color)
    
    return frame

def process_single_frame_with_court(args):
    """Process a single frame by creating a new 768x768x768 frame with court and adding the basketball player"""
    frame_index, tmp_frame_path, resolution, temp_dir = args
    
    try:
        # Load the basketball player frame from tmp file
        player_frame = splv.Frame.load(tmp_frame_path)
        
        # Create a new 768x768x768 frame with the rectangular court
        court_frame = splv.Frame(resolution, resolution, resolution)
        
        # Fill the entire floor level (y=0) with basketball court color
        # Using frame.fill(xMin, yMin, zMin, xMax, yMax, zMax, (r, g, b))
        basketball_color = (158, 117, 76)  # Saddle brown basketball court color
        # court_frame.fill(0, 0, 0, 512, 0, resolution-1, basketball_color)
        
        # Temporarily comment out the complex hardwood floor creation
        floor_voxels = create_rectangular_hardwood_floor(resolution)
        floor_voxels = cast_shadows_from_human(player_frame, floor_voxels, resolution)
        for (x, y, z), color in floor_voxels.items():
            if 0 <= x < resolution and 0 <= y < resolution and 0 <= z < resolution:
                court_frame.set_voxel(x, y, z, color)
        
        # Add the basketball player frame to the court frame at position (0, 0, 0)
        court_frame.add(player_frame, 0, 0, 256)
        
        # Save the combined frame to temporary file
        temp_file = os.path.join(temp_dir, f"frame_{frame_index:06d}.splv")
        court_frame.save(temp_file)
        
        # Clean up memory
        del court_frame
        del player_frame
        
        # Return metadata instead of the frame object
        return {
            'frame_index': frame_index,
            'temp_file': temp_file,
            'status': 'success'
        }
    
    except Exception as e:
        print(f"❌ Error processing frame {frame_index}: {e}")
        # Save original frame if processing fails
        try:
            player_frame = splv.Frame.load(tmp_frame_path)
            temp_file = os.path.join(temp_dir, f"frame_{frame_index:06d}.splv")
            player_frame.save(temp_file)
            del player_frame
            
            return {
                'frame_index': frame_index,
                'temp_file': temp_file,
                'status': 'error'
            }
        except:
            return {
                'frame_index': frame_index,
                'temp_file': None,
                'status': 'failed'
            }

def create_basketball_court_video(tmp_frames_dir=None, audio_path=None):
    """Add basketball court floor to existing video using tmp frames with optional audio"""
    if tmp_frames_dir is None:
        tmp_frames_dir = "tmp/owlii_basketball_frames"
    
    # Output path depends on whether audio is included
    if audio_path:
        output_path = "outputs/owlii_basketball_with_court_512_with_audio.splv"
    else:
        output_path = "outputs/owlii_basketball_with_court_512.splv"
    
    resolution = 768
    max_frames = 600  # Process more frames since it's faster now
    
    print(f"🏀 Adding rectangular basketball court floor to video (PARALLEL)...")
    print(f"📁 Input frames: {tmp_frames_dir}")
    print(f"📁 Output: {output_path}")
    print(f"📊 Resolution: {resolution}x{resolution}x{resolution}")
    print(f"🎬 Processing frames with rectangular court (1.5x back extension)")
    if audio_path:
        print(f"🎵 Audio: {audio_path}")
    
    # Check if tmp frames directory exists
    if not os.path.exists(tmp_frames_dir):
        print(f"❌ Tmp frames directory not found: {tmp_frames_dir}")
        print("Please run the basketball video generation script first to create tmp frames.")
        return
    
    # Find all tmp frame files
    tmp_frame_files = sorted(glob.glob(os.path.join(tmp_frames_dir, "basketball_frame_*.splv")))
    
    if not tmp_frame_files:
        print(f"❌ No tmp frame files found in: {tmp_frames_dir}")
        print("Please run the basketball video generation script first.")
        return
    
    # Limit to max_frames
    tmp_frame_files = tmp_frame_files[:max_frames]
    total_frames = len(tmp_frame_files)
    
    print(f"📁 Found {total_frames} tmp frame files")
    
    start_time = perf_counter()
    
    # Configure multiprocessing
    # Use fewer processes for high resolution to avoid memory issues
    if resolution >= 512:
        num_processes = min(4, multiprocessing.cpu_count())
    else:
        num_processes = min(8, multiprocessing.cpu_count())
    
    print(f"🔧 Using {num_processes} processes for parallel processing")
    
    # Create temporary directory for enhanced frame files
    temp_dir = tempfile.mkdtemp(prefix="basketball_court_")
    print(f"📁 Using temporary directory: {temp_dir}")
    
    # Prepare arguments for multiprocessing
    frame_args = []
    for i, tmp_frame_path in enumerate(tmp_frame_files):
        frame_args.append((i, tmp_frame_path, resolution, temp_dir))
    
    # Process frames in parallel
    with multiprocessing.Pool(processes=num_processes) as pool:
        results = []
        
        print("🎨 Processing frames with rectangular basketball court floor...")
        with tqdm(desc="Adding court", total=len(frame_args)) as pbar:
            try:
                for result in pool.imap(process_single_frame_with_court, frame_args):
                    results.append(result)
                    pbar.update(1)
                    
                    # Show progress every 50 frames
                    if len(results) % 50 == 0:
                        pbar.set_postfix({'completed': len(results)})
            except Exception as e:
                print(f"❌ Error processing frames: {e}")
                return
    
    # Sort results by frame index to maintain order
    results.sort(key=lambda x: x['frame_index'])
    
    # Load enhanced frames from temp files
    enhanced_frames = []
    failed_frames = 0
    
    print("📂 Loading enhanced frames from temporary files...")
    for result in tqdm(results, desc="Loading frames"):
        if result['status'] == 'success' and result['temp_file'] is not None:
            try:
                # Load frame from temporary file
                frame = splv.Frame.load(result['temp_file'])
                enhanced_frames.append(frame)
            except Exception as e:
                print(f"❌ Error loading frame {result['frame_index']}: {e}")
                failed_frames += 1
                enhanced_frames.append(None)
        else:
            failed_frames += 1
            enhanced_frames.append(None)
    
    # Clean up temporary files
    try:
        shutil.rmtree(temp_dir)
        print(f"🗑️  Cleaned up temporary directory: {temp_dir}")
    except:
        print(f"⚠️  Could not clean up temporary directory: {temp_dir}")
    
    # Filter out None frames
    enhanced_frames = [frame for frame in enhanced_frames if frame is not None]
    
    if failed_frames > 0:
        print(f"⚠️  {failed_frames} frames failed to process")
    
    print(f"✅ Successfully processed {len(enhanced_frames)} frames")
    
    process_time = perf_counter()
    print(f"⚡ Frame processing completed in {process_time - start_time:.2f} seconds")
    
    if len(enhanced_frames) > 0:
        print(f"📊 Average time per frame: {(process_time - start_time) / len(enhanced_frames):.3f} seconds")
    else:
        print("❌ No frames were successfully processed")
        return
    
    # Setup audio processor if audio path is provided
    audio_processor = None
    if audio_path:
        print(f"🎵 Setting up audio...")
        if not os.path.exists(audio_path):
            print(f"❌ Audio file not found: {audio_path}")
            print("Proceeding without audio...")
            audio_path = None
        else:
            audio_processor = AudioProcessor(audio_path)
            audio_duration = len(audio_processor.rawAudio) / audio_processor.sampleRate
            print(f"🎵 Audio loaded: {audio_duration:.1f} seconds")
    
    # Create SPLV video (with or without audio)
    framerate = 30.0
    
    if audio_processor:
        print("🎬 Creating SPLV video with basketball court and audio...")
        
        # Calculate audio timing
        samples_per_frame = int(audio_processor.sampleRate / framerate)
        
        # Create encoder with audio support
        encoder = splv.Encoder(
            resolution, resolution, resolution,
            framerate,
            audioParams=(1, audio_processor.sampleRate, 2),  # Mono, sample rate, 16-bit
            outputPath=output_path
        )
        
        # Encode frames with audio
        print("🎞️ Encoding frames with audio...")
        for i, frame in enumerate(tqdm(enhanced_frames, desc="Encoding with audio")):
            try:
                # Get audio frame
                audio_sample_start = int(i * samples_per_frame)
                audio_frame_data = audio_processor.get_audio_frame(audio_sample_start, samples_per_frame)
                audio_bytes = np.frombuffer(audio_frame_data, dtype=np.uint8)
                
                # Encode frame and audio
                encoder.encode(frame)
                encoder.encode_audio(audio_bytes)
                
            except Exception as e:
                print(f"❌ Error encoding frame {i}: {e}")
                break
        
        # Finalize encoding
        encoder.finish()
        
    else:
        print("🎬 Creating SPLV video with basketball court...")
        vv_to_splv(enhanced_frames, output_path, resolution=resolution, framerate=framerate)
    
    end_time = perf_counter()
    print(f"🎉 Rectangular basketball court video created successfully!")
    print(f"📁 Output: {output_path}")
    print(f"⏱️ Total time: {end_time - start_time:.2f} seconds")
    print(f"📊 Processed {len(enhanced_frames)} frames")
    print(f"🚀 Performance: {len(enhanced_frames)/(end_time - start_time):.2f} frames/second")
    print(f"⚡ Speedup: ~{num_processes}x faster with parallel processing!")
    print(f"🏀 Court shape: Rectangular (1.5x back extension)")
    if audio_processor:
        print(f"🎵 Audio: {audio_path}")
        print(f"🎬 Ready to play basketball with sound!")

if __name__ == "__main__":
    multiprocessing.freeze_support()
    
    # Default paths
    tmp_frames_dir = "tmp/owlii_basketball_frames"
    audio_path = "audio/basketball2.mp3"  # Set to audio file path for audio support
    
    # Check if tmp frames exist
    if os.path.exists(tmp_frames_dir):
        tmp_files = glob.glob(os.path.join(tmp_frames_dir, "basketball_frame_*.splv"))
        if tmp_files:
            print(f"✅ Found {len(tmp_files)} tmp frames, adding rectangular court...")
            
            # Uncomment the line below to add audio support
            # audio_path = "audio/basketball2.mp3"
            
            create_basketball_court_video(tmp_frames_dir, audio_path)
        else:
            print(f"❌ No tmp frames found in: {tmp_frames_dir}")
            print("Please run the basketball video generation script first.")
    else:
        print(f"❌ Tmp frames directory not found: {tmp_frames_dir}")
        print("Please run the basketball video generation script first to create tmp frames.") 

owlii_basketball.py

import spatialstudio as splv
from utils import vv_to_splv
from time import perf_counter
import numpy as np
import os
import glob
from tqdm import tqdm
import multiprocessing
import tempfile
import shutil

def load_ply_file(file_path):
    """Load a PLY file and return vertices with colors (handles both ASCII and binary formats)"""
    import struct
    
    vertices = []
    colors = []
    
    with open(file_path, 'rb') as f:
        # Read header
        header_lines = []
        while True:
            line = f.readline().decode('ascii').strip()
            header_lines.append(line)
            if line == 'end_header':
                break
        
        # Parse header to get vertex count and format
        vertex_count = 0
        is_binary = False
        
        for line in header_lines:
            if line.startswith('format'):
                if 'binary' in line:
                    is_binary = True
            elif line.startswith('element vertex'):
                vertex_count = int(line.split()[2])
        
        # Read vertex data
        if is_binary:
            # Binary format - float32 for xyz and normals, uint8 for rgba
            for _ in range(vertex_count):
                # Read 24 bytes for x,y,z,nx,ny,nz (6 * 4 bytes each for float32)
                x, y, z, nx, ny, nz = struct.unpack('<ffffff', f.read(24))
                # Read 4 bytes for r,g,b,a (4 * 1 byte each for uint8)
                r, g, b, a = struct.unpack('<BBBB', f.read(4))
                vertices.append((x, y, z))
                colors.append((r, g, b))  # Ignore alpha and normals for now
        else:
            # ASCII format (fallback)
            remaining_data = f.read().decode('ascii')
            lines = remaining_data.strip().split('\n')
            for line in lines:
                parts = line.strip().split()
                if len(parts) >= 6:
                    x, y, z = float(parts[0]), float(parts[1]), float(parts[2])
                    r, g, b = int(parts[3]), int(parts[4]), int(parts[5])
                    vertices.append((x, y, z))
                    colors.append((r, g, b))
    
    return vertices, colors

def process_point_cloud(vertices, colors, resolution, scale_factor, global_bounds=None):
    """Process point cloud data with average pooling using global normalization"""
    
    if not vertices:
        return {}
    
    # Convert to numpy arrays for easier manipulation
    vertices = np.array(vertices)
    colors = np.array(colors)
    
    # Fix orientation for Owlii basketball player
    # Try different transformations to get player standing upright and facing forward
    vertices = vertices[:, [0, 1, 2]]  # Keep original x, y, z for now
    # vertices[:, 1] = -vertices[:, 1]  # Don't flip Y axis - keep right side up
    vertices[:, 2] = -vertices[:, 2]  # Flip Z axis to face forward
    
    voxel_colors = {}
    
    if global_bounds is not None:
        # Use global bounds for consistent normalization across all frames
        global_min_coords, global_max_coords = global_bounds
        range_coords = global_max_coords - global_min_coords
        max_range = np.max(range_coords)
        
        if max_range > 0:
            # Scale to fit in the voxel grid with some padding, then apply downscaling
            scale = (resolution - 40) / max_range * scale_factor if scale_factor > 0 else (resolution - 40) / max_range
            
            # Center the point cloud using global bounds
            center_offset = (resolution - 1) / 2
            global_center = (global_min_coords + global_max_coords) / 2
            scaled_vertices = (vertices - global_center) * scale + center_offset
            
            # Convert to voxel coordinates
            voxel_coords = np.round(scaled_vertices).astype(int)
            
            # Anchor to floor - shift all voxels down so the lowest one is at Y=0
            # Use global minimum Y for consistent floor level
            global_min_y_scaled = int(((global_min_coords[1] - global_center[1]) * scale + center_offset))
            voxel_coords[:, 1] -= global_min_y_scaled
            
            # Accumulate colors for overlapping voxels (average pooling)
            for i, (x, y, z) in enumerate(voxel_coords):
                if 0 <= x < resolution and 0 <= y < resolution and 0 <= z < resolution:
                    voxel_pos = (x, y, z)
                    if voxel_pos not in voxel_colors:
                        voxel_colors[voxel_pos] = []
                    voxel_colors[voxel_pos].append(colors[i])
            
            # Average the colors for each voxel (average pooling)
            for voxel_pos, color_list in voxel_colors.items():
                avg_color = np.mean(color_list, axis=0).astype(int)
                voxel_colors[voxel_pos] = tuple(avg_color)
    
    return voxel_colors

def calculate_frame_bounds(args):
    """Calculate bounding box for a single frame (first pass for global normalization)"""
    ply_file, frame_index = args
    
    try:
        # Load PLY file
        vertices, colors = load_ply_file(ply_file)
        
        if not vertices:
            return None
        
        # Apply same transformations as in process_point_cloud
        vertices = np.array(vertices)
        vertices = vertices[:, [0, 1, 2]]  # Keep original x, y, z
        vertices[:, 2] = -vertices[:, 2]  # Flip Z axis to face forward
        
        # Calculate bounds
        min_coords = np.min(vertices, axis=0)
        max_coords = np.max(vertices, axis=0)
        
        return {
            'frame_index': frame_index,
            'min_coords': min_coords,
            'max_coords': max_coords,
            'vertex_count': len(vertices),
            'filename': os.path.basename(ply_file)
        }
    
    except Exception as e:
        print(f"❌ Error calculating bounds for frame {frame_index}: {e}")
        return None

def process_single_frame(args):
    """Process a single frame and save to temporary file"""
    ply_file, resolution, scale_factor, frame_index, temp_dir, global_bounds = args
    
    # Load PLY file
    vertices, colors = load_ply_file(ply_file)
    
    # Process point cloud
    voxel_data = process_point_cloud(vertices, colors, resolution, scale_factor, global_bounds)
    
    # Create frame and save to temporary file
    frame = splv.SPLVframe(resolution, resolution, resolution)
    for (x, y, z), color in voxel_data.items():
        frame.set_voxel(x, y, z, color)
    
    # Save frame to permanent tmp file
    temp_file = os.path.join(temp_dir, f"basketball_frame_{frame_index:06d}.splv")
    frame.save(temp_file)
    
    # Store counts before cleanup
    vertex_count = len(vertices)
    voxel_count = len(voxel_data)
    
    # Clean up memory
    del frame
    del voxel_data
    del vertices
    del colors
    
    return {
        'frame_index': frame_index,
        'temp_file': temp_file,
        'vertex_count': vertex_count,
        'voxel_count': voxel_count,
        'filename': os.path.basename(ply_file)
    }

def create_video_from_temp_files(temp_dir, resolution=512):
    """Create SPLV video from existing temporary files"""
    output_path = f"outputs/owlii_basketball_video_{resolution}.splv"
    
    print(f"🔄 Loading existing temp files from: {temp_dir}")
    
    # Find all tmp files
    temp_files = sorted(glob.glob(os.path.join(temp_dir, "basketball_frame_*.splv")))
    
    if not temp_files:
        print(f"❌ No temp files found in {temp_dir}")
        return
    
    print(f"📁 Found {len(temp_files)} temp files")
    
    # Load frames from tmp files
    print("Loading frames from permanent tmp files...")
    frames = []
    for temp_file in tqdm(temp_files, desc="Loading frames"):
        try:
            frame = splv.frame_load(temp_file)
            frames.append(frame)
        except Exception as e:
            print(f"❌ Error loading {temp_file}: {e}")
            continue
    
    print(f"✅ Successfully loaded {len(frames)} frames")
    
    # Create SPLV video
    print("🎬 Creating SPLV video...")
    start_time = perf_counter()
    vv_to_splv(frames, output_path, resolution=resolution, framerate=30.0)
    end_time = perf_counter()
    
    print(f"🎉 Video created successfully: {output_path}")
    print(f"📊 {len(frames)} frames with resolution {resolution}x{resolution}x{resolution}")
    print(f"⏱️  Video creation time: {end_time - start_time:.2f} seconds")
    
    # Keep tmp files permanently
    print(f"📁 Tmp files preserved permanently at: {temp_dir}")
    print("💾 These are kept for manual cleanup only")
    
    return output_path

def create_owlii_basketball_video(missing_frames_only=None):
    """Create SPLV video from Owlii basketball player PLY files"""
    
    # Get all PLY files in order
    ply_dir = "data/Owlii/basketball_player_vox11"
    ply_files = sorted(glob.glob(os.path.join(ply_dir, "basketball_player_vox11_*.ply")))
    
    print(f"Found {len(ply_files)} PLY files")
    
    # Limit to first 60 frames for testing
    max_frames = 600
    ply_files = ply_files[:max_frames]
    
    # If only generating missing frames, filter the PLY files
    if missing_frames_only is not None:
        print(f"🎯 Processing only {len(missing_frames_only)} missing frames")
        filtered_ply_files = []
        filtered_indices = []
        for frame_idx in missing_frames_only:
            if frame_idx < len(ply_files):
                filtered_ply_files.append(ply_files[frame_idx])
                filtered_indices.append(frame_idx)
        ply_files = filtered_ply_files
        frame_indices = filtered_indices
    else:
        frame_indices = list(range(len(ply_files)))
    
    resolution = 512
    scale_factor = 1.0  # Default scaling (no additional scaling)
    output_path = f"outputs/owlii_basketball_video_{resolution}.splv"
    
    start_time = perf_counter()
    
    # Create permanent tmp directory for frame files
    temp_dir = "tmp/owlii_basketball_frames"
    os.makedirs(temp_dir, exist_ok=True)
    print(f"Using permanent tmp directory: {temp_dir}")
    
    try:
        # Configure multiprocessing
        # Reduce processes for high resolution to avoid memory issues
        if resolution >= 1024:
            num_processes = 4  # Very conservative for 1024+
        elif resolution >= 512:
            num_processes = min(10, multiprocessing.cpu_count())
        else:
            num_processes = multiprocessing.cpu_count()
        
        print(f"🎯 Using two-pass global normalization approach:")
        print(f"   1️⃣ First pass: Calculate global bounding box across all frames")
        print(f"   2️⃣ Second pass: Process frames with consistent global normalization")
        print(f"Processing frames in parallel using {num_processes} processes (resolution: {resolution})...")
        print(f"Estimated memory per process: ~{(resolution**3 * 4) // (1024**3):.1f}GB for SPLVframe")
        
        # FIRST PASS: Calculate global bounding box
        print(f"\n1️⃣ FIRST PASS: Calculating global bounding box from {len(ply_files)} frames...")
        
        bounds_args = []
        for i, ply_file in enumerate(ply_files):
            frame_index = frame_indices[i]  # Use original frame index
            bounds_args.append((ply_file, frame_index))
        
        bounds_results = []
        with multiprocessing.Pool(processes=num_processes) as pool:
            with tqdm(desc="Calculating bounds", total=len(bounds_args)) as pbar:
                for result in pool.imap(calculate_frame_bounds, bounds_args):
                    if result is not None:
                        bounds_results.append(result)
                    pbar.update(1)
        
        if not bounds_results:
            print("❌ No valid bounds calculated!")
            return
        
        # Calculate global bounds
        print(f"\n📊 Calculating global bounds from {len(bounds_results)} valid frames...")
        all_min_coords = np.array([r['min_coords'] for r in bounds_results])
        all_max_coords = np.array([r['max_coords'] for r in bounds_results])
        
        global_min_coords = np.min(all_min_coords, axis=0)
        global_max_coords = np.max(all_max_coords, axis=0)
        global_bounds = (global_min_coords, global_max_coords)
        
        # Report global bounds
        print(f"🌍 Global bounds calculated:")
        print(f"   X: {global_min_coords[0]:.3f} to {global_max_coords[0]:.3f} (range: {global_max_coords[0] - global_min_coords[0]:.3f})")
        print(f"   Y: {global_min_coords[1]:.3f} to {global_max_coords[1]:.3f} (range: {global_max_coords[1] - global_min_coords[1]:.3f})")
        print(f"   Z: {global_min_coords[2]:.3f} to {global_max_coords[2]:.3f} (range: {global_max_coords[2] - global_min_coords[2]:.3f})")
        
        # SECOND PASS: Process frames with global normalization
        print(f"\n2️⃣ SECOND PASS: Processing frames with global normalization...")
        
        # Prepare arguments for each frame
        frame_args = []
        for i, ply_file in enumerate(ply_files):
            frame_index = frame_indices[i]  # Use original frame index for tmp file naming
            frame_args.append((ply_file, resolution, scale_factor, frame_index, temp_dir, global_bounds))
        
        # Process frames in parallel with progress reporting
        with multiprocessing.Pool(processes=num_processes) as pool:
            # Use imap to get results as they complete
            results = []
            total_vertices = 0
            total_voxels = 0
            
            with tqdm(desc="Processing frames (global normalization)", total=len(frame_args)) as pbar:
                for result in pool.imap(process_single_frame, frame_args):
                    results.append(result)
                    total_vertices += result['vertex_count']
                    total_voxels += result['voxel_count']
                    
                    # Update progress bar with current frame info
                    pbar.set_postfix({
                        'frame': result['frame_index'] + 1,
                        'vertices': f"{result['vertex_count']:,}",
                        'voxels': f"{result['voxel_count']:,}"
                    })
                    pbar.update(1)
                    
                    # Show detailed info for first few frames
                    if result['frame_index'] < 3:
                        print(f"\nFrame {result['frame_index']+1} ({result['filename']}): {result['vertex_count']:,} vertices -> {result['voxel_count']:,} voxels")
        
        # Sort results by frame index
        results.sort(key=lambda x: x['frame_index'])
        
        print(f"\nLoading frames from temporary files...")
        
        # If we only processed missing frames, we need to load ALL frames from tmp directory
        if missing_frames_only is not None:
            print(f"🔄 Loading complete frame set from tmp directory (including existing frames)...")
            all_tmp_files = sorted(glob.glob(os.path.join(temp_dir, "basketball_frame_*.splv")))
            
            # Create a mapping of frame indices to tmp files
            frame_file_map = {}
            for tmp_file in all_tmp_files:
                filename = os.path.basename(tmp_file)
                try:
                    frame_num = int(filename.split('_')[2].split('.')[0])
                    if 0 <= frame_num < max_frames:
                        frame_file_map[frame_num] = tmp_file
                except (IndexError, ValueError):
                    continue
            
            # Load frames in order
            frames = []
            for i in range(max_frames):
                if i in frame_file_map:
                    try:
                        frame = splv.frame_load(frame_file_map[i])
                        frames.append(frame)
                    except Exception as e:
                        print(f"❌ Error loading frame {i}: {e}")
                        break
                else:
                    print(f"❌ Missing frame {i} after processing!")
                    break
            
            print(f"⚠️  Note: Frames processed with global normalization may require regeneration")
            print(f"    for optimal consistency. Consider regenerating all frames if needed.")
        else:
            # We processed all frames sequentially
            frames = []
            for result in tqdm(results, desc="Loading frames"):
                # Load frame from temporary file using correct API
                frame = splv.frame_load(result['temp_file'])
                frames.append(frame)
        
        print(f"Total vertices processed: {total_vertices:,}")
        print(f"Total voxels created: {total_voxels:,}")
        print(f"Average voxels per frame: {total_voxels / len(frames):,.0f}")
        
        process_time = perf_counter()
        print(f"Processing time: {process_time - start_time:.2f} seconds")
        
        # Keep tmp files permanently
        print(f"📁 Tmp files preserved permanently at: {temp_dir}")
        print("💾 These are kept for manual cleanup only")
        
    except Exception as e:
        print(f"❌ Error occurred: {e}")
        print(f"🗂️  Tmp files preserved at: {temp_dir}")
        print("You can inspect them manually for debugging.")
        raise
    
    # Create SPLV video
    print("Creating SPLV video...")
    vv_to_splv(frames, output_path, resolution=resolution, framerate=30.0)
    
    end_time = perf_counter()
    print(f"Owlii basketball video created successfully: {output_path}")
    print(f"Total time taken: {end_time - start_time:.2f} seconds")
    print(f"Average time per frame: {(end_time - start_time) / len(frames):.2f} seconds")
    print(f"Processed {len(frames)} frames with resolution {resolution}x{resolution}x{resolution}")
    print(f"Scale factor: {scale_factor}")
    print(f"🎯 Global normalization: All frames normalized with same global bounds")
    print(f"🏀 This should eliminate swaying caused by per-frame normalization!")

if __name__ == "__main__":
    multiprocessing.freeze_support()
    
    # Check for existing tmp files and determine what needs to be generated
    temp_dir = "tmp/owlii_basketball_frames"
    max_frames = 600  # Should match the max_frames in create_owlii_basketball_video
    
    if os.path.exists(temp_dir):
        existing_files = glob.glob(os.path.join(temp_dir, "basketball_frame_*.splv"))
        existing_frame_numbers = set()
        
        for file_path in existing_files:
            filename = os.path.basename(file_path)
            # Extract frame number from filename like "basketball_frame_000123.splv"
            try:
                frame_num = int(filename.split('_')[2].split('.')[0])
                if 0 <= frame_num < max_frames:
                    existing_frame_numbers.add(frame_num)
            except (IndexError, ValueError):
                continue
        
        missing_frames = []
        for i in range(max_frames):
            if i not in existing_frame_numbers:
                missing_frames.append(i)
        
        print(f"📁 Found {len(existing_frame_numbers)} existing tmp files")
        print(f"🔍 Missing {len(missing_frames)} frames out of {max_frames}")
        
        if len(missing_frames) == 0:
            print("✅ All frames already exist, creating video from tmp files...")
            create_video_from_temp_files(temp_dir, resolution=512)
        elif len(missing_frames) < max_frames:
            print(f"⚡ Generating only missing frames: {missing_frames[:5]}{'...' if len(missing_frames) > 5 else ''}")
            create_owlii_basketball_video(missing_frames_only=missing_frames)
        else:
            print("🆕 No valid tmp files found, generating complete video...")
            create_owlii_basketball_video()
    else:
        print("🆕 No tmp directory found, generating complete video...")
        create_owlii_basketball_video()