"""

Proposed PL-Seg training

"""

import argparse

import logging

import os

import tqdm

import random

import time

import torch

import numpy as np

import torch.optim as optim

from torchvision import transforms

from torch import nn

from torch.nn import DataParallel

import SimpleITK as sitk

from dataloader.Dataloader3d import AbdomenOrgan

from torchvision.utils import make_grid

from torch.utils.data import DataLoader

from torch.autograd import Variable

from models.Proposed.unet3d_ProgCon import UNet3D_ProgCon

from models.Proposed.sing_class_loss import DFLoss, HADFLoss, HADFLossHybrid

from models.weight_init import initialize_weights

from dataloader import transforms as tr

from utils.average_meter import AverageMeter

from utils.evaluation_seg import *

from utils.losses import *

from utils.ramps import *

from utils.val3D import test_single_case_fourpre

from tensorboardX import SummaryWriter

# -------------------- reproduction ------------------------ #

# torch.cuda.current_device() # no problem in server

# torch.cuda._initialized = True # no problem in server

torch.manual_seed(42)

torch.cuda.manual_seed(42)

torch.cuda.manual_seed_all(42)

np.random.seed(42) # Numpy module.

random.seed(42) # Python random module.

torch.manual_seed(42)

torch.backends.cudnn.benchmark = False

torch.backends.cudnn.deterministic = True

torch.backends.cudnn.enabled = True

torch.set_default_tensor_type('torch.FloatTensor')

def get_current_consistency_weight(epoch, max_epoch):

# Consistency ramp-up from https://arxiv.org/abs/1610.02242

return sigmoid_rampup(epoch, max_epoch)

def parse_args():

desc = "Pytorch implementation of PL-Seg"

parser = argparse.ArgumentParser(description=desc)

# dir config

parser.add_argument('--exp_dir', type=str, default='./exp/WORD/PLSeg')

parser.add_argument('--data_dir', type=str, default='./datasets/WORD')

parser.add_argument('--workspace', type=str, default='./exp/WORD/PLSeg/checkpoint')

# GPU config

parser.add_argument('--gpu', type=str, default='0')

parser.add_argument('--gpu_grop', type=int, default=[0, 1])

# training config

parser.add_argument('--resume', default=None, help='checkpoint path')

parser.add_argument('--in_channel', type=int, default=1)

parser.add_argument('--batch_size', type=int, default=4)

parser.add_argument(

"--patch_size",

type=int,

nargs="+",

required=True,

help="WORD [128, 128, 96]; FLARE2023[128, 128, 64]"

)

parser.add_argument('--num_classes', type=int, default=17)

parser.add_argument('--start_epoch', type=int, default=0)

parser.add_argument('--epoches', type=int, default=500)

parser.add_argument('--learning_rate', type=float, default=0.01)

parser.add_argument('-wi', '--weight_init', type=str, default="xavier",

help='Weight initialization method, or path to weights file'

'(for fine-tuning or continuing training)')

parser.add_argument('--print_interval', type=int, default=1)

parser.add_argument('--val_interval', type=int, default=2)

parser.add_argument('--save_interval', type=int, default=25)

# Ours

parser.add_argument('--consistency', type=float, default=0.1)

parser.add_argument('--temperature', type=float, default=10)

parser.add_argument('--alpha', type=float, default=0.1)

return parser.parse_args()

def validate_slice(model, dataloader, args, writer, epoch):

training = model.training

model.eval()

val_dice = AverageMeter()

with torch.no_grad():

for sample in tqdm.tqdm(dataloader, total=len(dataloader), ncols=80, leave=False):

image = Variable(sample['image'].squeeze(dim=0).squeeze(dim=0).cuda())

target = sample['label'].cuda()

target = Variable(target)

pred_seg = test_single_case_fourpre(model, image, stride_xy=args.patch_size[0], stride_z=args.patch_size[2],

patch_size=args.patch_size, num_classes=args.num_classes)

gt_volumn = target.squeeze(dim=0)

if args.num_classes == 17:

dice_score = get_multi_class_evaluation_score(pred_seg, gt_volumn.cpu().numpy(),

[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16], False, 'dice')

elif args.num_classes == 13:

dice_score = get_multi_class_evaluation_score(pred_seg, gt_volumn.cpu().numpy(),

[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12], False, 'dice')

else:

raise ValueError(f"Unknown dataset")

val_dice.update(torch.tensor(dice_score))

if (epoch + 1) % (args.print_interval + 19) == 0:

image_v = image.unsqueeze(0).permute(0, 2, 3, 1)

image_v = image_v[0:1, :, :, 20:61:10].permute(3, 0, 1, 2).repeat(1, 3, 1, 1)

grid_image = make_grid(image_v, 5, normalize=True)

writer.add_image('train/Image', grid_image, epoch)

seg_pred_v = torch.from_numpy(pred_seg).permute(1, 2, 0)

seg_pred_v = seg_pred_v[:, :, 20:61:10].permute(2, 0, 1)

pre_v = seg_pred_v.unsqueeze(dim=1).repeat(1, 3, 1, 1)

grid_image = make_grid(pre_v, 5, normalize=False)

writer.add_image('train/Predicted_label', grid_image, epoch)

label_v = gt_volumn.permute(1, 2, 0)

label_v = label_v[:, :, 20:61:10].permute(2, 0, 1)

label_v = label_v.unsqueeze(dim=1).repeat(1, 3, 1, 1)

grid_image = make_grid(label_v, 5, normalize=False)

writer.add_image('train/Groundtruth_label', grid_image, epoch)

if training:

model.train()

return val_dice

def train(model, train_loader, val_loader, writer, args):

# define hardness-aware decoupled foreground loss

# df_loss = DFLoss()

# hadfHb_loss = HADFLossHybrid(num_classes=args.num_classes)

hadf_loss_1 = HADFLoss(num_classes=args.num_classes)

hadf_loss_2 = HADFLoss(num_classes=args.num_classes)

hadf_loss_3 = HADFLoss(num_classes=args.num_classes)

hadf_loss_4 = HADFLoss(num_classes=args.num_classes)

# define the distillation loss

distill_loss = DistillationLoss(T = args.temperature)

# define the optimizer

optimizer = optim.Adam(model.parameters(), betas=(0.9, 0.99), lr=args.learning_rate)

scheduler = optim.lr_scheduler.ReduceLROnPlateau(optimizer, mode='max', factor=0.95, patience=4, verbose=True,

min_lr=1e-4)

best_dice = torch.zeros(args.num_classes)

best_epoch = 1

for epoch in range(args.start_epoch, args.epoches):

seg_loss_epoch = AverageMeter()

con_loss_epoch = AverageMeter()

mpc_loss_epoch = AverageMeter()

# train in each epoch

for batch_idx, sample in tqdm.tqdm(

enumerate(train_loader), total=len(train_loader),

desc='Train epoch=%d' % epoch, ncols=80, leave=False):

model.train()

image = sample['image'].cuda()

label = sample['onehot_label'].cuda()

cur_task = sample['cur_task'].cuda()

seg_pred_1, seg_pred_2, seg_pred_3, seg_pred_4 = model(image)

optimizer.zero_grad()

seg_pred_1_soft = torch.softmax(seg_pred_1, dim=1)

seg_pred_2_soft = torch.softmax(seg_pred_2, dim=1)

seg_pred_3_soft = torch.softmax(seg_pred_3, dim=1)

seg_pred_4_soft = torch.softmax(seg_pred_4, dim=1)

seg_pred_1_soft = torch.clamp(seg_pred_1_soft, min=1e-10, max=1)

seg_pred_2_soft = torch.clamp(seg_pred_2_soft, min=1e-10, max=1)

seg_pred_3_soft = torch.clamp(seg_pred_3_soft, min=1e-10, max=1)

seg_pred_4_soft = torch.clamp(seg_pred_4_soft, min=1e-10, max=1)

# Pyramid Partial Label Supervision

# seg_loss = hadfHb_loss(seg_pred_1_soft, label, cur_task)

seg_loss_1 = hadf_loss_1(seg_pred_1_soft, label)

seg_loss_2 = hadf_loss_2(seg_pred_2_soft, F.interpolate(label, size=seg_pred_2.shape[2:], mode='nearest'))

seg_loss_3 = hadf_loss_3(seg_pred_3_soft, F.interpolate(label, size=seg_pred_3.shape[2:], mode='nearest'))

seg_loss_4 = hadf_loss_4(seg_pred_4_soft, F.interpolate(label, size=seg_pred_4.shape[2:], mode='nearest'))

supervised_loss = (seg_loss_1 + seg_loss_2 + seg_loss_3 + seg_loss_4) / 4

# Progressive Self-Distillation

pseudo_outputs_1 = F.interpolate(seg_pred_1.detach(), size=seg_pred_2.shape[2:], mode="trilinear", align_corners=True)

pseudo_outputs_2 = F.interpolate(seg_pred_2.detach(), size=seg_pred_3.shape[2:], mode="trilinear", align_corners=True)

pseudo_outputs_3 = F.interpolate(seg_pred_3.detach(), size=seg_pred_4.shape[2:], mode="trilinear", align_corners=True)

pseudo_supervision_loss1 = distill_loss(seg_pred_2, pseudo_outputs_1)

pseudo_supervision_loss2 = distill_loss(seg_pred_3, pseudo_outputs_2)

pseudo_supervision_loss3 = distill_loss(seg_pred_4, pseudo_outputs_3)

pseudo_loss = (pseudo_supervision_loss1 + pseudo_supervision_loss2 + pseudo_supervision_loss3) / 3

consistency_weight = args.consistency * get_current_consistency_weight(epoch, args.epoches)

# Classwise Orthogonal Contrastive Regularization

COCR_loss = CO_Contrastive(seg_pred_1_soft)

total_loss = supervised_loss + consistency_weight * pseudo_loss + args.alpha * COCR_loss

seg_loss_epoch.update(supervised_loss.cpu())

con_loss_epoch.update(pseudo_loss.cpu())

mpc_loss_epoch.update(COCR_loss.cpu())

# backward the gradient

total_loss.backward()

optimizer.step()

# print training result

logging.info(

'\n Epoch[%4d/%4d]-Lr: %.6f -->for ct_array in ct_array_list: Train...' % (epoch + 1, args.epoches, optimizer.param_groups[0]['lr']))

logging.info(

'\t Seg Loss = %.4f, Con Loss = %.4f, MPC Loss = %.4f' % (seg_loss_epoch.avg, con_loss_epoch.avg, mpc_loss_epoch.avg))

# tensorboard

if (epoch + 1) % args.print_interval == 0:

writer.add_scalar('Learning_Rate', optimizer.param_groups[0]['lr'], epoch)

writer.add_scalars('Train/Losses',{'seg': seg_loss_epoch.avg, 'con': con_loss_epoch.avg, 'mpc': mpc_loss_epoch.avg}, epoch)

# validate and visualization

if not os.path.exists(args.workspace):

os.mkdir(args.workspace)

model_dir = os.path.join(args.workspace, 'models')

if not os.path.exists(model_dir):

os.mkdir(model_dir)

if (epoch + 1) % args.val_interval == 0:

val_dice = validate_slice(model, val_loader, args, writer, epoch)

logging.info('\n Epoch[%4d/%4d] --> Valid...' % (epoch + 1, args.epoches))

if args.num_classes == 17:

logging.info(

'\t [Dice Coef: mean=%.4f, BG=%.4f, Liver=%.4f, Spleen=%.4f, LK=%.4f, RK=%.4f, Stomach=%.4f, Gallb=%.4f, Esopha=%.4f, Pancreas=%.4f, Duode=%.4f, Colon=%.4f, Intes=%.4f, Adrenal=%.4f, Rectum=%.4f, Bladder=%.4f, LH=%.4f, RH=%.4f]' %

(torch.mean(val_dice.avg), val_dice.avg[0], val_dice.avg[1], val_dice.avg[2], val_dice.avg[3], val_dice.avg[4], val_dice.avg[5], val_dice.avg[6],

val_dice.avg[7], val_dice.avg[8], val_dice.avg[9], val_dice.avg[10], val_dice.avg[11], val_dice.avg[12], val_dice.avg[13], val_dice.avg[14], val_dice.avg[15], val_dice.avg[16]))

writer.add_scalars('Val/Dice',

{'Liver': val_dice.avg[1], 'Spleen': val_dice.avg[2],

'LK': val_dice.avg[3], 'RK': val_dice.avg[4], 'Stomach' :val_dice.avg[5],

'Gallb': val_dice.avg[6], 'Esopha': val_dice.avg[7], 'Pancreas': val_dice.avg[8],

'Duode': val_dice.avg[9], 'Colon': val_dice.avg[10], 'Intes': val_dice.avg[11],

'Adrenal': val_dice.avg[12], 'Rectum': val_dice.avg[13], 'Bladder': val_dice.avg[14],

'LH': val_dice.avg[15], 'ROH': val_dice.avg[16],'BG': val_dice.avg[0],

'mean': torch.mean(val_dice.avg)}, epoch)

elif args.num_classes == 13:

logging.info(

'\t [Dice Coef: mean=%.4f, BG=%.4f, Liver=%.4f, R_Kidney=%.4f, Spleen=%.4f, Pancreas=%.4f, Aorta=%.4f, IVC=%.4f, R_AdGland=%.4f, L_AdGland=%.4f, Esophagus=%.4f, Stomach=%.4f, Duodenum=%.4f, L_Kidney=%.4f]' %

(torch.mean(val_dice.avg), val_dice.avg[0], val_dice.avg[1], val_dice.avg[2], val_dice.avg[3], val_dice.avg[4], val_dice.avg[5], val_dice.avg[6],

val_dice.avg[7], val_dice.avg[8], val_dice.avg[9], val_dice.avg[10], val_dice.avg[11], val_dice.avg[12]))

writer.add_scalars('Val/Dice',

{'Liver': val_dice.avg[1], 'R_Kidney': val_dice.avg[2],

'Spleen': val_dice.avg[3], 'Pancreas': val_dice.avg[4], 'Aorta' :val_dice.avg[5],

'IVC': val_dice.avg[6], 'R_AdGland': val_dice.avg[7], 'L_AdGland': val_dice.avg[8],

'Esophagus': val_dice.avg[9], 'Stomach': val_dice.avg[10], 'Duodenum': val_dice.avg[11],

'L_Kidney': val_dice.avg[12], 'BG': val_dice.avg[0],

'mean': torch.mean(val_dice.avg)}, epoch)

# save best model

if torch.mean(val_dice.avg) >= torch.mean(best_dice):

best_model_path = os.path.join(model_dir, 'best_model.pth')

torch.save(model.state_dict(), best_model_path)

logging.info(

'\n [Epoch[%4d/%4d] --> Dice improved from %.4f in epoch %4d to %.4f]' %

(epoch + 1, args.epoches, torch.mean(best_dice), best_epoch, torch.mean(val_dice.avg)))

best_dice, best_epoch = val_dice.avg, epoch + 1

else:

logging.info('\n [Epoch[%4d/%4d] --> Dice did not improved with %.4f in epoch %d)]' %

(epoch + 1, args.epoches, torch.mean(best_dice), best_epoch))

# check for plateau

dice_sum = 0

dice_sum += torch.mean(val_dice.avg)

scheduler.step(dice_sum)

# save final model

final_model_path = os.path.join(model_dir, 'final_model.pth')

torch.save({'model': model.state_dict(), 'optim': optimizer.state_dict()}, final_model_path)

logging.info('\t [Save Final Model] to %s' % (final_model_path))

def main():

args = parse_args()

# GPU Setting

# single GPU

os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu

# GPU Parallel

# os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"

# os.environ["CUDA_VISIBLE_DEVICES"] = "2, 3"

# define logger

os.makedirs(args.exp_dir, exist_ok=True)

logging.basicConfig(filename=os.path.join(args.exp_dir, 'train.log'), level=logging.DEBUG,

format='%(asctime)s %(message)s')

logging.getLogger().addHandler(logging.StreamHandler())

# print all parameters

for name, v in vars(args).items():

logging.info(name + ': ' + str(v))

# dataset

composed_transforms_tr = transforms.Compose([

tr.LabeledClass(args.num_classes),

tr.TrainerCrop(args.patch_size),

tr.CreateOnehotLabel(args.num_classes),

tr.ToTensor()

])

composed_transforms_ts = transforms.Compose([

tr.Test_ToTensor()

])

# dataloader config

train_set = AbdomenOrgan(nii_dir=args.data_dir, mode='train', transform=composed_transforms_tr)

train_loader = DataLoader(train_set, batch_size=args.batch_size, shuffle=True, num_workers=4, drop_last=True,

pin_memory=True)

valid_set = AbdomenOrgan(nii_dir=args.data_dir, mode='val', transform=composed_transforms_ts)

valid_loader = DataLoader(valid_set, batch_size=1, shuffle=False, num_workers=1, pin_memory=True)

# init model

params={'in_chns': args.in_channel,

'class_num': args.num_classes,

'feature_chns': [16, 32, 64, 128],

'dropout' : [0, 0, 0.1, 0.2],

'trilinear': True}

model = UNet3D_ProgCon(params).cuda()

print('parameter numer:', sum([p.numel() for p in model.parameters()]))

# GPU Parallel

# if torch.cuda.device_count() > 1:

# model = DataParallel(model, device_ids=args.gpu_grop)

if args.resume:

checkpoint = torch.load(args.resume)

pretrained_dict = checkpoint['model']

model_dict = model.state_dict()

# 1. filter out unnecessary keys

pretrained_dict = {k: v for k, v in pretrained_dict.items() if k in model_dict}

# 2. overwrite entries in the existing state dict

model_dict.update(pretrained_dict)

# 3. load the new state dict

model.load_state_dict(model_dict)

print('Resume finished!')

else:

initialize_weights(model, args.weight_init)

# summary writer config

writer = SummaryWriter(log_dir=args.exp_dir, comment=args.exp_dir.split('/')[-1])

# train

train(model, train_loader, valid_loader, writer, args)

if __name__ == '__main__':

main()