all_TELT_mix.py

'''
2024.5.9: consider the consistency of transformed images (colorjitter)
2024.5.23: set random transformation from flip, rotate, color, scale.
           set the selection mode to mix, i.e., mixing all the transformations
'''
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
# from torchvision import transforms
# color_trans = transforms.ColorJitter(0.5,0.5,0.5,0.2)
# imagenet_mean = np.array([0.485, 0.456, 0.406])
# imagenet_std = np.array([0.229, 0.224, 0.225])
# imagenet_mean_torch = torch.from_numpy(imagenet_mean).float().reshape((1, 3, 1, 1))
# imagenet_std_torch = torch.from_numpy(imagenet_std).float().reshape((1, 3, 1, 1))
import torchvision.transforms.functional as TF
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, images_trans = src_data
        pb_ori = model(images.cuda())

        # Segmentation Loss
        labels = labels.cuda().long()
        ori_img = (ori_img.float()/255.0) # .cuda()
        croppings =croppings.cuda()
        L_seg_value = L_seg(pb_ori[0], ori_img.cuda(), pb_ori[1], croppings, labels) # preds, low_feats, high_feats, unlabeled_ROIs, target

        ########################## randomly select augmentation composition ##################
        # t: transform times, p: transform list, v: transform parameters
        trans_t = np.random.randint(1, 5, size=1)[0] # [1, 4]
        trans_p = np.random.choice(['flip', 'scale', 'color', 'rotate'], size=trans_t)

        labels_trans = labels.clone().unsqueeze(1) # N 1 H W 
        
        preds = nn.functional.interpolate(pb_ori[0], size=(h, w), mode='bilinear', align_corners=True)
        preds = preds.softmax(dim=1)
        preds_ori_trans = preds.clone()

        # color transform
        if 'color' not in trans_p:
            images_trans = images.clone() # original images, grad sharing, but memory not

        # flip consistency
        if 'flip' in trans_p:
            flip_dim = np.random.randint(2,4) # N C H W
            images_trans = torch.flip(images_trans, dims=[flip_dim])
            preds_ori_trans = torch.flip(preds_ori_trans, dims=[flip_dim])
            labels_trans = torch.flip(labels_trans, dims=[flip_dim])

        # scale consistency: change scale to [0.75, 1.5]
        if 'scale' in trans_p:
            scale_factor = np.random.choice([0.75, 1.5], size=1)[0]
            images_trans = nn.functional.interpolate(images_trans, scale_factor=scale_factor, mode='bilinear', align_corners=True)

        # rotate consistency: random select angle from  [-90, 90]
        if 'rotate' in trans_p:
            angle = np.random.randint(-90, 90, size=1)[0]
            images_trans = TF.rotate(images_trans, angle=float(angle), interpolation=TF.InterpolationMode.BILINEAR, fill=0)
            # # rotate original image, labels
            preds_ori_trans = TF.rotate(preds_ori_trans, angle=float(angle), interpolation=TF.InterpolationMode.BILINEAR, fill=0)
            labels_trans = TF.rotate(labels_trans, angle=float(angle), interpolation=TF.InterpolationMode.NEAREST, fill=255)

        # calculate consistency loss
        preds_trans = model(images_trans.cuda())[0]
        preds_trans = nn.functional.interpolate(preds_trans, size=(h, w), mode='bilinear', align_corners=True)
        L_con_value = L_con(preds_trans.softmax(dim=1), preds_ori_trans)
        L_ce_value = L_ce(preds_trans, labels_trans.squeeze(1))

        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()