PatchTrain.py

import torch
import torch.optim as optim
from torchvision import transforms 
from copy import deepcopy
import numpy as np
from torch.utils.data import Dataset, DataLoader
from PromptLearner import CustomCONCH, OriginCONCH, CustomKEEP, OriginKEEP
from PatchDataset import PatchCLSDataset, PtDataset
import open_clip_custom as conch_clip
import random
import utils
import params
from transformers import AutoModel, AutoTokenizer
from PIL import Image
import os
from sklearn.metrics import roc_curve, roc_auc_score, balanced_accuracy_score, classification_report, confusion_matrix
from datetime import datetime

INIT_LR = 2e-4
DEVICE = "cuda:0"

### set seed ###
SEED = 2025
random.seed(SEED)
os.environ['PYTHONHASHSEED'] = str(SEED)
np.random.seed(SEED)
torch.manual_seed(SEED)
torch.cuda.manual_seed(SEED)
# torch.cuda.manual_seed_all(SEED)
######

TEMPLATES = ['CLASSNAME.',
        'a photomicrograph showing CLASSNAME.',
        'a photomicrograph of CLASSNAME.',
        'an image of CLASSNAME.',
        'an image showing CLASSNAME.',
        'an example of CLASSNAME.',
        'CLASSNAME is shown.',
        'this is CLASSNAME.',
        'there is CLASSNAME.',
        'a histopathological image showing CLASSNAME.',
        'a histopathological image of CLASSNAME.',
        'a histopathological photograph of CLASSNAME.',
        'a histopathological photograph showing CLASSNAME.',
        'shows CLASSNAME.',
        'presence of CLASSNAME.',
        'CLASSNAME is present.',
        'an H&E stained image of CLASSNAME.',
        'an H&E stained image showing CLASSNAME.',
        'an H&E image showing CLASSNAME.',
        'an H&E image of CLASSNAME.',
        'CLASSNAME, H&E stain.',
        'CLASSNAME, H&E.'
        ]



def eval_model(model, valloader, device, zeroshot=False):
    model.to(device)
    gts = []
    preds = []
    model.eval()
    for img_embs, label in valloader:
        img_embs = img_embs.to(device)
        with torch.no_grad():
            sim = model(img_embs).squeeze()
        pred = torch.argmax(sim, dim=1)
        # 将每个batch的标签和预测结果展平并转换为numpy数组
        gts.extend(label.cpu().numpy().flatten().tolist())
        preds.extend(pred.cpu().numpy().flatten().tolist())
    
    # 转换为整数类型以确保分类正确
    gts = np.array(gts, dtype=int)
    preds = np.array(preds, dtype=int)
    
    # 检查数据维度
    if gts.ndim > 1 or preds.ndim > 1:
        gts = gts.flatten()
        preds = preds.flatten()
    
    report = classification_report(gts, preds, output_dict=True, zero_division=0)
    weighted_f1 = report['weighted avg']['f1-score']
    return weighted_f1
    


def train(model, img_encoder, loss_fn, dataloader, valloader, epoch=10, lr=INIT_LR, device="cuda:0"):
    best_val = 0
    ### Freeze all parameters except prompt_learner.ctx ###
    for name, parameter in model.named_parameters():
        if name != "prompt_learner.ctx":
            parameter.requires_grad = False

    optimizer = optim.Adam(model.parameters(), lr=lr)
    img_encoder.to(device)

    warm_up_iter = int(epoch * 0.1)
    lambda0 = lambda cur_iter: cur_iter / warm_up_iter if  cur_iter < warm_up_iter else \
            0.5 * (1.0 + np.cos(np.pi * (cur_iter - warm_up_iter) / (epoch - warm_up_iter)))
    scheduler = optim.lr_scheduler.LambdaLR(optimizer, lr_lambda=lambda0)

    for i in range(epoch):
        ### train the model ###
        model.train()
        running_loss = torch.zeros(1, device=device)
        for img, label in dataloader: #(1024, 3, 224, 224), TODO:3是RGB？
            img = img.to(device)
            
            # get image embeddings
            with torch.inference_mode():
                #img_emb = img_encoder.encode_image(img, proj_contrast=False, normalize=False)
                img_embs = img_encoder.encode_image(img)
            
            # change classnames
            #model.change_classnames()
            
            optimizer.zero_grad()
            sim = model(img_embs.squeeze()) #(1024, 4)
            label = label.to(device)

            loss = loss_fn(sim, label)
            loss.backward()
            optimizer.step()
            
            running_loss += loss.item()
        scheduler.step()
        
        ### evaluate the model ###
        model.eval()
        with torch.no_grad():
            #model.change_classnames(ens=True)
            val_result = eval_model(model, valloader, device)
        #print(f"Epoch {i} loss: {running_loss.item() / len(dataloader) :.3f}, val:{val_result.tolist()}, dist: {(loss_distribution / loss_distribution.sum()).tolist()}, cancer_ratio: {cancer_ratio.tolist()}")
        print(f"Epoch {i} loss: {running_loss.item() / len(dataloader)}, val:{val_result}")
        
        if val_result > best_val:
            timestamp = datetime.now().strftime("%y%m%d%H%M%S")
            torch.save(model.prompt_learner.ctx, os.path.join(params.SAVE_DIR, f'prompt_{timestamp}.pt'))
            print(f'best epoch find: {i}, prompt saved as prompt_{timestamp}.pt')
            best_val = val_result
            
    return model


def multiple_trains_and_eval(cfg, traindata_lst, valset, testset, ckpt_path, zeroshot_prompt_lst, classnames_lst, epoch=30, times=30, device="cuda:0", mode='binary'):
    #conch_model, preprocess = conch_clip.create_model_from_pretrained("conch_ViT-B-16", checkpoint_path=ckpt_path)
    #conch_model.to(device)
    keep_model = AutoModel.from_pretrained(ckpt_path, trust_remote_code=True)
    keep_tokenizer = AutoTokenizer.from_pretrained(ckpt_path, trust_remote_code=True)
    keep_model.to(device)
    results = []
    print(device)

    for i in range(times):
        print("Experiment No.{}".format(i))
        result_per_time = []
        trainset_lst = [PatchCLSDataset(data_dir=traindata['data_dir'], label_csv=traindata['label_csv'], divide_json=params.TRAINSET_DIVISION, aug=True, shot=traindata['shot'], dataset=traindata['dataset']) for traindata in traindata_lst]
        train_loader_lst = [DataLoader(trainset, batch_size=256, shuffle=True) for trainset in trainset_lst]
        val_loader = DataLoader(valset, batch_size=64, shuffle=False)
        test_loader = DataLoader(testset, batch_size=64, shuffle=False)
        loss_fn = torch.nn.CrossEntropyLoss()

        index = random.randint(0, len(zeroshot_prompt_lst) - 1)
        prompt = zeroshot_prompt_lst[index]
        print(prompt)
        #zero_shot_model = OriginCONCH(prompt, conch_model, device)
        zero_shot_model = OriginKEEP(prompt, keep_model, keep_tokenizer, device)
        result = eval_model(zero_shot_model, test_loader, device, zeroshot=True)
        print(f'zeroshot result{result}')
        result_per_time.append(result)
        
        unique_classnames_lst = utils.unique_classnames(classnames_lst) #原本的classnames_lst的每个元素是一种组合，unique_classnames_lst的每个元素是单类
        #template_model = CustomCONCH(cfg, unique_classnames_lst, conch_model, device, mode)
        template_model = CustomKEEP(cfg, unique_classnames_lst, keep_model, keep_tokenizer, device)

        for loader in train_loader_lst:
            model = deepcopy(template_model).to(device)
            model = train(model, keep_model, loss_fn, loader, val_loader, epoch=epoch, lr=INIT_LR, device=device)
            #model = train(model, conch_model, preprocess, loss_fn, loader, val_loader, epoch=epoch, lr=INIT_LR, device=device)

            model.eval()
            with torch.no_grad():
                #model.change_classnames(ens=True)
                result = eval_model(model, test_loader, device)
            print(f'result{result}')
            result_per_time.append(result)
            timestamp = datetime.now().strftime("%y%m%d%H%M%S")
            torch.save(model.prompt_learner.ctx, os.path.join(params.SAVE_DIR, f'prompt_{timestamp}.pt'))
            print(f'Experiment {i}, prompt saved as prompt_{timestamp}.pt')
        results.append(result_per_time)
    return results


def main_pathcamelyon_conch():
    cfg = params.PromptLearnerConfig(n_ctx=32, input_size=96)
    ckpt_path = params.CKPT_PATH
    param = params.patchcamelyon_cls
    zeroshot_prompt_path = param["zero_shot_template_file"]
    zeroshot_prompt_lst, class_lst = utils.load_prompts_from_template(zeroshot_prompt_path, classnames_example=['lymph', 'tumor'])

    print('loading trainset')
    trainset_10shot = PatchCLSDataset(data_dir=param['train_data_dir'],
                            label_csv=param['train_label_csv'],
                            aug=True,
                            shot=10
                            )
    trainset_5shot = PatchCLSDataset(data_dir=param['train_data_dir'],
                            label_csv=param['train_label_csv'],
                            aug=True,
                            shot=5
                            )
    trainset_1shot = PatchCLSDataset(data_dir=param['train_data_dir'],
                            label_csv=param['train_label_csv'],
                            aug=True,
                            shot=1
                            )
    
    print('loading testset')
    test_data = torch.load(param['test_data_file'], weights_only=True).squeeze()
    test_label = torch.load(param['test_label_file'], weights_only=True)
    print(test_data.shape, test_label.shape)
    testset = PtDataset(test_data, test_label)
    
    print('loading valset')
    valid_data = torch.load(param['valid_data_file'], weights_only=True).squeeze()
    valid_label = torch.load(param['valid_label_file'], weights_only=True)
    print(valid_data.shape, valid_label.shape)
    valset = PtDataset(valid_data, valid_label)
    
    result = multiple_trains_and_eval(cfg, [trainset_10shot, trainset_5shot, trainset_1shot], valset, testset, ckpt_path, zeroshot_prompt_lst, class_lst,
                                    epoch=30, times=30, device=DEVICE)
    
    np_result = result.cpu.numpy()
    
    print(result)
    
    np.save('patch_cls_result.npy', np_result)
    
def main_keep():
    cfg = params.PromptLearnerConfig(n_ctx=32, input_size=224)
    keep_path = params.KEEP_PATH
    dataset = 'Skincancer'
    param = params.skincancer_cls
    templates = TEMPLATES
    #class_lst = [['breast non-malignant benign tissue','breast malignant in-situ carcinoma','breast malignant invasive carcinoma','normal breast tissue']]
    #class_lst = [['lung adenocarcinoma','benign lung','lung squamous cell carcinoma','colon adenocarcinoma','benign colon']]
    #class_lst = [['red blood cells','renal cancer','non-tumor','torn adipose necrotic tissue','muscle fibrous stroma blood vessels']]
    class_lst = [['necrosis','skeletal muscle','eccrine sweat glands','vessels','elastosis','chondral tissue','hair follicle','epidermis','nerves','subcutis','dermis','sebaceous','squamous-cell carcinoma','melanoma in-situ','basal-cell carcinoma','naevus']]
    #class_lst = [['normal non-tumor','necrosis','viable tumor']]
    zeroshot_prompt_lst = []
    for prompt in templates:
        prompts = [prompt.replace('CLASSNAME', class_name) for class_name in class_lst[0]]
        zeroshot_prompt_lst.append(prompts)

    print('loading trainset')
    traindata10shot = {'data_dir': param['train_data_dir'], 'label_csv': param['label_csv'], 'shot': 10,  'dataset': dataset}
    traindata5shot = {'data_dir': param['train_data_dir'], 'label_csv': param['label_csv'], 'shot': 5,  'dataset': dataset}

    print('loading testset')
    test_data = torch.load(param['test_data_file'])
    test_feats = []
    test_labels = []
    for test_label, test_feat in test_data.items():
        test_feats.append(test_feat)
        for i in range(len(test_feat)):
            test_labels.append(test_label)
    test_feats = torch.concatenate(test_feats, dim=0)
    test_labels = torch.tensor(test_labels)
    print(test_feats.shape, test_labels.shape)
    testset = PtDataset(test_feats, test_labels)
    
    print('loading valset')
    valid_ids = torch.randperm(len(test_feats))
    valid_feats = test_feats[valid_ids[:200]]
    valid_labels = test_labels[valid_ids[:200]]
    print(valid_feats.shape, valid_labels.shape)
    valset = PtDataset(valid_feats, valid_labels)
    
    result = multiple_trains_and_eval(cfg, [traindata10shot, traindata5shot], valset, testset, keep_path, zeroshot_prompt_lst, class_lst,
                                    epoch=20, times=10, device=DEVICE, mode='multi')
    
    np_result = np.array(result)
    
    np.save(f'{dataset}_cls_{INIT_LR}.npy', np_result)

if __name__ == "__main__":
    main_keep()