inference.py

"""
PathLUPI Inference Script
========================
This script performs inference on WSI features using a trained PathLUPI model.
Only WSI features are required at inference time - no transcriptomic data needed.

Usage:
    python inference.py \
        --checkpoint /path/to/model.pth.tar \
        --wsi_features /path/to/features.pt \
        --task BRCA_ERSub \
        --signatures labels/signatures/hallmarks_signatures.csv
"""

import argparse
import torch
import torch.nn.functional as F
import pandas as pd

from models.PathLUPI.network import PathLUPI


def parse_args():
    parser = argparse.ArgumentParser(description="PathLUPI Inference")
    parser.add_argument("--checkpoint", type=str, required=True,
                        help="Path to trained model checkpoint (.pth.tar)")
    parser.add_argument("--wsi_features", type=str, required=True,
                        help="Path to WSI features (.pt file)")
    parser.add_argument("--task", type=str, required=True,
                        choices=["survival", "BRCA_ERSub", "BRCA_PRSub", "BRCA_HER2Sub",
                                 "BRCA_TNBCSub", "BRCA_MolSub", "BRCA_PIK3CASub", "BRCA_TP53Sub",
                                 "GBMLGG_IDH1Sub", "GBMLGG_MolSub", "NSCLC_TMBSub", "CRC_CMSSub",
                                 "BLCA_MolSub", "BLCA_FGFR3Sub", "CRC_BRAFSub", "CRC_TP53Sub",
                                 "CRC_KRASSub", "HNSC_MolSub", "LIHC_TP53Sub",
                                 "LUAD_TP53Sub", "LUAD_KRASSub", "LUAD_EGFRSub", "KIRC_MolSub",
                                 "SKCM_BRAFSub", "SKCM_KRASSub", "PanGI_GISub", "UCEC_MolSub"],
                        help="Prediction task type")
    parser.add_argument("--signatures", type=str, required=True,
                        help="Path to pathway signatures CSV file")
    parser.add_argument("--region_num", type=int, default=50,
                        help="Number of pathway regions (default: 50)")
    parser.add_argument("--output", type=str, default=None,
                        help="Output file path for predictions (optional)")
    parser.add_argument("--device", type=str, default="cuda",
                        help="Device to run inference on (cuda/cpu)")
    return parser.parse_args()


def get_task_config(task):
    """Get number of classes and class labels for each task.

    Labels are ordered by their numeric class index (0, 1, 2, ...).
    These mappings are derived from datasets/TCGA_Subtype.py process_label().
    """
    task_configs = {
        # Survival: 4 discrete time bins (0-3), outputs risk score
        "survival": {"n_classes": 4, "labels": ["Bin0", "Bin1", "Bin2", "Bin3"]},

        # BRCA tasks
        "BRCA_ERSub": {"n_classes": 2, "labels": ["Negative", "Positive"]},  # {0: 0, 1: 1}
        "BRCA_PRSub": {"n_classes": 2, "labels": ["Negative", "Positive"]},  # {0: 0, 1: 1}
        "BRCA_HER2Sub": {"n_classes": 2, "labels": ["Negative", "Positive"]},  # {"negative": 0, "positive": 1}
        "BRCA_TNBCSub": {"n_classes": 2, "labels": ["Non-TNBC", "TNBC"]},  # {"NO": 0, "YES": 1}
        "BRCA_MolSub": {"n_classes": 4, "labels": ["LumA", "LumB", "Basal", "Her2"]},  # {"LumA": 0, "LumB": 1, "Basal": 2, "Her2": 3}
        "BRCA_PIK3CASub": {"n_classes": 2, "labels": ["Wild-type", "Mutant"]},  # {0: 0, 1: 1}
        "BRCA_TP53Sub": {"n_classes": 2, "labels": ["Wild-type", "Mutant"]},  # {0: 0, 1: 1}

        # GBMLGG tasks
        "GBMLGG_IDH1Sub": {"n_classes": 2, "labels": ["Wild-type", "Mutant"]},  # {0: 0, 1: 1}
        "GBMLGG_MolSub": {"n_classes": 5, "labels": ["G-CIMP-high", "Codel", "Mesenchymal-like", "Classic-like", "Other"]},
        # {'G-CIMP-high': 0, 'Codel': 1, 'Mesenchymal-like': 2, 'Classic-like': 3, 'PA-like/LGm6-GBM/G-CIMP-low': 4}

        # NSCLC task
        "NSCLC_TMBSub": {"n_classes": 2, "labels": ["TMB-Low", "TMB-High"]},  # {"TMB_L": 0, "TMB_H": 1}

        # CRC tasks
        "CRC_CMSSub": {"n_classes": 4, "labels": ["CMS1", "CMS2", "CMS3", "CMS4"]},
        "CRC_BRAFSub": {"n_classes": 2, "labels": ["Wild-type", "Mutant"]},
        "CRC_KRASSub": {"n_classes": 2, "labels": ["Wild-type", "Mutant"]},
        "CRC_TP53Sub": {"n_classes": 2, "labels": ["Wild-type", "Mutant"]},

        # BLCA tasks
        "BLCA_MolSub": {"n_classes": 4, "labels": ["Luminal", "Luminal_infiltrated", "Luminal_papillary", "Basal_squamous"]},
        "BLCA_FGFR3Sub": {"n_classes": 2, "labels": ["Wild-type", "Mutant"]},  # {"wt": 0, "mut": 1}

        # LUAD tasks
        "LUAD_EGFRSub": {"n_classes": 2, "labels": ["Wild-type", "Mutant"]},
        "LUAD_KRASSub": {"n_classes": 2, "labels": ["Wild-type", "Mutant"]},
        "LUAD_TP53Sub": {"n_classes": 2, "labels": ["Wild-type", "Mutant"]},

        # Other cancer tasks
        "LIHC_TP53Sub": {"n_classes": 2, "labels": ["Wild-type", "Mutant"]},
        "SKCM_BRAFSub": {"n_classes": 2, "labels": ["Wild-type", "Mutant"]},
        "SKCM_KRASSub": {"n_classes": 2, "labels": ["Wild-type", "Mutant"]},
        "HNSC_MolSub": {"n_classes": 4, "labels": ["Classical", "Basal", "Mesenchymal", "Atypical"]},
        # {'HNSC.Classical': 0, 'HNSC.Basal': 1, 'HNSC.Mesenchymal': 2, 'HNSC.Atypical': 3}

        # Multi-cancer tasks
        "PanGI_GISub": {"n_classes": 5, "labels": ["MSI", "CIN", "EBV", "HM-SNV", "GS"]},
        # {'GI.MSI': 0, 'GI.CIN': 1, 'GI.EBV': 2, 'GI.HM-SNV': 3, 'GI.GS': 4}
        "KIRC_MolSub": {"n_classes": 4, "labels": ["KIRC.1", "KIRC.2", "KIRC.3", "KIRC.4"]},
        "UCEC_MolSub": {"n_classes": 4, "labels": ["CN_HIGH", "CN_LOW", "MSI", "POLE"]},
        # {"UCEC.CN_HIGH": 0, "UCEC.CN_LOW": 1, "UCEC.MSI": 2, "UCEC.POLE": 3}
    }
    return task_configs.get(task, {"n_classes": 2, "labels": ["Class0", "Class1"]})


def load_omic_sizes(signatures_path):
    """Load pathway signature sizes from CSV file."""
    csv = pd.read_csv(signatures_path)
    omic_sizes = []
    for col in csv.columns:
        omic = csv[col].dropna().unique()
        omic_sizes.append(len(omic))
    return omic_sizes


def load_model(checkpoint_path, task, omic_sizes, region_num, device):
    """Load trained PathLUPI model from checkpoint."""
    checkpoint = torch.load(checkpoint_path, map_location=device)
    
    task_config = get_task_config(task)
    n_classes = task_config["n_classes"]
    survival = (task == "survival")
    
    # Infer path_size from checkpoint
    state_dict = checkpoint["state_dict"]
    path_size = state_dict["wsi_net.0.weight"].shape[1]
    
    model = PathLUPI(
        omic_sizes=omic_sizes,
        n_classes=n_classes,
        path_size=path_size,
        path_proj_size=path_size // 2,
        num_tokens=region_num,
        region_num=region_num,
        survival=survival
    )
    
    model.load_state_dict(state_dict)
    model = model.to(device)
    model.eval()
    
    print(f"[Model] Loaded checkpoint: {checkpoint_path}")
    print(f"[Model] Task: {task}, Classes: {n_classes}, Path size: {path_size}")
    
    return model, task_config


def run_inference(model, wsi_features, task, task_config, device):
    """Run inference on WSI features."""
    wsi_features = wsi_features.to(device)
    
    if wsi_features.dim() == 2:
        wsi_features = wsi_features.unsqueeze(0)
    
    with torch.no_grad():
        if task == "survival":
            _, S, _, _, _ = model(x_path=wsi_features)
            risk_score = -torch.sum(S, dim=1).cpu().numpy()[0]
            return {"risk_score": float(risk_score), "survival_probability": S.cpu().numpy()[0].tolist()}
        else:
            logits, _, _, _, _ = model(x_path=wsi_features)
            probs = F.softmax(logits, dim=1)
            pred_class = torch.argmax(probs, dim=1).item()
            pred_label = task_config["labels"][pred_class]
            return {
                "predicted_class": pred_class,
                "predicted_label": pred_label,
                "probabilities": {label: float(probs[0, i]) for i, label in enumerate(task_config["labels"])}
            }


def main():
    args = parse_args()
    device = torch.device(args.device if torch.cuda.is_available() else "cpu")
    print(f"[Device] Using: {device}")
    
    # Load pathway signature sizes
    omic_sizes = load_omic_sizes(args.signatures)
    print(f"[Signatures] Loaded {len(omic_sizes)} pathways")
    
    # Load model
    model, task_config = load_model(args.checkpoint, args.task, omic_sizes, args.region_num, device)
    
    # Load WSI features
    wsi_features = torch.load(args.wsi_features, map_location=device)
    print(f"[WSI] Loaded features: {wsi_features.shape}")
    
    # Run inference
    result = run_inference(model, wsi_features, args.task, task_config, device)
    
    # Print results
    print("\n" + "="*50)
    print("PREDICTION RESULTS")
    print("="*50)
    if args.task == "survival":
        print(f"Risk Score: {result['risk_score']:.4f}")
    else:
        print(f"Predicted: {result['predicted_label']} (class {result['predicted_class']})")
        print("\nClass Probabilities:")
        for label, prob in result['probabilities'].items():
            print(f"  {label}: {prob:.4f}")
    
    # Save results if output path specified
    if args.output:
        import json
        with open(args.output, 'w') as f:
            json.dump(result, f, indent=2)
        print(f"\n[Output] Results saved to: {args.output}")
    
    return result


if __name__ == "__main__":
    main()