File size: 2,038 Bytes
db46ad7 c5a246d db46ad7 c5a246d db46ad7 c5a246d db46ad7 c5a246d db46ad7 c5a246d db46ad7 c5a246d | 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 45 46 47 48 49 50 51 52 53 54 55 56 57 58 | """Score a sentence against a chest CT volume and save its similarity map as NIfTI.
python inference.py --volume volume.nii.gz --text "There is a lung nodule."
"""
import argparse
import os
import warnings
import torch
from transformers import AutoModel, AutoTokenizer
from utils import (
load_ct_volume,
model_inference,
save_similarity_map,
similarity_scores_to_volume,
)
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)
model = AutoModel.from_pretrained(REPO_DIR, trust_remote_code=True, dtype=dtype).to(device).eval()
return {"tokenizer": tokenizer, "model": model}
def main():
parser = argparse.ArgumentParser(description=__doc__.splitlines()[0])
parser.add_argument("--volume", default="volume.nii.gz", help="NIfTI CT volume or DICOM series directory")
parser.add_argument("--text", default="There is a lung nodule.")
parser.add_argument("--output", default="similarity_map.nii.gz", help="similarity map on the input CT grid (NIfTI)")
parser.add_argument(
"--raw_spacing",
type=float,
nargs=3,
default=None,
help="voxel spacing along the reoriented I, P, L axes (e.g. 1.5 0.7 0.7) if the input affine is unreliable",
)
args = parser.parse_args()
models = load_model()
pixel_values, geometry = load_ct_volume(args.volume, raw_spacing=args.raw_spacing, return_geometry=True)
prob, scores = model_inference(pixel_values, args.text, **models)
print(f"{args.text!r}: probability {prob.item():.3f}")
similarity_map = similarity_scores_to_volume(scores, models["model"], geometry)
save_similarity_map(similarity_map, geometry["affine"], args.output)
print(f"similarity map {tuple(similarity_map.shape)} saved to {args.output}")
if __name__ == "__main__":
main()
|