Vision-Language Models for Clinical Applications

[Chest X-Ray] ↔ "No acute cardiopulmonary process. Heart size normal."

[Tissue Sample] ↔ "Adenocarcinoma, moderately differentiated"

[Skin Lesion] ↔ "2cm pigmented lesion, irregular borders, asymmetric"

[Symptom Sketch] ↔ "Sharp chest pain radiating to left arm"

# Image → Text generation
chest_xray = load_image("chest_xray.jpg")
report = vlm.generate_report(chest_xray)
# Output: "Cardiomediastinal silhouette within normal limits.
#          No acute infiltrate or effusion. No pneumothorax."

# No training examples needed - just text description
query = "Ground-glass opacities in bilateral lungs consistent with viral pneumonia"
 
# Compare image to text description
image_emb = vlm.encode_image(chest_xray)
text_emb = vlm.encode_text(query)
similarity = cosine_similarity(image_emb, text_emb)
 
# High similarity → likely match
if similarity > threshold:
    print("Consistent with COVID-19 pneumonia")

# Text query → Find matching images
query = "pulmonary embolism with right heart strain"
matching_images = vlm.search_images_by_text(query, image_database)
 
# Image query → Find similar reports
query_image = load_image("patient_xray.jpg")
similar_reports = vlm.search_text_by_image(query_image, report_database)

# Patient provides:
symptom_sketch = patient_drawing  # 3D body sketch with symptom locations
symptom_text = "Burning sensation in upper abdomen after meals"
 
# VLM predicts outcome
outcome_prob = vlm.predict(symptom_sketch, symptom_text)
# Predicts: Likely gastroesophageal reflux disease (GERD)

import torch
import torch.nn as nn
from transformers import AutoModel
from torchvision.models import resnet50
 
class ClinicalVLM(nn.Module):
    """
    CLIP-style vision-language model for healthcare
 
    Pre-trains on (medical_image, clinical_text) pairs
    using contrastive learning
    """
    def __init__(self, embed_dim=512):
        super().__init__()
 
        # ===== Vision Encoder =====
        # Option 1: CNN (ResNet)
        self.vision_encoder = resnet50(pretrained=True)
        self.vision_encoder.fc = nn.Identity()
        vision_dim = 2048
 
        # Option 2: Vision Transformer (ViT)
        # from transformers import ViTModel
        # self.vision_encoder = ViTModel.from_pretrained('google/vit-base-patch16-224')
        # vision_dim = 768
 
        # Project to shared embedding space
        self.vision_projection = nn.Linear(vision_dim, embed_dim)
 
        # ===== Text Encoder =====
        # ClinicalBERT for medical text understanding
        self.text_encoder = AutoModel.from_pretrained(
            'emilyalsentzer/Bio_ClinicalBERT'
        )
        text_dim = 768  # BERT hidden size
 
        # Project to shared embedding space
        self.text_projection = nn.Linear(text_dim, embed_dim)
 
        # ===== Learnable Temperature =====
        # Temperature for contrastive loss (like CLIP)
        self.temperature = nn.Parameter(torch.ones([]) * np.log(1 / 0.07))
 
    def encode_image(self, image):
        """
        Encode medical image to embedding
 
        Args:
            image: (batch, 3, H, W) - medical image
 
        Returns:
            image_emb: (batch, embed_dim) - normalized embedding
        """
        # Extract visual features
        features = self.vision_encoder(image)  # (batch, vision_dim)
 
        # Project to shared space
        emb = self.vision_projection(features)  # (batch, embed_dim)
 
        # Normalize (important for cosine similarity)
        emb = F.normalize(emb, dim=-1)
 
        return emb
 
    def encode_text(self, text_tokens):
        """
        Encode clinical text to embedding
 
        Args:
            text_tokens: dict - tokenized text
 
        Returns:
            text_emb: (batch, embed_dim) - normalized embedding
        """
        # Extract text features
        output = self.text_encoder(**text_tokens)
        features = output.pooler_output  # (batch, text_dim)
 
        # Project to shared space
        emb = self.text_projection(features)  # (batch, embed_dim)
 
        # Normalize
        emb = F.normalize(emb, dim=-1)
 
        return emb
 
    def forward(self, images, text_tokens):
        """
        Compute contrastive loss (CLIP-style)
 
        Args:
            images: (batch, 3, H, W) - medical images
            text_tokens: dict - tokenized clinical text
 
        Returns:
            loss: scalar - bidirectional contrastive loss
        """
        batch_size = images.shape[0]
 
        # Encode both modalities
        image_emb = self.encode_image(images)  # (batch, embed_dim)
        text_emb = self.encode_text(text_tokens)  # (batch, embed_dim)
 
        # Compute similarity matrix (scaled by temperature)
        logits = image_emb @ text_emb.T / self.temperature.exp()
        # logits: (batch, batch) - similarity between all pairs
 
        # Labels: diagonal elements are positive pairs
        labels = torch.arange(batch_size, device=logits.device)
 
        # Bidirectional contrastive loss
        loss_i2t = F.cross_entropy(logits, labels)      # Image → Text
        loss_t2i = F.cross_entropy(logits.T, labels)    # Text → Image
        loss = (loss_i2t + loss_t2i) / 2
 
        return loss

def pretrain_clinical_vlm(model, dataset, epochs=50):
    """
    Pre-train VLM on (medical_image, clinical_text) pairs
 
    Dataset: MIMIC-CXR (377,000 chest X-rays + reports),
             OpenI, PadChest, etc.
    """
    optimizer = torch.optim.AdamW(model.parameters(), lr=1e-4)
 
    for epoch in range(epochs):
        for images, texts in dataset:
            # Contrastive loss
            loss = model(images, texts)
 
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()
 
        # Evaluate zero-shot performance
        if epoch % 5 == 0:
            zero_shot_acc = evaluate_zero_shot(model, test_dataset)
            print(f"Epoch {epoch}: Zero-shot accuracy = {zero_shot_acc:.3f}")

@torch.no_grad()
def zero_shot_diagnosis(model, image, candidate_diagnoses):
    """
    Classify image using text descriptions (no labeled examples)
 
    Args:
        image: Medical image tensor
        candidate_diagnoses: List of text descriptions
 
    Returns:
        predicted_diagnosis: Most likely match
        probabilities: Confidence for each diagnosis
    """
    # Encode image
    image_emb = model.encode_image(image.unsqueeze(0))  # (1, embed_dim)
 
    # Encode all candidate diagnoses
    diagnosis_embs = []
    for diagnosis_text in candidate_diagnoses:
        tokens = tokenizer(diagnosis_text, return_tensors='pt')
        text_emb = model.encode_text(tokens)
        diagnosis_embs.append(text_emb)
 
    diagnosis_embs = torch.cat(diagnosis_embs, dim=0)  # (num_diagnoses, embed_dim)
 
    # Compute similarities
    similarities = image_emb @ diagnosis_embs.T  # (1, num_diagnoses)
 
    # Convert to probabilities
    probabilities = F.softmax(similarities / temperature, dim=-1)
 
    # Predict
    predicted_idx = probabilities.argmax()
    predicted_diagnosis = candidate_diagnoses[predicted_idx]
 
    return predicted_diagnosis, probabilities
 
# Example usage
image = load_chest_xray("patient.jpg")
candidates = [
    "Normal chest X-ray with no acute findings",
    "Pneumonia with consolidation in right lower lobe",
    "Congestive heart failure with pulmonary edema",
    "Pneumothorax on the left side"
]
 
diagnosis, probs = zero_shot_diagnosis(model, image, candidates)
print(f"Predicted: {diagnosis} (confidence: {probs.max():.2f})")

# Good prompts: Specific, medical terminology, structured
good_prompts = [
    "Chest X-ray demonstrating acute pulmonary edema with bilateral infiltrates",
    "CT scan revealing acute ischemic stroke in left MCA territory",
    "Dermatology image showing melanoma with irregular borders and variegated color"
]
 
# Poor prompts: Vague, non-medical language
poor_prompts = [
    "Lung problem",  # Too vague
    "Heart issue",   # Non-specific
    "Bad skin"       # Lacks clinical detail
]
 
# Template ensembles (like CLIP)
def create_prompt_ensemble(diagnosis):
    """Generate multiple prompt variations"""
    templates = [
        f"This is a case of {diagnosis}",
        f"Findings consistent with {diagnosis}",
        f"Radiographic appearance of {diagnosis}",
        f"Clinical presentation suggests {diagnosis}"
    ]
    return templates
 
# Average across prompts for robustness
ensemble_probs = []
for template in create_prompt_ensemble("pneumonia"):
    prob = model.predict(image, template)
    ensemble_probs.append(prob)
 
final_prob = torch.stack(ensemble_probs).mean()

# Start with ImageNet + general BERT
# Fine-tune on medical domain first
vision_encoder = resnet50(pretrained=True)  # ImageNet
text_encoder = ClinicalBERT  # Pre-trained on MIMIC notes
 
# Then: contrastive pre-training on (image, report) pairs

class MomentumClinicalVLM(nn.Module):
    """
    Use momentum encoders for large effective batch size
    (when GPU memory limits actual batch size)
    """
    def __init__(self):
        self.online_encoder = ClinicalVLM()
        self.momentum_encoder = copy.deepcopy(self.online_encoder)
 
        # Freeze momentum encoder (updated via EMA)
        for param in self.momentum_encoder.parameters():
            param.requires_grad = False
 
        # Queue of previous embeddings (acts as larger batch)
        self.register_buffer("queue", torch.randn(65536, embed_dim))
 
    @torch.no_grad()
    def update_momentum_encoder(self, momentum=0.999):
        """Exponential moving average update"""
        for param_q, param_k in zip(
            self.online_encoder.parameters(),
            self.momentum_encoder.parameters()
        ):
            param_k.data = param_k.data * momentum + param_q.data * (1 - momentum)

# Medical images often need higher resolution
high_res_transforms = T.Compose([
    T.Resize(512),  # vs 224 for CLIP
    T.CenterCrop(512),
    T.ToTensor(),
    T.Normalize(mean=[0.485], std=[0.229])
])

from sklearn.metrics import roc_auc_score, accuracy_score
 
# Zero-shot classification accuracy
zero_shot_acc = accuracy_score(y_true, y_pred)
 
# Image-text retrieval
# - Image→Text: Given image, retrieve correct report
# - Text→Image: Given report, retrieve correct image
recall_at_k = compute_recall_at_k(image_emb, text_emb, k=[1, 5, 10])

# Clinical utility (beyond accuracy)
def evaluate_clinical_utility(model, test_data):
    """
    Metrics that matter for clinical deployment
    """
    metrics = {}
 
    # 1. AUROC (discrimination ability)
    metrics['auroc'] = roc_auc_score(y_true, y_pred_prob)
 
    # 2. Sensitivity at high specificity
    # (Catch critical cases without too many false alarms)
    fpr, tpr, thresholds = roc_curve(y_true, y_pred_prob)
    sens_at_95spec = tpr[np.where(fpr <= 0.05)[0][-1]]
    metrics['sensitivity_at_95spec'] = sens_at_95spec
 
    # 3. Calibration (are probabilities accurate?)
    prob_true, prob_pred = calibration_curve(y_true, y_pred_prob, n_bins=10)
    metrics['calibration_error'] = np.abs(prob_true - prob_pred).mean()
 
    # 4. Fairness across demographics
    for demographic in ['age', 'sex', 'race']:
        group_aurocs = {}
        for group in np.unique(test_data[demographic]):
            mask = test_data[demographic] == group
            group_aurocs[group] = roc_auc_score(y_true[mask], y_pred_prob[mask])
 
        metrics[f'{demographic}_disparity'] = max(group_aurocs.values()) - min(group_aurocs.values())
 
    return metrics