-
Notifications
You must be signed in to change notification settings - Fork 1
/
Copy pathmain_2_mgcl_test.py
144 lines (121 loc) · 6.32 KB
/
main_2_mgcl_test.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
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
import os
import torch
import pickle
import random
import argparse
import numpy as np
import torch.nn as nn
import PIL.Image as Image
import torch.optim as optim
import torch.nn.functional as F
from torchvision.models import vgg
from alisuretool.Tools import Tools
from torch.utils.data import Dataset
from torchvision.models import resnet
from torch.utils.data import DataLoader
from net.net_tools import MGCLNetwork
from util.util_tools import MyCommon, MyOptim, AverageMeter, Logger
from dataset.dataset_tools import FSSDataset, Evaluator, DatasetISAID
class Runner(object):
def __init__(self, args):
self.args = args
self.device = MyCommon.gpu_setup(use_gpu=self.args.use_gpu, gpu_id=args.gpuid)
self.model = MGCLNetwork(args, False).to(self.device)
self.model.eval()
weights = torch.load(args.load, map_location=None if self.args.use_gpu else torch.device('cpu'))
weights = {one.replace("module.", ""): weights[one] for one in weights.keys()}
weights = {one.replace("hpn_learner.", "mgcd."): weights[one] for one in weights.keys()}
self.model.load_state_dict(weights)
FSSDataset.initialize(img_size=args.img_size, datapath=args.datapath,
use_original_imgsize=args.use_original_imgsize)
self.dataloader_val = FSSDataset.build_dataloader(
args.benchmark, args.bsz, args.nworker, args.fold, 'val', args.shot,
use_mask=args.mask, mask_num=args.mask_num)
pass
@torch.no_grad()
def test(self):
dataloader = self.dataloader_val
average_meter = AverageMeter(dataloader.dataset, device=self.device)
for idx, batch in enumerate(dataloader):
# 1. forward pass
batch = MyCommon.to_cuda(batch, device=self.device)
pred = self.model.predict_nshot(batch)
# 2. Evaluate prediction
area_inter, area_union = Evaluator.classify_prediction(pred, batch)
average_meter.update(area_inter, area_union, batch['class_id'], loss=None)
average_meter.write_process(idx, len(dataloader), 0, write_batch_idx=5)
pass
miou, fb_iou = average_meter.compute_iou()
return miou, fb_iou
@torch.no_grad()
def test_class(self):
dataloader = self.dataloader_val
average_meter = AverageMeter(dataloader.dataset, device=self.device)
for idx, batch in enumerate(dataloader):
# 1. forward pass
batch = MyCommon.to_cuda(batch, device=self.device)
pred = self.model.predict_nshot(batch)
# 2. Evaluate prediction
area_inter, area_union = Evaluator.classify_prediction(pred, batch)
average_meter.update(area_inter, area_union, batch['class_id'], loss=None)
average_meter.write_process(idx, len(dataloader), 0, write_batch_idx=5)
pass
miou, fb_iou, iou = average_meter.compute_iou_class()
return miou, fb_iou, iou
pass
def my_parser(fold=0, shot=1, backbone='resnet50', benchmark="isaid",
load='./logs/demo/best_model.pt', use_gpu=False, gpu_id=0,
bsz=2, mask=True, mask_num=128):
datapath = None
if benchmark == "isaid":
datapath = '/mnt/4T/ALISURE/FSS-RS/remote_sensing/iSAID_patches'
elif benchmark == "dlrsd":
datapath = '/mnt/4T/ALISURE/FSS-RS/DLRSD'
parser = argparse.ArgumentParser(description='MGCL Pytorch Implementation')
parser.add_argument('--logpath', type=str, default='demo')
parser.add_argument('--use_gpu', type=bool, default=use_gpu)
parser.add_argument('--gpuid', type=int, default=gpu_id)
parser.add_argument('--bsz', type=int, default=bsz)
parser.add_argument('--fold', type=int, default=fold, choices=[0, 1, 2])
parser.add_argument('--shot', type=int, default=shot, choices=[1, 5])
parser.add_argument('--load', type=str, default=load)
parser.add_argument('--backbone', type=str, default=backbone,
choices=['vgg16', 'resnet50', 'resnet101'])
parser.add_argument('--mask', type=bool, default=mask)
parser.add_argument('--mask_num', type=int, default=mask_num)
parser.add_argument('--datapath', type=str, default=datapath)
parser.add_argument('--benchmark', type=str, default=benchmark, choices=['isaid', 'dlrsd'])
parser.add_argument('--nworker', type=int, default=8)
parser.add_argument('--img_size', type=int, default=256)
parser.add_argument('--use_original_imgsize', type=bool, default=False)
args = parser.parse_args()
return args
if __name__ == '__main__':
MyCommon.fix_randseed(0)
# isaid resnet50
# 2023-10-24 09:54:36 fold=0, shot=1, mIoU: 42.77439880371094 FB-IoU: 62.60266876220703
# 2023-11-03 14:01:02 iou: [0.3375, 0.4660, 0.4952, 0.5550, 0.2850]
# 2023-10-24 10:04:32 fold=0, shot=5, mIoU: 49.13707733154297 FB-IoU: 66.159912109375
args = my_parser(fold=0, shot=1, backbone='resnet50', benchmark="isaid",
use_gpu=True, gpu_id=0, bsz=8, mask=True, mask_num=128,
load='./logs/resnet50_fold0/best_model.pt')
# 2023-10-24 09:57:40 fold=1, shot=1, mIoU: 30.592941284179688 FB-IoU: 55.60072708129883
# 2023-11-03 14:01:54 iou: [0.3214, 0.4551, 0.3579, 0.1319, 0.2634]
# 2023-10-24 10:06:42 fold=1, shot=5, mIoU: 32.57774353027344 FB-IoU: 57.06192398071289
# args = my_parser(fold=1, shot=1, backbone='resnet50', benchmark="isaid",
# use_gpu=True, gpu_id=0, bsz=8, mask=True, mask_num=128,
# load='./logs/resnet50_fold1/best_model.pt')
# 2023-10-24 09:57:02 fold=2, shot=1, mIoU: 46.38703918457031 FB-IoU: 58.5185661315918
# 2023-11-03 14:02:02 iou: [0.4573, 0.6166, 0.4611, 0.5389, 0.2456],
# 2023-10-24 10:25:51 fold=2, shot=5, mIoU: 52.74596405029297 FB-IoU: 62.11357116699219
# args = my_parser(fold=2, shot=1, backbone='resnet50', benchmark="isaid",
# use_gpu=True, gpu_id=0, bsz=8, mask=True, mask_num=128,
# load='./logs/resnet50_fold2/best_model.pt')
Logger.initialize(args, training=False)
runner = Runner(args=args)
Logger.log_params(runner.model)
# miou, fb_iou = runner.test()
# Tools.print("mIoU: {} FB-IoU: {}".format(miou, fb_iou))
miou, fb_iou, iou = runner.test_class()
Tools.print("mIoU: {} FB-IoU: {} iou: {}".format(miou, fb_iou, iou))
pass