FPGA & FreeNet
12 |Fast Patch-Free Global Learning Framework for Fully End-to-End Hyperspectral Image Classification
13 | 14 |by Zhuo Zheng, Yanfei Zhong, Ailong Ma and Liangpei Zhang
15 | 16 |
17 | 
18 |
19 |
20 | This is an official implementation of FPGA framework and FreeNet in our TGRS 2020 paper ["FPGA: Fast Patch-Free Global Learning Framework for Fully End-to-End Hyperspectral Image Classification"](https://ieeexplore.ieee.org/document/9007624).
21 |
22 | We hope the FPGA framework can become a stronger and cleaner baseline for hyperspectral image classification research in the future.
23 |
24 | ## News
25 | 1. 2020/05/28, We release the code of FreeNet and FPGA framework.
26 |
27 |
28 | ## Features
29 | 1. Patch-free training and inference
30 | 2. Fully end-to-end (w/o preprocess technologies, such as dimension reduction)
31 |
32 |
33 | ## Citation
34 | If you use FPGA framework or FreeNet in your research, please cite the following paper:
35 | ```text
36 | @article{zheng2020fpga,
37 | title={FPGA: Fast Patch-Free Global Learning Framework for Fully End-to-End Hyperspectral Image Classification},
38 | author={Zheng, Zhuo and Zhong, Yanfei and Ma, Ailong and Zhang, Liangpei},
39 | journal={IEEE Transactions on Geoscience and Remote Sensing},
40 | year={2020},
41 | publisher={IEEE},
42 | note={doi: {10.1109/TGRS.2020.2967821}}
43 | }
44 | ```
45 |
46 |
47 | ## Getting Started
48 | ### 1. Install SimpleCV
49 |
50 | ```bash
51 | pip install --upgrade git+https://github.com/Z-Zheng/SimpleCV.git
52 | ```
53 | ### 2. Prepare datasets
54 |
55 | It is recommended to symlink the dataset root to `$FreeNet`.
56 |
57 | The project should be organized as:
58 | ```text
59 | FreeNet
60 | ├── configs // configure files
61 | ├── data // dataset and dataloader class
62 | ├── module // network arch.
63 | ├── scripts
64 | ├── pavia // data 1
65 | │ ├── PaviaU.mat
66 | │ ├── PaviaU_gt.mat
67 | ├── salinas // data 2
68 | │ ├── Salinas_corrected.mat
69 | │ ├── Salinas_gt.mat
70 | ├── GRSS2013 // data 3
71 | │ ├── 2013_IEEE_GRSS_DF_Contest_CASI.tif
72 | │ ├── train_roi.tif
73 | │ ├── val_roi.tif
74 | ```
75 |
76 | ### 3. run experiments
77 |
78 | #### 1. PaviaU
79 | ```bash
80 | bash scripts/freenet_1_0_pavia.sh
81 | ```
82 |
83 | #### 2. Salinas
84 | ```bash
85 | bash scripts/freenet_1_0_salinas.sh
86 | ```
87 |
88 | #### 3. GRSS2013
89 | ```bash
90 | bash scripts/freenet_1_0_grss.sh
91 | ```
92 |
93 | ### License
94 | This source code is released under [GPLv3](http://www.gnu.org/licenses/) license.
95 |
96 | For commercial use, please contact Prof. Zhong (zhongyanfei@whu.edu.cn).
97 |
98 |
--------------------------------------------------------------------------------
/configs/freenet/freenet_1_0_grss2013.py:
--------------------------------------------------------------------------------
1 | config = dict(
2 | model=dict(
3 | type='FreeNet',
4 | params=dict(
5 | in_channels=144,
6 | num_classes=15,
7 | block_channels=(96, 128, 192, 256),
8 | inner_dim=128,
9 | reduction_ratio=1.0,
10 | )
11 | ),
12 | data=dict(
13 | train=dict(
14 | type='NewGRSS2013Loader',
15 | params=dict(
16 | training=True,
17 | num_workers=0,
18 | image_path='./GRSS2013/2013_IEEE_GRSS_DF_Contest_CASI.tif',
19 | gt_path='./GRSS2013/train_roi.tif',
20 | sub_minibatch=20
21 | )
22 | ),
23 | test=dict(
24 | type='NewGRSS2013Loader',
25 | params=dict(
26 | training=False,
27 | num_workers=0,
28 | image_path='./GRSS2013/2013_IEEE_GRSS_DF_Contest_CASI.tif',
29 | gt_path='./GRSS2013/val_roi.tif',
30 | sub_minibatch=20
31 | )
32 | )
33 | ),
34 | optimizer=dict(
35 | type='sgd',
36 | params=dict(
37 | momentum=0.9,
38 | weight_decay=0.001
39 | )
40 | ),
41 | learning_rate=dict(
42 | type='poly',
43 | params=dict(
44 | base_lr=0.001,
45 | power=0.9,
46 | max_iters=1000),
47 | ),
48 | train=dict(
49 | forward_times=1,
50 | num_iters=1000,
51 | eval_per_epoch=True,
52 | summary_grads=False,
53 | summary_weights=False,
54 | eval_after_train=True,
55 | resume_from_last=True,
56 | ),
57 | test=dict(
58 | draw=dict(
59 | image_size=(349, 1905),
60 | palette=[
61 | 0, 0, 0,
62 | 139, 67, 45,
63 | 0, 0, 255,
64 | 255, 100, 0,
65 | 0, 255, 123,
66 | 164, 75, 155,
67 | 101, 173, 255,
68 | 118, 254, 172,
69 | 60, 91, 112,
70 | 255, 255, 0,
71 | 255, 255, 125,
72 | 255, 0, 255,
73 | 100, 0, 255,
74 | 0, 172, 254,
75 | 0, 255, 0,
76 | 171, 175, 80
77 | ]
78 | )
79 | ),
80 |
81 | )
82 |
--------------------------------------------------------------------------------
/configs/freenet/freenet_1_0_pavia.py:
--------------------------------------------------------------------------------
1 | config = dict(
2 | model=dict(
3 | type='FreeNet',
4 | params=dict(
5 | in_channels=103,
6 | num_classes=9,
7 | block_channels=(96, 128, 192, 256),
8 | inner_dim=128,
9 | reduction_ratio=1.0,
10 | )
11 | ),
12 | data=dict(
13 | train=dict(
14 | type='NewPaviaLoader',
15 | params=dict(
16 | training=True,
17 | num_workers=0,
18 | image_mat_path='./pavia/PaviaU.mat',
19 | gt_mat_path='./pavia/PaviaU_gt.mat',
20 | num_train_samples_per_class=200,
21 | sub_minibatch=20
22 | )
23 | ),
24 | test=dict(
25 | type='NewPaviaLoader',
26 | params=dict(
27 | training=False,
28 | num_workers=0,
29 | image_mat_path='./pavia/PaviaU.mat',
30 | gt_mat_path='./pavia/PaviaU_gt.mat',
31 | num_train_samples_per_class=200,
32 | sub_minibatch=20
33 | )
34 | )
35 | ),
36 | optimizer=dict(
37 | type='sgd',
38 | params=dict(
39 | momentum=0.9,
40 | weight_decay=0.001
41 | )
42 | ),
43 | learning_rate=dict(
44 | type='poly',
45 | params=dict(
46 | base_lr=0.001,
47 | power=0.9,
48 | max_iters=1000),
49 | ),
50 | train=dict(
51 | forward_times=1,
52 | num_iters=1000,
53 | eval_per_epoch=True,
54 | summary_grads=False,
55 | summary_weights=False,
56 | eval_after_train=True,
57 | resume_from_last=False,
58 | ),
59 | test=dict(
60 | draw=dict(
61 | image_size=(610, 340),
62 | palette=[
63 | 0, 0, 0,
64 | 192, 192, 192,
65 | 0, 255, 1,
66 | 0, 255, 255,
67 | 0, 128, 1,
68 | 255, 0, 254,
69 | 165, 82, 40,
70 | 129, 0, 127,
71 | 255, 0, 0,
72 | 255, 255, 0, ]
73 | )
74 | ),
75 | )
76 |
--------------------------------------------------------------------------------
/configs/freenet/freenet_1_0_salinas.py:
--------------------------------------------------------------------------------
1 | config = dict(
2 | model=dict(
3 | type='FreeNet',
4 | params=dict(
5 | in_channels=204,
6 | num_classes=16,
7 | block_channels=(96, 128, 192, 256),
8 | inner_dim=128,
9 | reduction_ratio=1.0,
10 | )
11 | ),
12 | data=dict(
13 | train=dict(
14 | type='NewSalinasLoader',
15 | params=dict(
16 | num_workers=0,
17 | image_mat_path='./salinas/Salinas_corrected.mat',
18 | gt_mat_path='./salinas/Salinas_gt.mat',
19 | training=True,
20 | num_train_samples_per_class=200,
21 | sub_minibatch=20
22 | )
23 | ),
24 | test=dict(
25 | type='NewSalinasLoader',
26 | params=dict(
27 | num_workers=0,
28 | image_mat_path='./salinas/Salinas_corrected.mat',
29 | gt_mat_path='./salinas/Salinas_gt.mat',
30 | training=False,
31 | num_train_samples_per_class=200,
32 | sub_minibatch=20
33 | )
34 | )
35 | ),
36 | optimizer=dict(
37 | type='sgd',
38 | params=dict(
39 | momentum=0.9,
40 | weight_decay=0.001
41 | )
42 | ),
43 | learning_rate=dict(
44 | type='poly',
45 | params=dict(
46 | base_lr=0.001,
47 | power=0.9,
48 | max_iters=1000),
49 | ),
50 | train=dict(
51 | forward_times=1,
52 | num_iters=1000,
53 | eval_per_epoch=True,
54 | summary_grads=False,
55 | summary_weights=False,
56 | eval_after_train=True,
57 | resume_from_last=False,
58 | ),
59 | test=dict(
60 | draw=dict(
61 | image_size=(610, 340),
62 | palette=[
63 | 0, 0, 0,
64 | 192, 192, 192,
65 | 0, 255, 1,
66 | 0, 255, 255,
67 | 0, 128, 1,
68 | 255, 0, 254,
69 | 165, 82, 40,
70 | 129, 0, 127,
71 | 255, 0, 0,
72 | 255, 255, 0, ]
73 | )
74 | ),
75 | )
76 |
--------------------------------------------------------------------------------
/data/__init__.py:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/Z-Zheng/FreeNet/ac0f5a35262ec867529d279a6d9c53074348767b/data/__init__.py
--------------------------------------------------------------------------------
/data/base.py:
--------------------------------------------------------------------------------
1 | from torch.utils.data import dataset
2 | import numpy as np
3 | from simplecv.data.preprocess import divisible_pad
4 | import torch
5 | from torch.utils import data
6 |
7 | SEED = 2333
8 |
9 |
10 | class FullImageDataset(dataset.Dataset):
11 | def __init__(self,
12 | image,
13 | mask,
14 | training,
15 | np_seed=2333,
16 | num_train_samples_per_class=200,
17 | sub_minibatch=10,
18 | ):
19 | self.image = image
20 | self.mask = mask
21 | self.training = training
22 | self.num_train_samples_per_class = num_train_samples_per_class
23 | self.sub_minibatch = sub_minibatch
24 | self._seed = np_seed
25 | self._rs = np.random.RandomState(np_seed)
26 | # set list lenght = 9999 to make sure seeds enough
27 | self.seeds_for_minibatchsample = [e for e in self._rs.randint(low=2 << 31 - 1, size=9999)]
28 | self.preset()
29 |
30 | def preset(self):
31 | train_indicator, test_indicator = fixed_num_sample(self.mask, self.num_train_samples_per_class,
32 | self.num_classes, self._seed)
33 |
34 | blob = divisible_pad([np.concatenate([self.image.transpose(2, 0, 1),
35 | self.mask[None, :, :],
36 | train_indicator[None, :, :],
37 | test_indicator[None, :, :]], axis=0)], 16, False)
38 | im = blob[0, :self.image.shape[-1], :, :]
39 |
40 | mask = blob[0, -3, :, :]
41 | self.train_indicator = blob[0, -2, :, :]
42 | self.test_indicator = blob[0, -1, :, :]
43 |
44 | if self.training:
45 | self.train_inds_list = minibatch_sample(mask, self.train_indicator, self.sub_minibatch,
46 | seed=self.seeds_for_minibatchsample.pop())
47 |
48 | self.pad_im = im
49 | self.pad_mask = mask
50 |
51 | def resample_minibatch(self):
52 | self.train_inds_list = minibatch_sample(self.pad_mask, self.train_indicator, self.sub_minibatch,
53 | seed=self.seeds_for_minibatchsample.pop())
54 |
55 | @property
56 | def num_classes(self):
57 | return 9
58 |
59 | def __getitem__(self, idx):
60 |
61 | if self.training:
62 | return self.pad_im, self.pad_mask, self.train_inds_list[idx]
63 |
64 | else:
65 | return self.pad_im, self.pad_mask, self.test_indicator
66 |
67 | def __len__(self):
68 | if self.training:
69 | return len(self.train_inds_list)
70 | else:
71 | return 1
72 |
73 |
74 | class MinibatchSampler(data.Sampler):
75 | def __init__(self, dataset: FullImageDataset):
76 | super(MinibatchSampler, self).__init__(None)
77 | self.dataset = dataset
78 | self.g = torch.Generator()
79 | self.g.manual_seed(SEED)
80 |
81 | def __iter__(self):
82 | self.dataset.resample_minibatch()
83 | n = len(self.dataset)
84 | return iter(torch.randperm(n, generator=self.g).tolist())
85 |
86 | def __len__(self):
87 | return len(self.dataset)
88 |
89 |
90 | def fixed_num_sample(gt_mask: np.ndarray, num_train_samples, num_classes, seed=2333):
91 | """
92 |
93 | Args:
94 | gt_mask: 2-D array of shape [height, width]
95 | num_train_samples: int
96 | num_classes: scalar
97 | seed: int
98 |
99 | Returns:
100 | train_indicator, test_indicator
101 | """
102 | rs = np.random.RandomState(seed)
103 |
104 | gt_mask_flatten = gt_mask.ravel()
105 | train_indicator = np.zeros_like(gt_mask_flatten)
106 | test_indicator = np.zeros_like(gt_mask_flatten)
107 | for i in range(1, num_classes + 1):
108 | inds = np.where(gt_mask_flatten == i)[0]
109 | rs.shuffle(inds)
110 |
111 | train_inds = inds[:num_train_samples]
112 | test_inds = inds[num_train_samples:]
113 |
114 | train_indicator[train_inds] = 1
115 | test_indicator[test_inds] = 1
116 |
117 | train_indicator = train_indicator.reshape(gt_mask.shape)
118 | test_indicator = test_indicator.reshape(gt_mask.shape)
119 |
120 | return train_indicator, test_indicator
121 |
122 |
123 | def minibatch_sample(gt_mask: np.ndarray, train_indicator: np.ndarray, minibatch_size, seed):
124 | """
125 |
126 | Args:
127 | gt_mask: 2-D array of shape [height, width]
128 | train_indicator: 2-D array of shape [height, width]
129 | minibatch_size:
130 |
131 | Returns:
132 |
133 | """
134 | rs = np.random.RandomState(seed)
135 | # split into N classes
136 | cls_list = np.unique(gt_mask)
137 | inds_dict_per_class = dict()
138 | for cls in cls_list:
139 | train_inds_per_class = np.where(gt_mask == cls, train_indicator, np.zeros_like(train_indicator))
140 | inds = np.where(train_inds_per_class.ravel() == 1)[0]
141 | rs.shuffle(inds)
142 |
143 | inds_dict_per_class[cls] = inds
144 |
145 | train_inds_list = []
146 | cnt = 0
147 | while True:
148 | train_inds = np.zeros_like(train_indicator).ravel()
149 | for cls, inds in inds_dict_per_class.items():
150 | left = cnt * minibatch_size
151 | if left >= len(inds):
152 | continue
153 | # remain last batch though the real size is smaller than minibatch_size
154 | right = min((cnt + 1) * minibatch_size, len(inds))
155 | fetch_inds = inds[left:right]
156 | train_inds[fetch_inds] = 1
157 | cnt += 1
158 | if train_inds.sum() == 0:
159 | return train_inds_list
160 | train_inds_list.append(train_inds.reshape(train_indicator.shape))
161 |
--------------------------------------------------------------------------------
/data/dataloader.py:
--------------------------------------------------------------------------------
1 | from torch.utils.data.dataloader import DataLoader
2 | from simplecv import registry
3 | from data.pavia import NewPaviaDataset
4 | from data.base import MinibatchSampler
5 | from data.grss2013 import NewGRSS2013Dataset
6 | from data.salinas import NewSalinasDataset
7 |
8 |
9 | @registry.DATALOADER.register('NewPaviaLoader')
10 | class NewPaviaLoader(DataLoader):
11 | def __init__(self, config):
12 | self.config = dict()
13 | self.set_defalut()
14 | self.config.update(config)
15 | for k, v in self.config.items():
16 | self.__dict__[k] = v
17 |
18 | dataset = NewPaviaDataset(self.image_mat_path, self.gt_mat_path, self.training,
19 | self.num_train_samples_per_class, self.sub_minibatch)
20 | sampler = MinibatchSampler(dataset)
21 | super(NewPaviaLoader, self).__init__(dataset,
22 | batch_size=1,
23 | shuffle=False,
24 | sampler=sampler,
25 | batch_sampler=None,
26 | num_workers=self.num_workers,
27 | pin_memory=True,
28 | drop_last=False,
29 | timeout=0,
30 | worker_init_fn=None)
31 |
32 | def set_defalut(self):
33 | self.config.update(dict(
34 | num_workers=0,
35 | image_mat_path='',
36 | gt_mat_path='',
37 | training=True,
38 | num_train_samples_per_class=200,
39 | # mini-batch per class, if there are 10 categories, the total mini-batch is sub_minibatch * num_classes (10)
40 | sub_minibatch=10
41 | ))
42 |
43 |
44 | @registry.DATALOADER.register('NewSalinasLoader')
45 | class NewSalinasLoader(DataLoader):
46 | def __init__(self, config):
47 | self.config = dict()
48 | self.set_defalut()
49 | self.config.update(config)
50 | for k, v in self.config.items():
51 | self.__dict__[k] = v
52 |
53 | dataset = NewSalinasDataset(self.image_mat_path, self.gt_mat_path, self.training,
54 | self.num_train_samples_per_class, self.sub_minibatch)
55 | sampler = MinibatchSampler(dataset)
56 | super(NewSalinasLoader, self).__init__(dataset,
57 | batch_size=1,
58 | shuffle=False,
59 | sampler=sampler,
60 | batch_sampler=None,
61 | num_workers=self.num_workers,
62 | pin_memory=True,
63 | drop_last=False,
64 | timeout=0,
65 | worker_init_fn=None)
66 |
67 | def set_defalut(self):
68 | self.config.update(dict(
69 | num_workers=0,
70 | image_mat_path='',
71 | gt_mat_path='',
72 | training=True,
73 | num_train_samples_per_class=200,
74 | # mini-batch per class, if there are 10 categories, the total mini-batch is sub_minibatch * num_classes (10)
75 | sub_minibatch=10
76 | ))
77 |
78 |
79 | @registry.DATALOADER.register('NewGRSS2013Loader')
80 | class NewGRSS2013Loader(DataLoader):
81 | def __init__(self, config):
82 | self.config = dict()
83 | self.set_defalut()
84 | self.config.update(config)
85 | for k, v in self.config.items():
86 | self.__dict__[k] = v
87 |
88 | dataset = NewGRSS2013Dataset(self.image_path, self.gt_path, self.training, self.sub_minibatch)
89 | sampler = MinibatchSampler(dataset)
90 | super(NewGRSS2013Loader, self).__init__(dataset,
91 | batch_size=1,
92 | shuffle=False,
93 | sampler=sampler,
94 | batch_sampler=None,
95 | num_workers=self.num_workers,
96 | pin_memory=True,
97 | drop_last=False,
98 | timeout=0,
99 | worker_init_fn=None)
100 |
101 | def set_defalut(self):
102 | self.config.update(dict(
103 | num_workers=0,
104 | image_path='',
105 | gt_path='',
106 | training=True,
107 | # mini-batch per class, if there are 10 categories, the total mini-batch is sub_minibatch * num_classes (10)
108 | sub_minibatch=10
109 | ))
110 |
--------------------------------------------------------------------------------
/data/grss2013.py:
--------------------------------------------------------------------------------
1 | from simplecv.data import preprocess
2 | import numpy as np
3 | from data.base import FullImageDataset
4 | import tifffile
5 | from data.base import minibatch_sample
6 | from simplecv.data.preprocess import divisible_pad
7 |
8 | SEED = 2333
9 |
10 |
11 | class NewGRSS2013Dataset(FullImageDataset):
12 | def __init__(self,
13 | image_path,
14 | gt_path,
15 | training=True,
16 | sub_minibatch=10):
17 | self.im_mat_path = image_path
18 | self.gt_mat_path = gt_path
19 |
20 | image = tifffile.imread(image_path)
21 | mask = tifffile.imread(gt_path)
22 |
23 | im_cmean = image.reshape((-1, image.shape[-1])).mean(axis=0)
24 | im_cstd = image.reshape((-1, image.shape[-1])).std(axis=0)
25 | self.vanilla_image = image
26 | image = preprocess.mean_std_normalize(image, im_cmean, im_cstd)
27 | self.training = training
28 | self.sub_minibatch = sub_minibatch
29 | super(NewGRSS2013Dataset, self).__init__(image, mask, training, np_seed=SEED,
30 | num_train_samples_per_class=None,
31 | sub_minibatch=sub_minibatch)
32 |
33 | def preset(self):
34 | indicator = np.where(self.mask != 0, np.ones_like(self.mask), np.zeros_like(self.mask))
35 |
36 | blob = divisible_pad([np.concatenate([self.image.transpose(2, 0, 1),
37 | self.mask[None, :, :],
38 | indicator[None, :, :]], axis=0)], 16, False)
39 | im = blob[0, :self.image.shape[-1], :, :]
40 |
41 | mask = blob[0, -2, :, :]
42 | self.indicator = blob[0, -1, :, :]
43 |
44 | if self.training:
45 | self.train_inds_list = minibatch_sample(mask, self.indicator, self.sub_minibatch,
46 | seed=self.seeds_for_minibatchsample.pop())
47 |
48 | self.pad_im = im
49 | self.pad_mask = mask
50 |
51 | def resample_minibatch(self):
52 | self.train_inds_list = minibatch_sample(self.pad_mask, self.indicator, self.sub_minibatch,
53 | seed=self.seeds_for_minibatchsample.pop())
54 |
55 | @property
56 | def num_classes(self):
57 | return 15
58 |
59 | def __getitem__(self, idx):
60 |
61 | if self.training:
62 | return self.pad_im, self.pad_mask, self.train_inds_list[idx]
63 |
64 | else:
65 | return self.pad_im, self.pad_mask, self.indicator
66 |
67 | def __len__(self):
68 | if self.training:
69 | return len(self.train_inds_list)
70 | else:
71 | return 1
72 |
--------------------------------------------------------------------------------
/data/pavia.py:
--------------------------------------------------------------------------------
1 | from scipy.io import loadmat
2 | from simplecv.data import preprocess
3 |
4 | from data.base import FullImageDataset
5 |
6 | SEED = 2333
7 |
8 |
9 | class NewPaviaDataset(FullImageDataset):
10 | def __init__(self,
11 | image_mat_path,
12 | gt_mat_path,
13 | training=True,
14 | num_train_samples_per_class=200,
15 | sub_minibatch=10):
16 | self.im_mat_path = image_mat_path
17 | self.gt_mat_path = gt_mat_path
18 |
19 | im_mat = loadmat(image_mat_path)
20 | image = im_mat['paviaU']
21 | gt_mat = loadmat(gt_mat_path)
22 | mask = gt_mat['paviaU_gt']
23 |
24 | im_cmean = image.reshape((-1, image.shape[-1])).mean(axis=0)
25 | im_cstd = image.reshape((-1, image.shape[-1])).std(axis=0)
26 | self.vanilla_image = image
27 | image = preprocess.mean_std_normalize(image, im_cmean, im_cstd)
28 | self.training = training
29 | self.num_train_samples_per_class = num_train_samples_per_class
30 | self.sub_minibatch = sub_minibatch
31 | super(NewPaviaDataset, self).__init__(image, mask, training, np_seed=SEED,
32 | num_train_samples_per_class=num_train_samples_per_class,
33 | sub_minibatch=sub_minibatch)
34 |
--------------------------------------------------------------------------------
/data/salinas.py:
--------------------------------------------------------------------------------
1 | from scipy.io import loadmat
2 | from simplecv.data import preprocess
3 |
4 | from data.base import FullImageDataset
5 |
6 | SEED = 2333
7 |
8 |
9 | class NewSalinasDataset(FullImageDataset):
10 | def __init__(self,
11 | image_mat_path,
12 | gt_mat_path,
13 | training=True,
14 | num_train_samples_per_class=200,
15 | sub_minibatch=10):
16 | self.im_mat_path = image_mat_path
17 | self.gt_mat_path = gt_mat_path
18 |
19 | im_mat = loadmat(image_mat_path)
20 | image = im_mat['salinas_corrected']
21 | gt_mat = loadmat(gt_mat_path)
22 | mask = gt_mat['salinas_gt']
23 |
24 | im_cmean = image.reshape((-1, image.shape[-1])).mean(axis=0)
25 | im_cstd = image.reshape((-1, image.shape[-1])).std(axis=0)
26 | self.vanilla_image = image
27 | image = preprocess.mean_std_normalize(image, im_cmean, im_cstd)
28 | self.training = training
29 | self.num_train_samples_per_class = num_train_samples_per_class
30 | self.sub_minibatch = sub_minibatch
31 | super(NewSalinasDataset, self).__init__(image, mask, training, np_seed=SEED,
32 | num_train_samples_per_class=num_train_samples_per_class,
33 | sub_minibatch=sub_minibatch)
34 |
35 | @property
36 | def num_classes(self):
37 | return 16
38 |
--------------------------------------------------------------------------------
/module/__init__.py:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/Z-Zheng/FreeNet/ac0f5a35262ec867529d279a6d9c53074348767b/module/__init__.py
--------------------------------------------------------------------------------
/module/freenet.py:
--------------------------------------------------------------------------------
1 | import torch.nn as nn
2 | import torch.nn.functional as F
3 | from simplecv.interface import CVModule
4 | from simplecv.module import SEBlock
5 | from simplecv import registry
6 | import torch
7 | import math
8 |
9 |
10 | def conv3x3_gn_relu(in_channel, out_channel, num_group):
11 | return nn.Sequential(
12 | nn.Conv2d(in_channel, out_channel, 3, 1, 1),
13 | nn.GroupNorm(num_group, out_channel),
14 | nn.ReLU(inplace=True),
15 | )
16 |
17 |
18 | def downsample2x(in_channel, out_channel):
19 | return nn.Sequential(
20 | nn.Conv2d(in_channel, out_channel, 3, 2, 1),
21 | nn.ReLU(inplace=True)
22 | )
23 |
24 |
25 | def repeat_block(block_channel, r, n):
26 | layers = [
27 | nn.Sequential(
28 | SEBlock(block_channel, r),
29 | conv3x3_gn_relu(block_channel, block_channel, r)
30 | )
31 | for _ in range(n)]
32 | return nn.Sequential(*layers)
33 |
34 |
35 | @registry.MODEL.register('FreeNet')
36 | class FreeNet(CVModule):
37 | def __init__(self, config):
38 | super(FreeNet, self).__init__(config)
39 | r = int(16 * self.config.reduction_ratio)
40 | block1_channels = int(self.config.block_channels[0] * self.config.reduction_ratio / r) * r
41 | block2_channels = int(self.config.block_channels[1] * self.config.reduction_ratio / r) * r
42 | block3_channels = int(self.config.block_channels[2] * self.config.reduction_ratio / r) * r
43 | block4_channels = int(self.config.block_channels[3] * self.config.reduction_ratio / r) * r
44 |
45 | self.feature_ops = nn.ModuleList([
46 | conv3x3_gn_relu(self.config.in_channels, block1_channels, r),
47 |
48 | repeat_block(block1_channels, r, self.config.num_blocks[0]),
49 | nn.Identity(),
50 | downsample2x(block1_channels, block2_channels),
51 |
52 | repeat_block(block2_channels, r, self.config.num_blocks[1]),
53 | nn.Identity(),
54 | downsample2x(block2_channels, block3_channels),
55 |
56 | repeat_block(block3_channels, r, self.config.num_blocks[2]),
57 | nn.Identity(),
58 | downsample2x(block3_channels, block4_channels),
59 |
60 | repeat_block(block4_channels, r, self.config.num_blocks[3]),
61 | nn.Identity(),
62 | ])
63 | inner_dim = int(self.config.inner_dim * self.config.reduction_ratio)
64 | self.reduce_1x1convs = nn.ModuleList([
65 | nn.Conv2d(block1_channels, inner_dim, 1),
66 | nn.Conv2d(block2_channels, inner_dim, 1),
67 | nn.Conv2d(block3_channels, inner_dim, 1),
68 | nn.Conv2d(block4_channels, inner_dim, 1),
69 | ])
70 | self.fuse_3x3convs = nn.ModuleList([
71 | nn.Conv2d(inner_dim, inner_dim, 3, 1, 1),
72 | nn.Conv2d(inner_dim, inner_dim, 3, 1, 1),
73 | nn.Conv2d(inner_dim, inner_dim, 3, 1, 1),
74 | nn.Conv2d(inner_dim, inner_dim, 3, 1, 1),
75 | ])
76 | self.cls_pred_conv = nn.Conv2d(inner_dim, self.config.num_classes, 1)
77 |
78 | def top_down(self, top, lateral):
79 | top2x = F.interpolate(top, scale_factor=2.0, mode='nearest')
80 | return lateral + top2x
81 |
82 | def forward(self, x, y=None, w=None, **kwargs):
83 | feat_list = []
84 | for op in self.feature_ops:
85 | x = op(x)
86 | if isinstance(op, nn.Identity):
87 | feat_list.append(x)
88 |
89 | inner_feat_list = [self.reduce_1x1convs[i](feat) for i, feat in enumerate(feat_list)]
90 | inner_feat_list.reverse()
91 |
92 | out_feat_list = [self.fuse_3x3convs[0](inner_feat_list[0])]
93 | for i in range(len(inner_feat_list) - 1):
94 | inner = self.top_down(out_feat_list[i], inner_feat_list[i + 1])
95 | out = self.fuse_3x3convs[i + 1](inner)
96 | out_feat_list.append(out)
97 |
98 | final_feat = out_feat_list[-1]
99 |
100 | logit = self.cls_pred_conv(final_feat)
101 | if self.training:
102 | loss_dict = {
103 | 'cls_loss': self.loss(logit, y, w)
104 | }
105 | return loss_dict
106 |
107 | return torch.softmax(logit, dim=1)
108 |
109 | def loss(self, x, y, weight):
110 | losses = F.cross_entropy(x, y.long() - 1, weight=None,
111 | ignore_index=-1, reduction='none')
112 |
113 | v = losses.mul_(weight).sum() / weight.sum()
114 | return v
115 |
116 | def set_defalut_config(self):
117 | self.config.update(dict(
118 | in_channels=204,
119 | num_classes=16,
120 | block_channels=(96, 128, 192, 256),
121 | num_blocks=(1, 1, 1, 1),
122 | inner_dim=128,
123 | reduction_ratio=1.0,
124 | ))
125 |
--------------------------------------------------------------------------------
/scripts/freenet_1_0_grss.sh:
--------------------------------------------------------------------------------
1 | #!/usr/bin/env bash
2 |
3 | export CUDA_VISIBLE_DEVICES=0
4 | export PYTHONPATH=${PYTHONPATH}:`pwd`
5 | config_path='freenet.freenet_1_0_grss2013'
6 |
7 | model_dir='./log/grss2013/freenet/1.0_poly'
8 |
9 |
10 | python train.py \
11 | --config_path=${config_path} \
12 | --model_dir=${model_dir} \
13 | train.save_ckpt_interval_epoch 9999
--------------------------------------------------------------------------------
/scripts/freenet_1_0_pavia.sh:
--------------------------------------------------------------------------------
1 | #!/usr/bin/env bash
2 |
3 | export CUDA_VISIBLE_DEVICES=0
4 | export PYTHONPATH=${PYTHONPATH}:`pwd`
5 | config_path='freenet.freenet_1_0_pavia'
6 |
7 | model_dir='./log/pavia/freenet/1.0_poly'
8 |
9 |
10 | python train.py \
11 | --config_path=${config_path} \
12 | --model_dir=${model_dir} \
13 | train.save_ckpt_interval_epoch 9999
--------------------------------------------------------------------------------
/scripts/freenet_1_0_salinas.sh:
--------------------------------------------------------------------------------
1 | #!/usr/bin/env bash
2 |
3 | export CUDA_VISIBLE_DEVICES=0
4 | export PYTHONPATH=${PYTHONPATH}:`pwd`
5 | config_path='freenet.freenet_1_0_salinas'
6 |
7 | model_dir='./log/salinas/freenet/1.0_poly'
8 |
9 |
10 | python train.py \
11 | --config_path=${config_path} \
12 | --model_dir=${model_dir} \
13 | train.save_ckpt_interval_epoch 9999
--------------------------------------------------------------------------------
/train.py:
--------------------------------------------------------------------------------
1 | from simplecv import dp_train as train
2 | import torch
3 | from simplecv.util.logger import eval_progress, speed
4 | import time
5 | from module import freenet
6 | from simplecv.util import metric
7 | from data import dataloader
8 |
9 |
10 | def fcn_evaluate_fn(self, test_dataloader, config):
11 | if self.checkpoint.global_step < 0:
12 | return
13 | self._model.eval()
14 | total_time = 0.
15 | with torch.no_grad():
16 | for idx, (im, mask, w) in enumerate(test_dataloader):
17 | start = time.time()
18 | y_pred = self._model(im).squeeze()
19 | torch.cuda.synchronize()
20 | time_cost = round(time.time() - start, 3)
21 | y_pred = y_pred.argmax(dim=0).cpu() + 1
22 | w.unsqueeze_(dim=0)
23 |
24 | w = w.byte()
25 | mask = torch.masked_select(mask.view(-1), w.view(-1))
26 | y_pred = torch.masked_select(y_pred.view(-1), w.view(-1))
27 |
28 | oa = metric.th_overall_accuracy_score(mask.view(-1), y_pred.view(-1))
29 | aa, acc_per_class = metric.th_average_accuracy_score(mask.view(-1), y_pred.view(-1),
30 | self._model.module.config.num_classes,
31 | return_accuracys=True)
32 | kappa = metric.th_cohen_kappa_score(mask.view(-1), y_pred.view(-1), self._model.module.config.num_classes)
33 | total_time += time_cost
34 | speed(self._logger, time_cost, 'im')
35 |
36 | eval_progress(self._logger, idx + 1, len(test_dataloader))
37 |
38 | speed(self._logger, round(total_time / len(test_dataloader), 3), 'batched im (avg)')
39 |
40 | metric_dict = {
41 | 'OA': oa.item(),
42 | 'AA': aa.item(),
43 | 'Kappa': kappa.item()
44 | }
45 | for i, acc in enumerate(acc_per_class):
46 | metric_dict['acc_{}'.format(i + 1)] = acc.item()
47 | self._logger.eval_log(metric_dict=metric_dict, step=self.checkpoint.global_step)
48 |
49 |
50 | def register_evaluate_fn(launcher):
51 | launcher.override_evaluate(fcn_evaluate_fn)
52 |
53 |
54 | if __name__ == '__main__':
55 | torch.backends.cudnn.benchmark = True
56 | args = train.parser.parse_args()
57 | SEED = 2333
58 | torch.manual_seed(SEED)
59 | torch.cuda.manual_seed(SEED)
60 | train.run(config_path=args.config_path,
61 | model_dir=args.model_dir,
62 | cpu_mode=args.cpu,
63 | after_construct_launcher_callbacks=[register_evaluate_fn],
64 | opts=args.opts)
65 |
--------------------------------------------------------------------------------
18 |