all_baseline.py

'''
2023.12.27: change base model to unetformer (wang libo), cyx
use baseline model
'''
import argparse
import time
import os
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils import data
import torch.backends.cudnn as cudnn
from utils.tools import *
import random
from tensorboardX import SummaryWriter
from metrics import AverageMeter
from tqdm import tqdm
from allconfig import get_arguments, get_model, get_train_testloader, init_seeds


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

    L_seg = 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, _, name = src_data
        images = images.cuda()      
        pb_ori = model(images)

        # Segmentation Loss
        labels = labels.cuda().long()
        L_seg_value = L_seg(pb_ori, labels)

        _, predict_labels = torch.max(pb_ori, 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)

        total_loss = L_seg_value
        total_loss.backward()
        optimizer.step()

        batch_size = images.shape[0]
        loss_hist[0].update(L_seg_value.item(), batch_size)
        loss_hist[1].update(0, batch_size)
        loss_hist[2].update(0,  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_exp = %.3f L_con = %.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}. Lexp {Lexp:.3f}. Lcon {Lcon:.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, Lexp=loss_hist[1].avg, Lcon=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/lexp', loss_hist[1].avg, batch_index)
        writer.add_scalar('train/lcon', 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:
                F1_best = mF1
                # save the 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))
 
    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()