| """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() |
|
|
