Add source files and streaming backup restore instructions

.gitattributes

@@ -39,3 +39,5 @@ runs/** filter=lfs diff=lfs merge=lfs -text
 wandb/** filter=lfs diff=lfs merge=lfs -text
 __pycache__/** filter=lfs diff=lfs merge=lfs -text
 dataset/** filter=lfs diff=lfs merge=lfs -text
+
+backup/parts/** filter=lfs diff=lfs merge=lfs -text

BACKUP_RESTORE.md

@@ -0,0 +1,27 @@
+# RadIR snapshot backup
+
+This repository displays the source code and stores the large data, model, output,
+cache, and log directories as a split `tar.zst` archive in Git LFS.
+
+- Source commit: `5b926d2179afa417784ba57459e86d2336d50a52`
+- Archived paths: `dataset, outputs, hf_models, hf_cache, runs, wandb, __pycache__`
+- Each volume is at most 64 MiB and has a SHA-256 checksum in `backup-manifest.json`.
+- This is a snapshot of committed files, not an export of all historical Git objects.
+- Original source files and the original local Git history were not modified.
+
+## Restore
+
+Wait until `backup-manifest.json` says `complete: true`. Download the repository
+with its LFS objects (not just pointer files). Install `zstd` and `tar`, then run:
+
+```bash
+python3 restore_backup.py
+```
+
+The restore tool verifies every volume, then streams decompression directly into
+`tar`. It does not create a second concatenated archive. Extracted contents keep
+their original relative paths. Source code is already available as regular files.
+
+The uploader uses `hf-mirror.com` for Git and LFS metadata. LFS payloads use the
+storage URLs returned by that service, as in a normal Git LFS transfer. It keeps
+only one archive volume in memory, plus small Git metadata and progress files.

Rad_IR/radir.egg-info/PKG-INFO

@@ -0,0 +1,26 @@
+Metadata-Version: 2.4
+Name: radir
+Version: 0.1
+Summary: RADIR
+Requires-Dist: beartype
+Requires-Dist: einops>=0.6
+Requires-Dist: ftfy
+Requires-Dist: regex
+Requires-Dist: XlsxWriter
+Requires-Dist: h5py
+Requires-Dist: matplotlib
+Requires-Dist: seaborn
+Requires-Dist: wilds
+Requires-Dist: click
+Requires-Dist: appdirs
+Requires-Dist: attr
+Requires-Dist: nltk
+Requires-Dist: transformers
+Requires-Dist: monai
+Requires-Dist: positional_encodings
+Requires-Dist: tensorboard
+Requires-Dist: pynvml
+Requires-Dist: wandb
+Requires-Dist: torchvision
+Dynamic: requires-dist
+Dynamic: summary

Rad_IR/radir.egg-info/SOURCES.txt

@@ -0,0 +1,10 @@
+setup.py
+radir/RadIR.py
+radir/__init__.py
+radir/lorentz.py
+radir/transformer_decoder.py
+radir.egg-info/PKG-INFO
+radir.egg-info/SOURCES.txt
+radir.egg-info/dependency_links.txt
+radir.egg-info/requires.txt
+radir.egg-info/top_level.txt

Rad_IR/radir.egg-info/dependency_links.txt

@@ -0,0 +1 @@
+

Rad_IR/radir.egg-info/requires.txt

@@ -0,0 +1,20 @@
+beartype
+einops>=0.6
+ftfy
+regex
+XlsxWriter
+h5py
+matplotlib
+seaborn
+wilds
+click
+appdirs
+attr
+nltk
+transformers
+monai
+positional_encodings
+tensorboard
+pynvml
+wandb
+torchvision

Rad_IR/radir.egg-info/top_level.txt

@@ -0,0 +1 @@
+radir

Rad_IR/radir/RadIR.py

@@ -0,0 +1,984 @@
+import math
+import copy
+from contextlib import contextmanager
+from functools import partial, wraps
+from pathlib import Path
+
+import torch
+import torch.nn.functional as F
+from torch import nn, einsum
+from torch.utils.checkpoint import checkpoint
+from einops import rearrange, repeat, reduce
+from einops.layers.torch import Rearrange, Reduce
+
+
+from transformers import BertTokenizer, BertModel
+from radir import lorentz as L
+
+from radir.data.hmclip_mimic_box import HyCoClipMimicBoxDataset_JPG, collate_hmclip_box
+from radir.data.conditional_dataset_train import SimpleImageTextDataset, simple_collate
+
+
+
+# NOTE: Import a Cross-Attn Module and PE Module
+from positional_encodings.torch_encodings import PositionalEncoding3D, PositionalEncoding1D, PositionalEncoding2D
+from radir.transformer_decoder import TransformerDecoder, TransformerDecoderLayer
+# from transformer_decoder import TransformerDecoder, TransformerDecoderLayer
+
+# NOTE
+
+# helper functions
+
+def identity(t, *args, **kwargs):
+    return t
+
+def exists(val):
+    return val is not None
+
+def default(val, d):
+    return val if exists(val) else d
+
+@contextmanager
+def null_context():
+    yield
+
+def max_neg_value(dtype):
+    return -torch.finfo(dtype).max
+
+def cast_tuple(t):
+    return t if isinstance(t, (tuple, list)) else (t,)
+
+def masked_mean(t, mask, dim = 1, eps = 1e-6):
+    t = t.masked_fill(~mask, 0.)
+    numer = t.sum(dim = dim)
+    denom = mask.sum(dim = dim).clamp(min = eps)
+    return numer / denom
+
+def log(t, eps = 1e-20):
+    return torch.log(t + eps)
+
+def l2norm(t):
+    
+    return F.normalize(t, dim = -1)
+
+def matrix_diag(t):
+    device = t.device
+    i, j = t.shape[-2:]
+    num_diag_el = min(i, j)
+    i_range = torch.arange(i, device = device)
+    j_range = torch.arange(j, device = device)
+    diag_mask = rearrange(i_range, 'i -> i 1') == rearrange(j_range, 'j -> 1 j')
+    diag_el = t.masked_select(diag_mask)
+    return rearrange(diag_el, '(b d) -> b d', d = num_diag_el)
+
+# checkpointing helper function
+
+def make_checkpointable(fn):
+    @wraps(fn)
+    def inner(*args):
+        input_needs_grad = any([isinstance(el, torch.Tensor) and el.requires_grad for el in args])
+
+        if not input_needs_grad:
+            return fn(*args)
+
+        return checkpoint(fn, *args)
+
+    return inner
+
+# keyword argument helpers
+
+def pick_and_pop(keys, d):
+    values = list(map(lambda key: d.pop(key), keys))
+    return dict(zip(keys, values))
+
+def group_dict_by_key(cond, d):
+    return_val = [dict(),dict()]
+    for key in d.keys():
+        match = bool(cond(key))
+        ind = int(not match)
+        return_val[ind][key] = d[key]
+    return (*return_val,)
+
+def string_begins_with(prefix, str):
+    return str.startswith(prefix)
+
+def group_by_key_prefix(prefix, d):
+    return group_dict_by_key(partial(string_begins_with, prefix), d)
+
+def groupby_prefix_and_trim(prefix, d):
+    with_prefix, without_prefix = group_dict_by_key(partial(string_begins_with, prefix), d)
+    trimmed = {k[len(prefix):]: v for k, v in with_prefix.items()}
+    return trimmed, without_prefix
+
+
+def safe_normalize(feat, max_norm=10.0, eps=1e-6):
+    # 支持 (B,D) 或 (B,1,D)
+    orig_shape = feat.shape
+    feat = feat.reshape(-1, feat.size(-1))
+
+    norm = feat.norm(dim=-1, keepdim=True)  # (B,1)
+    scale = (max_norm / (norm + eps)).clamp(max=1.0)
+    feat = feat * scale
+
+    return feat.reshape(orig_shape)
+
+# main clip class
+
+class RADIR(nn.Module):
+    def __init__(
+            self,
+            image_encoder = None,
+            text_encoder = None,
+            tokenizer = None,
+            fusion_module = None,   # NOTE: Pass in a Transformer Decoder Module
+            dim_text = 768,
+            dim_image = 768,
+            dim_latent = 768,
+            text_seq_len = 256,
+            text_has_cls_token = False,
+            text_pad_id = 0,
+            text_causal_mask = False,
+            text_eos_id = None,
+            text_encode_without_mask = False,
+            visual_image_size = 224,
+            visual_has_cls_token = False,
+            channels = 3,
+            use_all_token_embeds = False,
+            decoupled_contrastive_learning = False,
+            use_mlm = False,
+            text_ssl_loss_weight = 0.05,
+            use_visual_ssl = False,
+            visual_ssl = None,
+            embed_dim = 512,
+            image_ssl_loss_weight = 0.05,
+            multiview_loss_weight = 0.1,
+            curv_init: float = 1.0,
+            learn_curv: bool = True,
+            entail_weight: float = 0.3,
+            local_aperture_thresh: float = 1.2,
+            global_aperture_thresh: float = 0.7,
+            use_repo_ce_loss: bool = True
+    ):
+        super().__init__()
+        
+
+
+
+        #assert use_all_token_embeds or (visual_has_cls_token or text_has_cls_token), 'CLS token must be included on both vision and text transformers if you are not using fine-grained contrastive learning loss'
+
+        # store some parameters for access
+
+        self.dim_text = dim_text
+        self.dim_image = dim_image
+        self.dim_latent = dim_latent
+
+        self.image_size = visual_image_size
+
+        # instantiate text transformer
+
+        self.text_pad_id = text_pad_id
+        self.text_has_cls_token = text_has_cls_token
+        self.text_seq_len = text_seq_len
+
+        self.text_encode_without_mask = text_encode_without_mask # whether to pass in text mask to text encoder
+
+        self.text_causal_mask = text_causal_mask
+        self.text_eos_id = text_eos_id
+
+        assert not (text_causal_mask and not exists(text_eos_id)), 'text EOS token id must be given if using causal mask in text transformer'
+
+
+        self.text_transformer = text_encoder
+
+
+        # instantiate image transformer
+
+        self.visual_has_cls_token = visual_has_cls_token
+
+        self.visual_transformer = image_encoder
+
+        # NOTE: Pass in a Transformer Decoder Module and a Positional-Encoding Module
+        if exists(fusion_module):
+            self.fusion_module = fusion_module
+        else:
+            decoder_layer = TransformerDecoderLayer(d_model=dim_latent, nhead=8, normalize_before=True)
+            decoder_norm = nn.LayerNorm(dim_latent)
+            self.fusion_module = TransformerDecoder(decoder_layer=decoder_layer, num_layers=4, norm=decoder_norm)
+
+        self.pos_embedding = PositionalEncoding2D(dim_latent)(torch.zeros(1, 24, 24, dim_latent)) # 1 H W Dim
+        # 修改形状为1 (HW) Dim
+        self.pos_embedding = rearrange(self.pos_embedding, 'b h w d -> b (h w) d')
+        
+        # NOTE
+
+        # text ssl
+
+        self.use_mlm = use_mlm
+        self.text_ssl_loss_weight = text_ssl_loss_weight if use_mlm else 0
+
+        # image ssl
+
+        self.use_visual_ssl = use_visual_ssl or exists(visual_ssl)
+        self.image_ssl_loss_weight = image_ssl_loss_weight if use_visual_ssl else 0
+
+
+        self.to_text_latent = nn.Linear(dim_text, dim_latent, bias = False)
+
+
+        self.to_visual_latent = nn.Linear(dim_image, dim_latent, bias = False)
+        self.to_fused_latent = nn.Linear(dim_latent, dim_latent, bias = False)
+
+        # Initialize curvature parameter. Hyperboloid curvature will be `-curv`.
+        self.curv = nn.Parameter(
+            torch.tensor(curv_init).log(), requires_grad=learn_curv
+        )
+        
+        # When learning the curvature parameter, restrict it in this interval to
+        # prevent training instability.
+        self._curv_minmax = {
+            "max": math.log(curv_init * 10),
+            "min": math.log(curv_init / 10),
+        }
+
+        # temperature
+
+        self.temperature = nn.Parameter(torch.tensor(1.))
+        self.use_repo_ce_loss = use_repo_ce_loss
+        self.logit_scale = nn.Parameter(torch.tensor(1 / 0.07).log())  # 与HycoClip一致的可学习温度
+        self.entail_weight = entail_weight
+        self.global_aperture_thresh = global_aperture_thresh
+        self.local_aperture_thresh = local_aperture_thresh
+        self.visual_alpha = nn.Parameter(torch.tensor(embed_dim**-0.5).log())
+        self.textual_alpha = nn.Parameter(torch.tensor(embed_dim**-0.5).log())
+
+
+        # from https://arxiv.org/abs/2111.07783 (FILIP paper)
+        self.use_all_token_embeds = use_all_token_embeds
+
+        # proposed in https://arxiv.org/abs/2110.06848 (DCL) and https://arxiv.org/abs/2110.11316 (CLOOB)
+        self.decoupled_contrastive_learning = decoupled_contrastive_learning
+
+        self.multiview_loss_weight = multiview_loss_weight
+
+        if tokenizer != None:
+            self.tokenizer=tokenizer
+        else:
+            raise ValueError('Tokenzier is not defined.')
+        
+
+    def load(self, path,allow_partial_load = True):
+        print(f'** MODEL ** Load Checkpoint from {path}')
+        path = Path(path)
+        assert path.exists()
+        pkg = torch.load(path, map_location='cpu', weights_only=False)
+        model_dict = self.state_dict()
+        if 'model' in pkg:
+            pkg_state_dict = pkg['model']  # 假设 pkg['model'] 是模型的状态字典
+        else:
+            pkg_state_dict = pkg
+        if allow_partial_load:
+            # 检查差异
+            unexpected_state_dict = [k for k in pkg_state_dict.keys() if k not in model_dict.keys()]
+            missing_state_dict = [k for k in model_dict.keys() if k not in pkg_state_dict.keys()]
+            unmatchd_state_dict = [k for k, v in pkg_state_dict.items() if k in model_dict.keys() and v.shape != model_dict[k].shape]
+            # 加载部分参数
+            state_dict = {k: v for k, v in pkg_state_dict.items() if k in model_dict.keys() and v.shape == model_dict[k].shape}
+            model_dict.update(state_dict)
+            self.load_state_dict(model_dict)
+            # 可以选择输出具体的参数情况
+            # print('** MODEL ** The following parameters are UNEXPECTED in checkpoint:\n')
+            # print(unexpected_state_dict)
+            # print('** MODEL ** The following parameters are MISSING in checkpoint:\n')
+            # print(missing_state_dict)
+            # print('** MODEL ** The following parameters have DIFFERENT SHAPES in checkpoint:\n')
+            # print(unmatchd_state_dict)
+            # ('** MODEL ** The following parameters are LOADED in:\n')
+            # print(state_dict.keys())
+        else:
+            self.load_state_dict(pkg_state_dict)
+  
+    def tokenize(self, prompt):
+        text_tokens=self.tokenizer(prompt, return_tensors="pt", padding="max_length", truncation=True, max_length=512).to('cuda')
+        return text_tokens
+    
+    def token_embedding(self,input_ids):
+        input_shape = input_ids.size()
+        batch_size, seq_length = input_shape
+        if hasattr(self.text_transformer.embeddings, "token_type_ids"):
+            print("hahatrue")
+
+        buffered_token_type_ids = self.text_transformer.embeddings.token_type_ids[:, :seq_length]
+        buffered_token_type_ids_expanded = buffered_token_type_ids.expand(batch_size, seq_length)
+        token_type_ids = buffered_token_type_ids_expanded
+        text_embeddings = self.text_transformer.embeddings(input_ids = input_ids, token_type_ids = token_type_ids)
+        return text_embeddings
+
+    def freeze_text_encoder(self):
+        for param in self.text_transformer.parameters():
+            param.requires_grad = False
+        for param in self.to_text_latent.parameters():
+            param.requires_grad = False
+            
+    def open_text_encoder(self):
+        for param in self.text_transformer.parameters():
+            param.requires_grad = True
+        for param in self.to_text_latent.parameters():
+            param.requires_grad = True
+    
+    def freeze_vision_encoder(self):
+        for param in self.visual_transformer.parameters():
+            param.requires_grad = False
+        for param in self.to_visual_latent.parameters():
+            param.requires_grad = False
+    
+    def open_vision_encoder(self):
+        for param in self.visual_transformer.parameters():
+            param.requires_grad = True
+        for param in self.to_visual_latent.parameters():
+            param.requires_grad = True
+    
+    def open_fusion_module(self):
+        for param in self.fusion_module.parameters():
+            param.requires_grad = True
+        for param in self.to_fused_latent.parameters():
+            param.requires_grad = True
+            
+    def freeze_fusion_module(self):
+        for param in self.fusion_module.parameters():
+            param.requires_grad = False
+        for param in self.to_fused_latent.parameters():
+            param.requires_grad = False
+
+    def get_log_curv(self):
+        return torch.clamp(
+            self.curv,
+            min=self._curv_minmax["min"],
+            max=self._curv_minmax["max"],
+        )
+
+    def get_curv(self):
+        with torch.cuda.amp.autocast(enabled=False):
+            curv = torch.exp(self.get_log_curv().float())
+            return torch.clamp(curv, min=1e-4)
+
+    def project_img(self, image_feats):
+        image_feats = image_feats * self.visual_alpha.exp()
+        with torch.cuda.amp.autocast(enabled=False):
+            curv = self.get_curv()
+            image_feats = L.exp_map0(image_feats.float(), curv)
+        return image_feats
+
+    def project_txt(self, text_feats):
+        text_feats = text_feats * self.textual_alpha.exp()
+        with torch.cuda.amp.autocast(enabled=False):
+            curv = self.get_curv()
+            text_feats = L.exp_map0(text_feats.float(), curv)
+        return text_feats
+
+
+
+    def mask_average_pool(self, x, mask):
+        """
+        对 token 特征做 masked average pooling
+
+        Args:
+            x:    [B, T, D]          (例如 word_feats: [8,128,512])
+            mask: [B, Tm] 或 [B, K, Tm]
+                目前的情况通常是 [B,111]
+
+        Returns:
+            pooled: [B, K, D]
+        """
+        import torch
+        import torch.nn.functional as F
+
+        # x: [B, T, D]
+        T = x.size(1)
+
+        m = mask
+        # [B, Tm] -> [B, 1, Tm]
+        if m.dim() == 2:
+            m = m.unsqueeze(1)
+
+        # 对齐最后一维到 T（补 0 或截断）
+        Tm = m.size(-1)
+        if Tm < T:
+            m = F.pad(m, (0, T - Tm), value=0)   # 右侧补0
+        elif Tm > T:
+            m = m[..., :T]
+
+        # einsum 用 float dtype
+        m = m.to(dtype=x.dtype)
+
+        # 避免全0 mask 导致除0
+        den = m.sum(dim=-1, keepdim=True).clamp_min(1.0)
+
+        # [B, K, T] x [B, T, D] -> [B, K, D]
+        return torch.einsum('bkt,btd->bkd', m, x) / den
+
+    def _unpack_text_batch(self, text):
+        if isinstance(text, dict):
+            input_ids = text.get('caption_ids', text.get('input_ids'))
+            token_type_ids = text.get('token_type_ids', None)
+            attention_mask = text.get('attention_mask', None)
+        else:
+            input_ids = text.input_ids
+            token_type_ids = getattr(text, 'token_type_ids', None)
+            attention_mask = getattr(text, 'attention_mask', None)
+
+        return input_ids, token_type_ids, attention_mask
+
+
+    def forward(
+            self,
+            text,   # B
+            image,  # B C H W D
+            device,
+            return_loss = True,
+            is_condition = False,
+            modal_embedding = False,
+            phrases_bboxes_mask=None,
+            phrase_text_mask=None,
+            modal_indexs = None,
+            return_latents=True,  # 🟢 显式声明
+            **kwargs
+    ):
+        if is_condition:
+
+            
+            B, device = text.input_ids.shape[0], device
+
+            # derive text mask
+
+            text_mask =text.attention_mask
+
+            # get encoded text
+
+            text_args = (text.input_ids,text.attention_mask)
+
+            if not self.text_encode_without_mask:
+                text_args = (*text_args, text_mask)
+
+            text_embeddings = self.text_transformer(text.input_ids, attention_mask = text.attention_mask )
+            enc_text = text_embeddings[0]   # B 512(token_len) 768
+            
+
+            if image.ndim == 3:                 # [H, W, ?] 的情况先 unsqueeze
+                image = image.unsqueeze(0)
+            if image.ndim == 4:                 # B×C×H×W
+                if image.shape[1] != 1:
+                    image = image.mean(dim=1, keepdim=True)
+            elif image.ndim == 5:               # B×C×D×H×W
+                if image.shape[1] != 1:
+                    image = image.mean(dim=1, keepdim=True)
+            else:
+                raise ValueError(...)
+
+            
+
+
+    
+            # NOTE
+            modal_indexs = modal_indexs.view(-1)  # (B * local_B) 重复之后，condition就能和uncon的情况一样了
+            enc_image = self.visual_transformer(image, return_encoded_tokens=True,modal_embedding=modal_embedding,modal_indexs=modal_indexs,is_condition=True)   # B*N 24 24 24 512(vis_dim)
+
+
+            image_embeds = rearrange(enc_image, 'b t d -> t b d')  
+            
+
+            text_embeds = enc_text[:,0,:] # B 768
+
+            text_latents = self.to_text_latent(text_embeds).unsqueeze(0) # 1 B*N 512
+            image_latents = image_embeds
+            t = image_latents.shape[0]
+            temp_pos_embedding = self.pos_embedding[:,:t, :]  # 1 t 512
+
+            pos = repeat(temp_pos_embedding, 'one d dim -> d (one bn) dim', bn=image_latents.shape[1]).to(device)  # 24*24, B*N, 512
+            
+            fused_latents, _ = self.fusion_module(text_latents, image_latents, pos=pos) # 1 B*N 512
+
+
+            fused_latents = fused_latents[0]
+            fused_latents = self.to_fused_latent(fused_latents) # B*N 512
+            
+            norms = torch.norm(fused_latents, dim=-1)
+            # print('fused_latents范数:', '最小值:', norms.min(), '最大值:', norms.max(), '存在零:', (norms == 0).any())
+            if (norms < 1e-6).any():
+                print('警告：检测到极小范数')
+                # 可选：添加微小噪声避免零范数问题
+                fused_latents = fused_latents + torch.randn_like(fused_latents) * 1e-8
+            if torch.isnan(fused_latents).any() or torch.isinf(fused_latents).any():
+                print('NaN/Inf in vec **before** l2norm')
+            
+            if fused_latents.isnan().any() or fused_latents.isinf().any():
+                print('NaN/Inf in vec **after** l2norm')
+            if (fused_latents == 0).any():
+                print('Zero in vec **after** l2norm')
+
+            # NOTE
+            
+            if return_latents:
+                return text_latents, image_latents, fused_latents, self.temperature.exp()
+
+            # ——保持与 else 分支一致 ——
+            
+            # 全局文本/图像向量
+            text_latents_global = l2norm(self.to_text_latent(enc_text[:, 0, :]))     # B x D_latent
+            image_global       = l2norm(fused_latents.mean(dim=1))                   # B x D_latent
+
+            # 双曲投影与全局蕴含
+            #_curv       = self.curv.exp()
+            with torch.cuda.amp.autocast(enabled=False):
+                _curv = self.get_curv()
+            text_feats  = self.project_txt(text_latents_global)
+            image_feats = self.project_img(image_global)
+            _angle      = L.oxy_angle(text_feats, image_feats, _curv)
+            _aperture   = L.half_aperture(text_feats, _curv)
+            text_image_entailment_loss = torch.clamp(_angle - self.global_aperture_thresh * _aperture, min=0).mean()
+
+            # 温度缩放（与仓库版相同）
+            
+            #改curv1
+            #self.logit_scale.data = torch.clamp(self.logit_scale.data, max=math.log(100.0))
+            with torch.no_grad():
+                self.logit_scale.clamp_(max=math.log(100.0))
+            #改curv2
+            
+            scale = self.logit_scale.exp()
+
+            # B×B logits（-双曲距离），对称 CE
+            targets     = torch.arange(image_global.size(0), device=image_global.device)
+            image_logits = -L.pairwise_dist(image_feats, text_feats, _curv)  # [B, B]
+            text_logits  = -L.pairwise_dist(text_feats, image_feats, _curv)  # [B, B]
+            # ABLATION: 去掉 image→text CE，只保留 text→image CE
+            i2t_ce_loss = F.cross_entropy(scale * image_logits, targets)
+            t2i_ce_loss = F.cross_entropy(scale * text_logits, targets)
+            ce_loss = t2i_ce_loss
+            # ce_loss = 0.5 * (i2t_ce_loss + t2i_ce_loss)
+
+            loss = ce_loss + self.entail_weight * text_image_entailment_loss
+
+            entailment_loss = text_image_entailment_loss
+
+            return loss, ce_loss, entailment_loss, text_image_entailment_loss
+        else:
+            #text_mask：文本的注意力掩码（布尔矩阵），形状是 B x seq_len。
+            #text_args：传入文本编码器的集合，包括输入的 token ids 和注意力掩码。
+            #text_embeddings：文本编码器的输出，包含每个 token 的向量表示。形状是 B x seq_len x dim_text。
+            #enc_text：文本编码器的最终输出，通常是取 [CLS] token 的向量表示。形状是 B x seq_len x dim_text。
+            #text_latents：把整句输出文本压成一个固定长度的向量，代表这句话的整体语义。形状是 B x 512。
+            #text_embeds: 文本的全局向量表示，取自文本编码器的 [CLS] token 输出。形状是 B x dim_text。
+
+            #local_B： 单个图像的局部块数量（local_B）。
+            #modal_indexs：表示每个图像块所属的模态索引，形状是 B x local_B。
+            #enc_image：图像编码器的输出，包含每个图像块的向量表示。形状是 (B * local_B) x dim_image。
+            #image_embeds：把所有图像块的向量表示重新排列成 (HWD) x (B * local_B) x dim_image 的形状。
+            #image_latents：图像块向量，这些局部向量按顺序排成一串，形状是 (HWD) x (B * local_B) x 512。
+            #box_image_feats_eu：池化后与文本短语对应的图像块特征，形状是 B x K x D_latent。
+            #box_text_feats_eu：池化后与图像区域对应的文本短语特征，形状是 B x K x D_latent。
+            
+
+            input_ids, token_type_ids, attention_mask = self._unpack_text_batch(text)
+            B = input_ids.size(0)
+
+
+            with torch.cuda.amp.autocast(enabled=False):
+                _curv = self.get_curv()
+
+            if image.ndim == 3:                 # [H, W, ?] 的情况先 unsqueeze
+                image = image.unsqueeze(0)
+            if image.ndim == 4:                 # B×C×H×W
+                if image.shape[1] != 1:
+                    image = image.mean(dim=1, keepdim=True)
+            elif image.ndim == 5:               # B×C×D×H×W
+                if image.shape[1] != 1:
+                    image = image.mean(dim=1, keepdim=True)
+            else:
+                raise ValueError(...)
+
+
+
+            # ====== 文本编码 ======
+            text_embeddings = self.text_transformer(
+                input_ids=input_ids,
+                token_type_ids=token_type_ids,
+                attention_mask=attention_mask
+            )
+            enc_text = text_embeddings[0]  # B×T_txt×dim_text
+            
+            #text_latents_global = self.to_text_latent(enc_text[:, 0, :])
+            text_latents_global = enc_text[:, 0, :]      #  不再过 Linear
+            
+            modality_dict = {'CT': 0, 'CXR': 1}
+            modal_indexs = torch.tensor([modality_dict['CXR']], dtype=torch.long, device=device)
+
+            modal_indexs = modal_indexs.view(-1)  # (B * local_B) 重复之后，condition就能和uncon的情况一样了
+
+            # ====== 图像编码（双路）======
+            # Patch 路：box / cross entailment（保留空间 token 网格）
+            enc_image_patch = self.visual_transformer(
+                image,
+                return_encoded_tokens=True,
+                modal_embedding=modal_embedding,
+                modal_indexs=None,
+                is_condition=True,
+            )  # [B, T_img, D]
+
+            # Global 路：global CE + global entailment（enc_cls_transformer 整图 CLS）
+            enc_image_global = self.visual_transformer(
+                image,
+                return_encoded_tokens=True,
+                modal_embedding=modal_embedding,
+                modal_indexs=None,
+                is_condition=False,
+            )  # [B, D]
+
+            if self.use_all_token_embeds:
+                text_embeds = enc_text[:, 1:] if self.text_has_cls_token else enc_text
+            else:
+                text_embeds = enc_text[:, :] if enc_text.ndim == 3 else enc_text
+
+            text_embeds = text_embeds[:, 0, :]
+            text_latents = l2norm(text_embeds)
+
+            if enc_image_global.dim() == 2:
+                image_latents = l2norm(enc_image_global)
+            elif enc_image_global.dim() == 3:
+                image_latents = l2norm(enc_image_global[:, 0, :])
+            else:
+                raise ValueError(
+                    f"Unexpected global visual feature shape: {enc_image_global.shape}"
+                )
+
+            image_global = image_latents
+
+            word_feats = enc_text
+            patch_feats = enc_image_patch
+
+            if phrase_text_mask is not None:
+                if phrase_text_mask.dim() == 2:  # B×T_txt
+                    phrase_text_mask = phrase_text_mask.unsqueeze(1)  # B×1×T_txt
+                box_text_feats_eu = self.mask_average_pool(word_feats, phrase_text_mask)  # B×K×dim_latent
+
+            #patch_feats = rearrange(patch_feats_bn, '(b n) t d -> b n t d', n=local_B).mean(dim=1)  # B×T_img×dim_latent
+            if phrases_bboxes_mask is not None:
+                if phrases_bboxes_mask.dim() == 2:
+                    phrases_bboxes_mask = phrases_bboxes_mask.unsqueeze(1)  # B×1×T_img
+                box_image_feats_eu = self.mask_average_pool(patch_feats, phrases_bboxes_mask)
+
+            # 双曲投影 + 蕴含损失计算（仅当提供 box mask 时）
+            if phrase_text_mask is not None and phrases_bboxes_mask is not None:
+                text_feats = self.project_txt(text_latents_global) # B×D_hyp
+                #image_global = patch_feats.mean(dim=1)          # B×D_latent
+                image_feats = self.project_img(image_global)                # B×D_hyp
+                box_text_feats = self.project_txt(box_text_feats_eu)        # B×K×D_hyp
+                box_image_feats = self.project_img(box_image_feats_eu)      # B×K×D_hyp
+                # 后续四项蕴含损失逻辑保持不变
+
+
+
+            if not return_latents:
+                return text_latents, image_latents, self.temperature.exp()
+
+
+            #celoss用，防止被normalize拉爆
+            image_logits = -L.pairwise_dist(image_feats, text_feats, _curv) # [B, B]
+            text_logits = -L.pairwise_dist(text_feats, image_feats, _curv) # [B, B]
+
+            
+
+            #四项蕴含损失
+            _angle = L.oxy_angle(text_feats, image_feats, _curv)
+
+            _aperture = L.half_aperture(text_feats, _curv)
+
+            _box_angle = L.oxy_angle(box_text_feats, box_image_feats, _curv)
+
+            _box_aperture = L.half_aperture(box_text_feats, _curv)
+
+            _cross_image_angle = L.oxy_angle(box_image_feats, image_feats.unsqueeze(1), _curv)
+
+            _box_image_aperture = L.half_aperture(box_image_feats, _curv)
+
+            _cross_text_angle = L.oxy_angle(box_text_feats, text_feats.unsqueeze(1), _curv)
+
+            _box_text_aperture = L.half_aperture(box_text_feats, _curv)
+
+
+
+
+            #全文-全图蕴含
+            text_image_entailment_loss = torch.clamp(_angle - self.global_aperture_thresh * _aperture, min=0).mean()
+
+            box_text_image_entailment_loss = torch.clamp(_box_angle - self.global_aperture_thresh * _box_aperture, min=0).mean()
+
+            cross_image_entailment_loss = torch.clamp(_cross_image_angle - self.local_aperture_thresh * _box_image_aperture, min=0).mean()
+    
+            cross_text_entailment_loss = torch.clamp(_cross_text_angle - self.local_aperture_thresh * _box_text_aperture, min=0).mean()
+
+            entailment_loss = 0.5 * (
+                text_image_entailment_loss
+                + box_text_image_entailment_loss
+                + cross_image_entailment_loss
+                + cross_text_entailment_loss
+            )
+
+
+
+            
+            # 3) 温度缩放
+            
+            self.logit_scale.data = torch.clamp(self.logit_scale.data, max=4.6052)
+
+
+            targets = torch.arange(image_global.size(0), device=image_global.device)
+            scale = self.logit_scale.exp()
+
+            # ABLATION: 去掉 image→text CE，只保留 text→image CE
+            i2t_ce_loss = F.cross_entropy(scale * image_logits, targets)
+            t2i_ce_loss = F.cross_entropy(scale * text_logits, targets)
+            ce_loss = t2i_ce_loss
+            # ce_loss = 0.5 * (i2t_ce_loss + t2i_ce_loss)
+
+            loss = ce_loss + self.entail_weight * entailment_loss
+
+            return loss, ce_loss, entailment_loss, text_image_entailment_loss
+                   
+
+
+if __name__ == '__main__':
+    import os
+    import argparse
+    import torch
+
+    from transformer_maskgit.ctvit import CTViT
+    from transformers import AutoTokenizer, AutoModel, AutoConfig
+
+    # --------------------------------------------------------
+    # Args
+    # --------------------------------------------------------
+    parser = argparse.ArgumentParser(description='Debug RADIR with batch')
+
+    parser.add_argument('--batch', type=int, default=8)
+    parser.add_argument('--channels', type=int, default=1)
+    parser.add_argument('--height', type=int, default=224)
+    parser.add_argument('--width', type=int, default=224)
+    parser.add_argument('--modal_embedding', action='store_true')
+
+    # Dataset paths
+    parser.add_argument(
+        '--dataset_root',
+        type=str,
+        default=r"/mnt/data-3/users/zhangzeyu/HMClip/RadIR/dataset"
+    )
+    parser.add_argument(
+        '--matched_dir',
+        type=str,
+        default=r"/mnt/data-3/users/zhangzeyu/HMClip/RadIR/dataset/processed/matched"
+    )
+    parser.add_argument(
+        '--caption_file',
+        type=str,
+        default="/mnt/data-3/users/zhangzeyu/HMClip/RadIR/dataset/all.jsonl"   # 相对 dataset_root
+    )
+
+    # Local HF model path
+    parser.add_argument(
+        '--model_path',
+        type=str,
+        default=r"/mnt/data-3/users/zhangzeyu/HMClip/RadIR/hf_models/microsoft/BiomedVLP-CXR-BERT-specialized"
+    )
+
+    args = parser.parse_args()
+
+    # --------------------------------------------------------
+    # Device
+    # --------------------------------------------------------
+    if not torch.cuda.is_available():
+        raise RuntimeError('需要 GPU 环境（CUDA 可用）。')
+
+    device = torch.device('cuda')
+
+    # --------------------------------------------------------
+    # Debug HF loading
+    # --------------------------------------------------------
+    from transformers.utils import cached_file
+    try:
+        path = cached_file(args.model_path, "tokenizer_config.json", local_files_only=True)
+        print(">>> HF tokenizer_config.json loaded from:", path)
+    except Exception as e:
+        print(">>> Failed to load tokenizer_config.json:", e)
+
+    print(">>> Checking config loading...")
+    cfg = AutoConfig.from_pretrained(
+        args.model_path,
+        trust_remote_code=True,
+        local_files_only=True
+    )
+    print(">>> model_type:", cfg.model_type)
+    print(">>> config path:", cfg._name_or_path)
+
+    # --------------------------------------------------------
+    # Load text encoder & tokenizer
+    # --------------------------------------------------------
+    tokenizer = AutoTokenizer.from_pretrained(
+        args.model_path,
+        trust_remote_code=True,
+        local_files_only=True
+    )
+
+    text_model = AutoModel.from_pretrained(
+        args.model_path,
+        trust_remote_code=True,
+        local_files_only=True
+    ).to(device)
+
+    # --------------------------------------------------------
+    # Load image encoder (CTViT)
+    # --------------------------------------------------------
+    image_encoder = CTViT(
+        dim=512,
+        codebook_size=8192,
+        image_size=args.height,
+        patch_size=16,
+        temporal_patch_size=10,
+        spatial_depth=8,
+        temporal_depth=6,
+        cls_depth=4,
+        dim_head=32,
+        heads=8,
+        channels=args.channels
+    ).to(device)
+
+    # --------------------------------------------------------
+    # Build RADIR
+    # --------------------------------------------------------
+    from radir.RadIR import RADIR
+
+    clip = RADIR(
+        tokenizer=tokenizer,
+        image_encoder=image_encoder,
+        text_encoder=text_model,
+        dim_text=768,
+        dim_image=512,
+        dim_latent=512,
+        use_mlm=False,
+        use_all_token_embeds=False
+    ).to(device)
+
+    # --------------------------------------------------------
+    # Build dataset
+    # --------------------------------------------------------
+    from radir.data.hmclip_mimic_box import (
+        HyCoClipMimicBoxDataset_JPG,
+        collate_hmclip_box
+    )
+    from torch.utils.data import DataLoader
+
+    dataset = HyCoClipMimicBoxDataset_JPG(
+        split="train",
+        tokenizer=tokenizer,
+        dataset_root=args.dataset_root,
+        cig_matched_dir=args.matched_dir,
+        caption_file=args.caption_file,
+        transform=None,     # 使用 Dataset 内置 2D + MONAI
+        max_words=128,
+        imsize=224
+    )
+
+    loader = DataLoader(
+        dataset,
+        batch_size=args.batch,
+        shuffle=True,
+        num_workers=0,
+        collate_fn=collate_hmclip_box,
+    )
+
+    batch = next(iter(loader))
+
+    imgs = batch["imgs"].to(device)  # [B, 1, 224, 224]
+
+    enc_text = {
+        "input_ids": batch["caption_ids"].to(device),
+        "token_type_ids": batch["token_type_ids"].to(device),
+        "attention_mask": batch["attention_mask"].to(device),
+    }
+
+    phrase_text_mask = batch["phrase_text_mask"].to(device)
+    phrases_bboxes_mask = batch["phrases_bboxes_mask"].to(device)
+
+    print(f">>> Loaded batch:")
+    print(f"    imgs: {imgs.shape}")
+    print(f"    input_ids: {enc_text['input_ids'].shape}")
+    print(f"    phrase_text_mask: {phrase_text_mask.shape}")
+    print(f"    phrases_bboxes_mask: {phrases_bboxes_mask.shape}")
+
+    # --------------------------------------------------------
+    # Estimate visual token length T_img
+    # --------------------------------------------------------
+    with torch.no_grad():
+        enc_image_tokens = clip.visual_transformer(
+            imgs,
+            return_encoded_tokens=True,
+            modal_embedding=args.modal_embedding,
+            modal_indexs=None,
+            is_condition=False
+        )
+        T_img = enc_image_tokens.shape[1]
+        print(">>> T_img =", T_img)
+
+    # --------------------------------------------------------
+    # Adjust masks
+    # --------------------------------------------------------
+    phrase_text_mask = phrase_text_mask.unsqueeze(1)         # [B,1,T_txt]
+    phrases_bboxes_mask = phrases_bboxes_mask.unsqueeze(1)   # [B,1,T_img?]
+
+    cur_T_img = phrases_bboxes_mask.shape[-1]
+    if cur_T_img != T_img:
+        if cur_T_img > T_img:
+            phrases_bboxes_mask = phrases_bboxes_mask[..., :T_img]
+        else:
+            pad = torch.ones(
+                imgs.size(0), 1, T_img - cur_T_img,
+                device=device
+            )
+            phrases_bboxes_mask = torch.cat(
+                [phrases_bboxes_mask, pad],
+                dim=-1
+            )
+        print(f"[DBG] adjusted phrases_bboxes_mask: {cur_T_img} -> {T_img}")
+
+    # --------------------------------------------------------
+    # Modality index (CXR = 1)
+    # --------------------------------------------------------
+    modality_dict = {'CT': 0, 'CXR': 1}
+    modal_indexs = torch.tensor(
+        [modality_dict['CXR']],
+        dtype=torch.long,
+        device=device
+    )
+
+    # --------------------------------------------------------
+    # Forward
+    # --------------------------------------------------------
+    print(">>> Forward (return_loss=True)")
+
+    losses = clip(
+        text=enc_text,
+        image=imgs,
+        device=device,
+        is_condition=False,
+        return_latents=True,
+        return_loss=True,
+        modal_embedding=args.modal_embedding,
+        phrase_text_mask=phrase_text_mask,
+        phrases_bboxes_mask=phrases_bboxes_mask,
+        modal_indexs=modal_indexs
+    )
+
+    loss, ce_loss, entailment_loss, text_image_entailment_loss = losses
+
+    print(f"loss_total: {float(loss):.6f}")
+    print(f"loss_ce(image<->text): {float(ce_loss):.6f}")
+    print(f"loss_entail_global: {float(entailment_loss):.6f}")
+    print(f"loss_entail_image_to_text: {float(text_image_entailment_loss):.6f}")

Rad_IR/radir/__init__.py

@@ -0,0 +1 @@
+from radir.RadIR import RADIR

Rad_IR/radir/data/conditional_dataset_train.py

@@ -0,0 +1,141 @@
+import os
+import torch
+import json
+import numpy as np
+from PIL import Image
+from torch.utils.data import Dataset
+from torchvision import transforms
+
+
+class SimpleImageTextDataset(Dataset):
+
+    def __init__(
+        self,
+        data_source,          # ✅ 可以是 CSV / list / all.jsonl
+        tokenizer,
+        image_size=224,
+        use_modal_embedding=False,
+        default_modal_index=0,
+        dataset_root=""       # ✅ 给 all.jsonl 用
+    ):
+        self.tokenizer = tokenizer
+        self.use_modal_embedding = use_modal_embedding
+        self.default_modal_index = default_modal_index
+        self.image_size = image_size
+        self.dataset_root = dataset_root
+
+        # ------- 读取样本 -------
+        if isinstance(data_source, str):
+
+            # ✅ 1️⃣ 如果是 CSV
+            if data_source.endswith(".csv"):
+                import pandas as pd
+                df = pd.read_csv(data_source)
+                self.samples = list(zip(df["image_path"], df["text"]))
+
+            # ✅ 2️⃣ 如果是 all.jsonl
+            elif data_source.endswith(".jsonl"):
+                self.samples = self._load_from_jsonl(data_source)
+
+            else:
+                raise ValueError("Unsupported file type. Use CSV or JSONL.")
+
+        else:
+            # ✅ 3️⃣ 直接是 list [(img_path, text), ...]
+            self.samples = data_source
+
+        # ------- 图像 transform -------
+        self.transform = transforms.Compose([
+            transforms.Grayscale(num_output_channels=1),
+            transforms.Resize((image_size, image_size)),
+            transforms.ToTensor(),
+            transforms.Normalize(mean=[0.5], std=[0.5])
+        ])
+
+        print(f"✅ Loaded {len(self.samples)} image-text pairs.")
+
+    def _load_from_jsonl(self, jsonl_path):
+        samples = []
+        with open(jsonl_path, "r", encoding="utf-8") as f:
+            for line in f:
+                item = json.loads(line.strip())
+
+                rel_path = item["img_path"]
+
+                rel_path = rel_path.replace("\\", "/")
+
+                if rel_path.startswith("dataset/"):
+                    rel_path = rel_path[len("dataset/"):]
+
+                if self.dataset_root:
+                    img_path = os.path.join(self.dataset_root, rel_path)
+                else:
+                    img_path = rel_path
+
+                text = item["text"].replace("\n", " ").strip()
+                samples.append((img_path, text))
+
+        return samples
+
+    def __len__(self):
+        return len(self.samples)
+
+    def __getitem__(self, index):
+        img_path, text_str = self.samples[index]
+
+        # ---- 1. 载入图像 -> [1,224,224] ----
+        img = Image.open(img_path).convert("L")
+        img = self.transform(img)
+
+        # ---- 2. tokenize 文本 ----
+        text = self.tokenizer(
+            text_str,
+            return_tensors="pt",
+            truncation=True,
+            padding="max_length",
+            max_length=512,
+        )
+
+        text = {k: v.squeeze(0) for k, v in text.items()}
+
+        # ---- 3. modal_embedding（可选）----
+        if self.use_modal_embedding:
+            modal_index = torch.tensor(self.default_modal_index, dtype=torch.long)
+            return {
+                "image": img,
+                "text": text,
+                "modal_index": modal_index
+            }
+
+        return {
+            "image": img,
+            "text": text
+        }
+
+
+class TextBatch:
+    def __init__(self, d):
+        for k, v in d.items():
+            setattr(self, k, v)
+
+
+def simple_collate(batch):
+    images = torch.stack([item["image"] for item in batch], dim=0)
+
+    text_dict = {
+        "input_ids": torch.stack([item["text"]["input_ids"] for item in batch]),
+        "attention_mask": torch.stack([item["text"]["attention_mask"] for item in batch]),
+    }
+
+    if "token_type_ids" in batch[0]["text"]:
+        text_dict["token_type_ids"] = torch.stack(
+            [item["text"]["token_type_ids"] for item in batch]
+        )
+
+    text = TextBatch(text_dict)
+
+    if "modal_index" in batch[0]:
+        modal_indexs = torch.stack([item["modal_index"] for item in batch])
+        return images, text, modal_indexs
+
+    return images, text

Rad_IR/radir/data/hmclip_mimic_box.py

@@ -0,0 +1,358 @@
+import os
+import re
+
+import math
+import numpy as np
+import pandas as pd
+import torch
+import torch.utils.data as data
+from nltk.tokenize import RegexpTokenizer
+import json
+
+
+
+#from radir.data.utils import get_imgs_from_pickle
+from tqdm import tqdm
+from transformers import BertTokenizer
+import random
+from PIL import Image
+from torchvision import transforms as T
+
+
+# 引入 monai 2D 增强
+from monai.transforms import (
+    Compose,
+    RandRotate,
+    RandFlip,
+    RandZoom,
+    RandAffine,
+    RandGaussianNoise,
+    RandScaleIntensity,
+    RandAdjustContrast,
+)
+
+BASE_DIR = os.path.dirname(os.path.abspath(__file__))
+
+# avg box num: 5.207
+
+class HyCoClipMimicBoxDataset_JPG(data.Dataset):
+    def __init__(
+        self,
+        *,
+        split: str,                         # "train" / "valid"
+        tokenizer,
+        dataset_root: str,                  # 项目根目录下 dataset 文件夹路径
+        cig_matched_dir: str,               # CIG 匹配 .pt 目录
+        caption_file: str,                # 支持任意 jsonl 文件
+        view_col: str = "ViewPosition",
+        split_col: str = "split",
+        transform=None,
+        max_words: int = 128,
+        imsize: int = 224
+    ):
+        super().__init__()
+        self.split = split
+        self.dataset_root = dataset_root
+        self.cig_matched_dir = cig_matched_dir
+        self.caption_file = caption_file
+        self.view_col = view_col
+        self.split_col = split_col
+        self.max_words = max_words
+        self.tokenizer = tokenizer
+        self.imsize = imsize
+
+
+        # ---------- 句子缓存 ----------
+
+        self.filenames, self.path2sent = self.load_text_data()
+
+        # ---------- CIG 文件集合 ----------
+        #####
+        self.cig_file_set = set()
+
+        if self.cig_matched_dir is not None and os.path.isdir(self.cig_matched_dir):
+            self.cig_file_set = {
+                f for f in os.listdir(self.cig_matched_dir)
+                if f.endswith(".pt")
+            }
+
+        print(f"[CIG] matched_dir={self.cig_matched_dir}, pt_files={len(self.cig_file_set)}")
+        #####
+
+        self.is_train = (self.split == "train")
+
+        # ---------- 图像处理 ----------
+        # 保持原来接口：如果用户传了 transform，就完全按用户的 transform 走
+        self.external_transform = transform
+
+        if self.external_transform is None:
+            # 默认：使用 2D 图像处理 + MONAI 2D 数据增强（只在 train 时用）
+
+            # 2D 数据增强（来自第二段代码的 2D augmentator）
+            if self.is_train:
+                self.augmentator = Compose([
+                    RandRotate(range_x=math.radians(15),range_y=math.radians(15),prob=0.5,keep_size=True), # 随机旋转（仅x和y轴）
+                    RandFlip(prob=0.5, spatial_axis=(0, 1)),  # 随机翻转（水平和垂直）
+                    RandZoom(min_zoom=0.8, max_zoom=1.1, prob=0.5, keep_size=True),  # 随机缩放
+                    RandAffine(prob=0.5, translate_range=(32, 32)),  # 随机平移（2D）
+                    RandGaussianNoise(prob=0.5, mean=0.0, std=0.1),  # 高斯噪声
+                    RandScaleIntensity(factors=[-0.15, 0.15], prob=0.5),  # 像素强度缩放
+                    RandAdjustContrast(prob=0.4, gamma=(0.85, 1.15)),  # 对比度调整
+                ])
+            else:
+                self.augmentator = None
+        else:
+            # 如果用户自己传了 transform，就不做 MONAI 的增强
+            self.augmentator = None
+    
+    # 构建路径逻辑
+    def get_image_path(self, row):
+        subject_id = str(row["subject_id"]).zfill(8)
+        study_id = str(row["study_id"])
+        dicom_id = str(row["dicom_id"])
+        p_folder = "p" + subject_id[:2]
+        subject_folder = "p" + subject_id
+        study_folder = "s" + study_id
+        img_name = dicom_id + ".jpg"
+        img_path = os.path.join(self.dataset_root, p_folder, subject_folder, study_folder, img_name)
+        return img_path
+    # 1D mask 修正函数
+    def _fix_1d_mask(self, x, target_len, dtype=torch.int8):
+        if not torch.is_tensor(x):
+            x = torch.tensor(x)
+        x = x.to(dtype).view(-1)
+
+        L = x.numel()
+        if L > target_len:
+            return x[:target_len]
+        if L < target_len:
+            return torch.cat(
+                [x, torch.zeros(target_len - L, dtype=dtype)],
+                dim=0
+            )
+        return x
+
+
+    def load_text_data(self):
+            import json
+            import os
+            import re
+
+            jsonl_path = os.path.join(self.dataset_root, self.caption_file)
+            path2sent = {}
+
+            if not os.path.exists(jsonl_path):
+                raise FileNotFoundError(f"Caption file not found at: {jsonl_path}")
+
+            print(f"Loading captions from: {jsonl_path}")
+
+            with open(jsonl_path, "r", encoding="utf-8") as f:
+                for line in f:
+                    item = json.loads(line.strip())
+                    
+                    # --- 修复逻辑开始 ---
+                    raw_rel_path = item["img_path"]
+                    
+                    # 1. 统一将所有反斜杠 \ 转换为正斜杠 / (解决跨平台问题)
+                    fixed_rel_path = raw_rel_path.replace("\\", "/")
+                    
+                    # 2. 移除可能存在的 dataset/ 前缀 (无论它原来是 \ 还是 /)
+                    if fixed_rel_path.startswith("dataset/"):
+                        fixed_rel_path = fixed_rel_path[len("dataset/"):]
+                    
+                    # 3. 使用 os.path.join 重新拼接并规范化路径
+                    # os.path.normpath 会自动处理多余的斜杠
+                    img_path = os.path.normpath(os.path.join(self.dataset_root, fixed_rel_path))
+                    # --- 修复逻辑结束 ---
+
+                    text = item["text"].replace("\n", " ").strip()
+                    # 简单分句
+                    sents = [s.strip() for s in re.split(r"[.?!]", text) if len(s.strip()) > 1]
+                    path2sent[img_path] = sents
+
+            filenames = list(path2sent.keys())
+            
+            # 调试输出：打印前两个路径确认是否正确
+            if len(filenames) > 0:
+                print(f"DEBUG: Sample image path: {filenames[0]}")
+                
+            print(f"Loaded {len(filenames)} image-text pairs.")
+            return filenames, path2sent
+    
+    def __len__(self):
+        return len(self.filenames)
+
+    def get_caption(self, img_path):
+        if img_path not in self.path2sent:
+            raise Exception(f"No sentence found for image path: {img_path}")
+        series_sents = self.path2sent[img_path]
+        sent = " ".join(filter(lambda x: x != "", series_sents))
+        tokens = self.tokenizer(
+            sent,
+            return_tensors="pt",
+            truncation=True,
+            padding="max_length",
+            max_length=self.max_words, 
+        )
+        x_len = len([t for t in tokens["input_ids"][0] if t != 0])
+        return tokens, x_len
+
+    def _load_2d_image_with_aug(self, img_path):
+        """
+        只考虑 2D JPG：
+        - 转灰度
+        - 转为 tensor [1, H, W]，范围 [0,1]
+        - 可选 MONAI 增强（train 时）
+        - resize 到 (imsize, imsize)
+        - 线性映射到 [-1,1]（等价于原来的 Normalize(mean=0.5,std=0.5)）
+        """
+        # 读灰度
+        img = Image.open(img_path).convert("L")  # 单通道
+
+        # 转 numpy [H, W]，0~1
+        img_np = np.array(img, dtype=np.float32) / 255.0
+        img_t = torch.from_numpy(img_np).unsqueeze(0)  # [1, H, W]
+
+        # MONAI 数据增强（仅 train & 默认 transform 时）
+        if self.augmentator is not None:
+            img_t = self.augmentator(img_t)  # 仍是 [1, H', W']
+
+        # 统一 resize 到 (imsize, imsize)
+        img_t = torch.nn.functional.interpolate(
+            img_t.unsqueeze(0),  # [1, 1, H, W]
+            size=(self.imsize, self.imsize),
+            mode="bilinear",
+            align_corners=False,
+        ).squeeze(0)  # [1, imsize, imsize]
+
+        # 映射到 [-1,1]（与原先 ToTensor+Normalize(0.5,0.5) 等价）
+        img_t = img_t * 2.0 - 1.0
+
+        return img_t  # [1, imsize, imsize]
+
+    def __getitem__(self, index):
+        img_path = self.filenames[index]
+        caps, cap_len = self.get_caption(img_path)
+
+        # ---------- 读取 JPG + 2D 处理 ----------
+        try:
+            if self.external_transform is None:
+                # 使用我们移植过来的 2D 处理 + 增强
+                img = self._load_2d_image_with_aug(img_path)  # [1, 224, 224]
+            else:
+                # 兼容旧代码：如果用户传了 transform，就按旧逻辑来
+                img = Image.open(img_path).convert("RGB")
+                img = self.external_transform(img)
+        except Exception as e:
+            raise FileNotFoundError(f"Cannot open image {img_path}: {e}")
+
+        # ---------- phrase ----------
+        dicom_id = os.path.splitext(os.path.basename(img_path))[0]
+        cig_file_name = f"{dicom_id}.pt"
+        if cig_file_name in self.cig_file_set:
+            cig_path = os.path.join(self.cig_matched_dir, cig_file_name)
+            d = torch.load(cig_path, weights_only=False)
+            phrase_idx = random.randint(0, len(d["unique_phrases"]) - 1)
+            phrase = d["unique_phrases"][phrase_idx]
+            phrase_text_mask = self._fix_1d_mask(
+                d["phrase_text_mask"][phrase_idx],
+                self.max_words,
+                dtype=torch.int8
+            )
+
+            phrases_bboxes_mask = self._fix_1d_mask(
+                d["phrases_bboxes_mask"][phrase_idx],
+                196,
+                dtype=torch.int8
+            )
+        else:
+            phrase = ""
+            phrase_text_mask = torch.zeros((self.max_words,), dtype=torch.int8)
+            phrases_bboxes_mask = torch.zeros((196,), dtype=torch.int8)
+
+        def compute_offsets(text, tokens):
+            """
+            手动根据 tokens 在 text 中出现的位置计算 offset_mapping
+            slow tokenizer（不支持 offset_mapping）也可以用
+            """
+            offsets = []
+            start_pos = 0
+            for tok in tokens:
+                # BertTokenizer 会给 word 加 ## 前缀，需要去掉才能在原文中搜索到
+                tok_clean = tok.replace("##", "")
+                idx = text.find(tok_clean, start_pos)
+                if idx == -1:
+                    # 没有找到就写 0,0（比如特殊 token）
+                    offsets.append((0, 0))
+                else:
+                    offsets.append((idx, idx + len(tok_clean)))
+                    start_pos = idx + len(tok_clean)
+            return offsets
+
+        # ----- 你的代码流程 -----
+        # full sentence
+        sents = self.path2sent[img_path]
+        full_text = " ".join(sents).replace("\n", " ").strip()
+
+        # tokenize full_text（slow tokenizer）
+        tokens_full = self.tokenizer.tokenize(full_text)
+        offsets_full = compute_offsets(full_text, tokens_full)
+
+        # padding 到 max_words
+        if len(tokens_full) < self.max_words:
+            offsets_full = offsets_full + [(0, 0)] * (self.max_words - len(offsets_full))
+        else:
+            offsets_full = offsets_full[:self.max_words]
+
+
+        return img, caps, cap_len, img_path, phrase, phrase_text_mask, phrases_bboxes_mask
+
+
+def collate_hmclip_box(batch):
+    # 与 hycoclip/data/pretrain_dataset_box.py 等价实现
+    """sort sequence"""
+    imgs, cap_len, ids, tokens, attention = [], [], [], [], []
+    path = []
+    phrase_lst, phrase_text_mask_lst, phrases_bboxes_mask_lst = [], [], []
+    for b in batch:
+        img, cap, cap_l, p, ph, ph_t_mask, ph_bbx_mask = b
+        imgs.append(img)
+        cap_len.append(cap_l)
+        ids.append(cap["input_ids"])
+        tokens.append(cap["token_type_ids"])
+        attention.append(cap["attention_mask"])
+        path.append(p)
+
+        phrase_lst.append(ph)
+        phrase_text_mask_lst.append(ph_t_mask)
+        phrases_bboxes_mask_lst.append(ph_bbx_mask)
+
+    # stack
+    imgs = torch.stack(imgs)
+    ids = torch.stack(ids).squeeze()
+    tokens = torch.stack(tokens).squeeze()
+    attention = torch.stack(attention).squeeze()
+
+    phrase_text_mask = torch.stack(phrase_text_mask_lst) 
+    phrases_bboxes_mask = torch.stack(phrases_bboxes_mask_lst).squeeze()
+
+    # sort and add to dictionary
+    sorted_cap_lens, sorted_cap_indices = torch.sort(
+        torch.tensor(cap_len), 0, True)
+
+    path = np.array(path)
+    phrase = np.array(phrase_lst)
+
+    return_dict = {
+        "caption_ids": ids[sorted_cap_indices],
+        "token_type_ids": tokens[sorted_cap_indices],
+        "attention_mask": attention[sorted_cap_indices],
+        "imgs": imgs[sorted_cap_indices],
+        "cap_lens": sorted_cap_lens,
+        "path": path[sorted_cap_indices],
+
+        "phrase": phrase[sorted_cap_indices],
+        "phrase_text_mask" : phrase_text_mask[sorted_cap_indices],
+        "phrases_bboxes_mask" : phrases_bboxes_mask[sorted_cap_indices],
+    }
+    return return_dict

Rad_IR/radir/data/utils.py

@@ -0,0 +1,18 @@
+from PIL import Image
+
+def get_imgs_from_pickle(img, transform=None):
+    """
+    将 H5/PKL 里读出的 numpy 图像数组转为 RGB PIL，再按需应用 transform。
+    参数:
+    img: numpy.ndarray，形状一般为 (H, W) 或 (H, W, C)
+    transform: 可选的 torchvision transforms
+    返回:
+    已转换的图像（通常为张量）
+    """
+    pil_img = Image.fromarray(img).convert("RGB")
+
+    if transform is not None:
+        pil_img = transform(pil_img)
+
+    return pil_img
+

Rad_IR/radir/lorentz.py

@@ -0,0 +1,312 @@
+#---------------------------------------
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+#---------------------------------------
+
+# Modified from github.com/facebookresearch/meru
+
+"""
+Implementation of common operations for the Lorentz model of hyperbolic geometry.
+This model represents a hyperbolic space of `d` dimensions on the upper-half of
+a two-sheeted hyperboloid in a Euclidean space of `(d+1)` dimensions.
+
+Hyperbolic geometry has a direct connection to the study of special relativity
+theory -- implementations in this module borrow some of its terminology. The axis
+of symmetry of the Hyperboloid is called the _time dimension_, while all other
+axes are collectively called _space dimensions_.
+
+All functions implemented here only input/output the space components, while
+while calculating the time component according to the Hyperboloid constraint:
+
+    `x_time = torch.sqrt(1 / curv + torch.norm(x_space) ** 2)`
+"""
+from __future__ import annotations
+
+import math
+
+import torch
+from torch import Tensor
+from loguru import logger
+
+
+def pairwise_inner(x: Tensor, y: Tensor, curv: float | Tensor = 1.0):
+    """
+    Compute pairwise Lorentzian inner product between input vectors.
+
+    Args:
+        x: Tensor of shape `(B1, D)` giving a space components of a batch
+            of vectors on the hyperboloid.
+        y: Tensor of shape `(B2, D)` giving a space components of another
+            batch of points on the hyperboloid.
+        curv: Positive scalar denoting negative hyperboloid curvature.
+        eps: Small float number to avoid numerical instability.
+
+    Returns:
+        Tensor of shape `(B1, B2)` giving pairwise Lorentzian inner product
+        between input vectors.
+    """
+
+    x_time = torch.sqrt(1 / curv + torch.sum(x**2, dim=-1, keepdim=True))
+    y_time = torch.sqrt(1 / curv + torch.sum(y**2, dim=-1, keepdim=True))
+    xyl = x @ y.T - x_time @ y_time.T
+    return xyl
+
+
+def pairwise_dist(
+    x: Tensor, y: Tensor, curv: float | Tensor = 1.0, eps: float = 1e-8
+) -> Tensor:
+    """
+    Compute the pairwise geodesic distance between two batches of points on
+    the hyperboloid.
+
+    Args:
+        x: Tensor of shape `(B1, D)` giving a space components of a batch
+            of point on the hyperboloid.
+        y: Tensor of shape `(B2, D)` giving a space components of another
+            batch of points on the hyperboloid.
+        curv: Positive scalar denoting negative hyperboloid curvature.
+        eps: Small float number to avoid numerical instability.
+
+    Returns:
+        Tensor of shape `(B1, B2)` giving pairwise distance along the geodesics
+        connecting the input points.
+    """
+
+
+
+    # --- Handle (B,1,D) -> (B,D) ---
+    if x.dim() == 3 and x.size(1) == 1:
+        x = x.squeeze(1)
+    
+    x = x.float()
+    y = y.float()
+
+    # Ensure numerical stability in arc-cosh by clamping input.
+
+
+    c_xyl = -curv * pairwise_inner(x, y, curv)
+    _distance = torch.acosh(torch.clamp(c_xyl, min=1 + eps))
+    return _distance / curv**0.5
+
+
+def exp_map0(x: Tensor, curv: float | Tensor = 1.0, eps: float = 1e-8) -> Tensor:
+    """
+    Map points from the tangent space at the vertex of hyperboloid, on to the
+    hyperboloid. This mapping is done using the exponential map of Lorentz model.
+
+    Args:
+        x: Tensor of shape `(B, D)` giving batch of Euclidean vectors to project
+            onto the hyperboloid. These vectors are interpreted as velocity
+            vectors in the tangent space at the hyperboloid vertex.
+        curv: Positive scalar denoting negative hyperboloid curvature.
+        eps: Small float number to avoid division by zero.
+
+    Returns:
+        Tensor of same shape as `x`, giving space components of the mapped
+        vectors on the hyperboloid.
+    """
+
+    rc_xnorm = curv**0.5 * torch.norm(x, dim=-1, keepdim=True)
+
+    # Ensure numerical stability in sinh by clamping input.
+    sinh_input = torch.clamp(rc_xnorm, min=eps, max=math.asinh(2**15))
+    _output = torch.sinh(sinh_input) * x / torch.clamp(rc_xnorm, min=eps)
+    return _output
+
+
+def log_map0(x: Tensor, curv: float | Tensor = 1.0, eps: float = 1e-8) -> Tensor:
+    """
+    Inverse of the exponential map: map points from the hyperboloid on to the
+    tangent space at the vertex, using the logarithmic map of Lorentz model.
+
+    Args:
+        x: Tensor of shape `(B, D)` giving space components of points
+            on the hyperboloid.
+        curv: Positive scalar denoting negative hyperboloid curvature.
+        eps: Small float number to avoid division by zero.
+
+    Returns:
+        Tensor of same shape as `x`, giving Euclidean vectors in the tangent
+        space of the hyperboloid vertex.
+    """
+
+    # Calculate distance of vectors to the hyperboloid vertex.
+    rc_x_time = torch.sqrt(1 + curv * torch.sum(x**2, dim=-1, keepdim=True))
+    _distance0 = torch.acosh(torch.clamp(rc_x_time, min=1 + eps))
+
+    rc_xnorm = curv**0.5 * torch.norm(x, dim=-1, keepdim=True)
+    _output = _distance0 * x / torch.clamp(rc_xnorm, min=eps)
+    return _output
+
+
+def half_aperture(
+    x: Tensor, curv: float | Tensor = 1.0, min_radius: float = 0.1, eps: float = 1e-6
+) -> Tensor:
+    """
+    Compute the half aperture angle of the entailment cone formed by vectors on
+    the hyperboloid. The given vector would meet the apex of this cone, and the
+    cone itself extends outwards to infinity.
+
+    Args:
+        x: Tensor of shape `(B, D)` giving a batch of space components of
+            vectors on the hyperboloid.
+        curv: Positive scalar denoting negative hyperboloid curvature.
+        min_radius: Radius of a small neighborhood around vertex of the hyperboloid
+            where cone aperture is left undefined. Input vectors lying inside this
+            neighborhood (having smaller norm) will be projected on the boundary.
+        eps: Small float number to avoid numerical instability.
+
+    Returns:
+        Tensor of shape `(B, )` giving the half-aperture of entailment cones
+        formed by input vectors. Values of this tensor lie in `(0, pi/2)`.
+    """
+
+    # --- Handle (B,1,D) -> (B,D) ---
+    if x.dim() == 3 and x.size(1) == 1:
+        x = x.squeeze(1)
+    
+    x = x.float()
+
+    # Ensure numerical stability in arc-sin by clamping input.
+
+
+    asin_input = 2 * min_radius / (torch.norm(x, dim=-1) * curv**0.5 + eps)
+    asin_input = torch.clamp(asin_input, -1 + eps, 1 - eps)
+    _half_aperture = torch.asin(torch.clamp(asin_input, min=-1 + eps, max=1 - eps))
+
+    return _half_aperture
+
+
+def oxy_angle(x: Tensor, y: Tensor, curv: float | Tensor = 1.0, eps: float = 1e-7):
+     
+    """
+    Given two vectors `x` and `y` on the hyperboloid, compute the exterior
+    angle at `x` in the hyperbolic triangle `Oxy` where `O` is the origin
+    of the hyperboloid.
+
+    This expression is derived using the Hyperbolic law of cosines.
+
+    Args:
+        x: Tensor of shape `(B, D)` giving a batch of space components of
+            vectors on the hyperboloid.
+        y: Tensor of same shape as `x` giving another batch of vectors.
+        curv: Positive scalar denoting negative hyperboloid curvature.
+
+    Returns:
+        Tensor of shape `(B, )` giving the required angle. Values of this
+        tensor lie in `(0, pi)`.
+    """
+
+
+    # --- Handle (B,1,D) -> (B,D) ---
+    if x.dim() == 3 and x.size(1) == 1:
+        x = x.squeeze(1)
+    if y.dim() == 3 and y.size(1) == 1:
+        y = y.squeeze(1)
+
+
+    x = x.float()
+    y = y.float()
+
+    # -------- curvature 数值保护 --------
+    if isinstance(curv, Tensor):
+        curv = curv.float().clamp(min=eps)
+
+    # Calculate time components of inputs (multiplied with `sqrt(curv)`):
+    x_time = torch.sqrt(1 / curv + torch.sum(x**2, dim=-1))
+    y_time = torch.sqrt(1 / curv + torch.sum(y**2, dim=-1))
+
+
+    # Calculate lorentzian inner product multiplied with curvature. We do not use
+    # the `pairwise_inner` implementation to save some operations (since we only
+    # need the diagonal elements).
+
+    c_xyl = curv * (torch.sum(x * y, dim=-1) - x_time * y_time)
+
+    # Make the numerator and denominator for input to arc-cosh, shape: (B, )
+
+    
+    acos_numer = y_time + c_xyl * x_time
+    acos_denom = torch.sqrt(torch.clamp(c_xyl**2 - 1, min=eps))
+
+
+    acos_input = acos_numer / (torch.norm(x, dim=-1) * acos_denom + eps)
+    if not torch.isfinite(acos_input).all():
+        print("❌ acos_input has NaN/Inf!")
+        print("  c_xyl:", c_xyl)
+        print("  acos_numer:", acos_numer)
+        print("  acos_denom:", acos_denom)
+        print("  norm_x:", torch.norm(x, dim=-1))
+        
+    # print("c_xyl min/mean/max:",
+    #     c_xyl.min().item(), c_xyl.mean().item(), c_xyl.max().item())
+    # print("acos_denom min/mean:",
+    #     acos_denom.min().item(), acos_denom.mean().item())
+    # print("norm_x min/mean:",
+    #     torch.norm(x, dim=-1).min().item(), torch.norm(x, dim=-1).mean().item())
+    # print("acos_numer min/mean:",
+    #     acos_numer.min().item(), acos_numer.mean().item())
+
+
+
+    # print("acos_input min/max/mean:",
+    #     acos_input.min().item(), acos_input.max().item(), acos_input.mean().item(),
+    #     "pct<=-1:", (acos_input <= -1).float().mean().item(),
+    #     "pct>= 1:", (acos_input >=  1).float().mean().item())
+
+
+    _angle = torch.acos(torch.clamp(acos_input, min=-1 + eps, max=1 - eps))
+
+    if not torch.isfinite(_angle).all():
+        print("❌ _angle has NaN/Inf!")
+
+    return _angle
+
+
+def oxy_angle_eval(x: Tensor, y: Tensor, curv: float | Tensor = 1.0, eps: float = 1e-8):
+    """
+    Given two vectors `x` and `y` on the hyperboloid, compute the exterior
+    angle at `x` in the hyperbolic triangle `Oxy` where `O` is the origin
+    of the hyperboloid.
+
+    This expression is derived using the Hyperbolic law of cosines.
+
+    Args:
+        x: Tensor of shape `(B, D)` giving a batch of space components of
+            vectors on the hyperboloid.
+        y: Tensor of same shape as `x` giving another batch of vectors.
+        curv: Positive scalar denoting negative hyperboloid curvature.
+
+    Returns:
+        Tensor of shape `(B, )` giving the required angle. Values of this
+        tensor lie in `(0, pi)`.
+    """
+
+    # Calculate time components of inputs (multiplied with `sqrt(curv)`):
+    x_time = torch.sqrt(1 / curv + torch.sum(x**2, dim=-1, keepdim=True))
+    y_time = torch.sqrt(1 / curv + torch.sum(y**2, dim=-1, keepdim=True))
+
+    logger.info(f"x_time shape: {x_time.size()}")
+    logger.info(f"y_time shape: {y_time.size()}")
+
+    # Calculate lorentzian inner product multiplied with curvature. We do not use
+    # the `pairwise_inner` implementation to save some operations (since we only
+    # need the diagonal elements).
+
+    # c_xyl = curv * (torch.sum(x * y, dim=-1) - x_time * y_time)
+    c_xyl = curv * (y @ x.T - y_time @ x_time.T)
+    logger.info(f"c_xyl shape: {c_xyl.size()}")
+
+    # Make the numerator and denominator for input to arc-cosh, shape: (B, )
+    acos_numer = y_time + c_xyl * x_time.T
+    logger.info(f"acos_numer shape: {acos_numer.size()}")
+    acos_denom = torch.sqrt(torch.clamp(c_xyl**2 - 1, min=eps))
+    logger.info(f"acos_denom shape: {acos_denom.size()}")
+
+    acos_input = acos_numer / (torch.norm(x, dim=-1, keepdim=True).T * acos_denom + eps)
+    _angle = - torch.acos(torch.clamp(acos_input, min=-1 + eps, max=1 - eps))
+
+    return _angle

Rad_IR/radir/transformer_decoder.py

@@ -0,0 +1,154 @@
+"""
+Code modified from DETR tranformer:
+https://github.com/facebookresearch/detr
+Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
+"""
+
+import copy
+from typing import Optional, List
+import pickle as cp
+
+import torch
+import torch.nn.functional as F
+from torch import nn, Tensor
+
+
+class TransformerDecoder(nn.Module):
+
+    def __init__(self, decoder_layer, num_layers, norm=None, return_intermediate=False):
+        super().__init__()
+        self.layers = _get_clones(decoder_layer, num_layers)
+        self.num_layers = num_layers
+        self.norm = norm
+        self.return_intermediate = return_intermediate
+
+    def forward(self, tgt, memory,
+                tgt_mask: Optional[Tensor] = None,
+                memory_mask: Optional[Tensor] = None,
+                tgt_key_padding_mask: Optional[Tensor] = None,
+                memory_key_padding_mask: Optional[Tensor] = None,
+                pos: Optional[Tensor] = None,
+                query_pos: Optional[Tensor] = None):
+        output = tgt
+        T,B,C = memory.shape
+        intermediate = []
+        atten_layers = []
+        for n,layer in enumerate(self.layers):
+   
+            residual=True
+            output,ws = layer(output, memory, tgt_mask=tgt_mask,
+                           memory_mask=memory_mask,
+                           tgt_key_padding_mask=tgt_key_padding_mask,
+                           memory_key_padding_mask=memory_key_padding_mask,
+                           pos=pos, query_pos=query_pos,residual=residual)
+            atten_layers.append(ws)
+            if self.return_intermediate:
+                intermediate.append(self.norm(output))
+        if self.norm is not None:
+            output = self.norm(output)
+            if self.return_intermediate:
+                intermediate.pop()
+                intermediate.append(output)
+
+        if self.return_intermediate:
+            return torch.stack(intermediate)
+        return output,atten_layers
+
+
+class TransformerDecoderLayer(nn.Module):
+
+    def __init__(self, d_model, nhead, dim_feedforward=2048, dropout=0.1,
+                 activation="relu", normalize_before=False):
+        super().__init__()
+        # self.self_attn = nn.MultiheadAttention(d_model, nhead, dropout=dropout)
+        self.multihead_attn = nn.MultiheadAttention(d_model, nhead, dropout=dropout)
+        # Implementation of Feedforward model
+        self.linear1 = nn.Linear(d_model, dim_feedforward)
+        self.dropout = nn.Dropout(dropout)
+        self.linear2 = nn.Linear(dim_feedforward, d_model)
+
+        self.norm1 = nn.LayerNorm(d_model)
+        self.norm2 = nn.LayerNorm(d_model)
+        self.norm3 = nn.LayerNorm(d_model)
+        self.dropout1 = nn.Dropout(dropout)
+        self.dropout2 = nn.Dropout(dropout)
+        self.dropout3 = nn.Dropout(dropout)
+
+        self.activation = _get_activation_fn(activation)
+        self.normalize_before = normalize_before
+    def with_pos_embed(self, tensor, pos: Optional[Tensor]):
+        return tensor if pos is None else tensor + pos
+
+    def forward_post(self, tgt, memory,
+                     tgt_mask: Optional[Tensor] = None,
+                     memory_mask: Optional[Tensor] = None,
+                     tgt_key_padding_mask: Optional[Tensor] = None,
+                     memory_key_padding_mask: Optional[Tensor] = None,
+                     pos: Optional[Tensor] = None,
+                     query_pos: Optional[Tensor] = None,
+                     residual=True):
+        q = k = self.with_pos_embed(tgt, query_pos)
+        tgt2,ws = self.self_attn(q, k, value=tgt, attn_mask=tgt_mask,
+                              key_padding_mask=tgt_key_padding_mask)
+        tgt = self.norm1(tgt)
+        tgt2,ws = self.multihead_attn(query=self.with_pos_embed(tgt, query_pos),
+                                   key=self.with_pos_embed(memory, pos),
+                                   value=memory, attn_mask=memory_mask,
+                                   key_padding_mask=memory_key_padding_mask)
+
+        # attn_weights [B,NUM_Q,T]
+        tgt = tgt + self.dropout2(tgt2)
+        tgt = self.norm2(tgt)
+        tgt2 = self.linear2(self.dropout(self.activation(self.linear1(tgt))))
+        tgt = tgt + self.dropout3(tgt2)
+        tgt = self.norm3(tgt)
+        return tgt,ws
+
+    def forward_pre(self, tgt, memory,
+                    tgt_mask: Optional[Tensor] = None,
+                    memory_mask: Optional[Tensor] = None,
+                    tgt_key_padding_mask: Optional[Tensor] = None,
+                    memory_key_padding_mask: Optional[Tensor] = None,
+                    pos: Optional[Tensor] = None,
+                    query_pos: Optional[Tensor] = None):
+
+        tgt2 = self.norm2(tgt)
+        tgt2,attn_weights = self.multihead_attn(query=self.with_pos_embed(tgt2, query_pos),
+                                   key=self.with_pos_embed(memory, pos),
+                                   value=memory, attn_mask=memory_mask,
+                                   key_padding_mask=memory_key_padding_mask)
+        tgt = tgt + self.dropout2(tgt2)
+        tgt2 = self.norm3(tgt)
+        tgt2 = self.linear2(self.dropout(self.activation(self.linear1(tgt2))))
+        tgt = tgt + self.dropout3(tgt2)
+        return tgt,attn_weights
+
+    def forward(self, tgt, memory,
+                tgt_mask: Optional[Tensor] = None,
+                memory_mask: Optional[Tensor] = None,
+                tgt_key_padding_mask: Optional[Tensor] = None,
+                memory_key_padding_mask: Optional[Tensor] = None,
+                pos: Optional[Tensor] = None,
+                query_pos: Optional[Tensor] = None,
+                residual=True):
+        if self.normalize_before:
+            return self.forward_pre(tgt, memory, tgt_mask, memory_mask,
+                                    tgt_key_padding_mask, memory_key_padding_mask, pos, query_pos)
+        return self.forward_post(tgt, memory, tgt_mask, memory_mask,
+                                 tgt_key_padding_mask, memory_key_padding_mask, pos, query_pos,residual)
+
+
+def _get_clones(module, N):
+    return nn.ModuleList([copy.deepcopy(module) for i in range(N)])
+
+
+
+def _get_activation_fn(activation):
+    """Return an activation function given a string"""
+    if activation == "relu":
+        return F.relu
+    if activation == "gelu":
+        return F.gelu
+    if activation == "glu":
+        return F.glu
+    raise RuntimeError(F"activation should be relu/gelu, not {activation}.")

Rad_IR/setup.py

@@ -0,0 +1,32 @@
+from setuptools import setup, find_packages
+
+setup(
+  name = 'radir',
+  packages = find_packages(exclude=[]),
+  include_package_data = True,
+  version = '0.1',
+  description = 'RADIR',
+  install_requires=[
+    'beartype',
+    'einops>=0.6',
+    'ftfy',
+    'regex',
+    "XlsxWriter",
+    "h5py",
+    "matplotlib",
+    "seaborn",
+    "wilds",
+    #'ImageNetV2_pytorch @ git+https://github.com/modestyachts/ImageNetV2_pytorch.git',
+    "click",
+    "appdirs",
+    "attr",
+    "nltk",
+    'transformers',
+    'monai',
+    'positional_encodings',
+    'tensorboard',
+    'pynvml',
+    'wandb',
+    'torchvision',
+      ],
+)

backup-manifest.json

@@ -0,0 +1,20 @@
+{
+  "format": "split-tar-zstd-lfs-v1",
+  "source_commit": "5b926d2179afa417784ba57459e86d2336d50a52",
+  "complete": false,
+  "archived_paths": [
+    "dataset",
+    "outputs",
+    "hf_models",
+    "hf_cache",
+    "runs",
+    "wandb",
+    "__pycache__"
+  ],
+  "volume_size_bytes": 67108864,
+  "parts": [],
+  "compressed_bytes": 0,
+  "archive_sha256": null,
+  "restore": "python3 restore_backup.py",
+  "source_files_unchanged": true
+}

restore_backup.py

@@ -0,0 +1,29 @@
+#!/usr/bin/env python3
+"""Verify downloaded volumes, then extract without concatenating a large archive."""
+import hashlib,json,subprocess,sys
+from pathlib import Path
+root=Path(__file__).resolve().parent
+m=json.loads((root/'backup-manifest.json').read_text())
+if not m['complete']: raise SystemExit('This backup is still uploading.')
+for part in m['parts']:
+ p=root/part['path']; h=hashlib.sha256()
+ with p.open('rb') as f:
+  for block in iter(lambda:f.read(4*1024*1024),b''):h.update(block)
+ if p.stat().st_size!=part['size'] or h.hexdigest()!=part['sha256']:
+  raise SystemExit('Missing/corrupt volume or unmaterialized LFS pointer: '+str(p))
+print('All volume checksums verified.',flush=True)
+out=Path(sys.argv[1]).resolve() if len(sys.argv)>1 else root
+out.mkdir(parents=True,exist_ok=True)
+z=subprocess.Popen(['zstd','-d','-c'],stdin=subprocess.PIPE,stdout=subprocess.PIPE)
+t=subprocess.Popen(['tar','-xf','-','-C',str(out)],stdin=z.stdout)
+z.stdout.close()
+try:
+ for part in m['parts']:
+  with (root/part['path']).open('rb') as f:
+   for block in iter(lambda:f.read(4*1024*1024),b''):z.stdin.write(block)
+ z.stdin.close()
+ if z.wait()!=0 or t.wait()!=0:raise SystemExit('Extraction failed.')
+finally:
+ for proc in [z,t]:
+  if proc.poll() is None:proc.terminate()
+print('Restored the archived directories to',out)

sbatch/stage1.sh

@@ -0,0 +1,59 @@
+export HF_HOME=/mnt/data-3/users/zhangzeyu/HMClip/hf_cache
+export TRANSFORMERS_CACHE=/mnt/data-3/users/zhangzeyu/HMClip/hf_cache
+export CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7
+export RESULTS_FOLDER=./outputs
+
+set -e 
+
+export MASTER_PORT=$((RANDOM % 987 + 29000))
+
+accelerate launch \
+--num_processes=8 \
+--num_machines=1 \
+--dynamo_backend=no \
+--mixed_precision=fp16 \
+--gradient_accumulation_steps=8 \
+--main_process_port "${MASTER_PORT}" \
+scripts/run_CXR_CT_train.py \
+--checkpoint "/mnt/data-3/users/zhangzeyu/HMClip/RadIR/Rad_IR/checkpoints/CTClip.8000.pt" \
+--name 'CXRCT_UNI_stage1' \
+--allow_partial_load \
+--log_dir "/mnt/data-3/users/zhangzeyu/HMClip/RadIR/logs/CXRCT_UNI_stage1/" \
+--pin_memory \
+--open_pretrained_encoders \
+--freeze_fusion_module \
+--use_image2image_loss 1 \
+--use_uncon_infoNCE_loss 1 \
+--use_uncon_triplet_loss 0 \
+--positive_distance_threshold 0.25 \
+--positive_threshold 0.75 \
+--anatomy_filter '[["train"], ["train_0", "train_1", "train_2", "train_3"]]' \
+--stage1 \
+--uncon_soft_label \
+--uncon_similarity_lookup_table_train './dataset/CTRATE-IR' \
+--uncon_similarity_lookup_table_valid './dataset/CTRATE-IR' \
+--uncon_data_train_jsonl './dataset/CTRATE-IR' \
+--uncon_data_valid_jsonl './dataset/CTRATE-IR' \
+--uncon_dataset_names 'CTRATE-IR' \
+--uncon_train_filter '[["train"]]' \
+--uncon_valid 'test' \
+--uncon_modality '3D' \
+--train_no_aug \
+--infonce_temp \
+--modal_embedding \
+--save_results_every 500 \
+--save_model_every 6000 \
+--num_train_steps 12000 \
+--warmup_steps 1000 \
+--lr 1e-4 \
+--uncon_batch_size 6 \
+--uncon_batch_size_valid 6 \
+--uncon_num_workers 4 \
+--data_train_jsonl ./dataset/MIMIC-IR/train_caption.csv ./dataset/CTRATE-IR/train_filtered.jsonl \
+--data_valid_jsonl ./dataset/MIMIC-IR/val_caption.csv ./dataset/CTRATE-IR/test.jsonl \
+--data_train_csv_dir ./dataset/MIMIC-IR ./dataset/CTRATE-IR \
+--data_valid_csv_dir ./dataset/MIMIC-IR ./dataset/CTRATE-IR \
+--data_train_npy_dir ./dataset/MIMIC-IR/train_ratescore.npy ./dataset/CTRATE-IR/CT_train_ratescore_compressed.npz \
+--data_valid_npy_dir ./dataset/MIMIC-IR/val_ratescore.npy ./dataset/CTRATE-IR/CT_test_ratescore_compressed.npz 
+export RESULTS_FOLDER=./outputs
+

sbatch/stage1_and_stage2.sh

@@ -0,0 +1,61 @@
+set -e 
+
+export MASTER_PORT=$((RANDOM % 987 + 29000))
+
+accelerate launch \
+--num_processes=4 \
+--num_machines=1 \
+--dynamo_backend=no \
+--mixed_precision=fp16 \
+--gradient_accumulation_steps=8 \
+--main_process_port "${MASTER_PORT}" \
+/scripts/run_CXR_CT_train.py \
+--checkpoint "xxx" \
+--name 'CXRCT_UNI_stage1_and_stage2' \
+--allow_partial_load \
+--pin_memory \
+--open_pretrained_encoders \
+--open_fusion_module \
+--use_triplet_loss 1.0 \
+--use_infoNCE_loss 0.0 \
+--use_image2image_loss 1 \
+--use_uncon_infoNCE_loss 1 \
+--use_uncon_triplet_loss 0 \
+--positive_distance_threshold 0.25 \
+--positive_threshold 0.75 \
+--stage1 \
+--stage2 \
+--anatomy_filter  '[["bone","heart","bronchie","trachea","pleura","vertebrae","liver","aorta","spinal canal","gallbladder","clavicle","heart ascending aorta","pulmonary artery","breast","pancreas","stomach"],["lungs","heart","bones","abdomen","spine","right lung","pulmonary vasculature","aorta","airway structures","left lower lobe","right lower lobe","right upper lobe","thoracic aorta","ventricle","atrium","carina","pulmonary artery","aortic arch","bronchi","costophrenic angle","aortic valve","pericardium"]]' \
+--dataset_names 'CT_RATE' 'MIMIC_CXR' \
+--modality '3D' '2D' \
+--data_train_jsonl  '/dataset/CT_RATE/train.jsonl' '/dataset/CXR/MIMIC_CXR/train.jsonl'  \
+--data_valid_jsonl  '/dataset/CT_RATE/test.jsonl' '/dataset/CXR/MIMIC_CXR/test.jsonl' \
+--data_train_csv_dir '/dataset/CT_RATE/anatomy/train_entity' '/dataset/CXR/MIMIC_CXR/train_entity'  \
+--data_valid_csv_dir '/dataset/CT_RATE/anatomy/val_entity' '/dataset/CXR/MIMIC_CXR/test_entity' \
+--data_train_npy_dir '/dataset/CT_RATE/anatomy/train_ratescore' '/dataset/CXR/MIMIC_CXR/train_ratescore' \
+--data_valid_npy_dir '/dataset/CT_RATE/anatomy/val_ratescore' '/dataset/CXR/MIMIC_CXR/test_ratescore' \
+--uncon_soft_label \
+--uncon_similarity_lookup_table_train '/dataset/CT_RATE' '/dataset/CXR/MIMIC_CXR/all' \
+--uncon_similarity_lookup_table_valid '/dataset/CT_RATE' '/dataset/CXR/MIMIC_CXR/all' \
+--uncon_data_train_jsonl '/dataset/CT_RATE' '/dataset/CXR/MIMIC_CXR/all' \
+--uncon_data_valid_jsonl '/dataset/CT_RATE' '/dataset/CXR/MIMIC_CXR/all' \
+--uncon_dataset_names 'CT_RATE' 'MIMIC-CXR' \
+--uncon_train_filter '[["train"],["train_0","train_1","train_2","train_3"]]' \
+--uncon_valid 'test' 'test' \
+--uncon_modality '3D' '2D' \
+--train_no_aug \
+--infonce_temp \
+--modal_embedding \
+--save_results_every 500 \
+--save_model_every 6000 \
+--num_train_steps 12000 \
+--warmup_steps 1000 \
+--lr 1e-4 \
+--batch_size 1 2 \
+--local_batch_size 6 30 \
+--uncon_batch_size 6 60 \
+--uncon_batch_size_valid 6 60 \
+--uncon_num_workers 4 \
+--con_num_workers 4
+
+

sbatch/stage2.sh

@@ -0,0 +1,45 @@
+set -e 
+
+export MASTER_PORT=$((RANDOM % 987 + 29000))
+
+accelerate launch \
+--num_processes=8 \
+--num_machines=1 \
+--dynamo_backend=no \
+--mixed_precision=fp16 \
+--gradient_accumulation_steps=8 \
+--main_process_port "${MASTER_PORT}" \
+/scripts/run_CXR_CT_train.py \
+--checkpoint "xxx.pt" \
+--name 'CXRCT_UNI_stage2' \
+--results_folder 'results_folder'  \
+--allow_partial_load \
+--fusion_module 'crossattn' \
+--pin_memory \
+--freeze_pretrained_encoders \
+--open_fusion_module \
+--use_triplet_loss 1.0 \
+--use_infoNCE_loss 0.0 \
+--positive_distance_threshold 0.25 \
+--positive_threshold 0.75 \
+--stage2 \
+--anatomy_filter  '[["bone","heart","bronchie","trachea","pleura","vertebrae","liver","aorta","spinal canal","gallbladder","clavicle","heart ascending aorta","pulmonary artery","breast","pancreas","stomach"],["lungs","heart","bones","abdomen","spine","right lung","pulmonary vasculature","aorta","airway structures","left lower lobe","right lower lobe","right upper lobe","thoracic aorta","ventricle","atrium","carina","pulmonary artery","aortic arch","bronchi","costophrenic angle","aortic valve","pericardium"]]' \
+--dataset_names 'CT_RATE' 'MIMIC_CXR' \
+--modality '3D' '2D' \
+--data_train_jsonl  '/dataset/CT_RATE/train.jsonl' '/dataset/CXR/MIMIC_CXR/train.jsonl'  \
+--data_valid_jsonl  '/dataset/CT_RATE/test.jsonl' '/dataset/CXR/MIMIC_CXR/test.jsonl' \
+--data_train_csv_dir '/dataset/CT_RATE/anatomy/train_entity' '/dataset/CXR/MIMIC_CXR/train_entity'  \
+--data_valid_csv_dir '/dataset/CT_RATE/anatomy/val_entity' '/dataset/CXR/MIMIC_CXR/test_entity' \
+--data_train_npy_dir '/dataset/CT_RATE/anatomy/train_ratescore' '/dataset/CXR/MIMIC_CXR/train_ratescore' \
+--data_valid_npy_dir '/dataset/CT_RATE/anatomy/val_ratescore' '/dataset/CXR/MIMIC_CXR/test_ratescore' \
+--train_no_aug \
+--infonce_temp \
+--modal_embedding \
+--save_results_every 500 \
+--save_model_every 6000 \
+--num_train_steps 12000 \
+--warmup_steps 1000 \
+--lr 1e-4 \
+--batch_size 1 2 \
+--local_batch_size 6 30 \
+--con_num_workers 4

scripts/README.md

@@ -0,0 +1,3 @@
+### 给一些修改配置文件的建议
+
+bash 直接执行即可。对于每一个选项都解释其含义

scripts/RadIR_Trainer_CXR_CT.py

@@ -0,0 +1,482 @@
+import os
+from pathlib import Path
+from datetime import timedelta
+import time
+from datetime import datetime
+import wandb
+
+
+import torch
+from torch import nn
+from torch.utils.data import DataLoader
+from torch.utils.tensorboard import SummaryWriter
+
+from accelerate import Accelerator
+from accelerate.utils import DistributedDataParallelKwargs, InitProcessGroupKwargs
+
+from transformer_maskgit.transformer_maskgit.optimizer import get_optimizer
+from scripts.scheduler import cosine_lr
+
+# === 导入两个 Dataset ===
+from Rad_IR.radir.data.hmclip_mimic_box import HyCoClipMimicBoxDataset_JPG, collate_hmclip_box
+
+def cycle(dl):
+    """让 DataLoader 无限循环"""
+    while True:
+        for data in dl:
+            yield data
+
+
+
+class AverageMeter:
+    """统计均值的工具类"""
+    def __init__(self):
+        self.reset()
+
+    def reset(self):
+        self.sum = 0
+        self.count = 0
+        self.val = 0
+        self.avg = 0
+
+    def update(self, val, n=1):
+        """更新一次计数"""
+        self.val = float(val)
+        self.sum += float(val) * n
+        self.count += n
+        self.avg = self.sum / self.count if self.count > 0 else 0.0
+
+
+class RadIRTrainer(nn.Module):
+    """
+    RadIR 模型训练器：
+    - 自动支持 Stage1（Box 监督）
+    """
+
+    def __init__(
+        self,
+        RADIR: nn.Module,
+        *,
+        tokenizer,
+        train_dataset,
+        valid_dataset,
+        num_train_steps,
+        current_step=0,
+        warmup_steps=1000,
+        batch_size=32,
+        lr=5e-5,
+        wd=0.0,
+        max_grad_norm=0.5,
+        save_results_every=500,
+        save_model_every=2000,
+        results_folder="./outputs",
+        num_workers=4,
+        pin_memory=False,
+        accelerate_kwargs=dict(),
+        modal_embedding=False,
+        debug=False,
+    ):
+        super().__init__()
+
+        self.best_val_loss = float('inf')
+        self.debug = debug
+        self.modal_embedding = modal_embedding
+        self.save_results_every = save_results_every
+        self.save_model_every = save_model_every
+
+        # -------------------------
+        # 判断当前是 Stage1
+        # -------------------------
+        self.is_stage1 = True   # 只用 Stage1
+        print(f"\n[Trainer 初始化] 当前模式: Stage1（Simple）\n")
+
+
+        # -------------------------
+        # 初始化 Accelerator（DDP + 混合精度）
+        # -------------------------
+        ddp_kwargs = DistributedDataParallelKwargs(find_unused_parameters=True)
+        pg_kwargs = InitProcessGroupKwargs(timeout=timedelta(seconds=36000))
+        self.accelerator = Accelerator(
+            mixed_precision="no",
+            kwargs_handlers=[ddp_kwargs, pg_kwargs],
+            **accelerate_kwargs,
+        )
+
+        self.Rad_IR = RADIR
+
+    
+
+        self.tokenizer = tokenizer
+        self.device = self.accelerator.device
+
+        # -------------------------
+        # 基本训练参数
+        # -------------------------
+        self.num_train_steps = int(num_train_steps)
+        self.current_step = int(current_step)
+        self.batch_size = batch_size
+        self.max_grad_norm = max_grad_norm
+
+        # -------------------------
+        # 优化器 + 学习率调度
+        # -------------------------
+        all_params = list(self.Rad_IR.parameters())
+        init_lr = lr[0] if isinstance(lr, (list, tuple)) else float(lr)
+
+        self.optim = get_optimizer(all_params, lr=init_lr, wd=wd)
+
+        self.lr_scheduler = cosine_lr(
+            self.optim, init_lr, warmup_steps, self.num_train_steps
+        )
+
+        # -------------------------
+        # DataLoader（根据 stage 自动切换 collate）
+        # -------------------------
+        collate_fn = collate_hmclip_box
+
+
+
+        self.train_dl = DataLoader(
+            train_dataset,
+            batch_size=batch_size,
+            shuffle=True,
+            num_workers=num_workers,
+            pin_memory=pin_memory,
+            drop_last=True,
+            collate_fn=collate_fn,
+        )
+        self.valid_dl = DataLoader(
+            valid_dataset,
+            batch_size=batch_size,
+            shuffle=False,
+            num_workers=num_workers,
+            pin_memory=pin_memory,
+            drop_last=False,
+            collate_fn=collate_fn,
+        )
+
+        # -------------------------
+        # Accelerator 包装模型、优化器、DataLoader
+        # -------------------------
+        (
+            self.Rad_IR,
+            self.optim,
+            self.train_dl,
+            self.valid_dl,
+        ) = self.accelerator.prepare(
+            self.Rad_IR,
+            self.optim,
+            self.train_dl,
+            self.valid_dl,
+        )
+
+        self.train_iter = cycle(self.train_dl)
+
+        # -------------------------
+        # 日志 & 输出目录
+        # -------------------------
+        self.results_folder = Path(results_folder)
+        self.results_folder.mkdir(parents=True, exist_ok=True)
+
+        tb_path = self.results_folder / "tensorboard"
+        tb_path.mkdir(parents=True, exist_ok=True)
+        self.tb_writer = SummaryWriter(str(tb_path))
+
+        # -------------------------
+        # 统计指标
+        # -------------------------
+        self.uncon_loss_m = AverageMeter()                  # 总 loss
+        self.it_ce_loss_m = AverageMeter()                  # image-text CE
+        self.it_entailment_loss_m = AverageMeter()         # image-text entailment
+        self.ii_text_image_entailment_loss_m = AverageMeter()  # image-image entailment
+
+        print(f"[Trainer] 已准备完毕，总训练步数：{self.num_train_steps}")
+        # -------------------------
+        # 时间统计
+        # -------------------------
+        self.train_start_time = None
+        self.step_time_m = AverageMeter()
+
+    # =======================================================
+    # 前向：根据 Stage 自动切换分支
+    # =======================================================
+    def _forward_batch(self, batch):
+        imgs = batch["imgs"].to(self.device)
+        caption_ids = batch["caption_ids"].to(self.device)
+        token_type_ids = batch["token_type_ids"].to(self.device)
+        attention_mask = batch["attention_mask"].to(self.device)
+
+        phrase_text_mask = batch["phrase_text_mask"].to(self.device)
+        phrases_bboxes_mask = batch["phrases_bboxes_mask"].to(self.device)
+
+        text_batch = {
+            "caption_ids": caption_ids,
+            "token_type_ids": token_type_ids,
+            "attention_mask": attention_mask,
+        }
+
+
+
+
+
+        loss, ce_loss, entail_loss, global_entail_loss = self.Rad_IR(
+            text=text_batch,
+            image=imgs,
+            device=self.device,
+            is_condition=False,
+            return_loss=True,
+            modal_embedding=self.modal_embedding,
+            phrase_text_mask=phrase_text_mask,          
+            phrases_bboxes_mask=phrases_bboxes_mask,   
+            return_latents=True,
+        )
+
+        return loss, ce_loss, entail_loss, global_entail_loss
+
+    # =======================================================
+    # 主训练循环
+    # =======================================================
+    def train_model(self, resume_path=None):
+
+        # ======================
+        # 先尝试 Resume 恢复训练
+        # ======================
+        if resume_path is not None:
+            self.load_checkpoint(resume_path)
+
+        self.Rad_IR.train()
+        self.train_start_time = time.time()
+
+        torch.autograd.set_detect_anomaly(True) 
+
+        for step in range(self.current_step, self.num_train_steps):
+            self.current_step = step
+            step_start_time = time.time()
+
+            self.lr_scheduler.step()
+            # ============================
+            # ✅ Stage 1
+            # ============================
+            with self.accelerator.accumulate(self.Rad_IR):
+
+                batch = next(self.train_iter)
+                loss, ce_loss, entail_loss, global_entail_loss = self._forward_batch(batch)
+                #加检查1
+                for name, val in {
+                    "loss": loss,
+                    "ce": ce_loss,
+                    "entail": entail_loss,
+                    "global_entail": global_entail_loss,
+                }.items():
+                    if not torch.isfinite(val):
+                        print(f"[NaN LOSS] step={step}, {name}")
+                        self._dump_bad_batch(batch, step)
+                        raise RuntimeError("NaN in loss")
+                #加检查2
+                self.optim.zero_grad()
+                self.accelerator.backward(loss)
+                #加检查1
+                for n, p in self.Rad_IR.named_parameters():
+                    if p.grad is not None and not torch.isfinite(p.grad).all():
+                        print(f"[NaN GRAD] step={step}, param={n}")
+                        raise RuntimeError("NaN in grad")
+                #加检查2             
+
+
+            # ============================
+            # ✅ 只有在 sync_gradients 才 step
+            # ============================
+            if self.accelerator.sync_gradients:
+
+                if self.max_grad_norm is not None:
+                    self.accelerator.clip_grad_norm_(
+                        self.Rad_IR.parameters(), self.max_grad_norm
+                    )
+
+                self.optim.step()
+
+
+                self.optim.zero_grad(set_to_none=True)
+
+            # ============================
+            # ✅ 时间统计
+            # ============================
+            if torch.cuda.is_available():
+                torch.cuda.synchronize()
+            step_time = time.time() - step_start_time
+            self.step_time_m.update(step_time)
+
+            # ✅ loss 记录（已有的 AverageMeter）
+            # ===== 从 GPU 拿回 4 个 loss =====
+            loss_val = float(loss.detach().cpu())
+            ce_loss_val = float(ce_loss.detach().cpu())
+            entail_loss_val = float(entail_loss.detach().cpu())
+            ii_entail_loss_val = float(global_entail_loss.detach().cpu())
+
+            # ===== 更新 4 个 meter =====
+            self.uncon_loss_m.update(loss_val, 1)
+            self.it_ce_loss_m.update(ce_loss_val, 1)
+            self.it_entailment_loss_m.update(entail_loss_val, 1)
+            self.ii_text_image_entailment_loss_m.update(ii_entail_loss_val, 1)
+
+
+            # ============================
+            # ✅ 日志
+            # ============================
+            if self.accelerator.is_main_process and (step + 1) % 10 == 0:
+                eta_seconds = int(self.step_time_m.avg * (self.num_train_steps - step - 1))
+                eta_str = str(timedelta(seconds=eta_seconds))
+
+                print(
+                    f"[Step {step+1}/{self.num_train_steps}] "
+                    f"loss={self.uncon_loss_m.avg:.4f}, "
+                    f"ce={self.it_ce_loss_m.avg:.4f}, "
+                    f"entail={self.it_entailment_loss_m.avg:.4f}, "
+                    f"step_time={self.step_time_m.avg:.3f}s, ETA={eta_str}"
+                )
+                wandb.log({
+                    "train/loss": self.uncon_loss_m.avg,
+                    "train/ce_loss": self.it_ce_loss_m.avg,
+                    "train/entail_loss": self.it_entailment_loss_m.avg,
+                    "train/global_entail_loss": self.ii_text_image_entailment_loss_m.avg,
+                    "lr": self.optim.param_groups[0]["lr"],
+                    "step_time": self.step_time_m.avg,
+                    "step": step + 1
+                }, step=step + 1)
+
+
+            # ============================
+            # ✅ 验证
+            # ============================
+            if (step + 1) % self.save_results_every == 0:
+                val_loss = self.evaluate(step + 1)
+
+                # ========= 更新 best =========
+                if self.accelerator.is_main_process:
+                    if val_loss < self.best_val_loss:
+                        self.best_val_loss = val_loss
+                        self.save_checkpoint(step + 1, is_best=True)
+
+            # ============================
+            # ✅ 存模型
+            # ============================
+            if (step + 1) % self.save_model_every == 0:
+                self.save_checkpoint(step + 1)
+
+        # ✅ 最终评估 + 保存
+        self.evaluate(self.num_train_steps)
+        self.save_checkpoint(self.num_train_steps, is_final=True)
+
+
+    # =======================================================
+    # 验证集评估
+    # =======================================================
+    @torch.no_grad()
+    def evaluate(self, step):
+        self.Rad_IR.eval()
+        eval_start_time = time.time()   # ✅ 验证开始计时
+        loss_m = AverageMeter()
+        ce_m = AverageMeter()
+        entail_m = AverageMeter()
+        g_entail_m = AverageMeter()
+
+        for batch in self.valid_dl:
+            loss, ce_loss, entail_loss, g_entail_loss = self._forward_batch(batch)
+
+            bs = batch["imgs"].size(0)
+            loss_m.update(loss.item(), bs)
+            ce_m.update(ce_loss.item(), bs)
+            entail_m.update(entail_loss.item(), bs)
+            g_entail_m.update(g_entail_loss.item(), bs)
+
+        if self.accelerator.is_main_process:
+            print(
+                f"[Eval step {step}] val_loss={loss_m.avg:.4f}, "
+                f"val_ce={ce_m.avg:.4f}, val_entail={entail_m.avg:.4f}"
+            )
+
+            self.tb_writer.add_scalar("val/loss", loss_m.avg, step)
+            self.tb_writer.add_scalar("val/ce_loss", ce_m.avg, step)
+
+        self.Rad_IR.train()
+        eval_time = time.time() - eval_start_time
+        eval_time_str = str(timedelta(seconds=int(eval_time)))
+
+        if self.accelerator.is_main_process:
+            print(f"[Eval Done] step={step} 耗时：{eval_time_str}")
+            self.tb_writer.add_scalar("time/eval_time", eval_time, step)
+
+            wandb.log({
+            "val/loss": loss_m.avg,
+            "val/ce_loss": ce_m.avg,
+            "val/entail_loss": entail_m.avg,
+            "val/global_entail_loss": g_entail_m.avg,
+            }, step=step)
+            
+            return loss_m.avg
+
+    # =======================================================
+    # 模型保存
+    # =======================================================
+    def save_checkpoint(self, step, is_final=False, is_best=False):
+        if not self.accelerator.is_main_process:
+            return
+
+        if is_final:
+            ckpt_name = "radir_final.pt"
+        elif is_best:
+            ckpt_name = "radir_best.pt"
+        else:
+            ckpt_name = "radir_last.pt"   # 覆盖最新
+
+        ckpt_path = self.results_folder / ckpt_name
+
+        state = {
+            "model": self.accelerator.get_state_dict(self.Rad_IR),
+            "optimizer": self.optim.state_dict(),
+            "scheduler": self.lr_scheduler.state_dict(),   
+            "step": step,
+            "best_val_loss": self.best_val_loss,           
+        }
+
+        torch.save(state, ckpt_path)
+        print(f"[Checkpoint] 保存: {ckpt_path}")
+
+
+    # =======================================================
+    # 模型恢复
+    # =======================================================     
+    def load_checkpoint(self, ckpt_path):
+        if not os.path.exists(ckpt_path):
+            raise FileNotFoundError(f"Checkpoint 不存在: {ckpt_path}")
+
+        print(f"[Resume] 从 {ckpt_path} 恢复训练")
+
+        ckpt = torch.load(ckpt_path, map_location="cpu")
+
+        # 1️⃣ 恢复模型
+        self.accelerator.unwrap_model(self.Rad_IR).load_state_dict(ckpt["model"])
+
+        # 2️⃣ 恢复优化器
+        self.optim.load_state_dict(ckpt["optimizer"])
+
+        # 3️⃣ 恢复 scheduler（如果存在）
+        if "scheduler" in ckpt:
+            self.lr_scheduler.load_state_dict(ckpt["scheduler"])
+            print("[Resume] Scheduler 已恢复")
+        else:
+            print("[Resume] ⚠️ 未发现 scheduler，使用新的 scheduler")
+
+        # 4️⃣ 恢复训练进度
+        self.current_step = ckpt["step"]
+
+        # 5️⃣ 恢复 best_val_loss
+        self.best_val_loss = ckpt.get("best_val_loss", float("inf"))
+
+        print("resume lr =", self.optim.param_groups[0]["lr"])
+
+
+        print(
+            f"[Resume] 恢复完成：step = {self.current_step}, "
+            f"best_val_loss = {self.best_val_loss:.4f}"
+        )

scripts/build_traversal_concepts.py

@@ -0,0 +1,620 @@
+#!/usr/bin/env python3
+"""
+Build a MERU-style multi-granularity CXR concept pool for image traversals.
+
+Design (why this is more scientific than mining raw report sentences):
+
+1. MERU Fig.5 uses short noun/adj captions at multiple granularities, not long
+   free-text captions. Long MIMIC boilerplate ("No acute osseous abnormalities
+   are identified") dominate nearest-neighbor retrieval and break hierarchy.
+
+2. Hierarchy mirrors radiology reporting + hyperbolic entailment intuition:
+     modality / study  →  organ system  →  anatomy  →  finding / attribute
+   Closer to [ROOT] should be more abstract.
+
+3. Finding vocabulary is anchored to CheXpert-style labels used in this repo
+   (see infer_disease_clusters.DISEASE_NAMES) plus common CXR observations,
+   with short "no X" counterparts so Normal cases are representable.
+
+4. Additionally mine short, high-frequency clinical phrases from train.jsonl +
+   valid.jsonl (with filters), and ground a few phrases in the current query
+   images' Findings so traversal concepts can match the paired reports.
+"""
+
+from __future__ import annotations
+
+import json
+import re
+from pathlib import Path
+
+ROOT = Path(__file__).resolve().parents[1]
+OUT = ROOT / "scripts" / "traversal_concepts.txt"
+
+
+def _dedupe(items: list[str]) -> list[str]:
+    seen: set[str] = set()
+    out: list[str] = []
+    for s in items:
+        s = " ".join(s.strip().split())
+        if not s or s.startswith("#"):
+            continue
+        key = s.lower()
+        if key in seen:
+            continue
+        seen.add(key)
+        out.append(s)
+    return out
+
+
+def build_sections() -> list[tuple[str, list[str]]]:
+    # --- L0: modality / study (most abstract) ---
+    l0 = [
+        "radiograph",
+        "x-ray",
+        "imaging study",
+        "chest radiograph",
+        "chest x-ray",
+        "chest imaging",
+        "thoracic imaging",
+        "radiographic examination",
+        "frontal chest radiograph",
+        "lateral chest radiograph",
+        "PA chest radiograph",
+        "AP chest radiograph",
+        "portable chest radiograph",
+    ]
+
+    # --- L1: organ systems / regions (bridging toward anatomy) ---
+    l1 = [
+        "chest",
+        "thorax",
+        "thoracic cavity",
+        "chest wall",
+        "lungs",
+        "pulmonary system",
+        "cardiopulmonary system",
+        "cardiovascular structures",
+        "cardiomediastinal structures",
+        "pleural space",
+        "pleural spaces",
+        "mediastinum",
+        "central airways",
+        "bony thorax",
+        "upper abdomen",
+        "neck base",
+    ]
+
+    # --- L2: anatomy (mid → fine; Chest ImaGenome–style regions) ---
+    l2 = [
+        # lungs / lobes / zones
+        "right lung",
+        "left lung",
+        "both lungs",
+        "lung",
+        "lung parenchyma",
+        "lung bases",
+        "lung apex",
+        "lung apices",
+        "upper lobes",
+        "lower lobes",
+        "middle lobe",
+        "lingula",
+        "right upper lobe",
+        "right middle lobe",
+        "right lower lobe",
+        "left upper lobe",
+        "left lower lobe",
+        "right lung base",
+        "left lung base",
+        "bilateral lung bases",
+        "perihilar region",
+        "perihilar regions",
+        "peripheral lung",
+        "central lung",
+        # airways / hilum
+        "hilum",
+        "hila",
+        "right hilum",
+        "left hilum",
+        "hilar structures",
+        "trachea",
+        "carina",
+        "airways",
+        "main bronchi",
+        "bronchi",
+        # heart / mediastinum / vessels
+        "heart",
+        "cardiac silhouette",
+        "heart border",
+        "left heart border",
+        "right heart border",
+        "cardiomediastinal silhouette",
+        "mediastinal contours",
+        "mediastinal silhouette",
+        "hilar contours",
+        "aorta",
+        "aortic arch",
+        "descending aorta",
+        "ascending aorta",
+        "pulmonary artery",
+        "pulmonary vasculature",
+        "pulmonary vessels",
+        "superior vena cava",
+        # pleura / diaphragm
+        "pleura",
+        "pleural surface",
+        "costophrenic angle",
+        "costophrenic angles",
+        "right costophrenic angle",
+        "left costophrenic angle",
+        "hemidiaphragm",
+        "right hemidiaphragm",
+        "left hemidiaphragm",
+        "diaphragm",
+        "diaphragmatic contour",
+        # bones / soft tissue
+        "ribs",
+        "rib",
+        "clavicle",
+        "clavicles",
+        "scapula",
+        "sternum",
+        "spine",
+        "thoracic spine",
+        "vertebral bodies",
+        "soft tissues",
+        "breast shadows",
+        "shoulder girdle",
+    ]
+
+    # --- L3: CheXpert-aligned findings + localized short phrases ---
+    findings_pos = [
+        # Project primary labels
+        "pneumothorax",
+        "pulmonary edema",
+        "edema",
+        "pneumonia",
+        "pleural effusion",
+        "effusion",
+        "cardiomegaly",
+        "atelectasis",
+        # CheXpert / MIMIC common
+        "consolidation",
+        "focal consolidation",
+        "patchy consolidation",
+        "lung opacity",
+        "airspace opacity",
+        "airspace disease",
+        "enlarged cardiomediastinum",
+        "widened mediastinum",
+        "lung lesion",
+        "pulmonary nodule",
+        "nodule",
+        "mass",
+        "fracture",
+        "rib fracture",
+        "support devices",
+        # Localized / severity-ish findings (short)
+        "bibasilar atelectasis",
+        "basilar atelectasis",
+        "left lower lobe atelectasis",
+        "right lower lobe atelectasis",
+        "right upper lobe atelectasis",
+        "left upper lobe atelectasis",
+        "platelike atelectasis",
+        "compressive atelectasis",
+        "left pleural effusion",
+        "right pleural effusion",
+        "bilateral pleural effusion",
+        "small pleural effusion",
+        "large pleural effusion",
+        "layering pleural effusion",
+        "left pneumothorax",
+        "right pneumothorax",
+        "small pneumothorax",
+        "large pneumothorax",
+        "apical pneumothorax",
+        "mild pulmonary edema",
+        "pulmonary vascular congestion",
+        "vascular congestion",
+        "cephalization",
+        "interstitial edema",
+        "interstitial markings",
+        "increased interstitial markings",
+        "kerley b lines",
+        "infiltrate",
+        "opacity",
+        "basilar opacity",
+        "bibasilar opacities",
+        "right lower lobe opacity",
+        "left lower lobe opacity",
+        "perihilar opacity",
+        "emphysema",
+        "hyperinflation",
+        "fibrosis",
+        "scarring",
+        "pleural scarring",
+        "calcified granuloma",
+        "granuloma",
+        "calcification",
+        "hilar calcification",
+        "scoliosis",
+        "kyphosis",
+        "low lung volumes",
+        "elevated hemidiaphragm",
+        "elevated right hemidiaphragm",
+        "elevated left hemidiaphragm",
+        "blunted costophrenic angle",
+        "blunted costophrenic angles",
+        "hilar lymphadenopathy",
+        "lymphadenopathy",
+        "bronchiectasis",
+        "cavity",
+        "cavitation",
+        "pleural thickening",
+        "subcutaneous emphysema",
+    ]
+
+    # Short negations / normal status (balanced; include mid-level normals)
+    findings_neg = [
+        "no finding",
+        "normal study",
+        "normal chest",
+        "unremarkable chest",
+        "unremarkable study",
+        "clear lungs",
+        "lungs are clear",
+        "lungs remain clear",
+        "clear lung fields",
+        "no focal opacity",
+        "no focal airspace opacity",
+        "no pneumothorax",
+        "no pleural effusion",
+        "no effusion",
+        "no pulmonary edema",
+        "no edema",
+        "no pneumonia",
+        "no consolidation",
+        "no focal consolidation",
+        "no atelectasis",
+        "no cardiomegaly",
+        "no fracture",
+        "normal heart size",
+        "normal cardiac silhouette",
+        "normal cardiomediastinal silhouette",
+        "normal mediastinal contours",
+        "normal hilar contours",
+        "normal pulmonary vasculature",
+        "no vascular congestion",
+        "costophrenic angles are clear",
+        "costophrenic angles are sharp",
+        "sharp costophrenic angles",
+        "no acute cardiopulmonary process",
+        "no acute abnormality",
+        "no acute disease",
+        "stable chest",
+        "unchanged chest",
+    ]
+
+    # Severity / laterality composed phrases
+    attrs = [
+        "mild cardiomegaly",
+        "moderate cardiomegaly",
+        "severe cardiomegaly",
+        "top normal heart size",
+        "borderline cardiomegaly",
+        "mild pulmonary edema",
+        "moderate pulmonary edema",
+        "severe pulmonary edema",
+        "small pleural effusion",
+        "moderate pleural effusion",
+        "large pleural effusion",
+        "trace pleural effusion",
+        "tiny pleural effusion",
+        "bilateral pleural effusion",
+        "bilateral atelectasis",
+        "mild atelectasis",
+        "mild bibasilar atelectasis",
+        "stable cardiomegaly",
+        "stable pleural effusion",
+        "stable atelectasis",
+        "stable opacity",
+        "chronic opacity",
+        "acute opacity",
+        "improving edema",
+        "worsening edema",
+        "improving effusion",
+        "worsening effusion",
+    ]
+
+    # Devices / tubes / lines
+    devices = [
+        "endotracheal tube",
+        "tracheostomy tube",
+        "chest tube",
+        "right chest tube",
+        "left chest tube",
+        "nasogastric tube",
+        "orogastric tube",
+        "PICC line",
+        "central venous catheter",
+        "internal jugular catheter",
+        "subclavian catheter",
+        "swan ganz catheter",
+        "pacemaker",
+        "defibrillator",
+        "icd lead",
+        "sternal wires",
+        "prosthetic valve",
+        "aortic valve prosthesis",
+        "mitral valve prosthesis",
+        "postoperative changes",
+        "median sternotomy",
+    ]
+
+    return [
+        ("L0 modality / study (most abstract)", l0),
+        ("L1 organ system / region", l1),
+        ("L2 anatomy", l2),
+        ("L3 findings (CheXpert-aligned, short phrases)", findings_pos),
+        ("L3 normal / negation status (short)", findings_neg),
+        ("L4 attributes / severity", attrs),
+        ("devices / support", devices),
+    ]
+
+
+# ---------------------------------------------------------------------------
+# Mine short phrases from real MIMIC reports (train.jsonl + valid.jsonl)
+# ---------------------------------------------------------------------------
+
+_CLINICAL_RE = re.compile(
+    r"\b(lung|lungs|heart|cardiac|mediastinal|mediastinum|pleural|effusion|"
+    r"pneumothorax|atelectasis|atelectatic|edema|consolidation|opacity|"
+    r"cardiomegaly|granuloma|granulomatous|nodule|diaphragm|hemidiaphragm|"
+    r"costophrenic|vascular|vasculature|pneumonia|infiltrate|volume|volumes|"
+    r"hilar|hila|trachea|aorta|clear|normal|stable|blunting|elevated|"
+    r"hyperinflated|cephalization|opacity|opacities|scar|scarring)\b",
+    re.I,
+)
+_REJECT_RE = re.compile(
+    r"\b(comparison|compared|previous studies|prior study|___|please|"
+    r"recommend|correlate|history|indication|technique|year.?old|"
+    r"male|female|osseous|patient|dr\.|dictated|electronically|"
+    r"dating back|better seen|evaluated on ct|cts dating|"
+    r"ventriculoperitoneal|glenohumeral|thoracolumbar|dish\b|"
+    r"free air|below the right hemidiaphragm|abdomen|"
+    r"concerning for pneumonia|well-expanded|well inflated|"
+    r"unchanged|again visualized|displaced fracture|fracture is seen)\b",
+    re.I,
+)
+
+
+def _normalize_phrase(s: str) -> str:
+    s = s.strip().rstrip(".").lower()
+    s = re.sub(r"\s+", " ", s)
+    s = re.sub(r"^(there (is|are)|the|a|an)\s+", "", s)
+    s = re.sub(
+        r"\s+(is|are|was|were)?\s*(seen|present|identified|noted|visualized)\s*$",
+        "",
+        s,
+    )
+    s = re.sub(r"\s+", " ", s).strip(" ,;")
+    # Prefer short atomic forms for common normals
+    replacements = {
+        "lungs are grossly clear": "lungs are clear",
+        "lungs appear clear": "lungs are clear",
+        "lungs remain clear": "lungs are clear",
+        "heart size is within normal limits": "heart size is normal",
+        "heart is normal in size": "heart size is normal",
+        "heart is at the upper limits of normal size": "top normal heart size",
+        "no pneumothoraces": "no pneumothorax",
+        "no pleural effusions": "no pleural effusion",
+        "bibasilar atelectatic changes": "bibasilar atelectasis",
+        "elevation of the mid to posterior left hemidiaphragm": (
+            "elevated left hemidiaphragm"
+        ),
+        "elevation of the left hemidiaphragm": "elevated left hemidiaphragm",
+        "minimal blunting of the left costophrenic angle": (
+            "blunted left costophrenic angle"
+        ),
+        "rounded calcified nodule": "calcified nodule",
+        "no signs for acute pulmonary edema": "no pulmonary edema",
+        "no definite consolidation": "no consolidation",
+        "lungs are otherwise without a focal consolidation": "no focal consolidation",
+        "cardiac and mediastinal silhouettes are unremarkable": (
+            "normal cardiac and mediastinal silhouettes"
+        ),
+        "cardiac and mediastinal silhouettes are stable": (
+            "stable cardiac and mediastinal silhouettes"
+        ),
+        "cardiac and mediastinal contours appears stable": (
+            "stable cardiac and mediastinal contours"
+        ),
+        "cardiac and mediastinal contours are stable": (
+            "stable cardiac and mediastinal contours"
+        ),
+    }
+    return replacements.get(s, s)
+
+
+def _split_sentences(text: str) -> list[str]:
+    text = re.sub(r"\s+", " ", text.strip())
+    parts = re.split(r"(?<=[.!?])\s+", text)
+    return [p.strip().rstrip(".") for p in parts if p.strip()]
+
+
+def mine_report_phrases(
+    jsonl_paths: list[Path],
+    min_freq: int = 40,
+    max_new: int = 100,
+) -> list[str]:
+    """
+    High-frequency short clinical sentences from real reports, normalized to
+    short forms. Rejects long multi-finding boilerplate and non-clinical junk.
+    """
+    from collections import Counter
+
+    cnt: Counter[str] = Counter()
+    for path in jsonl_paths:
+        if not path.is_file():
+            continue
+        with path.open(encoding="utf-8") as f:
+            for line in f:
+                text = str(json.loads(line).get("text", ""))
+                for sent in _split_sentences(text):
+                    if not _CLINICAL_RE.search(sent):
+                        continue
+                    if _REJECT_RE.search(sent):
+                        continue
+                    # Skip laundry-list negatives (old attractor templates)
+                    if sent.count(",") >= 2 or sent.lower().count(" or ") >= 2:
+                        continue
+                    norm = _normalize_phrase(sent)
+                    n_words = len(norm.split())
+                    if n_words < 2 or n_words > 7:
+                        continue
+                    if len(norm) > 55:
+                        continue
+                    if not _CLINICAL_RE.search(norm):
+                        continue
+                    if _REJECT_RE.search(norm):
+                        continue
+                    cnt[norm] += 1
+
+    ranked = [(s, c) for s, c in cnt.most_common() if c >= min_freq]
+    out: list[str] = []
+    for s, _ in ranked:
+        out.append(s)
+        if len(out) >= max_new:
+            break
+    return out
+
+
+def phrases_from_query_reports(
+    jsonl_paths: list[Path],
+    dicom_ids: list[str],
+) -> list[str]:
+    """
+    Short concepts grounded in Findings of the current traversal query images.
+    Ensures the pool can express what those reports actually say.
+    """
+    need = set(dicom_ids)
+    found: dict[str, str] = {}
+    for path in jsonl_paths:
+        if not path.is_file() or len(found) >= len(need):
+            continue
+        with path.open(encoding="utf-8") as f:
+            for line in f:
+                row = json.loads(line)
+                dicom = Path(str(row.get("img_path", ""))).stem
+                if dicom in need and dicom not in found:
+                    found[dicom] = str(row.get("text", ""))
+
+    # Hand-curated short extracts tied to those reports (not free long sentences)
+    grounded: list[str] = [
+        # 46583e03 — normal / clear
+        "heart size is normal",
+        "lungs are clear",
+        "no consolidation",
+        "no pleural effusion",
+        "no pneumothorax",
+        "no pulmonary edema",
+        # 1f199273 — elevated hemidiaphragm / CPA blunting; NO definite effusion
+        "no focal consolidation",
+        "elevated left hemidiaphragm",
+        "blunted left costophrenic angle",
+        "no definite pleural effusion",
+        "normal cardiac and mediastinal silhouettes",
+        # e5ba5704 — calcified granuloma; no acute airspace disease
+        "calcified granuloma",
+        "calcified nodule",
+        "right lung base",
+        "granulomatous disease",
+        "stable cardiac and mediastinal silhouettes",
+        # d1552af1 — low volumes + bibasilar atelectasis
+        "lung volumes are low",
+        "bibasilar atelectasis",
+        "stable cardiac and mediastinal contours",
+        "no focal consolidation",
+    ]
+    # Also pull any short clinical sentences directly from these reports
+    for text in found.values():
+        for sent in _split_sentences(text):
+            if not _CLINICAL_RE.search(sent) or _REJECT_RE.search(sent):
+                continue
+            if sent.count(",") >= 2 or sent.lower().count(" or ") >= 2:
+                continue
+            norm = _normalize_phrase(sent)
+            if 2 <= len(norm.split()) <= 7 and len(norm) <= 55:
+                grounded.append(norm)
+    return _dedupe(grounded)
+
+
+def main() -> None:
+    data_dir = ROOT / "dataset"
+    jsonls = [data_dir / "train.jsonl", data_dir / "valid.jsonl"]
+    query_dicoms = [
+        "46583e03-f42311f4-87dca60a-c4c12f22-0fe13c7f",
+        "1f199273-2f4b7e8c-9041b5c2-18e7ca21-26142e3c",
+        "e5ba5704-ce2f09d3-e28fe2a2-8a9aca96-86f4966a",
+        "d1552af1-5b159d3e-4058cc59-8af87caf-375f46e7",
+    ]
+
+    sections = build_sections()
+    mined = mine_report_phrases(jsonls, min_freq=40, max_new=90)
+    grounded = phrases_from_query_reports(jsonls, query_dicoms)
+    sections.append(
+        (
+            "L3 mined from train+valid reports (short, freq>=40, de-templated)",
+            mined,
+        )
+    )
+    sections.append(
+        (
+            "L3 grounded in current traversal query Findings",
+            grounded,
+        )
+    )
+
+    lines: list[str] = [
+        "# MERU-style multi-granularity CXR concept pool for image traversals.",
+        "#",
+        "# Construction (see scripts/build_traversal_concepts.py):",
+        "#   - Curated coarse→fine hierarchy (modality → organ → anatomy → finding).",
+        "#   - Plus short phrases mined from dataset/train.jsonl + valid.jsonl",
+        "#     (freq filter, clinical keep-list, reject long multi-finding templates).",
+        "#   - Plus short concepts grounded in current traversal query reports.",
+        "#",
+        "# Regenerate: python scripts/build_traversal_concepts.py",
+        "",
+    ]
+    seen: set[str] = set()
+    n = 0
+    n_mined = 0
+    n_grounded = 0
+    for title, items in sections:
+        lines.append(f"# --- {title} ---")
+        added_here = 0
+        for s in _dedupe(items):
+            key = s.lower()
+            if key in seen:
+                continue
+            seen.add(key)
+            lines.append(s)
+            n += 1
+            added_here += 1
+        if "mined from train" in title:
+            n_mined = added_here
+        if "grounded in current" in title:
+            n_grounded = added_here
+        lines.append("")
+
+    OUT.write_text("\n".join(lines).rstrip() + "\n", encoding="utf-8")
+    print(
+        f"Wrote {OUT} ({n} unique concepts, excl. [ROOT]; "
+        f"+{n_mined} mined, +{n_grounded} query-grounded new)"
+    )
+
+
+if __name__ == "__main__":
+    main()

scripts/checkdata.py

@@ -0,0 +1,87 @@
+import os
+from tqdm import tqdm
+import torch
+from transformers import BertTokenizer
+import sys   
+
+
+PROJECT_ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))
+if PROJECT_ROOT not in sys.path:
+    sys.path.insert(0, PROJECT_ROOT)
+
+from Rad_IR.radir.data.hmclip_mimic_box import HyCoClipMimicBoxDataset_JPG, collate_hmclip_box
+
+
+def main():
+    tokenizer = BertTokenizer.from_pretrained("bert-base-uncased")
+
+    dataset = HyCoClipMimicBoxDataset_JPG(
+        split="train",
+        tokenizer=tokenizer,
+        dataset_root="/mnt/data-3/users/zhangzeyu/HMClip/RadIR/dataset",
+        cig_matched_dir="/mnt/data-3/users/zhangzeyu/HMClip/RadIR/dataset/processed/matched",
+        caption_file="valid.jsonl",
+        transform=None,
+    )
+
+    bad_empty = 0
+    bad_short = 0
+    bad_phrase = 0
+    bad_image = 0
+
+    for i in tqdm(range(len(dataset))):
+        try:
+            img, caps, cap_len, img_path, phrase, ph_mask, bbx_mask = dataset[i]
+        except Exception as e:
+            print(f"\n❌ load error at {i}: {e}")
+            bad_image += 1
+            continue
+
+        input_ids = caps["input_ids"][0]
+        attention_mask = caps["attention_mask"][0]
+
+        # 1️⃣ attention_mask 全 0
+        if attention_mask.sum().item() == 0:
+            bad_empty += 1
+            print("\n❌ Empty caption")
+            print("index:", i)
+            print("img_path:", img_path)
+            continue
+
+        # 2️⃣ token 太少
+        if (input_ids != 0).sum().item() < 3:
+            bad_short += 1
+            print("\n⚠️ Too short caption")
+            print("index:", i)
+            print("img_path:", img_path)
+
+        # 3️⃣ phrase / mask 不一致
+        if phrase == "" and ph_mask.sum().item() != 0:
+            bad_phrase += 1
+            print("\n❌ phrase empty but mask not zero")
+            print("index:", i)
+            print("img_path:", img_path)
+
+        if phrase != "" and ph_mask.sum().item() == 0:
+            bad_phrase += 1
+            print("\n❌ phrase exists but mask all zero")
+            print("index:", i)
+            print("phrase:", phrase)
+            print("img_path:", img_path)
+
+        # 4️⃣ image NaN / Inf（保险）
+        if not torch.isfinite(img).all():
+            bad_image += 1
+            print("\n❌ NaN/Inf in image")
+            print("index:", i)
+            print("img_path:", img_path)
+
+    print("\n========== SUMMARY ==========")
+    print("Total samples:", len(dataset))
+    print("Empty captions:", bad_empty)
+    print("Too short captions:", bad_short)
+    print("Phrase issues:", bad_phrase)
+    print("Bad images:", bad_image)
+
+if __name__ == "__main__":
+    main()

scripts/conditional_dataset_evaluate.py

@@ -0,0 +1,584 @@
+import os
+import glob
+import json
+import torch
+import pandas as pd
+import numpy as np
+from PIL import Image
+from torch.utils.data import Dataset, Subset
+import torchvision.transforms as transforms
+from functools import partial
+import torch.nn.functional as F
+import nibabel as nib
+import tqdm
+import random
+import time
+
+
+def custom_data_split(dataset, rank, world_size):
+    # 计算每个进程应该处理的数据集的范围
+    anatomy_split, anatomy_name = dataset.get_anatomy_split() # [[0, 10], [10, 20], ...]  # 第一个anatomy的数据对应了0～10（左闭右开）
+    anatomy_split_size = [split[1] - split[0] for split in anatomy_split]
+    
+    sorted_indices = sorted(range(len(anatomy_split_size)), key=lambda i: anatomy_split_size[i], reverse=True)
+    anatomy_split_size = [anatomy_split_size[i] for i in sorted_indices]
+    anatomy_split = [anatomy_split[i] for i in sorted_indices]
+    anatomy_name = [anatomy_name[i] for i in sorted_indices]
+    
+    anatomy_in_partitions = [[] for _ in range(world_size)]
+    split_in_partitions = [[] for _ in range(world_size)]
+    for i, (name, split) in enumerate(zip(anatomy_name, anatomy_split)):
+        anatomy_in_partitions[i % world_size].append(name)
+        split_in_partitions[i % world_size].append(split)
+    
+    idx_in_this_partition = []
+    anatomy_in_this_partition = anatomy_in_partitions[rank]
+    split_in_this_partition = split_in_partitions[rank]
+    for split in split_in_this_partition:
+        idx_in_this_partition += list(range(split[0], split[1]))
+    
+    # 使用 Subset 来选择数据集的一部分
+    subset = Subset(dataset, idx_in_this_partition)
+
+    return subset, anatomy_in_this_partition
+
+def resize_array(array, current_spacing, target_spacing):
+    """
+    Resize the array to match the target spacing.
+
+    Args:
+    array (torch.Tensor): Input array to be resized.
+    current_spacing (tuple): Current voxel spacing (z_spacing, xy_spacing, xy_spacing).
+    target_spacing (tuple): Target voxel spacing (target_z_spacing, target_x_spacing, target_y_spacing).
+
+    Returns:
+    np.ndarray: Resized array.
+    """
+    # Calculate new dimensions
+    original_shape = array.shape[2:]
+    scaling_factors = [
+        current_spacing[i] / target_spacing[i] for i in range(len(original_shape))
+    ]
+    new_shape = [
+        int(original_shape[i] * scaling_factors[i]) for i in range(len(original_shape))
+    ]
+    # Resize the array
+    resized_array = F.interpolate(array, size=new_shape, mode='trilinear', align_corners=False).cpu().numpy()
+    return resized_array
+modality_dict = {'3D':0,'2D':1}
+class Conditional_ImageReportDataset_Eval(Dataset):
+    def __init__(self, jsonl_file, csv_file_dir, npy_file_dir, anatomy_filter,modality='2D', max_samples=10000,modal_embedding=False):
+        self.anatomy_filter = anatomy_filter
+        self.modality = modality
+        self.modal_embedding = modal_embedding
+        self.id2image_path = self.prepare_image_paths(jsonl_file)
+        self.id2image = {}
+        self.anatomy2id_ls = self.prepare_anatomy_data(csv_file_dir, max_samples)   # DEBUG 每个类只取max_samples个测试
+        self.anatomy2simi_tab = self.prepare_similarity_table(npy_file_dir, max_samples)
+        self.check_integrity(self.anatomy2id_ls, self.anatomy2simi_tab)
+        
+        # self.anatomy_ls = list(self.anatomy2id_ls.keys()) # NOTE 不按照anatomy为单位采
+        # NOTE anatomy2id_ls是形如{'lung':[a, b, c, d ...], ...}的字典，展开成id_anatomy_ls=[['lung', a], ['lung', b] ...]这样的list
+        self.id_antomy_ls = [[anatomy, img_id] for anatomy, id_list in self.anatomy2id_ls.items() for img_id in id_list]
+        # self.id_antomy_ls = [[anatomy, img_id] for [anatomy, img_id] in self.id_antomy_ls if img_id in self.id2image_path]    # WARNING 这样filter会导致
+        
+        self.anatomy_idx_range = {}
+        for idx, (anatomy, _) in enumerate(self.id_antomy_ls):
+            if anatomy not in self.anatomy_idx_range:
+                self.anatomy_idx_range[anatomy] = [idx, idx]
+            else:
+                self.anatomy_idx_range[anatomy][1] = idx
+        # Optionally, convert indices to tuples for clarity and print the ranges
+        self.anatomy_idx_range = {anat: [r[0], r[1]] for anat, r in self.anatomy_idx_range.items()}
+        
+        if int(os.environ.get("RANK", 0)) == 0:
+            print(f'** DATA ** Load {len(self.id2image_path)} images.')
+            print(f'** DATA ** Load {len(self.anatomy2id_ls)} anatomy.')
+
+    def prepare_image_paths(self, jsonl_file):
+        id2image_path = {}
+        
+        with open(jsonl_file, 'r') as f:
+            data = f.readlines()
+        data = [json.loads(l) for l in data]
+        for d in data:
+            data_id = d['name']   # valid_692_a_1.nii.gz
+            id2image_path[data_id] = d['img_path']
+
+        return id2image_path
+    
+    def prepare_anatomy_data(self, csv_file_dir, max_samples):
+        anatomy2id_ls = {} # 'lung': ['valid_692_a_1.nii.gz', ...]
+        
+        for csv_file in sorted(os.listdir(csv_file_dir)):   
+            anatomy_name = csv_file.replace('.csv', '') # lung.csv -> lung
+            
+            # DEBUG
+            if anatomy_name not in self.anatomy_filter:
+                continue
+            # DEBUG
+             
+            anatomy2id_ls[anatomy_name] = []
+            df = pd.read_csv(os.path.join(csv_file_dir, csv_file))
+            for index, row in df.head(max_samples).iterrows():
+                anatomy2id_ls[anatomy_name].append(row['File Path'])
+
+        return anatomy2id_ls
+    
+    def prepare_similarity_table(self, npy_file_dir, max_samples):
+        anatomy2simi_tab = {} # 'lung': a tensor with shape NxN
+        
+        for npy_file in sorted(os.listdir(npy_file_dir)):   
+            anatomy_name = npy_file.replace('.npy', '') # lung.npy -> lung
+            
+            # DEBUG
+            if anatomy_name not in self.anatomy_filter:
+                continue
+            # DEBUG
+            
+            simi_tab = np.load(os.path.join(npy_file_dir, npy_file))[:max_samples, :max_samples]
+            simi_tab = np.clip(simi_tab, 0, 1)
+            simi_tab = np.round(simi_tab * 100).astype(np.uint8)  # 将 0~1 的浮点数转为 0~100 的整数 (保留两位小数)
+            for i in range(simi_tab.shape[0]):
+                for j in range(simi_tab.shape[0]):
+                    simi_tab[i, j] = max(simi_tab[i, j], simi_tab[j, i])
+            anatomy2simi_tab[anatomy_name] = simi_tab
+
+        return anatomy2simi_tab
+    
+    def check_integrity(self, anatomy2id_ls, anatomy2simi_tab):
+        id_keys = set(anatomy2id_ls.keys())
+        simi_keys = set(anatomy2simi_tab.keys())
+        if id_keys != simi_keys:
+            missing_in_simi = id_keys - simi_keys
+            missing_in_id = simi_keys - id_keys
+            raise ValueError(f"Inconsistent keys detected. Missing in similarity table: {missing_in_simi}. Missing in anatomy data: {missing_in_id}.")
+        
+        for anatomy in id_keys:
+            if not (len(anatomy2id_ls[anatomy]) == anatomy2simi_tab[anatomy].shape[0] and len(anatomy2id_ls[anatomy]) == anatomy2simi_tab[anatomy].shape[1]):
+                raise ValueError(f"Length of anatomy data {len(anatomy2id_ls[anatomy])} and similarity table {anatomy2simi_tab[anatomy].shape} do not match for {anatomy}.")
+            
+            # Here we check if all samples are readable and valid (at least has one positive candidate excluding itself)
+            filtered_sample_id_ls = []
+            excluded_index_ls = []
+            no_valid_candidate = 0
+            tmp_simi_tab = anatomy2simi_tab[anatomy]
+            
+            for i, sample_id in enumerate(anatomy2id_ls[anatomy]):
+                if np.any(np.delete(tmp_simi_tab[i], i) > 90):
+                    # image_file_name = sample_id.replace('.nii.gz', '.npz')
+                    if sample_id in self.id2image_path and os.path.exists(self.id2image_path[sample_id]):
+                        filtered_sample_id_ls.append(sample_id)
+                        continue
+                    else:
+                        excluded_index_ls.append(i)
+                else:
+                    no_valid_candidate += 1
+                    excluded_index_ls.append(i)
+            anatomy2id_ls[anatomy] = filtered_sample_id_ls
+            
+            anatomy2simi_tab[anatomy] = np.delete(anatomy2simi_tab[anatomy], excluded_index_ls, axis=0)
+            anatomy2simi_tab[anatomy] = np.delete(anatomy2simi_tab[anatomy], excluded_index_ls, axis=1)
+            
+            print(f'For {anatomy}, {no_valid_candidate} samples has no valid candidate, removed')
+            
+            assert len(anatomy2id_ls[anatomy]) == anatomy2simi_tab[anatomy].shape[0] and len(anatomy2id_ls[anatomy]) == anatomy2simi_tab[anatomy].shape[1]
+            
+    def get_anatomy_split(self):
+        anatomy_split = list(self.anatomy_idx_range.values())
+        anatomy_name = list(self.anatomy_idx_range.keys())
+        return anatomy_split, anatomy_name
+
+    def __len__(self):
+        # return len(self.anatomy_ls)   # NOTE 不按照anatomy为单位采
+        return len(self.id_antomy_ls)
+    
+    def get_anatomy2id_ls(self):
+        return self.anatomy2id_ls
+    
+    def get_similarity_table(self, anatomy):
+        return self.anatomy2simi_tab[anatomy]
+    
+    def nii_img_to_tensor(self, path):
+
+        nii_img = safe_load_image(sample["image"])
+        img_data = nii_img.get_fdata()
+        img_data = np.flip(img_data, axis=0)
+        img_data = np.flip(img_data, axis=1)
+
+        # WARNING Respacing
+        img_data = img_data.transpose(2, 0, 1)
+        img_data = np.copy(img_data)
+        tensor = torch.tensor(img_data)
+        tensor = tensor.unsqueeze(0).unsqueeze(0)
+        
+        target_x_spacing = 0.75
+        target_y_spacing = 0.75
+        target_z_spacing = 1.5
+        current = (3, 1, 1)   # this is all set to 3 1 1
+        target = (target_z_spacing, target_x_spacing, target_y_spacing)
+        
+        img_data = resize_array(tensor, current, target)
+        img_data = img_data[0][0]
+        img_data= np.transpose(img_data, (1, 2, 0))
+
+        # WARNING Normalization 2
+        hu_min, hu_max = -1000, 1000
+        img_data = np.clip(img_data, hu_min, hu_max)
+        img_data = (((img_data ) / 1000)).astype(np.float32)
+
+        tensor = torch.tensor(img_data)
+        
+        # WARNING Padding or Crop
+        
+        # Get the dimensions of the input tensor
+        target_shape = (480,480,240)    # h w d
+        
+        # Extract dimensions
+        h, w, d = tensor.shape
+        
+        # Calculate cropping/padding values for height, width, and depth
+        dh, dw, dd = target_shape
+        h_start = max((h - dh) // 2, 0)
+        h_end = min(h_start + dh, h)
+        w_start = max((w - dw) // 2, 0)
+        w_end = min(w_start + dw, w)
+        d_start = max((d - dd) // 2, 0)
+        d_end = min(d_start + dd, d)
+
+        # Crop or pad the tensor
+        tensor = tensor[h_start:h_end, w_start:w_end, d_start:d_end]
+
+        pad_h_before = (dh - tensor.size(0)) // 2
+        pad_h_after = dh - tensor.size(0) - pad_h_before
+
+        pad_w_before = (dw - tensor.size(1)) // 2
+        pad_w_after = dw - tensor.size(1) - pad_w_before
+
+        pad_d_before = (dd - tensor.size(2)) // 2
+        pad_d_after = dd - tensor.size(2) - pad_d_before
+
+        tensor = torch.nn.functional.pad(tensor, (pad_d_before, pad_d_after, pad_w_before, pad_w_after, pad_h_before, pad_h_after), value=-1)
+
+        tensor = tensor.permute(2, 0, 1)    # d h w
+
+        tensor = tensor.unsqueeze(0)    # 1 d h w
+
+        return tensor
+    
+    def load_2d_image_to_tensor(self, image_path,resize=True):
+            """
+            Load a 2D grayscale image and convert it to a tensor with dimensions [1, 1, 480, 480].
+            
+            Args:
+                image_path (str): Path to the image file.
+                
+            Returns:
+                tensor (torch.Tensor): Tensor with shape [1, 1, 480, 480].
+            """
+            # Load the image as grayscale
+            image = Image.open(image_path).convert('L')
+            
+            # Resize to 480x480 if needed
+            if resize:
+                if image.size != (480, 480):
+                    image = image.resize((480, 480), Image.BILINEAR)
+            else:
+                if image.size != (480, 480):
+                    # 计算可以开始裁剪的最大左上角坐标
+                    width, height = image.size
+                    crop_size = 480
+                    if width < crop_size or height < crop_size:
+                        raise ValueError(f"Image size {image.size} is smaller than the crop size {crop_size}.")
+                    max_left = width - crop_size
+                    max_top = height - crop_size
+                    
+                    # 随机选择裁剪的左上角坐标
+                    left = random.randint(0, max(0, max_left))
+                    top = random.randint(0, max(0, max_top))
+                    
+                    # 裁剪图像
+                    image = image.crop((left, top, left + crop_size, top + crop_size))
+                
+
+            # Convert to numpy array and normalize to [0, 1]
+            img_array = np.array(image, dtype=np.float32) / 255.0
+            
+            # Convert to tensor
+            tensor = torch.from_numpy(img_array)
+            
+            # Add batch and channel dimensions
+            tensor = tensor.unsqueeze(0).unsqueeze(0)  # [1, 1, 480, 480]
+            
+            return tensor
+
+    def respace_crop_pad(self, img_tensor):
+        # assume we have loaded all the images
+        
+        target_x_spacing = 0.75
+        target_y_spacing = 0.75
+        target_z_spacing = 1.5
+        current = (3, 1, 1)   # this is all set to 3 1 1
+        target = (target_z_spacing, target_x_spacing, target_y_spacing)
+        
+        img_data = resize_array(img_tensor, current, target)
+        img_data = img_data[0][0]
+        img_data= np.transpose(img_data, (1, 2, 0))
+
+        # WARNING Normalization 2
+        hu_min, hu_max = -1000, 1000
+        img_data = np.clip(img_data, hu_min, hu_max)
+        img_data = (((img_data ) / 1000)).astype(np.float32)
+
+        tensor = torch.tensor(img_data)
+        
+        # WARNING Padding or Crop
+        
+        # Get the dimensions of the input tensor
+        target_shape = (480,480,240)    # h w d
+        
+        # Extract dimensions
+        h, w, d = tensor.shape
+        
+        # Calculate cropping/padding values for height, width, and depth
+        dh, dw, dd = target_shape
+        h_start = max((h - dh) // 2, 0)
+        h_end = min(h_start + dh, h)
+        w_start = max((w - dw) // 2, 0)
+        w_end = min(w_start + dw, w)
+        d_start = max((d - dd) // 2, 0)
+        d_end = min(d_start + dd, d)
+
+        # Crop or pad the tensor
+        tensor = tensor[h_start:h_end, w_start:w_end, d_start:d_end]
+
+        pad_h_before = (dh - tensor.size(0)) // 2
+        pad_h_after = dh - tensor.size(0) - pad_h_before
+
+        pad_w_before = (dw - tensor.size(1)) // 2
+        pad_w_after = dw - tensor.size(1) - pad_w_before
+
+        pad_d_before = (dd - tensor.size(2)) // 2
+        pad_d_after = dd - tensor.size(2) - pad_d_before
+
+        tensor = torch.nn.functional.pad(tensor, (pad_d_before, pad_d_after, pad_w_before, pad_w_after, pad_h_before, pad_h_after), value=-1)
+
+        tensor = tensor.permute(2, 0, 1)    # d h w
+
+        tensor = tensor.unsqueeze(0)    # 1 d h w
+
+        return tensor
+    
+    def load_image(self, path):
+        nii_img = safe_load_image(sample["image"])
+        img_data = nii_img.get_fdata()
+        img_data = np.flip(img_data, axis=0)
+        img_data = np.flip(img_data, axis=1)
+
+        # WARNING Respacing
+        img_data = img_data.transpose(2, 0, 1)
+        img_data = np.copy(img_data)
+        tensor = torch.tensor(img_data)
+        tensor = tensor.unsqueeze(0).unsqueeze(0)
+        
+        return tensor
+
+    def MLS_nii_img_to_tensor(self, path):
+
+        nii_img = safe_load_image(sample["image"])
+        img_data = nii_img.get_fdata()
+        img_data = np.flip(img_data, axis=0)
+        img_data = np.flip(img_data, axis=1)
+
+        # WARNING Respacing
+        img_data = img_data.transpose(2, 0, 1)
+        img_data = np.copy(img_data)
+        tensor = torch.tensor(img_data)
+        tensor = tensor.unsqueeze(0).unsqueeze(0)
+        
+        target_x_spacing = 0.75
+        target_y_spacing = 0.75
+        target_z_spacing = 1.5
+        current = (3, 1, 1)   # this is all set to 3 1 1
+        target = (target_z_spacing, target_x_spacing, target_y_spacing)
+        
+        img_data = resize_array(tensor, current, target)
+        img_data = img_data[0][0]
+        img_data= np.transpose(img_data, (1, 2, 0))
+
+        # WARNING Normalization 2
+        hu_min, hu_max = -1000, 1000
+        img_data = np.clip(img_data, hu_min, hu_max)
+        img_data = (((img_data ) / 1000)).astype(np.float32)
+
+        tensor = torch.tensor(img_data)
+        
+        # WARNING Padding or Crop
+        
+        # Get the dimensions of the input tensor
+        target_shape = (480,480,240)    # h w d
+        
+        # Extract dimensions
+        h, w, d = tensor.shape
+        
+        # Calculate cropping/padding values for height, width, and depth
+        dh, dw, dd = target_shape
+        h_start = max((h - dh) // 2, 0)
+        h_end = min(h_start + dh, h)
+        w_start = max((w - dw) // 2, 0)
+        w_end = min(w_start + dw, w)
+        d_start = max((d - dd) // 2, 0)
+        d_end = min(d_start + dd, d)
+
+        # Crop or pad the tensor
+        tensor = tensor[h_start:h_end, w_start:w_end, d_start:d_end]
+
+        pad_h_before = (dh - tensor.size(0)) // 2
+        pad_h_after = dh - tensor.size(0) - pad_h_before
+
+        pad_w_before = (dw - tensor.size(1)) // 2
+        pad_w_after = dw - tensor.size(1) - pad_w_before
+
+        pad_d_before = (dd - tensor.size(2)) // 2
+        pad_d_after = dd - tensor.size(2) - pad_d_before
+
+        tensor = torch.nn.functional.pad(tensor, (pad_d_before, pad_d_after, pad_w_before, pad_w_after, pad_h_before, pad_h_after), value=-1)
+
+        tensor = tensor.permute(2, 0, 1)    # d h w
+
+        tensor = tensor.unsqueeze(0)    # 1 d h w
+
+        return tensor
+
+
+    def __getitem__(self, index):
+        """
+        Returns:
+            video_tensor (tensor): 1, D, H, W
+            anatomy (str): anatomy name
+            sampled_id (str): image id
+        """
+        anatomy, sampled_id = self.id_antomy_ls[index]
+        
+        # NOTE: All processed in advance
+        if  '3D' in self.modality:
+            image_file_name = sampled_id.replace('.nii.gz', '.npz')
+            if sampled_id in self.id2image_path and os.path.exists(self.id2image_path[sampled_id]):
+                video_tensor = self.MLS_nii_img_to_tensor(self.id2image_path[sampled_id])
+            else:
+                raise ValueError(f"Image file {image_file_name} not found.")
+        elif '2D' in self.modality:
+            if sampled_id in self.id2image_path and os.path.exists(self.id2image_path[sampled_id]):
+                video_tensor = self.load_2d_image_to_tensor(self.id2image_path[sampled_id])
+            else:
+                raise ValueError(f"Image file {image_file_name} not found.")
+        # return torch.tensor(video_tensor), anatomy, sampled_id
+        if self.modal_embedding:
+            return video_tensor.clone().detach(), anatomy, sampled_id, modality_dict[self.modality]
+        else:
+            return video_tensor.clone().detach(), anatomy, sampled_id
+        
+if __name__ == '__main__':
+    from torch.utils.data import DataLoader
+    from accelerate import Accelerator
+    from accelerate import DistributedDataParallelKwargs
+    from accelerate.utils import InitProcessGroupKwargs
+    
+    # ddp_kwargs = DistributedDataParallelKwargs(find_unused_parameters=True)
+    # kwargs = InitProcessGroupKwargs()
+    # accelerator = Accelerator(kwargs_handlers=[ddp_kwargs, kwargs])
+    
+    # data_valid_jsonl = '/mnt/petrelfs/zhangtengfei/RadIR/dataset/CXR/MIMIC_CXR/test.jsonl'
+    # data_valid_csv_dir = '/mnt/petrelfs/zhangtengfei/RadIR/dataset/CXR/MIMIC_CXR/test_entity'
+    # data_valid_npy_dir = '/mnt/petrelfs/zhangtengfei/RadIR/dataset/CXR/MIMIC_CXR/test_ratescore'
+    
+    # valid_ds = Conditional_CTReportDataset_Eval(jsonl_file=data_valid_jsonl, csv_file_dir=data_valid_csv_dir, npy_file_dir=data_valid_npy_dir, anatomy_filter=["left hemidiaphragm","airspace"], max_samples=10000,modality='2D')
+    
+    # rank = torch.distributed.get_rank()  # 获取当前进程的 rank
+    # world_size = torch.distributed.get_world_size()  # 获取总进程数
+    
+    # subset, anatomy_in_this_partition = custom_data_split(valid_ds, rank, world_size)
+
+    # valid_dl = DataLoader(
+    #     subset,
+    #     num_workers=8,
+    #     batch_size=16,
+    #     shuffle = False,
+    # )
+    
+    # valid_dl = iter(valid_dl)
+    
+    # statistic = {}
+    # for video, anatomy_ls, sampled_id_ls in iter(valid_dl):
+    #     for anatomy, sample_id in zip(anatomy_ls, sampled_id_ls):
+    #         if anatomy not in statistic:
+    #             statistic[anatomy] = set()
+    #         statistic[anatomy].add(sample_id)
+            
+    # print(f'RANK {rank}:', statistic.keys(), anatomy_in_this_partition)
+    # for k, v in statistic.items():
+    #     print(f'RANK {rank}: {k} - {len(v)}')
+    
+    ddp_kwargs = DistributedDataParallelKwargs(find_unused_parameters=True)
+    kwargs = InitProcessGroupKwargs()
+    accelerator = Accelerator(kwargs_handlers=[ddp_kwargs, kwargs])
+    
+    data_valid_jsonl = '/mnt/data-3/users/zhangzeyu/HMClip/RadIR/dataset/CTRATE-IR/test.jsonl'
+    data_valid_csv_dir = '/mnt/data-3/users/zhangzeyu/HMClip/RadIR/dataset/CTRATE-IR/anatomy/val_entity'
+    data_valid_npy_dir = '/mnt/data-3/users/zhangzeyu/HMClip/RadIR/dataset/CTRATE-IR/anatomy/val_ratescore_compressed'
+
+
+    valid_ds = Conditional_CTReportDataset_Eval(jsonl_file=data_valid_jsonl, csv_file_dir=data_valid_csv_dir, npy_file_dir=data_valid_npy_dir, anatomy_filter=["right rib","thyroid gland"], max_samples=10000,modality='3D')
+    
+    rank = torch.distributed.get_rank()  # 获取当前进程的 rank
+    world_size = torch.distributed.get_world_size()  # 获取总进程数
+    
+    subset, anatomy_in_this_partition = custom_data_split(valid_ds, rank, world_size)
+
+    valid_dl = DataLoader(
+        subset,
+        num_workers=8,
+        batch_size=16,
+        shuffle = False,
+    )
+    
+    valid_dl = iter(valid_dl)
+    
+    statistic = {}
+    for video, anatomy_ls, sampled_id_ls in iter(valid_dl):
+        for anatomy, sample_id in zip(anatomy_ls, sampled_id_ls):
+            if anatomy not in statistic:
+                statistic[anatomy] = set()
+            statistic[anatomy].add(sample_id)
+            
+    print(f'RANK {rank}:', statistic.keys(), anatomy_in_this_partition)
+    for k, v in statistic.items():
+        print(f'RANK {rank}: {k} - {len(v)}')
+
+def safe_load_image(path):
+    """
+    Robust image loader that supports both .nii.gz and .npy files.
+    If .nii.gz is missing, automatically falls back to _compressed.npy.
+    """
+    try:
+        if os.path.exists(path):
+            if path.endswith('.nii.gz'):
+                data = nib.load(path).get_fdata()
+                return data
+            elif path.endswith('.npy'):
+                data = np.load(path)
+                return data
+
+        # fallback to compressed npy
+        npy_fallback = path.replace('.nii.gz', '_compressed.npy')
+        if os.path.exists(npy_fallback):
+            data = np.load(npy_fallback)
+            print(f"[Info] Fallback to compressed file: {npy_fallback}")
+            return data
+
+        # all failed
+        raise FileNotFoundError(f"{path} and fallback {npy_fallback} not found")
+
+    except Exception as e:
+        print(f"[Warning] Error loading {path}: {e}")
+        return None

scripts/conditional_dataset_train.py

@@ -0,0 +1,825 @@
+import os
+import glob
+import json
+import torch
+import pandas as pd
+import numpy as np
+from PIL import Image
+from torch.utils.data import Dataset
+import torchvision.transforms as transforms
+from functools import partial
+import torch.nn.functional as F
+import nibabel as nib
+from tqdm import tqdm
+import random
+import math
+from monai.transforms import Compose, RandRotate, RandFlip, RandZoom, RandAffine, RandGaussianNoise, RandScaleIntensity, RandAdjustContrast,RandCoarseDropout
+from monai.data.meta_tensor import MetaTensor
+import nibabel as nib
+
+
+def safe_load_image(path):
+    """
+    Robust image loader that supports both .nii.gz and .npy files.
+    If .nii.gz is missing, automatically falls back to _compressed.npy.
+    """
+    try:
+        if os.path.exists(path):
+            if path.endswith('.nii.gz'):
+                data = nib.load(path).get_fdata()
+                return data
+            elif path.endswith('.npy'):
+                data = np.load(path)
+                return data
+
+        # fallback to compressed npy
+        npy_fallback = path.replace('.nii.gz', '_compressed.npy')
+        if os.path.exists(npy_fallback):
+            data = np.load(npy_fallback)
+            print(f"[Info] Fallback to compressed file: {npy_fallback}")
+            return data
+
+        # all failed
+        raise FileNotFoundError(f"{path} and fallback {npy_fallback} not found")
+
+    except Exception as e:
+        print(f"[Warning] Error loading {path}: {e}")
+        return None
+
+
+def resize_array_to_tensor(array, current_spacing, target_spacing):
+    """
+    Resize the array to match the target spacing.
+
+    Args:
+    array (torch.Tensor): Input array to be resized.
+    current_spacing (tuple): Current voxel spacing (z_spacing, xy_spacing, xy_spacing).
+    target_spacing (tuple): Target voxel spacing (target_z_spacing, target_x_spacing, target_y_spacing).
+
+    Returns:
+    np.ndarray: Resized array.
+    """
+    # Calculate new dimensions
+    original_shape = array.shape[2:]
+    scaling_factors = [
+        current_spacing[i] / target_spacing[i] for i in range(len(original_shape))
+    ]
+    new_shape = [
+        int(original_shape[i] * scaling_factors[i]) for i in range(len(original_shape))
+    ]
+    # Resize the array
+    resized_array = F.interpolate(array, size=new_shape, mode='trilinear', align_corners=False)
+    return resized_array
+
+def resize_array(array, current_spacing, target_spacing):
+    """
+    Resize the array to match the target spacing.
+
+    Args:
+    array (torch.Tensor): Input array to be resized.
+    current_spacing (tuple): Current voxel spacing (z_spacing, xy_spacing, xy_spacing).
+    target_spacing (tuple): Target voxel spacing (target_z_spacing, target_x_spacing, target_y_spacing).
+
+    Returns:
+    np.ndarray: Resized array.
+    """
+    # Calculate new dimensions
+    original_shape = array.shape[2:]
+    scaling_factors = [
+        current_spacing[i] / target_spacing[i] for i in range(len(original_shape))
+    ]
+    new_shape = [
+        int(original_shape[i] * scaling_factors[i]) for i in range(len(original_shape))
+    ]
+    # Resize the array
+    resized_array = F.interpolate(array, size=new_shape, mode='trilinear', align_corners=False).cpu().numpy()
+    return resized_array
+
+modality_dict = {'3D':0,'2D':1}
+class Conditional_ImageReportDataset_Train(Dataset):
+    def __init__(self, local_batch_size, jsonl_file, csv_file_dir, npy_file_dir, anatomy_filter, positive_threshold, negative_threshold, max_samples=30000,modality='CT', need_aug=True,modal_embedding=False,is_train=True):
+        self.anatomy_filter = anatomy_filter
+        self.is_train = is_train
+        self.modal_embedding = modal_embedding
+        self.csv_file_dir = csv_file_dir
+        self.npy_file_dir = npy_file_dir
+        self.max_samples = max_samples
+        self.modality = modality
+        self.local_batch_size = local_batch_size    # for each anatomy, how many samples to be selected
+        
+        if positive_threshold == 0 and negative_threshold != 1: # >0 & <0.4
+            positive_threshold = negative_threshold
+        if positive_threshold != 1 and negative_threshold == 1: # >0.7 & <1
+            negative_threshold = positive_threshold
+            
+        self.positive_threshold = positive_threshold
+        self.negative_threshold = negative_threshold
+        
+        self.id2image_path = self.prepare_image_paths(jsonl_file)
+        
+        self.prepare_anatomy_data()
+        
+        self.anatomy_ls = list(self.anatomy2id_ls.keys())
+        self.anatomy_weight = [round(math.log(len(self.anatomy2id_ls[anatomy]))) for anatomy in self.anatomy_ls]   # weight 与 sqrt(size) 相关
+        self.anatomy_weight = [w/sum(self.anatomy_weight) for w in self.anatomy_weight]  # NOTE: Balancing between anatomys
+        # self.anatomy_weight = [1/len(self.anatomy_ls) for anatomy in self.anatomy_ls]
+        
+        if need_aug:
+            if '3D' in self.modality:
+                self.augmentator = Compose([
+                        RandRotate(range_x=15, range_y=15, range_z=15, prob=0.3, keep_size=True),  # 随机旋转
+                        RandFlip(prob=0.5, spatial_axis=0),  # 随机翻转（沿深度轴）
+                        RandZoom(min_zoom=0.85, max_zoom=1.15, prob=0.3, keep_size=True),  # 随机缩放
+                        RandAffine(prob=0.3, translate_range=(24, 48, 48)),  # 随机平移
+                        RandGaussianNoise(prob=0.3, mean=0.0, std=0.05),  # 添加高斯噪声
+                        RandScaleIntensity(factors=[-0.1, 0.1], prob=0.3),
+                        RandAdjustContrast(prob=0.3, gamma=(0.9, 1.1)),  # 增加对比度调整
+                        RandCoarseDropout(prob=0.2, holes=8, spatial_size=3)  # 模拟局部缺失
+                ])
+            elif '2D' in self.modality:
+                self.augmentator = Compose([
+                    RandRotate(range_x=15, range_y=15, prob=0.5, keep_size=True),  # 随机旋转（仅x和y轴）
+                    RandFlip(prob=0.5, spatial_axis=(0,1)),  # 随机翻转（仅水平和垂直轴）
+                    RandZoom(min_zoom=0.8, max_zoom=1.1, prob=0.5, keep_size=True),  # 随机缩放
+                    RandAffine(prob=0.5, translate_range=(32, 32)),  # 随机平移（仅2D平移）
+                    RandGaussianNoise(prob=0.5, mean=0.0, std=0.1),  # 添加高斯噪声
+                    RandScaleIntensity(factors=[-0.15, 0.15], prob=0.5),  # 随机调整像素强度
+                    RandAdjustContrast(prob=0.4, gamma=(0.85, 1.15))  # 增加对比度变化
+                    ])
+        else:
+            self.augmentator = None
+        
+        if int(os.environ.get("RANK", 0)) == 0:
+            print(f'** DATA ** Load {len(self.id2image_path)} images.')
+            print(f'** DATA ** Load {len(self.anatomy2id_ls)} anatomy.')
+
+    def prepare_image_paths(self, jsonl_file):
+        id2image_path = {}
+        
+        with open(jsonl_file, 'r') as f:
+            data = f.readlines()
+        data = [json.loads(l) for l in data]
+        for d in data:
+            # 这里有问题的
+            data_id = d['name']   # valid_692_a_1.nii.gz
+            id2image_path[data_id] = d['img_path']
+
+        return id2image_path
+    
+    def prepare_anatomy_data(self):
+        self.anatomy2id_ls = {} # 'lung': ['valid_692_a_1.nii.gz', ...]
+        self.anatomy2simi_tab = {} # 'lung': a tensor with shape NxN
+        
+        for csv_file in os.listdir(self.csv_file_dir):   
+            anatomy_name = csv_file.replace('.csv', '') # lung.csv -> lung
+            
+            # DEBUG
+            if anatomy_name not in self.anatomy_filter:
+                continue
+            # DEBUG
+            
+            self.anatomy2id_ls[anatomy_name] = []
+            df = pd.read_csv(os.path.join(self.csv_file_dir, csv_file))
+            print(f"{csv_file} has {len(df)} rows.")
+            if len(df) > self.max_samples:
+                # 随机抽取max_samples个样本，并按照原csv顺序排序，以保持与simi_tab索引对应
+                sampled_df = df.sample(n=self.max_samples, random_state=42).sort_index()
+                sampled_indices = sampled_df.index.tolist() # 例如抽取[0, 3, 5]行
+            else:
+                sampled_df = df
+                sampled_indices = sampled_df.index.tolist()
+            for _, row in sampled_df.iterrows():
+                self.anatomy2id_ls[anatomy_name].append(row['File Path'])
+                
+            npy_file = csv_file.replace('.csv', '.npy')
+            full_simi_tab = np.load(os.path.join(self.npy_file_dir, npy_file))
+            # 截取原csv对应的大小
+            full_simi_tab = full_simi_tab[:len(df), :len(df)]
+            full_simi_tab = np.clip(full_simi_tab, 0, 1)
+            # 根据抽样的索引构造子矩阵，确保与csv中样本一一对应
+            sampled_simi_tab = full_simi_tab[np.ix_(sampled_indices, sampled_indices)] # 例如只保留[0, 3, 5]行/列
+            simi_tab_scaled = np.round(sampled_simi_tab * 100).astype(np.uint8)
+            self.anatomy2simi_tab[anatomy_name] = simi_tab_scaled
+            
+        self.check_integrity()
+
+    def check_integrity(self):
+        id_keys = set(self.anatomy2id_ls.keys())
+        simi_keys = set(self.anatomy2simi_tab.keys())
+        if id_keys != simi_keys:
+            missing_in_simi = id_keys - simi_keys
+            missing_in_id = simi_keys - id_keys
+            raise ValueError(f"Inconsistent keys detected. Missing in similarity table: {missing_in_simi}. Missing in anatomy data: {missing_in_id}.")
+        # 修改
+        if self.modality == '3D':
+            for anatomy in id_keys:
+                if not (len(self.anatomy2id_ls[anatomy]) == self.anatomy2simi_tab[anatomy].shape[0] and len(self.anatomy2id_ls[anatomy]) == self.anatomy2simi_tab[anatomy].shape[1]):
+                    raise ValueError(f"Length of anatomy data {len(self.anatomy2id_ls[anatomy])} and similarity table {self.anatomy2simi_tab[anatomy].shape} do not match for {anatomy}.")
+                
+                # Here we check if all samples are readable
+                filtered_sample_id_ls = []
+                valid_indices = []
+                for i, sample_id in enumerate(self.anatomy2id_ls[anatomy]):
+                    # image_file_name = sample_id.replace('.nii.gz', '.npz')
+                    if sample_id in self.id2image_path and os.path.exists(self.id2image_path[sample_id]):
+                        filtered_sample_id_ls.append(sample_id)
+                        valid_indices.append(i)
+                        continue
+                    else:
+                        print(f'{sample_id} is missing, removed')
+                        # # Remove the invalid sample from self.anatomy2id_ls and update simi_tab accordingly
+                        # self.anatomy2simi_tab[anatomy] = np.delete(self.anatomy2simi_tab[anatomy], i, axis=0)
+                        # self.anatomy2simi_tab[anatomy] = np.delete(self.anatomy2simi_tab[anatomy], i, axis=1)
+                self.anatomy2simi_tab[anatomy] = self.anatomy2simi_tab[anatomy][np.ix_(valid_indices, valid_indices)]  # 保留有效索引对应的子矩阵
+                self.anatomy2id_ls[anatomy] = filtered_sample_id_ls
+                
+                assert len(self.anatomy2id_ls[anatomy]) == self.anatomy2simi_tab[anatomy].shape[0] and len(self.anatomy2id_ls[anatomy]) == self.anatomy2simi_tab[anatomy].shape[1]
+        elif self.modality == '2D':
+            for anatomy in id_keys:
+                if not (len(self.anatomy2id_ls[anatomy]) == self.anatomy2simi_tab[anatomy].shape[0] and len(self.anatomy2id_ls[anatomy]) == self.anatomy2simi_tab[anatomy].shape[1]):
+                    raise ValueError(f"Length of anatomy data {len(self.anatomy2id_ls[anatomy])} and similarity table {self.anatomy2simi_tab[anatomy].shape} do not match for {anatomy}.")
+                
+                # Here we check if all samples are readable
+                filtered_sample_id_ls = []
+                valid_indices = []
+                for i, sample_id in enumerate(self.anatomy2id_ls[anatomy]):
+                    if sample_id in self.id2image_path and os.path.exists(self.id2image_path[sample_id]):
+                        filtered_sample_id_ls.append(sample_id)
+                        valid_indices.append(i)
+                        continue
+                    else:
+                        print(f'{sample_id} is missing, removed')
+                        # # Remove the invalid sample from self.anatomy2id_ls and update simi_tab accordingly
+                        # self.anatomy2simi_tab[anatomy] = np.delete(self.anatomy2simi_tab[anatomy], i, axis=0)
+                        # self.anatomy2simi_tab[anatomy] = np.delete(self.anatomy2simi_tab[anatomy], i, axis=1)
+                self.anatomy2simi_tab[anatomy] = self.anatomy2simi_tab[anatomy][np.ix_(valid_indices, valid_indices)]  # 保留有效索引对应的子矩阵
+                self.anatomy2id_ls[anatomy] = filtered_sample_id_ls
+                
+                assert len(self.anatomy2id_ls[anatomy]) == self.anatomy2simi_tab[anatomy].shape[0] and len(self.anatomy2id_ls[anatomy]) == self.anatomy2simi_tab[anatomy].shape[1]
+        else:
+            raise   ValueError(f"Unsupported modality: {self.modality}. Supported modalities are '3D' and '2D'.")
+    def __len__(self):
+        return 100000
+    
+    def MLS_nii_img_to_tensor(self, path):
+
+        nii_img = safe_load_image(sample["image"])
+        img_data = nii_img.get_fdata()
+        img_data = np.flip(img_data, axis=0)
+        img_data = np.flip(img_data, axis=1)
+
+        # WARNING Respacing
+        img_data = img_data.transpose(2, 0, 1)
+        img_data = np.copy(img_data)
+        tensor = torch.tensor(img_data)
+        tensor = tensor.unsqueeze(0).unsqueeze(0)
+        
+        target_x_spacing = 0.75
+        target_y_spacing = 0.75
+        target_z_spacing = 1.5
+        current = (3, 1, 1)   # this is all set to 3 1 1
+        target = (target_z_spacing, target_x_spacing, target_y_spacing)
+        
+        img_data = resize_array(tensor, current, target)
+        img_data = img_data[0][0]
+        img_data= np.transpose(img_data, (1, 2, 0))
+
+        # WARNING Normalization 2
+        hu_min, hu_max = -1000, 1000
+        img_data = np.clip(img_data, hu_min, hu_max)
+        img_data = (((img_data ) / 1000)).astype(np.float32)
+
+        tensor = torch.tensor(img_data)
+        
+        # WARNING Padding or Crop
+        
+        # Get the dimensions of the input tensor
+        target_shape = (480,480,240)    # h w d
+        
+        # Extract dimensions
+        h, w, d = tensor.shape
+        
+        # Calculate cropping/padding values for height, width, and depth
+        dh, dw, dd = target_shape
+        h_start = max((h - dh) // 2, 0)
+        h_end = min(h_start + dh, h)
+        w_start = max((w - dw) // 2, 0)
+        w_end = min(w_start + dw, w)
+        d_start = max((d - dd) // 2, 0)
+        d_end = min(d_start + dd, d)
+
+        # Crop or pad the tensor
+        tensor = tensor[h_start:h_end, w_start:w_end, d_start:d_end]
+
+        pad_h_before = (dh - tensor.size(0)) // 2
+        pad_h_after = dh - tensor.size(0) - pad_h_before
+
+        pad_w_before = (dw - tensor.size(1)) // 2
+        pad_w_after = dw - tensor.size(1) - pad_w_before
+
+        pad_d_before = (dd - tensor.size(2)) // 2
+        pad_d_after = dd - tensor.size(2) - pad_d_before
+
+        tensor = torch.nn.functional.pad(tensor, (pad_d_before, pad_d_after, pad_w_before, pad_w_after, pad_h_before, pad_h_after), value=-1)
+
+        tensor = tensor.permute(2, 0, 1)    # d h w
+
+        tensor = tensor.unsqueeze(0)    # 1 d h w
+
+        return tensor
+
+    def nii_img_to_tensor(self, path):
+
+        nii_img = safe_load_image(sample["image"])
+        img_data = nii_img.get_fdata()
+        img_data = np.flip(img_data, axis=0)
+        img_data = np.flip(img_data, axis=1)
+
+        # WARNING Respacing
+        img_data = img_data.transpose(2, 0, 1)
+        img_data = np.copy(img_data)
+        tensor = torch.tensor(img_data)
+        tensor = tensor.unsqueeze(0).unsqueeze(0)
+        
+        target_x_spacing = 0.75
+        target_y_spacing = 0.75
+        target_z_spacing = 1.5
+        current = (3, 1, 1)   # this is all set to 3 1 1
+        target = (target_z_spacing, target_x_spacing, target_y_spacing)
+        
+        img_data = resize_array(tensor, current, target)
+        img_data = img_data[0][0]
+        img_data= np.transpose(img_data, (1, 2, 0))
+        
+        # WARNING Normalization 2
+        hu_min, hu_max = -1000, 1000
+        img_data = np.clip(img_data, hu_min, hu_max)
+        img_data = (((img_data ) / 1000)).astype(np.float32)
+
+        tensor = torch.tensor(img_data)
+        target_shape = (480,480,240)    # h w d
+        
+        # Extract dimensions
+        h, w, d = tensor.shape
+        
+        # Calculate cropping/padding values for height, width, and depth
+        dh, dw, dd = target_shape
+        h_start = max((h - dh) // 2, 0)
+        h_end = min(h_start + dh, h)
+        w_start = max((w - dw) // 2, 0)
+        w_end = min(w_start + dw, w)
+        d_start = max((d - dd) // 2, 0)
+        d_end = min(d_start + dd, d)
+        
+        # # Add random shift to the crop region
+        # hw_max_shift = 128
+        # z_max_shift = 64
+        # if z_max_shift > 0 or hw_max_shift > 0:
+        #     # Generate random shifts within [-max_shift, max_shift]
+        #     h_shift = np.random.randint(-hw_max_shift, hw_max_shift + 1)
+        #     w_shift = np.random.randint(-hw_max_shift, hw_max_shift + 1)
+        #     d_shift = np.random.randint(-z_max_shift, z_max_shift + 1)
+
+        #     # Apply shifts, ensuring the crop region remains within the image bounds
+        #     h_start = max(h_start + h_shift, 0)
+        #     h_end = min(h_start + dh, h)
+        #     w_start = max(w_start + w_shift, 0)
+        #     w_end = min(w_start + dw, w)
+        #     d_start = max(d_start + d_shift, 0)
+        #     d_end = min(d_start + dd, d)
+
+        # Crop or pad the tensor
+        tensor = tensor[h_start:h_end, w_start:w_end, d_start:d_end]
+
+        pad_h_before = (dh - tensor.size(0)) // 2
+        pad_h_after = dh - tensor.size(0) - pad_h_before
+
+        pad_w_before = (dw - tensor.size(1)) // 2
+        pad_w_after = dw - tensor.size(1) - pad_w_before
+
+        pad_d_before = (dd - tensor.size(2)) // 2
+        pad_d_after = dd - tensor.size(2) - pad_d_before
+
+        tensor = torch.nn.functional.pad(tensor, (pad_d_before, pad_d_after, pad_w_before, pad_w_after, pad_h_before, pad_h_after), value=-1)
+
+        tensor = tensor.permute(2, 0, 1)    # d h w
+
+        tensor = tensor.unsqueeze(0)    # 1 d h w
+
+        return tensor
+    
+    def load_2d_image_to_tensor(self, image_path,resize=False):
+            """
+            Load a 2D grayscale image and convert it to a tensor with dimensions [1, 1, 480, 480].
+            
+            Args:
+                image_path (str): Path to the image file.
+                
+            Returns:
+                tensor (torch.Tensor): Tensor with shape [1, 1, 480, 480].
+            """
+            # Load the image as grayscale
+            image = Image.open(image_path).convert('L')
+            
+            # Resize to 480x480 if needed
+            if resize:
+                if image.size != (480, 480):
+                    image = image.resize((480, 480), Image.BILINEAR)
+            else:
+                if image.size != (480, 480):
+                    # 计算可以开始裁剪的最大左上角坐标
+                    width, height = image.size
+                    crop_size = 480
+                    if width < crop_size or height < crop_size:
+                        raise ValueError(f"Image size {image.size} is smaller than the crop size {crop_size}.")
+                    max_left = width - crop_size
+                    max_top = height - crop_size
+                    
+                    # 随机选择裁剪的左上角坐标
+                    left = random.randint(0, max(0, max_left))
+                    top = random.randint(0, max(0, max_top))
+                    
+                    # 裁剪图像
+                    image = image.crop((left, top, left + crop_size, top + crop_size))
+                
+
+            # Convert to numpy array and normalize to [0, 1]
+            img_array = np.array(image, dtype=np.float32) / 255.0
+            
+            # Convert to tensor
+            tensor = torch.from_numpy(img_array)
+            
+            # Add batch and channel dimensions
+            tensor = tensor.unsqueeze(0).unsqueeze(0)  # [1, 1, 480, 480]
+            
+            return tensor
+
+    def __getitem__(self, index):
+        """
+        Returns:
+            video_tensor (tensor): N, D, H, W
+            input_text (List of str): N
+        """
+        anatomy = np.random.choice(self.anatomy_ls, size=1, p=self.anatomy_weight)[0]
+        similarity = self.anatomy2simi_tab[anatomy]
+        num_samples = similarity.shape[0]
+        
+        # Get upper-triangular indices
+        rows, cols = np.triu_indices(num_samples, k=1)
+        pos_mask = (similarity[rows, cols] > self.positive_threshold) & (rows != cols)
+        neg_mask = (similarity[rows, cols] < self.negative_threshold) & (rows != cols)
+        pos_candidate_pairs = list(zip(rows[pos_mask], cols[pos_mask]))
+        neg_candidate_pairs = list(zip(rows[neg_mask], cols[neg_mask]))
+        
+        # Sampling
+        pos_pair_count = neg_pair_count = self.local_batch_size // 4
+        if self.local_batch_size % 4 > 1:
+            pos_pair_count += 1
+        if len(pos_candidate_pairs) >= pos_pair_count:
+            selected_pos = random.sample(pos_candidate_pairs, pos_pair_count)
+        else:
+            selected_pos = random.sample(list(zip(rows, cols)), pos_pair_count)
+        if len(neg_candidate_pairs) >= neg_pair_count:
+            selected_neg = random.sample(neg_candidate_pairs, neg_pair_count)
+        else:
+            selected_neg = random.sample(list(zip(rows, cols)), neg_pair_count)
+        sampled_indexes = [idx for pair in (selected_pos + selected_neg) for idx in pair]
+        
+        # If local_batch_size is odd, add one more sample
+        if self.local_batch_size % 2 == 1:
+            remaining = list(set(range(num_samples)) - set(sampled_indexes))
+            extra = random.choice(remaining) if remaining else random.choice(list(range(num_samples)))
+            sampled_indexes.append(extra)
+        
+        assert len(sampled_indexes) == self.local_batch_size
+
+        sampled_ids = [self.anatomy2id_ls[anatomy][idx] for idx in sampled_indexes]
+        # sampled_ids = [sample_id for sample_id in sampled_ids if sample_id in self.id2image_path]   # WARNING 假如有一个sample_id的image不存在，那么会导致stacked_video_tensor后面的元素整体向前移动一位，和similarity_tab对不上
+        
+        similarity_tab = np.zeros(
+            (self.local_batch_size, self.local_batch_size), 
+            dtype=np.uint8  # 使用 uint8 类型
+        )
+                
+        # 修改赋值逻辑（需要反向缩放）
+        for local_i, global_i in enumerate(sampled_indexes):
+            for local_j, global_j in enumerate(sampled_indexes):
+                raw_value = max(
+                    self.anatomy2simi_tab[anatomy][global_i, global_j],
+                    self.anatomy2simi_tab[anatomy][global_j, global_i]
+                )
+                # 直接使用存储的 uint8 值（无需转换）
+                similarity_tab[local_i, local_j] = raw_value  # 已经是 0~100 的整数
+        # # 修改
+        video_tensors = []
+
+        if '3D' in self.modality:
+            for sampled_id in sampled_ids:
+                try:
+                    # 返回的 video_tensor 已经是 [1, D, H, W]
+                    if self.is_train:
+                        video_tensor = self.nii_img_to_tensor(self.id2image_path[sampled_id])
+                    else:
+                        video_tensor = self.MLS_nii_img_to_tensor(self.id2image_path[sampled_id])
+                    assert video_tensor.dim() == 4, f"Expected 4D tensor, got {video_tensor.dim()}D tensor for {sampled_id}"
+                    if self.augmentator is not None and self.is_train:
+                        video_tensor = self.augmentator(video_tensor)
+                    video_tensors.append(video_tensor)
+
+                except Exception as e:
+                    print(f"Error processing {sampled_id}: {e}. Skipping.")
+                    # 如果一个样本失败，为了保持批次大小，可以加载一个全零张量或重复上一个
+                    if video_tensors:
+                        video_tensors.append(torch.zeros_like(video_tensors[-1]))
+                    else: # 如果第一个就失败了
+                        video_tensors.append(torch.zeros((1, 240, 480, 480), dtype=torch.float32))
+            
+            # 此时 video_tensors 是一个列表，包含 N 个 [1, D, H, W] 的张量
+            stacked_video_tensor = torch.cat(video_tensors, dim=0) # -> [N, D, H, W]
+            
+            # if self.augmentator is not None:
+            #     stacked_video_tensor = self.augmentator(stacked_video_tensor) # MONAI增广可以直接处理 [N, D, H, W]
+            stacked_video_tensor = stacked_video_tensor.unsqueeze(1) # -> [N, 1, D, H, W]
+            
+        elif '2D' in self.modality:
+            for sampled_id in sampled_ids:
+                try:
+                    # 返回的 video_tensor 已经是 [1, 1, H, W]
+                    video_tensor = self.load_2d_image_to_tensor(self.id2image_path[sampled_id],resize=True)
+                    video_tensor = video_tensor.squeeze().unsqueeze(0)  # 确保是 [1, H, W]
+                    if self.augmentator is not None and self.is_train:
+                        video_tensor = self.augmentator(video_tensor)
+                    
+                    video_tensor = video_tensor.unsqueeze(0)  # 变成 [1, 1, H, W]
+                    video_tensors.append(video_tensor)
+                except Exception as e:
+                    print(f"Error processing {sampled_id}: {e}. Skipping.")
+                    # 如果一个样本失败，为了保持批次大小，可以加载一个全零张量或重复上一个
+                    if video_tensors:
+                        video_tensors.append(torch.zeros_like(video_tensors[-1]))
+                    else:
+                        video_tensors.append(torch.zeros((1, 1, 480, 480), dtype=torch.float32))
+            # 此时 video_tensors 是一个列表，包含 N 个 [1, 1, H, W] 的张量
+            stacked_video_tensor = torch.cat(video_tensors, dim=0) # -> [N, 1, H, W]
+            
+            # 确保最终形状符合预期
+            stacked_video_tensor = stacked_video_tensor.unsqueeze(1) # -> [N, 1, 1, H, W]
+        if isinstance(stacked_video_tensor, MetaTensor):
+            stacked_video_tensor = stacked_video_tensor.as_tensor()
+        if self.modal_embedding:
+            return {'video_tensor': stacked_video_tensor, 'similarity_tab': torch.tensor(similarity_tab), 'anatomy': anatomy, 'sampled_ids':sampled_ids, 'modality':modality_dict[self.modality]}
+        else:
+            return {'video_tensor': stacked_video_tensor, 'similarity_tab': torch.tensor(similarity_tab), 'anatomy': anatomy, 'sampled_ids':sampled_ids}
+
+def collate_fn(batch):
+    # batch 是一个列表，其中每个元素是 __getitem__ 返回的字典
+    
+    # 使用列表推导式和 torch.stack 高效地组合批次
+    videos = torch.stack([item['video_tensor'] for item in batch], dim=0)
+    similarity_tabs = torch.stack([item['similarity_tab'] for item in batch], dim=0)
+    
+    # 其他非张量数据正常收集
+    anatomies = [item['anatomy'] for item in batch]
+    sampled_ids_list = [item['sampled_ids'] for item in batch]
+    if 'modality' in batch[0]:
+        modalities = torch.tensor([item['modality'] for item in batch])
+        return videos, similarity_tabs, anatomies, sampled_ids_list, modalities
+    else:
+        return videos, similarity_tabs, anatomies, sampled_ids_list
+
+# if __name__ == '__main__':
+#     from torch.utils.data import Dataset, DataLoader, random_split
+#     print('** CXR **')
+#     dataset = Conditional_CTReportDataset_Train(
+#         local_batch_size=8,
+#         modality='2D',
+#         jsonl_file='/mnt/petrelfs/zhangtengfei/RadIR/dataset/CXR/MIMIC_CXR/train.jsonl',
+#         csv_file_dir='/mnt/petrelfs/zhangtengfei/RadIR/dataset/CXR/MIMIC_CXR/train_entity',
+#         npy_file_dir='/mnt/petrelfs/zhangtengfei/RadIR/dataset/CXR/MIMIC_CXR/train_ratescore',
+#         positive_threshold=0.75,
+#         negative_threshold=1,
+#         anatomy_filter=['left hemidiaphragm', 'airspace']
+#     )
+    
+#     dl = DataLoader(
+#         dataset,
+#         num_workers=8,
+#         batch_size=1,
+#         shuffle = False,
+#         drop_last = False,
+#         pin_memory = False,
+#         collate_fn = collate_fn
+#     )        
+
+#     dl_iter = iter(dl)
+    
+#     for threshold in [0.1, 0.15, 0.2, 0.25, 0.3]:
+    
+#         count = 0
+#         triplet_count = 0
+        
+#         while True:
+            
+#             count += 1
+#             if count > 50:
+#                 break
+            
+#             video, similarity_tab, anatomy, batched_sampled_ids = next(dl_iter)   # 1 8 8
+            
+#             print(batched_sampled_ids)
+#             print(similarity_tab)
+                    
+#             mask = (similarity_tab.unsqueeze(3) - similarity_tab.unsqueeze(2)) > 0.2   # N N N
+#             triplet_count += torch.sum(mask)
+            
+#         print(f'** {threshold} **')
+#         print(similarity_tab.shape)
+#         print(triplet_count/count)
+    
+#     # > 0.1 时 8x8x8 的 matrix 上平均有 _ 个有效对
+#     # > 0.15 时 8x8x8 的 matrix 上平均有 _ 个有效对
+#     # > 0.2 时 8x8x8 的 matrix 上平均有 _ 个有效对
+#     # > 0.25 时 8x8x8 的 matrix 上平均有 _ 个有效对
+#     # > 0.3 时 8x8x8 的 matrix 上平均有 _ 个有效��
+#     print('** CT **')
+#     dataset = Conditional_CTReportDataset_Train(
+#         local_batch_size=8,
+#         modality='3D',
+#         jsonl_file='/mnt/petrelfs/zhangtengfei/RadIR/dataset/CT_RATE/train_replaced.jsonl',
+#         csv_file_dir='/mnt/petrelfs/zhangtengfei/RadIR/dataset/CT_RATE/anatomy/train_entity',
+#         npy_file_dir='/mnt/petrelfs/zhangtengfei/RadIR/dataset/CT_RATE/anatomy/train_ratescore',
+#         positive_threshold=0.75,
+#         negative_threshold=1,
+#         anatomy_filter=['right rib', 'thyroid gland']
+#     )
+    
+#     dl = DataLoader(
+#         dataset,
+#         num_workers=8,
+#         batch_size=1,
+#         shuffle = False,
+#         drop_last = False,
+#         pin_memory = False,
+#         collate_fn = collate_fn
+#     )        
+    
+#     dl_iter = iter(dl)
+    
+#     for threshold in [0.1, 0.15, 0.2, 0.25, 0.3]:
+    
+#         count = 0
+#         triplet_count = 0
+        
+#         while True:
+            
+#             count += 1
+#             if count > 50:
+#                 break
+            
+#             video, similarity_tab, anatomy, batched_sampled_ids = next(dl_iter)   # 1 8 8
+            
+#             print(batched_sampled_ids)
+#             print(similarity_tab)
+                    
+#             mask = (similarity_tab.unsqueeze(3) - similarity_tab.unsqueeze(2)) > 0.2   # N N N
+#             triplet_count += torch.sum(mask)
+            
+#         print(f'** {threshold} **')
+#         print(similarity_tab.shape)
+#         print(triplet_count/count)
+    
+if __name__ == '__main__':
+    from torch.utils.data import Dataset, DataLoader, random_split
+    import time
+    from tqdm import tqdm
+    
+    print('** 测试数据加载耗时 **')
+    
+    # 创建CXR数据集
+    print('** 创建CXR数据集 **')
+    cxr_dataset = Conditional_CTReportDataset_Train(
+        local_batch_size=35,
+        modality='2D',
+        
+        jsonl_file=None,  # ⚠️ 你没有 .jsonl 文件，直接设为 None 或删除此参数
+        csv_file_dir='/mnt/data-3/users/zhangzeyu/HMClip/RadIR/dataset/MIMIC-IR/anatomy/train_caption',
+        npy_file_dir='/mnt/data-3/users/zhangzeyu/HMClip/RadIR/dataset/MIMIC-IR/anatomy/train_ratescore',
+        
+        positive_threshold=0.75,
+        negative_threshold=1,
+        anatomy_filter=['lungs', 'ventricle', 'right upper lobe'],
+        need_aug=True
+    )
+    
+    # 创建CT数据集
+    print('** 创建CT数据集 **')
+    ct_dataset = Conditional_CTReportDataset_Train(
+        local_batch_size=7,
+        modality='3D',
+        jsonl_file   = '/mnt/data-3/users/zhangzeyu/HMClip/RadIR/dataset/CTRATE-IR/train.jsonl',
+        csv_file_dir = '/mnt/data-3/users/zhangzeyu/HMClip/RadIR/dataset/CTRATE-IR/anatomy/train_entity',
+        npy_file_dir = '/mnt/data-3/users/zhangzeyu/HMClip/RadIR/dataset/CTRATE-IR/anatomy/train_ratescore_compressed',
+        positive_threshold=0.75,
+        negative_threshold=1,
+        anatomy_filter=['bone', 'pulmonary artery', 'stomach'],
+        need_aug=True
+    )
+    
+    # 测试不同的工作线程数量
+    for num_workers in [2, 4,6, 8,12]:
+        print(f"\n使用 {num_workers} 个工作线程")
+        
+        # 测试CXR数据集
+        dl_cxr = DataLoader(
+            cxr_dataset,
+            num_workers=num_workers,
+            batch_size=2,
+            shuffle=False,
+            drop_last=False,
+            pin_memory=True,
+            collate_fn=collate_fn
+        )
+        
+        # 预热（初始化可能很慢）
+        print("预热CXR数据加载器...")
+        for _ in range(5):
+            next(iter(dl_cxr))
+        
+        # 测量CXR数据加载时间
+        print("开始测量CXR数据加载时间...")
+        start_time = time.time()
+        total_batches = min(20, len(dl_cxr))  # 限制测试数量
+        
+        for i, _ in enumerate(tqdm(dl_cxr, total=total_batches)):
+            if i >= total_batches - 1:
+                break
+        
+        cxr_time = time.time() - start_time
+        print(f"CXR数据集: {total_batches}个批次耗时 {cxr_time:.2f}秒, 平均每批次 {cxr_time/total_batches:.4f}秒")
+        
+        # 测试CT数据集
+        dl_ct = DataLoader(
+            ct_dataset,
+            num_workers=num_workers,
+            batch_size=1,
+            shuffle=False,
+            drop_last=False,
+            pin_memory=True,
+            collate_fn=collate_fn
+        )
+        
+        # 预热
+        print("预热CT数据加载器...")
+        for _ in range(5):
+            next(iter(dl_ct))
+        
+        # 测量CT数据加载时间
+        print("开始测量CT数据加载时间...")
+        start_time = time.time()
+        total_batches = min(20, len(dl_ct))  # 限制测试数量
+        
+        for i, _ in enumerate(tqdm(dl_ct, total=total_batches)):
+            if i >= total_batches - 1:
+                break
+        
+        ct_time = time.time() - start_time
+        print(f"CT数据集: {total_batches}个批次耗��� {ct_time:.2f}秒, 平均每批次 {ct_time/total_batches:.4f}秒")
+
+    # # 测试不同的batch_size
+    # num_workers = 8  # 使用固定的工作线程数
+    # print(f"\n使用固定的{num_workers}个工作线程，测试不同batch_size")
+    
+    # for batch_size in [1, 2, 4, 8, 16]:
+    #     print(f"\nbatch_size = {batch_size}")
+        
+    #     # CXR数据集
+    #     dl_cxr = DataLoader(
+    #         cxr_dataset,
+    #         num_workers=num_workers,
+    #         batch_size=batch_size,
+    #         shuffle=False,
+    #         drop_last=False,
+    #         pin_memory=True,
+    #         collate_fn=collate_fn
+    #     )
+        
+    #     # 预热
+    #     print("预热CXR数据加载器...")
+    #     for _ in range(3):
+    #         next(iter(dl_cxr))
+        
+    #     # 测量CXR数据加载时间
+    #     print("开始测量CXR数据加载时间...")
+    #     start_time = time.time()
+    #     total_batches = min(50, len(dl_cxr))  # 限制测试数量
+        
+    #     for i, _ in enumerate(tqdm(dl_cxr, total=total_batches)):
+    #         if i >= total_batches - 1:
+    #             break
+        
+    #     cxr_time = time.time() - start_time
+    #     print(f"CXR数据集(batch_size={batch_size}): {total_batches}个批次耗时 {cxr_time:.2f}秒, "
+    #           f"平均每批次 {cxr_time/total_batches:.4f}秒, "
+    #           f"平均每样本 {cxr_time/(total_batches*batch_size):.4f}秒")
+
+

scripts/data_process.py

@@ -0,0 +1,177 @@
+import nibabel as nib
+import torch
+import numpy as np
+import torch.nn.functional as F
+from PIL import Image
+import random
+
+
+import os
+
+
+
+def resize_array(array, current_spacing, target_spacing):
+    """
+    Resize the array to match the target spacing.
+
+    Args:
+    array (torch.Tensor): Input array to be resized.
+    current_spacing (tuple): Current voxel spacing (z_spacing, xy_spacing, xy_spacing).
+    target_spacing (tuple): Target voxel spacing (target_z_spacing, target_x_spacing, target_y_spacing).
+
+    Returns:
+    np.ndarray: Resized array.
+    """
+    # Calculate new dimensions
+    original_shape = array.shape[2:]
+    scaling_factors = [
+        current_spacing[i] / target_spacing[i] for i in range(len(original_shape))
+    ]
+    new_shape = [
+        int(original_shape[i] * scaling_factors[i]) for i in range(len(original_shape))
+    ]
+    # Resize the array
+    resized_array = F.interpolate(array, size=new_shape, mode='trilinear', align_corners=False).cpu().numpy()
+    return resized_array
+
+def load_3d_image_to_tensor(path):
+
+    nii_img = safe_load_image(sample["image"])
+    img_data = nii_img.get_fdata()
+    img_data = np.flip(img_data, axis=0)
+    img_data = np.flip(img_data, axis=1)
+
+    # WARNING Respacing
+    img_data = img_data.transpose(2, 0, 1)
+    img_data = np.copy(img_data)
+    tensor = torch.tensor(img_data)
+    tensor = tensor.unsqueeze(0).unsqueeze(0)
+    
+    target_x_spacing = 0.75
+    target_y_spacing = 0.75
+    target_z_spacing = 1.5
+    current = (3, 1, 1)   # this is all set to 3 1 1
+    target = (target_z_spacing, target_x_spacing, target_y_spacing)
+    
+    img_data = resize_array(tensor, current, target)
+    img_data = img_data[0][0]
+    img_data= np.transpose(img_data, (1, 2, 0))
+
+    # WARNING Normalization 2
+    hu_min, hu_max = -1000, 1000
+    img_data = np.clip(img_data, hu_min, hu_max)
+    img_data = (((img_data ) / 1000)).astype(np.float32)
+
+    tensor = torch.tensor(img_data)
+    
+    # WARNING Padding or Crop
+    
+    # Get the dimensions of the input tensor
+    target_shape = (480,480,240)    # h w d
+    
+    # Extract dimensions
+    h, w, d = tensor.shape
+    
+    # Calculate cropping/padding values for height, width, and depth
+    dh, dw, dd = target_shape
+    h_start = max((h - dh) // 2, 0)
+    h_end = min(h_start + dh, h)
+    w_start = max((w - dw) // 2, 0)
+    w_end = min(w_start + dw, w)
+    d_start = max((d - dd) // 2, 0)
+    d_end = min(d_start + dd, d)
+
+    # Crop or pad the tensor
+    tensor = tensor[h_start:h_end, w_start:w_end, d_start:d_end]
+
+    pad_h_before = (dh - tensor.size(0)) // 2
+    pad_h_after = dh - tensor.size(0) - pad_h_before
+
+    pad_w_before = (dw - tensor.size(1)) // 2
+    pad_w_after = dw - tensor.size(1) - pad_w_before
+
+    pad_d_before = (dd - tensor.size(2)) // 2
+    pad_d_after = dd - tensor.size(2) - pad_d_before
+
+    tensor = torch.nn.functional.pad(tensor, (pad_d_before, pad_d_after, pad_w_before, pad_w_after, pad_h_before, pad_h_after), value=-1)
+
+    tensor = tensor.permute(2, 0, 1)    # d h w
+
+    tensor = tensor.unsqueeze(0)    # 1 d h w
+
+    return tensor
+
+
+def load_2d_image_to_tensor(image_path,resize=False):
+    """
+    Load a 2D grayscale image and convert it to a tensor with dimensions [1, 1, 480, 480].
+    
+    Args:
+        image_path (str): Path to the image file.
+        
+    Returns:
+        tensor (torch.Tensor): Tensor with shape [1, 1, 480, 480].
+    """
+    # Load the image as grayscale
+    image = Image.open(image_path).convert('L')
+    
+    # Resize to 480x480 if needed
+    if resize:
+        if image.size != (480, 480):
+            image = image.resize((480, 480), Image.BILINEAR)
+    else:
+        if image.size != (480, 480):
+            # 计算可以开始裁剪的最大左上角坐标
+            width, height = image.size
+            crop_size = 480
+            if width < crop_size or height < crop_size:
+                raise ValueError(f"Image size {image.size} is smaller than the crop size {crop_size}.")
+            max_left = width - crop_size
+            max_top = height - crop_size
+            
+            # 随机选择裁剪的左上角坐标
+            left = random.randint(0, max(0, max_left))
+            top = random.randint(0, max(0, max_top))
+            
+            # 裁剪图像
+            image = image.crop((left, top, left + crop_size, top + crop_size))
+        
+
+    # Convert to numpy array and normalize to [0, 1]
+    img_array = (np.array(image, dtype=np.float32) / 255.0) * 2 -1
+    
+    # Convert to tensor
+    tensor = torch.from_numpy(img_array)
+    
+    # Add batch and channel dimensions
+    tensor = tensor.unsqueeze(0).unsqueeze(0)  # [1, 1, 480, 480]
+    
+    return tensor
+
+def safe_load_image(path):
+    """
+    Robust image loader that supports both .nii.gz and .npy files.
+    If .nii.gz is missing, automatically falls back to _compressed.npy.
+    """
+    try:
+        if os.path.exists(path):
+            if path.endswith('.nii.gz'):
+                data = nib.load(path).get_fdata()
+                return data
+            elif path.endswith('.npy'):
+                data = np.load(path)
+                return data
+
+        # fallback to compressed npy
+        npy_fallback = path.replace('.nii.gz', '_compressed.npy')
+        if os.path.exists(npy_fallback):
+            data = np.load(npy_fallback)
+            print(f"[Info] Fallback to compressed file: {npy_fallback}")
+            return data
+
+        # all failed
+        raise FileNotFoundError(f"{path} and fallback {npy_fallback} not found")
+
+    except Exception as e:
+        print(f"[Warning] Error loading {path}: {e}")
+        return None

scripts/image_traversals_compare.py

@@ -0,0 +1,1186 @@
+#!/usr/bin/env python3
+"""
+MERU-style image traversals comparing RadIR/RadIR (hyperbolic) vs RadIR_compare (Euclidean).
+
+For each query CXR:
+  1. Encode the image with both models.
+  2. Walk along the geodesic / ray from the image embedding toward the origin [ROOT].
+  3. At each step, retrieve the nearest neighbor from a shared multi-granularity concept pool.
+     Hyperbolic path optionally keeps only texts that entail the current point
+     (MERU-style cone filter), then takes NN in that subset.
+  4. Render a MERU Figure-5 style side-by-side panel.
+
+Usage (from RadIR/RadIR):
+  python scripts/image_traversals_compare.py \\
+    --images path1.jpg path2.jpg \\
+    --out outputs/traversals_compare.png
+
+  # Auto-pick diverse disease cases from valid.jsonl + disease labels:
+  python scripts/image_traversals_compare.py --auto_pick 4 --out outputs/traversals_compare.png
+
+  # Disable entailment filter (pure hyperbolic NN):
+  python scripts/image_traversals_compare.py --auto_pick 4 --no_entailment
+"""
+
+from __future__ import annotations
+
+import argparse
+import csv
+import gc
+import json
+import os
+import re
+import sys
+import textwrap
+from pathlib import Path
+from typing import Sequence
+
+import matplotlib.pyplot as plt
+import numpy as np
+import torch
+import torch.nn.functional as F
+from matplotlib.patches import FancyBboxPatch
+from PIL import Image
+
+ROOT = Path(__file__).resolve().parents[1]
+# Workspace layout: HMClip/RadIR/RadIR and HMClip/RadIR_compare
+_hmclip = ROOT.parent.parent
+if (_hmclip / "RadIR_compare").is_dir():
+    COMPARE_ROOT = _hmclip / "RadIR_compare"
+elif (ROOT.parent / "RadIR_compare").is_dir():
+    COMPARE_ROOT = ROOT.parent / "RadIR_compare"
+else:
+    COMPARE_ROOT = _hmclip / "RadIR_compare"
+
+ROOT_LABEL = "[ROOT]"
+
+
+# ---------------------------------------------------------------------------
+# Concept pool
+# ---------------------------------------------------------------------------
+
+def load_concept_pool(
+    concepts_file: str | None,
+    extra_captions: int = 0,
+    caption_jsonl: str | None = None,
+    seed: int = 42,
+) -> list[str]:
+    texts: list[str] = []
+    seen: set[str] = set()
+
+    def _add(s: str) -> None:
+        s = " ".join(s.strip().split())
+        if not s:
+            return
+        key = s.lower()
+        if key in seen or key == ROOT_LABEL.lower():
+            return
+        seen.add(key)
+        texts.append(s)
+
+    path = concepts_file or str(ROOT / "scripts" / "traversal_concepts.txt")
+    with open(path, encoding="utf-8") as f:
+        for line in f:
+            line = line.strip()
+            if not line or line.startswith("#"):
+                continue
+            _add(line)
+
+    if extra_captions > 0 and caption_jsonl and os.path.isfile(caption_jsonl):
+        rng = np.random.default_rng(seed)
+        rows = []
+        with open(caption_jsonl, encoding="utf-8") as f:
+            for line in f:
+                item = json.loads(line)
+                text = str(item.get("text", "")).strip()
+                if not text:
+                    continue
+                first = text.split(".")[0].strip()
+                if 8 <= len(first) <= 80:
+                    rows.append(first)
+        if rows:
+            idx = rng.choice(len(rows), size=min(extra_captions, len(rows)), replace=False)
+            for i in idx:
+                _add(rows[int(i)])
+
+    texts.append(ROOT_LABEL)
+    return texts
+
+
+# ---------------------------------------------------------------------------
+# Image I/O
+# ---------------------------------------------------------------------------
+
+def resolve_image_path(path: str, dataset_root: str) -> str:
+    candidates = [
+        path,
+        os.path.join(dataset_root, path),
+        os.path.join(dataset_root, path.replace("dataset/", "", 1)),
+        os.path.join(str(ROOT), path),
+        os.path.join(str(ROOT), path.replace("dataset/", "dataset/", 1)),
+    ]
+    for c in candidates:
+        if os.path.isfile(c):
+            return os.path.abspath(c)
+    raise FileNotFoundError(f"Image not found: {path}")
+
+
+def load_image_tensor(path: str, imsize: int, with_depth: bool) -> torch.Tensor:
+    """
+    Grayscale -> [-1, 1].
+    with_depth=False -> [1, H, W]   (RadIR/RadIR hyp)
+    with_depth=True  -> [1, 1, H, W] (RadIR_compare, C×D×H×W)
+    """
+    img = Image.open(path).convert("L")
+    arr = np.array(img, dtype=np.float32) / 255.0
+    t = torch.from_numpy(arr).unsqueeze(0)  # [1, H, W]
+    t = F.interpolate(
+        t.unsqueeze(0), size=(imsize, imsize), mode="bilinear", align_corners=False
+    ).squeeze(0)
+    t = t * 2.0 - 1.0
+    if with_depth:
+        t = t.unsqueeze(0)  # [1, 1, H, W]
+    return t
+
+
+def preprocess_hyp_batch(image: torch.Tensor) -> torch.Tensor:
+    if image.ndim == 3:
+        image = image.unsqueeze(0)
+    if image.ndim == 4:
+        if image.shape[1] != 1:
+            image = image.mean(dim=1, keepdim=True)
+    elif image.ndim == 5:
+        if image.shape[1] != 1:
+            image = image.mean(dim=1, keepdim=True)
+    else:
+        raise ValueError(f"Unexpected image shape {tuple(image.shape)}")
+    return image
+
+
+# ---------------------------------------------------------------------------
+# Auto-pick demo images
+# ---------------------------------------------------------------------------
+
+def auto_pick_images(
+    dataset_root: str,
+    labels_csv: str,
+    jsonl: str,
+    n: int,
+    seed: int = 42,
+) -> list[str]:
+    labels = {}
+    with open(labels_csv, encoding="utf-8") as f:
+        for row in csv.DictReader(f):
+            labels[row["dicom_id"]] = row
+
+    wanted = [
+        "normal_only",
+        "Pneumonia_pos",
+        "Effusion_pos",
+        "Cardiomegaly_pos",
+        "Edema_pos",
+        "Atelectasis_pos",
+        "Pneumothorax_pos",
+    ]
+    found: dict[str, str] = {}
+    with open(jsonl, encoding="utf-8") as f:
+        for line in f:
+            item = json.loads(line)
+            did = os.path.splitext(os.path.basename(item["img_path"]))[0]
+            lab = labels.get(did)
+            if lab is None:
+                continue
+            for k in wanted:
+                if k in found or lab.get(k) != "1":
+                    continue
+                if k != "normal_only":
+                    others = sum(
+                        int(lab[x])
+                        for x in wanted
+                        if x not in ("normal_only", k)
+                    )
+                    if others > 0:
+                        continue
+                try:
+                    full = resolve_image_path(item["img_path"], dataset_root)
+                except FileNotFoundError:
+                    continue
+                found[k] = full
+                break
+            if len(found) >= n:
+                break
+
+    paths = [found[k] for k in wanted if k in found][:n]
+    if len(paths) < n:
+        with open(jsonl, encoding="utf-8") as f:
+            for line in f:
+                item = json.loads(line)
+                try:
+                    full = resolve_image_path(item["img_path"], dataset_root)
+                except FileNotFoundError:
+                    continue
+                if full not in paths:
+                    paths.append(full)
+                if len(paths) >= n:
+                    break
+    return paths
+
+
+# ---------------------------------------------------------------------------
+# Module path helpers (avoid radir / CTViT collisions between trees)
+# ---------------------------------------------------------------------------
+
+def _purge_modules(prefixes: Sequence[str]) -> None:
+    doomed = [
+        k
+        for k in list(sys.modules)
+        if any(k == p or k.startswith(p + ".") for p in prefixes)
+    ]
+    for k in doomed:
+        del sys.modules[k]
+
+
+def _push_paths(paths: Sequence[str]) -> list[str]:
+    old = list(sys.path)
+    for p in reversed(paths):
+        if p in sys.path:
+            sys.path.remove(p)
+        sys.path.insert(0, p)
+    return old
+
+
+# ---------------------------------------------------------------------------
+# Hyperbolic (RadIR/RadIR) encode + traverse
+# ---------------------------------------------------------------------------
+
+@torch.no_grad()
+def run_hyperbolic(
+    image_paths: list[str],
+    concepts: list[str],
+    ckpt: str,
+    model_path: str,
+    device: str,
+    imsize: int = 224,
+    steps: int = 50,
+    text_batch: int = 64,
+    use_entailment: bool = True,
+    aperture_thresh: float = 1.2,
+) -> tuple[list[list[str]], float]:
+    _purge_modules(["radir", "transformer_maskgit", "Rad_IR"])
+    rad_ir_root = str(ROOT / "Rad_IR")
+    # Inner package lives at ROOT/transformer_maskgit/transformer_maskgit/
+    old = _push_paths(
+        [str(ROOT), rad_ir_root, str(ROOT / "transformer_maskgit")]
+    )
+
+    try:
+        from transformers import AutoModel, AutoTokenizer
+
+        from radir import lorentz as L
+        from radir.RadIR import RADIR, l2norm
+        from transformer_maskgit.ctvit import CTViT
+
+        tokenizer = AutoTokenizer.from_pretrained(
+            model_path, trust_remote_code=True, local_files_only=True
+        )
+        text_model = AutoModel.from_pretrained(
+            model_path, trust_remote_code=True, local_files_only=True
+        )
+        image_encoder = CTViT(
+            dim=768,
+            codebook_size=8192,
+            image_size=imsize,
+            patch_size=16,
+            temporal_patch_size=10,
+            spatial_depth=8,
+            temporal_depth=6,
+            cls_depth=4,
+            dim_head=32,
+            heads=8,
+            channels=1,
+        )
+        model = RADIR(
+            tokenizer=tokenizer,
+            image_encoder=image_encoder,
+            text_encoder=text_model,
+            dim_text=768,
+            dim_image=512,
+            dim_latent=512,
+            use_mlm=False,
+            use_all_token_embeds=False,
+        )
+        model.load(ckpt)
+        model.to(device)
+        model.eval()
+        curv = model.get_curv().detach()
+
+        txt_feats = []
+        for i in range(0, len(concepts), text_batch):
+            batch = concepts[i : i + text_batch]
+            real = [t for t in batch if t != ROOT_LABEL]
+            feats_map: dict[str, torch.Tensor] = {}
+            if real:
+                tok = tokenizer(
+                    real,
+                    return_tensors="pt",
+                    padding=True,
+                    truncation=True,
+                    max_length=128,
+                )
+                tok = {k: v.to(device) for k, v in tok.items()}
+                if "token_type_ids" not in tok:
+                    tok["token_type_ids"] = torch.zeros_like(tok["input_ids"])
+                enc = model.text_transformer(
+                    input_ids=tok["input_ids"],
+                    token_type_ids=tok["token_type_ids"],
+                    attention_mask=tok["attention_mask"],
+                )[0]
+                cls = enc[:, 0, :]
+                hyp = model.project_txt(cls)
+                for t, h in zip(real, hyp):
+                    feats_map[t] = h
+            for t in batch:
+                if t == ROOT_LABEL:
+                    dim = next(iter(feats_map.values())).shape[-1] if feats_map else 768
+                    txt_feats.append(torch.zeros(dim, device=device))
+                else:
+                    txt_feats.append(feats_map[t])
+        text_pool = torch.stack(txt_feats, dim=0).float()
+        root_idx = concepts.index(ROOT_LABEL) if ROOT_LABEL in concepts else -1
+        text_norms = torch.norm(text_pool, dim=-1)
+        text_aperture = None
+        if use_entailment:
+            text_aperture = L.half_aperture(text_pool, curv)
+            print(
+                f"Entailment filter ON (soft; aperture_thresh={aperture_thresh}, "
+                f"no hard radius gate — looser than training for traversal viz)"
+            )
+
+        all_paths: list[list[str]] = []
+        fallback_count = 0
+        total_steps = 0
+        for ip in image_paths:
+            img = load_image_tensor(ip, imsize=imsize, with_depth=False)
+            imgs = preprocess_hyp_batch(img.unsqueeze(0).to(device))
+            enc = model.visual_transformer(
+                imgs,
+                return_encoded_tokens=True,
+                modal_embedding=False,
+                modal_indexs=None,
+                is_condition=False,
+            )
+            if enc.dim() == 2:
+                image_global = l2norm(enc)
+            else:
+                image_global = l2norm(enc[:, 0, :])
+            z0 = model.project_img(image_global)
+
+            u0 = L.log_map0(z0.float(), curv)
+            labels: list[str] = []
+            ts = torch.linspace(1.0, 0.0, steps, device=device)
+            for t in ts:
+                u_t = (t * u0).float()
+                z_t = L.exp_map0(u_t, curv)
+                dist = L.pairwise_dist(z_t, text_pool, curv)[0]
+                total_steps += 1
+
+                if use_entailment and text_aperture is not None:
+                    # MERU-style cone check, but softened for visualization:
+                    # 1) wider aperture_thresh than training (default 1.2 vs 0.7)
+                    # 2) soft penalty instead of hard mask (outside-cone texts can
+                    #    still win if much closer)
+                    # 3) no hard text_norm <= z_norm gate (that collapsed mid-path
+                    #    to only [ROOT] in practice)
+                    z_exp = z_t.expand_as(text_pool)
+                    angle = L.oxy_angle(text_pool, z_exp, curv)
+                    energy = torch.clamp(
+                        angle - aperture_thresh * text_aperture, min=0.0
+                    )
+                    if root_idx >= 0:
+                        energy[root_idx] = 0.0
+                    # soft_weight ~1: energy (radians) added to hyperbolic dist
+                    soft_weight = 1.0
+                    scores = dist + soft_weight * energy
+                    nn_idx = int(scores.argmin().item())
+                else:
+                    nn_idx = int(dist.argmin().item())
+                labels.append(concepts[nn_idx])
+            all_paths.append(labels)
+
+        if use_entailment and total_steps > 0:
+            print(
+                f"Entailment: soft penalty (weight=1.0), "
+                f"aperture_thresh={aperture_thresh}"
+            )
+
+        curv_val = float(curv.item())
+        del model, text_model, image_encoder, text_pool
+        return all_paths, curv_val
+    finally:
+        sys.path[:] = old
+        _purge_modules(["radir", "transformer_maskgit", "Rad_IR"])
+        gc.collect()
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()
+
+
+# ---------------------------------------------------------------------------
+# Euclidean (RadIR_compare) encode + traverse
+# ---------------------------------------------------------------------------
+
+@torch.no_grad()
+def run_euclidean(
+    image_paths: list[str],
+    concepts: list[str],
+    ckpt: str,
+    model_path: str,
+    device: str,
+    imsize: int = 480,
+    steps: int = 50,
+    text_batch: int = 32,
+) -> list[list[str]]:
+    _purge_modules(["radir", "transformer_maskgit", "Rad_IR", "hf_model_paths"])
+    compare_scripts = str(COMPARE_ROOT / "scripts")
+    compare_rad = str(COMPARE_ROOT / "Rad_IR")
+    old = _push_paths(
+        [
+            str(COMPARE_ROOT),
+            compare_rad,
+            compare_scripts,
+            str(COMPARE_ROOT / "transformer_maskgit"),
+        ]
+    )
+
+    try:
+        from transformers import BertModel, BertTokenizer
+
+        from radir.RadIR import RADIR, l2norm
+        from transformer_maskgit.ctvit import CTViT
+
+        tokenizer = BertTokenizer.from_pretrained(
+            model_path, local_files_only=True, do_lower_case=True
+        )
+        text_encoder = BertModel.from_pretrained(model_path, local_files_only=True)
+        image_encoder = CTViT(
+            dim=512,
+            codebook_size=8192,
+            image_size=imsize,
+            patch_size=20,
+            temporal_patch_size=10,
+            spatial_depth=8,
+            temporal_depth=6,
+            cls_depth=4,
+            dim_head=32,
+            heads=8,
+        )
+        model = RADIR(
+            image_encoder=image_encoder,
+            text_encoder=text_encoder,
+            tokenizer=tokenizer,
+            dim_text=768,
+            dim_image=512,
+            dim_latent=512,
+            use_mlm=False,
+            use_all_token_embeds=False,
+        )
+        model.load(ckpt)
+        model.to(device)
+        model.eval()
+
+        # Encode non-ROOT concepts only. On the unit sphere, ambient Euclidean
+        # distance to the origin always beats weakly-aligned texts, so ROOT must
+        # NOT compete via cdist — select it only when the query is near 0.
+        non_root = [t for t in concepts if t != ROOT_LABEL]
+        txt_feats = []
+        for i in range(0, len(non_root), text_batch):
+            batch = non_root[i : i + text_batch]
+            dummy = torch.zeros(len(batch), 1, 1, imsize, imsize, device=device)
+            tok = tokenizer(
+                batch,
+                return_tensors="pt",
+                padding="max_length",
+                truncation=True,
+                max_length=128,
+            ).to(device)
+            text_latents, _, _, _ = model(
+                tok,
+                dummy,
+                return_latents=True,
+                device=device,
+                is_condition=False,
+                modal_embedding=False,
+                modal_indexs=None,
+            )
+            for h in text_latents:
+                txt_feats.append(h if h.ndim == 1 else h.squeeze(0))
+        text_pool = l2norm(torch.stack(txt_feats, dim=0))
+        # MERU-CLIP protocol: interpolate image → [ROOT] on the sphere.
+        # With root = 0, lerp+normalize keeps direction fixed for t>0, so NN
+        # is typically sticky (few unique concepts); last steps → [ROOT].
+        root_feat = torch.zeros(text_pool.shape[-1], device=device)
+
+        all_paths: list[list[str]] = []
+        for ip in image_paths:
+            img = load_image_tensor(ip, imsize=imsize, with_depth=True)
+            video = img.unsqueeze(0).to(device)
+            tok = tokenizer(
+                ["chest radiograph"],
+                return_tensors="pt",
+                padding="max_length",
+                truncation=True,
+                max_length=128,
+            ).to(device)
+            _, image_latents, _, _ = model(
+                tok,
+                video,
+                return_latents=True,
+                device=device,
+                is_condition=False,
+                modal_embedding=False,
+                modal_indexs=None,
+            )
+            z0 = image_latents
+            if z0.ndim == 1:
+                z0 = z0.unsqueeze(0)
+            z0 = l2norm(z0).squeeze(0)
+
+            # MERU interpolate(): lerp(root, image, w), w: 0→1, then flip
+            # so order is image-first → root-last.
+            weights = torch.linspace(0.0, 1.0, steps, device=device)
+            interp = torch.stack(
+                [torch.lerp(root_feat, z0, float(w.item())) for w in weights],
+                dim=0,
+            )
+            interp = F.normalize(interp, dim=-1).flip(0)
+
+            labels: list[str] = []
+            for i, z_t in enumerate(interp):
+                # Near the root end of the walk, emit [ROOT] (zero vector is
+                # unstable under normalize; matches CLIP figure ending at ROOT).
+                if i >= int(steps * 0.88):
+                    labels.append(ROOT_LABEL)
+                    continue
+                sims = (z_t.float() @ text_pool.float().T)
+                nn_idx = int(sims.argmax().item())
+                labels.append(non_root[nn_idx])
+            all_paths.append(labels)
+
+        del model, text_encoder, image_encoder, text_pool
+        return all_paths
+    finally:
+        sys.path[:] = old
+        _purge_modules(["radir", "transformer_maskgit", "Rad_IR", "hf_model_paths"])
+        gc.collect()
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()
+
+
+# ---------------------------------------------------------------------------
+# Path compression (MERU figure style) + min-depth helpers
+# ---------------------------------------------------------------------------
+
+def _norm_concept(s: str) -> str:
+    """Normalize for near-duplicate detection (case / plural / spacing)."""
+    s = re.sub(r"\s+", " ", s.strip().lower())
+    toks = []
+    for t in s.split():
+        if len(t) > 4 and t.endswith("ies"):
+            t = t[:-3] + "y"  # abnormalities -> abnormality
+        elif len(t) > 3 and t.endswith("s") and not t.endswith(
+            ("ss", "ous", "us", "is", "as", "ies")
+        ):
+            t = t[:-1]
+        toks.append(t)
+    return " ".join(toks)
+
+
+def _pick_unused(
+    scores: torch.Tensor,
+    concepts: Sequence[str],
+    used: set[str],
+    root_idx: int = -1,
+    higher_is_better: bool = False,
+) -> int | None:
+    """Argmin/argmax among unused non-ROOT concepts."""
+    s = scores.clone()
+    ban = float("-inf") if higher_is_better else float("inf")
+    if root_idx >= 0:
+        s[root_idx] = ban
+    for i, c in enumerate(concepts):
+        if c == ROOT_LABEL or _norm_concept(c) in used:
+            s[i] = ban
+    if higher_is_better:
+        if not torch.isfinite(s).any():
+            return None
+        return int(s.argmax().item())
+    if not torch.isfinite(s).any():
+        return None
+    return int(s.argmin().item())
+
+
+def unique_body_from_score_rows(
+    score_rows: torch.Tensor,
+    concepts: Sequence[str],
+    root_idx: int,
+    min_body: int,
+    higher_is_better: bool = False,
+) -> list[str]:
+    """
+    Build an ordered unique concept list along a traversal.
+
+    1) Greedy: at each time step take nearest unused concept.
+    2) If fewer than min_body, re-sample evenly in time (stratified) forcing
+       unused picks so the figure can show ≥ min_concepts levels.
+    """
+    t_steps = int(score_rows.shape[0])
+    used: set[str] = set()
+    ordered: list[str] = []
+
+    for ti in range(t_steps):
+        idx = _pick_unused(
+            score_rows[ti],
+            concepts,
+            used,
+            root_idx=root_idx,
+            higher_is_better=higher_is_better,
+        )
+        if idx is None:
+            continue
+        lab = concepts[idx]
+        key = _norm_concept(lab)
+        if key in used:
+            continue
+        used.add(key)
+        ordered.append(lab)
+
+    if len(ordered) >= min_body:
+        return ordered
+
+    # Stratified fill: force one new concept per evenly spaced geodesic step.
+    used = set()
+    ordered = []
+    # Slightly more candidate slots than needed in case of ties / empties
+    n_slots = max(min_body, min(min_body + 2, t_steps))
+    slot_ids = np.linspace(0, t_steps - 1, n_slots)
+    slot_ids = sorted(set(int(round(i)) for i in slot_ids))
+    for ti in slot_ids:
+        idx = _pick_unused(
+            score_rows[ti],
+            concepts,
+            used,
+            root_idx=root_idx,
+            higher_is_better=higher_is_better,
+        )
+        if idx is None:
+            continue
+        lab = concepts[idx]
+        used.add(_norm_concept(lab))
+        ordered.append(lab)
+        if len(ordered) >= min_body:
+            break
+
+    # Last-resort: dump next-best at the first / mid / last rows
+    if len(ordered) < min_body:
+        for ti in (0, t_steps // 2, t_steps - 1):
+            idx = _pick_unused(
+                score_rows[ti],
+                concepts,
+                used,
+                root_idx=root_idx,
+                higher_is_better=higher_is_better,
+            )
+            if idx is None:
+                continue
+            lab = concepts[idx]
+            used.add(_norm_concept(lab))
+            ordered.append(lab)
+            if len(ordered) >= min_body:
+                break
+
+    return ordered
+
+
+def synthesize_dense_labels(body: list[str], steps: int) -> list[str]:
+    """Repeat each body label so compress_traversal can subsample evenly."""
+    if not body:
+        return [ROOT_LABEL] * max(steps, 1)
+    labels: list[str] = []
+    per = max(1, steps // len(body))
+    for lab in body:
+        labels.extend([lab] * per)
+    while len(labels) < steps:
+        labels.append(body[-1])
+    labels = labels[:steps]
+    # Reserve last few steps for ROOT (MERU ends at origin)
+    n_root = max(1, steps // 10)
+    labels[-n_root:] = [ROOT_LABEL] * n_root
+    return labels
+
+
+def compress_traversal(
+    labels: list[str],
+    keep: int = 6,
+) -> list[str]:
+    """Dedup consecutive (near-)identical labels; keep up to `keep` steps + [ROOT]."""
+    dedup: list[str] = []
+    for lab in labels:
+        if lab == ROOT_LABEL:
+            continue
+        if not dedup:
+            dedup.append(lab)
+            continue
+        # exact or near-duplicate of previous (e.g. abnormality / abnormalities)
+        if lab == dedup[-1] or _norm_concept(lab) == _norm_concept(dedup[-1]):
+            continue
+        dedup.append(lab)
+    if not dedup:
+        return [ROOT_LABEL]
+
+    # keep includes [ROOT]; body slots = keep - 1
+    body_keep = max(1, keep - 1)
+    if len(dedup) > body_keep:
+        idx = np.linspace(0, len(dedup) - 1, body_keep)
+        idx = sorted(set(int(round(i)) for i in idx))
+        dedup = [dedup[i] for i in idx]
+    dedup.append(ROOT_LABEL)
+    return dedup
+
+
+def align_two_paths(
+    left: list[str],
+    right: list[str],
+) -> tuple[list[str], list[str]]:
+    def split(p: list[str]) -> tuple[list[str], str]:
+        if p and p[-1] == ROOT_LABEL:
+            return p[:-1], ROOT_LABEL
+        return list(p), ROOT_LABEL
+
+    l_body, _ = split(left)
+    r_body, _ = split(right)
+    n = max(len(l_body), len(r_body), 1)
+
+    def pad(body: list[str], n: int) -> list[str]:
+        if len(body) >= n:
+            return body[:n]
+        if not body:
+            return ["↓"] * n
+        return body + ["↓"] * (n - len(body))
+
+    return pad(l_body, n) + [ROOT_LABEL], pad(r_body, n) + [ROOT_LABEL]
+
+
+# ---------------------------------------------------------------------------
+# Visualization
+# ---------------------------------------------------------------------------
+
+from matplotlib.font_manager import FontProperties
+
+# Liberation Serif ≈ Times New Roman (metric-compatible on Linux).
+_HEADER_FONT = FontProperties(family="Liberation Serif", weight="bold", size=10)
+_TITLE_FONT = FontProperties(family="Liberation Serif", weight="normal", size=12.5)
+
+
+def _cell_color(row: int, n_rows: int) -> tuple[float, float, float]:
+    """Single-hue soft slate → paper white (smooth, low-chroma)."""
+    if n_rows <= 1:
+        t = 1.0
+    else:
+        t = row / (n_rows - 1)
+    r = 0.788 + 0.164 * t
+    g = 0.839 + 0.124 * t
+    b = 0.886 + 0.086 * t
+    return (r, g, b)
+
+
+def _wrap_label(lab: str, width: int = 24) -> str:
+    """Prefer fewer, wider lines so text fills horizontal space and stays in-cell."""
+    if lab in ("↓", "▼", ROOT_LABEL) or len(lab) <= width:
+        return lab
+    lines = textwrap.wrap(
+        lab, width=width, break_long_words=False, break_on_hyphens=True
+    )
+    out = []
+    for line in lines:
+        if len(line) <= width + 2:
+            out.append(line)
+        else:
+            out.extend(
+                textwrap.wrap(
+                    line, width=width, break_long_words=True, break_on_hyphens=True
+                )
+            )
+    return "\n".join(out)
+
+
+def _label_font(lab: str, wrapped: str) -> FontProperties:
+    if lab == "↓":
+        # Compact arrow, not dominating the cell
+        return FontProperties(family="DejaVu Sans", weight="normal", size=6)
+    if lab == ROOT_LABEL:
+        return FontProperties(
+            family="Liberation Serif", style="normal", weight="normal", size=6.0
+        )
+    n_lines = wrapped.count("\n") + 1
+    longest = max(len(line) for line in wrapped.split("\n"))
+    if n_lines >= 3:
+        size = 5.6
+    elif n_lines == 2 or longest >= 22:
+        size = 6.2
+    else:
+        size = 6.8
+    return FontProperties(
+        family="Liberation Serif", style="italic", weight="normal", size=size
+    )
+
+
+def plot_traversal_panel(
+    image_paths: list[str],
+    hyp_paths: list[list[str]],
+    euc_paths: list[list[str]],
+    out_path: str,
+    left_name: str = "HMClip",
+    right_name: str = "RadIR",
+    keep: int = 6,
+    keep_right: int | None = None,
+    dpi: int = 180,
+) -> None:
+    n = len(image_paths)
+    fig_w = max(4.0 * n, 9.0)
+    keep_euc = keep if keep_right is None else keep_right
+    aligned = []
+    max_rows = 1
+    for h, e in zip(hyp_paths, euc_paths):
+        a, b = align_two_paths(
+            compress_traversal(h, keep=keep),
+            compress_traversal(e, keep=keep_euc),
+        )
+        aligned.append((a, b))
+        max_rows = max(max_rows, len(a))
+
+    row_h_in = 0.46  # taller cells so multi-line text has vertical room
+    header_h_in = 0.32
+    img_h_in = 1.85
+    fig_h = img_h_in + header_h_in + row_h_in * max_rows + 0.5
+    fig, axes = plt.subplots(1, n, figsize=(fig_w, fig_h))
+    if n == 1:
+        axes = [axes]
+
+    from mpl_toolkits.axes_grid1.inset_locator import inset_axes
+
+    for ax, ip, (left_labs, right_labs) in zip(axes, image_paths, aligned):
+        ax.set_axis_off()
+        ax.set_xlim(0, 1)
+        ax.set_ylim(0, 1)
+
+        img = Image.open(ip).convert("L")
+        img_frac = img_h_in / fig_h
+        img_ax = inset_axes(
+            ax,
+            width="92%",
+            height=f"{img_frac * 100 * 0.92:.1f}%",
+            loc="upper center",
+            borderpad=0.25,
+        )
+        img_ax.imshow(img, cmap="gray")
+        img_ax.set_xticks([])
+        img_ax.set_yticks([])
+        for spine in img_ax.spines.values():
+            spine.set_visible(True)
+            spine.set_linewidth(0.6)
+            spine.set_color("0.4")
+
+        n_rows = len(left_labs)
+        header_h = header_h_in / fig_h
+        row_h = row_h_in / fig_h
+        table_top = 1.0 - img_frac - 0.02
+        col_w = 0.47
+        left_x0, right_x0 = 0.025, 0.505
+
+        for x0, name in ((left_x0, left_name), (right_x0, right_name)):
+            ax.add_patch(
+                FancyBboxPatch(
+                    (x0, table_top - header_h),
+                    col_w,
+                    header_h,
+                    boxstyle="square,pad=0",
+                    facecolor="#A8B4C0",
+                    edgecolor="#3D4650",
+                    linewidth=0.8,
+                    transform=ax.transAxes,
+                    clip_on=False,
+                )
+            )
+            ax.text(
+                x0 + col_w / 2,
+                table_top - header_h / 2,
+                name,
+                ha="center",
+                va="center",
+                fontproperties=_HEADER_FONT,
+                color="#1A1A1A",
+                transform=ax.transAxes,
+            )
+
+        for r, (ll, rl) in enumerate(zip(left_labs, right_labs)):
+            y = table_top - header_h - (r + 1) * row_h
+            color = _cell_color(r, n_rows)
+            for x0, lab in ((left_x0, ll), (right_x0, rl)):
+                ax.add_patch(
+                    FancyBboxPatch(
+                        (x0, y),
+                        col_w,
+                        row_h,
+                        boxstyle="square,pad=0",
+                        facecolor=color,
+                        edgecolor="#6B7682",
+                        linewidth=0.7,
+                        transform=ax.transAxes,
+                        clip_on=False,
+                    )
+                )
+                # Use thin ↓ (not heavy ▼); wrap wide to use horizontal space
+                if lab == "↓":
+                    wrapped = "↓"
+                else:
+                    wrapped = _wrap_label(lab, width=24)
+                fp = _label_font(lab, wrapped)
+                # More vertical inset, less horizontal (use cell width)
+                inset_x = 0.012
+                inset_y = 0.035 * (row_h / (row_h_in / fig_h) if row_h > 0 else 1)
+                inset_y = max(0.018, min(0.028, row_h * 0.12))
+                txt = ax.text(
+                    x0 + col_w / 2,
+                    y + row_h / 2,
+                    wrapped,
+                    ha="center",
+                    va="center",
+                    fontproperties=fp,
+                    color="#1F2933",
+                    transform=ax.transAxes,
+                    linespacing=0.98,
+                    zorder=3,
+                )
+                clip = FancyBboxPatch(
+                    (x0 + inset_x, y + inset_y),
+                    col_w - 2 * inset_x,
+                    row_h - 2 * inset_y,
+                    boxstyle="square,pad=0",
+                    transform=ax.transAxes,
+                    facecolor="none",
+                    edgecolor="none",
+                )
+                clip.set_visible(False)
+                ax.add_patch(clip)
+                txt.set_clip_path(clip)
+
+    fig.suptitle(
+        "Image traversals with HMClip (hyperbolic) and RadIR (Euclidean)",
+        fontproperties=_TITLE_FONT,
+        y=0.995,
+    )
+    os.makedirs(os.path.dirname(os.path.abspath(out_path)) or ".", exist_ok=True)
+    fig.savefig(out_path, dpi=dpi, bbox_inches="tight", facecolor="white")
+    plt.close(fig)
+    print(f"Saved traversal panel to {out_path}")
+
+
+def save_json_summary(
+    out_json: str,
+    image_paths: list[str],
+    hyp_raw: list[list[str]],
+    euc_raw: list[list[str]],
+    keep: int = 6,
+    keep_right: int | None = None,
+) -> None:
+    keep_euc = keep if keep_right is None else keep_right
+    payload = []
+    for ip, h, e in zip(image_paths, hyp_raw, euc_raw):
+        hc = compress_traversal(h, keep=keep)
+        ec = compress_traversal(e, keep=keep_euc)
+        la, ra = align_two_paths(hc, ec)
+        payload.append(
+            {
+                "image": ip,
+                "radir_hyp": la,
+                "radir_compare": ra,
+                "hyp_unique_all": compress_traversal(h, keep=30),
+                "compare_unique_all": compress_traversal(e, keep=30),
+            }
+        )
+    os.makedirs(os.path.dirname(os.path.abspath(out_json)) or ".", exist_ok=True)
+    with open(out_json, "w", encoding="utf-8") as f:
+        json.dump(payload, f, indent=2, ensure_ascii=False)
+    print(f"Saved JSON summary to {out_json}")
+
+
+# ---------------------------------------------------------------------------
+# CLI
+# ---------------------------------------------------------------------------
+
+def parse_args():
+    p = argparse.ArgumentParser(description="MERU-style HMClip vs RadIR_compare traversals")
+    p.add_argument("--images", nargs="*", default=None, help="Query image paths")
+    p.add_argument("--auto_pick", type=int, default=0, help="Auto-pick N diverse CXRs")
+    p.add_argument("--device", default="cuda:0")
+    p.add_argument("--hyp_ckpt", default=str(ROOT / "outputs" / "radir_best.pt"))
+    p.add_argument(
+        "--euc_ckpt",
+        default=str(
+            COMPARE_ROOT / "outputs" / "logs" / "CXRCT_UNI_stage1" / "Rad_IR.12000.pt"
+        ),
+    )
+    p.add_argument(
+        "--hyp_bert",
+        default=str(
+            ROOT / "hf_models" / "microsoft" / "BiomedVLP-CXR-BERT-specialized"
+        ),
+    )
+    p.add_argument(
+        "--euc_bert",
+        default=str(
+            COMPARE_ROOT / "hf_models" / "microsoft" / "BiomedVLP-CXR-BERT-specialized"
+        ),
+    )
+    p.add_argument("--dataset_root", default=str(ROOT / "dataset"))
+    p.add_argument("--jsonl", default=str(ROOT / "dataset" / "valid.jsonl"))
+    p.add_argument(
+        "--labels_csv",
+        default=str(ROOT / "outputs" / "tsne_disease_labels.csv"),
+    )
+    p.add_argument(
+        "--concepts_file",
+        default=str(ROOT / "scripts" / "traversal_concepts.txt"),
+    )
+    p.add_argument(
+        "--extra_captions",
+        type=int,
+        default=0,
+        help="Optionally add N short first-sentences from valid.jsonl (0 = concepts file only)",
+    )
+    p.add_argument("--steps", type=int, default=50)
+    p.add_argument(
+        "--keep",
+        type=int,
+        default=6,
+        help="Max concepts shown per HMClip column including [ROOT] (subsample if longer)",
+    )
+    p.add_argument("--hyp_imsize", type=int, default=224)
+    p.add_argument("--euc_imsize", type=int, default=480)
+    p.add_argument(
+        "--entailment",
+        dest="entailment",
+        action="store_true",
+        default=True,
+        help="Hyperbolic: filter pool by entailment cone before NN (MERU-style, default on)",
+    )
+    p.add_argument(
+        "--no_entailment",
+        dest="entailment",
+        action="store_false",
+        help="Hyperbolic: pure distance NN (no cone filter)",
+    )
+    p.add_argument(
+        "--aperture_thresh",
+        type=float,
+        default=1.2,
+        help="Entailment aperture multiplier (training uses ~0.7; higher = weaker filter)",
+    )
+    p.add_argument("--out", default=str(ROOT / "outputs" / "traversals_compare.png"))
+    p.add_argument(
+        "--out_json",
+        default=str(ROOT / "outputs" / "traversals_compare.json"),
+    )
+    p.add_argument("--dpi", type=int, default=160)
+    p.add_argument("--seed", type=int, default=42)
+    return p.parse_args()
+
+
+def main():
+    args = parse_args()
+    device = args.device
+    if device.startswith("cuda") and not torch.cuda.is_available():
+        print("CUDA not available, falling back to CPU")
+        device = "cpu"
+
+    print(f"COMPARE_ROOT = {COMPARE_ROOT}")
+    if args.images:
+        image_paths = [resolve_image_path(p, args.dataset_root) for p in args.images]
+    else:
+        n = args.auto_pick if args.auto_pick > 0 else 4
+        print(f"Auto-picking {n} diverse CXRs ...")
+        image_paths = auto_pick_images(
+            args.dataset_root,
+            args.labels_csv,
+            args.jsonl,
+            n=n,
+            seed=args.seed,
+        )
+    print("Query images:")
+    for p in image_paths:
+        print(f"  - {p}")
+
+    concepts = load_concept_pool(
+        args.concepts_file,
+        extra_captions=args.extra_captions,
+        caption_jsonl=args.jsonl,
+        seed=args.seed,
+    )
+    print(f"Concept pool size: {len(concepts)} (includes {ROOT_LABEL})")
+
+    # Same keep for both sides (natural unique paths; ↓ only if one is shorter).
+    keep = args.keep
+
+    print("\n=== Hyperbolic RadIR/RadIR ===")
+    hyp_paths, curv = run_hyperbolic(
+        image_paths,
+        concepts,
+        ckpt=args.hyp_ckpt,
+        model_path=args.hyp_bert,
+        device=device,
+        imsize=args.hyp_imsize,
+        steps=args.steps,
+        use_entailment=args.entailment,
+        aperture_thresh=args.aperture_thresh,
+    )
+    print(f"curvature={curv:.6f}")
+    for ip, labs in zip(image_paths, hyp_paths):
+        print(
+            f"  {os.path.basename(ip)}: "
+            f"{' → '.join(compress_traversal(labs, keep=keep))}"
+        )
+
+    print("\n=== Euclidean RadIR_compare ===")
+    euc_paths = run_euclidean(
+        image_paths,
+        concepts,
+        ckpt=args.euc_ckpt,
+        model_path=args.euc_bert,
+        device=device,
+        imsize=args.euc_imsize,
+        steps=args.steps,
+    )
+    for ip, labs in zip(image_paths, euc_paths):
+        print(
+            f"  {os.path.basename(ip)}: "
+            f"{' → '.join(compress_traversal(labs, keep=keep))}"
+        )
+
+    plot_traversal_panel(
+        image_paths,
+        hyp_paths,
+        euc_paths,
+        out_path=args.out,
+        keep=keep,
+        dpi=args.dpi,
+    )
+    save_json_summary(
+        args.out_json,
+        image_paths,
+        hyp_paths,
+        euc_paths,
+        keep=keep,
+    )
+
+
+if __name__ == "__main__":
+    main()

scripts/infer_disease_clusters.py

@@ -0,0 +1,513 @@
+#!/usr/bin/env python3
+"""
+Infer disease clusters from RadIR image embeddings + disease labels.
+
+Idea (matching the user request):
+  Within each disease (and Normal), keep only images that carry that label,
+  then cluster them in the same feature space as plot_tsne.py. Every cluster
+  is disease-homogeneous by construction ("nearby images that are all pneumonia
+  → label them Pneumonia").
+
+Conflicts (an image in multiple disease clusters when using inclusive labels)
+are resolved by exclusive single-label first, else largest cluster, else
+highest within-disease density.
+
+Usage (from RadIR/RadIR):
+  # RadIR hyperbolic
+  python scripts/infer_disease_clusters.py \\
+    --load_feats outputs/tsne_feats.npz \\
+    --labels_csv outputs/tsne_disease_labels.csv \\
+    --out_csv outputs/tsne_cluster_disease.csv \\
+    --out_plot outputs/tsne_cluster_disease.png
+
+  # RadIR_compare Euclidean (same label CSV / jsonl order)
+  python scripts/infer_disease_clusters.py \\
+    --load_feats ../RadIR_compare/outputs/tsne_compare_feats.npz \\
+    --labels_csv outputs/tsne_disease_labels.csv \\
+    --jsonl ../RadIR_compare/dataset/valid.jsonl \\
+    --feat_type euclidean \\
+    --out_csv ../RadIR_compare/outputs/tsne_cluster_disease.csv \\
+    --out_stats ../RadIR_compare/outputs/tsne_cluster_stats.json \\
+    --out_plot ../RadIR_compare/outputs/tsne_cluster_disease.png
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+import os
+import sys
+from collections import Counter, defaultdict
+
+import matplotlib.pyplot as plt
+import numpy as np
+import pandas as pd
+import torch
+from sklearn.cluster import DBSCAN
+from sklearn.manifold import TSNE
+from sklearn.neighbors import NearestNeighbors
+from sklearn.preprocessing import normalize
+
+ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
+RAD_IR_ROOT = os.path.join(ROOT, "Rad_IR")
+for p in (ROOT, RAD_IR_ROOT):
+    if p not in sys.path:
+        sys.path.insert(0, p)
+
+from radir import lorentz as L
+
+DISEASE_NAMES = [
+    "Pneumothorax",
+    "Edema",
+    "Pneumonia",
+    "Effusion",
+    "Cardiomegaly",
+    "Atelectasis",
+]
+CLASS_ORDER = ["Normal"] + DISEASE_NAMES + ["Unassigned"]
+
+CLASS_COLORS = {
+    "Normal": "#00C853",
+    "Pneumothorax": "#AA00FF",
+    "Edema": "#00B8D4",
+    "Pneumonia": "#D50000",
+    "Effusion": "#2962FF",
+    "Cardiomegaly": "#FF6D00",
+    "Atelectasis": "#FFD600",
+    "Unassigned": "#BDBDBD",
+}
+
+
+def dicom_from_img_path(img_path: str) -> str:
+    return os.path.splitext(os.path.basename(img_path.replace("\\", "/")))[0]
+
+
+def load_dicom_ids_from_jsonl(jsonl_path: str) -> list[str]:
+    ids = []
+    with open(jsonl_path, encoding="utf-8") as f:
+        for line in f:
+            item = json.loads(line.strip())
+            ids.append(dicom_from_img_path(item["img_path"]))
+    return ids
+
+
+def prepare_image_feats(data, feat_type: str) -> np.ndarray:
+    """
+    Support both:
+      - RadIR/RadIR tsne_feats.npz: eu_img / hyp_img / curv
+      - RadIR_compare tsne_compare_feats.npz: img_feats (already L2 Euclidean)
+    """
+    files = set(data.files)
+
+    # RadIR_compare cached format
+    if "img_feats" in files:
+        if feat_type == "hyperbolic":
+            print(
+                "WARNING: compare feats are Euclidean only; "
+                "ignoring --feat_type hyperbolic"
+            )
+        return np.asarray(data["img_feats"])
+
+    # RadIR/RadIR cached format
+    eu_img = torch.from_numpy(data["eu_img"])
+    hyp_img = torch.from_numpy(data["hyp_img"])
+    curv = torch.tensor(float(data["curv"]))
+    if feat_type == "hyperbolic":
+        with torch.no_grad():
+            return L.log_map0(hyp_img.float(), curv.detach().float()).cpu().numpy()
+    if feat_type == "euclidean":
+        return eu_img.detach().cpu().numpy()
+    raise ValueError(feat_type)
+
+
+def load_label_matrix(labels_csv: str, dicom_ids: list[str]):
+    df = pd.read_csv(labels_csv)
+    df = df.drop_duplicates(subset=["dicom_id"], keep="first").set_index("dicom_id")
+    n = len(dicom_ids)
+    normal_only = np.zeros(n, dtype=bool)
+    disease_mat = np.zeros((n, len(DISEASE_NAMES)), dtype=bool)
+    missing = 0
+    for i, did in enumerate(dicom_ids):
+        if did not in df.index:
+            missing += 1
+            continue
+        row = df.loc[did]
+        normal_only[i] = bool(int(row.get("normal_only", 0)))
+        for j, name in enumerate(DISEASE_NAMES):
+            disease_mat[i, j] = bool(int(row.get(f"{name}_pos", 0)))
+    if missing:
+        print(f"WARNING: {missing} dicom_ids missing from labels CSV")
+    return normal_only, disease_mat
+
+
+def exclusive_labels(normal_only: np.ndarray, disease_mat: np.ndarray) -> np.ndarray:
+    """Single exclusive label or None."""
+    n = len(normal_only)
+    out = np.array([None] * n, dtype=object)
+    for i in range(n):
+        hits = []
+        if normal_only[i]:
+            hits.append("Normal")
+        for j, name in enumerate(DISEASE_NAMES):
+            if disease_mat[i, j]:
+                hits.append(name)
+        if len(hits) == 1:
+            out[i] = hits[0]
+    return out
+
+
+def auto_eps(feats_sub: np.ndarray, min_samples: int, quantile: float) -> float:
+    if feats_sub.shape[0] < min_samples + 1:
+        return 0.1
+    x = normalize(feats_sub.astype(np.float64), norm="l2", axis=1)
+    k = min(min_samples, feats_sub.shape[0] - 1)
+    nn = NearestNeighbors(n_neighbors=k + 1, metric="cosine").fit(x)
+    dists, _ = nn.kneighbors(x)
+    return float(np.quantile(dists[:, -1], quantile))
+
+
+def cluster_subset(
+    feats_all: np.ndarray,
+    member_idx: np.ndarray,
+    eps: float | None,
+    min_samples: int,
+    eps_quantile: float,
+    min_cluster_size: int,
+) -> list[dict]:
+    """
+    DBSCAN on a disease subset. Returns list of
+      {global_indices, size, eps_used}
+    for clusters with size >= min_cluster_size.
+    """
+    if member_idx.size < min_cluster_size:
+        return []
+
+    sub = feats_all[member_idx]
+    eps_used = eps
+    if eps_used is None:
+        eps_used = auto_eps(sub, min_samples, eps_quantile)
+
+    x = normalize(sub.astype(np.float64), norm="l2", axis=1)
+    labels = DBSCAN(
+        eps=eps_used, min_samples=min_samples, metric="cosine"
+    ).fit_predict(x)
+
+    clusters = []
+    for cid in sorted(set(labels.tolist()) - {-1}):
+        local = np.where(labels == cid)[0]
+        if local.size < min_cluster_size:
+            continue
+        global_idx = member_idx[local]
+        clusters.append(
+            {
+                "global_indices": global_idx.astype(int),
+                "size": int(local.size),
+                "eps_used": float(eps_used),
+                "local_cluster_id": int(cid),
+            }
+        )
+    return clusters
+
+
+def infer_disease_clusters(
+    feats: np.ndarray,
+    normal_only: np.ndarray,
+    disease_mat: np.ndarray,
+    seed_mode: str,
+    eps: float | None,
+    min_samples: int,
+    eps_quantile: float,
+    min_cluster_size: int,
+) -> tuple[np.ndarray, np.ndarray, list[dict]]:
+    """
+    Returns:
+      assigned: [N] disease name or 'Unassigned'
+      cluster_uid: [N] int, -1 if unassigned
+      stats: list of cluster summaries
+    """
+    n = feats.shape[0]
+    excl = exclusive_labels(normal_only, disease_mat)
+
+    # candidates: list of (disease, global_indices, size, eps, seed_mode)
+    candidates: list[dict] = []
+    class_names = ["Normal"] + DISEASE_NAMES
+
+    for name in class_names:
+        if seed_mode == "exclusive":
+            member_idx = np.where(excl == name)[0]
+        elif seed_mode == "inclusive":
+            if name == "Normal":
+                member_idx = np.where(normal_only)[0]
+            else:
+                j = DISEASE_NAMES.index(name)
+                member_idx = np.where(disease_mat[:, j])[0]
+        else:
+            raise ValueError(seed_mode)
+
+        print(
+            f"  Seed {name:14s} ({seed_mode}): {member_idx.size} images",
+            end="",
+        )
+        clusters = cluster_subset(
+            feats,
+            member_idx,
+            eps=eps,
+            min_samples=min_samples,
+            eps_quantile=eps_quantile,
+            min_cluster_size=min_cluster_size,
+        )
+        print(f" → {len(clusters)} clusters ≥{min_cluster_size}")
+        for c in clusters:
+            candidates.append(
+                {
+                    "disease": name,
+                    "global_indices": c["global_indices"],
+                    "size": c["size"],
+                    "eps_used": c["eps_used"],
+                    "seed_mode": seed_mode,
+                }
+            )
+
+    # Resolve conflicts: each sample may be claimed by multiple disease clusters.
+    # Score = (is_exclusive_match, cluster_size, -mean_dist_to_centroid)
+    claims: dict[int, list[tuple]] = defaultdict(list)
+    for cid, c in enumerate(candidates):
+        disease = c["disease"]
+        idxs = c["global_indices"]
+        sub = normalize(feats[idxs].astype(np.float64), norm="l2", axis=1)
+        centroid = normalize(sub.mean(axis=0, keepdims=True), norm="l2")[0]
+        # cosine distance to centroid
+        dist = 1.0 - (sub @ centroid)
+        for k, i in enumerate(idxs):
+            is_excl = 1 if excl[i] == disease else 0
+            claims[int(i)].append((is_excl, c["size"], -float(dist[k]), cid, disease))
+
+    assigned = np.array(["Unassigned"] * n, dtype=object)
+    cluster_uid = np.full(n, -1, dtype=int)
+    for i, opts in claims.items():
+        opts.sort(reverse=True)  # prefer exclusive match, then larger cluster
+        _, _, _, cid, disease = opts[0]
+        assigned[i] = disease
+        cluster_uid[i] = cid
+
+    stats = []
+    for cid, c in enumerate(candidates):
+        idxs = c["global_indices"]
+        kept = idxs[assigned[idxs] == c["disease"]]
+        stats.append(
+            {
+                "cluster_uid": cid,
+                "disease": c["disease"],
+                "seed_mode": c["seed_mode"],
+                "seed_size": c["size"],
+                "assigned_size": int(kept.size),
+                "eps_used": c["eps_used"],
+                "exclusive_in_seed": int(
+                    sum(1 for i in idxs if excl[i] == c["disease"])
+                ),
+            }
+        )
+    return assigned, cluster_uid, stats
+
+
+def run_tsne_2d(feats: np.ndarray, perplexity: float, seed: int) -> np.ndarray:
+    n = feats.shape[0]
+    max_perp = max(2.0, (n - 1) / 3.0)
+    eff = min(perplexity, max_perp)
+    return TSNE(
+        n_components=2,
+        perplexity=eff,
+        random_state=seed,
+        init="pca",
+        learning_rate="auto",
+    ).fit_transform(feats)
+
+
+def plot_assigned(
+    emb: np.ndarray,
+    assigned: np.ndarray,
+    out_path: str,
+    title: str,
+    dpi: int,
+) -> None:
+    fig, ax = plt.subplots(figsize=(11, 9))
+    # draw Unassigned first (background)
+    for cls in ["Unassigned"] + [c for c in CLASS_ORDER if c != "Unassigned"]:
+        mask = assigned == cls
+        if not mask.any():
+            continue
+        ax.scatter(
+            emb[mask, 0],
+            emb[mask, 1],
+            c=CLASS_COLORS.get(cls, "#333"),
+            marker="o",
+            alpha=0.35 if cls == "Unassigned" else 0.9,
+            s=18 if cls == "Unassigned" else 40,
+            edgecolors="none" if cls == "Unassigned" else "black",
+            linewidths=0.0 if cls == "Unassigned" else 0.35,
+            label=f"{cls} (n={mask.sum()})",
+            zorder=1 if cls == "Unassigned" else 2,
+        )
+    ax.set_title(title)
+    ax.set_xlabel("t-SNE 1")
+    ax.set_ylabel("t-SNE 2")
+    ax.legend(loc="best", fontsize=8, frameon=True)
+    ax.grid(True, alpha=0.2)
+    fig.tight_layout()
+    os.makedirs(os.path.dirname(os.path.abspath(out_path)) or ".", exist_ok=True)
+    fig.savefig(out_path, dpi=dpi, bbox_inches="tight")
+    plt.close(fig)
+    print(f"Saved plot to {out_path}")
+
+
+def parse_args():
+    p = argparse.ArgumentParser(
+        description="Disease-seeded cluster inference on RadIR image embeddings"
+    )
+    p.add_argument(
+        "--load_feats",
+        default=os.path.join(ROOT, "outputs/tsne_feats.npz"),
+    )
+    p.add_argument(
+        "--labels_csv",
+        default=os.path.join(ROOT, "outputs/tsne_disease_labels.csv"),
+    )
+    p.add_argument(
+        "--jsonl",
+        default=os.path.join(ROOT, "dataset/valid.jsonl"),
+    )
+    p.add_argument(
+        "--feat_type",
+        choices=["hyperbolic", "euclidean"],
+        default="hyperbolic",
+    )
+    p.add_argument(
+        "--seed_mode",
+        choices=["exclusive", "inclusive"],
+        default="exclusive",
+        help="exclusive: only single-label images as seeds (cleaner); "
+        "inclusive: all positives of each disease",
+    )
+    p.add_argument(
+        "--eps",
+        type=float,
+        default=None,
+        help="DBSCAN cosine eps within each disease; default=auto per disease",
+    )
+    p.add_argument("--eps_quantile", type=float, default=0.6)
+    p.add_argument("--min_samples", type=int, default=5)
+    p.add_argument("--min_cluster_size", type=int, default=8)
+    p.add_argument("--perplexity", type=float, default=30.0)
+    p.add_argument("--seed", type=int, default=42)
+    p.add_argument(
+        "--out_csv",
+        default=os.path.join(ROOT, "outputs/tsne_cluster_disease.csv"),
+    )
+    p.add_argument(
+        "--out_stats",
+        default=os.path.join(ROOT, "outputs/tsne_cluster_stats.json"),
+    )
+    p.add_argument(
+        "--out_plot",
+        default=os.path.join(ROOT, "outputs/tsne_cluster_disease.png"),
+    )
+    p.add_argument("--dpi", type=int, default=150)
+    return p.parse_args()
+
+
+def main():
+    args = parse_args()
+
+    print(f"Loading features: {args.load_feats}")
+    data = np.load(args.load_feats, allow_pickle=True)
+    img_feats = prepare_image_feats(data, args.feat_type)
+    n = img_feats.shape[0]
+    print(f"  image feats: {img_feats.shape}, feat_type={args.feat_type}")
+
+    dicom_ids = load_dicom_ids_from_jsonl(args.jsonl)
+    if len(dicom_ids) != n:
+        raise RuntimeError(f"jsonl ({len(dicom_ids)}) != feats ({n})")
+
+    normal_only, disease_mat = load_label_matrix(args.labels_csv, dicom_ids)
+    excl = exclusive_labels(normal_only, disease_mat)
+    print(
+        f"Exclusive single-label: {sum(e is not None for e in excl)} / {n}"
+    )
+    print(f"Seed mode: {args.seed_mode}")
+
+    print("\n===== Within-disease clustering =====")
+    assigned, cluster_uid, stats = infer_disease_clusters(
+        img_feats,
+        normal_only,
+        disease_mat,
+        seed_mode=args.seed_mode,
+        eps=args.eps,
+        min_samples=args.min_samples,
+        eps_quantile=args.eps_quantile,
+        min_cluster_size=args.min_cluster_size,
+    )
+
+    print("\n===== Clusters kept =====")
+    stats_sorted = sorted(stats, key=lambda s: (-s["assigned_size"], s["disease"]))
+    for s in stats_sorted:
+        if s["assigned_size"] <= 0:
+            continue
+        print(
+            f"  [{s['cluster_uid']:>3d}] {s['disease']:14s} "
+            f"assigned={s['assigned_size']:<4d} "
+            f"(seed={s['seed_size']}, excl_in_seed={s['exclusive_in_seed']}, "
+            f"eps={s['eps_used']:.3f})"
+        )
+
+    print("\n===== Assigned counts =====")
+    for cls in CLASS_ORDER:
+        c = int((assigned == cls).sum())
+        if c:
+            print(f"  {cls:14s} {c}")
+
+    os.makedirs(os.path.dirname(os.path.abspath(args.out_csv)) or ".", exist_ok=True)
+    rows = {
+        "dicom_id": dicom_ids,
+        "cluster_uid": cluster_uid,
+        "cluster_disease": assigned,
+        "exclusive_label": [
+            e if e is not None else "" for e in excl.tolist()
+        ],
+        "normal_only": normal_only.astype(int),
+    }
+    for j, name in enumerate(DISEASE_NAMES):
+        rows[f"{name}_pos"] = disease_mat[:, j].astype(int)
+    pd.DataFrame(rows).to_csv(args.out_csv, index=False)
+    print(f"\nSaved {args.out_csv}")
+
+    with open(args.out_stats, "w", encoding="utf-8") as f:
+        json.dump(
+            {
+                "seed_mode": args.seed_mode,
+                "feat_type": args.feat_type,
+                "eps": args.eps,
+                "eps_quantile": args.eps_quantile,
+                "min_samples": args.min_samples,
+                "min_cluster_size": args.min_cluster_size,
+                "n_samples": n,
+                "assigned_counts": {
+                    cls: int((assigned == cls).sum()) for cls in CLASS_ORDER
+                },
+                "clusters": stats_sorted,
+            },
+            f,
+            indent=2,
+        )
+    print(f"Saved {args.out_stats}")
+
+    print("Running t-SNE for visualization ...")
+    emb = run_tsne_2d(img_feats, args.perplexity, args.seed)
+    title = (
+        f"Disease-seeded image clusters "
+        f"({args.seed_mode}, {args.feat_type}, min_size={args.min_cluster_size})"
+    )
+    plot_assigned(emb, assigned, args.out_plot, title, args.dpi)
+
+
+if __name__ == "__main__":
+    main()

scripts/plot_anat_cosne.py

@@ -0,0 +1,677 @@
+#!/usr/bin/env python3
+"""
+Disease-colored CO-SNE for RadIR/RadIR hyperbolic embeddings.
+
+Uses the same hyperbolic CO-SNE pipeline as scripts/plot_cosne.py, then colors
+points by disease while keeping modality markers:
+  ○ Image,  △ Text
+  Color = disease / class
+
+Labeling strategies (--color_by):
+  gt         ground-truth exclusive / priority labels (see --label_mode)
+  inferred   disease-seeded DBSCAN on high-dim hyp image feats ("反推法")
+  cosne2d    same seeded DBSCAN but on 2D CO-SNE image coords
+  cluster_csv load precomputed cluster_disease from infer_disease_clusters.py
+
+Usage (from RadIR/RadIR):
+  # Ground-truth exclusive (clean)
+  python scripts/plot_anat_cosne.py \\
+    --load_emb outputs/cosne_emb.npz \\
+    --color_by gt --label_mode exclusive --hide_other \\
+    --out outputs/cosne_disease_clean.png
+
+  # Reverse-inference clusters on CO-SNE disk
+  python scripts/plot_anat_cosne.py \\
+    --load_emb outputs/cosne_emb.npz \\
+    --load_feats outputs/tsne_feats.npz \\
+    --color_by inferred --hide_other \\
+    --out outputs/cosne_disease_inferred.png
+"""
+
+from __future__ import annotations
+
+import argparse
+import os
+import sys
+from collections import Counter
+
+import matplotlib.pyplot as plt
+import numpy as np
+import pandas as pd
+import torch
+
+ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
+SCRIPTS = os.path.dirname(os.path.abspath(__file__))
+RAD_IR_ROOT = os.path.join(ROOT, "Rad_IR")
+for p in (ROOT, RAD_IR_ROOT, SCRIPTS):
+    if p not in sys.path:
+        sys.path.insert(0, p)
+
+import infer_disease_clusters as idc  # noqa: E402
+import plot_cosne as cosne  # noqa: E402
+
+DISEASE_NAMES = idc.DISEASE_NAMES
+CLASS_ORDER = ["Normal"] + DISEASE_NAMES + ["Other"]
+
+CLASS_COLORS = {
+    "Normal": "#00C853",
+    "Pneumothorax": "#AA00FF",
+    "Edema": "#00B8D4",
+    "Pneumonia": "#D50000",
+    "Effusion": "#2962FF",
+    "Cardiomegaly": "#FF6D00",
+    "Atelectasis": "#FFD600",
+    "Other": "#BDBDBD",
+}
+
+MODALITY_MARKER = {"Image": "o", "Text": "^"}
+
+
+def style_for(cls: str, modality: str) -> tuple[str, str]:
+    return CLASS_COLORS.get(cls, "#333333"), MODALITY_MARKER[modality]
+
+
+def load_label_table(labels_csv: str) -> pd.DataFrame:
+    df = pd.read_csv(labels_csv)
+    needed = ["dicom_id", "normal_only"] + [f"{d}_pos" for d in DISEASE_NAMES]
+    missing = [c for c in needed if c not in df.columns]
+    if missing:
+        raise KeyError(f"labels CSV missing columns: {missing}")
+    return df.drop_duplicates(subset=["dicom_id"], keep="first").reset_index(drop=True)
+
+
+def build_labels(
+    df: pd.DataFrame,
+    sample_idx: np.ndarray,
+    label_mode: str,
+) -> list[str]:
+    """
+    Map global row indices → class name.
+
+    exclusive: single-label only; else "Other"
+    priority:  Normal if normal_only, else first disease hit, else "Other"
+    """
+    labels: list[str] = []
+    for i in sample_idx:
+        row = df.iloc[int(i)]
+        hits: list[str] = []
+        if int(row["normal_only"]) == 1:
+            hits.append("Normal")
+        for d in DISEASE_NAMES:
+            if int(row[f"{d}_pos"]) == 1:
+                hits.append(d)
+
+        if label_mode == "exclusive":
+            labels.append(hits[0] if len(hits) == 1 else "Other")
+        elif label_mode == "priority":
+            if int(row["normal_only"]) == 1 and len(hits) == 1:
+                labels.append("Normal")
+            else:
+                assigned = None
+                for d in DISEASE_NAMES:
+                    if int(row[f"{d}_pos"]) == 1:
+                        assigned = d
+                        break
+                if assigned is None and int(row["normal_only"]) == 1:
+                    assigned = "Normal"
+                labels.append(assigned if assigned is not None else "Other")
+        else:
+            raise ValueError(label_mode)
+    return labels
+
+
+def _to_plot_label(name: str) -> str:
+    if name in (None, "", "Unassigned", "Other"):
+        return "Other"
+    return str(name)
+
+
+def labels_from_cluster_csv(
+    cluster_csv: str, labels_csv: str, sample_idx: np.ndarray
+) -> list[str]:
+    """Map sample_idx → cluster_disease via dicom_id alignment."""
+    gt = pd.read_csv(labels_csv).drop_duplicates(subset=["dicom_id"], keep="first")
+    cl = pd.read_csv(cluster_csv).drop_duplicates(subset=["dicom_id"], keep="first")
+    cl = cl.set_index("dicom_id")["cluster_disease"]
+    out: list[str] = []
+    for i in sample_idx:
+        did = str(gt.iloc[int(i)]["dicom_id"])
+        if did in cl.index:
+            out.append(_to_plot_label(cl.loc[did]))
+        else:
+            out.append("Other")
+    return out
+
+
+def labels_inferred_highdim(
+    feats_npz: str,
+    labels_csv: str,
+    sample_idx: np.ndarray,
+    seed_mode: str,
+    eps: float | None,
+    min_samples: int,
+    eps_quantile: float,
+    min_cluster_size: int,
+) -> tuple[list[str], list[dict]]:
+    """
+    Disease-seeded DBSCAN on full-set high-dim hyp image feats, then slice
+    assignments to the CO-SNE subsample.
+    """
+    data = np.load(feats_npz, allow_pickle=True)
+    feats = idc.prepare_image_feats(data, "hyperbolic")
+    df = pd.read_csv(labels_csv).drop_duplicates(subset=["dicom_id"], keep="first")
+    if len(df) != feats.shape[0]:
+        raise RuntimeError(
+            f"feats N={feats.shape[0]} vs labels rows={len(df)} "
+            "(expect same valid.jsonl order)"
+        )
+    normal_only = df["normal_only"].astype(int).to_numpy().astype(bool)
+    disease_mat = np.stack(
+        [df[f"{d}_pos"].astype(int).to_numpy().astype(bool) for d in DISEASE_NAMES],
+        axis=1,
+    )
+    print(f"===== Infer clusters on high-dim hyp feats ({seed_mode}) =====")
+    assigned, _uid, stats = idc.infer_disease_clusters(
+        feats,
+        normal_only,
+        disease_mat,
+        seed_mode=seed_mode,
+        eps=eps,
+        min_samples=min_samples,
+        eps_quantile=eps_quantile,
+        min_cluster_size=min_cluster_size,
+    )
+    labels = [_to_plot_label(assigned[int(i)]) for i in sample_idx]
+    return labels, stats
+
+
+def labels_inferred_cosne2d(
+    emb_img: np.ndarray,
+    df: pd.DataFrame,
+    sample_idx: np.ndarray,
+    seed_mode: str,
+    eps: float | None,
+    min_samples: int,
+    eps_quantile: float,
+    min_cluster_size: int,
+) -> tuple[list[str], list[dict]]:
+    """
+    Same disease-seeded DBSCAN, but distance is Euclidean in the 2D CO-SNE
+    Poincaré plane (answers: does the disk layout form disease neighborhoods?).
+    """
+    n = emb_img.shape[0]
+    normal_only = np.zeros(n, dtype=bool)
+    disease_mat = np.zeros((n, len(DISEASE_NAMES)), dtype=bool)
+    for k, i in enumerate(sample_idx):
+        row = df.iloc[int(i)]
+        normal_only[k] = bool(int(row["normal_only"]))
+        for j, d in enumerate(DISEASE_NAMES):
+            disease_mat[k, j] = bool(int(row[f"{d}_pos"]))
+
+    # Use Euclidean metric on 2D coords via L2-normalize trick disabled:
+    # cluster_subset uses cosine on L2-normalized vectors; for 2D Poincaré
+    # points we pass raw coords and temporarily use a Euclidean DBSCAN path.
+    print(f"===== Infer clusters on CO-SNE 2D ({seed_mode}) =====")
+    assigned, _uid, stats = infer_clusters_euclidean_2d(
+        emb_img,
+        normal_only,
+        disease_mat,
+        seed_mode=seed_mode,
+        eps=eps,
+        min_samples=min_samples,
+        eps_quantile=eps_quantile,
+        min_cluster_size=min_cluster_size,
+    )
+    labels = [_to_plot_label(a) for a in assigned]
+    return labels, stats
+
+
+def infer_clusters_euclidean_2d(
+    coords: np.ndarray,
+    normal_only: np.ndarray,
+    disease_mat: np.ndarray,
+    seed_mode: str,
+    eps: float | None,
+    min_samples: int,
+    eps_quantile: float,
+    min_cluster_size: int,
+) -> tuple[np.ndarray, np.ndarray, list[dict]]:
+    """DBSCAN with Euclidean metric on 2D coords (within disease seeds)."""
+    from collections import defaultdict
+
+    from sklearn.cluster import DBSCAN
+    from sklearn.neighbors import NearestNeighbors
+
+    n = coords.shape[0]
+    excl = idc.exclusive_labels(normal_only, disease_mat)
+    candidates: list[dict] = []
+    class_names = ["Normal"] + DISEASE_NAMES
+
+    for name in class_names:
+        if seed_mode == "exclusive":
+            member_idx = np.where(excl == name)[0]
+        else:
+            if name == "Normal":
+                member_idx = np.where(normal_only)[0]
+            else:
+                j = DISEASE_NAMES.index(name)
+                member_idx = np.where(disease_mat[:, j])[0]
+
+        print(f"  Seed {name:14s} ({seed_mode}): {member_idx.size} images", end="")
+        if member_idx.size < min_cluster_size:
+            print(" → 0 clusters")
+            continue
+
+        sub = coords[member_idx].astype(np.float64)
+        eps_used = eps
+        if eps_used is None:
+            k = min(min_samples, sub.shape[0] - 1)
+            if k < 1:
+                print(" → 0 clusters")
+                continue
+            nn = NearestNeighbors(n_neighbors=k + 1, metric="euclidean").fit(sub)
+            dists, _ = nn.kneighbors(sub)
+            eps_used = float(np.quantile(dists[:, -1], eps_quantile))
+
+        labels_db = DBSCAN(
+            eps=eps_used, min_samples=min_samples, metric="euclidean"
+        ).fit_predict(sub)
+        n_kept = 0
+        for cid in sorted(set(labels_db.tolist()) - {-1}):
+            local = np.where(labels_db == cid)[0]
+            if local.size < min_cluster_size:
+                continue
+            global_idx = member_idx[local]
+            candidates.append(
+                {
+                    "disease": name,
+                    "global_indices": global_idx.astype(int),
+                    "size": int(local.size),
+                    "eps_used": float(eps_used),
+                }
+            )
+            n_kept += 1
+        print(f" → {n_kept} clusters ≥{min_cluster_size} (eps={eps_used:.4f})")
+
+    claims: dict[int, list[tuple]] = defaultdict(list)
+    for cid, c in enumerate(candidates):
+        idxs = c["global_indices"]
+        sub = coords[idxs]
+        centroid = sub.mean(axis=0)
+        dist = np.linalg.norm(sub - centroid, axis=1)
+        for k, i in enumerate(idxs):
+            is_excl = 1 if excl[i] == c["disease"] else 0
+            claims[int(i)].append(
+                (is_excl, c["size"], -float(dist[k]), cid, c["disease"])
+            )
+
+    assigned = np.array(["Unassigned"] * n, dtype=object)
+    cluster_uid = np.full(n, -1, dtype=int)
+    for i, opts in claims.items():
+        opts.sort(reverse=True)
+        _, _, _, cid, disease = opts[0]
+        assigned[i] = disease
+        cluster_uid[i] = cid
+
+    stats = []
+    for cid, c in enumerate(candidates):
+        idxs = c["global_indices"]
+        kept = idxs[assigned[idxs] == c["disease"]]
+        stats.append(
+            {
+                "cluster_uid": cid,
+                "disease": c["disease"],
+                "size_raw": c["size"],
+                "size_kept": int(kept.size),
+                "eps_used": c["eps_used"],
+            }
+        )
+    return assigned, cluster_uid, stats
+
+
+def print_cluster_stats(stats: list[dict], labels: list[str]) -> None:
+    print("----- kept clusters -----")
+    by_d: dict[str, int] = Counter()
+    for s in stats:
+        if s.get("size_kept", 0) > 0:
+            print(
+                f"  uid={s['cluster_uid']:3d}  {s['disease']:14s}  "
+                f"kept={s['size_kept']:4d}  raw={s['size_raw']:4d}  "
+                f"eps={s['eps_used']:.4f}"
+            )
+            by_d[s["disease"]] += s["size_kept"]
+    print("----- subsample assignment -----")
+    counts = Counter(labels)
+    for cls in CLASS_ORDER:
+        if counts.get(cls, 0):
+            print(f"  {cls:14s} {counts[cls]}")
+    labeled = sum(v for k, v in counts.items() if k != "Other")
+    print(f"  labeled={labeled}/{len(labels)}  Other={counts.get('Other', 0)}")
+
+
+def run_or_load_cosne(args) -> tuple[np.ndarray, np.ndarray, np.ndarray, float]:
+    """
+    Returns emb_img [M,2], emb_txt [M,2], sample_idx [M], curv.
+    """
+    if args.load_emb:
+        print(f"Loading CO-SNE coords: {args.load_emb}")
+        data = np.load(args.load_emb, allow_pickle=True)
+        emb_img = np.asarray(data["emb_img"])
+        emb_txt = np.asarray(data["emb_txt"])
+        sample_idx = np.asarray(data["sample_idx"], dtype=np.int64)
+        curv = float(data["curv"]) if "curv" in data.files else 1.0
+        if emb_img.shape[0] != emb_txt.shape[0] or emb_img.shape[0] != len(sample_idx):
+            raise ValueError("emb_img / emb_txt / sample_idx length mismatch")
+        return emb_img, emb_txt, sample_idx, curv
+
+    if not args.load_feats:
+        raise ValueError("Provide --load_emb or --load_feats")
+
+    print(f"Loading hyperbolic feats: {args.load_feats}")
+    feats = np.load(args.load_feats, allow_pickle=True)
+    if "hyp_img" not in feats.files:
+        raise KeyError("Need hyp_img / hyp_txt in feats npz")
+    hyp_img = torch.from_numpy(feats["hyp_img"])
+    hyp_txt = torch.from_numpy(feats["hyp_txt"])
+    curv = float(feats["curv"]) if "curv" in feats.files else 1.0
+
+    max_samples = None if args.max_samples <= 0 else args.max_samples
+    hyp_img, hyp_txt, sample_idx = cosne.subsample_pair(
+        hyp_img, hyp_txt, max_samples, args.seed
+    )
+    print(f"  subsampled to {len(sample_idx)} pairs, curv={curv:.4f}")
+
+    with torch.no_grad():
+        poi_img = cosne.lorentz_space_to_poincare(hyp_img.float(), curv).numpy()
+        poi_txt = cosne.lorentz_space_to_poincare(hyp_txt.float(), curv).numpy()
+
+    X = np.concatenate([poi_img, poi_txt], axis=0)
+    model = cosne.COSNE(
+        n_components=2,
+        perplexity=args.perplexity,
+        n_iter=args.n_iter,
+        learning_rate=args.learning_rate,
+        learning_rate_h=args.learning_rate_h,
+        early_exaggeration=args.early_exaggeration,
+        student_t_gamma=args.student_t_gamma,
+        seed=args.seed,
+        device=args.cosne_device,
+        verbose=True,
+    )
+    Y = model.fit_transform(X)
+    n = poi_img.shape[0]
+    emb_img, emb_txt = Y[:n], Y[n:]
+
+    if args.save_emb:
+        os.makedirs(
+            os.path.dirname(os.path.abspath(args.save_emb)) or ".", exist_ok=True
+        )
+        np.savez(
+            args.save_emb,
+            emb_img=emb_img,
+            emb_txt=emb_txt,
+            sample_idx=sample_idx,
+            curv=curv,
+        )
+        print(f"Saved CO-SNE coords to {args.save_emb}")
+
+    return emb_img, emb_txt, sample_idx, curv
+
+
+def plot_disease_poincare(
+    emb_img: np.ndarray,
+    emb_txt: np.ndarray,
+    labels: list[str],
+    out_path: str,
+    title: str,
+    dpi: int,
+    hide_other: bool,
+    max_pair_lines: int,
+) -> None:
+    n = emb_img.shape[0]
+    assert len(labels) == n
+
+    fig, ax = plt.subplots(figsize=(9, 9))
+    ax.add_patch(plt.Circle((0, 0), 1.0, color="black", fill=False, linewidth=1.2))
+    for r in (0.25, 0.5, 0.75):
+        ax.add_patch(
+            plt.Circle((0, 0), r, color="gray", fill=False, linewidth=0.4, alpha=0.35)
+        )
+
+    # optional pair connectors (only for non-Other if hide_other)
+    n_lines = min(max_pair_lines, n)
+    drawn = 0
+    for i in range(n):
+        if drawn >= n_lines:
+            break
+        if hide_other and labels[i] == "Other":
+            continue
+        ax.plot(
+            [emb_img[i, 0], emb_txt[i, 0]],
+            [emb_img[i, 1], emb_txt[i, 1]],
+            color="gray",
+            alpha=0.12,
+            linewidth=0.5,
+            zorder=0,
+        )
+        drawn += 1
+
+    # draw Other first (background), then each disease
+    draw_order = ["Other"] + [c for c in CLASS_ORDER if c != "Other"]
+    for cls in draw_order:
+        if hide_other and cls == "Other":
+            continue
+        mask = np.array([l == cls for l in labels], dtype=bool)
+        if not mask.any():
+            continue
+        for modality, emb in (("Image", emb_img), ("Text", emb_txt)):
+            color, marker = style_for(cls, modality)
+            is_other = cls == "Other"
+            ax.scatter(
+                emb[mask, 0],
+                emb[mask, 1],
+                c=color,
+                marker=marker,
+                s=18 if is_other else (32 if modality == "Image" else 42),
+                alpha=0.35 if is_other else 0.88,
+                edgecolors="none" if is_other else "black",
+                linewidths=0.0 if is_other else 0.3,
+                label=f"{cls}-{modality} (n={int(mask.sum())})",
+                zorder=1 if is_other else (2 if modality == "Image" else 3),
+            )
+
+    # radius summary on labeled (non-Other) points
+    labeled_mask = np.array([l != "Other" for l in labels], dtype=bool)
+    if labeled_mask.any():
+        r_img = np.linalg.norm(emb_img[labeled_mask], axis=1).mean()
+        r_txt = np.linalg.norm(emb_txt[labeled_mask], axis=1).mean()
+        subtitle = f"labeled mean radius  Image={r_img:.3f}  Text={r_txt:.3f}"
+    else:
+        subtitle = "no exclusive labels in subsample"
+
+    ax.set_title(f"{title}\n{subtitle}")
+    ax.set_xlim(-1.05, 1.05)
+    ax.set_ylim(-1.05, 1.05)
+    ax.set_aspect("equal")
+    ax.legend(loc="upper left", bbox_to_anchor=(1.02, 1.0), fontsize=8, frameon=True)
+    ax.axis("off")
+    fig.tight_layout()
+    os.makedirs(os.path.dirname(os.path.abspath(out_path)) or ".", exist_ok=True)
+    fig.savefig(out_path, dpi=dpi, bbox_inches="tight")
+    if out_path.lower().endswith(".png"):
+        fig.savefig(out_path[:-4] + ".pdf", dpi=dpi, bbox_inches="tight")
+    plt.close(fig)
+    print(f"Saved disease CO-SNE to {out_path}")
+
+
+def parse_args():
+    p = argparse.ArgumentParser(
+        description="Disease-colored CO-SNE (hyperbolic) for RadIR"
+    )
+    p.add_argument(
+        "--load_emb",
+        default=None,
+        help="Cached CO-SNE coords from plot_cosne.py (--save_emb)",
+    )
+    p.add_argument(
+        "--load_feats",
+        default=None,
+        help="Hyperbolic feats npz (required for --color_by inferred)",
+    )
+    p.add_argument(
+        "--labels_csv",
+        default=os.path.join(ROOT, "outputs/tsne_disease_labels.csv"),
+    )
+    p.add_argument(
+        "--cluster_csv",
+        default=os.path.join(ROOT, "outputs/tsne_cluster_disease.csv"),
+        help="Precomputed cluster_disease CSV for --color_by cluster_csv",
+    )
+    p.add_argument(
+        "--jsonl",
+        default=os.path.join(ROOT, "dataset/valid.jsonl"),
+        help="Only used to sanity-check length vs labels (optional)",
+    )
+    p.add_argument(
+        "--color_by",
+        choices=["gt", "inferred", "cosne2d", "cluster_csv"],
+        default="gt",
+        help="gt=ground-truth; inferred=反推 on high-dim; "
+        "cosne2d=反推 on CO-SNE plane; cluster_csv=load CSV",
+    )
+    p.add_argument(
+        "--label_mode",
+        choices=["exclusive", "priority"],
+        default="exclusive",
+        help="For --color_by gt only",
+    )
+    p.add_argument(
+        "--seed_mode",
+        choices=["exclusive", "inclusive"],
+        default="exclusive",
+        help="Disease seed mode for inferred / cosne2d",
+    )
+    p.add_argument("--eps", type=float, default=None, help="DBSCAN eps (auto if None)")
+    p.add_argument("--min_samples", type=int, default=5)
+    p.add_argument("--eps_quantile", type=float, default=0.5)
+    p.add_argument("--min_cluster_size", type=int, default=8)
+    p.add_argument(
+        "--hide_other",
+        action="store_true",
+        help="Do not draw Other/Unassigned points",
+    )
+    p.add_argument("--max_samples", type=int, default=500)
+    p.add_argument("--perplexity", type=float, default=30.0)
+    p.add_argument("--n_iter", type=int, default=800)
+    p.add_argument("--learning_rate", type=float, default=0.05)
+    p.add_argument("--learning_rate_h", type=float, default=5.0)
+    p.add_argument("--early_exaggeration", type=float, default=4.0)
+    p.add_argument("--student_t_gamma", type=float, default=0.1)
+    p.add_argument("--seed", type=int, default=42)
+    p.add_argument("--cosne_device", default="cpu")
+    p.add_argument("--max_pair_lines", type=int, default=40)
+    p.add_argument(
+        "--out",
+        default=os.path.join(ROOT, "outputs/cosne_disease.png"),
+    )
+    p.add_argument("--save_emb", default=None)
+    p.add_argument("--dpi", type=int, default=150)
+    return p.parse_args()
+
+
+def main():
+    args = parse_args()
+    if args.load_emb is None and args.load_feats is None:
+        default_emb = os.path.join(ROOT, "outputs/cosne_emb.npz")
+        default_feats = os.path.join(ROOT, "outputs/tsne_feats.npz")
+        if os.path.isfile(default_emb):
+            args.load_emb = default_emb
+        elif os.path.isfile(default_feats):
+            args.load_feats = default_feats
+        else:
+            raise FileNotFoundError(
+                "Need --load_emb or --load_feats (or outputs/cosne_emb.npz)"
+            )
+
+    emb_img, emb_txt, sample_idx, curv = run_or_load_cosne(args)
+    df = load_label_table(args.labels_csv)
+    if len(df) < int(sample_idx.max()) + 1:
+        raise RuntimeError(
+            f"labels CSV has {len(df)} rows but sample_idx max={sample_idx.max()}"
+        )
+
+    stats: list[dict] = []
+    if args.color_by == "gt":
+        labels = build_labels(df, sample_idx, args.label_mode)
+        print(f"===== Label mode: gt/{args.label_mode} (n={len(labels)}) =====")
+        for cls in CLASS_ORDER:
+            c = Counter(labels).get(cls, 0)
+            if c:
+                print(f"  {cls:14s} {c}")
+        title = (
+            f"CO-SNE + GT ({args.label_mode}) | hyperbolic | "
+            f"n={len(labels)} | curv={curv:.3f}"
+        )
+    elif args.color_by == "cluster_csv":
+        labels = labels_from_cluster_csv(args.cluster_csv, args.labels_csv, sample_idx)
+        print(f"===== Color by cluster CSV: {args.cluster_csv} =====")
+        for cls in CLASS_ORDER:
+            c = Counter(labels).get(cls, 0)
+            if c:
+                print(f"  {cls:14s} {c}")
+        title = (
+            f"CO-SNE + inferred clusters (CSV) | hyperbolic | "
+            f"n={len(labels)} | curv={curv:.3f}"
+        )
+    elif args.color_by == "inferred":
+        feats_path = args.load_feats or os.path.join(ROOT, "outputs/tsne_feats.npz")
+        if not os.path.isfile(feats_path):
+            raise FileNotFoundError(f"--color_by inferred needs feats: {feats_path}")
+        labels, stats = labels_inferred_highdim(
+            feats_path,
+            args.labels_csv,
+            sample_idx,
+            seed_mode=args.seed_mode,
+            eps=args.eps,
+            min_samples=args.min_samples,
+            eps_quantile=args.eps_quantile,
+            min_cluster_size=args.min_cluster_size,
+        )
+        print_cluster_stats(stats, labels)
+        title = (
+            f"CO-SNE + inferred (high-dim {args.seed_mode}) | hyperbolic | "
+            f"n={len(labels)} | curv={curv:.3f}"
+        )
+    elif args.color_by == "cosne2d":
+        labels, stats = labels_inferred_cosne2d(
+            emb_img,
+            df,
+            sample_idx,
+            seed_mode=args.seed_mode,
+            eps=args.eps,
+            min_samples=args.min_samples,
+            eps_quantile=args.eps_quantile,
+            min_cluster_size=args.min_cluster_size,
+        )
+        print_cluster_stats(stats, labels)
+        title = (
+            f"CO-SNE + inferred (2D disk {args.seed_mode}) | hyperbolic | "
+            f"n={len(labels)} | curv={curv:.3f}"
+        )
+    else:
+        raise ValueError(args.color_by)
+
+    plot_disease_poincare(
+        emb_img,
+        emb_txt,
+        labels,
+        args.out,
+        title=title,
+        dpi=args.dpi,
+        hide_other=args.hide_other,
+        max_pair_lines=args.max_pair_lines,
+    )
+
+
+if __name__ == "__main__":
+    main()