inference.py

import os
import torch
import numpy as np
from PIL import Image
from gemmasight.config import MODELS_DIR, DIM_FUSED, LABEL_MAP
from gemmasight.models.feature_extractor import DualEncoderFeatureExtractor
from gemmasight.models.classifier import MSIClassifier
from gemmasight.models.retriever import CaseRetriever
from gemmasight.models.report_generator import MedGemmaReportGenerator
from gemmasight.utils.occlusion import generate_occlusion_heatmap

class GemmaSightInferencePipeline:
    def __init__(self, force_simulation=None):
        # Allow override of force_simulation
        if force_simulation is not None:
            self.force_simulation = force_simulation
        else:
            from gemmasight.config import FORCE_SIMULATION
            self.force_simulation = FORCE_SIMULATION
            
        print("=== Initializing GemmaSight Inference Pipeline ===")
        # 1. Feature Extractor
        self.extractor = DualEncoderFeatureExtractor(force_simulation=self.force_simulation)
        
        # 2. Classifier
        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        self.classifier = MSIClassifier(input_dim=DIM_FUSED).to(self.device)
        
        weights_path = os.path.join(MODELS_DIR, "best_classifier.pt")
        if os.path.exists(weights_path):
            try:
                self.classifier.load_state_dict(torch.load(weights_path, map_location=self.device))
                print("GemmaSight Pipeline: Loaded custom trained classifier weights.")
            except Exception as e:
                print(f"GemmaSight Pipeline Warning: Error loading custom weights ({e}). Running default weights.")
        else:
            print("GemmaSight Pipeline Warning: Trained weights not found. Initializing with default model weights.")
        self.classifier.eval()
        
        # 3. Case Retriever
        self.retriever = CaseRetriever(index_dim=DIM_FUSED)
        loaded = self.retriever.load_index()
        if not loaded:
            print("GemmaSight Pipeline Warning: FAISS database empty. Creating standard clinical database baseline...")
            self._build_dummy_retriever_database()
            
        # 4. Report Generator
        self.report_generator = MedGemmaReportGenerator(force_simulation=self.force_simulation)
        print("=== GemmaSight Inference Pipeline Initialized successfully ===")

    def _build_dummy_retriever_database(self):
        """Builds a temporary clinical FAISS database if none exists, ensuring retriever works."""
        dummy_embeddings = np.random.randn(5, DIM_FUSED).astype(np.float32)
        # Normalize
        norms = np.linalg.norm(dummy_embeddings, axis=1, keepdims=True)
        norms[norms == 0] = 1e-12
        dummy_embeddings = dummy_embeddings / norms
        
        labels = [1, 0, 1, 0, 1]
        descriptions = [
            "Crohn's-like reaction, poor differentiation, medullary shape features",
            "regular architecture, well-preserved crypt margins, lack of lymphocytic focus",
            "prominent tumor-infiltrating lymphocytes (TILs), high mucin content, signet ring segments",
            "intact tubular arrangement, mild atypical nuclear elongation, no dMMR indications",
            "focal lymphoid aggregation, extensive glandular architectural disruption, high inflammatory cells"
        ]
        patient_ids = ["GS-A045", "GS-B112", "GS-C902", "GS-D334", "GS-E501"]
        self.retriever.build_index(dummy_embeddings, labels, descriptions, patient_ids)

    def run_diagnosis(self, pil_image: Image.Image) -> tuple:
        """
        Runs the complete 5-phase end-to-end multimodal clinical diagnosis.
        Input: PIL Image (H&E colorectal patch)
        Returns:
            probability (float): MSI-High probability
            label (str): MSI-High or MSS label string
            saliency_heatmap (PIL.Image): Heatmap overlay visualization
            retrieved_cases (list): Top-3 matching cases from FAISS with confidence score
            report (str): Clinical pathology report text
        """
        # Phase 1: Dual-Encoder Feature Extraction
        features_tensor = self.extractor(pil_image) # (1, 1536)
        features_np = features_tensor.cpu().numpy()
        
        # Phase 2: Classification MLP
        with torch.no_grad():
            prob = self.classifier(features_tensor).item()
            
        predicted_class_idx = 1 if prob >= 0.5 else 0
        predicted_class_label = LABEL_MAP[predicted_class_idx]
        
        # Phase 3: Case Retrieval (FAISS)
        retrieved_cases = self.retriever.retrieve_top_k(features_np, k=3)
        
        # Phase 4: Spatial Explainability (Sliding-window Occlusion)
        print("GemmaSight Pipeline: Computing spatial explainability saliency maps...")
        saliency_heatmap = generate_occlusion_heatmap(pil_image, self.extractor, self.classifier)
        
        # Phase 5: Clinical Report Synthesis (MedGemma)
        print("GemmaSight Pipeline: Synthesizing clinical pathology report...")
        report = self.report_generator.generate_report(prob, retrieved_cases)
        
        return prob, predicted_class_label, saliency_heatmap, retrieved_cases, report
if __name__ == "__main__":
    # Self-test code
    pipeline = GemmaSightInferencePipeline(force_simulation=True)
    test_img = Image.fromarray((np.random.rand(224, 224, 3) * 255).astype(np.uint8))
    prob, label, heatmap, cases, report = pipeline.run_diagnosis(test_img)
    print(f"Prob: {prob:.4f} | Label: {label}")
    print(f"Retrieved Cases Count: {len(cases)}")
    print(f"Report Summary:\n{report[:150]}...")