-
Notifications
You must be signed in to change notification settings - Fork 9
/
hsi_denoising_urban_test.py
executable file
·61 lines (42 loc) · 1.54 KB
/
hsi_denoising_urban_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
import torch
import torch.nn as nn
import torch.optim as optim
import torchvision
import torchvision.transforms as transforms
from torchvision.transforms import Compose
import os
import argparse
from utility import *
from hsi_setup import Engine, train_options, make_dataset
if __name__ == '__main__':
"""Training settings"""
parser = argparse.ArgumentParser(
description='Hyperspectral Image Denoising (Complex noise)')
opt = train_options(parser)
print(opt)
"""Setup Engine"""
engine = Engine(opt)
"""Dataset Setting"""
HSI2Tensor = partial(HSI2Tensor, use_2dconv=engine.net.use_2dconv)
target_transform = HSI2Tensor()
"""Test-Dev"""
basefolder = '/data/HSI_Data/Hyperspectral_Project/'
mat_datasets = [MatDataFromFolder(
basefolder, size=1,fns=['Urban_304.mat']) ]
if not engine.get_net().use_2dconv:
mat_transform = Compose([
LoadMatHSI(input_key='input', gt_key='input',
transform=lambda x:x[ ...][None], needsigma=False),
])
else:
mat_transform = Compose([
LoadMatHSI(input_key='input', gt_key='input', needsigma=False),
])
mat_datasets = [TransformDataset(mat_dataset, mat_transform)
for mat_dataset in mat_datasets]
mat_loaders = [DataLoader(
mat_dataset,
batch_size=1, shuffle=False,
num_workers=1, pin_memory=opt.no_cuda
) for mat_dataset in mat_datasets]
engine.test(mat_loaders[0], basefolder)