HMCLIP / stage1_eva_t2t.py

Start resumable backup upload: root source files

1cc4fea 15 days ago

7.17 kB

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

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