Start resumable backup upload: root source files

.gitattributes

@@ -33,3 +33,9 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+hf_models/** filter=lfs diff=lfs merge=lfs -text
+outputs/** filter=lfs diff=lfs merge=lfs -text
+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

CXR_CT_inference.py

@@ -0,0 +1,117 @@
+from transformer_maskgit import CTViT
+from transformers import BertTokenizer, BertModel
+from radir import RADIR
+from data_process import load_2d_image_to_tensor, load_3d_image_to_tensor
+import torch
+
+MODALITY_DICT = {'CT': 0, 'CXR': 1}
+
+if __name__ == '__main__':
+    # Determine best available device
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+    # Load tokenizer and text encoder
+    tokenizer = BertTokenizer.from_pretrained('microsoft/BiomedVLP-CXR-BERT-specialized', do_lower_case=True)
+    text_encoder = BertModel.from_pretrained("microsoft/BiomedVLP-CXR-BERT-specialized").to(device)
+    
+    # Initialize image encoder
+    image_encoder = CTViT(
+        dim=512,
+        codebook_size=8192,
+        image_size=480,
+        patch_size=20,
+        temporal_patch_size=10,
+        spatial_depth=8,
+        temporal_depth=6,
+        cls_depth=4,
+        dim_head=32,
+        heads=8
+    ).to(device)
+    
+    # Initialize RadIR model
+    rad_ir = RADIR(
+        image_encoder=image_encoder,
+        text_encoder=text_encoder,
+        tokenizer=tokenizer,
+        dim_text=768,
+        dim_image=512,
+        dim_latent=512,
+        extra_latent_projection=False,
+        use_mlm=False,
+        downsample_image_embeds=False,
+        use_all_token_embeds=False
+    ).to(device)
+
+    # Load checkpoint
+    checkpoint_path = 'checkpoints/RadIR.pt'
+    rad_ir.load(checkpoint_path)
+    rad_ir.eval()  # Set to evaluation mode
+    
+    # Test examples
+    cxr_images = ['example/1.jpg', 'example/2.jpg']
+    cxr_reports = [
+        "Lateral view somewhat limited due to overlying motion artifact. The lungs are low in volume. There is no focal airspace consolidation to suggest pneumonia. A 1.2-cm calcified granuloma just below the medial aspect of the right hemidiaphragm is unchanged from prior study. No pleural effusions or pulmonary edema. There is no pneumothorax. The inferior sternotomy wire is fractured but unchanged. Surgical clips and vascular markers in the thorax are related to prior CABG surgery. No evidence of acute cardiopulmonary process.",
+        "Left PICC tip is seen terminating in the region of the distal left brachiocephalic vein. Tracheostomy tube is in unchanged standard position. The heart is moderately enlarged. Marked calcification of the aortic knob is again present. Mild pulmonary vascular congestion is similar. Bibasilar streaky airspace opacities are minimally improved. Previously noted left pleural effusion appears to have resolved. No pneumothorax is identified. Percutaneous gastrostomy tube is seen in the left upper quadrant. 1. Left PICC tip appears to terminate in the distal left brachiocephalic vein. 2. Mild pulmonary vascular congestion. 3. Interval improvement in aeration of the lung bases with residual streaky opacity likely reflective of atelectasis. Interval resolution of the left pleural effusion."
+    ]
+    anatomy_condition = ['right lung']
+    
+    # Extract unconditional image-text embeddings
+    with torch.no_grad():
+        # Process images
+        image_tensors = [load_2d_image_to_tensor(image_path) for image_path in cxr_images]
+        batched_images = torch.stack(image_tensors, dim=0).to(device)  # Shape: [B, C, D, H, W] -> [2, 1, 1, 480, 480]
+        
+        # Process text
+        text_tokens = tokenizer(
+            cxr_reports,
+            return_tensors="pt", 
+            padding="max_length", 
+            truncation=True, 
+            max_length=512
+        ).to(device)
+        print(text_tokens)
+        
+        # Set modality indicator for CXR
+        modal_indices = torch.tensor([MODALITY_DICT['CXR']] * len(cxr_images)).to(device)
+        
+        # Forward pass to get embeddings
+        image_embeddings, text_embeddings, _, _ = rad_ir(
+            text_tokens, 
+            image=batched_images,
+            device=device,
+            is_condition=False,  # Indicate unconditional scenario
+            return_latents=True,
+            modal_indexs=modal_indices,
+            modal_embedding=True
+        )
+        
+        print(f"Image embedding shape: {image_embeddings.shape}")  # [B, 512]
+        print(f"Text embedding shape: {text_embeddings.shape}")    # [B, 512]
+        
+        # Extract condition-specific features (focused on right lung)
+        con_batch_image = batched_images.unsqueeze(0)  # Shape: [B, local_B, C, D, H, W] -> [1, 2, 1, 1, 480, 480]
+        
+        # Set modality indicator for conditions
+        con_modal_indices = torch.tensor([MODALITY_DICT['CXR']] * len(con_batch_image)).to(device)
+        
+        # Tokenize anatomical condition
+        anatomy_condition_tokens = tokenizer(
+            anatomy_condition,
+            return_tensors="pt", 
+            padding="max_length", 
+            truncation=True, 
+            max_length=512
+        ).to(device)
+        
+        # Extract conditional features
+        _, _, condition_feature, _ = rad_ir(
+            anatomy_condition_tokens,
+            image=con_batch_image,
+            device=device,
+            is_condition=True,  # Indicate conditional scenario
+            return_latents=True,
+            modal_indexs=con_modal_indices,
+            modal_embedding=True
+        )
+        
+        print(f"Condition feature shape: {condition_feature.shape}")  # [1, 2, 512]

README.md

@@ -0,0 +1 @@
+# HMCLIP

check.py

@@ -0,0 +1,10 @@
+import numpy as np
+
+# 加载文件
+data = np.load('dataset/train_ratescore.npy')
+
+# 查看维度（例如返回 (100, 224, 224, 3)）
+print(f"数据的形状: {data.shape}")
+
+# 查看维度的数量（例如返回 4）
+print(f"维度的数量: {data.ndim}")

data_process.py

@@ -0,0 +1,143 @@
+import nibabel as nib
+import torch
+import numpy as np
+import torch.nn.functional as F
+from PIL import Image
+import random
+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 = nib.load(str(path))
+    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

retrieval_metric.py

@@ -0,0 +1,463 @@
+import torch
+import numpy as np
+
+def hard_retrieval(logits, gt_logits, k):
+    # treat ratescore > x as gt and calculate recall
+    
+    max_indices = torch.argmax(gt_logits, dim=1)
+    gt_matrix = torch.zeros_like(gt_logits)
+    gt_matrix[torch.arange(gt_logits.size(0)), max_indices] = 1 # for each row, the position of max element is 1, others 0
+    
+    # gt_matrix = torch.eye(logits.shape[0]).to(logits.device)
+    
+    metrics = {}    
+    top10 = {}
+    gt_ranks = {}
+
+    # a sample may have multiple gt, find the one with highest logits
+    gt_logit = gt_matrix * logits  # [0.1, 0.3, 0.2] * [1, 0, 1] --> [0.1, 0, 0.2]
+    gt_rank = torch.argmax(gt_logit, dim=1).view(-1, 1)  # --> 2
+
+    ranking = torch.argsort(logits, descending=True) # [0.1, 0.3, 0.2] --> [1, 2, 0] ranking logits along each row
+    preds = torch.where(ranking == gt_rank)[1] # [1, 2, 0] == [0] --> [2] where the gt lies in ranks
+    preds = preds.detach().cpu().numpy()
+    
+    if isinstance(k, int):
+        return np.mean(preds < k)
+    else:
+        results = []
+        for current_k in k:
+            results.append(np.mean(preds < current_k))
+
+    return results
+
+def hard_retrieval_exclude_self(logits, gt_logits, index_list, k, return_detail=False):
+    # treat ratescore > x as gt and calculate recall
+    
+    if return_detail:
+        assert len(k)==1, 'must be a single k when return_detail is True'
+    
+    # Exclude self elements using the provided index_list
+    logits = delete_elements_by_row(logits, index_list)
+    gt_logits = delete_elements_by_row(gt_logits, index_list)
+    logits = logits.detach().cpu()
+    gt_logits = gt_logits.detach().cpu()
+
+    # Ranking the logits in descending order for each row
+    ranking = torch.argsort(logits, descending=True)
+
+    # Define GT as all positions (excluding self) with gt_logits > 0.9
+    gt_mask = gt_logits > 0.9
+
+    # 过滤无正样本的行，只统计有GT的sample
+    valid_rows = gt_mask.any(dim=1)
+    if not valid_rows.any():
+        return -1  # 或无正样本时的处理
+
+    ranking = ranking[valid_rows]
+    gt_mask = gt_mask[valid_rows]
+
+    if return_detail:
+        k = k[0]
+        topk = ranking[:, :k]
+        # Check for each row if any of the topk indices is a valid GT
+        
+        topk_gt = torch.gather(gt_mask, 1, topk)
+        hits = topk_gt.any(dim=1).float().numpy()
+        pred_topk_indices = ranking[:, :k].detach().cpu().numpy()   # we delete them selves, so need to add 1
+        pred_topk_scores = logits.gather(1, ranking[:, :k]).detach().cpu().numpy()
+        
+        gt_topk = torch.topk(gt_logits, k, dim=1)
+        gt_topk_indices = gt_topk.indices.detach().cpu().numpy()
+        gt_topk_scores = gt_topk.values.detach().cpu().numpy()
+        
+        detail = {
+            "recall": [np.mean(hits)],
+            "prediction_topk_indices": pred_topk_indices,
+            "prediction_topk_scores": pred_topk_scores,
+            "gt_topk_indices": gt_topk_indices,
+            "gt_topk_scores": gt_topk_scores,
+            "row_recall": hits
+        }
+        return detail
+    else:
+        results = []
+        for current_k in k:
+            topk = ranking[:, :current_k]
+            topk_gt = torch.gather(gt_mask, 1, topk)
+            hits = topk_gt.any(dim=1).float().numpy()
+            results.append(np.mean(hits))
+        return results
+    
+def hard_retrieval_exclude_self_return_details(logits, gt_logits, index_list, k, return_detail=False):
+    # treat ratescore > x as gt and calculate recall
+    # to visualization, visualize retrieval details
+    
+    assert len(k)==1, 'must be a single k when return_detail is True'
+    
+    # Exclude self elements using the provided index_list
+    logits = delete_elements_by_row(logits, index_list)
+    gt_logits = delete_elements_by_row(gt_logits, index_list)
+    logits = logits.detach().cpu()
+    gt_logits = gt_logits.detach().cpu()
+
+    # Ranking the logits in descending order for each row
+    ranking = torch.argsort(logits, descending=True)
+
+    # Define GT as all positions (excluding self) with gt_logits > 0.9
+    gt_mask = gt_logits > 0.9
+
+    k = k[0]
+    topk = ranking[:, :k]
+    # Check for each row if any of the topk indices is a valid GT
+    
+    topk_gt = torch.gather(gt_mask, 1, topk)
+    hits = topk_gt.any(dim=1).float().numpy()
+    pred_topk_indices = ranking[:, :k].detach().cpu().numpy()   # we delete them selves, so need to add 1
+    pred_topk_scores = gt_logits.gather(1, ranking[:, :k]).detach().cpu().numpy()
+    
+    gt_topk = torch.topk(gt_logits, k, dim=1)
+    gt_topk_indices = gt_topk.indices.detach().cpu().numpy()
+    gt_topk_scores = gt_topk.values.detach().cpu().numpy()
+    
+    detail = {
+        "recall": [np.mean(hits)],
+        "prediction_topk_indices": pred_topk_indices,
+        "prediction_topk_scores": pred_topk_scores,
+        "gt_topk_indices": gt_topk_indices,
+        "gt_topk_scores": gt_topk_scores,
+        "row_recall": hits
+    }
+    return detail
+
+def soft_retrieval(logits, similarity_matrix, k=[1, 5, 10]):
+    # calculate avg and upperbound ratescore of image 2 image retrieval results
+    # NOTE: do not have to exclude items themselves in logits or similarity_matrix
+    
+    ranking = torch.argsort(logits, descending=True) # [0.1, 0.3, 0.2] --> [1, 2, 0] ranking logits along each row
+    pred_score = []
+    upperbound_score = []
+
+    for current_k in k:
+        top_ranking = ranking[:, 1:current_k+1] # [n, k] exclude themselves (top 1)
+        avg_similarities = torch.zeros(top_ranking.shape[0])
+        for i in range(top_ranking.shape[0]):
+            top_idx = top_ranking[i].tolist()
+            avg_similarities[i] = similarity_matrix[i, top_idx].mean()
+        pred_score.append(avg_similarities.mean().item())
+
+    for current_k in k:
+        actual_k = min(current_k + 1, similarity_matrix.size(1))
+        topk_values = torch.topk(similarity_matrix, actual_k, dim=1).values  # shape: (n_samples, actual_k)
+        # 如果实际取出的topk不足以排除 self，则直接使用所有
+        if actual_k > 1:
+            topk_without_self = topk_values[:, 1:]
+        else:
+            topk_without_self = topk_values
+        upperbound_score.append(topk_without_self.mean().item())
+
+    return pred_score, upperbound_score
+
+def soft_retrieval_uncon(logits, gt_logits, k):
+    # treat ratescore > x as gt and calculate recall
+    
+    logits = logits.detach().cpu()
+    gt_logits = gt_logits.detach().cpu()
+
+    # Ranking the logits in descending order for each row
+    ranking = torch.argsort(logits, descending=True)
+
+    # Define GT as all positions (excluding self) with gt_logits > 0.9
+    gt_mask = gt_logits > 0.9
+
+    # 过滤无正样本的行，只统计有GT的sample
+    valid_rows = gt_mask.any(dim=1)
+    if not valid_rows.any():
+        return -1  # 或无正样本时的处理
+
+    ranking = ranking[valid_rows]
+    gt_mask = gt_mask[valid_rows]
+
+    if isinstance(k, int):
+        topk = ranking[:, :k]
+        # Check for each row if any of the topk indices is a valid GT
+        topk_gt = torch.gather(gt_mask, 1, topk)
+        hits = topk_gt.any(dim=1).float().numpy()
+        return np.mean(hits)
+    else:
+        results = []
+        for current_k in k:
+            topk = ranking[:, :current_k]
+            topk_gt = torch.gather(gt_mask, 1, topk)
+            hits = topk_gt.any(dim=1).float().numpy()
+            results.append(np.mean(hits))
+        return results
+
+def soft_retrieval_exclude_self(logits, gt_logits, index_list, k):
+    # treat ratescore > x as gt and calculate recall
+    logits = delete_elements_by_row(logits, index_list)
+    gt_logits = delete_elements_by_row(gt_logits, index_list)
+    logits = logits.detach().cpu()
+    gt_logits = gt_logits.detach().cpu()
+
+    # Ranking the logits in descending order for each row
+    ranking = torch.argsort(logits, descending=True)
+
+    # Define GT as all positions (excluding self) with gt_logits > 0.9
+    gt_mask = gt_logits > 0.9
+
+    # 过滤无正样本的行，只统计有GT的sample
+    valid_rows = gt_mask.any(dim=1)
+    if not valid_rows.any():
+        return -1  # 或无正样本时的处理
+    ranking = ranking[valid_rows]
+    gt_mask = gt_mask[valid_rows]
+
+    if isinstance(k, int):
+        topk = ranking[:, :k]
+        # Check for each row if any of the topk indices is a valid GT
+        topk_gt = torch.gather(gt_mask, 1, topk)
+        hits = topk_gt.any(dim=1).float().numpy()
+        return np.mean(hits)
+    else:
+        results = []
+        for current_k in k:
+            topk = ranking[:, :current_k]
+            topk_gt = torch.gather(gt_mask, 1, topk)
+            hits = topk_gt.any(dim=1).float().numpy()
+            results.append(np.mean(hits))
+        print('results:', results)
+        return results
+
+
+def RateScore_retrieval(logits, similarity_matrix, k=[1, 5, 10]):
+    # calculate avg and upperbound ratescore of retrieval
+    # NOTE: do not have to exclude items themselves in logits or similarity_matrix
+    
+    ranking = torch.argsort(logits, descending=True) # [0.1, 0.3, 0.2] --> [1, 2, 0] ranking logits along each row
+    pred_score = []
+    upperbound_score = []
+
+    for current_k in k:
+        top_ranking = ranking[:, 1:current_k+1] # [n, k] exclude themselves (top 1)
+        avg_similarities = torch.zeros(top_ranking.shape[0])
+        for i in range(top_ranking.shape[0]):
+            top_idx = top_ranking[i].tolist()
+            avg_similarities[i] = similarity_matrix[i, top_idx].mean()
+        pred_score.append(avg_similarities.mean().item())
+
+    for current_k in k:
+        actual_k = min(current_k + 1, similarity_matrix.size(1))
+        topk_values = torch.topk(similarity_matrix, actual_k, dim=1).values  # shape: (n_samples, actual_k)
+        # 如果实际取出的topk不足以排除 self，则直接使用所有
+        if actual_k > 1:
+            topk_without_self = topk_values[:, 1:]
+        else:
+            topk_without_self = topk_values
+        upperbound_score.append(topk_without_self.mean().item())
+
+    return pred_score, upperbound_score
+
+
+def compute_ndcg_exclude_self(pred_scores, true_scores, index_list, k=None):
+    """
+    计算NDCG。
+    当 k 是一个整数时，计算对应 k 值的 ndcg；当 k 是列表时，分别计算每个 k 的 ndcg，返回一个与 k 列表等长的列表，每个元素是对应的 score。
+    输入的 true_scores 和 pred_scores 为 tensor 类型。
+    :param true_scores: 真实相似度矩阵 (n_query, n_samples)
+    :param pred_scores: 模型预测的相似度矩阵 (n_query, n_samples)
+    :param index_list: 每行需要删除的元素的索引列表
+    :param k: 只考虑前 k 个结果。如果为 None，则考虑所有结果；如果为 list，则分别计算每个 k 的 ndcg
+    :return: 当 k 为 list 时返回一个 ndcg 值的列表，否则返回单个 ndcg 值
+    """
+    true_scores = delete_elements_by_row(true_scores, index_list)
+    pred_scores = delete_elements_by_row(pred_scores, index_list)
+    # 转换为 numpy 数组以便后续排序和计算
+    true_scores = true_scores.detach().cpu().numpy()
+    pred_scores = pred_scores.detach().cpu().numpy()
+    n_query = true_scores.shape[0]
+
+    if k is None:
+        k_list = [true_scores.shape[1]]
+    elif isinstance(k, list):
+        k_list = k
+    else:
+        k_list = [k]
+
+    ndcg_results = []
+    for current_k in k_list:
+        ndcg_scores = []
+        for i in range(n_query):
+            true = true_scores[i]
+            pred = pred_scores[i]
+            pred_rank = np.argsort(pred)[::-1]
+            true_sorted_by_pred = true[pred_rank]
+
+            # 计算 DCG
+            if current_k is not None:
+                true_sorted_by_pred_k = true_sorted_by_pred[:current_k]
+            else:
+                true_sorted_by_pred_k = true_sorted_by_pred
+            dcg = compute_dcg(true_sorted_by_pred_k)
+
+            # 计算 IDCG
+            ideal_sorted_true = np.sort(true)[::-1]
+            if current_k is not None:
+                ideal_sorted_true = ideal_sorted_true[:current_k]
+            idcg = compute_dcg(ideal_sorted_true)
+
+            ndcg = dcg / idcg if idcg != 0 else 0
+            ndcg_scores.append(ndcg)
+        ndcg_results.append(np.mean(ndcg_scores))
+
+    # 如果原始 k 不是 list，则返回单个值
+    if not isinstance(k, list):
+        return ndcg_results[0]
+    return ndcg_results
+
+def compute_ndcg_uncon(pred_scores, true_scores, k=None):
+    """
+    计算NDCG。
+    当 k 是一个整数时，计算对应 k 值的 ndcg；当 k 是列表时，分别计算每个 k 的 ndcg，返回一个与 k 列表等长的列表，每个元素是对应的 score。
+    输入的 true_scores 和 pred_scores 为 tensor 类型。
+    :param true_scores: 真实相似度矩阵 (n_query, n_samples)
+    :param pred_scores: 模型预测的相似度矩阵 (n_query, n_samples)
+    :param k: 只考虑前 k 个结果。如果为 None，则考虑所有结果；如果为 list，则分别计算每个 k 的 ndcg
+    :return: 当 k 为 list 时返回一个 ndcg 值的列表，否则返回单个 ndcg 值
+    """
+    # 转换为 numpy 数组以便后续排序和计算
+    true_scores = true_scores.detach().cpu().numpy()
+    pred_scores = pred_scores.detach().cpu().numpy()
+    n_query = true_scores.shape[0]
+
+    if k is None:
+        k_list = [true_scores.shape[1]]
+    elif isinstance(k, list):
+        k_list = k
+    else:
+        k_list = [k]
+
+    ndcg_results = []
+    for current_k in k_list:
+        ndcg_scores = []
+        for i in range(n_query):
+            true = true_scores[i]
+            pred = pred_scores[i]
+            pred_rank = np.argsort(pred)[::-1]
+            true_sorted_by_pred = true[pred_rank]
+
+            # 计算 DCG
+            if current_k is not None:
+                true_sorted_by_pred_k = true_sorted_by_pred[:current_k]
+            else:
+                true_sorted_by_pred_k = true_sorted_by_pred
+            dcg = compute_dcg(true_sorted_by_pred_k)
+
+            # 计算 IDCG
+            ideal_sorted_true = np.sort(true)[::-1]
+            if current_k is not None:
+                ideal_sorted_true = ideal_sorted_true[:current_k]
+            idcg = compute_dcg(ideal_sorted_true)
+
+            ndcg = dcg / idcg if idcg != 0 else 0
+            ndcg_scores.append(ndcg)
+        ndcg_results.append(np.mean(ndcg_scores))
+
+    # 如果原始 k 不是 list，则返回单个值
+    if not isinstance(k, list):
+        return ndcg_results[0]
+    return ndcg_results
+
+
+def delete_elements_by_row(matrix, delete_list):
+    """
+    根据给定的索引列表，删除每行中指定位置的元素
+    :param matrix: 输入的 n*m 维 tensor
+    :param delete_list: 一个 n 大小的 list，表示每行需要删除的元素的索引
+    :return: 处理后的 n*(m-1) 维 tensor
+    """
+    n, m = matrix.size()
+    result = []
+
+    for i in range(n):
+        row = matrix[i]
+        delete_index = delete_list[i]
+        # 删除指定位置的元素
+        new_row = torch.cat([row[:delete_index], row[delete_index+1:]])
+        result.append(new_row)
+
+    return torch.stack(result)
+
+def compute_dcg(scores):
+    """计算DCG"""
+    return np.sum(scores / np.log2(np.arange(2, len(scores) + 2)))
+
+def compute_ndcg(pred_scores, true_scores, index_list, k=None):
+    """
+    计算NDCG。
+    当 k 是一个整数时，计算对应 k 值的 ndcg；当 k 是列表时，分别计算每个 k 的 ndcg，返回一个与 k 列表等长的列表，每个元素是对应的 score。
+    输入的 true_scores 和 pred_scores 为 tensor 类型。
+    :param true_scores: 真实相似度矩阵 (n_query, n_samples)
+    :param pred_scores: 模型预测的相似度矩阵 (n_query, n_samples)
+    :param index_list: 每行需要删除的元素的索引列表
+    :param k: 只考虑前 k 个结果。如果为 None，则考虑所有结果；如果为 list，则分别计算每个 k 的 ndcg
+    :return: 当 k 为 list 时返回一个 ndcg 值的列表，否则返回单个 ndcg 值
+    """
+    true_scores = delete_elements_by_row(true_scores, index_list)
+    pred_scores = delete_elements_by_row(pred_scores, index_list)
+    # 转换为 numpy 数组以便后续排序和计算
+    true_scores = true_scores.detach().cpu().numpy()
+    pred_scores = pred_scores.detach().cpu().numpy()
+    n_query = true_scores.shape[0]
+
+    if isinstance(k, list):
+        k_list = k
+    else:
+        k_list = [k]
+
+    ndcg_results = []
+    for current_k in k_list:
+        ndcg_scores = []
+        for i in range(n_query):
+            true = true_scores[i]
+            pred = pred_scores[i]
+            pred_rank = np.argsort(pred)[::-1]
+            true_sorted_by_pred = true[pred_rank]
+
+            # 计算 DCG
+            if current_k is not None:
+                true_sorted_by_pred_k = true_sorted_by_pred[:current_k]
+            else:
+                true_sorted_by_pred_k = true_sorted_by_pred
+            dcg = compute_dcg(true_sorted_by_pred_k)
+
+            # 计算 IDCG
+            ideal_sorted_true = np.sort(true)[::-1]
+            if current_k is not None:
+                ideal_sorted_true = ideal_sorted_true[:current_k]
+            idcg = compute_dcg(ideal_sorted_true)
+
+            ndcg = dcg / idcg if idcg != 0 else 0
+            ndcg_scores.append(ndcg)
+        ndcg_results.append(np.mean(ndcg_scores))
+
+    # 如果原始 k 不是 list，则返回单个值
+    if not isinstance(k, list):
+        return ndcg_results[0]
+    return ndcg_results
+
+
+
+if __name__ == '__main__':
+    logits = torch.rand(4, 4)
+    logits.fill_diagonal_(1)
+    gt_logits = torch.rand(4, 4)
+    gt_logits.fill_diagonal_(1)
+    index = [0, 1, 2, 3]
+    
+    print('logits:', logits)
+    print('gt_logits:', gt_logits)
+    
+    ndcg = compute_ndcg(gt_logits, logits, index, k=[1, 3])
+    
+    print(ndcg)

run_CXR_CT_train.py

@@ -0,0 +1,231 @@
+import os
+import argparse
+import torch
+import wandb
+from datetime import datetime
+
+from transformer_maskgit.transformer_maskgit.ctvit import CTViT
+from transformers import AutoTokenizer, AutoModel
+
+from Rad_IR.radir import RADIR
+from Rad_IR.radir.data.hmclip_mimic_box import HyCoClipMimicBoxDataset_JPG
+
+from scripts.RadIR_Trainer_CXR_CT import RadIRTrainer
+
+
+# 当前 ablation 标识，改 loss 时同步更新此处
+ABLATION_TAG = "no_i2t_ce"
+
+
+def build_run_name(args):
+    if args.run_name:
+        return args.run_name
+    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
+    return f"CXRCT_stage1_{ABLATION_TAG}_{timestamp}"
+
+
+def build_results_folder(args, run_name):
+    if args.results_folder:
+        return args.results_folder
+    return os.path.join("./outputs/ablation", run_name)
+
+
+# --------------------------------------------------------
+# Args
+# --------------------------------------------------------
+def get_args():
+    parser = argparse.ArgumentParser("RADIR Training with Trainer")
+
+    parser.add_argument('--resume', type=str, default=None, help="checkpoint 路径，填了就断点续训")
+    parser.add_argument('--batch', type=int, default=16)
+    parser.add_argument('--num_train_steps', type=int, default=100000)
+    parser.add_argument('--warmup_steps', type=int, default=1000)
+
+    parser.add_argument('--lr', type=float, default=5e-5)
+    parser.add_argument('--wd', type=float, default=0.0)
+
+    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')
+
+    parser.add_argument(
+        '--run_name',
+        type=str,
+        default=None,
+        help='实验名称；默认 CXRCT_stage1_{ablation_tag}_{timestamp}',
+    )
+    parser.add_argument(
+        '--results_folder',
+        type=str,
+        default=None,
+        help='输出目录；默认 ./outputs/ablation/{run_name}',
+    )
+    parser.add_argument(
+        '--wandb_project',
+        type=str,
+        default='RadIR_Ablation',
+        help='wandb project 名称',
+    )
+    parser.add_argument(
+        '--wandb_group',
+        type=str,
+        default='ablation_no_i2t_ce',
+        help='wandb group 名称',
+    )
+
+    # Dataset
+    parser.add_argument(
+        '--dataset_root',
+        type=str,
+        default="/testessfs10/users/zeyu.zhang/HMClip/RadIR/RadIR/dataset",
+    )
+    parser.add_argument(
+        '--matched_dir',
+        type=str,
+        default="/testessfs10/users/zeyu.zhang/HMClip/RadIR/RadIR/dataset/processed/matched",
+    )
+
+    # Local HF model
+    parser.add_argument(
+        '--model_path',
+        type=str,
+        default="/testessfs10/users/zeyu.zhang/HMClip/RadIR/RadIR/hf_models/microsoft/BiomedVLP-CXR-BERT-specialized",
+    )
+
+    parser.add_argument('--num_workers', type=int, default=4)
+
+    return parser.parse_args()
+
+
+# --------------------------------------------------------
+# Main
+# --------------------------------------------------------
+def main():
+    args = get_args()
+    run_name = build_run_name(args)
+    args.results_folder = build_results_folder(args, run_name)
+    args.run_name = run_name
+
+    if not torch.cuda.is_available():
+        raise RuntimeError("需要 GPU（CUDA）")
+
+    os.makedirs(args.results_folder, exist_ok=True)
+
+    print(f">>> Run name: {run_name}")
+    print(f">>> Results folder: {args.results_folder}")
+
+    # --------------------------------------------------------
+    # Load HF Text Encoder
+    # --------------------------------------------------------
+    print(">>> Loading tokenizer & text encoder...")
+    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
+    )
+
+    # --------------------------------------------------------
+    # Load Image Encoder
+    # --------------------------------------------------------
+    print(">>> Loading CTViT...")
+    image_encoder = CTViT(
+        dim=768,
+        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,
+    )
+
+    # --------------------------------------------------------
+    # Build RADIR
+    # --------------------------------------------------------
+    print(">>> Building RADIR...")
+    radir = 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,
+    )
+
+    # --------------------------------------------------------
+    # Dataset
+    # --------------------------------------------------------
+    print(">>> Loading datasets...")
+    train_dataset = HyCoClipMimicBoxDataset_JPG(
+        split="train",
+        tokenizer=tokenizer,
+        dataset_root=args.dataset_root,
+        cig_matched_dir=args.matched_dir,
+        caption_file="train.jsonl",
+        transform=None,
+    )
+
+    valid_dataset = HyCoClipMimicBoxDataset_JPG(
+        split="valid",
+        tokenizer=tokenizer,
+        dataset_root=args.dataset_root,
+        cig_matched_dir=args.matched_dir,
+        caption_file="valid.jsonl",
+        transform=None,
+    )
+
+    if int(os.environ.get("RANK", 0)) == 0:
+        wandb.init(
+            project=args.wandb_project,
+            name=run_name,
+            group=args.wandb_group,
+            tags=["stage1", "CXRCT", "ablation", ABLATION_TAG],
+            dir=args.results_folder,
+            config=vars(args),
+        )
+
+    # --------------------------------------------------------
+    # Trainer
+    # --------------------------------------------------------
+    print(">>> Building Trainer...")
+    trainer = RadIRTrainer(
+        RADIR=radir,
+        tokenizer=tokenizer,
+        train_dataset=train_dataset,
+        valid_dataset=valid_dataset,
+        num_train_steps=args.num_train_steps,
+        warmup_steps=args.warmup_steps,
+        batch_size=args.batch,
+        lr=args.lr,
+        wd=args.wd,
+        results_folder=args.results_folder,
+        num_workers=args.num_workers,
+        modal_embedding=args.modal_embedding,
+        debug=False,
+    )
+
+    # --------------------------------------------------------
+    # Start Training
+    # --------------------------------------------------------
+    resume_path = None
+    if args.resume is not None:
+        if args.resume == "auto":
+            resume_path = os.path.join(args.results_folder, "radir_last.pt")
+        else:
+            resume_path = args.resume
+
+    print("✅ Start Training...")
+    trainer.train_model(resume_path=resume_path)
+
+
+if __name__ == "__main__":
+    main()

stage1_eva_i2t.py

@@ -0,0 +1,393 @@
+from __future__ import annotations
+
+"""
+I2T / T2I evaluation with RaTEScore-aligned GT.
+
+Uses valid.jsonl for image-text pairs and maps each sample to val_caption.csv
+row indices, then slices val_ratescore.npy into a sub-matrix aligned with the
+eval set order (same logic as RadIR evaluate_uncon soft_label path).
+"""
+
+import argparse
+import json
+import os
+
+import numpy as np
+import pandas as pd
+import torch
+from torch.utils.data import DataLoader, Subset
+from transformers import AutoModel, AutoTokenizer
+
+from Rad_IR.radir import lorentz as L
+from Rad_IR.radir.RadIR import RADIR, l2norm
+from Rad_IR.radir.data.hmclip_mimic_box import HyCoClipMimicBoxDataset_JPG, collate_hmclip_box
+from retrieval_metric import compute_ndcg_uncon, soft_retrieval_uncon
+from transformer_maskgit.transformer_maskgit.ctvit import CTViT
+
+ROOT = os.path.dirname(os.path.abspath(__file__))
+
+
+# ---------------------------------------------------------------------------
+# RaTEScore index alignment (valid.jsonl  ->  val_caption.csv)
+# ---------------------------------------------------------------------------
+
+def study_key_from_csv_path(file_path: str) -> str:
+    """/p10/p10190940/s51351116.txt -> p10/p10190940/s51351116"""
+    return os.path.splitext(file_path.lstrip("/"))[0]
+
+
+def study_key_from_jsonl_img_path(img_path: str) -> str:
+    """dataset/p10/p10190940/s51351116/xxx.jpg -> p10/p10190940/s51351116"""
+    p = img_path.replace("\\", "/")
+    if p.startswith("dataset/"):
+        p = p[len("dataset/") :]
+    parts = p.split("/")
+    if len(parts) < 3:
+        raise ValueError(f"Cannot parse study key from img_path: {img_path}")
+    return "/".join(parts[:3])
+
+
+def build_eval_index_mapping(
+    jsonl_path: str,
+    csv_path: str,
+    dedupe_study: bool = False,
+):
+    """
+    Map each eval row in jsonl to a row index in val_caption.csv / val_ratescore.npy.
+
+    Returns:
+        keep_row_indices: indices into jsonl/dataset to evaluate (in jsonl order)
+        csv_indices: parallel list of val_caption.csv row indices for RaTEScore slice
+        stats: dict with alignment diagnostics
+    """
+    df = pd.read_csv(csv_path)
+    csv_keys = [study_key_from_csv_path(p) for p in df["File Path"]]
+    key_to_csv_idx = {k: i for i, k in enumerate(csv_keys)}
+
+    jsonl_keys = []
+    with open(jsonl_path, encoding="utf-8") as f:
+        for line in f:
+            item = json.loads(line.strip())
+            jsonl_keys.append(study_key_from_jsonl_img_path(item["img_path"]))
+
+    keep_row_indices = []
+    csv_indices = []
+    seen_study = set()
+    unmatched = []
+
+    for row_i, key in enumerate(jsonl_keys):
+        if key not in key_to_csv_idx:
+            unmatched.append((row_i, key))
+            continue
+        if dedupe_study and key in seen_study:
+            continue
+        seen_study.add(key)
+        keep_row_indices.append(row_i)
+        csv_indices.append(key_to_csv_idx[key])
+
+    stats = {
+        "jsonl_rows": len(jsonl_keys),
+        "csv_rows": len(csv_keys),
+        "eval_rows": len(keep_row_indices),
+        "unique_studies": len(seen_study),
+        "unmatched_jsonl": len(unmatched),
+        "csv_not_in_jsonl": sum(1 for k in csv_keys if k not in set(jsonl_keys)),
+        "dedupe_study": dedupe_study,
+    }
+    if unmatched:
+        stats["unmatched_examples"] = unmatched[:5]
+
+    return keep_row_indices, csv_indices, stats
+
+
+def load_aligned_ratescore_gt(gt_path: str, csv_indices: list[int]) -> torch.Tensor:
+    """Slice full RaTEScore matrix to match eval sample order."""
+    gt = np.load(gt_path)
+    idx = np.asarray(csv_indices, dtype=np.int64)
+    gt_sub = gt[np.ix_(idx, idx)].astype(np.float32)
+    return torch.from_numpy(gt_sub / 100.0)
+
+
+# ---------------------------------------------------------------------------
+# Feature extraction (aligned with RadIR.forward is_condition=False global path)
+# ---------------------------------------------------------------------------
+
+def preprocess_image_like_forward(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 ndim={image.ndim}, shape={tuple(image.shape)}")
+    return image
+
+
+def extract_hyperbolic_features(model, dataloader, device, modal_embedding=False):
+    """Extract hyperbolic features (project_img/txt after l2norm global encodings)."""
+    all_hyp_img, all_hyp_txt = [], []
+
+    model.eval()
+    with torch.no_grad():
+        for batch in dataloader:
+            imgs = preprocess_image_like_forward(batch["imgs"].to(device, non_blocking=True))
+
+            enc_image = model.visual_transformer(
+                imgs,
+                return_encoded_tokens=True,
+                modal_embedding=modal_embedding,
+                modal_indexs=None,
+                is_condition=False,
+            )
+            if enc_image.dim() == 2:
+                image_global = l2norm(enc_image)
+            elif enc_image.dim() == 3:
+                image_global = l2norm(enc_image[:, 0, :])
+            else:
+                raise ValueError(f"Unexpected visual feature shape: {enc_image.shape}")
+
+            text_embeddings = model.text_transformer(
+                input_ids=batch["caption_ids"].to(device, non_blocking=True),
+                token_type_ids=batch["token_type_ids"].to(device, non_blocking=True),
+                attention_mask=batch["attention_mask"].to(device, non_blocking=True),
+            )
+            enc_text = text_embeddings[0]
+            text_global = enc_text[:, 0, :]
+
+            all_hyp_img.append(model.project_img(image_global).cpu())
+            all_hyp_txt.append(model.project_txt(text_global).cpu())
+
+    return torch.cat(all_hyp_img, dim=0), torch.cat(all_hyp_txt, dim=0)
+
+
+# ---------------------------------------------------------------------------
+# Metrics
+# ---------------------------------------------------------------------------
+
+def paired_recall(sim_matrix: torch.Tensor, ks):
+    """Hard paired recall: query i should retrieve candidate i (diagonal)."""
+    ranking = torch.argsort(sim_matrix, dim=1, descending=True)
+    gt = torch.arange(sim_matrix.size(0), device=sim_matrix.device)
+    results = {}
+    for k in ks:
+        hits = (ranking[:, :k] == gt.unsqueeze(1)).any(dim=1).float()
+        results[k] = hits.mean().item()
+    return results
+
+
+def print_metric_block(title, ks, paired_r, soft_r, ndcg, scale_pct=True):
+    print(f"\n--- {title} ---")
+    for k in ks:
+        sr = soft_r[k] if isinstance(soft_r, dict) else soft_r[ks.index(k)]
+        if sr == -1:
+            sr_str = "N/A"
+        else:
+            sr_str = f"{sr * 100:.2f}%" if scale_pct else f"{sr:.4f}"
+        nd = ndcg[ks.index(k)]
+        nd_str = f"{nd * 100:.2f}" if scale_pct else f"{nd:.4f}"
+        print(
+            f"@{k:>3d}  Paired_R={paired_r[k] * 100:.2f}%  "
+            f"Soft_R={sr_str}  NDCG={nd_str}"
+        )
+
+
+def evaluate_i2t_t2i(sim_i2t, gt_ratescore, ks):
+    """Run RadIR-style I2T metrics with RaTEScore GT matrix."""
+    sim_t2i = sim_i2t.t()
+    gt = gt_ratescore.to(sim_i2t.device)
+
+    paired_i2t = paired_recall(sim_i2t, ks)
+    paired_t2i = paired_recall(sim_t2i, ks)
+
+    soft_i2t = soft_retrieval_uncon(sim_i2t, gt, k=ks)
+    soft_t2i = soft_retrieval_uncon(sim_t2i, gt, k=ks)
+    if not isinstance(soft_i2t, list):
+        soft_i2t = [soft_i2t]
+        soft_t2i = [soft_t2i]
+
+    ndcg_i2t = compute_ndcg_uncon(sim_i2t, gt, k=ks)
+    ndcg_t2i = compute_ndcg_uncon(sim_t2i, gt, k=ks)
+
+    soft_i2t = dict(zip(ks, soft_i2t))
+    soft_t2i = dict(zip(ks, soft_t2i))
+
+    return {
+        "i2t": {"paired_r": paired_i2t, "soft_r": soft_i2t, "ndcg": ndcg_i2t},
+        "t2i": {"paired_r": paired_t2i, "soft_r": soft_t2i, "ndcg": ndcg_t2i},
+    }
+
+
+# ---------------------------------------------------------------------------
+# Main
+# ---------------------------------------------------------------------------
+
+def main():
+    parser = argparse.ArgumentParser(
+        description="I2T/T2I evaluation with RaTEScore-aligned GT (valid.jsonl subset)"
+    )
+    parser.add_argument("--device", default="cuda", type=str)
+    parser.add_argument(
+        "--model_path",
+        default=os.path.join(
+            ROOT, "hf_models/microsoft/BiomedVLP-CXR-BERT-specialized"
+        ),
+    )
+    parser.add_argument("--radir_ckpt", default=os.path.join(ROOT, "outputs/ablation/CXRCT_stage1_no_cross_text_entail_20260703_153255/radir_best.pt"))
+    parser.add_argument("--dataset_root", default=os.path.join(ROOT, "dataset"))
+    parser.add_argument("--caption_file", default="valid.jsonl")
+    parser.add_argument(
+        "--csv_file",
+        default=os.path.join(ROOT, "dataset/val_caption.csv"),
+        help="Original RadIR caption list (defines val_ratescore row order)",
+    )
+    parser.add_argument(
+        "--gt_path",
+        default=os.path.join(ROOT, "dataset/val_ratescore.npy"),
+    )
+    parser.add_argument("--cig_matched_dir", default="dataset/processed/matched")
+    parser.add_argument("--split", default="valid", type=str)
+    parser.add_argument("--batch_size", default=16, type=int)
+    parser.add_argument("--num_workers", default=8, type=int)
+    parser.add_argument("--imsize", default=224, type=int)
+    parser.add_argument("--max_words", default=128, type=int)
+    parser.add_argument(
+        "--dedupe_study",
+        action="store_true",
+        help="Keep one image per study (~2808). Default: all jsonl rows (~3120).",
+    )
+    parser.add_argument(
+        "--modal_embedding",
+        action="store_true",
+        help="Match training when --modal_embedding was enabled",
+    )
+    parser.add_argument(
+        "--ks",
+        type=int,
+        nargs="+",
+        default=[1, 5, 10, 20, 50, 100],
+    )
+    args = parser.parse_args()
+    device = args.device
+
+    jsonl_path = os.path.join(args.dataset_root, args.caption_file)
+    keep_rows, csv_indices, stats = build_eval_index_mapping(
+        jsonl_path, args.csv_file, dedupe_study=args.dedupe_study
+    )
+    if stats["unmatched_jsonl"]:
+        raise RuntimeError(
+            f"{stats['unmatched_jsonl']} jsonl rows have no match in {args.csv_file}. "
+            f"Examples: {stats.get('unmatched_examples')}"
+        )
+
+    print("===== RaTEScore alignment =====")
+    for k, v in stats.items():
+        if k != "unmatched_examples":
+            print(f"  {k}: {v}")
+
+    gt_ratescore = load_aligned_ratescore_gt(args.gt_path, csv_indices)
+    print(f"  gt_sub shape: {tuple(gt_ratescore.shape)}")
+
+    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
+    )
+    image_encoder = CTViT(
+        dim=768,
+        codebook_size=8192,
+        image_size=args.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=768,
+        dim_latent=768,
+        use_mlm=False,
+        use_all_token_embeds=False,
+    )
+    print(f"\nLoading checkpoint: {args.radir_ckpt}")
+    model.load(args.radir_ckpt)
+    model.to(device)
+
+    full_dataset = HyCoClipMimicBoxDataset_JPG(
+        split=args.split,
+        tokenizer=tokenizer,
+        dataset_root=args.dataset_root,
+        cig_matched_dir=args.cig_matched_dir,
+        caption_file=args.caption_file,
+        imsize=args.imsize,
+        max_words=args.max_words,
+        transform=None,
+    )
+    dataset = Subset(full_dataset, keep_rows)
+    dataloader = DataLoader(
+        dataset,
+        batch_size=args.batch_size,
+        shuffle=False,
+        num_workers=args.num_workers,
+        pin_memory=True,
+        collate_fn=collate_hmclip_box,
+    )
+    print(f"\nEval samples = {len(dataset)} (from {len(full_dataset)} jsonl rows)")
+
+    hyp_img, hyp_txt = extract_hyperbolic_features(
+        model,
+        dataloader,
+        device,
+        modal_embedding=args.modal_embedding,
+    )
+
+    gt = gt_ratescore.to(device)
+    ks = args.ks
+
+    curv = model.get_curv()
+    sim_i2t = -L.pairwise_dist(
+        hyp_img.to(device), hyp_txt.to(device), curv
+    ).float()
+
+    print(f"\n===== Hyperbolic I2T/T2I (curv={curv.item():.4f}) =====")
+    res = evaluate_i2t_t2i(sim_i2t, gt, ks)
+    print_metric_block(
+        "Image -> Text",
+        ks,
+        res["i2t"]["paired_r"],
+        res["i2t"]["soft_r"],
+        res["i2t"]["ndcg"],
+    )
+    print_metric_block(
+        "Text -> Image",
+        ks,
+        res["t2i"]["paired_r"],
+        res["t2i"]["soft_r"],
+        res["t2i"]["ndcg"],
+    )
+    print("\n  Summary (mean I2T+T2I):")
+    for i, k in enumerate(ks):
+        pr = (res["i2t"]["paired_r"][k] + res["t2i"]["paired_r"][k]) / 2
+        nd = (res["i2t"]["ndcg"][i] + res["t2i"]["ndcg"][i]) / 2
+        si = res["i2t"]["soft_r"][k]
+        st = res["t2i"]["soft_r"][k]
+        sr = (si + st) / 2 if si != -1 and st != -1 else -1
+        sr_str = f"{sr * 100:.2f}%" if sr != -1 else "N/A"
+        print(
+            f"  @{k:>3d}  Paired_R={pr * 100:.2f}%  "
+            f"Soft_R={sr_str}  NDCG={nd * 100:.2f}"
+        )
+
+
+if __name__ == "__main__":
+    main()

stage1_eva_t2t.py

@@ -0,0 +1,238 @@
+import os
+import torch
+import torch.nn.functional as F
+from torch.utils.data import Dataset, DataLoader
+from PIL import Image
+import numpy as np
+import pandas as pd
+from tqdm import tqdm
+
+from Rad_IR.radir.RadIR import RADIR
+from transformer_maskgit.transformer_maskgit.ctvit import CTViT
+from transformers import AutoTokenizer, AutoModel
+
+from retrieval_metric import hard_retrieval_exclude_self, compute_ndcg_exclude_self
+
+
+# =========================================================
+#  Dataset — strictly follow CSV order
+# =========================================================
+class CXRStage1EvalDataset(Dataset):
+    def __init__(self, csv_file, dataset_root, imsize=224):
+        df = pd.read_csv(csv_file)
+
+        # ⚠️ MUST MATCH RaTEScore ORDER
+        self.paths = df["File Path"].tolist()
+        self.dataset_root = dataset_root
+        self.imsize = imsize
+
+    def __len__(self):
+        return len(self.paths)
+
+    def __getitem__(self, idx):
+        path = self.paths[idx]
+
+        # 不管是不是绝对路径，都按相对路径处理
+        rel_path = path.lstrip("/")
+
+        # 拼 dataset_root
+        full_path = os.path.join(self.dataset_root, rel_path)
+
+        # 去掉 .txt，得到目录路径
+        folder = os.path.splitext(full_path)[0]
+
+        # 找目录里图片（png/jpg/jpeg）
+        imgs = [
+            f for f in os.listdir(folder)
+            if f.lower().endswith((".png", ".jpg", ".jpeg"))
+        ]
+        assert len(imgs) > 0, f"No image found in: {folder}"
+
+        img_path = os.path.join(folder, imgs[0])
+
+        img = Image.open(img_path).convert("L")
+        img = np.array(img, dtype=np.float32) / 255.0
+        img = torch.from_numpy(img).unsqueeze(0)  # [1, H, W]
+
+        img = torch.nn.functional.interpolate(
+            img.unsqueeze(0),          # [1, 1, H, W]
+            size=(self.imsize, self.imsize),
+            mode="bilinear",
+            align_corners=False
+        ).squeeze(0)                  # [1, imsize, imsize]
+
+        img = img * 2.0 - 1.0         # [-1, 1]
+
+        return img, idx
+
+
+# =========================================================
+#  Extract Stage-1 features
+# =========================================================
+def extract_image_features(model, dataloader, device):
+    all_feats = []
+
+    model.eval()
+    with torch.no_grad():
+        for imgs, _ in tqdm(dataloader):
+            imgs = imgs.to(device)  # [B, 1, H, W]
+
+            enc = model.visual_transformer(
+                imgs,
+                return_encoded_tokens=True,
+                modal_embedding=False,
+                modal_indexs=None,
+                is_condition=False
+            )
+
+            latent = model.to_visual_latent(enc)
+
+            # 保证拿到的是 [B, D] 特征
+            if latent.dim() == 3:
+                # [B, T, D] -> token 维平均
+                feat = latent.mean(dim=1)      # [B, D]
+            elif latent.dim() == 2:
+                feat = latent                  # [B, D]
+            else:
+                raise ValueError(f"Unexpected latent shape: {latent.shape}")
+
+            # 按特征维做归一化（cosine）
+            feat = F.normalize(feat, dim=-1)   # [B, D]
+
+            all_feats.append(feat)
+
+    # cat 之后应该是 [N, D]
+    feats = torch.cat(all_feats, dim=0)
+    return feats
+
+
+# =========================================================
+#                       MAIN
+# =========================================================
+def main():
+    device = "cuda"
+
+    # ================= Model Components ==================
+    MODEL_PATH = "/testessfs10/users/zeyu.zhang/HMClip/RadIR/RadIR/hf_models/microsoft/BiomedVLP-CXR-BERT-specialized"
+
+    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=224,
+        patch_size=16,
+        temporal_patch_size=10,
+        spatial_depth=8,
+        temporal_depth=6,
+        cls_depth=4,
+        dim_head=32,
+        heads=8,
+        channels=1
+    )
+
+    print("Loading RADIR...")
+    model = RADIR(
+        tokenizer=tokenizer,
+        image_encoder=image_encoder,
+        text_encoder=text_model,
+        dim_text=768,
+        dim_image=768,
+        dim_latent=768,
+        use_mlm=False,
+        use_all_token_embeds=False
+    )
+
+    model.load("outputs/ablation/CXRCT_stage1_no_cross_text_entail_20260703_153255")
+    model.to(device)
+
+    # ================= Dataset ==================
+    dataset = CXRStage1EvalDataset(
+        csv_file="dataset/val_caption.csv",
+        dataset_root="/testessfs10/users/zeyu.zhang/HMClip/RadIR/RadIR/dataset",
+        imsize=224
+    )
+
+    dataloader = DataLoader(
+        dataset,
+        batch_size=16,
+        shuffle=False,
+        num_workers=8,
+        pin_memory=True
+    )
+
+    print("Eval Images =", len(dataset))
+
+    # ================= Features ==================
+    feats = extract_image_features(model, dataloader, device)
+    print("Feature shape:", feats.shape)  # 期望: [3185, 512]
+
+    # 使用 mT 而不是 T（避免 warning）
+    logits = feats @ feats.mT           # [N, N]
+    print("Logits shape:", logits.shape)
+
+    # ================= Load GT ==================
+    GT_PATH = "dataset/val_ratescore.npy"
+    gt_logits = torch.from_numpy(np.load(GT_PATH)).float().to(device)
+
+    # 关键：RaTEScore 是 0–100，这里归一化到 0–1，
+    # 这样 retrieval_metric 里用的阈值 >0.9 才等价于 “score >= 90”
+    gt_logits = gt_logits / 100.0
+
+    print("GT shape:", gt_logits.shape)
+
+    assert gt_logits.shape == logits.shape, \
+        f"GT {gt_logits.shape} != logits {logits.shape}"
+
+    # ================= Metrics ==================
+    Ks = [5, 10, 50, 100]
+    index_list = list(range(len(dataset)))     # 每行自己的 index，用于排除 self
+
+    # 论文对应的 Image2Image Retrieval：排除 self + RaTEScore>0.9 为正样本
+    
+    recall = hard_retrieval_exclude_self(
+        logits, gt_logits, index_list, k=Ks
+    )
+
+    
+    ndcg = compute_ndcg_exclude_self(
+        
+        logits, gt_logits, index_list, k=Ks
+    )
+
+    print("\n===== Stage-1 (Unconditional CXR Retrieval) =====")
+    for k, r in zip(Ks, recall):
+        print(f"Recall@{k}: {r:.4f}")
+
+    for k, n in zip(Ks, ndcg):
+        print(f"NDCG@{k}:  {n:.4f}")
+
+    # ============ Random baseline sanity check ============
+    rand_logits = torch.randn_like(logits)
+
+    recall_rand = hard_retrieval_exclude_self(
+        rand_logits, gt_logits, index_list, k=Ks
+    )
+    ndcg_rand = compute_ndcg_exclude_self(
+        rand_logits, gt_logits, index_list, k=Ks
+    )
+
+    i = 10
+    a = np.sort(gt_logits[i].cpu().numpy())[::-1]
+    print(a[:20])
+    print(a[-20:])
+
+    print("\n===== Random logits sanity check =====")
+    for k, r in zip(Ks, recall_rand):
+        print(f"Random Recall@{k}: {r:.4f}")
+    for k, n in zip(Ks, ndcg_rand):
+        print(f"Random NDCG@{k}:  {n:.4f}")
+
+
+if __name__ == "__main__":
+    main()