Views
No views yet
1pip install transformers
2pip install torch1import torch
2from PIL import Image
3from transformers import BertTokenizer, ViTImageProcessor, VisionEncoderDecoderModel, GenerationConfig
4import requests
5import re
6
7model_name = "StanfordAIMI/chexpert-plus-srrg_impression"
8model = VisionEncoderDecoderModel.from_pretrained(model_name).eval()
9tokenizer = BertTokenizer.from_pretrained(model_name)
10image_processor = ViTImageProcessor.from_pretrained(model_name)
11generation_args = {
12 "bos_token_id": model.config.bos_token_id,
13 "eos_token_id": model.config.eos_token_id,
14 "pad_token_id": model.config.pad_token_id,
15 "num_return_sequences": 1,
16 "max_length": 128,
17 "use_cache": True,
18 "beam_width": 2,
19}
20
21# Inference
22with torch.no_grad():
23 url = "https://huggingface.co/IAMJB/interpret-cxr-impression-baseline/resolve/main/effusions-bibasal.jpg"
24 image = Image.open(requests.get(url, stream=True).raw)
25 pixel_values = image_processor(image, return_tensors="pt").pixel_values
26 # Generate predictions
27 generated_ids = model.generate(
28 pixel_values,
29 generation_config=GenerationConfig(
30 **{**generation_args, "decoder_start_token_id": tokenizer.cls_token_id})
31 )
32 generated_texts = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)
33
34impression = generated_texts[0]
35print("Output raw impression:\n", impression)
36
37# format impression
38def process_impression_line(line):
39 """
40 Process one report line.
41
42 - First, deduce from the first valid impression header the expected style:
43 e.g. "1." vs. "1 ."
44 - Then, only when a candidate match’s number equals the expected impression number,
45 treat it as an impression header. For the first valid header, leave it unchanged.
46 For subsequent valid headers, insert a newline before it.
47 - Any candidate that does not match the expected number is assumed not to be an impression header.
48 """
49 impression_pattern = re.compile(r'(\d+)(\s*\.)(?!\d)')
50 expected_impression = 1
51 first_valid = True
52 deduced_style = None
53
54 def replacement(match):
55 nonlocal expected_impression, first_valid, deduced_style
56 candidate_num = int(match.group(1))
57 candidate_style = match.group(2) # may be " ." or "."
58
59 # Only consider a candidate a valid impression header if it equals the expected number.
60 if candidate_num == expected_impression:
61 if first_valid:
62 # This is our first valid impression header.
63 first_valid = False
64 deduced_style = candidate_style # record style (with or without whitespace)
65 expected_impression += 1
66 return match.group(0) # leave unchanged (i.e. no preceding newline)
67 else:
68 # For subsequent valid headers, we require the style to be the same as deduced.
69 if candidate_style == deduced_style:
70 expected_impression += 1
71 return "\n" + match.group(0)
72 else:
73 # If the style does not match the deduced style, leave it unchanged.
74 return match.group(0)
75 else:
76 # This candidate does not match the expected number; likely it's not a header.
77 return match.group(0)
78
79 processed = impression_pattern.sub(replacement, line)
80 return processed
81
82
83impression = impression.strip()
84impression = process_impression_line(impression)
85print("Formatted impression:\n", impression)Output raw impression:
1. moderate bilateral pleural effusions with associated bibasilar opacities, which may represent atelectasis or consolidation. 2. mild pulmonary edema.
Formatted impression:
1. moderate bilateral pleural effusions with associated bibasilar opacities, which may represent atelectasis or consolidation.
2. mild pulmonary edema.