droid_pipeline_episode.py

import os
import sys
import time
import json
import shutil
import warnings
import numpy as np
import pyzed.sl as sl
import torch
from torch.utils.data import Dataset, DataLoader
import glob
import gc
import h5py

from optical_flow_processor import OpticalFlowProcessor
from cpu_processors import CPUProcessor
from image_saver import AsyncImageSaver

# --- Data Loading Class ---
class SVOEpisodeDataset(Dataset):
    """Handles SVO data extraction, initializing ZED resources within each worker."""
    def __init__(self, svo_path, resize_width=None, resize_height=None):
        self.svo_path = svo_path
        self.resize_width = resize_width
        self.resize_height = resize_height
        temp_zed = sl.Camera()
        init_params = self._setup_init_params()
        if temp_zed.open(init_params) != sl.ERROR_CODE.SUCCESS:
            raise IOError(f"Failed to open SVO file for metadata inspection: {self.svo_path}")
        
        self.num_frames = temp_zed.get_svo_number_of_frames()
        cam_params = temp_zed.get_camera_information().camera_configuration.calibration_parameters.left_cam
        self.original_intrinsics = [cam_params.fx, cam_params.fy, cam_params.cx, cam_params.cy]
        
        # Get original image dimensions
        temp_zed.retrieve_image(sl.Mat(), sl.VIEW.LEFT)
        original_height, original_width = temp_zed.get_camera_information().camera_configuration.resolution.height, temp_zed.get_camera_information().camera_configuration.resolution.width
        
        # Calculate scaling factors
        if self.resize_width is not None or self.resize_height is not None:
            scale_x = (self.resize_width / original_width) if self.resize_width is not None else 1.0
            scale_y = (self.resize_height / original_height) if self.resize_height is not None else 1.0
            self.intrinsics = [cam_params.fx * scale_x, cam_params.fy * scale_y, cam_params.cx * scale_x, cam_params.cy * scale_y]
            print(f"📏 Intrinsics scaled for resize {self.resize_width or original_width}x{self.resize_height or original_height}: {self.intrinsics}")
        else:
            self.intrinsics = self.original_intrinsics
        
        temp_zed.close()

        self.zed = None
        self.left_image = None
        self.depth_image = None
        self.runtime_params = None

    def _initialize_worker(self):
        self.zed = sl.Camera()
        init_params = self._setup_init_params()
        err = self.zed.open(init_params)
        if err != sl.ERROR_CODE.SUCCESS:
            raise IOError(f"[Worker PID: {os.getpid()}] Failed to open SVO file: {self.svo_path}, Error: {err}")
        
        self.left_image = sl.Mat()
        self.depth_image = sl.Mat()
        self.runtime_params = self._setup_runtime_params()

    def _setup_init_params(self):
        params = sl.InitParameters()
        params.set_from_svo_file(self.svo_path)
        params.svo_real_time_mode = False
        params.depth_mode = sl.DEPTH_MODE.NEURAL
        params.coordinate_units = sl.UNIT.METER
        params.depth_stabilization = 100
        params.depth_maximum_distance = 2
        return params

    def _setup_runtime_params(self):
        params = sl.RuntimeParameters()
        params.confidence_threshold = 60
        params.texture_confidence_threshold = 100
        return params
        
    def __len__(self):
        if self.num_frames > 0:
            return self.num_frames - 1
        return 0

    def __getitem__(self, index):
        if self.zed is None: self._initialize_worker()
        self.zed.set_svo_position(index)
        
        if self.zed.grab(self.runtime_params) == sl.ERROR_CODE.SUCCESS:
            self.zed.retrieve_image(self.left_image, sl.VIEW.LEFT)
            self.zed.retrieve_measure(self.depth_image, sl.MEASURE.DEPTH)
            rgb_frame = self.left_image.get_data().copy()[..., :3]
            depth_frame = self.depth_image.get_data().copy()
            
            # Resize images if requested
            if self.resize_width is not None or self.resize_height is not None:
                import cv2
                h, w = rgb_frame.shape[:2]
                new_w = self.resize_width if self.resize_width is not None else w
                new_h = self.resize_height if self.resize_height is not None else h
                rgb_frame = cv2.resize(rgb_frame, (new_w, new_h), interpolation=cv2.INTER_LINEAR)
                depth_frame = cv2.resize(depth_frame, (new_w, new_h), interpolation=cv2.INTER_NEAREST)
            
            return rgb_frame, depth_frame
        return None, None
            
    def get_intrinsics(self):
        return self.intrinsics

    def close(self):
        if self.zed is not None:
            self.zed.close()

def collate_fn(batch):
    batch = [b for b in batch if b[0] is not None]
    if not batch: return None
    return torch.utils.data.dataloader.default_collate(batch)

def _process_single_camera(svo_path, camera_save_dir, cpu_processor, physical_gpu_id, num_io_workers, resize_width=None, resize_height=None, is_wrist_camera=False, wrist_extrinsics=None, reference_extrinsics=None, save_depth_and_flow=False):
    """
    Processes a single camera with accurate logging of the physical GPU ID.
    """
    all_rgbs, all_depths = [], []
    intrinsic_params = []
    
    # The local GPU ID for PyTorch is always 0 because of CUDA_VISIBLE_DEVICES
    local_gpu_id = 0
    
    try:
        # --- PHASE 1: Data Loading ---
        log_prefix = f"[{os.getpid()}, GPU {physical_gpu_id}]"
        print(f"{log_prefix} Phase 1: Loading all frames from {os.path.basename(svo_path)}...")
        dataset = SVOEpisodeDataset(svo_path, resize_width=resize_width, resize_height=resize_height)
        data_loader = DataLoader(dataset, batch_size=8, shuffle=False, num_workers=num_io_workers, collate_fn=collate_fn, pin_memory=True)
        intrinsic_params = dataset.get_intrinsics()

        for batch_data in data_loader:
            if batch_data is None: continue
            rgb_batch, depth_batch = batch_data
            all_rgbs.extend(list(rgb_batch.numpy()))
            all_depths.extend(list(depth_batch.numpy()))
        
        frame_count = len(all_rgbs)
        print(f"{log_prefix} Loaded {frame_count} frames into CPU RAM.")

        # --- PHASE 2: Purge ZED SDK ---
        print(f"{log_prefix} Phase 2: Shutting down data loader to clear ZED SDK from GPU...")
        del data_loader
        del dataset
        gc.collect()
        torch.cuda.empty_cache()

        # --- PHASE 3: Optical Flow ---
        if frame_count <= 1:
            print(f"{log_prefix} Not enough frames for optical flow. Skipping.")
            return {'success': True, 'frame_count': frame_count, 'intrinsics': intrinsic_params}

        print(f"{log_prefix} Phase 3: Processing optical flow with PyTorch...")
        optical_flow_processor = OpticalFlowProcessor(use_fp16=False)
        optical_flow_processor.setup_gpu(local_gpu_id, physical_gpu_id) # Use local ID for torch, physical ID for logging
        
        # Use larger batch size if images are resized
        is_resized = resize_width is not None or resize_height is not None
        scale = 1 if not is_resized else (1280 * 720) // (resize_width * resize_height)
        batch_size = 8 * scale
        print(f"{log_prefix} Using batch size {batch_size} for optical flow processing")
        flows = optical_flow_processor.process_batch_on_gpu(np.array(all_rgbs), local_gpu_id, batch_size=batch_size, physical_gpu_id=physical_gpu_id)
        
        # --- PHASE 4: Purge PyTorch ---
        print(f"{log_prefix} Phase 4: Clearing PyTorch model from GPU...")
        del optical_flow_processor
        gc.collect()
        torch.cuda.empty_cache()

        # --- PHASE 5: CPU & Saving ---
        if flows is not None and len(flows) > 0:
            print(f"{log_prefix} Phase 5: Processing masks, scene flow, and saving...")
            data_saver = AsyncImageSaver(camera_save_dir)
            data_saver.start()

            depths1, depths2 = all_depths[:-1], all_depths[1:]
            masks = cpu_processor.process_masks_parallel(list(flows), depths1, depths2)
            
            # Use specialized scene flow calculation for wrist camera
            if is_wrist_camera and wrist_extrinsics is not None and reference_extrinsics is not None:
                scene_flows = cpu_processor.process_wrist_scene_flows_parallel(
                    list(flows), masks, depths1, depths2, intrinsic_params, 
                    wrist_extrinsics, reference_extrinsics
                )
            else:
                scene_flows = cpu_processor.process_scene_flows_parallel(list(flows), masks, depths1, depths2, intrinsic_params)

            # Always save RGB images and scene flow
            for i, rgb in enumerate(all_rgbs): 
                data_saver.add_to_queue({'type': 'rgb', 'index': i, 'data': rgb})
            
            for i in range(len(flows)):
                data_saver.add_to_queue({'type': 'scene_flow', 'index': i, 'data': scene_flows[i]})
            
            # Conditionally save depth and optical flow+mask data
            if save_depth_and_flow:
                for i, depth in enumerate(all_depths): 
                    data_saver.add_to_queue({'type': 'depth', 'index': i, 'data': depth})
                for i in range(len(flows)):
                    data_saver.add_to_queue({'type': 'flow_and_mask', 'index': i, 'flow_data': flows[i], 'mask_data': masks[i]})
            else:
                # Still process depth and flows for scene flow computation, but don't save them
                print(f"{log_prefix} Skipping depth and optical flow+mask saving (storage optimization enabled)")
            
            data_saver.stop()
        
        print(f"{log_prefix} Finished processing for {camera_save_dir}")
        return {'success': True, 'frame_count': frame_count, 'intrinsics': intrinsic_params}

    except Exception as e:
        import traceback
        print(f"❌ [{os.getpid()}, GPU {physical_gpu_id}] Error in camera processing for {camera_save_dir}: {e}\n{traceback.format_exc()}", flush=True)
        return {'success': False}

def run_episode(base_dir, physical_gpu_id, resize_width=None, resize_height=None, include_wrist_camera=False, save_depth_and_flow=False):
    """ Main processing pipeline for a single episode, now aware of its physical GPU ID for logging. """
    pipeline_start_time = time.time()
    log_prefix = f"[{os.getpid()}, GPU {physical_gpu_id}]"
    
    raw_dir = '../droid_raw/droid/1.0.1/'
    droid_dir = '../droid_processed/'
    dir_split = base_dir[len(raw_dir):].split('/')
    save_dir = os.path.join(droid_dir, dir_split[0], '+'.join(dir_split).replace('+success', ''))
    os.makedirs(save_dir, exist_ok=True)

    try:
        json_files = glob.glob(f'{base_dir}/metadata_*.json')
        if not json_files: raise FileNotFoundError(f"No metadata JSON found in {base_dir}")
        with open(json_files[0]) as f: metadata = json.load(f)
        
        left_cam_serial = os.path.basename(metadata['left_mp4_path']).split('.')[0]
        right_cam_serial = os.path.basename(metadata['right_mp4_path']).split('.')[0]

        svo_path_left = f"{base_dir}/recordings/SVO/{left_cam_serial}.svo"
        svo_path_right = f"{base_dir}/recordings/SVO/{right_cam_serial}.svo"

        save_dir_left = os.path.join(save_dir, 'camera_left')
        save_dir_right = os.path.join(save_dir, 'camera_right')
        
        # Add wrist camera if enabled
        if include_wrist_camera:
            wrist_cam_serial = os.path.basename(metadata['wrist_mp4_path']).split('.')[0]
            svo_path_wrist = f"{base_dir}/recordings/SVO/{wrist_cam_serial}.svo"
            save_dir_wrist = os.path.join(save_dir, 'camera_wrist')

        total_cpus = os.cpu_count()
        num_io_workers = min(8, max(4, total_cpus // 8))
        num_cpu_workers = min(64, max(1, total_cpus - num_io_workers - 2))
        
        cpu_processor = CPUProcessor(num_workers=num_cpu_workers, physical_gpu_id=physical_gpu_id)

        print(f"{log_prefix} Processing LEFT camera for episode {os.path.basename(base_dir)}")
        result_left = _process_single_camera(svo_path_left, save_dir_left, cpu_processor, physical_gpu_id, num_io_workers, resize_width, resize_height, save_depth_and_flow=save_depth_and_flow)

        print(f"{log_prefix} Processing RIGHT camera for episode {os.path.basename(base_dir)}")
        result_right = _process_single_camera(svo_path_right, save_dir_right, cpu_processor, physical_gpu_id, num_io_workers, resize_width, resize_height, save_depth_and_flow=save_depth_and_flow)
        
        # Process wrist camera if enabled
        result_wrist = None
        if include_wrist_camera:
            print(f"{log_prefix} Processing WRIST camera for episode {os.path.basename(base_dir)}")
            
            # Load trajectory.h5 to get extrinsics for wrist camera scene flow
            with h5py.File(os.path.join(base_dir, 'trajectory.h5'), "r") as f:
                action_dict, observation_dict, state_dict = f['action'], f['observation'], f['observation']['robot_state']
                wrist_extrinsics = np.array(observation_dict['camera_extrinsics'][f'{wrist_cam_serial}_left'])
                
                # For each pair (i, i+1), use frame i as reference and frame i+1 as current
                # So for pair (i, i+1): reference is frame i, current is frame i+1
                reference_extrinsics = [wrist_extrinsics[i] for i in range(len(wrist_extrinsics) - 1)]  # Frame i
                wrist_extrinsics_list = [wrist_extrinsics[i+1] for i in range(len(wrist_extrinsics) - 1)]  # Frame i+1
            
            result_wrist = _process_single_camera(svo_path_wrist, save_dir_wrist, cpu_processor, physical_gpu_id, num_io_workers, resize_width, resize_height, 
                                                is_wrist_camera=True, wrist_extrinsics=wrist_extrinsics_list, reference_extrinsics=reference_extrinsics, save_depth_and_flow=save_depth_and_flow)
        
        # Check if all required cameras processed successfully
        if not (result_left['success'] and result_right['success']):
            raise RuntimeError("Processing failed for one or both main cameras.")
        if include_wrist_camera and result_wrist and not result_wrist['success']:
            raise RuntimeError("Processing failed for wrist camera.")

        # --- METADATA GENERATION ---
        print(f"{log_prefix} Generating and saving final metadata.json...")
        
        shutil.copyfile(os.path.join(base_dir, 'trajectory.h5'), os.path.join(save_dir, 'trajectory.h5'))
        
        number_of_frames = min(result_left['frame_count'], result_right['frame_count'])
        if include_wrist_camera and result_wrist:
            number_of_frames = min(number_of_frames, result_wrist['frame_count'])
        if number_of_frames == 0:
            warnings.warn(f"Episode {base_dir} resulted in 0 frames.")
        
        metadata['left_cam_serial'] = left_cam_serial
        metadata['right_cam_serial'] = right_cam_serial
        metadata["left_cam_intrinsics"] = result_left['intrinsics']
        metadata["right_cam_intrinsics"] = result_right['intrinsics']
        metadata["number_of_frames"] = number_of_frames
        
        # Add wrist camera metadata if enabled
        if include_wrist_camera and result_wrist:
            metadata['wrist_cam_serial'] = wrist_cam_serial
            metadata["wrist_cam_intrinsics"] = result_wrist['intrinsics']
            
            # Load trajectory.h5 to get wrist camera extrinsics
            with h5py.File(os.path.join(save_dir, 'trajectory.h5'), "r") as f:
                action_dict, observation_dict, state_dict = f['action'], f['observation'], f['observation']['robot_state']
                wrist_extrinsics = np.array(observation_dict['camera_extrinsics'][f'{wrist_cam_serial}_left'])
                metadata["h5_wrist_extrinsics"] = wrist_extrinsics.tolist()
        
        with open(os.path.join(save_dir, 'metadata.json'), 'w') as f:
            json.dump(metadata, f, indent=4)

        print(f"🎉 {log_prefix} Full episode completed in {time.time() - pipeline_start_time:.2f}s")
        return True, "Success"

    except Exception as e:
        import traceback
        error_msg = f"Error in episode pipeline for {base_dir}: {e}\n{traceback.format_exc()}"
        print(error_msg, flush=True)
        return False, str(e)

    finally:
        gc.collect()
        if torch.cuda.is_available():
            torch.cuda.empty_cache()