all_UTELT_flip.py

'''
2024.5.9: consider the consistency of transformed images (image flip)
'''
import argparse
import time
import os
import torch
import torch.nn as nn
import torch.optim as optim
import torch.backends.cudnn as cudnn
from utils.tools import *
from tensorboardX import SummaryWriter
from metrics import AverageMeter
from tqdm import tqdm
from allconfig import get_arguments, get_model, get_train_testloader, init_seeds
from treefilter.tree_energy_loss import CEandTreeEnergyLoss
from utils import ramps
import copy

def main():
    # Fixed
    cudnn.enabled = True
    cudnn.benchmark = True
    init_seeds()

    # Setup parameters
    args = get_arguments()
    model = get_model(args)
    src_loader, test_loader = get_train_testloader(args)

    # Setup writers
    f = open(args.snapshot_dir + f'{args.city}Seg_log.txt', 'w')
    # save args
    argsDict = args.__dict__
    f.writelines('------------------ start ------------------' + '\n')
    for eachArg, value in argsDict.items():
        f.writelines(eachArg + ' : ' + str(value) + '\n')
    f.writelines('------------------- end -------------------')
    f.flush()

    writer = SummaryWriter(logdir=args.snapshot_dir)

    w, h = map(int, args.input_size_test.split(','))
    input_size_test = (w, h)
    w, h = map(int, args.input_size_train.split(','))
    input_size_train = (w, h)

    # resume
    init_iter_index = 0
    resume = os.path.join(args.snapshot_dir, f'{args.city}_batch_checkpoint.pth')
    if os.path.exists(resume):
        print('restore weight')
        resume_weight = torch.load(resume)
        model.load_state_dict(resume_weight['state_dict'])
        init_iter_index = resume_weight['batch_index']

    model.train()
    model = model.cuda()

    optimizer = optim.SGD(model.parameters(),
                          lr=args.learning_rate, momentum=args.momentum, weight_decay=args.weight_decay)
    optimizer.zero_grad()
    
    # interpolation for the probability maps and labels 
    interp_train = nn.Upsample(size=(input_size_train[1], input_size_train[0]), mode='bilinear')
    interp_test = nn.Upsample(size=(input_size_test[1], input_size_test[0]), mode='bilinear')

    loss_hist = [AverageMeter() for _ in range(5)]  # np.zeros((args.num_steps_stop,5))
    F1_best = 0.6
    iter_best = 0

    L_seg = CEandTreeEnergyLoss(ignore_index=255, sigma=args.tel_sigma)
    L_con = nn.MSELoss()
    L_ce = nn.CrossEntropyLoss(ignore_index=255)

    pbar = tqdm(range(args.num_steps_stop), disable=False)
    for batch_index, src_data in enumerate(src_loader):
        if batch_index==args.num_steps_stop:
            break
        tem_time = time.time()
        model.train()
        optimizer.zero_grad()
        
        lr = adjust_learning_rate(optimizer,args.learning_rate,batch_index,args.num_steps
                                  , power=args.learning_power)

        images, labels, ori_img, croppings = src_data
        #images = images.cuda()
        pb_ori = model(images.cuda())

        # Segmentation Loss
        labels = labels.cuda().long()
        ori_img = (ori_img.float()/255.0).cuda()
        croppings =croppings.cuda()

        # TODO: change log(2) to log(num_classes)
        # calculate uncertainty
        preds =  nn.functional.interpolate(pb_ori[0], size=(h, w), mode='bilinear', align_corners=True)
        preds = torch.softmax(preds, dim=1) # N C H W
        uncertainty = -1.0 * \
                      torch.sum(preds * torch.log(preds + 1e-6), dim=1) # , keepdim=True)
        threshold = (args.Tuncert + (1-args.Tuncert) * ramps.sigmoid_rampup(batch_index,
                                                        args.num_steps)) * np.log(args.num_classes)
        mask = (uncertainty<threshold) #.float()
        croppings = (mask & croppings)

        L_seg_value = L_seg(pb_ori[0], ori_img, pb_ori[1], croppings, labels) # preds, low_feats, high_feats, unlabeled_ROIs, target

        # flip consistency
        flip_dim = np.random.randint(2,4) # N C H W
        images_flip = torch.flip(images, dims=[flip_dim])
        pred_flip = model(images_flip.cuda())[0]
        pred_flip = nn.functional.interpolate(pred_flip.flip(dims=[flip_dim]), size=(h, w), mode='bilinear', align_corners=True)
        pred_flip_prob = torch.softmax(pred_flip, dim=1)
        L_con_value = L_con(preds, pred_flip_prob)
        L_ce_value = L_ce(pred_flip, labels)

        total_loss = L_seg_value + L_con_value + L_ce_value
        total_loss.backward()
        optimizer.step()

        pb_output_pred = nn.functional.interpolate(pb_ori[0], size=(h, w), mode='bilinear', align_corners=True)
        _, predict_labels = torch.max(pb_output_pred, 1)
        lbl_pred = predict_labels.detach().cpu().numpy()
        lbl_true = labels.detach().cpu().numpy()
        metrics_batch = []
        for lt, lp in zip(lbl_true, lbl_pred):
            _,_,mean_iu,_ = label_accuracy_score(lt, lp, n_class=args.num_classes)
            metrics_batch.append(mean_iu)
        miou = np.nanmean(metrics_batch, axis=0)

        batch_size = images.shape[0]
        loss_hist[0].update(L_seg_value.item(), batch_size)
        loss_hist[1].update(L_con_value.item(), batch_size)
        loss_hist[2].update(L_ce_value.item(),  batch_size)
        loss_hist[3].update(miou, batch_size)
        loss_hist[4].update(total_loss.item(), batch_size)

        if (batch_index+1) % 10 == 0: 
            #print('Iter %d/%d time: %.2f miou = %.1f L_seg = %.3f L_exp = %.3f L_con = %.3f'%(batch_index+1,args.num_steps,np.mean(loss_hist[batch_index-9:batch_index+1,-1]),np.mean(loss_hist[batch_index-9:batch_index+1,3])*100,np.mean(loss_hist[batch_index-9:batch_index+1,0]),np.mean(loss_hist[batch_index-9:batch_index+1,1]),np.mean(loss_hist[batch_index-9:batch_index+1,2])))
            f.write('Iter %d/%d Loss = %.3f miou = %.1f  L_seg = %.3f L_con = %.3f L_ce = %.3f\n'%
                    (batch_index+1,args.num_steps,loss_hist[4].avg, loss_hist[3].avg*100,loss_hist[0].avg,loss_hist[1].avg,loss_hist[2].avg))
            f.flush()

        pbar.set_description(
            'Train Iter:{batch:4}|{iter:4}. Loss {loss:.3f}. miou {miou:.3f}. Lseg {Lseg:.3f}. Lcon {Lcon:.3f}. Lce {Lce:.3f}.'.format(
                batch=batch_index, iter=args.num_steps_stop, loss=loss_hist[4].avg, miou=loss_hist[3].avg,
                Lseg=loss_hist[0].avg, Lcon=loss_hist[1].avg, Lce=loss_hist[2].avg))
        pbar.update()
        writer.add_scalar('lr', lr, batch_index)
        writer.add_scalar('train/loss', loss_hist[4].avg, batch_index)
        writer.add_scalar('train/miou', loss_hist[3].avg, batch_index)
        writer.add_scalar('train/lseg', loss_hist[0].avg, batch_index)
        writer.add_scalar('train/lcon', loss_hist[1].avg, batch_index)
        writer.add_scalar('train/lce', loss_hist[2].avg, batch_index)

        # evaluation per 100 iterations
        if (batch_index+1) % 100 == 0:            
            model.eval()
            TP_all = np.zeros((args.num_classes, 1))
            FP_all = np.zeros((args.num_classes, 1))
            TN_all = np.zeros((args.num_classes, 1))
            FN_all = np.zeros((args.num_classes, 1))
            n_valid_sample_all = 0
            F1 = np.zeros((args.num_classes, 1))
            IoU = np.zeros((args.num_classes, 1))
        
            for index, batch in enumerate(test_loader):  
                image, label,_, name = batch
                label = label.squeeze().numpy()

                img_size = image.shape[2:] 
                block_size = input_size_test
                min_overlap = 40

                # crop the test images into 128×128 patches
                y_end,x_end = np.subtract(img_size, block_size)
                x = np.linspace(0, x_end, int(np.ceil(x_end/np.float64(block_size[1]-min_overlap)))+1, endpoint=True).astype('int')
                y = np.linspace(0, y_end, int(np.ceil(y_end/np.float64(block_size[0]-min_overlap)))+1, endpoint=True).astype('int')

                test_pred = np.zeros(img_size)
                 
                for j in range(len(x)):    
                    for k in range(len(y)):            
                        r_start,c_start = (y[k],x[j])
                        r_end,c_end = (r_start+block_size[0],c_start+block_size[1])
                        image_part = image[0,:,r_start:r_end, c_start:c_end].unsqueeze(0).cuda()

                        with torch.no_grad():
                            pb = model(image_part)

                        # _,pred = torch.max(interp_test(nn.functional.softmax(pb,dim=1)+nn.functional.softmax(pe,dim=1)).detach(), 1)
                        pred = torch.argmax(
                            interp_test(nn.functional.softmax(pb, dim=1)).detach(),
                            1)
                        pred = pred.squeeze().data.cpu().numpy()
                        
                        
                        if (j==0)and(k==0):
                            test_pred[r_start:r_end, c_start:c_end] = pred
                        elif (j==0)and(k!=0):
                            test_pred[r_start+int(min_overlap/2):r_end, c_start:c_end] = pred[int(min_overlap/2):,:]
                        elif (j!=0)and(k==0):
                            test_pred[r_start:r_end, c_start+int(min_overlap/2):c_end] = pred[:,int(min_overlap/2):]
                        elif (j!=0)and(k!=0):
                            test_pred[r_start+int(min_overlap/2):r_end, c_start+int(min_overlap/2):c_end] = pred[int(min_overlap/2):,int(min_overlap/2):]
            
                
                #print(index+1, '/', len(test_loader), ': Testing ', name)

                # evaluate one image
                TP,FP,TN,FN,n_valid_sample = eval_image(test_pred.reshape(-1),label.reshape(-1),args.num_classes)
                TP_all += TP
                FP_all += FP
                TN_all += TN
                FN_all += FN
                n_valid_sample_all += n_valid_sample

            OA = np.sum(TP_all)*1.0 / n_valid_sample_all
            for i in range(args.num_classes):
                P = TP_all[i]*1.0 / (TP_all[i] + FP_all[i] + args.epsilon)
                R = TP_all[i]*1.0 / (TP_all[i] + FN_all[i] + args.epsilon)
                F1[i] = 2.0*P*R / (P + R + args.epsilon)
                IoU[i] = TP_all[i]*1.0 / (TP_all[i] + FP_all[i] + FN_all[i] + args.epsilon)
            

            for i in range(args.num_classes):
                f.write('===>' + args.name_classes[i] + ': %.2f\n'%(float(F1[i]) * 100))
                print('===>' + args.name_classes[i] + ': %.2f'%(float(F1[i]) * 100))
            mF1 = np.mean(F1)
            mIoU = np.mean(IoU)
                        
            f.write('===> mean F1: %.2f mean IoU: %.2f OA: %.2f\n'%(mF1*100,mIoU*100,OA*100))
            print('===> mean F1: %.2f mean IoU: %.2f OA: %.2f'%(mF1*100,mIoU*100,OA*100))
            writer.add_scalar('test/f1', mF1, batch_index)
            writer.add_scalar('test/miou', miou, batch_index)
            writer.add_scalar('test/oa', OA, batch_index)

            # save every validation
            model_name = f'{args.city}_batch_checkpoint.pth'
            torch.save({'state_dict': model.state_dict(),
                        'batch_index': batch_index+1}, os.path.join(
                args.snapshot_dir, model_name))

            if mF1>F1_best:
                # save the current models
                f.write('Save Model\n')
                print('Save Model')
                model_name = f'{args.city}_batch'+repr(batch_index+1)+'mF1_'+repr(int(mF1*10000))+'.pth'
                torch.save(model.state_dict(), os.path.join(
                    args.snapshot_dir, model_name))
                # delete the previous weights
                oldfile = os.path.join(args.snapshot_dir, f'{args.city}_batch'+repr(iter_best+1)+'mF1_'+repr(int(F1_best*10000))+'.pth')
                if os.path.exists(oldfile):
                    os.remove(oldfile)
                F1_best = copy.deepcopy(mF1)
                iter_best = copy.deepcopy(batch_index)
 
    f.close()
    pbar.close()

    # save the last one
    model_name = f'{args.city}_batch' + repr(batch_index + 1) + 'mF1_' + repr(int(mF1 * 10000)) + '.pth'
    torch.save(model.state_dict(), os.path.join(
        args.snapshot_dir, model_name))


if __name__ == '__main__':
    main()