Resample similarity maps onto the input CT grid

- similarity_scores_to_volume(scores, model, geometry) returns the map with the input CT's shape and affine (geometry from load_ct_volume(..., return_geometry=True)); without geometry it stays on the 240x384x384 preprocessed grid
- load_ct_volume: return_affine -> return_geometry; accepts a NIfTI file or a DICOM series directory
- inference.py: saves only similarity_map.nii.gz (no --ct_output); drops the 18 CT-RATE finding scores
- README: new API, geometry, raw_spacing axis order (I, P, L), pydicom for DICOM input

README.md

@@ -43,11 +43,11 @@ Both examples use the bundled `inference.py` / `utils.py`; run them from the dow
 ```bash
 hf download Deepnoid/GLINT-CT --local-dir glint-ct && cd glint-ct
 python inference.py --volume volume.nii.gz --text "There is a lung nodule."
-# prints the sentence probability and the 18 CT-RATE findings,
-# saves similarity_map.nii.gz and ct_preprocessed.nii.gz
+# prints the sentence probability and saves similarity_map.nii.gz
 ```
 
-Both files share one affine, so the map overlays the CT in ITK-SNAP / 3D Slicer.
+`--volume` takes a NIfTI file or a DICOM series directory. The map has the input CT's shape and affine, so it
+overlays the input in ITK-SNAP / 3D Slicer.
 
 ### AutoModel
 
@@ -62,17 +62,23 @@ device = "cuda" if torch.cuda.is_available() else "cpu"
 model = AutoModel.from_pretrained(repo, trust_remote_code=True, dtype=torch.float32).to(device).eval()
 tokenizer = AutoTokenizer.from_pretrained(repo)
 
-pixel_values, affine = load_ct_volume("volume.nii.gz", return_affine=True)   # [1, 80, 3, 384, 384]
+pixel_values, geometry = load_ct_volume("volume.nii.gz", return_geometry=True)   # [1, 80, 3, 384, 384]
 prob, scores = model_inference(pixel_values, "There is a lung nodule.", model=model, tokenizer=tokenizer)
-volume = similarity_scores_to_volume(scores, model)                          # (240, 384, 384)
-save_similarity_map(volume, affine, "similarity_map.nii.gz", pixel_values, "ct_preprocessed.nii.gz")
+similarity_map = similarity_scores_to_volume(scores, model, geometry)           # input CT shape, e.g. (512, 512, 205)
+save_similarity_map(similarity_map, geometry["affine"], "similarity_map.nii.gz")
 ```
 
 - **Outputs:** `prob` is the sentence probability (sigmoid); `scores` is the per-token similarity in [0, 1]
-  (40 × 24 × 24 tokens).
-- **Inputs:** `utils.load_ct_volume` applies the training preprocessing (IPL, HU [-1000, 1000] → [-1, 1],
-  1.5 × 0.75 × 0.75 mm, crop/pad to 240×480×480, resize to 240×384×384) and returns `[1, 80, 3, 384, 384]`.
-- **Unreliable NIfTI affine:** pass `raw_spacing=(sx, sy, sz)` to `load_ct_volume` (`--raw_spacing` in the script).
+  (40 × 24 × 24 tokens); `similarity_scores_to_volume` resamples it onto the input CT grid (0 outside the 360 mm
+  preprocessing crop), or keeps the 240×384×384 preprocessed grid without `geometry`.
+- **`geometry`:** returned by `load_ct_volume(..., return_geometry=True)`: `shape` and `affine` of the input CT (the
+  map's NIfTI header) and `to_model`, the 4×4 input-voxel → preprocessed-voxel mapping used for the resampling; for
+  array input `affine` is `None` and the map stays on the preprocessed grid.
+- **Inputs:** a NIfTI file or a DICOM series directory (one series); `utils.load_ct_volume` applies the training
+  preprocessing (IPL, HU [-1000, 1000] → [-1, 1], 1.5 × 0.75 × 0.75 mm, crop/pad to 240×480×480, resize to
+  240×384×384) and returns `[1, 80, 3, 384, 384]`.
+- **Unreliable affine:** pass the voxel spacing along the reoriented I, P, L axes, e.g. `raw_spacing=(1.5, 0.7, 0.7)`,
+  to `load_ct_volume` (`--raw_spacing` in the script).
 - **Underlying call:** `model.compute_logits(pixel_values, [encoded_text])` returns `logits` and `similarity_scores`.
 
 Vision encoder only (`vision_only=True`):
@@ -93,7 +99,7 @@ print(out.last_hidden_state.shape)               # (1, 23040, 1024)
 
 ## Requirements
 
-`transformers==4.56.1`, `torch`, `entmax`, `monai`, `nibabel`.
+`transformers==4.56.1`, `torch`, `entmax`, `monai`, `nibabel`; `pydicom` for DICOM series input.
 
 > ⚠️ Keep `transformers==4.56.1`, the version the bundled code was validated with.
 

inference.py

@@ -1,4 +1,4 @@
-"""Score a sentence against a chest CT volume, save its similarity map as NIfTI and score the 18 CT-RATE findings.
+"""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."
 """
@@ -13,7 +13,6 @@ from transformers import AutoModel, AutoTokenizer
 from utils import (
     load_ct_volume,
     model_inference,
-    model_inference_multiple_text,
     save_similarity_map,
     similarity_scores_to_volume,
 )
@@ -22,27 +21,6 @@ warnings.filterwarnings("ignore", category=UserWarning)
 
 REPO_DIR = os.path.dirname(os.path.abspath(__file__))
 
-CT_RATE_LABELS = [
-    "Medical material",
-    "Arterial wall calcification",
-    "Cardiomegaly",
-    "Pericardial effusion",
-    "Coronary artery wall calcification",
-    "Hiatal hernia",
-    "Lymphadenopathy",
-    "Emphysema",
-    "Atelectasis",
-    "Lung nodule",
-    "Lung opacity",
-    "Pulmonary fibrotic sequela",
-    "Pleural effusion",
-    "Mosaic attenuation pattern",
-    "Peribronchial thickening",
-    "Consolidation",
-    "Bronchiectasis",
-    "Interlobular septal thickening",
-]
-
 
 def load_model(device=None, dtype=torch.float32):
     device = device or ("cuda" if torch.cuda.is_available() else "cpu")
@@ -53,30 +31,26 @@ def load_model(device=None, dtype=torch.float32):
 
 def main():
     parser = argparse.ArgumentParser(description=__doc__.splitlines()[0])
-    parser.add_argument("--volume", default="volume.nii.gz", help="NIfTI CT volume")
+    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 (NIfTI)")
-    parser.add_argument(
-        "--ct_output", default="ct_preprocessed.nii.gz", help="preprocessed CT on the same grid (NIfTI)"
-    )
+    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 if the NIfTI affine is unreliable"
+        "--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, affine = load_ct_volume(args.volume, raw_spacing=args.raw_spacing, return_affine=True)
+    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}")
-    volume = similarity_scores_to_volume(scores, models["model"])
-    save_similarity_map(volume, affine, args.output, pixel_values, args.ct_output)
-    print(f"similarity map saved to {args.output} (CT: {args.ct_output})")
-
-    prompts = [f"There is {label}" for label in CT_RATE_LABELS]  # zero-shot evaluation prompts
-    probs, _ = model_inference_multiple_text(pixel_values, prompts, **models)
-    for label, p in zip(CT_RATE_LABELS, probs.flatten().tolist()):
-        print(f"  {label:42s} {p:.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__":

utils.py

@@ -83,21 +83,41 @@ def normalize_pixel_values_videos(pixel_values_videos):
     return pixel_values_videos
 
 
-def load_ct_volume(volume, orientd=True, raw_spacing=None, return_affine=False):
+def _input_geometry(image, orientd, raw_spacing):
+    """Input voxel grid (shape, affine) and its 4x4 voxel mapping onto the preprocessed grid."""
+    import nibabel as nib
+
+    affine = np.asarray(image.meta["original_affine"], dtype=np.float64)
+    shape = tuple(int(s) for s in image.meta["spatial_shape"])
+    ref_affine = affine
+    if raw_spacing is not None:  # the pipeline replaced the affine with diag(raw_spacing) after reorientation
+        to_oriented = np.eye(4)
+        if orientd:
+            ornt = nib.orientations.ornt_transform(
+                nib.orientations.io_orientation(affine),
+                nib.orientations.axcodes2ornt(CT_MONAI_CONFIG["orientation"]),
+            )
+            to_oriented = np.linalg.inv(nib.orientations.inv_ornt_aff(ornt, shape))
+        ref_affine = np.diag([*map(float, raw_spacing), 1.0]) @ to_oriented
+    to_model = np.linalg.inv(np.asarray(image.affine, dtype=np.float64)) @ ref_affine
+    return {"shape": shape, "affine": affine, "to_model": to_model}
+
+
+def load_ct_volume(volume, orientd=True, raw_spacing=None, return_geometry=False):
     """Load and preprocess a CT volume.
 
-    volume : path to a NIfTI (.nii/.nii.gz) file, or a (240, 384, 384) tensor/ndarray in HU that is
-             already on the preprocessed grid (IPL, 1.5/0.75/0.75 mm resampled, cropped/padded, resized);
-             only the HU window is applied to arrays. Returns pixel_values of shape
-             (1, num_frames, 3, H, W), plus the 4x4 world affine of the preprocessed grid if
-             ``return_affine`` (None for array input).
+    volume : path to a NIfTI (.nii/.nii.gz) file or a DICOM series directory, or a (240, 384, 384)
+             tensor/ndarray in HU that is already on the preprocessed grid (IPL, 1.5/0.75/0.75 mm
+             resampled, cropped/padded, resized); only the HU window is applied to arrays. Returns
+             pixel_values of shape (1, num_frames, 3, H, W), plus the input grid if ``return_geometry``
+             (``shape``, ``affine`` (None for array input) and ``to_model``, its voxel mapping onto the
+             preprocessed grid), which ``similarity_scores_to_volume`` maps the scores back onto.
     """
-    affine = None
     if isinstance(volume, str):
         transform = build_ct_transform(orientd=orientd, raw_spacing=raw_spacing)
         out = transform({"image": volume})
         image = out["image"].squeeze(0)  # (1, D, H, W) -> (D, H, W)
-        affine = np.asarray(out["image"].affine, dtype=np.float64)
+        geometry = _input_geometry(out["image"], orientd, raw_spacing)
     else:
         image = torch.as_tensor(np.asarray(volume), dtype=torch.float32)
         if tuple(image.shape) != tuple(CT_MONAI_CONFIG["final_size"]):
@@ -106,45 +126,61 @@ def load_ct_volume(volume, orientd=True, raw_spacing=None, return_affine=False):
             )
         (a_min, a_max), (b_min, b_max) = CT_MONAI_CONFIG["hu_range"], CT_MONAI_CONFIG["normalize_range"]
         image = (image.clamp(a_min, a_max) - a_min) / (a_max - a_min) * (b_max - b_min) + b_min
+        geometry = {"shape": tuple(image.shape), "affine": None, "to_model": np.eye(4)}
 
     pixel_values_videos = normalize_pixel_values_videos(image).unsqueeze(0)  # (1, num_frames, 3, H, W)
-    return (pixel_values_videos, affine) if return_affine else pixel_values_videos
+    return (pixel_values_videos, geometry) if return_geometry else pixel_values_videos
 
 
-def similarity_scores_to_volume(similarity_scores, model, target_shape=CT_MONAI_CONFIG["final_size"]):
-    """Per-token scores ``[1, 1, num_tokens]`` -> (D, H, W) map on the preprocessed grid (trilinear)."""
+def _resample_to_input(volume, to_model, shape, chunk=32):
+    """Trilinearly sample a preprocessed-grid volume at every input voxel; 0 outside the preprocessing crop."""
+    if np.allclose(to_model, np.eye(4)) and tuple(shape) == tuple(volume.shape):
+        return volume
+    device = volume.device
+    M = torch.as_tensor(to_model, dtype=torch.float32, device=device)
+    size = torch.tensor(volume.shape, dtype=torch.float32, device=device)
+    ni, nj, nk = shape
+    jk = torch.stack(torch.meshgrid(torch.arange(nj), torch.arange(nk), indexing="ij"), -1).reshape(-1, 2).to(device)
+    out = torch.empty(tuple(shape), dtype=torch.float32, device=device)
+    for i0 in range(0, ni, chunk):
+        ii = torch.arange(i0, min(i0 + chunk, ni), device=device)
+        ijk = torch.cat([ii.repeat_interleave(len(jk))[:, None], jk.repeat(len(ii), 1)], 1).float()
+        src = ijk @ M[:3, :3].T + M[:3, 3]  # input voxel -> preprocessed voxel
+        grid = (2 * src / (size - 1) - 1).flip(-1)  # (d, h, w) -> grid_sample's (x, y, z)
+        out[i0 : i0 + len(ii)] = torch.nn.functional.grid_sample(
+            volume[None, None], grid.view(1, -1, 1, 1, 3), mode="bilinear", padding_mode="zeros", align_corners=True
+        ).view(len(ii), nj, nk)
+    return out
+
+
+def similarity_scores_to_volume(similarity_scores, model, geometry=None):
+    """Per-token scores ``[1, 1, num_tokens]`` -> similarity map on the input CT grid.
+
+    The token grid is upsampled (trilinear) to the 240x384x384 preprocessed grid and, with ``geometry`` from
+    ``load_ct_volume``, resampled onto the input voxels so the map has the input's shape and affine (0 outside
+    the 360 mm preprocessing crop). Without ``geometry`` the map stays on the preprocessed grid.
+    """
     vc = model.config.vision_config
     patch_size = getattr(vc, "patch_size", None) or 16
     frame_length = vc.num_channels * (getattr(vc, "tubelet_size", None) or 2)
     grid = (vc.input_slices // frame_length, vc.img_size // patch_size, vc.img_size // patch_size)
 
-    volume = similarity_scores.float().reshape(grid)
-    if target_shape is not None:
-        volume = torch.nn.functional.interpolate(
-            volume[None, None], size=tuple(target_shape), mode="trilinear", align_corners=False
-        )[0, 0]
-    return volume
-
+    similarity_map = similarity_scores.float().reshape(grid)
+    similarity_map = torch.nn.functional.interpolate(
+        similarity_map[None, None], size=CT_MONAI_CONFIG["final_size"], mode="trilinear", align_corners=False
+    )[0, 0]
+    if geometry is None:
+        return similarity_map
+    return _resample_to_input(similarity_map, geometry["to_model"], geometry["shape"])
 
-def save_similarity_map(similarity_volume, affine, save_path, pixel_values=None, ct_save_path=None):
-    """Save the similarity map (and optionally the preprocessed CT) as NIfTI in world space.
 
-    With the affine from ``load_ct_volume(..., return_affine=True)`` the map overlays the
-    original CT in ITK-SNAP / 3D Slicer. ``pixel_values`` + ``ct_save_path`` also save the
-    preprocessed volume (HU window [-1000, 1000]) on the same grid.
-    """
+def save_similarity_map(similarity_map, affine, save_path):
+    """Save the similarity map as NIfTI; with ``geometry["affine"]`` it overlays the input CT voxel for voxel."""
     import nibabel as nib
 
-    sim = similarity_volume.detach().float().cpu().numpy() if torch.is_tensor(similarity_volume) else similarity_volume
+    sim = similarity_map.detach().float().cpu().numpy() if torch.is_tensor(similarity_map) else similarity_map
     nib.save(nib.Nifti1Image(sim.astype(np.float32), affine), save_path)
 
-    if pixel_values is not None and ct_save_path is not None:
-        mean = torch.tensor([0.485, 0.456, 0.406]).view(1, 3, 1, 1)
-        std = torch.tensor([0.229, 0.224, 0.225]).view(1, 3, 1, 1)
-        frames = pixel_values[0].detach().float().cpu() * std + mean  # (num_frames, 3, H, W) in [0, 1]
-        hu = (frames.reshape(-1, *frames.shape[-2:]) * 2 - 1) * 1000.0
-        nib.save(nib.Nifti1Image(hu.numpy().astype(np.float32), affine), ct_save_path)
-
 
 def _as_pixel_values(volume, raw_spacing):
     if torch.is_tensor(volume) and volume.dim() == 5:  # already preprocessed by load_ct_volume
@@ -159,7 +195,7 @@ def model_inference(volume, text, model, tokenizer, raw_spacing=None):
     similarity_prob   : scalar sigmoid probability that the phrase applies.
     similarity_scores : per-token similarity in [0, 1] over the V-JEPA video tokens
                         (shape [1, 1, num_tokens]); ``similarity_scores_to_volume``
-                        maps it onto the preprocessed CT grid.
+                        maps it onto the input CT grid.
 
     volume may also be the ``pixel_values`` returned by ``load_ct_volume``.
     """