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/SRRG-BERT-Upper-with-Statuses"
9MAPPING_URL = "https://raw.githubusercontent.com/jbdel/StructEval/refs/heads/main/structeval/upper_with_statuses_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: ['Consolidation (Uncertain)', 'Air space opacity (Uncertain)']{
"Pleural Effusion (Present)": 0,
"Upper abdominal finding (Present)": 1,
"Widened cardiac silhouette (Present)": 2,
"Lung Finding (Present)": 3,
"Widened aortic contour (Present)": 4,
"Pleural Thickening (Present)": 5,
"Vascular finding (Present)": 6,
"Consolidation (Present)": 7,
"Pneumothorax (Present)": 8,
"Subdiaphragmatic gas (Present)": 9,
"Masslike opacity (Present)": 10,
"Chest wall finding (Present)": 11,
"Focal air space opacity (Present)": 12,
"Segmental collapse (Present)": 13,
"Fracture (Present)": 14,
"Mediastinal mass (Present)": 15,
"Solitary masslike opacity (Present)": 16,
"Support Devices (Present)": 17,
"Mediastinal finding (Present)": 18,
"Pleural finding (Present)": 19,
"Air space opacity (Present)": 20,
"Diffuse air space opacity (Present)": 21,
"Multiple masslike opacities (Present)": 22,
"Musculoskeletal finding (Present)": 23,
"Pleural Effusion (Uncertain)": 24,
"Upper abdominal finding (Uncertain)": 25,
"Widened cardiac silhouette (Uncertain)": 26,
"Lung Finding (Uncertain)": 27,
"Widened aortic contour (Uncertain)": 28,
"Pleural Thickening (Uncertain)": 29,
"Vascular finding (Uncertain)": 30,
"Consolidation (Uncertain)": 31,
"Pneumothorax (Uncertain)": 32,
"Subdiaphragmatic gas (Uncertain)": 33,
"Masslike opacity (Uncertain)": 34,
"Chest wall finding (Uncertain)": 35,
"Focal air space opacity (Uncertain)": 36,
"Segmental collapse (Uncertain)": 37,
"Fracture (Uncertain)": 38,
"Mediastinal mass (Uncertain)": 39,
"Solitary masslike opacity (Uncertain)": 40,
"Support Devices (Uncertain)": 41,
"Mediastinal finding (Uncertain)": 42,
"Pleural finding (Uncertain)": 43,
"Air space opacity (Uncertain)": 44,
"Diffuse air space opacity (Uncertain)": 45,
"Multiple masslike opacities (Uncertain)": 46,
"Musculoskeletal finding (Uncertain)": 47,
"Pleural Effusion (Absent)": 48,
"Upper abdominal finding (Absent)": 49,
"Widened cardiac silhouette (Absent)": 50,
"Lung Finding (Absent)": 51,
"Widened aortic contour (Absent)": 52,
"Pleural Thickening (Absent)": 53,
"Vascular finding (Absent)": 54,
"Consolidation (Absent)": 55,
"Pneumothorax (Absent)": 56,
"Subdiaphragmatic gas (Absent)": 57,
"Masslike opacity (Absent)": 58,
"Chest wall finding (Absent)": 59,
"Focal air space opacity (Absent)": 60,
"Segmental collapse (Absent)": 61,
"Fracture (Absent)": 62,
"Mediastinal mass (Absent)": 63,
"Solitary masslike opacity (Absent)": 64,
"Support Devices (Absent)": 65,
"Mediastinal finding (Absent)": 66,
"Pleural finding (Absent)": 67,
"Air space opacity (Absent)": 68,
"Diffuse air space opacity (Absent)": 69,
"Multiple masslike opacities (Absent)": 70,
"Musculoskeletal finding (Absent)": 71,
"No Finding": 72
} precision recall f1-score support
Pleural Effusion (Present) 0.97 0.98 0.97 12112
Upper abdominal finding (Present) 0.00 0.00 0.00 0
Widened cardiac silhouette (Present) 0.96 0.98 0.97 4125
Lung Finding (Present) 0.87 0.87 0.87 5159
Widened aortic contour (Present) 0.92 0.94 0.93 1512
Pleural Thickening (Present) 0.84 0.92 0.88 1729
Vascular finding (Present) 0.92 0.95 0.94 1676
Consolidation (Present) 0.90 0.90 0.90 10395
Pneumothorax (Present) 0.97 0.99 0.98 2706
Subdiaphragmatic gas (Present) 0.91 0.97 0.94 217
Masslike opacity (Present) 0.00 0.00 0.00 0
Chest wall finding (Present) 0.98 0.99 0.98 1049
Focal air space opacity (Present) 0.85 0.68 0.75 815
Segmental collapse (Present) 0.84 0.87 0.85 789
Fracture (Present) 0.89 0.93 0.91 1849
Mediastinal mass (Present) 0.84 0.83 0.83 133
Solitary masslike opacity (Present) 0.88 0.91 0.90 2072
Support Devices (Present) 0.75 0.73 0.74 8674
Mediastinal finding (Present) 0.88 0.93 0.90 1065
Pleural finding (Present) 0.90 0.85 0.88 455
Air space opacity (Present) 0.77 0.72 0.74 2557
Diffuse air space opacity (Present) 0.94 0.95 0.94 5544
Multiple masslike opacities (Present) 0.93 0.62 0.75 40
Musculoskeletal finding (Present) 0.80 0.92 0.85 38
Pleural Effusion (Uncertain) 0.92 0.91 0.91 3853
Upper abdominal finding (Uncertain) 0.00 0.00 0.00 0
Widened cardiac silhouette (Uncertain) 0.79 0.79 0.79 1232
Lung Finding (Uncertain) 0.79 0.69 0.74 3880
Widened aortic contour (Uncertain) 0.84 0.61 0.71 268
Pleural Thickening (Uncertain) 0.83 0.87 0.85 1358
Vascular finding (Uncertain) 0.76 0.67 0.71 329
Consolidation (Uncertain) 0.91 0.92 0.91 14858
Pneumothorax (Uncertain) 0.91 0.94 0.92 589
Subdiaphragmatic gas (Uncertain) 0.77 0.84 0.80 82
Masslike opacity (Uncertain) 0.00 0.00 0.00 0
Chest wall finding (Uncertain) 0.86 0.63 0.73 19
Focal air space opacity (Uncertain) 0.74 0.35 0.48 614
Segmental collapse (Uncertain) 0.78 0.72 0.75 256
Fracture (Uncertain) 0.76 0.78 0.77 419
Mediastinal mass (Uncertain) 0.74 0.83 0.78 428
Solitary masslike opacity (Uncertain) 0.86 0.87 0.87 1831
Support Devices (Uncertain) 0.45 0.22 0.29 220
Mediastinal finding (Uncertain) 0.83 0.73 0.78 256
Pleural finding (Uncertain) 0.75 0.59 0.66 296
Air space opacity (Uncertain) 0.67 0.55 0.61 2218
Diffuse air space opacity (Uncertain) 0.91 0.88 0.90 3337
Multiple masslike opacities (Uncertain) 0.60 0.64 0.62 14
Musculoskeletal finding (Uncertain) 0.67 0.46 0.55 13
Pleural Effusion (Absent) 0.77 0.66 0.71 1008
Upper abdominal finding (Absent) 0.00 0.00 0.00 0
Widened cardiac silhouette (Absent) 0.82 0.57 0.67 145
Lung Finding (Absent) 0.72 0.55 0.62 228
Widened aortic contour (Absent) 0.00 0.00 0.00 0
Pleural Thickening (Absent) 0.00 0.00 0.00 1
Vascular finding (Absent) 0.00 0.00 0.00 8
Consolidation (Absent) 0.89 0.81 0.85 1114
Pneumothorax (Absent) 0.75 0.71 0.73 2198
Subdiaphragmatic gas (Absent) 0.84 0.72 0.78 43
Masslike opacity (Absent) 0.00 0.00 0.00 0
Chest wall finding (Absent) 0.95 0.87 0.91 46
Focal air space opacity (Absent) 0.47 0.89 0.62 9
Segmental collapse (Absent) 0.65 0.43 0.52 35
Fracture (Absent) 0.77 0.79 0.78 925
Mediastinal mass (Absent) 0.50 0.08 0.14 12
Solitary masslike opacity (Absent) 0.73 0.77 0.75 145
Support Devices (Absent) 0.34 0.06 0.11 250
Mediastinal finding (Absent) 0.88 0.78 0.83 68
Pleural finding (Absent) 0.93 0.72 0.81 18
Air space opacity (Absent) 0.56 0.29 0.38 34
Diffuse air space opacity (Absent) 0.81 0.84 0.83 1119
Multiple masslike opacities (Absent) 0.00 0.00 0.00 1
Musculoskeletal finding (Absent) 0.00 0.00 0.00 4
No Finding 0.92 0.92 0.92 59962
micro avg 0.89 0.88 0.89 168454
macro avg 0.68 0.64 0.65 168454
weighted avg 0.89 0.88 0.88 168454
samples avg 0.88 0.88 0.88 168454