File size: 1,614 Bytes
32e07d6 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 | """Score a sentence against a chest X-ray and save its similarity map.
python inference.py --image cxr.png --text "There is a right pleural effusion."
"""
import argparse
import os
import warnings
import torch
from transformers import AutoImageProcessor, AutoModel, AutoTokenizer
from utils import model_inference, save_similarity_map
warnings.filterwarnings("ignore", category=UserWarning)
REPO_DIR = os.path.dirname(os.path.abspath(__file__))
def load_model(device=None, dtype=torch.float32):
device = device or ("cuda" if torch.cuda.is_available() else "cpu")
tokenizer = AutoTokenizer.from_pretrained(REPO_DIR)
image_processor = AutoImageProcessor.from_pretrained(REPO_DIR)
model = AutoModel.from_pretrained(REPO_DIR, trust_remote_code=True, dtype=dtype).to(device).eval()
return {"tokenizer": tokenizer, "image_processor": image_processor, "model": model}
def main():
parser = argparse.ArgumentParser(description="Score a sentence against a chest X-ray and save its similarity map.")
parser.add_argument("--image", default="cxr.png", help="PNG / JPG or DICOM file")
parser.add_argument("--text", default="There is effusion.")
parser.add_argument("--output", default="similarity_map.png", help="heat-map overlay")
args = parser.parse_args()
models = load_model()
prob, similarity_map = model_inference(args.image, args.text, **models)
print(f"{args.text!r}: probability {prob.item():.3f}")
save_similarity_map(args.image, similarity_map, args.output)
print(f"similarity map saved to {args.output}")
if __name__ == "__main__":
main()
|