infer_geometrics.py

import os
import importlib
import torch
from monai.data import DataLoader
import nrrd
import matplotlib
import numpy as np
from tqdm import tqdm
from models.fpnseg import FPN
torch.autograd.set_detect_anomaly(True)
matplotlib.use('Agg')
import pandas as pd


def save_dict_to_excel(data_dict, excel_path):

    def make_sheet_name(seed, backbone, metric):
        name = f"seed{seed}_backbone_{backbone}_{metric}"
        # Replace invalid Excel sheet-name characters.
        for ch in ['\\', '/', '*', '[', ']', ':', '?']:
            name = name.replace(ch, '_')
        return name[:31]  # Excel limits sheet names to 31 characters.

    os.makedirs(os.path.dirname(excel_path), exist_ok=True)

    with pd.ExcelWriter(excel_path, engine='openpyxl') as writer:
        sheet_count = 0
        for (seed, backbone), metrics_dict in data_dict.items():
            for metric, models_dict in metrics_dict.items():
                sheet_name = make_sheet_name(seed, backbone, metric)
                
                # Build a DataFrame where each model is one column.
                df = pd.DataFrame(dict(models_dict))  # Align lengths automatically and pad missing values with NaN.
                
                df.to_excel(writer, sheet_name=sheet_name, index=False)
                sheet_count += 1

    print(f"Successfully saved {sheet_count} worksheets to: {excel_path}")


class Test():
    def __init__(self, config):
        self.config = config
        self.model_name = config['model_name']
         # load test_dataset
        test_dataset = getattr(importlib.import_module("datasets." + config['dataset']['name'].lower()), config['dataset']['name'])
        
        self.test_target_dataset  = test_dataset(**config['dataset']['config'], stage = 'test') 
        # load model
        self.network = FPN([2,4,23,3], num_classes=1, in_channel=1,  back_bone=config['backbone']) # todo
        self.print_allow = True



    def inference(self):
        self.test_target_loader = DataLoader(self.test_target_dataset, batch_size=1, shuffle=False, num_workers=self.config['num_workers'])

        pred_frames_list, masks_list, imgs_list = [], [], []


        self.load()
        self.network.cuda()
        self.network.eval()

        with torch.no_grad():
            progress_bar = tqdm(self.test_target_loader)
            for self.step, (imgs, masks, _, _) in enumerate(progress_bar):
                imgs = imgs.float().cuda()
                masks = masks.cuda() / 1.0
                b, c, f, h, w = imgs.shape

                logits, _, _ = self.network(imgs.reshape(-1, c, h, w))

                pred_frames_list.append(logits.reshape(-1, h, w))
                masks_list.append(masks.reshape(-1, h, w))
                imgs_list.append(imgs.reshape(-1, h, w))



            pred_frames_list = torch.cat(pred_frames_list, dim=0)
            masks_list = torch.cat(masks_list, dim=0)


            pred_seg = torch.where(torch.sigmoid(pred_frames_list) > 0.5, 1, 0)

            hd95 = self.compute_hausdorff_distance_95(
                masks_list.cpu().numpy(),
                pred_seg.cpu().numpy(),)
            
            dice = self.compute_dice(masks_list.cpu().numpy(), pred_seg.cpu().numpy())

        return {
            "hd95": hd95,
            "dice": dice,
        }
    
    def compute_dice(self, gt, pred):
        b = gt.shape[0]
        dice = []
        for i in range(b):
            dice.append(self._calculate_overlap_metrics(gt[i], pred[i]))
        return dice


    def _calculate_overlap_metrics(self, gt, pred, eps=1e-5):
        output = pred.reshape(-1, )
        target = gt.reshape(-1, )

        tp = np.sum(output * target)  # TP
        fp = np.sum(output * (1 - target))  # FP
        fn = np.sum((1 - output) * target)  # FN
        tn = np.sum((1 - output) * (1 - target))  # TN
        dice = (2 * tp + eps) / (2 * tp + fp + fn + eps)

        return  dice


    def compute_hausdorff_distance_95(self, gt, pred):
        from medpy.metric.binary import hd95
        assert gt.shape == pred.shape, "Ground truth and prediction shapes must match."
        b = gt.shape[0]  # Number of images.
        hd_95s = []
        
        for i in range(b):
            if not np.any(gt[i]) or not np.any(pred[i]):
                if  np.array_equal(gt[i], pred[i]):
                    hd_95s.append(0)
                else:
                    hd_95s.append(gt.shape[1]*1.414) # max distance
            else:
                hd_95s.append(hd95(pred[i].astype(bool), gt[i].astype(bool)))
        return hd_95s


    
    def load(self):
        net_path = self.config['save_dir']
        data = torch.load(net_path)
        data['network'] = {k.replace('module.', ''): v for k, v in data['network'].items() if k.replace('module.', '') in self.network.state_dict()}
        self.network.load_state_dict(data['network'])

    
    def save(self, volume, save_path):
        volume = volume.cpu().numpy()
        nrrd.write(save_path, volume)
        
def main(config):
    dataset_setting = 'Cardiac_uda2Echo'
    # dataset_setting = 'Cardiac_uda2Camus'
    # dataset_setting = 'Camus2Cardiac_uda'

    excel_path = os.path.join("/Share8/fangxinyan/STAffEcho_jn_results/geo-metrics/", dataset_setting + '.xlsx')

    seeds = [42, 3407, 114514]

    model_names = ['woDA','cycleGAN', 'fda', 'histomatch', 'histo_eqa','ASANet', 'TIST','sepico','TPS_finetune', 'GraphEcho','SLCL', 'CRI-TSI', 'STAffEcho']

    backbones = ['VGG16','resnet50'] 

    all_results = {}

    for seed in seeds:
        set_seed(seed)
        for backbone in backbones:
            all_results[(seed, backbone)] = {
                "hd95": {},
                "dice": {}
            }

            for model_name in model_names:
                config['model_name'] = model_name
                config['backbone'] = backbone
                config['save_dir'] = os.path.join(
                    "/Share8/fangxinyan/STAffEcho_jn_results/",
                    f'seed{seed}', "model", model_name, backbone, dataset_setting, "opt_Target Domain - Valid.pth" # opt_Target Domain - Valid.pth, opt_Inner-Val.pth
                )
                Test_ = Test(config)
                results = Test_.inference()

                all_results[(seed, backbone)]["hd95"][model_name] = results["hd95"]
                all_results[(seed, backbone)]["dice"][model_name] = results["dice"]


    save_dict_to_excel(all_results, excel_path)


if __name__ == "__main__":
    from utils.tools import set_seed 
   
    os.environ['CUDA_VISIBLE_DEVICES'] = '2'

    config = {
            #    "dataset":{
            #     "name":"Camus",
            #     "config":
            #      {
            #         "root": "/Share8/fangxinyan/fangxinyan/Echo_dataset/CAMUS/",
            #         "clip_length": -1, 
            #         "sample_rate": 1
            #      }

            #    },
            
            "dataset":{
                "name":"Echo",
                "config":
                 {
                    "root": "/Share8/fangxinyan/fangxinyan/Echo_dataset/EchoNet-Dynamic/",
                    "clip_length": -1, 
                    "sample_rate": 1
                 }

               },
        

            #     "dataset":{
            #     "name":"Cardiac_uda",
            #     "config":
            #      {
            #         "root": "/Share8/fangxinyan/fangxinyan/Echo_dataset/cardiac_dataset/",
            #         "clip_length": 1, 
            #         "sample_rate": 1
            #      }

            #    },

                        "num_workers": 2,
              }
    
    main(config)