-
Notifications
You must be signed in to change notification settings - Fork 3
/
main_cd.py
92 lines (71 loc) · 3.48 KB
/
main_cd.py
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
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
from argparse import ArgumentParser
import torch
from models.trainer import *
print(torch.cuda.is_available())
"""
the main function for training the CD networks
"""
def train(args):
dataloaders = utils.get_loaders(args)
model = CDTrainer(args=args, dataloaders=dataloaders)
model.train_models()
def test(args):
from models.evaluator import CDEvaluator
dataloader = utils.get_loader(args.data_name, img_size=args.img_size,
batch_size=args.batch_size, is_train=False,
split='test', dataset=args.dataset)
model = CDEvaluator(args=args, dataloader=dataloader)
model.eval_models()
def test2(args):
from models.evaluator import CDEvaluator
model = CDEvaluator(args=args, dataloader='')
model.pred_gdal_blocks_write(r'E:\bianhuajiance\cq\fengjiexian\nanbu\2020\500115_clip4.tif',
r'E:\bianhuajiance\cq\fengjiexian\nanbu\2021\500115_clip4.tif')
if __name__ == '__main__':
# ------------
# args
# ------------
parser = ArgumentParser()
parser.add_argument('--gpu_ids', type=str, default='-1', help='gpu ids: e.g. 0 0,1,2, 0,2. use -1 for CPU')
parser.add_argument('--project_name', default='chongqing_EGCTNet_eas_bs16_0.0001', type=str)
parser.add_argument('--checkpoint_root', default='checkpoints', type=str)
parser.add_argument('--vis_root', default='vis', type=str)
# data
parser.add_argument('--num_workers', default=2, type=int)
parser.add_argument('--dataset', default='CDDataset', type=str)
parser.add_argument('--data_name', default='WHU-512-100', type=str)
parser.add_argument('--batch_size', default=1, type=int)
parser.add_argument('--split', default="train", type=str)
parser.add_argument('--split_val', default="val", type=str)
parser.add_argument('--img_size', default=512, type=int)
parser.add_argument('--shuffle_AB', default=False, type=str)
# model
parser.add_argument('--n_class', default=2, type=int)
parser.add_argument('--embed_dim', default=32, type=int)
parser.add_argument('--pretrain', default=None, type=str)
parser.add_argument('--multi_scale_train', default=False, type=str)
parser.add_argument('--multi_scale_infer', default=False, type=str)
parser.add_argument('--multi_pred_weights', nargs='+', type=float, default=[0.5, 0.5, 0.5, 0.8, 1.0])
parser.add_argument('--net_G', default='EGCTNet', type=str,
help='base_resnet18 | base_transformer_pos_s4 | '
'base_transformer_pos_s4_dd8 | '
'base_transformer_pos_s4_dd8_dedim8|ChangeFormerV5|SiamUnet_diff')
parser.add_argument('--loss', default='eas', type=str)
# optimizer
parser.add_argument('--optimizer', default='adamw', type=str)
parser.add_argument('--lr', default=0.002, type=float)
parser.add_argument('--max_epochs', default=100, type=int)
parser.add_argument('--lr_policy', default='linear', type=str,
help='linear | step')
parser.add_argument('--lr_decay_iters', default=100, type=int)
args = parser.parse_args()
utils.get_device(args)
print(args.gpu_ids)
# checkpoints dir
args.checkpoint_dir = os.path.join(args.checkpoint_root, args.project_name)
os.makedirs(args.checkpoint_dir, exist_ok=True)
# visualize dir
args.vis_dir = os.path.join(args.vis_root, args.project_name)
os.makedirs(args.vis_dir, exist_ok=True)
# train(args)
test2(args)