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