train.py

import os
import importlib
import torch
import torch.nn as nn
import torch.nn.functional as F
import json
from monai.data import DataLoader
import matplotlib
from tqdm import tqdm
from torch.utils.tensorboard import SummaryWriter
from utils.losses import  DiceLoss, WeightedMSE, spatloss
from models.affinity_generator import AffinityMapGenerator
from models.fpnseg import FPN
import numpy as np
import cv2
from utils.tools import *
torch.autograd.set_detect_anomaly(True)
matplotlib.use('Agg')


class Trainer():
    def __init__(self, config):

        self.best_valid = {}
        self.spatial_start_epoch = -1
        self.temporal_start_epoch = -1
        self.spatial_coff = 0.01
        self.temporal_coff = 0.01
 
        self.size_crop = {
        'Echo2Echo': [(128,128), (128,128)],
        'Camus2Cardiac_uda': [(272,256), (328, 256)],
        'Cardiac_uda2Camus':[(328, 256),(272,256),],
        'Cardiac_uda2Echo':[(164,128), (128,128)]
        }
 
        self.config = config
        logdir = os.path.join(config['train']["log_dir"], config['train']['backbone'], config['source']['dataset'] + '2' + config['target']['dataset'])
        self.savedir = os.path.join(config['train']["save_dir"], config['train']['backbone'], config['source']['dataset'] + '2' + config['target']['dataset'])
        self.imgdir = os.path.join(config['train']["img_dir"], config['train']['backbone'], config['source']['dataset'] + '2' + config['target']['dataset'])
        if not os.path.exists(logdir):
            os.makedirs(logdir)
        if not os.path.exists(self.savedir):
            os.makedirs(self.savedir)
        if not os.path.exists(self.imgdir):
            os.makedirs(self.imgdir)

        os.environ['CUDA_VISIBLE_DEVICES'] = config['train']['enable_GPU_id']
        self.STAff_uda = config['train']['STAff_uda']

        self.logger = logger_config(log_path = os.path.join(logdir, "batch_{}log.txt".format(str(config['train']['batch_size']))), logging_name='experiment') # todo
        self.cls = 1

        self.network = FPN([2,4,23,3], num_classes=self.cls, in_channel=1,  back_bone=config['train']['backbone']) # todo
        self.spatial_loss_crit = spatloss
    

        spatial_size_src, crop_size_src, spatial_size_tgt, crop_size_tgt = self.get_size(config['source']['dataset'], config['target']['dataset'])
        dataset_src = getattr(importlib.import_module("datasets." + config['source']['dataset'].lower()), config['source']['dataset'])
        dataset_tgt = getattr(importlib.import_module("datasets." + config['target']['dataset'].lower()), config['target']['dataset'])
   
        if self.STAff_uda:
            self.crit = WeightedMSE().cuda()
            augset = config['source']['aug']
            self.affinity_gen = AffinityMapGenerator(shifts = config['train']['shifts'], neighbor = config['train']['neighbor'])
            self.train_dataset_src = dataset_src(**config['source']['config'], stage = 'train',  spatial_size = spatial_size_src, crop_size = crop_size_src,  affinity_gen = self.affinity_gen,  aug = augset)
        else:
            self.train_dataset_src = dataset_src(**config['source']['config'], stage = 'train',  spatial_size = spatial_size_src, crop_size = crop_size_src, )

        self.train_dataset_tgt = dataset_tgt(**config['target']['config'], stage  = 'train',   spatial_size = spatial_size_tgt, crop_size = crop_size_tgt,  )

        self.val_dataset_src  = dataset_src(**config['source']['config'], stage = 'val',  spatial_size = spatial_size_src, crop_size = crop_size_src, )  
        self.val_dataset_tgt  = dataset_tgt(**config['target']['config'], stage = 'val',  spatial_size = spatial_size_tgt, crop_size = crop_size_tgt, ) 

        self.test_dataset_src  = dataset_src(**config['source']['config'], stage = 'test',  spatial_size = spatial_size_src, crop_size = crop_size_src, )  
        self.test_dataset_tgt  = dataset_tgt(**config['target']['config'], stage = 'test',  spatial_size = spatial_size_tgt, crop_size = crop_size_tgt, ) 


        self.optimizer_dict, self.scheduler_dict = dict(), dict()
        self.optimizer_dict['Net'] = set_optimizer(self.network, config['net']['opt'])
        config['net']['sch']['STEPS'] = [config['train']['num_epochs'] * 0.75]
        self.scheduler_dict['Net'] = set_scheduler(self.optimizer_dict['Net'], config['net']['sch'])
        self.network = self.network.cuda()


        self.bce_loss = nn.BCEWithLogitsLoss(reduction="mean").cuda()
        self.ce_loss  = nn.CrossEntropyLoss().cuda()
        self.dice_loss = DiceLoss().cuda()


        self.print_allow = True 
        self.sampler = None
        
        self.train_loader_src = DataLoader(self.train_dataset_src, batch_size=config['train']['batch_size'], shuffle=True, num_workers=config['train']['num_workers'], drop_last=True) 
        self.train_loader_tgt = DataLoader(self.train_dataset_tgt, batch_size=config['train']['batch_size'], shuffle=True, num_workers=config['train']['num_workers'], drop_last=True)
        
      
        self.val_loader_src  = DataLoader(self.val_dataset_src, batch_size=1, shuffle=False, num_workers=config['train']['num_workers']) 
        self.val_loader_tgt  = DataLoader(self.val_dataset_tgt, batch_size=1, shuffle=False, num_workers=config['train']['num_workers']) 
        self.test_loader_tgt  = DataLoader(self.test_dataset_tgt, batch_size=1, shuffle=False, num_workers=config['train']['num_workers']) 
      
        
        if self.print_allow:
            self.writer = SummaryWriter(logdir)
    
    def get_size(self, source_dataset, target_dataset):
        transform_size = self.size_crop["{}2{}".format(source_dataset,target_dataset)]
        return transform_size[0][0],transform_size[0][1], transform_size[1][0], transform_size[1][1], 
        

    def train(self):
        count = 0
        losses = {}
        self.logger.info(str(self.config))

        for self.epoch in range(self.config['train']['num_epochs']):
            if self.print_allow:
                print('Start Epoch / Total Epoch: {} / {}'.format(self.epoch, self.config['train']['num_epochs']))

            train_loader_tgt = iter(self.train_loader_tgt)
            self.network.train()

            progress_bar = tqdm(self.train_loader_src) if self.print_allow else self.train_loader_src 
      
            for self.step, (img_src, mask_src, affinity_src) in enumerate(progress_bar):

                for name in self.optimizer_dict:
                    self.optimizer_dict[name].zero_grad()
                
                try:
                    img_tgt, _, _ = next(train_loader_tgt)

                except StopIteration:
                    train_loader_tgt = iter(self.train_loader_tgt)
                    img_tgt, _, _ = next(train_loader_tgt)
                

                b, c, f, h, w = img_src.shape
                img_src = img_src.cuda()
                mask_src = mask_src.cuda().reshape(b,-1,h,w)/ 1.0 
                img_tgt = img_tgt.cuda()  
            
                logits_src, feat_src, _ = self.network(img_src.reshape(-1, c, h, w))
                logits_tgt, _, enc_feat_tgt = self.network(img_tgt.reshape(-1, c, h, w))
                logits_src = logits_src.reshape(b,-1,h,w)                
                seg_loss =  (self.dice_loss(logits_src, mask_src) + self.bce_loss(logits_src, mask_src)) / 2
                losses.update({'seg_loss': seg_loss})


                if self.STAff_uda and self.epoch > self.spatial_start_epoch:
                    spatial_loss= self.spatial_loss_crit(feat_src, affinity_src, self.crit,  self.affinity_gen) * self.spatial_coff
                    losses.update({'spatial_loss': spatial_loss})
                    

                if self.STAff_uda and self.epoch > self.temporal_start_epoch:
                    b, c, f, h, w = img_tgt.shape
                    logits_tgt = logits_tgt.reshape(b,f,h,w)
                    offsets = self.affinity_gen.offsets_inter
                    feat_tgt_layer = enc_feat_tgt[0]

                    h, w = feat_tgt_layer.shape[-2:]
                    
                    feat_tgt_layer = feat_tgt_layer.reshape(b, f, -1, h, w)
                    logits_tgt = F.interpolate(logits_tgt, size=(h, w), mode='bilinear')


                    prob_tgt = nn.Sigmoid()(logits_tgt)
                    prob1 = prob_tgt[:,::2,:].reshape(b,h,w)
                    prob2 =  prob_tgt[:,1::2,:].reshape(b,h,w)

                    prob2_temp = self.affinity_gen.temporal_propagation(prob1, feat_tgt_layer[:,::2,:].reshape(b,-1,h,w),feat_tgt_layer[:,1::2,:].reshape(b,-1,h,w), offsets[:9])
                    
                    propagate_lables = self.extract_largest_connected_component(prob2_temp > 0.5).float()

                    temporal_loss = (
                        self.bce_loss(prob2, propagate_lables)  
                    ) * self.temporal_coff
 
                    
                    losses.update({'temporal_loss': temporal_loss})
                

                total_loss = sum(loss for loss in losses.values())
                total_loss.backward(retain_graph=False)


                for name in self.optimizer_dict:
                    self.optimizer_dict[name].step()
                
                if self.print_allow:
                    add_summary(self.writer, 'train/net_loss', total_loss.sum().item(), count)
                    add_summary(self.writer, 'train/lr', self.optimizer_dict['Net'].param_groups[0]['lr'], count)
                    if  self.STAff_uda:              
                        add_summary(self.writer, 'train/spatial_loss', spatial_loss, count)
                        add_summary(self.writer, 'train/spat_temp_loss', temporal_loss, count)
                    count += 1


            for name in self.scheduler_dict:
                self.scheduler_dict[name].step()
            

            if self.print_allow:                
                self.validation(self.val_loader_src, 'Inner-Val')
                self.validation(self.val_loader_tgt, 'Target Domain - Valid')
                # self.validation(self.test_loader_tgt, 'Target Domain - Test')
                print('End Training Epoch: {}'.format(self.epoch))

        self.writer.close()



    def extract_largest_connected_component(self, pred_mask: torch.Tensor) -> torch.Tensor:
        """
        Args:
            pred_mask: Binary mask tensor with shape (b, h, w).

        Returns:
            Tensor with shape (b, h, w), keeping only the largest connected component.
        """
        b, h, w = pred_mask.shape
        pred_np = pred_mask.cpu().numpy().astype(np.uint8)
        output = np.zeros_like(pred_np)
        for i in range(b):
            mask = pred_np[i]
            num_labels, labels, stats, centroids = cv2.connectedComponentsWithStats(mask, connectivity=8)
            if num_labels <= 1:
                # No connected component found; keep the original mask.
                output[i] = mask
                continue
            # stats[:, cv2.CC_STAT_AREA] stores the area of each connected component.
            # Find the largest connected component, excluding background label 0.
            largest_label = 1 + np.argmax(stats[1:, cv2.CC_STAT_AREA])
            largest_cc = (labels == largest_label).astype(np.uint8)
            output[i] = largest_cc
        return torch.from_numpy(output).to(pred_mask.device).type_as(pred_mask)



    def validation(self, datasets, dataset_type):
        count, pred_frames_list, masks_list, frames = 0, [], [],[]
        self.network.eval()
        with torch.no_grad():    
            progress_bar = tqdm(datasets) if self.print_allow else datasets
            for self.step, (imgs, masks, _, _) in enumerate(progress_bar):
                imgs = imgs.float().cuda()
                b,c,f,h,w = imgs.shape
                imgs = imgs.reshape(-1, c, h, w)
                masks = masks.cuda().reshape(-1, c, h, w)/1.0

                pred_frames, _, _ = self.network(imgs)
                loss = self.bce_loss(pred_frames, masks)

                pred_frames_list.append(pred_frames) 
                masks_list.append(masks)
                frames.append(imgs)
                count += 1

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

            if self.print_allow:
                if count == len(progress_bar):
    
                    pred_result = torch.where(nn.Sigmoid()(pred_frames_list) > 0.5, 1, 0)
                    pixel_acc, dice, precision, specificity, recall = self._calculate_overlap_metrics(masks_list, pred_result)

                print_info = '------Validation Result . {dataset_type} in |{current_epoch}/{total_epoch}| ------\n \
                       Loss : {loss:.4f} \
                       Pixel Acc : {pixel_acc:.4f} \
                       Dice : {dice:.4f} \
                       Precision : {pre:.4f} \
                       Specificity : {specificity:.4f} \
                       Recall : {recall:.4f}'.\
                       format(dataset_type=dataset_type, current_epoch=self.epoch, total_epoch=self.config['train']['num_epochs'],
                              loss=loss.item(), pixel_acc=pixel_acc, dice=dice, pre=precision, specificity=specificity, recall=recall)
                self.logger.info(print_info)
                add_summary(self.writer, dataset_type + "/dice", dice, self.epoch)

                if dataset_type in self.best_valid:
                    if self.best_valid[dataset_type] < dice:
                        self.best_valid[dataset_type] = dice
                        self.save(dataset_type)
                    
                else:
                    self.best_valid[dataset_type] = dice
                    self.save(dataset_type)
    
    def load(self):
        model_path = self.config['train']['save_dir']
        if os.path.isfile(os.path.join(model_path, 'latest.ckpt')):
            latest_epoch = open(os.path.join(
                model_path, 'latest.ckpt'), 'r').read().splitlines()[-1]
        else:
            ckpts = [os.path.basename(i).split('.pth')[0] for i in glob.glob(
                os.path.join(model_path, '*.pth'))]
            ckpts.sort()
            latest_epoch = ckpts[-1] if len(ckpts) > 0 else None

        if latest_epoch is not None:
            net_path = os.path.join(
                model_path, 'net_{}.pth'.format(str(latest_epoch).zfill(5)))
            print('Loading model from {}...'.format(net_path))

            data = torch.load(net_path, map_location=self.device)
            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'])

        else:
            print('Warnning: There is no trained model found. An initialized model will be used.')
    

    def calculate_overlap_metrics(self, gt, pred):
        metrics = [self._calculate_overlap_metrics(gt[i], pred[i]) for i in range(len(gt))]
        results = [torch.stack(m).mean().item() for m in zip(*metrics)]
        return tuple(results)  # pixel_acc, dice, precision, specificity, recall
            


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

        tp = torch.sum(output * target)  # TP
        fp = torch.sum(output * (1 - target))  # FP
        fn = torch.sum((1 - output) * target)  # FN
        tn = torch.sum((1 - output) * (1 - target))  # TN

        pixel_acc = (tp + tn + eps) / (tp + tn + fp + fn + eps)
        dice = (2 * tp + eps) / (2 * tp + fp + fn + eps)
        precision = (tp + eps) / (tp + fp + eps)
        recall = (tp + eps) / (tp + fn + eps)
        specificity = (tn + eps) / (tn + fp + eps)

        return pixel_acc, dice, precision, specificity, recall 

    def save(self, it):
        opt_path = os.path.join(
            self.savedir, 'opt_{}.pth'.format(it))
        print('\nsaving model to {} ...'.format(opt_path))
        if isinstance(self.network, torch.nn.DataParallel) or isinstance(self.network, torch.utils.data.distributed.DistributedSampler):
            network = self.network.module
        else:
            network = self.network
        torch.save({'network': network.state_dict()}, opt_path)


if __name__ == "__main__":
    from utils.tools import set_seed 
    from pathlib import Path
    basedir = "/home/fangxinyan/public4jbhi/config/STAffEcho"
    configs = os.listdir(basedir)
    seeds = [114514, 42, 3407]

    for seed in seeds:
        set_seed(seed)
        for config in configs:
            config_path = os.path.join(basedir, config)
            with open(config_path, "r") as file:
                config = json.load(file) 
            
            config['train']["log_dir"] = str(Path(*Path(config['train']["log_dir"]).parts[:-3], 'seed{}'.format(seed), *Path(config['train']["log_dir"]).parts[-2:]))
            config['train']["save_dir"] = str(Path(*Path(config['train']["save_dir"]).parts[:-3], 'seed{}'.format(seed), *Path(config['train']["save_dir"]).parts[-2:]))
            config['train']["img_dir"] = str(Path(*Path(config['train']["img_dir"]).parts[:-3], 'seed{}'.format(seed),  *Path(config['train']["img_dir"]).parts[-2:]))

            Train_ = Trainer(config)
            Train_.train()