code/scripts/inference.py

import os
os.environ['PYTORCH_CUDA_ALLOC_CONF'] = 'expandable_segments:True'
os.environ['NCCL_TIMEOUT'] = '120'
import torch
import torch.nn as nn
from dataset.hpadataset import HPADatasetMIL, HPADatasetDownsample, HPADatasetMIL_url, StanfordTMAInferenceMILDataset
from torch.utils.data import DataLoader
import pandas as pd
from tqdm import tqdm
import numpy as np
import time
import configparser
import argparse
import copy
from datetime import datetime
from os.path import join
import matplotlib.pyplot as plt
import seaborn as sns
from sklearn.metrics import confusion_matrix, accuracy_score, balanced_accuracy_score, f1_score
import torch.cuda.amp  # Import Automatic Mixed Precision
import re
import random
import json

"""
Be sure to init conda env:
conda activate hpa

Note: If running on a non-H100 GPU, there might be an error for pyarrow.
Reinstall pyarrow with following command will resolve the issue:
conda install -c conda-forge pyarrow
"""

# Define loss functions for evaluation
criterion_intensity = nn.CrossEntropyLoss()  # For intensity
criterion_location = nn.CrossEntropyLoss()   # For location
criterion_quantity = nn.CrossEntropyLoss()   # For quantity
criterion_tissue = nn.CrossEntropyLoss()     # For tissue type (58 classes)
criterion_malignancy = nn.CrossEntropyLoss() # For tumor vs non-tumor (binary)

def set_seed(seed):
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)

def custom_collate_fn(batch):
    # If no valid samples, return None
    batch_valid = [batch_item for batch_item in batch if batch_item is not None]
    
    if len(batch_valid) == 0:
        return None
    images, metadata, query_input, caption_output, snomed_text, snomed_code, image_url, staining_intensity, staining_location, staining_quantity, malignancy, tissue_one_hot, cell_type_one_hot = zip(*batch_valid)
    # Note: here images is a list of (N, 3, 224, 224), where N is the number of patches per image, and is different for each image.
    
    # Convert to tensors and stack
    staining_intensity = torch.stack([torch.tensor(sublist, dtype=torch.float) for sublist in staining_intensity])
    staining_location = torch.stack([torch.tensor(sublist, dtype=torch.float) for sublist in staining_location])
    staining_quantity = torch.stack([torch.tensor(sublist, dtype=torch.float) for sublist in staining_quantity])
    malignancy = torch.stack([torch.tensor(sublist, dtype=torch.float) for sublist in malignancy])
    tissue_one_hot = torch.stack([torch.tensor(sublist, dtype=torch.float) for sublist in tissue_one_hot])
    cell_type_one_hot = torch.stack([torch.tensor(sublist, dtype=torch.float) for sublist in cell_type_one_hot])
    
    return images, metadata, query_input, caption_output, snomed_text, snomed_code, image_url, staining_intensity, staining_location, staining_quantity, malignancy, tissue_one_hot, cell_type_one_hot


def run_inference(config):
    # Setup device
    device = "cuda" if torch.cuda.is_available() else "cpu"
    print(f"Using device: {device}")

    # Initialize model based on configuration
    base_model_name = config.base_model_name

    from model.patch_encoder_with_clam import create_clam_vit
    model, _, _ = create_clam_vit(
        base_model_name=base_model_name,
        training_config=config,
        device=device,
        lr=config.learning_rate,
        weight_decay=config.weight_decay,
        freeze_vit=config.freeze_vit,
        freeze_query_features_encoder=config.freeze_query_features_encoder,
        use_cell_type_embedding=config.use_cell_type_embedding
    )

    model.to(device)

    # Determine patch size based on model
    if base_model_name == "vinid/plip":
        patch_size = 224
    elif base_model_name == "openai/clip-vit-large-patch14-336":
        patch_size = 336
    else:
        raise ValueError(f"Model {base_model_name} not supported.")

    processor = model.patch_processor

    # Load checkpoint
    if config.checkpoint_path:
        print(f"Loading checkpoint from {config.checkpoint_path}")
        checkpoint = torch.load(config.checkpoint_path, map_location=device)
        model.load_state_dict(checkpoint['model_state_dict'])
        print("Model loaded successfully")
    else:
        raise ValueError("No checkpoint path provided. Please specify a checkpoint path.")

    # Load inference data
    print(f"Loading inference data from {config.inference_data_path}")
    if config.inference_data_path.endswith('.csv'):
        inference_data = pd.read_csv(config.inference_data_path)
    elif config.inference_data_path.endswith('.feather'):
        inference_data = pd.read_feather(config.inference_data_path)
    else:
        raise ValueError("Inference data must be a CSV or Feather file")
    
    print(f"Loaded {len(inference_data)} samples for inference")


    # Setup dataset and dataloader
    hdf5_base_dir = config.hdf5_base_dir
    rle_map_path = config.rle_map_path

    if config.model_version == "simple_downsample":
        dataset = HPADatasetDownsample(inference_data, None, data_split="test", target_size=patch_size, processor=processor)
    else:
        dataset = HPADatasetMIL_url(inference_data, rle_map_path, hdf5_base_dir, data_split="test", patch_size=patch_size, processor=processor)

    # Add config to dataset for accessing save_logits parameter
    dataset.config = config

    dataloader = DataLoader(
        dataset,
        batch_size=config.batch_size,
        shuffle=False,
        num_workers=config.num_workers,
        collate_fn=custom_collate_fn,
        pin_memory=False,
        prefetch_factor=8,
        persistent_workers=True
    )

    # Run inference
    results, metrics, metrics_df = run_model_inference(model, dataloader, device, config.output_dir)
    
    # Save results
    save_name = config.save_name if config.save_name else datetime.now().strftime("%Y%m%d_%H%M%S")
    output_file = os.path.join(config.output_dir, f"inference_results_{save_name}.csv")
    results.to_csv(output_file, index=False)
    print(f"Inference results saved to {output_file}")
    
    # Save metrics if available
    if metrics:
        with open(os.path.join(config.output_dir, f'metrics_{save_name}.json'), 'w') as f:
            json.dump(metrics, f, indent=4)
        
        metrics_df.to_csv(os.path.join(config.output_dir, f'metrics_{save_name}.csv'), index=False)
        print(f"Metrics saved to {config.output_dir}/metrics_{save_name}.csv")

    # Generate visualizations if requested
    if config.generate_visualizations:
        generate_visualizations(results, config.output_dir, save_name)
        print(f"Visualizations saved to {config.output_dir}")


def run_model_inference(model, dataloader, device, output_dir):
    model.eval()
    
    # Lists to store all outputs and metadata
    all_outputs = {
        'intensity': [],
        'location': [],
        'quantity': [],
        'tissue': [],
        'malignancy': []
    }
    
    all_predicted_classes = {
        'intensity': [],
        'location': [],
        'quantity': [],
        'tissue': [],
        'malignancy': []
    }
    
    all_ground_truth = {
        'intensity': [],
        'location': [],
        'quantity': [],
        'tissue': [],
        'malignancy': []
    }
    
    all_image_urls = []
    all_names = []
    all_metadata = []
    all_snomed_text = []
    all_snomed_code = []
    
    # Set deterministic behavior for reproducibility
    torch.backends.cudnn.deterministic = True
    torch.backends.cudnn.benchmark = False
    
    with torch.no_grad():
        for batch in tqdm(dataloader, desc="Running inference"):
            if batch is None:
                continue
                
            patches, metadata, query_input, caption_output, snomed_text, snomed_code, image_url, staining_intensity, staining_location, staining_quantity, malignancy, tissue_one_hot, cell_type_one_hot = batch
            
            # Extract names from metadata
            names = [meta.get('name', 'unknown') for meta in metadata]
            
            # Move data to device
            patches_device = [patch.to(device) for patch in patches]
            staining_intensity = staining_intensity.to(device)
            staining_location = staining_location.to(device)
            staining_quantity = staining_quantity.to(device)
            tissue_one_hot = tissue_one_hot.to(device)
            malignancy = malignancy.to(device)
            
            # Process each sample in the batch individually to ensure consistency
            for i in range(len(names)):
                try:
                    # Debug the shape
                    patch = patches_device[i]
                    
                    # Ensure patch has the right shape (N, 3, H, W)
                    if len(patch.shape) != 4:
                        print(f"Warning: Patch shape is {patch.shape}, expected 4D tensor")
                        continue
                    
                    # Mixed precision inference
                    with torch.cuda.amp.autocast():
                        intensity_out, location_out, quantity_out, tissue_out, malignancy_out, A_raw = model(
                            [patch], [query_input[i]], 
                            cell_type_one_hot[i:i+1].to(device) if cell_type_one_hot is not None else None, 
                            phase="test"
                        )
                    
                    # Store outputs
                    all_outputs['intensity'].append(intensity_out.cpu().numpy()[0])
                    all_outputs['location'].append(location_out.cpu().numpy()[0])
                    all_outputs['quantity'].append(quantity_out.cpu().numpy()[0])
                    all_outputs['tissue'].append(tissue_out.cpu().numpy()[0])
                    all_outputs['malignancy'].append(malignancy_out.cpu().numpy()[0])
                    
                    # Store predicted classes
                    all_predicted_classes['intensity'].append(torch.argmax(intensity_out, dim=1).cpu().numpy()[0])
                    all_predicted_classes['location'].append(torch.argmax(location_out, dim=1).cpu().numpy()[0])
                    all_predicted_classes['quantity'].append(torch.argmax(quantity_out, dim=1).cpu().numpy()[0])
                    all_predicted_classes['tissue'].append(torch.argmax(tissue_out, dim=1).cpu().numpy()[0])
                    all_predicted_classes['malignancy'].append(torch.argmax(malignancy_out, dim=1).cpu().numpy()[0])
                    
                    # Store ground truth if available
                    all_ground_truth['intensity'].append(staining_intensity[i].cpu().numpy())
                    all_ground_truth['location'].append(staining_location[i].cpu().numpy())
                    all_ground_truth['quantity'].append(staining_quantity[i].cpu().numpy())
                    all_ground_truth['tissue'].append(tissue_one_hot[i].cpu().numpy())
                    all_ground_truth['malignancy'].append(malignancy[i].cpu().numpy())
                    
                    # Store metadata
                    all_image_urls.append(image_url[i])
                    all_names.append(names[i])
                    all_metadata.append(metadata[i])
                    all_snomed_text.append(snomed_text[i])
                    all_snomed_code.append(snomed_code[i])
                except Exception as e:
                    print(f"Error processing sample {names[i]}: {e}")
                    print(f"Patch shape: {patches[i].shape}")
                    continue
    
    # Create results dataframe with metadata
    results_df = pd.DataFrame({
        'name': all_names,
        'image_url': all_image_urls,
    })
    
    # Add all metadata fields
    for i, meta in enumerate(all_metadata):
        for key, value in meta.items():
            if key not in results_df.columns:
                results_df[key] = [None] * len(all_metadata)
            results_df.at[i, key] = value
    
    # Add SNOMED information
    results_df['snomed_text'] = all_snomed_text
    results_df['snomed_code'] = all_snomed_code
    
    # Map numeric predictions to human-readable labels
    intensity_labels = ['negative', 'weak', 'moderate', 'strong']
    location_labels = ['none', 'cytoplasmic/membranous', 'nuclear', 'cytoplasmic/membranous,nuclear']
    quantity_labels = ['none', '<25%', '25%-75%', '>75%']
    malignancy_labels = ['normal', 'cancer']
    
    # Load tissue labels
    tissue_labels = ['adipose', 'adrenal gland', 'appendix', 'bone marrow', 'breast',
                    'bronchus', 'carcinoid', 'caudate', 'cerebellum',
                    'cerebral cortex', 'cervical', 'cervix', 'colon', 'colorectal',
                    'duodenum', 'endometrial', 'endometrium', 'epididymis',
                    'esophagus', 'fallopian tube', 'gallbladder', 'glioma',
                    'head and neck', 'heart muscle', 'hippocampus', 'kidney', 'liver',
                    'lung', 'lymph node', 'lymphoma', 'melanoma', 'nasopharynx',
                    'oral mucosa', 'ovarian', 'ovary', 'pancreas', 'pancreatic',
                    'parathyroid gland', 'placenta', 'prostate', 'rectum', 'renal',
                    'salivary gland', 'seminal vesicle', 'skeletal muscle', 'skin',
                    'small intestine', 'smooth muscle', 'soft', 'spleen', 'stomach',
                    'testis', 'thyroid', 'thyroid gland', 'tonsil', 'urinary bladder',
                    'urothelial', 'vagina']
    
    # Add human-readable labels only (no class indices)
    results_df['intensity_pred'] = [intensity_labels[i] for i in all_predicted_classes['intensity']]
    results_df['location_pred'] = [location_labels[i] for i in all_predicted_classes['location']]
    results_df['quantity_pred'] = [quantity_labels[i] for i in all_predicted_classes['quantity']]
    results_df['tissue_pred'] = [tissue_labels[i] if i < len(tissue_labels) else f"unknown_{i}" for i in all_predicted_classes['tissue']]
    results_df['malignancy_pred'] = [malignancy_labels[i] for i in all_predicted_classes['malignancy']]
    
    # Add raw logits if requested
    if hasattr(dataloader.dataset, 'config') and dataloader.dataset.config.save_logits:
        for task in ['intensity', 'location', 'quantity', 'tissue', 'malignancy']:
            for i, logits in enumerate(all_outputs[task]):
                for j, logit in enumerate(logits):
                    results_df.at[i, f'{task}_logit_{j}'] = logit
    
    # Add ground truth labels if available (no class indices)
    if all(len(gt) > 0 for gt in all_ground_truth.values()):
        # Convert one-hot encoded ground truth to class indices
        for task in ['intensity', 'location', 'quantity', 'tissue', 'malignancy']:
            gt_indices = [np.argmax(gt) for gt in all_ground_truth[task]]
            
            # Add human-readable ground truth labels only
            if task == 'intensity':
                results_df[f'{task}_true'] = [intensity_labels[i] for i in gt_indices]
            elif task == 'location':
                results_df[f'{task}_true'] = [location_labels[i] for i in gt_indices]
            elif task == 'quantity':
                results_df[f'{task}_true'] = [quantity_labels[i] for i in gt_indices]
            elif task == 'tissue':
                results_df[f'{task}_true'] = [tissue_labels[i] if i < len(tissue_labels) else f"unknown_{i}" for i in gt_indices]
            elif task == 'malignancy':
                results_df[f'{task}_true'] = [malignancy_labels[i] for i in gt_indices]
    
    # Calculate and print metrics if ground truth is available
    metrics = {}
    metrics_df = pd.DataFrame()
    
    if all(len(gt) > 0 for gt in all_ground_truth.values()):
        print("\nPerformance Metrics:")
        print("=" * 50)
        
        for task in ['intensity', 'location', 'quantity', 'tissue', 'malignancy']:
            gt_indices = [np.argmax(gt) for gt in all_ground_truth[task]]
            y_true = np.array(gt_indices)
            y_pred = np.array(all_predicted_classes[task])
            
            # Get unique classes in both true and predicted
            unique_classes = np.unique(np.concatenate([y_true, y_pred]))
            
            # Calculate metrics
            try:
                accuracy = accuracy_score(y_true, y_pred)
                balanced_acc = balanced_accuracy_score(y_true, y_pred)
                f1 = f1_score(y_true, y_pred, average='weighted', labels=np.unique(y_true), zero_division=0)
                
                # Store metrics
                metrics[task] = {
                    'accuracy': float(accuracy),  # Convert to float for JSON serialization
                    'balanced_accuracy': float(balanced_acc),
                    'f1_score': float(f1),
                    'num_classes_true': int(len(np.unique(y_true))),
                    'num_classes_pred': int(len(np.unique(y_pred))),
                    'num_classes_total': int(len(unique_classes))
                }
                
                # Add to metrics dataframe
                metrics_df = pd.concat([metrics_df, pd.DataFrame({
                    'task': [task],
                    'accuracy': [accuracy],
                    'balanced_accuracy': [balanced_acc],
                    'f1_score': [f1],
                    'num_classes_true': [len(np.unique(y_true))],
                    'num_classes_pred': [len(np.unique(y_pred))],
                    'num_classes_total': [len(unique_classes)]
                })], ignore_index=True)
                
                # Print metrics
                print(f"\n{task.capitalize()} Metrics:")
                print(f"  Accuracy: {accuracy:.4f}")
                print(f"  Balanced Accuracy: {balanced_acc:.4f}")
                print(f"  Weighted F1 Score: {f1:.4f}")
                print(f"  Classes in ground truth: {len(np.unique(y_true))}")
                print(f"  Classes in predictions: {len(np.unique(y_pred))}")
                print(f"  Total unique classes: {len(unique_classes)}")
                
                # Print warning if there are classes in predictions not in ground truth
                if len(np.unique(y_pred)) > len(np.unique(y_true)):
                    extra_classes = set(y_pred) - set(y_true)
                    print(f"  Warning: Predictions contain {len(extra_classes)} classes not in ground truth")
                
            except Exception as e:
                print(f"Error calculating metrics for {task}: {e}")
                metrics[task] = {
                    'error': str(e)
                }
    
    return results_df, metrics, metrics_df


def generate_visualizations(results_df, output_dir, save_name):
    """Generate visualizations from inference results"""
    os.makedirs(os.path.join(output_dir, 'visualizations', save_name), exist_ok=True)
    
    # Extract ground truth and predictions
    tasks = ['intensity', 'location', 'quantity', 'tissue', 'malignancy']
    
    # Check if ground truth is available
    has_ground_truth = all(f'{task}_true' in results_df.columns for task in tasks)
    
    if has_ground_truth:
        # Generate confusion matrices
        for task in tasks:
            # Define ordered class names for each task
            if task == 'intensity':
                class_names = ['negative', 'weak', 'moderate', 'strong']
            elif task == 'location':
                class_names = ['none', 'nuclear', 'cytoplasmic/membranous', 'cytoplasmic/membranous,nuclear']
            elif task == 'quantity':
                class_names = ['none', '<25%', '25%-75%', '>75%']
            elif task == 'malignancy':
                class_names = ['normal', 'cancer']
            elif task == 'tissue':
                # For tissue, handle separately due to large number of classes
                tissue_labels = ['adipose', 'adrenal gland', 'appendix', 'bone marrow', 'breast',
                                'bronchus', 'carcinoid', 'caudate', 'cerebellum',
                                'cerebral cortex', 'cervical', 'cervix', 'colon', 'colorectal',
                                'duodenum', 'endometrial', 'endometrium', 'epididymis',
                                'esophagus', 'fallopian tube', 'gallbladder', 'glioma',
                                'head and neck', 'heart muscle', 'hippocampus', 'kidney', 'liver',
                                'lung', 'lymph node', 'lymphoma', 'melanoma', 'nasopharynx',
                                'oral mucosa', 'ovarian', 'ovary', 'pancreas', 'pancreatic',
                                'parathyroid gland', 'placenta', 'prostate', 'rectum', 'renal',
                                'salivary gland', 'seminal vesicle', 'skeletal muscle', 'skin',
                                'small intestine', 'smooth muscle', 'soft', 'spleen', 'stomach',
                                'testis', 'thyroid', 'thyroid gland', 'tonsil', 'urinary bladder',
                                'urothelial', 'vagina']
                
                # For tissue, we'll create a confusion matrix for the top N most common classes
                # to keep the visualization manageable
                top_n = 15
                
                # Get unique tissue types in the data
                unique_tissues_true = results_df[f'{task}_true'].unique()
                unique_tissues_pred = results_df[f'{task}_pred'].unique()
                unique_tissues = np.unique(np.concatenate([unique_tissues_true, unique_tissues_pred]))
                
                # Count occurrences of each tissue type
                tissue_counts = {}
                for tissue in unique_tissues:
                    count_true = np.sum(results_df[f'{task}_true'] == tissue)
                    count_pred = np.sum(results_df[f'{task}_pred'] == tissue)
                    tissue_counts[tissue] = count_true + count_pred
                
                # Get top N tissues by occurrence
                top_tissues = sorted(tissue_counts.keys(), key=lambda x: tissue_counts[x], reverse=True)[:top_n]
                
                # Filter data to only include top tissues
                mask = (results_df[f'{task}_true'].isin(top_tissues)) & (results_df[f'{task}_pred'].isin(top_tissues))
                filtered_df = results_df[mask]
                
                # Create confusion matrix
                cm = confusion_matrix(
                    filtered_df[f'{task}_true'], 
                    filtered_df[f'{task}_pred'],
                    labels=top_tissues
                )
                
                plt.figure(figsize=(15, 15))
                sns.heatmap(cm, annot=True, fmt='d', cmap='Blues', 
                            xticklabels=top_tissues, yticklabels=top_tissues)
                plt.title(f'Confusion Matrix - Top {top_n} {task.capitalize()} Classes')
                plt.ylabel('True label')
                plt.xlabel('Predicted label')
                plt.xticks(rotation=90, ha='right')
                plt.yticks(rotation=0)
                plt.tight_layout()
                plt.savefig(os.path.join(output_dir, 'visualizations', save_name, f'{task}_confusion_matrix_{save_name}.png'))
                plt.close()
                continue
            
            # Create mapping from label to index for consistent ordering
            label_to_idx = {label: i for i, label in enumerate(class_names)}
            
            # Convert string labels to indices based on our defined order
            y_true = np.array([label_to_idx[label] for label in results_df[f'{task}_true']])
            y_pred = np.array([label_to_idx[label] for label in results_df[f'{task}_pred']])
            
            # Create confusion matrix with ordered classes
            cm = confusion_matrix(y_true, y_pred, labels=range(len(class_names)))
            
            # Normalize confusion matrix for better visualization
            with np.errstate(divide='ignore', invalid='ignore'):
                cm_normalized = cm.astype('float') / cm.sum(axis=1)[:, np.newaxis]
                cm_normalized = np.nan_to_num(cm_normalized, nan=0)
            
            # Plot raw counts
            plt.figure(figsize=(10, 8))
            sns.heatmap(cm, annot=True, fmt='d', cmap='Blues', 
                        xticklabels=class_names, yticklabels=class_names)
            plt.title(f'Confusion Matrix (Counts) - {task.capitalize()}')
            plt.ylabel('True label')
            plt.xlabel('Predicted label')
            plt.tight_layout()
            plt.savefig(os.path.join(output_dir, 'visualizations', save_name, f'{task}_confusion_matrix_counts_{save_name}.png'))
            plt.close()
            
            # Plot normalized (percentage)
            plt.figure(figsize=(10, 8))
            sns.heatmap(cm_normalized, annot=True, fmt='.2f', cmap='Blues', 
                        xticklabels=class_names, yticklabels=class_names)
            plt.title(f'Confusion Matrix (Normalized) - {task.capitalize()}')
            plt.ylabel('True label')
            plt.xlabel('Predicted label')
            plt.tight_layout()
            plt.savefig(os.path.join(output_dir, 'visualizations', save_name, f'{task}_confusion_matrix_normalized_{save_name}.png'))
            plt.close()


def parse_config():
    parser = argparse.ArgumentParser(description='Run inference with HPA-VLM model')
    parser.add_argument('--config', type=str, required=True, help='Path to the configuration file')
    parser.add_argument('--inference_data_path', type=str, required=True, help='Path to the inference data CSV')
    parser.add_argument('--checkpoint_path', type=str, required=True, help='Path to the model checkpoint')
    parser.add_argument('--output_dir', type=str, default='inference_results', help='Directory to save inference results')
    parser.add_argument('--hdf5_base_dir', type=str, required=True, help='Base directory for HDF5 files')
    parser.add_argument('--rle_map_path', type=str, required=True, help='Path to RLE map JSON file')
    parser.add_argument('--batch_size', type=int, default=4, help='Batch size for inference')
    parser.add_argument('--num_workers', type=int, default=4, help='Number of workers for data loading')
    parser.add_argument('--generate_visualizations', action='store_true', help='Generate visualizations of results')
    parser.add_argument('--save_name', type=str, default='', help='Suffix for saved files (instead of timestamp)')
    parser.add_argument('--save_logits', action='store_true', help='Save raw logits in the output')
    
    args = parser.parse_args()
    
    # Load config file
    config = configparser.ConfigParser()
    config.read(args.config)
    
    # Convert config values to desired data types
    def convert_value(convert_func, key, default=None):
        try:
            return convert_func(config['DEFAULT'][key])
        except (KeyError, ValueError):
            if default is not None:
                return default
            raise ValueError(f"Error converting {key}. Please check your config file.")
    
    # Create a dictionary to hold the converted values
    converted_config = {
        'base_model_name': config['DEFAULT']['base_model_name'],
        'freeze_vit': convert_value(lambda x: x.lower() == 'true', 'freeze_vit', True),
        'freeze_query_features_encoder': convert_value(lambda x: x.lower() == 'true', 'freeze_query_features_encoder', True),
        'learning_rate': convert_value(float, 'learning_rate', 1e-4),
        'weight_decay': convert_value(float, 'weight_decay', 1e-5),
        'use_cell_type_embedding': convert_value(lambda x: x.lower() == 'true', 'use_cell_type_embedding', False),
        'model_version': config['DEFAULT']['model_version'],
        'inference_data_path': args.inference_data_path,
        'checkpoint_path': args.checkpoint_path,
        'output_dir': args.output_dir,
        'hdf5_base_dir': args.hdf5_base_dir,
        'rle_map_path': args.rle_map_path,
        'batch_size': args.batch_size,
        'num_workers': args.num_workers,
        'generate_visualizations': args.generate_visualizations,
        'save_name': args.save_name,
        'save_logits': args.save_logits,
    }
    
    return argparse.Namespace(**converted_config)


if __name__ == "__main__":
    config = parse_config()
    set_seed(1001)  # Set seed for reproducibility
    
    # Create output directory
    os.makedirs(config.output_dir, exist_ok=True)
    
    # Save configuration
    config_dict = vars(config)
    with open(os.path.join(config.output_dir, 'inference_config.json'), 'w') as f:
        json.dump(config_dict, f, indent=4)
    
    # Run inference
    run_inference(config)