Upload extract-embeddings.py with huggingface_hub

extract-embeddings.py

@@ -0,0 +1,939 @@
+from huggingface_hub import snapshot_download
+import sys
+import os
+import torch
+import numpy as np
+import pandas as pd
+import json
+import re
+import pydicom
+from datetime import datetime
+from concurrent.futures import ThreadPoolExecutor, ProcessPoolExecutor
+from pathlib import Path
+import threading
+import multiprocessing as mp
+
+# Download and setup model
+model_path = snapshot_download(repo_id="Lab-Rasool/sybil")
+sys.path.append(model_path)
+
+from modeling_sybil_hf import SybilHFWrapper
+from configuration_sybil import SybilConfig
+
+def load_model(device_id=0):
+    """
+    Load and initialize the Sybil model once.
+    
+    Args:
+        device_id: GPU device ID to load model on
+    
+    Returns:
+        Initialized SybilHFWrapper model
+    """
+    print(f"Loading Sybil model on GPU {device_id}...")
+    config = SybilConfig()
+    model = SybilHFWrapper(config)
+    
+    # Move model to specific GPU
+    device = torch.device(f'cuda:{device_id}')
+    
+    # CRITICAL: Set the model's internal device attribute
+    # This ensures preprocessing moves data to the correct GPU
+    model.device = device
+    
+    # Move all ensemble models to the correct GPU
+    for m in model.models:
+        m.to(device)
+        m.eval()
+    
+    print(f"Model loaded successfully on GPU {device_id}!")
+    print(f"  Model internal device: {model.device}")
+    return model, device
+
+def is_localizer_scan(dicom_folder):
+    """
+    Check if a DICOM folder contains a localizer/scout scan.
+    Based on preprocessing.py logic.
+    
+    Returns:
+        Tuple of (is_localizer, reason)
+    """
+    folder_path = Path(dicom_folder)
+    folder_name = folder_path.name.lower()
+    localizer_keywords = ['localizer', 'scout', 'topogram', 'surview', 'scanogram']
+    
+    # Check folder name
+    if any(keyword in folder_name for keyword in localizer_keywords):
+        return True, f"Folder name contains localizer keyword: {folder_name}"
+    
+    try:
+        dcm_files = list(folder_path.glob("*.dcm"))
+        if not dcm_files:
+            return False, "No DICOM files found"
+        
+        # Check first few DICOM files for localizer metadata
+        sample_files = dcm_files[:min(3, len(dcm_files))]
+        for dcm_file in sample_files:
+            try:
+                dcm = pydicom.dcmread(str(dcm_file), stop_before_pixels=True)
+                
+                # Check ImageType field
+                if hasattr(dcm, 'ImageType'):
+                    image_type_str = ' '.join(str(val).lower() for val in dcm.ImageType)
+                    if any(keyword in image_type_str for keyword in localizer_keywords):
+                        return True, f"ImageType indicates localizer: {dcm.ImageType}"
+                
+                # Check SeriesDescription field
+                if hasattr(dcm, 'SeriesDescription'):
+                    if any(keyword in dcm.SeriesDescription.lower() for keyword in localizer_keywords):
+                        return True, f"SeriesDescription indicates localizer: {dcm.SeriesDescription}"
+            except Exception as e:
+                continue
+    except Exception as e:
+        pass
+    
+    return False, "Not a localizer scan"
+
+def extract_timepoint_from_path(scan_dir):
+    """
+    Extract timepoint from scan directory path based on year.
+    1999 -> T0, 2000 -> T1, 2001 -> T2, etc.
+    
+    Looks for year patterns in folder names in date format MM-DD-YYYY.
+    
+    Args:
+        scan_dir: Directory path string
+        
+    Returns:
+        Timepoint string (e.g., 'T0', 'T1', 'T2') or None if not found
+    """
+    # Split path into components
+    path_parts = scan_dir.split('/')
+    
+    # Look for date patterns like "01-02-2000-NLST-LSS"
+    # Pattern: Date format MM-DD-YYYY at the start of a folder name
+    date_pattern = r'^\d{2}-\d{2}-(19\d{2}|20\d{2})'
+    
+    base_year = 1999
+    
+    for part in path_parts:
+        # Check for date pattern (e.g., "01-02-2000-NLST-LSS-50335")
+        match = re.match(date_pattern, part)
+        if match:
+            year = int(match.group(1))
+            if 1999 <= year <= 2010:  # Reasonable range for NLST
+                timepoint_num = year - base_year
+                print(f"  DEBUG: Found year {year} in '{part}' -> T{timepoint_num}")
+                return f'T{timepoint_num}'
+    
+    return None
+
+def extract_embedding_single_model(model_idx, ensemble_model, pixel_values, device):
+    """
+    Extract embedding from a single ensemble model.
+    
+    Args:
+        model_idx: Index of the model in the ensemble
+        ensemble_model: Single model from the ensemble
+        pixel_values: Preprocessed pixel values tensor (already on correct device)
+        device: Device to run on (e.g., cuda:0, cuda:1)
+        
+    Returns:
+        numpy array of embeddings from this model
+    """
+    embeddings_buffer = []
+    
+    def create_hook(buffer):
+        def hook(module, input, output):
+            # Capture the output of ReLU layer (before dropout)
+            buffer.append(output.detach().cpu())
+        return hook
+    
+    # Register hook on the ReLU layer (this is AFTER pooling, BEFORE dropout/classification)
+    hook_handle = ensemble_model.relu.register_forward_hook(create_hook(embeddings_buffer))
+    
+    # Run forward pass on THIS model only with keyword argument
+    with torch.no_grad():
+        _ = ensemble_model(pixel_values=pixel_values)
+    
+    # Remove hook
+    hook_handle.remove()
+    
+    # Get the embeddings (should be shape [1, 512])
+    if embeddings_buffer:
+        embedding = embeddings_buffer[0].numpy().squeeze()
+        print(f"Model {model_idx + 1}: Embedding shape = {embedding.shape}")
+        return embedding
+    return None
+
+def extract_embeddings(model, dicom_paths, device, use_parallel=True):
+    """
+    Extract embeddings from the layer after ReLU, before Dropout.
+    Processes ensemble models in parallel for speed.
+
+    Args:
+        model: Pre-loaded SybilHFWrapper model
+        dicom_paths: List of DICOM file paths
+        device: Device to run on (e.g., cuda:0, cuda:1)
+        use_parallel: If True, process ensemble models in parallel
+
+    Returns:
+        numpy array of shape (512,) - averaged embeddings across ensemble
+    """
+    # Preprocess ONCE (not 5 times!)
+    # The model's preprocessing handles moving data to the correct device
+    with torch.no_grad():
+        # Get the preprocessed input by calling the wrapper's preprocess_dicom method
+        # This returns the tensor that would be fed to each ensemble model
+        pixel_values = model.preprocess_dicom(dicom_paths)
+    
+    if use_parallel:
+        # Process all ensemble models in parallel using ThreadPoolExecutor
+        all_embeddings = []
+        
+        with ThreadPoolExecutor(max_workers=len(model.models)) as executor:
+            # Submit all models for parallel processing with the SAME preprocessed input
+            futures = [
+                executor.submit(extract_embedding_single_model, model_idx, ensemble_model, pixel_values, device)
+                for model_idx, ensemble_model in enumerate(model.models)
+            ]
+            
+            # Collect results as they complete
+            for future in futures:
+                embedding = future.result()
+                if embedding is not None:
+                    all_embeddings.append(embedding)
+    else:
+        # Sequential processing (original implementation)
+        all_embeddings = []
+        for model_idx, ensemble_model in enumerate(model.models):
+            embedding = extract_embedding_single_model(model_idx, ensemble_model, pixel_values, device)
+            if embedding is not None:
+                all_embeddings.append(embedding)
+    
+    # Average embeddings across ensemble
+    averaged_embedding = np.mean(all_embeddings, axis=0)
+    return averaged_embedding
+
+def check_directory_for_dicoms(dirpath):
+    """
+    Check a single directory for valid DICOM files.
+    Returns (dirpath, num_files, subject_id, filter_reason) or None if invalid.
+    """
+    try:
+        # Quick check: does this directory have .dcm files?
+        dcm_files = [f for f in os.listdir(dirpath) 
+                     if f.endswith('.dcm') and os.path.isfile(os.path.join(dirpath, f))]
+        
+        if not dcm_files:
+            return None
+        
+        num_files = len(dcm_files)
+        
+        # Filter out scans with 1-2 DICOM files (likely localizers)
+        if num_files <= 2:
+            return (dirpath, num_files, None, 'too_few_slices')
+        
+        # Check if it's a localizer scan
+        is_loc, _ = is_localizer_scan(dirpath)
+        if is_loc:
+            return (dirpath, num_files, None, 'localizer')
+        
+        # Extract subject ID (PID) from path
+        # Path structure: /NLST/<PID>/<date-info>/<scan-info>
+        # Example: /NLST/106639/01-02-1999-NLST-LSS-45699/1.000000-0OPLGEHSQXAnullna...
+        path_parts = dirpath.rstrip('/').split('/')
+        
+        # Find the PID: it's the part after 'NLST' directory
+        try:
+            nlst_idx = path_parts.index('NLST')
+            subject_id = path_parts[nlst_idx + 1]  # PID is right after 'NLST'
+        except (ValueError, IndexError):
+            # Fallback to old logic if path structure is different
+            subject_id = path_parts[-3] if len(path_parts) >= 3 else path_parts[-1]
+        
+        return (dirpath, num_files, subject_id, 'valid')
+        
+    except Exception as e:
+        return None
+
+def save_directory_cache(dicom_dirs, cache_file):
+    """
+    Save the list of DICOM directories to a cache file.
+    
+    Args:
+        dicom_dirs: List of directory paths
+        cache_file: Path to cache file
+    """
+    print(f"\n💾 Saving directory cache to {cache_file}...")
+    cache_data = {
+        "timestamp": datetime.now().isoformat(),
+        "num_directories": len(dicom_dirs),
+        "directories": dicom_dirs
+    }
+    with open(cache_file, 'w') as f:
+        json.dump(cache_data, f, indent=2)
+    print(f"✓ Cache saved with {len(dicom_dirs)} directories\n")
+
+def load_directory_cache(cache_file):
+    """
+    Load the list of DICOM directories from a cache file.
+    
+    Args:
+        cache_file: Path to cache file
+        
+    Returns:
+        List of directory paths, or None if cache doesn't exist or is invalid
+    """
+    if not os.path.exists(cache_file):
+        return None
+    
+    try:
+        with open(cache_file, 'r') as f:
+            cache_data = json.load(f)
+        
+        dicom_dirs = cache_data.get("directories", [])
+        timestamp = cache_data.get("timestamp", "unknown")
+        
+        print(f"\n✓ Loaded directory cache from {cache_file}")
+        print(f"  Cache created: {timestamp}")
+        print(f"  Directories: {len(dicom_dirs)}\n")
+        
+        return dicom_dirs
+    except Exception as e:
+        print(f"⚠️  Failed to load cache: {e}")
+        return None
+
+def find_dicom_directories(root_dir, max_subjects=None, num_workers=12, cache_file=None, filter_pids=None):
+    """
+    Walk through directory tree and find all directories containing DICOM files.
+    Uses parallel processing for much faster scanning of large directory trees.
+    Only returns leaf directories (directories with .dcm files, not their parents).
+    Filters out localizer scans with 1-2 DICOM files.
+    
+    Args:
+        root_dir: Root directory to search
+        max_subjects: Optional maximum number of unique subjects to process (None = all)
+        num_workers: Number of parallel workers for directory scanning (default: 12)
+        cache_file: Optional path to cache file for saving/loading directory list
+        filter_pids: Optional set of PIDs to filter (only include these subjects)
+        
+    Returns:
+        List of directory paths containing .dcm files
+    """
+    # Try to load from cache first
+    if cache_file:
+        cached_dirs = load_directory_cache(cache_file)
+        if cached_dirs is not None:
+            print("✓ Using cached directory list (skipping scan)")
+            
+            # Apply PID filter if specified
+            if filter_pids:
+                print(f"  Filtering to {len(filter_pids)} PIDs from CSV...")
+                filtered_dirs = []
+                for d in cached_dirs:
+                    # Extract PID from path: /NLST/<PID>/<date>/<scan>
+                    path_parts = d.rstrip('/').split('/')
+                    try:
+                        nlst_idx = path_parts.index('NLST')
+                        subject_id = path_parts[nlst_idx + 1]
+                    except (ValueError, IndexError):
+                        subject_id = path_parts[-3] if len(path_parts) >= 3 else path_parts[-1]
+                    
+                    if subject_id in filter_pids:
+                        filtered_dirs.append(d)
+                print(f"  ✓ Found {len(filtered_dirs)} scans matching PIDs")
+                return filtered_dirs
+            
+            # Still apply max_subjects limit if specified
+            if max_subjects:
+                subjects_seen = set()
+                filtered_dirs = []
+                for d in cached_dirs:
+                    # Extract PID from path: /NLST/<PID>/<date>/<scan>
+                    path_parts = d.rstrip('/').split('/')
+                    try:
+                        nlst_idx = path_parts.index('NLST')
+                        subject_id = path_parts[nlst_idx + 1]
+                    except (ValueError, IndexError):
+                        subject_id = path_parts[-3] if len(path_parts) >= 3 else path_parts[-1]
+                    
+                    # Check if we should include this scan
+                    # Add scan if: (1) already collecting this subject, OR (2) under subject limit
+                    if subject_id in subjects_seen:
+                        # Already collecting this subject - add this scan
+                        filtered_dirs.append(d)
+                    elif len(subjects_seen) < max_subjects:
+                        # New subject and under limit - start collecting this subject
+                        subjects_seen.add(subject_id)
+                        filtered_dirs.append(d)
+                    
+                    # Stop once we have enough subjects
+                    if len(subjects_seen) >= max_subjects:
+                        # Count remaining scans from these subjects
+                        remaining_count = 0
+                        for remaining_d in cached_dirs[cached_dirs.index(d)+1:]:
+                            remaining_parts = remaining_d.rstrip('/').split('/')
+                            try:
+                                remaining_nlst_idx = remaining_parts.index('NLST')
+                                remaining_subject_id = remaining_parts[remaining_nlst_idx + 1]
+                            except (ValueError, IndexError):
+                                remaining_subject_id = remaining_parts[-3] if len(remaining_parts) >= 3 else remaining_parts[-1]
+                            if remaining_subject_id in subjects_seen:
+                                filtered_dirs.append(remaining_d)
+                        break
+                
+                print(f"  ✓ Limited to {len(subjects_seen)} subjects ({len(filtered_dirs)} total scans)")
+                return filtered_dirs
+            return cached_dirs
+    
+    print(f"Starting parallel directory scan with {num_workers} workers...")
+    print("This is much faster than the old sequential scan!")
+    
+    # Phase 1: Fast parallel scan to find all directories with DICOM files
+    print("\nPhase 1: Scanning filesystem for DICOM directories...")
+    start_time = datetime.now()
+    
+    # Collect all directories first (fast)
+    all_dirs = []
+    for dirpath, dirnames, filenames in os.walk(root_dir):
+        # Quick check: if directory has .dcm files, add to list
+        if any(f.endswith('.dcm') for f in filenames):
+            all_dirs.append(dirpath)
+    
+    print(f"Found {len(all_dirs)} potential DICOM directories in {(datetime.now() - start_time).total_seconds():.1f}s")
+    
+    # Phase 2: Parallel validation and filtering
+    print(f"\nPhase 2: Validating directories in parallel ({num_workers} workers)...")
+    
+    from concurrent.futures import ProcessPoolExecutor, as_completed
+    
+    dicom_dirs = []
+    subjects_found = set()
+    filtered_stats = {'localizers': 0, 'too_few_slices': 0}
+    
+    with ProcessPoolExecutor(max_workers=num_workers) as executor:
+        # Submit all directories for checking
+        future_to_dir = {executor.submit(check_directory_for_dicoms, d): d for d in all_dirs}
+        
+        # Process results as they complete
+        for i, future in enumerate(as_completed(future_to_dir), 1):
+            # Print progress every 1000 dirs (more frequent for visibility)
+            if i % 1000 == 0:
+                elapsed = (datetime.now() - start_time).total_seconds()
+                rate = i / elapsed if elapsed > 0 else 0
+                remaining = (len(all_dirs) - i) / rate if rate > 0 else 0
+                print(f"  [{i}/{len(all_dirs)}] Found: {len(dicom_dirs)} scans from {len(subjects_found)} PIDs | "
+                      f"Filtered: {filtered_stats['localizers'] + filtered_stats['too_few_slices']} | "
+                      f"ETA: {remaining/60:.1f} min")
+            
+            try:
+                result = future.result()
+                if result is None:
+                    continue
+                
+                dirpath, num_files, subject_id, status = result
+                
+                if status == 'too_few_slices':
+                    filtered_stats['too_few_slices'] += 1
+                elif status == 'localizer':
+                    filtered_stats['localizers'] += 1
+                elif status == 'valid':
+                    # Check PID filter
+                    if filter_pids is not None and subject_id not in filter_pids:
+                        continue
+                    
+                    # Check subject limit
+                    if max_subjects is not None and subject_id not in subjects_found and len(subjects_found) >= max_subjects:
+                        continue
+                    
+                    subjects_found.add(subject_id)
+                    dicom_dirs.append(dirpath)
+                    
+                    # Print when we find a new PID match (helpful for filtered runs)
+                    if filter_pids and len(dicom_dirs) % 100 == 1:
+                        print(f"  ✓ Found {len(dicom_dirs)} scans so far ({len(subjects_found)} unique PIDs)")
+                    
+                    # Stop if we've hit subject limit
+                    if max_subjects is not None and len(subjects_found) >= max_subjects:
+                        print(f"\n✓ Reached limit of {max_subjects} subjects. Stopping search.")
+                        # Cancel remaining futures
+                        for f in future_to_dir:
+                            f.cancel()
+                        break
+                        
+            except Exception as e:
+                continue
+    
+    scan_time = (datetime.now() - start_time).total_seconds()
+    
+    print(f"\n{'='*80}")
+    print(f"Directory Scan Complete in {scan_time:.1f}s ({scan_time/60:.1f} minutes)")
+    print(f"{'='*80}")
+    print(f"Filtering Summary:")
+    print(f"  ✅ Valid scans found: {len(dicom_dirs)}")
+    print(f"  🚫 Localizers filtered: {filtered_stats['localizers']}")
+    print(f"  ⏭️  Too few slices (≤2) filtered: {filtered_stats['too_few_slices']}")
+    print(f"  📊 Unique subjects: {len(subjects_found)}")
+    print(f"  ⚡ Speed: {len(all_dirs)/scan_time:.0f} dirs/second")
+    print(f"{'='*80}\n")
+    
+    # Save to cache if specified
+    if cache_file:
+        save_directory_cache(dicom_dirs, cache_file)
+    
+    return dicom_dirs
+
+def prepare_scan_metadata(scan_dir):
+    """
+    Prepare metadata for a scan without processing.
+    
+    Args:
+        scan_dir: Directory containing DICOM files for one scan
+        
+    Returns:
+        tuple: (dicom_file_paths, num_files, subject_id, scan_id)
+    """
+    # Count DICOM files (ensure they are actual files, not directories)
+    dicom_files = [f for f in os.listdir(scan_dir) 
+                   if f.endswith('.dcm') and os.path.isfile(os.path.join(scan_dir, f))]
+    num_dicom_files = len(dicom_files)
+    
+    if num_dicom_files == 0:
+        raise ValueError("No valid DICOM files found")
+    
+    # Create list of full paths to DICOM files
+    dicom_file_paths = [os.path.join(scan_dir, f) for f in dicom_files]
+    
+    # Parse directory path to extract identifiers
+    # Path structure: /NLST/<PID>/<date-info>/<scan-info>
+    path_parts = scan_dir.rstrip('/').split('/')
+    scan_id = path_parts[-1] if path_parts[-1] else path_parts[-2]
+    
+    # Extract PID from path
+    try:
+        nlst_idx = path_parts.index('NLST')
+        subject_id = path_parts[nlst_idx + 1]  # PID is right after 'NLST'
+    except (ValueError, IndexError):
+        # Fallback to old logic
+        subject_id = path_parts[-3] if len(path_parts) >= 3 else path_parts[-1]
+    
+    return dicom_file_paths, num_dicom_files, subject_id, scan_id
+
+def save_checkpoint(all_embeddings, all_metadata, failed, output_dir, checkpoint_num):
+    """
+    Save a checkpoint of embeddings and metadata.
+    
+    Args:
+        all_embeddings: List of embedding arrays
+        all_metadata: List of metadata dictionaries
+        failed: List of failed scans
+        output_dir: Output directory
+        checkpoint_num: Checkpoint number
+    """
+    print(f"\n💾 Saving checkpoint {checkpoint_num}...")
+    
+    # Convert embeddings to array
+    embeddings_array = np.array(all_embeddings)
+    embedding_dim = int(embeddings_array.shape[1]) if len(embeddings_array.shape) > 1 else int(embeddings_array.shape[0])
+    
+    # Create DataFrame
+    df_data = {
+        'case_number': [m['case_number'] for m in all_metadata],
+        'subject_id': [m['subject_id'] for m in all_metadata],
+        'scan_id': [m['scan_id'] for m in all_metadata],
+        'timepoint': [m.get('timepoint') for m in all_metadata],
+        'dicom_directory': [m['dicom_directory'] for m in all_metadata],
+        'num_dicom_files': [m['num_dicom_files'] for m in all_metadata],
+        'embedding_index': [m['embedding_index'] for m in all_metadata],
+        'embedding': list(embeddings_array)
+    }
+    df = pd.DataFrame(df_data)
+    
+    # Save checkpoint parquet
+    checkpoint_path = os.path.join(output_dir, f"checkpoint_{checkpoint_num}_embeddings.parquet")
+    df.to_parquet(checkpoint_path, index=False, compression='snappy')
+    print(f"  ✓ Saved embeddings checkpoint: {checkpoint_path}")
+    
+    # Save checkpoint metadata
+    checkpoint_metadata = {
+        "checkpoint_num": checkpoint_num,
+        "timestamp": datetime.now().isoformat(),
+        "total_scans": len(all_embeddings),
+        "failed_scans": len(failed),
+        "embedding_shape": list(embeddings_array.shape),
+        "scans": all_metadata,
+        "failed_scans": failed
+    }
+    metadata_path = os.path.join(output_dir, f"checkpoint_{checkpoint_num}_metadata.json")
+    with open(metadata_path, 'w') as f:
+        json.dump(checkpoint_metadata, f, indent=2)
+    print(f"  ✓ Saved metadata checkpoint: {metadata_path}")
+    print(f"💾 Checkpoint {checkpoint_num} complete!\n")
+
+def process_scan(model, device, scan_dir):
+    """
+    Process a single scan directory and extract embeddings.
+    
+    Args:
+        model: Pre-loaded SybilHFWrapper model
+        device: Device to run on (e.g., cuda:0, cuda:1)
+        scan_dir: Directory containing DICOM files for one scan
+        
+    Returns:
+        tuple: (embeddings, scan_metadata)
+    """
+    dicom_file_paths, num_dicom_files, subject_id, scan_id = prepare_scan_metadata(scan_dir)
+    
+    print(f"\nProcessing: {scan_dir}")
+    print(f"DICOM files: {num_dicom_files}")
+    
+    # Extract embeddings
+    embeddings = extract_embeddings(model, dicom_file_paths, device)
+    
+    print(f"Embedding shape: {embeddings.shape}")
+    
+    # Extract timepoint from path (e.g., 1999 -> T0, 2000 -> T1)
+    timepoint = extract_timepoint_from_path(scan_dir)
+    if timepoint:
+        print(f"Timepoint: {timepoint}")
+    else:
+        print(f"Timepoint: Not detected")
+    
+    # Create metadata for this scan
+    scan_metadata = {
+        "case_number": subject_id,  # Case number (e.g., 205749)
+        "subject_id": subject_id,
+        "scan_id": scan_id,
+        "timepoint": timepoint,  # T0, T1, T2, etc. or None
+        "dicom_directory": scan_dir,
+        "num_dicom_files": num_dicom_files,
+        "embedding_index": None,  # Will be set later
+        "statistics": {
+            "mean": float(np.mean(embeddings)),
+            "std": float(np.std(embeddings)),
+            "min": float(np.min(embeddings)),
+            "max": float(np.max(embeddings))
+        }
+    }
+    
+    return embeddings, scan_metadata
+
+# Main execution
+if __name__ == "__main__":
+    import argparse
+    
+    # Parse command line arguments
+    parser = argparse.ArgumentParser(description='Extract Sybil embeddings from DICOM scans')
+    parser.add_argument('--pid-csv', type=str, default=None,
+                        help='CSV file with "pid" column to filter subjects (e.g., subsets/hybridModels-train.csv)')
+    parser.add_argument('--output-dir', type=str, default='embeddings_output_full',
+                        help='Output directory for embeddings')
+    parser.add_argument('--max-subjects', type=int, default=None,
+                        help='Maximum number of subjects to process (for testing)')
+    parser.add_argument('--num-parallel', type=int, default=1,
+                        help='Number of parallel scans (default: 1 for sequential processing)')
+    parser.add_argument('--num-gpus', type=int, default=1,
+                        help='Number of GPUs to use (default: 1)')
+    args = parser.parse_args()
+    
+    # ==========================================
+    # CONFIGURATION - Edit these variables as needed
+    # ==========================================
+    root_dir = "/roshare/nlst_global/data_23Dec2024/manifest-NLST_allCT/NLST"
+    output_dir = args.output_dir
+    max_subjects = args.max_subjects
+    num_gpus = args.num_gpus  # Use specified number of GPUs
+    num_parallel_scans = args.num_parallel
+    num_scan_workers = 12  # Number of workers for parallel directory scanning (fast filesystem scan)
+    checkpoint_interval = 20000  # Save checkpoint every N scans
+    
+    # Always use the main cache file from the full run
+    main_cache = "embeddings_output_full/directory_cache.json"
+    if os.path.exists(main_cache):
+        cache_file = main_cache
+        print(f"✓ Found main directory cache: {main_cache}")
+    else:
+        cache_file = os.path.join(output_dir, "directory_cache.json")
+    
+    # Load PIDs from CSV if provided
+    filter_pids = None
+    if args.pid_csv:
+        print(f"Loading subject PIDs from: {args.pid_csv}")
+        import pandas as pd
+        csv_data = pd.read_csv(args.pid_csv)
+        filter_pids = set(str(pid) for pid in csv_data['pid'].unique())
+        print(f"  Found {len(filter_pids)} unique PIDs to extract")
+        print(f"  Examples: {list(filter_pids)[:5]}")
+    
+    # Create output directory
+    os.makedirs(output_dir, exist_ok=True)
+    
+    # Verify root directory exists
+    if not os.path.exists(root_dir):
+        raise ValueError(f"Root directory does not exist: {root_dir}")
+    
+    print(f"\nSearching for DICOM directories in: {root_dir}")
+    if filter_pids:
+        print(f"Filtering to {len(filter_pids)} PIDs from CSV")
+    if max_subjects:
+        print(f"Limiting to {max_subjects} subjects")
+    print("⚡ Using parallel directory scanning for speed!\n")
+    
+    # Find all directories containing DICOM files (FAST with parallel processing!)
+    # Will use cached directory list if available, otherwise scan and save cache
+    dicom_dirs = find_dicom_directories(root_dir, max_subjects=max_subjects, 
+                                       num_workers=num_scan_workers, cache_file=cache_file,
+                                       filter_pids=filter_pids)
+    
+    if len(dicom_dirs) == 0:
+        raise ValueError(f"No directories with DICOM files found in {root_dir}")
+    
+    print(f"\n{'='*80}")
+    print(f"Found {len(dicom_dirs)} directories containing DICOM files")
+    print(f"{'='*80}\n")
+    
+    # Detect and load models on multiple GPUs
+    print(f"🎮 Detected {num_gpus} GPU(s)")
+    print(f"🚀 Will process {num_parallel_scans} scans in parallel ({num_parallel_scans // num_gpus} per GPU)")
+    print(f"💾 Checkpoints will be saved every {checkpoint_interval} scans\n")
+    
+    # Load models on each GPU
+    models_and_devices = []
+    for gpu_id in range(num_gpus):
+        model, device = load_model(gpu_id)
+        models_and_devices.append((model, device, gpu_id))
+    
+    # Process each scan directory and collect all embeddings
+    all_embeddings = []
+    all_metadata = []
+    failed = []
+    checkpoint_counter = 0
+    
+    if num_parallel_scans > 1:
+        # Parallel processing of multiple scans across multiple GPUs
+        print(f"Processing {num_parallel_scans} scans in parallel across {num_gpus} GPU(s)...")
+        print(f"Note: This requires ~{(num_parallel_scans // num_gpus) * 10}GB VRAM per GPU.\n")
+        
+        from functools import partial
+        from concurrent.futures import as_completed
+        
+        # Process scans in batches for checkpoint saving
+        batch_size = checkpoint_interval
+        num_batches = (len(dicom_dirs) + batch_size - 1) // batch_size
+        
+        for batch_idx in range(num_batches):
+            start_idx = batch_idx * batch_size
+            end_idx = min(start_idx + batch_size, len(dicom_dirs))
+            batch_dirs = dicom_dirs[start_idx:end_idx]
+            
+            print(f"\n{'='*80}")
+            print(f"Processing batch {batch_idx + 1}/{num_batches} (scans {start_idx + 1} to {end_idx})")
+            print(f"{'='*80}\n")
+            
+            # Use ThreadPoolExecutor for parallel scan processing
+            # IMPORTANT: max_workers limits concurrent execution to prevent OOM
+            with ThreadPoolExecutor(max_workers=num_parallel_scans) as executor:
+                # Submit scans in controlled batches to avoid memory issues
+                # We submit only max_workers scans at once, then submit more as they complete
+                future_to_info = {}
+                scan_queue = list(enumerate(batch_dirs))
+                scans_submitted = 0
+                
+                # Submit initial batch (up to max_workers scans)
+                while scan_queue and scans_submitted < num_parallel_scans:
+                    i, scan_dir = scan_queue.pop(0)
+                    # Select GPU in round-robin fashion
+                    gpu_idx = i % num_gpus
+                    model, device, gpu_id = models_and_devices[gpu_idx]
+                    
+                    # Create partial function with model and device
+                    process_func = partial(process_scan, model, device)
+                    future = executor.submit(process_func, scan_dir)
+                    future_to_info[future] = (start_idx + i + 1, scan_dir, gpu_id)
+                    scans_submitted += 1
+                
+                # Process results as they complete and submit new scans
+                while future_to_info:
+                    # Wait for next completion
+                    done_futures = []
+                    for future in list(future_to_info.keys()):
+                        if future.done():
+                            done_futures.append(future)
+                    
+                    if not done_futures:
+                        import time
+                        time.sleep(0.1)
+                        continue
+                    
+                    # Process completed futures
+                    for future in done_futures:
+                        scan_num, scan_dir, gpu_id = future_to_info.pop(future)
+                        try:
+                            print(f"[{scan_num}/{len(dicom_dirs)}] Processing on GPU {gpu_id}...")
+                            embeddings, scan_metadata = future.result()
+                            
+                            # Set the index for this scan
+                            scan_metadata["embedding_index"] = len(all_embeddings)
+                            
+                            # Collect embeddings and metadata
+                            all_embeddings.append(embeddings)
+                            all_metadata.append(scan_metadata)
+                            
+                        except Exception as e:
+                            print(f"ERROR processing {scan_dir}: {e}")
+                            failed.append({"scan_dir": scan_dir, "error": str(e)})
+                        
+                        # Submit next scan from queue
+                        if scan_queue:
+                            i, next_scan_dir = scan_queue.pop(0)
+                            gpu_idx = i % num_gpus
+                            model, device, gpu_id = models_and_devices[gpu_idx]
+                            
+                            process_func = partial(process_scan, model, device)
+                            new_future = executor.submit(process_func, next_scan_dir)
+                            future_to_info[new_future] = (start_idx + i + 1, next_scan_dir, gpu_id)
+            
+            # Save checkpoint after each batch
+            checkpoint_counter += 1
+            save_checkpoint(all_embeddings, all_metadata, failed, output_dir, checkpoint_counter)
+            
+            print(f"Progress: {len(all_embeddings)}/{len(dicom_dirs)} scans completed "
+                  f"({len(all_embeddings)/len(dicom_dirs)*100:.1f}%)\n")
+    else:
+        # Sequential processing (original behavior)
+        model, device, gpu_id = models_and_devices[0]  # Use first GPU
+        
+        for i, scan_dir in enumerate(dicom_dirs, 1):
+            try:
+                print(f"\n[{i}/{len(dicom_dirs)}] Processing scan...")
+                
+                # Process scan and get results
+                embeddings, scan_metadata = process_scan(model, device, scan_dir)
+                
+                # Set the index for this scan
+                scan_metadata["embedding_index"] = len(all_embeddings)
+                
+                # Collect embeddings and metadata
+                all_embeddings.append(embeddings)
+                all_metadata.append(scan_metadata)
+                
+                # Save checkpoint every checkpoint_interval scans
+                if i % checkpoint_interval == 0:
+                    checkpoint_counter += 1
+                    save_checkpoint(all_embeddings, all_metadata, failed, output_dir, checkpoint_counter)
+                
+            except Exception as e:
+                print(f"ERROR processing {scan_dir}: {e}")
+                failed.append({"scan_dir": scan_dir, "error": str(e)})
+    
+    # Convert embeddings list to numpy array
+    # Shape will be (num_scans, embedding_dim)
+    embeddings_array = np.array(all_embeddings)
+    embedding_dim = int(embeddings_array.shape[1]) if len(embeddings_array.shape) > 1 else int(embeddings_array.shape[0])
+    
+    # Create DataFrame with embeddings and metadata for Parquet
+    # Store embeddings as a single array column
+    df_data = {
+        'case_number': [m['case_number'] for m in all_metadata],
+        'subject_id': [m['subject_id'] for m in all_metadata],
+        'scan_id': [m['scan_id'] for m in all_metadata],
+        'timepoint': [m.get('timepoint') for m in all_metadata],  # T0, T1, T2, etc.
+        'dicom_directory': [m['dicom_directory'] for m in all_metadata],
+        'num_dicom_files': [m['num_dicom_files'] for m in all_metadata],
+        'embedding_index': [m['embedding_index'] for m in all_metadata],
+        'embedding': list(embeddings_array)  # Store as list of arrays
+    }
+    
+    # Create DataFrame
+    df = pd.DataFrame(df_data)
+    
+    # Save final complete file as Parquet
+    embeddings_filename = "all_embeddings.parquet"
+    embeddings_path = os.path.join(output_dir, embeddings_filename)
+    df.to_parquet(embeddings_path, index=False, compression='snappy')
+    print(f"\n✅ Saved FINAL embeddings to Parquet: {embeddings_path}")
+    
+    # Create comprehensive metadata JSON
+    dataset_metadata = {
+        "dataset_info": {
+            "root_directory": root_dir,
+            "total_scans": len(all_embeddings),
+            "failed_scans": len(failed),
+            "embedding_shape": list(embeddings_array.shape),
+            "embedding_dim": embedding_dim,
+            "extraction_timestamp": datetime.now().isoformat(),
+            "file_format": "parquet"
+        },
+        "model_info": {
+            "model": "Lab-Rasool/sybil",
+            "layer": "after_relu_before_dropout",
+            "ensemble_averaged": True,
+            "num_ensemble_models": 5
+        },
+        "embeddings_file": embeddings_filename,
+        "parquet_schema": {
+            "metadata_columns": ["case_number", "subject_id", "scan_id", "timepoint", "dicom_directory", "num_dicom_files", "embedding_index"],
+            "embedding_column": "embedding",
+            "embedding_shape": f"({embedding_dim},)",
+            "total_columns": 8,
+            "timepoint_info": "T0=1999, T1=2000, T2=2001, etc. Extracted from year in path. Can be None if not detected."
+        },
+        "filtering_info": {
+            "localizer_detection": "Scans identified as localizers (by folder name or DICOM metadata) are filtered out",
+            "min_slices": "Scans with ≤2 DICOM files are filtered out (likely localizers)",
+            "accepted_scans": len(all_embeddings)
+        },
+        "scans": all_metadata,
+        "failed_scans": failed
+    }
+
+    metadata_filename = "dataset_metadata.json"
+    metadata_path = os.path.join(output_dir, metadata_filename)
+    with open(metadata_path, 'w') as f:
+        json.dump(dataset_metadata, f, indent=2)
+    print(f"✅ Saved FINAL metadata: {metadata_path}")
+    
+    # Summary
+    print(f"\n{'='*80}")
+    print(f"PROCESSING COMPLETE")
+    print(f"{'='*80}")
+    print(f"Successfully processed: {len(all_embeddings)}/{len(dicom_dirs)} scans")
+    print(f"Failed: {len(failed)}/{len(dicom_dirs)} scans")
+    print(f"\nEmbeddings array shape: {embeddings_array.shape}")
+    print(f"Saved embeddings to: {embeddings_path}")
+    print(f"Saved metadata to: {metadata_path}")
+    
+    # Timepoint summary
+    timepoint_counts = {}
+    for m in all_metadata:
+        tp = m.get('timepoint', 'Unknown')
+        timepoint_counts[tp] = timepoint_counts.get(tp, 0) + 1
+    
+    if timepoint_counts:
+        print(f"\n📅 Timepoint Distribution:")
+        for tp in sorted(timepoint_counts.keys(), key=lambda x: (x is None, x)):
+            count = timepoint_counts[tp]
+            if tp is None:
+                print(f"  Unknown/Not detected: {count} scans")
+            else:
+                print(f"  {tp}: {count} scans")
+    
+    if failed:
+        print(f"\nFailed scans: {len(failed)}")
+        for fail_info in failed[:5]:  # Show first 5 failures
+            print(f"  - {fail_info['scan_dir']}")
+            print(f"    Error: {fail_info['error']}")
+        if len(failed) > 5:
+            print(f"  ... and {len(failed) - 5} more failures")
+    
+    print(f"\n{'='*80}")
+    print(f"For downstream training, load embeddings with:")
+    print(f"  import pandas as pd")
+    print(f"  import numpy as np")
+    print(f"  df = pd.read_parquet('{embeddings_path}')")
+    print(f"  # Total rows: {len(df)}, Total columns: {len(df.columns)}")
+    print(f"  # Extract embeddings array: embeddings = np.stack(df['embedding'].values)")
+    print(f"  # Shape: {embeddings_array.shape}")
+    print(f"  # Access individual: df.loc[0, 'embedding'] -> array of shape ({embedding_dim},)")
+    print(f"{'='*80}")