Views
No views yet
1import json
2import torch
3from transformers import BertTokenizer, BertForSequenceClassification
4from datasets import load_dataset
5import requests
6
7# Configuration
8MODEL_PATH = "StanfordAIMI/SRR-BERT-Leaves"
9MAPPING_URL = "https://raw.githubusercontent.com/jbdel/StructEval/refs/heads/main/structeval/leaves_mapping.json"
10MAX_LENGTH = 128
11DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
12
13# Fetch mapping from GitHub
14resp = requests.get(MAPPING_URL)
15resp.raise_for_status()
16label_map = resp.json()
17idx2label = {v: k for k, v in label_map.items()}
18
19# Load tokenizer & model
20tokenizer = BertTokenizer.from_pretrained("microsoft/BiomedVLP-CXR-BERT-general")
21model = BertForSequenceClassification.from_pretrained(MODEL_PATH, num_labels=len(label_map))
22model.to(DEVICE).eval()
23
24# Grab one test sentence
25dataset = load_dataset("StanfordAIMI/StructUtterances", split="test_reviewed")
26sentence = dataset[35]["utterance"]
27
28# Tokenize and infer
29inputs = tokenizer(
30 sentence,
31 padding="max_length",
32 truncation=True,
33 max_length=MAX_LENGTH,
34 return_tensors="pt"
35).to(DEVICE)
36
37with torch.no_grad():
38 logits = model(**inputs).logits
39 preds = (torch.sigmoid(logits)[0].cpu().numpy() > 0.5).astype(int)
40
41pred_labels = [idx2label[i] for i, flag in enumerate(preds) if flag]
42
43print(f"Sentence: {sentence}")
44print("Predicted labels:", pred_labels)Sentence: Patchy consolidation in the left retrocardiac area, suggestive of atelectasis or early airspace disease.
Predicted labels: ['Atelectasis', 'Air space opacity–multifocal']{
"No Finding": 0,
"Lung Lesion": 1,
"Edema": 2,
"Pneumonia": 3,
"Atelectasis": 4,
"Aspiration": 5,
"Lung collapse": 6,
"Perihilar airspace opacity": 7,
"Air space opacity\u2013multifocal": 8,
"Mass/Solitary lung mass": 9,
"Nodule/Solitary lung nodule": 10,
"Cavitating mass with content": 11,
"Cavitating masses": 12,
"Emphysema": 13,
"Fibrosis": 14,
"Pulmonary congestion": 15,
"Hilar lymphadenopathy": 16,
"Bronchiectasis": 17,
"Simple pneumothorax": 18,
"Loculated pneumothorax": 19,
"Tension pneumothorax": 20,
"Simple pleural effusion": 21,
"Loculated pleural effusion": 22,
"Pleural scarring": 23,
"Hydropneumothorax": 24,
"Pleural Other": 25,
"Cardiomegaly": 26,
"Pericardial effusion": 27,
"Inferior mediastinal mass": 28,
"Superior mediastinal mass": 29,
"Tortuous Aorta": 30,
"Calcification of the Aorta": 31,
"Enlarged pulmonary artery": 32,
"Hernia": 33,
"Pneumomediastinum": 34,
"Tracheal deviation": 35,
"Acute humerus fracture": 36,
"Acute rib fracture": 37,
"Acute clavicle fracture": 38,
"Acute scapula fracture": 39,
"Compression fracture": 40,
"Shoulder dislocation": 41,
"Subcutaneous Emphysema": 42,
"Suboptimal central line": 43,
"Suboptimal endotracheal tube": 44,
"Suboptimal nasogastric tube": 45,
"Suboptimal pulmonary arterial catheter": 46,
"Pleural tube": 47,
"PICC line": 48,
"Port catheter": 49,
"Pacemaker": 50,
"Implantable defibrillator": 51,
"LVAD": 52,
"Intraaortic balloon pump": 53,
"Pneumoperitoneum": 54
} precision recall f1-score support
No Finding 0.93 0.91 0.92 59962
Lung Lesion 0.66 0.51 0.58 533
Edema 0.96 0.98 0.97 10049
Pneumonia 0.90 0.95 0.92 14164
Atelectasis 0.99 0.99 0.99 17551
Aspiration 0.98 1.00 0.99 2566
Lung collapse 0.93 0.83 0.88 1097
Perihilar airspace opacity 0.87 0.60 0.71 1442
Air space opacity-multifocal 0.81 0.73 0.77 4816
Mass/Solitary lung mass 0.92 0.89 0.91 1213
Nodule/Solitary lung nodule 0.92 0.96 0.94 2759
Cavitating mass with content 0.93 0.91 0.92 163
Cavitating masses 0.83 0.82 0.83 55
Emphysema 0.92 0.97 0.94 1931
Fibrosis 0.79 0.87 0.83 2432
Pulmonary congestion 0.88 0.87 0.87 2960
Hilar lymphadenopathy 0.85 0.91 0.87 1170
Bronchiectasis 0.98 0.96 0.97 480
Simple pneumothorax 0.86 0.92 0.89 5336
Loculated pneumothorax 0.90 0.87 0.88 159
Tension pneumothorax 0.86 0.87 0.86 84
Simple pleural effusion 0.97 0.98 0.97 16162
Loculated pleural effusion 0.98 0.90 0.94 1149
Pleural scarring 0.87 0.94 0.90 3098
Hydropneumothorax 0.99 0.96 0.98 395
Pleural Other 0.75 0.63 0.68 377
Cardiomegaly 0.96 0.96 0.96 5463
Pericardial effusion 0.98 0.94 0.96 334
Inferior mediastinal mass 0.71 0.51 0.59 83
Superior mediastinal mass 0.79 0.86 0.82 514
Tortuous Aorta 0.91 0.97 0.94 1782
Calcification of the Aorta 0.96 0.91 0.93 1542
Enlarged pulmonary artery 0.93 0.97 0.95 482
Hernia 0.94 0.97 0.95 657
Pneumomediastinum 0.91 0.93 0.92 270
Tracheal deviation 0.82 0.89 0.86 463
Acute humerus fracture 0.82 0.87 0.85 206
Acute rib fracture 0.87 0.93 0.90 1988
Acute clavicle fracture 0.87 0.93 0.90 335
Acute scapula fracture 0.94 0.80 0.86 80
Compression fracture 0.78 0.96 0.86 794
Shoulder dislocation 0.91 0.78 0.84 55
Subcutaneous Emphysema 0.98 0.99 0.99 1115
Suboptimal central line 0.62 0.20 0.30 393
Suboptimal endotracheal tube 0.64 0.54 0.58 1014
Suboptimal nasogastric tube 0.62 0.35 0.45 835
Suboptimal pulmonary arterial catheter 0.41 0.11 0.18 80
Pleural tube 0.69 0.55 0.61 1458
PICC line 0.74 0.88 0.80 2273
Port catheter 0.67 0.77 0.72 686
Pacemaker 0.82 0.92 0.87 2189
Implantable defibrillator 0.66 0.71 0.68 520
LVAD 0.73 0.87 0.79 127
Intraaortic balloon pump 0.77 0.93 0.84 120
Pneumoperitoneum 0.96 0.87 0.91 342
micro avg 0.91 0.92 0.92 178303
macro avg 0.85 0.83 0.83 178303
weighted avg 0.91 0.92 0.91 178303
samples avg 0.91 0.92 0.91 178303