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