StanfordAIMI / chexpert-plus-srrg_impression

Requirements:

pip install transformers
pip install torch

Adapted sample script for SRRG

import torch
from PIL import Image
from transformers import BertTokenizer, ViTImageProcessor, VisionEncoderDecoderModel, GenerationConfig
import requests
import re

model_name  = "StanfordAIMI/chexpert-plus-srrg_impression"
model = VisionEncoderDecoderModel.from_pretrained(model_name).eval()
tokenizer = BertTokenizer.from_pretrained(model_name)
image_processor = ViTImageProcessor.from_pretrained(model_name)
generation_args = {
   "bos_token_id": model.config.bos_token_id,
   "eos_token_id": model.config.eos_token_id,
   "pad_token_id": model.config.pad_token_id,
   "num_return_sequences": 1,
   "max_length": 128,
   "use_cache": True,
   "beam_width": 2,
}

# Inference
with torch.no_grad():
   url = "https://huggingface.co/IAMJB/interpret-cxr-impression-baseline/resolve/main/effusions-bibasal.jpg"
   image = Image.open(requests.get(url, stream=True).raw)
   pixel_values = image_processor(image, return_tensors="pt").pixel_values
   # Generate predictions
   generated_ids = model.generate(
       pixel_values,
       generation_config=GenerationConfig(
           **{**generation_args, "decoder_start_token_id": tokenizer.cls_token_id})
   )
   generated_texts = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)

impression = generated_texts[0]
print("Output raw impression:\n", impression)

# format impression
def process_impression_line(line):
    """
    Process one report line.

    - First, deduce from the first valid impression header the expected style:
      e.g. "1." vs. "1 ."
    - Then, only when a candidate match’s number equals the expected impression number,
      treat it as an impression header. For the first valid header, leave it unchanged.
      For subsequent valid headers, insert a newline before it.
    - Any candidate that does not match the expected number is assumed not to be an impression header.
    """
    impression_pattern = re.compile(r'(\d+)(\s*\.)(?!\d)')
    expected_impression = 1
    first_valid = True
    deduced_style = None

    def replacement(match):
        nonlocal expected_impression, first_valid, deduced_style
        candidate_num = int(match.group(1))
        candidate_style = match.group(2)  # may be " ." or "."

        # Only consider a candidate a valid impression header if it equals the expected number.
        if candidate_num == expected_impression:
            if first_valid:
                # This is our first valid impression header.
                first_valid = False
                deduced_style = candidate_style  # record style (with or without whitespace)
                expected_impression += 1
                return match.group(0)  # leave unchanged (i.e. no preceding newline)
            else:
                # For subsequent valid headers, we require the style to be the same as deduced.
                if candidate_style == deduced_style:
                    expected_impression += 1
                    return "\n" + match.group(0)
                else:
                    # If the style does not match the deduced style, leave it unchanged.
                    return match.group(0)
        else:
            # This candidate does not match the expected number; likely it's not a header.
            return match.group(0)

    processed = impression_pattern.sub(replacement, line)
    return processed


impression = impression.strip()
impression = process_impression_line(impression)
print("Formatted impression:\n", impression)

Output

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.