class TrainModel():
def name(self):
return 'Train Model'
def initialize(self, opt):
self.opt = opt
self.opt.imageSize = self.opt.imageSize if len(
self.opt.imageSize) == 2 else self.opt.imageSize * 2
self.gpu_ids = ''
self.batchSize = self.opt.batchSize
self.checkpoints_path = os.path.join(self.opt.checkpoints,
self.opt.name)
self.scheduler = None
self.create_save_folders()
# criterion to evaluate the val split
self.criterion_eval = MSEScaledError()
self.mse_scaled_error = MSEScaledError()
self.opt.print_freq = self.opt.display_freq
self.visualizer = Visualizer(opt)
if self.opt.resume and self.opt.display_id > 0:
self.load_plot_data()
elif opt.train:
self.start_epoch = 1
self.best_val_error = 999.9
# self.print_save_options()
# Logfile
self.logfile = open(os.path.join(self.checkpoints_path, 'logfile.txt'),
'a')
if opt.validate:
self.logfile_val = open(
os.path.join(self.checkpoints_path, 'logfile_val.txt'), 'a')
# Prepare a random seed that will be the same for everyone
# opt.manualSeed = random.randint(1, 10000) # fix seed
# print("Random Seed: ", opt.manualSeed)
# # random.seed(opt.manualSeed)
# torch.manual_seed(opt.manualSeed)
self.random_seed = 123
random.seed(self.random_seed)
torch.cuda.manual_seed_all(self.random_seed)
torch.manual_seed(self.random_seed)
if opt.cuda:
self.cuda = torch.device(
'cuda:0') # set externally. ToDo: set internally
torch.cuda.manual_seed(self.random_seed)
# uses the inbuilt cudnn auto-tuner to find the fastest convolution algorithms.
cudnn.benchmark = self.opt.use_cudnn_benchmark # using too much memory - use when not in astroboy
cudnn.enabled = True
if not opt.train and not opt.test and not opt.resume:
raise Exception("You have to set --train or --test")
if torch.cuda.is_available and not opt.cuda:
print(
"WARNING: You have a CUDA device, so you should run WITHOUT --cpu"
)
if not torch.cuda.is_available and opt.cuda:
raise Exception("No GPU found, run WITH --cpu")
def set_input(self, input):
self.input = input
def create_network(self):
netG = networks.define_G(input_nc=self.opt.input_nc,
output_nc=self.opt.output_nc,
ngf=64,
net_architecture=self.opt.net_architecture,
opt=self.opt,
gpu_ids='')
if self.opt.cuda:
netG = netG.cuda()
return netG
def get_optimizerG(self, network, lr, weight_decay=0.0):
generator_params = filter(lambda p: p.requires_grad,
network.parameters())
return optim.Adam(generator_params,
lr=lr,
betas=(self.opt.beta1, 0.999),
weight_decay=weight_decay)
def get_checkpoint(self, epoch):
pass
def train_batch(self):
"""Each method has a different implementation"""
pass
def display_gradients_norms(self):
return 'nothing yet'
def get_current_errors_display(self):
pass
def get_regression_criterion(self):
if self.opt.regression_loss == 'L1':
return nn.L1Loss()
def get_variable(self, tensor, requires_grad=False):
if self.opt.cuda:
tensor = tensor.cuda()
return Variable(tensor, requires_grad=requires_grad)
def restart_variables(self):
self.it = 0
self.rmse = 0
self.n_images = 0
def train(self, data_loader, val_loader=None):
self.data_loader = data_loader
self.len_data_loader = len(
self.data_loader) # check if gonna use elsewhere
self.total_iter = 0
for epoch in range(self.start_epoch, self.opt.nEpochs):
self.restart_variables()
self.data_iter = iter(self.data_loader)
# self.pbar = tqdm(range(self.len_data_loader))
self.pbar = range(self.len_data_loader)
# while self.it < self.len_data_loader:
for self.it in self.pbar:
if self.opt.optim == 'SGD':
self.scheduler.step()
self.total_iter += self.opt.batchSize
self.netG.train(True)
iter_start_time = time.time()
self.train_batch()
d_time = (time.time() - iter_start_time) / self.opt.batchSize
# print errors
self.print_current_errors(epoch, d_time)
# display errors
self.display_current_results(epoch)
# Validate
self.evaluate(val_loader, epoch)
# save checkpoint
self.save_checkpoint(epoch, is_best=0)
self.logfile.close()
if self.opt.validate:
self.logfile_val.close()
def get_next_batch(self):
# self.it += 1 # important for GANs
rgb_cpu, depth_cpu = self.data_iter.next()
# depth_cpu = depth_cpu[0]
self.input.data.resize_(rgb_cpu.size()).copy_(rgb_cpu)
# self.target.data.resize_(depth_cpu.size()).copy_(depth_cpu)
def apply_valid_pixels_mask(self, *data, value=0.0):
# self.nomask_outG = data[0].data # for displaying purposes
mask = (data[1].data > value).to(self.cuda, dtype=torch.float32)
masked_data = []
for d in data:
masked_data.append(d * mask)
return masked_data, mask.sum()
def update_learning_rate(self, epoch):
if epoch > self.opt.niter_decay and self.opt.use_cgan: # but independs if conditional or not
# Linear decay for discriminator
[self.opt.d_lr,
self.optimD] = self._update_learning_rate(self.opt.niter_decay,
self.opt.d_lr,
self.optimD)
[self.opt.lr,
self.optimG] = self._update_learning_rate(self.opt.niter_decay,
self.opt.lr,
self.optimG)
def _update_learning_rate(self, niter_decay, old_lr, optim):
lr = old_lr - old_lr / niter_decay
for param_group in optim.param_groups:
param_group['lr'] = lr
return lr, optim
# CONTROL FUNCTIONS OF THE ARCHITECTURE
def _get_plot_data_filename(self, phase):
return os.path.join(
self.checkpoints_path,
'plot_data' + ('' if phase == 'train' else '_' + phase) + '.p')
def save_static_plot_image():
return None
def save_interactive_plot_image():
return None
def _save_plot_data(self, plot_data, filename):
# save
pickle.dump(plot_data, open(filename, 'wb'))
def save_plot_data(self):
self._save_plot_data(self.visualizer.plot_data,
self._get_plot_data_filename('train'))
if self.opt.validate and self.total_iter > self.opt.val_freq:
self._save_plot_data(self.visualizer.plot_data_val,
self._get_plot_data_filename('val'))
def _load_plot_data(self, filename):
# verify if file exists
if not os.path.isfile(filename):
raise Exception(
'In _load_plot_data file {} doesnt exist.'.format(filename))
else:
return pickle.load(open(filename, "rb"))
def load_plot_data(self):
self.visualizer.plot_data = self._load_plot_data(
self._get_plot_data_filename('train'))
if self.opt.validate:
self.visualizer.plot_data_val = self._load_plot_data(
self._get_plot_data_filename('val'))
def save_checkpoint(self, epoch, is_best):
if epoch % self.opt.save_checkpoint_freq == 0 or is_best:
checkpoint = self.get_checkpoint(epoch)
checkpoint_filename = '{}/{:04}.pth.tar'.format(
self.checkpoints_path, epoch)
self._save_checkpoint(
checkpoint, is_best=is_best, filename=checkpoint_filename
) # standart is_best=0 here cause we didn' evaluate on validation data
# save plot data as well
def _save_checkpoint(self, state, is_best, filename):
print("Saving checkpoint...")
# uncomment next 2 lines if we still want per epoch
torch.save(state, filename)
shutil.copyfile(
filename, os.path.join(os.path.dirname(filename),
'latest.pth.tar'))
# comment next 2 lines if necessary if using last two lines
# filename = os.path.join(self.checkpoints_path, 'latest.pth.tar')
# torch.save(state, os.path.join(self.checkpoints_path, 'latest.pth.tar'))
if is_best:
shutil.copyfile(
filename, os.path.join(self.checkpoints_path, 'best.pth.tar'))
def create_save_folders(self):
if self.opt.train:
os.system('mkdir -p {0}'.format(self.checkpoints_path))
# if self.opt.save_samples:
# subfolders = ['input', 'target', 'results', 'output']
# self.save_samples_path = os.path.join('results/train_results/', self.opt.name)
# for subfolder in subfolders:
# path = os.path.join(self.save_samples_path, subfolder)
# os.system('mkdir -p {0}'.format(path))
# if self.opt.test:
# self.save_samples_path = os.path.join('results/test_results/', self.opt.name)
# self.save_samples_path = os.path.join(self.save_samples_path, self.opt.epoch)
# for subfolder in subfolders:
# path = os.path.join(self.save_samples_path, subfolder)
# os.system('mkdir -p {0}'.format(path))
def print_save_options(self):
options_file = open(os.path.join(self.checkpoints_path, 'options.txt'),
'w')
args = dict((arg, getattr(self.opt, arg)) for arg in dir(self.opt)
if not arg.startswith('_'))
print('---Options---')
for k, v in sorted(args.items()):
option = '{}: {}'.format(k, v)
# print options
print(option)
# save options in file
options_file.write(option + '\n')
options_file.close()
def mean_errors(self):
pass
def get_current_errors(self):
pass
def print_current_errors(self, epoch, d_time):
if self.total_iter % self.opt.print_freq == 0:
self.mean_errors()
errors = self.get_current_errors()
message = self.visualizer.print_errors(errors, epoch, self.it,
self.len_data_loader,
d_time)
# self.pbar.set_description(message)
print(message)
# self.pbar.refresh()
# def print_epoch_error(error):
# pass
def get_current_visuals(self):
pass
def display_current_results(self, epoch):
if self.opt.display_id > 0 and self.total_iter % self.opt.display_freq == 0:
errors = self.get_current_errors_display()
self.visualizer.display_errors(
errors, epoch,
float(self.it) / self.len_data_loader)
visuals = self.get_current_visuals()
self.visualizer.display_images(visuals, epoch)
# save printed errors to logfile
self.visualizer.save_errors_file(self.logfile)
def evaluate(self, data_loader, epoch):
if self.opt.validate and self.total_iter % self.opt.val_freq == 0:
val_error = self.get_eval_error(data_loader, self.netG,
self.criterion_eval, epoch)
# errors = OrderedDict([('LossL1', self.e_reg if self.opt.reg_type == 'L1' else self.L1error),
# ('ValError', val_error.item())])
errors = OrderedDict([('RMSE', self.rmse_epoch),
('RMSEVal', val_error)])
self.visualizer.display_errors(errors,
epoch,
float(self.it) /
self.len_data_loader,
phase='val')
message = self.visualizer.print_errors(errors, epoch, self.it,
len(data_loader), 0)
print('[Validation] ' + message)
self.visualizer.save_errors_file(self.logfile_val)
self.save_plot_data()
# save best models
is_best = self.best_val_error > val_error
if is_best: # and not self.opt.not_save_val_model:
print("Updating BEST model (epoch {}, iters {})\n".format(
epoch, self.total_iter))
self.best_val_error = val_error
self.save_checkpoint(epoch, is_best)
def get_eval_error(self, val_loader, model, criterion, epoch):
"""
Validate every self.opt.val_freq epochs
"""
# no need to switch to model.eval because we want to keep dropout layers. Do I gave to ignore batch norm layers?
cumulated_rmse = 0
batchSize = 1
input = self.get_variable(torch.FloatTensor(batchSize, 3,
self.opt.imageSize[0],
self.opt.imageSize[1]),
requires_grad=False)
mask = self.get_variable(torch.FloatTensor(batchSize, 1,
self.opt.imageSize[0],
self.opt.imageSize[1]),
requires_grad=False)
target = self.get_variable(
torch.FloatTensor(batchSize, 1, self.opt.imageSize[0],
self.opt.imageSize[1]))
# model.eval()
model.train(False)
pbar_val = tqdm(val_loader)
for i, (rgb_cpu, depth_cpu) in enumerate(pbar_val):
pbar_val.set_description('[Validation]')
input.data.resize_(rgb_cpu.size()).copy_(rgb_cpu)
target.data.resize_(depth_cpu.size()).copy_(depth_cpu)
if self.opt.use_padding:
from torch.nn import ReflectionPad2d
self.opt.padding = self.get_padding_image(input)
input = ReflectionPad2d(self.opt.padding)(input)
target = ReflectionPad2d(self.opt.padding)(target)
# get output of the network
with torch.no_grad():
outG = model.forward(input)
# apply mask
nomask_outG = outG.data # for displaying purposes
mask_ByteTensor = self.get_mask(target.data)
mask.data.resize_(mask_ByteTensor.size()).copy_(mask_ByteTensor)
outG = outG * mask
target = target * mask
cumulated_rmse += sqrt(criterion(outG, target, mask,
no_mask=False))
if (i == 1):
self.visualizer.display_images(OrderedDict([
('input', input.data), ('gt', target.data),
('output', nomask_outG)
]),
epoch='val {}'.format(epoch),
phase='val')
return cumulated_rmse / len(val_loader)
def get_mask(self, data, value=0.0):
return (target.data > 0.0)
def get_padding(self, dim):
final_dim = (dim // 32 + 1) * 32
return final_dim - dim
def get_padding_image(self, img):
# get tensor dimensions
h, w = img.size()[2:]
w_pad, h_pad = self.get_padding(w), self.get_padding(h)
pwr = w_pad // 2
pwl = w_pad - pwr
phb = h_pad // 2
phu = h_pad - phb
# pwl, pwr, phu, phb
return (pwl, pwr, phu, phb)
def adjust_learning_rate(self, initial_lr, optimizer, epoch):
"""Sets the learning rate to the initial LR decayed by 10 every 30 epochs"""
lr = initial_lr * (0.1**(epoch // self.opt.niter_decay))
if epoch % self.opt.niter_decay == 0:
print("LEARNING RATE DECAY HERE: lr = {}".format(lr))
for param_group in optimizer.param_groups:
param_group['lr'] = lr
class TestModel(GenericTestModel):
def initialize(self, opt):
# GenericTestModel.initialize(self, opt)
self.opt = opt
self.get_color_palette()
self.opt.imageSize = self.opt.imageSize if len(
self.opt.imageSize) == 2 else self.opt.imageSize * 2
self.gpu_ids = ''
self.batchSize = self.opt.batchSize
self.checkpoints_path = os.path.join(self.opt.checkpoints,
self.opt.name)
self.create_save_folders()
self.opt.use_semantics = (('multitask' in self.opt.model)
or ('semantics' in self.opt.model))
self.netG = self.load_network()
# self.opt.dfc_preprocessing = 2
# self.data_loader, _ = CreateDataLoader(opt, Dataset)
# visualizer
self.visualizer = Visualizer(self.opt)
if 'semantics' in self.opt.tasks:
from util.util import get_color_palette
self.opt.color_palette = np.array(
get_color_palette(self.opt.dataset_name))
# self.opt.color_palette = list(self.opt.color_palette.reshape(-1))
# st()
# def initialize(self, opt):
# GenericTestModel.initialize(self, opt)
# self.get_color_palette()
def name(self):
return 'Raster Test Model'
def get_color_palette(self):
if self.opt.dataset_name == 'dfc':
self.opt.color_palette = [[0, 0, 0], [0, 205, 0], [127, 255, 0],
[46, 139, 87], [0, 139, 0], [0, 70, 0],
[160, 82, 45], [0, 255, 255],
[255, 255, 255], [216, 191, 216],
[255, 0, 0], [170, 160, 150],
[128, 128, 128], [160, 0, 0], [80, 0, 0],
[232, 161, 24], [255, 255, 0],
[238, 154, 0], [255, 0, 255],
[0, 0, 255], [176, 196, 222]]
elif self.opt.dataset_name == 'isprs':
self.opt.color_palette = [
[0, 0, 0],
[255, 255, 255],
[0, 0, 255],
[0, 255, 255],
[0, 255, 0],
[255, 255, 0],
[255, 0, 0],
]
def load_network(self):
if self.opt.epoch is not 'latest' or self.opt.epoch is not 'best':
self.opt.epoch = self.opt.epoch.zfill(4)
checkpoint_file = os.path.join(self.checkpoints_path,
self.opt.epoch + '.pth.tar')
if os.path.isfile(checkpoint_file):
print("Loading {} checkpoint of model {} ...".format(
self.opt.epoch, self.opt.name))
checkpoint = torch.load(checkpoint_file)
self.start_epoch = int(checkpoint['epoch'])
self.opt.net_architecture = checkpoint['arch_netG']
try:
self.opt.d_block_type = checkpoint['d_block_type']
# Extra options for raster:
self.opt.which_raster = checkpoint['which_raster']
self.opt.model = checkpoint['model']
self.opt.tasks = checkpoint['tasks']
self.opt.outputs_nc = checkpoint['outputs_nc']
self.opt.n_classes = checkpoint['n_classes']
except:
pass
self.opt.use_skips = checkpoint['use_skips']
self.opt.model = checkpoint['model']
self.opt.dfc_preprocessing = checkpoint['dfc_preprocessing']
self.opt.mtl_method = checkpoint['mtl_method']
self.opt.tasks = checkpoint['tasks']
self.opt.outputs_nc = checkpoint['outputs_nc']
netG = self.create_G_network()
pretrained_dict = checkpoint['state_dictG']
pattern = re.compile(
r'^(.*denselayer\d+\.(?:norm|relu|conv))\.((?:[12])\.(?:weight|bias|running_mean|running_var))$'
)
for key in list(pretrained_dict.keys()):
res = pattern.match(key)
if res:
new_key = res.group(1) + res.group(2)
pretrained_dict[new_key] = pretrained_dict[key]
del pretrained_dict[key]
netG.load_state_dict(pretrained_dict)
if self.opt.cuda:
netG = netG.cuda()
self.best_val_error = checkpoint['best_pred']
print("Loaded model from epoch {}".format(self.start_epoch))
return netG
else:
print("Couldn't find checkpoint on path: {}".format(
self.checkpoints_path + '/' + self.opt.epoch))
def save_raster_png(self, data, filename):
if 'semantics' in filename:
from util.util import labels_to_colors
image_save = Image.fromarray(np.squeeze(
labels_to_colors(data,
self.opt.color_palette).astype(np.int8)),
mode='RGB').convert(
'P',
palette=Image.ADAPTIVE,
colors=256)
image_save.save(filename)
def save_merged_rasters(self, datatype, fileroot=None):
import rasterio
from rasterio.merge import merge
from rasterio.plot import show
from os.path import join
import argparse
import glob
if fileroot == None:
fileroot = datatype
root = '{}/{}/{}*.tif'.format(self.save_samples_path, datatype,
fileroot)
filename = '{}/{}/merged_{}.tif'.format(self.save_samples_path,
datatype, fileroot)
files = glob.glob(join(root))
mosaic_rasters = [rasterio.open(file) for file in files]
mosaic, out_transform = merge(mosaic_rasters)
meta = (rasterio.open(files[0])).meta
meta.update({
"driver": "GTiff",
"height": mosaic.shape[1],
"width": mosaic.shape[2],
"transform": out_transform
})
with rasterio.open(filename, "w", **meta) as dest:
dest.write(mosaic)
filename = '{}/{}/{}_merged.png'.format(self.save_samples_path,
datatype, datatype)
self.save_raster_png(mosaic, filename)
if 'output' in filename or 'target' in filename:
self.save_height_colormap(filename, mosaic)
def test_raster(self):
if 'semantics' in self.opt.model:
from dataloader.dataset_raster import load_rgb_and_label as load_data
self.test_raster_notarget(load_data)
else:
from dataloader.dataset_raster import load_rgb_and_labels as load_data
self.test_raster_target(load_data)
def initialize_test_bayesian(self, opt):
from dataloader.dataset_bank import dataset_dfc
print('Test phase using {} split.'.format(self.opt.test_split))
phase = 'test'
input_list, target_path = dataset_dfc(
self.opt.dataroot,
data_split=self.opt.test_split,
phase='test',
model=self.opt.model,
which_raster=self.opt.which_raster)
# Sanity check : raise an error if some files do not exist
for f in input_list + target_path:
if not os.path.isfile(f):
raise KeyError('{} is not a file !'.format(f))
from dataloader.dataset_raster import load_rgb_and_labels as load_data
self.data_loader = [
(rgb, depth, meta, depth_patch_shape)
for rgb, depth, meta, depth_patch_shape in load_data(
input_list,
target_path,
phase,
self.opt.dfc_preprocessing,
which_raster=self.opt.which_raster,
use_semantics=False,
save_semantics=self.opt.save_semantics)
] # false because we do not have the GT
# no error in save semantics, same value to both variables
self.netG.eval()
self.netG.apply(self.activate_dropout)
def activate_dropout(self, m):
if type(m) == nn.Dropout:
# print(m)
m.train()
def get_meta_data(self):
return self.meta_data
def get_shape(self):
return self.shape
def get_data_loader_size(self):
return len(self.data_loader)
def test_bayesian(self, it, n_iters):
error_list = []
outG_list = []
use_semantics = self.opt.use_semantics
self.opt.use_semantics = False
# self.augmentation = augmentation
imageSize = self.opt.imageSize if len(
self.opt.imageSize) == 2 else self.opt.imageSize * 2
test_stride = self.opt.test_stride if len(
self.opt.test_stride) == 2 else self.opt.test_stride * 2
# create a matrix with a gaussian distribution to be the weights during reconstruction
prob_matrix = self.gaussian_kernel(imageSize[0], imageSize[1])
# for it, (input, target, meta_data, depth_patch_shape) in enumerate(tqdm(self.data_loader)):
input, target, meta_data, depth_patch_shape = self.data_loader[it]
for it in (tqdm(range(n_iters))):
rgb_cache = []
depth_cache = []
self.meta_data = meta_data
self.shape = depth_patch_shape
# pred = np.zeros(input.shape[-2:])
# concatenate probability matrix
pred = np.zeros([input.shape[-2], input.shape[-1]])
if self.opt.reconstruction_method == 'gaussian':
pred = np.zeros([2, input.shape[-2], input.shape[-1]])
pred_sem = np.zeros(
[self.opt.n_classes, input.shape[-2], input.shape[-1]])
else:
pred_sem = np.zeros([input.shape[-2], input.shape[-1]])
target_reconstructed = np.zeros(input.shape[-2:])
# input is a tensor
rgb_cache = [
crop for crop in self.sliding_window_coords(
input, test_stride, imageSize)
]
depth_cache = [
crop for crop in self.sliding_window_coords(
target, test_stride, imageSize)
] # don't need both
for input_crop_tuple, target_crop_tuple in tqdm(
zip(rgb_cache, depth_cache), total=len(rgb_cache)):
input_crop, (x1, x2, y1, y2) = input_crop_tuple
input_crop = self.get_variable(input_crop)
# self.complete_padding = True
# ToDo: Deal with padding later
if self.opt.use_padding:
from torch.nn import ReflectionPad2d
self.opt.padding = self.get_padding_image_dims(input_crop)
input_crop = ReflectionPad2d(self.opt.padding)(input_crop)
(pwl, pwr, phu, phb) = self.opt.padding
# target_crop = ReflectionPad2d(self.opt.padding)(target_crop)
with torch.no_grad():
outG, _ = self.netG.forward(input_crop)
out_numpy = outG.data[0].cpu().float().numpy()
if self.opt.reconstruction_method == 'concatenation':
if self.opt.use_padding:
pred[y1:y2,
x1:x2] = (out_numpy[0])[phu:phu +
self.opt.imageSize[1],
pwl:pwl +
self.opt.imageSize[0]]
else:
pred[y1:y2, x1:x2] = out_numpy[0]
elif self.opt.reconstruction_method == 'gaussian':
pred[0, y1:y2,
x1:x2] += np.multiply(out_numpy[0], prob_matrix)
pred[1, y1:y2, x1:x2] += prob_matrix
target_reconstructed[y1:y2, x1:x2] = target_crop_tuple[0]
if self.opt.reconstruction_method == 'gaussian':
gaussian = pred[1]
pred = np.divide(pred[0], gaussian)
# pred_sem = np.divide(pred_sem, gaussian)
# st()
if self.opt.dfc_preprocessing == 0:
# resize outputs
pred = np.array(
Image.fromarray(pred).resize(
(pred.shape[1] // 10, pred.shape[0] // 10),
Image.BILINEAR))
target_reconstructed = np.array(
Image.fromarray(target_reconstructed).resize(
(target_reconstructed.shape[1] // 10,
target_reconstructed.shape[0] // 10),
Image.BILINEAR))
error_list.append(np.abs(pred - target_reconstructed))
outG_list.append(np.abs(pred))
return error_list, outG_list, target_reconstructed
def test_raster_notarget(self, load_data):
from dataloader.dataset_bank import dataset_dfc
print('Test phase using {} split.'.format(self.opt.test_split))
phase = 'test'
imageSize = self.opt.imageSize if len(
self.opt.imageSize) == 2 else self.opt.imageSize * 2
test_stride = self.opt.test_stride if len(
self.opt.test_stride) == 2 else self.opt.test_stride * 2
input_list = dataset_dfc(self.opt.dataroot,
data_split=self.opt.test_split,
phase='test',
model=self.opt.model)
# Sanity check : raise an error if some files do not exist
for f in input_list:
if not os.path.isfile(f):
raise KeyError('{} is not a file !'.format(f))
data_loader = [(rgb) for rgb in load_data(
input_list, phase, dfc_preprocessing=self.opt.dfc_preprocessing)]
self.netG.eval()
# create a matrix with a gaussian distribution to be the weights during reconstruction
prob_matrix = self.gaussian_kernel(imageSize[0], imageSize[1])
for it, input in enumerate(tqdm(data_loader)):
rgb_cache = []
pred_gaussian = np.zeros([input.shape[-2], input.shape[-1]])
if self.opt.reconstruction_method == 'gaussian':
pred_sem = np.zeros(
[self.opt.n_classes, input.shape[-2], input.shape[-1]])
else:
pred_sem = np.zeros([input.shape[-2], input.shape[-1]])
target_reconstructed = np.zeros(input.shape[-2:])
# input is a tensor
rgb_cache = [
crop for crop in self.sliding_window_coords(
input, test_stride, imageSize)
]
for input_crop_tuple in tqdm(rgb_cache, total=len(rgb_cache)):
input_crop, (x1, x2, y1, y2) = input_crop_tuple
input_crop = self.get_variable(input_crop)
# ToDo: Deal with padding later
if self.opt.use_padding:
from torch.nn import ReflectionPad2d
self.opt.padding = self.get_padding_image_dims(input_crop)
input_crop = ReflectionPad2d(self.opt.padding)(input_crop)
(pwl, pwr, phu, phb) = self.opt.padding
# target_crop = ReflectionPad2d(self.opt.padding)(target_crop)
with torch.no_grad():
outG_sem = self.netG.forward(input_crop)
if self.opt.reconstruction_method == 'gaussian':
outG_sem_prob = nn.Sigmoid()(outG_sem)
seg_map = outG_sem_prob.cpu().data[0].numpy()
pred_sem[:, y1:y2,
x1:x2] += np.multiply(seg_map, prob_matrix)
pred_gaussian[y1:y2, x1:x2] += prob_matrix
else:
pred_sem[y1:y2,
x1:x2] = np.argmax(outG_sem.cpu().data[0].numpy(),
axis=0)
# visualize
visuals = OrderedDict([
('input', input_crop.data),
('out_sem',
np.argmax(outG_sem.cpu().data[0].numpy(), axis=0))
])
self.display_test_results(visuals)
if self.opt.save_samples:
if self.opt.reconstruction_method == 'gaussian':
pred_sem = np.divide(pred_sem, pred_gaussian)
self.save_raster_images_semantics_only(input,
pred_sem,
index=it + 1,
phase='test')
def test_raster_target(self, load_data):
from dataloader.dataset_bank import dataset_dfc
print('Test phase using {} split.'.format(self.opt.test_split))
phase = 'test'
use_semantics = self.opt.use_semantics
self.opt.use_semantics = False
# self.augmentation = augmentation
imageSize = self.opt.imageSize if len(
self.opt.imageSize) == 2 else self.opt.imageSize * 2
test_stride = self.opt.test_stride if len(
self.opt.test_stride) == 2 else self.opt.test_stride * 2
input_list, target_path = dataset_dfc(
self.opt.dataroot,
data_split=self.opt.test_split,
phase='test',
model=self.opt.model,
which_raster=self.opt.which_raster)
# Sanity check : raise an error if some files do not exist
for f in input_list + target_path:
if not os.path.isfile(f):
raise KeyError('{} is not a file !'.format(f))
data_loader = [(rgb, depth, meta, depth_patch_shape)
for rgb, depth, meta, depth_patch_shape in load_data(
input_list,
target_path,
phase,
self.opt.dfc_preprocessing,
which_raster=self.opt.which_raster,
use_semantics=False,
save_semantics=self.opt.save_semantics)
] # false because we do not have the GT
# no error in save semantics, same value to both variables
self.netG.eval()
# if self.opt.normalize:
# from dataloader.dataset_raster import get_min_max
# import rasterio
# max_v, min_v = get_min_max([rasterio.open(path) for path in target_path], self.opt.which_raster)
# create a matrix with a gaussian distribution to be the weights during reconstruction
prob_matrix = self.gaussian_kernel(imageSize[0], imageSize[1])
# st()
time_array = np.zeros(len(data_loader))
for it, (input, target, meta_data,
depth_patch_shape) in enumerate(tqdm(data_loader)):
rgb_cache = []
depth_cache = []
start = time.time()
# pred = np.zeros(input.shape[-2:])
# concatenate probability matrix
pred = np.zeros([input.shape[-2], input.shape[-1]])
if self.opt.reconstruction_method == 'gaussian':
pred = np.zeros([2, input.shape[-2], input.shape[-1]])
pred_sem = np.zeros(
[self.opt.n_classes, input.shape[-2], input.shape[-1]])
else:
pred_sem = np.zeros([input.shape[-2], input.shape[-1]])
target_reconstructed = np.zeros(input.shape[-2:])
# input is a tensor
rgb_cache = [
crop for crop in self.sliding_window_coords(
input, test_stride, imageSize)
]
depth_cache = [
crop for crop in self.sliding_window_coords(
target, test_stride, imageSize)
] # don't need both
# import cProfile
# torch.cuda.synchronize()
for input_crop_tuple, target_crop_tuple in tqdm(
zip(rgb_cache, depth_cache), total=len(rgb_cache)):
# cp = cProfile.Profile()
# cp.enable()
# for input_crop_tuple, target_crop_tuple in zip(rgb_cache, depth_cache):
input_crop, (x1, x2, y1, y2) = input_crop_tuple
input_crop = self.get_variable(input_crop)
with torch.no_grad():
if 'multitask' in self.opt.model:
outG, outG_sem = self.netG.forward(input_crop)
else:
outG = self.netG.forward(input_crop)[0]
out_numpy = outG.data[0].cpu().float().numpy()
if self.opt.reconstruction_method == 'concatenation':
pred[y1:y2, x1:x2] = out_numpy[0]
elif self.opt.reconstruction_method == 'gaussian':
pred[0, y1:y2,
x1:x2] += np.multiply(out_numpy[0], prob_matrix)
pred[1, y1:y2, x1:x2] += prob_matrix
if self.opt.save_semantics:
# pred_sem[:,y1:y2,x1:x2] += outG_sem.cpu().data[0].numpy()
if self.opt.reconstruction_method == 'gaussian':
# seg_map = np.argmax(outG_sem.cpu().data[0].numpy(), axis=0)
# pred_sem[y1:y2,x1:x2] += np.multiply(seg_map, prob_matrix)
outG_sem_prob = nn.Sigmoid()(outG_sem)
seg_map = outG_sem_prob.cpu().data[0].numpy()
pred_sem[:, y1:y2,
x1:x2] += np.multiply(seg_map, prob_matrix)
else:
pred_sem[y1:y2, x1:x2] = np.argmax(
outG_sem.cpu().data[0].numpy(), axis=0)
# visualize takes a lot of time
# visuals = OrderedDict([('input', input_crop.data),
# # ('gt', target_crop.data),
# ('output', outG),
# # ('gt_sem', self.target_sem.data[0].cpu().float().numpy()),
# ('out_sem', np.argmax(outG_sem.cpu().data[0].numpy(), axis=0))
# ])
# self.display_test_results(visuals)
target_reconstructed[y1:y2, x1:x2] = target_crop_tuple[0]
# break
# st()
# cp.disable()
# cp.print_stats()
# st()
end = time.time()
time_array[it] = end - start
print('Time in seconds: {}'.format(end - start))
t_time = end - start
day = t_time // (24 * 3600)
t_time = t_time % (24 * 3600)
hour = t_time // 3600
t_time %= 3600
minutes = t_time // 60
t_time %= 60
seconds = t_time
print("d:h:m:s-> %d:%d:%d:%d" % (day, hour, minutes, seconds))
if self.opt.save_samples:
if self.opt.reconstruction_method == 'gaussian':
gaussian = pred[1]
pred = np.divide(pred[0], gaussian)
pred_sem = np.divide(pred_sem, gaussian)
# if self.opt.normalize:
# target_reconstructed = (target_reconstructed + 1) / 2
# target_reconstructed = (target_reconstructed * (max_v - min_v)) + min_v
# pred = (pred + 1) / 2
# pred = (pred * (max_v - min_v)) + min_v
# print('Target: max[{}] min[{}]'.format(target_reconstructed.max(),target_reconstructed.min()))
# print('Pred: max[{}] min[{}]'.format(pred.max(),pred.min()))
indexplus = 4
if self.opt.save_semantics:
self.save_raster_images_semantics(
input, pred, target_reconstructed, pred_sem, meta_data,
depth_patch_shape, it + 1 + indexplus, 'test')
else:
self.save_raster_images(input, pred, target_reconstructed,
meta_data, depth_patch_shape,
it + 1 + indexplus, 'test')
del input, pred, target_reconstructed, gaussian, pred_sem
print('Test statistics')
print('Mean and standard deviation in seconds {:.3f} {:.3f}'.format(
np.mean(time_array), np.std(time_array)))
print('Saving merged!')
self.save_merged_rasters('output')
self.save_merged_rasters('target')
self.save_merged_rasters('semantics')
def display_test_results(self, visuals):
self.visualizer.display_images(visuals, 1)
def get_padding(self, dim):
final_dim = (dim // 32 + 1) * 32
return final_dim - dim
def sliding_window_coords(self, data, step, window_size):
from dataloader.dataset_raster import sliding_window
# data = data.data[0].cpu().float().numpy()
for x1, x2, y1, y2 in sliding_window(data, step, window_size):
if len(data.shape) == 2:
yield (data[y1:y2, x1:x2], [x1, x2, y1, y2])
else:
yield (torch.from_numpy(data[:, y1:y2, x1:x2]).unsqueeze(0),
[x1, x2, y1, y2]) # why do I have to unsqueeze here?
def gaussian_kernel(self, width, height, sigma=0.2, mu=0.0):
x, y = np.meshgrid(np.linspace(-1, 1, height),
np.linspace(-1, 1, height))
d = np.sqrt(x * x + y * y)
gaussian_k = (np.exp(-(
(d - mu)**2 / (2.0 * sigma**2)))) / np.sqrt(2 * np.pi * sigma**2)
return gaussian_k # / gaussian_k.sum()
def get_padding_image_dims(self, img):
# get tensor dimensions
h, w = img.size()[2:]
# self.opt.imageSize = (w + 4, h + 4)
w_pad, h_pad = self.get_padding(w + 4) + 4, self.get_padding(h + 4) + 4
pwr = w_pad // 2
pwl = w_pad - pwr
phb = h_pad // 2
phu = h_pad - phb
# pwl, pwr, phu, phb
return (pwl, pwr, phu, phb)
def get_padding_image(self, img):
# get tensor dimensions
h, w = img.size()[2:]
self.opt.imageSize = (w, h)
w_pad, h_pad = self.get_padding(w), self.get_padding(h)
pwr = w_pad // 2
pwl = w_pad - pwr
phb = h_pad // 2
phu = h_pad - phb
# pwl, pwr, phu, phb
return (pwl, pwr, phu, phb)
def save_height_colormap(self, filename, data, cmap='jet'):
import matplotlib.pyplot as plt
plt.switch_backend('agg')
dpi = 80
data = data[0, :, :]
height, width = data.shape
figsize = width / float(dpi), height / float(dpi)
# change string
if 'output' in filename:
filename = filename.replace('merged_output', 'cmap_merged_output')
else:
filename = filename.replace('merged_target', 'cmap_merged_target')
fig = plt.figure(figsize=figsize)
ax = fig.add_axes([0, 0, 1, 1])
ax.axis('off')
cax = ax.imshow(data,
vmax=30,
vmin=0,
aspect='auto',
interpolation='spline16',
cmap=cmap)
ax.set(xlim=[0, width], ylim=[height, 0], aspect=1)
fig.savefig(filename, dpi=dpi)
del fig, data
def save_dsm_as_raster(self, data, filename, meta_data, shape):
import rasterio
filename = os.path.join(self.save_samples_path, filename)
depth_patch = np.expand_dims(np.array(
Image.fromarray(data).resize(shape, Image.BILINEAR)),
axis=0)
# if 'output' in filename or 'target' in filename:
# try:
# self.save_height_colormap(filename, depth_patch, cmap='jet')
# except:
# pass
with rasterio.open(filename, "w", **meta_data) as dest:
if dest.write(depth_patch) == False:
print('Couldnt save image, sorry')
def save_raster_images_semantics(self,
input,
output,
target,
semantics,
meta_data,
shape,
index,
phase='train',
out_type='png'):
from dataloader.dataset_raster import sliding_window
self.save_raster_images(input, output, target, meta_data[0], shape,
index, phase)
del input, output, target
import gc
gc.collect()
filename = '{}/semantics/semantics_{:04}.tif'.format(
self.save_samples_path, index)
if self.opt.reconstruction_method == 'gaussian':
semantics = np.argmax(semantics, axis=0)
semantics = np.array(semantics, dtype=np.uint8)
sem_patch = np.expand_dims(np.array(
Image.fromarray(semantics, mode='P').resize(shape, Image.NEAREST)),
axis=0)
del semantics
import rasterio
with rasterio.open(filename, "w", **meta_data[1]) as dest:
if dest.write(sem_patch) == False:
print('Couldnt save image, sorry')
# base_stride /= 2
del sem_patch
# def save_height_colormap(self, filename, data, cmap='jet'):
# import matplotlib.pyplot as plt
# plt.switch_backend('agg')
# dpi = 80
# data = data[0,:,:]
# height, width = data.shape
# figsize = width / float(dpi), height / float(dpi)
# # change string
# filename = filename.replace('output_', 'cmap_output_')
# fig = plt.figure(figsize=figsize)
# ax = fig.add_axes([0, 0, 1, 1])
# ax.axis('off')
# cax = ax.imshow(data, vmax=30, vmin=0, aspect='auto', interpolation='spline16', cmap=cmap)
# ax.set(xlim=[0, width], ylim=[height, 0], aspect=1)
# fig.savefig(filename, dpi=dpi)
def save_raster_images_semantics_only(self,
input,
semantics,
index,
shape=(1192, 1202),
phase='train',
out_type='png'):
print('Saving semantics...')
import gc
gc.collect()
filename = '{}/semantics/semantics_{:04}.tif'.format(
self.save_samples_path, index)
# 1192
# 1202
import time
start = time.time()
if self.opt.reconstruction_method == 'gaussian':
semantics = np.argmax(semantics, axis=0)
end = time.time()
print('Time to argmax: {}'.format(end - start))
semantics = np.array(semantics, dtype=np.uint8)
sem_patch = np.expand_dims(np.array(
Image.fromarray(semantics, mode='P').resize(shape, Image.NEAREST)),
axis=0)
import rasterio
with rasterio.open(filename, "w", **OUT_META_SEM[index - 1]) as dest:
if dest.write(sem_patch) == False:
print('Couldnt save image, sorry')
def save_raster_images(self,
input,
output,
target,
meta_data,
shape,
index,
phase='train',
out_type='png'):
# self.save_rgb_raster(input.data, '{}/input/input_{:04}.png'.format(self.save_samples_path, index))
self.save_dsm_as_raster(output,
'output/output_{:04}.tif'.format(index),
meta_data, shape)
self.save_dsm_as_raster(target,
'target/target_{:04}.tif'.format(index),
meta_data, shape)
# self.save_dsm_as_raster(target.data, '{}/target/target_{:04}.tif'.format(self.save_samples_path, index), meta_data)
def get_variable(self, tensor):
variable = Variable(tensor)
if self.opt.cuda:
return variable.cuda()
def create_save_folders(
self,
subfolders=['input', 'target', 'results', 'output', 'semantics']):
if self.opt.save_samples:
if self.opt.test:
self.save_samples_path = os.path.join(
'results/{}'.format(self.opt.dataset_name), self.opt.name,
self.opt.epoch)
for subfolder in subfolders:
path = os.path.join(self.save_samples_path, subfolder)
os.system('mkdir -p {0}'.format(path))
