SS-MAE/get_dat.py at main · zhoulifan/SS-MAE · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
import data.dataset as Dataset

def get_dataset(args):
    train_dataset = args.dataset
    Dataset.set_random_seed(0)

    if train_dataset == "Houston2018":

        pretrain_loader, train_loader, test_loader, trntst_loader, all_loader = Dataset.getHSData(
            datasetType="Houston2018",
            channels=args.pca_num,
            windowSize=args.crop_size,
            batch_size=args.batch_size,
            num_workers=0,args=args)

    elif train_dataset == "Berlin":

        pretrain_loader, train_loader, test_loader, trntst_loader, all_loader = Dataset.getHSData(
            datasetType="Berlin",
            channels=args.pca_num,
            windowSize=args.crop_size,
            batch_size=args.batch_size,
            num_workers=0,args=args)

    elif train_dataset == "Augsburg":

        pretrain_loader, train_loader, test_loader, trntst_loader, all_loader = Dataset.getHSData(
            datasetType="Augsburg",
            channels=args.pca_num,
            windowSize=args.crop_size,
            batch_size=args.batch_size,
            num_workers=0,args=args)

    print("completed!")

    return pretrain_loader, train_loader, test_loader, trntst_loader, all_loader