allconfig.py

'''
config for general models
'''
import argparse
import os
from torch.utils import data
import torch
import random
import numpy as np

DATA_PATH_HOME = '/localssd/yinxia/data'
SAVE_WEIGHT_PATH = '/localssd/yinxia/CRGNetdata/'

def get_arguments():
    modename = ["full", "point"]
    parser = argparse.ArgumentParser(description="CRGNet")

    # dataset
    parser.add_argument("--city", type=str, default='Vaihingen',
                        help="Vaihingen, Zurich")
    # add train/val data path, 2024.5.14
    parser.add_argument("--data_dir_train", type=str, default='/home/yinxia/data',
                        help="Vaihingen, Zurich")
    parser.add_argument("--data_dir_val", type=str, default='/home/yinxia/data',
                        help="Vaihingen, Zurich")
    parser.add_argument("--data_dir_test", type=str, default='/home/yinxia/data',
                        help="Vaihingen, Zurich")
    parser.add_argument("--cache", type=bool, default=True, help="cache in dataloader")

    parser.add_argument("--ignore_label", type=int, default=255,
                        help="the index of the label ignored in the training.")
    parser.add_argument("--input_size_train", type=str, default='256,256',
                        help="width and height of input training images.")
    parser.add_argument("--input_size_test", type=str, default='256,256',
                        help="width and height of input test images.")
    parser.add_argument("--mode", type=int, default=1,  # set to point
                        help="annotation type (0-full, 1-point).")
    parser.add_argument("--id", type=int, default=1,
                        help="annotator id).")

    # network
    parser.add_argument("--backbone", default="effiSAM", help="effiSAM, unetformer, vgg16")
    parser.add_argument("--method", default="CRGNet", help="CRGNet, TFCSD, FESTA, TEL, UTEL, UTELT"
                                                           "baseline, oracle")
    parser.add_argument("--device", type=str, default='cuda',
                        help="cuda, cpu")
    parser.add_argument("--epsilon", type=float, default=1e-14,
                        help="avoid zero")

    parser.add_argument("--batch_size", type=int, default=32,
                        help="number of images in each batch.")
    parser.add_argument("--num_workers", type=int, default=8,
                        help="number of workers for multithread dataloading.")
    parser.add_argument("--learning_rate", type=float, default=0.001,
                        help="base learning rate.")
    parser.add_argument("--learning_power", type=float, default=0.9, # add
                        help="base learning rate.")
    parser.add_argument("--num_steps", type=int, default=5000,
                        help="Number of training steps.")
    parser.add_argument("--num_steps_stop", type=int, default=5000,
                        help="Number of training steps for early stopping.")
    parser.add_argument("--restore_from", type=str, default='./weights/fcn8s_from_caffe.pth',
                        help="pretrained VGG-16 model.")
    parser.add_argument("--weight_decay", type=float, default=5e-4,
                        help="regularisation parameter for L2-loss.")
    parser.add_argument("--momentum", type=float, default=0.9,
                        help="Momentum component of the optimiser.")
    parser.add_argument("--tau", type=float, default=0.95,
                        help="prob tau.")
    parser.add_argument("--lambda_con", type=float, default=1,
                        help="consistency weight.")
    # TEL parameters
    parser.add_argument("--tel_sigma", type=float, default=0.002,
                        help="0.002 for TEL")
    # Uncertainty threshold
    parser.add_argument("--Tuncert", type=float, default=0.9,
                        help="threshold for uncertainty")
    # samplerate for CRGNet
    parser.add_argument("--samplerate", type=int, default=4,
                        help="the input v.s. the output ")

    # add a subdir
    parser.add_argument("--subdir", type=str, default='',
                        help="the additional subdir, e.g, scale/")

    # add restore_iter: for testing
    parser.add_argument("--restore_iter", type=int, default=5000,
                        help="the restore iter")
    parser.add_argument("--palette", default=[(0,0,0)])

    # additional savedir
    parser.add_argument("--suffix", default='')

    # color transform
    parser.add_argument("--iscolor", default=False, help="apply colortans to images")

    # restore path
    parser.add_argument("--restore_path", type=str, default=None,
                        help="the restore iter")

    # log_num_classes
    parser.add_argument("--log_num_classes", type=int, default=None,
                        help="the restore iter")
    
    # create dirs
    parser.add_argument("--makedirs", default=True)
    args = parser.parse_args()

    # dataset
    if args.city=="Vaihingen":
        args.data_dir_train = f'{DATA_PATH_HOME}/Vaihingen/'
        args.data_dir_val = args.data_dir_test = args.data_dir_train
        args.train_list = f'./dataset/{args.city.lower()}_train.txt'
        args.test_list = args.val_list =f'./dataset/{args.city.lower()}_test.txt'
        args.name_classes = ['impervious surfaces', 'buildings', 'low vegetation', 'trees', 'cars']
        # BGR
        args.palette = [(255, 255, 255), (255, 0, 0), (255, 255, 0), (0, 255, 0), (0, 255, 255), (0, 0, 255), (0, 0, 0)]

    elif args.city=="Zurich":
        args.data_dir_train = f'{DATA_PATH_HOME}/Zurich/'
        args.data_dir_val = args.data_dir_test = args.data_dir_train
        args.train_list = f'./dataset/{args.city.lower()}_train.txt'
        args.test_list = args.val_list = f'./dataset/{args.city.lower()}_test.txt'
        args.name_classes = ['Roads', 'Buildings', 'Trees', 'Grass', 'Bare Soil', 'Water', 'Rails', 'Pools']
        args.palette = [(0, 0, 0), (100, 100, 100), (0, 125, 0), (0, 255, 0), (0, 80, 150),
                        (150, 0, 0), (0, 255, 255), (255, 150, 150), (255, 255, 255)]

    elif args.city=="uavid":
        args.data_dir_train = f'{DATA_PATH_HOME}/uavid/uavid_train/' # for training
        args.data_dir_val = f'{DATA_PATH_HOME}/uavid/uavid_val/' # for validation: select 7 images, each for each video
        args.data_dir_test = f'{DATA_PATH_HOME}/uavid/uavid_test/'
        args.train_list = f'./dataset/{args.city.lower()}_train.txt'
        args.val_list = f'./dataset/{args.city.lower()}_val.txt'
        args.test_list = f'./dataset/{args.city.lower()}_test.txt'
        args.name_classes = ['Building', 'Road', 'Tree', 'LowVeg', 'Cars'] # only consider five classes
        args.input_size_train = '1024, 1024'
        args.input_size_test = '1024, 1024'
        if args.batch_size>=8:
            args.batch_size = 8
        args.num_workers = 4
        args.cache = True # False
        args.palette = [(0, 0, 128), (128, 64, 128), (0, 128, 0), (0, 128, 128), (128, 0, 64), (192, 0, 192),
                        (0, 64, 64), (0, 0, 0)]
        # if 'TELT' in args.method:
        #     args.batch_size = 4
        #     args.num_workers = 4
    else:
        raise ValueError('unknown city')
    args.num_classes = len(args.name_classes)
    # add 
    if args.log_num_classes == None:
        args.log_num_classes = args.num_classes

    # tel_sigma
    tel_sigma=''
    if args.tel_sigma>0.002:
        tel_sigma = '_sigma'+str(args.tel_sigma)

    # learning rate
    if 'SAM' in args.backbone:
        args.learning_rate = 0.1
        args.learning_power = 3.0

    # samplerate: 8 for BaseNet, the original network
    sr=''
    if (args.backbone =="FCNcls") and (args.method=="CRGNet"):
        args.samplerate = 8
        sr = '_sr'+str(args.samplerate)

    Tuncert = ''
    if args.Tuncert !=0.9:
        Tuncert = '_tu'+str(args.Tuncert)

    # Result
    args.snapshot_dir = SAVE_WEIGHT_PATH + (f'{args.city}/{args.backbone}{args.subdir}/{args.method}_{modename[args.mode]}_'
                                    f'{args.id}{tel_sigma}{sr}{Tuncert}') + f'{args.suffix}'+'/'
    if args.makedirs:
        if args.restore_path is None:
            if os.path.exists(args.snapshot_dir):
                print(args.snapshot_dir)
                raise ValueError('snapshot_dir exist! and stop')
            else:
                os.makedirs(args.snapshot_dir, exist_ok=True)
            #os.makedirs(args.snapshot_dir, exist_ok=True)

    return args


def get_model(args):
    w, h = map(int, args.input_size_train.split(','))
    if args.backbone=="effiSAM":
        model_name = 'l0'
        weight_url = f'{SAVE_WEIGHT_PATH}/weights/sam/l0.pt'
        if args.method in ["baseline", "oracle", "oracle_lovaz", "dCRF"]:
            from efficientvit.SAMseg import SAM_seg
            model = SAM_seg(model_name=model_name, weight_url=weight_url,
                   imgsz=h, device=args.device, num_classes=args.num_classes)
        elif args.method in ["TEL", "UTEL", "UTELT", "BUTELT", "TELT"]:
            from efficientvit.SAMseg import SAM_seg_tree
            model = SAM_seg_tree(model_name=model_name, weight_url=weight_url,
                    imgsz=h, device=args.device, num_classes=args.num_classes)
        elif args.method in ["CRGNet"]:
            from efficientvit.SAMseg import SAM_seg_2dec
            model = SAM_seg_2dec(model_name=model_name, weight_url=weight_url,
                    imgsz=h, device=args.device, num_classes=args.num_classes)
        elif args.method in ["TFCSD"]:
            from efficientvit.SAMseg import SAM_seg_tree_2dec_TFCSD
            model = SAM_seg_tree_2dec_TFCSD(model_name=model_name, weight_url=weight_url,
                    imgsz=h, device=args.device, num_classes=args.num_classes)
        elif args.method in ["FESTA", "AGMM"]:
            from efficientvit.SAMseg import SAM_seg_FESTA
            model = SAM_seg_FESTA(model_name=model_name, weight_url=weight_url,
                            imgsz=h, device=args.device, num_classes=args.num_classes)
        else:
            raise ValueError("invalid method")

    elif args.backbone=="unetformer":
        if args.method in ["baseline", "oracle", "oracle_lovaz", "dCRF"]:
            from SAM_RS.model.UNetFormer import UNetFormer
            model = UNetFormer(num_classes=args.num_classes)
        elif args.method in ["TEL", "UTEL", "UTELT", "BUTELT", "TELT"]:
            from SAM_RS.model.UNetFormer import UNetFormerTree
            model = UNetFormerTree(num_classes=args.num_classes)
        elif args.method in ["CRGNet"]:
            from SAM_RS.model.UNetFormer import UNetFormer_2dec
            model = UNetFormer_2dec(num_classes=args.num_classes)
        elif args.method in ["TFCSD"]:
            from SAM_RS.model.UNetFormer import UNetFormerTree_2dec_TFCSD
            model = UNetFormerTree_2dec_TFCSD(num_classes=args.num_classes)
        elif args.method in ["FESTA", "AGMM"]:
            from SAM_RS.model.UNetFormer import UNetFormer_FESTA
            model = UNetFormer_FESTA(num_classes=args.num_classes)
        else:
            raise ValueError("invalid method")

    elif args.backbone=="FCNunet":
        backbone = "vgg16_bn"
        if args.method in ["baseline", "oracle", "oracle_lovaz"]:
            from TFCSD.FCN_backbone import FCNs_VGG_ASPP_4conv1
            model = FCNs_VGG_ASPP_4conv1(in_ch=3, out_ch=args.num_classes, backbone=backbone, pretrained=True)
        elif args.method in ["TEL", "UTEL", "UTELT", "BUTELT", "TELT"]:
            from TFCSD.FCN_backbone import FCNs_VGG_ASPP_4conv1_tree
            model = FCNs_VGG_ASPP_4conv1_tree(in_ch=3, out_ch=args.num_classes, backbone=backbone, pretrained=True)
        # elif args.method in ["CRGNet"]:
        #     from TFCSD.FCN_backbone import FCNs_VGG_ASPP_4conv1_2dec
        #     model = FCNs_VGG_ASPP_4conv1_2dec(in_ch=3, out_ch=args.num_classes, backbone=backbone, pretrained=True)
        elif args.method in ["TFCSD"]:
            from TFCSD.FCN_backbone import FCNs_VGG_ASPP_4conv1_TFCSD
            model = FCNs_VGG_ASPP_4conv1_TFCSD(in_ch=3, out_ch=args.num_classes, backbone=backbone, pretrained=True)
        elif args.method in ["FESTA"]:
            from TFCSD.FCN_backbone import FCNs_VGG_ASPP_4conv1_FESTA
            model = FCNs_VGG_ASPP_4conv1_FESTA(in_ch=3, out_ch=args.num_classes, backbone=backbone, pretrained=True)
        else:
            raise ValueError("invalid method")

    elif args.backbone=="FCNcls":
        from model.FCNcls import load_predtrain
        weight_url = f'{SAVE_WEIGHT_PATH}/weights/fcn8s_from_caffe.pth'
        if args.method in ["baseline", "oracle", "oracle_lovaz", "dCRF"]:
            from model.FCNcls import BaseNet
            model = BaseNet(args.num_classes)
            load_predtrain(model, weight_url)
        elif args.method in ["TEL", "UTEL", "UTELT", "BUTELT", "TELT"]:
            from model.FCNcls import BaseNet_tree
            model = BaseNet_tree(args.num_classes)
            load_predtrain(model, weight_url)
        elif args.method in ["CRGNet"]:
            from model.FCNcls import BaseNet_2dec
            model = BaseNet_2dec(args.num_classes)
            load_predtrain(model, weight_url)
        elif args.method in ["TFCSD"]:
            from model.FCNcls import BaseNet_tree_2dec_TFCSD
            model = BaseNet_tree_2dec_TFCSD(args.num_classes)
            load_predtrain(model, weight_url)
        elif args.method in ["FESTA",]:
            from model.FCNcls import BaseNet_FESTA
            model = BaseNet_FESTA(args.num_classes)
            load_predtrain(model, weight_url)
        elif args.method in ["AGMM",]:
            from model.FCNcls import BaseNet_AGMM
            model = BaseNet_AGMM(args.num_classes)
            load_predtrain(model, weight_url)
        else:
            raise ValueError("invalid method")

    elif args.backbone=="effiSAMformer":
        model_name = 'l0'
        weight_url = f'{SAVE_WEIGHT_PATH}/weights/sam/l0.pt'
        if args.method in ["baseline", "oracle", "oracle_lovaz"]:
            from efficientvit.SAMunetformer import SAM_seg
            model = SAM_seg(model_name=model_name, weight_url=weight_url,
                   imgsz=h, device=args.device, num_classes=args.num_classes)
        elif args.method in ["TEL", "UTEL", "UTELT", "BUTELT", "TELT"]:
            from efficientvit.SAMunetformer import SAM_seg_tree
            model = SAM_seg_tree(model_name=model_name, weight_url=weight_url,
                    imgsz=h, device=args.device, num_classes=args.num_classes)
        elif args.method in ["CRGNet"]:
            from efficientvit.SAMunetformer import SAM_seg_2dec
            model = SAM_seg_2dec(model_name=model_name, weight_url=weight_url,
                    imgsz=h, device=args.device, num_classes=args.num_classes)
        elif args.method in ["TFCSD"]:
            from efficientvit.SAMunetformer import SAM_seg_tree_2dec_TFCSD
            model = SAM_seg_tree_2dec_TFCSD(model_name=model_name, weight_url=weight_url,
                    imgsz=h, device=args.device, num_classes=args.num_classes)
        elif args.method in ["FESTA"]:
            from efficientvit.SAMunetformer import SAM_seg_FESTA
            model = SAM_seg_FESTA(model_name=model_name, weight_url=weight_url,
                            imgsz=h, device=args.device, num_classes=args.num_classes)
        else:
            raise ValueError("invalid method")

    else:
        raise ValueError("invalid backbone")

    return model


def get_testloader(args):
    w, h = map(int, args.input_size_test.split(','))
    #input_size_test = (w, h)
    from dataset.vaihingen_dataset_cv2 import VaihingenDataSet
    test_loader = data.DataLoader(
    VaihingenDataSet(args.data_dir_test, args.test_list, set='test', cache=args.cache),
    batch_size=1, shuffle=False, num_workers=args.num_workers, pin_memory=True)
    return test_loader


def get_valloader(args):
    w, h = map(int, args.input_size_test.split(','))
    #input_size_test = (w, h)
    from dataset.vaihingen_dataset_cv2 import VaihingenDataSet
    test_loader = data.DataLoader(
    VaihingenDataSet(args.data_dir_val, args.val_list, set='test', cache=args.cache),
    batch_size=1, shuffle=False, num_workers=args.num_workers, pin_memory=True)
    return test_loader


def get_train_testloader(args):
    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)

    if args.method in ["baseline", "oracle", "oracle_lovaz", "FESTA", "CRGNet", "AGMM"]:
        iscls_label = False
        if args.method in ["AGMM"]:
            iscls_label = True
        from dataset.vaihingen_dataset_cv2 import VaihingenDataSet
        src_loader = data.DataLoader(
                        VaihingenDataSet(args.data_dir_train, args.train_list, max_iters=args.num_steps_stop*args.batch_size,
                        crop_size=input_size_train, set='train', mode=args.mode, id=args.id,
                        cache=args.cache, iscls_label=iscls_label, num_classes=args.num_classes),
                        batch_size=args.batch_size, shuffle=True, num_workers=args.num_workers, pin_memory=True)

        test_loader = data.DataLoader(
                        VaihingenDataSet(args.data_dir_val, args.val_list,set='test', cache=args.cache),
                        batch_size=1, shuffle=False, num_workers=args.num_workers, pin_memory=True)

    elif args.method in ["TEL", "UTEL", "TFCSD", "UTELT", "BUTELT", "TELT", "dCRF"]:
        unlabeled = True
        if args.method in ["dCRF"]:
            unlabeled = False
        # add cache by cyx, add unlabeled for the tree_energy_loss
        from dataset.vaihingen_dataset_cv2 import VaihingenDataSet_RRM
        src_loader = data.DataLoader(
            VaihingenDataSet_RRM(args.data_dir_train, args.train_list, max_iters=args.num_steps_stop * args.batch_size,
                                 crop_size=input_size_train, set='train', mode=args.mode, id=args.id,
                                 cache=args.cache, unlabeled=unlabeled, iscolor=args.iscolor),
            batch_size=args.batch_size, shuffle=True, num_workers=args.num_workers, pin_memory=True)

        test_loader = data.DataLoader(
            VaihingenDataSet_RRM(args.data_dir_val, args.val_list, set='test', cache=args.cache),
            batch_size=1, shuffle=False, num_workers=args.num_workers, pin_memory=True)

    else:
        raise ValueError("invalid method")

    return src_loader, test_loader


def init_seeds(seed=1234):
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)