eval_pcbm_fetal.py

import torch
from tqdm import tqdm
import metrics
import datasets
import albumentations as A
from torch.utils.data import DataLoader
import numpy as np
import argparse
import yaml
from misc import Logger
import os
import random
from sklearn.metrics import accuracy_score, balanced_accuracy_score, classification_report, confusion_matrix
from torchsampler import ImbalancedDatasetSampler
import math
import models
import torch.nn as nn
import matplotlib.pyplot as plt
import timm
import yaml
from pathlib import PurePath

#%%
## load configs
parser = argparse.ArgumentParser(description="PCBM")
parser.add_argument('--split', type=int, default=6)
parser.add_argument('--conceiver_path', type=str, default=f'logs/conceiver6/model.t7')
parser.add_argument('--observer_path', type=str, default=f'logs/nnn6/model.t7')
parser.add_argument('--predictor_path', type=str, default=f'logs/plain_predictor6/model.t7')

args = parser.parse_args()

device = "cuda" if torch.cuda.is_available() else "cpu"

batch_size = 32

seg_channel = 14
in_channels = 1
num_classes = 8
concept_num = 27
expand_dim = 256
head_num = 0

#%%
# # fix random seed
seed = 42
torch.manual_seed(seed)
# np.random.seed(seed)
# random.seed(seed)
torch.cuda.manual_seed(seed)
torch.cuda.manual_seed_all(seed)

reg_index = [0, 1, 5, 6, 7, 9, 10, 11, 12, 14, 15, 16, 18, 19, 20, 21, 23, 25, 26]
cat_index = [2, 3, 4, 8, 13, 17, 22, 24]
global_index = [[2, 3, 4, 8, 13, 17, 22, 24]]
local_index = [
                [5,6,7,14,15,16,23,25,26], # local concept indices, group 1 (organ quality) 
                [0,1,9,10,11,12,18,19,20,21] # local concept indices, group 2, (caliper quality)
                ]
relationship = {
                    0: [14],
                    1: [15],
                    5: [6],
                    6: [8],
                    7: [9],
                    9: [16,6],
                    10: [16,6],
                    11: [16,6],
                    12: [16,6],
                    14: [10],
                    15: [12],
                    16: [11],
                    18: [18,13],
                    19: [17,13],
                    20: [20,13],
                    21: [19,13],
                    23: [4],
                    25: [1,3],
                    26: [1,2]
                }


# fetal dataset
## build datasets
tfs = []
tfs.append(A.Resize(224, 288))
eval_transforms = A.Compose(
    tfs
)
testset = datasets.FetalSeg(
                            eval_transforms, 
                            split='test',
                            csv_dir= '/home/manli/progressive-concept-bottleneck-models/metas/trim3_sp.csv', # folder containing the data or name of the csv file 
                            meta_dir= '/data/proto/Zahra_Study1_Trials/trim3_sp.yaml', # name of the meta file
                            split_index= f'split{args.split}',
                            keep_trim3=True,
                            plane=['Femur', 'Head', 'Abdomen', 'Cervix'],
                            remove_calipers=False
                            )

## build the model
observer = models.DTUNet(in_channels, seg_channel)

conceiver = models.FetalConceiver(
                                in_channels=in_channels, 
                                seg_concept_num=seg_channel+7, # 7 for caliper concepts
                                global_index=global_index, 
                                local_index=local_index, 
                                relationship=relationship
                            )
predictor = models.Predictor(
                                num_classes=num_classes, 
                                concept_num=concept_num,
                                expand_dim=expand_dim,
                                head_num=head_num,
                                cat_index=cat_index 
                                )


observer = observer.to(device)
conceiver = conceiver.to(device)
predictor = predictor.to(device)

observer.load_state_dict(torch.load(args.observer_path))
conceiver.load_state_dict(torch.load(args.conceiver_path))
predictor.load_state_dict(torch.load(args.predictor_path))

exp_name = f'pcbm_fetal_eval{args.split}'
exp_folder = os.path.join("./logs", exp_name)
model_path = os.path.join(exp_folder, "model.t7")

if not os.path.exists(exp_folder):
    os.system(f"mkdir {exp_folder}")

# initialize logger
logger = Logger(os.path.join(exp_folder, 'logs.log'), 'a')
logger.fprint(f"Start experiment {exp_name}")
logger.fprint(f'Fix random seed at {seed}')
logger.fprint("Model")
#%%
testloader = DataLoader(testset, batch_size=batch_size, shuffle=False, num_workers=np.min([batch_size, 32]))

def validate_one_epoch(loader):
    observer.eval()
    conceiver.eval()
    predictor.eval()
    preds = []
    gts = []
    logits = []
    concept_preds = []
    concept_gts = []

    for x in tqdm(loader):
        image= x['gray_image'].to(device)
        label = x['label'].to(device)
        
        with torch.no_grad():
            assign_mtx = observer({'image':image})['logit']
            # seg_mask = torch.argmax(assign_mtx, dim=1)
            # seg_mask = x['mask'].to(device)
            # assign_mtx = torch.nn.functional.one_hot(seg_mask, num_classes=seg_channel).permute(0,3,1,2)

        x['image'] = image
        x['assign_mtx'] = assign_mtx
        concept_gt = x['concept'].to(device)


        with torch.no_grad():
            x = conceiver(x)
            concept_logit = x['concept_logit']
            concept_logit[:, [2, 4, 8, 13, 17, 24]] = torch.sigmoid(concept_logit[:, [2, 4, 8, 13, 17, 24]])
            concept_logit[:, 23] = 0 # cervix should not be considered according to ISUOG
            concept_logit = concept_logit * x['concept_mask']
            x['concept_logit'] = concept_logit
            x = predictor(x)
        
        logit = x['logit']
        batch_size = image.size(0)

        pred = torch.argmax(logit, dim=1)
        preds.append(pred.detach().cpu().numpy())
        gts.append(label.cpu().numpy())
        logits.append(logit.detach().cpu().numpy())
        concept_preds.append(concept_logit.detach().cpu().numpy())
        concept_gts.append(concept_gt.cpu().numpy())

    epoch_metrics = {}
    
    preds = np.concatenate(preds, axis=0)
    gts = np.concatenate(gts, axis=0)
    concept_preds = np.concatenate(concept_preds, axis=0)
    concept_gts = np.concatenate(concept_gts, axis=0)
    logits = np.concatenate(logits, axis=0)

    epoch_metrics['acc'] = accuracy_score(gts, preds)
    epoch_metrics['avg_acc'] = balanced_accuracy_score(gts, preds)

    np.save('pred.npy', preds)
    np.save('gt.npy', gts)
    np.save('concept_preds.npy', concept_preds)
    np.save('concept_gts.npy', concept_gts)
    np.save('logits.npy', logits)

    # print(classification_report(gts, preds, zero_division=0)) 
    print(confusion_matrix(gts, preds))   

    return epoch_metrics

def test():
    val_metrics = validate_one_epoch(testloader)
    log_info = f"Test on Testset"
    for k, v in val_metrics.items():
        if k == 'confusion_matrix' or k == "classification_report":
            log_info += f", eval_{k}: {v}"
        else:
            log_info += f", eval_{k}: {v: .4f}"
    logger.fprint(log_info)  

if __name__ == "__main__":
    test()