def baselineExperimentNIR(self, nirImages):
loadAndAgumentMasks = makeMasks.MaskClass(self.config,
rand_seed=None,
evaluation=True)
# If we load for each batchsize
# masks = loadAndAgumentMasks.returnTensorMasks(self.batchSize)
now = datetime.now()
dt_string = now.strftime("%d_%m_%Y_%H_%M_%S")
# Defect image over the same region in each color channel
self.output_path = self.imageOutputPath / 'Croatia' / dt_string / 'Data'
self.nir_output_path = self.imageOutputPath / 'Croatia' / dt_string / 'NIRData'
index = 0
start_time = datetime.now()
for real in self.testImages:
# Get masks and make tensor, set to GPU
mask = loadAndAgumentMasks.returnMask(82)
mask = mask[0, :, :]
# Get real and set to GPU
# Augment with masks
# Check if this applies to all three color channels?
real_masked = real.copy()
NIR_real_masked = nirImages[index]
for layer in range(real_masked.shape[-1]):
real_masked[np.where(mask)] = 0
for layer in range(NIR_real_masked.shape[-1]):
NIR_real_masked[np.where(mask)] = 0
NIR_results = inpaint_biharmonic(NIR_real_masked, mask)
results = inpaint_biharmonic(real_masked, mask, multichannel=True)
if self.config.test_mode:
self.storeArrayAsImageNIR(self.output_path, results,
self.nir_output_path, NIR_results,
index)
results = normalize_array(results)
real = normalize_array(real)
real_masked = normalize_array(real_masked)
self.show_images(results,
real,
real_masked,
self.config.run_TCI,
name=self.testImagesName[index])
else:
self.show_images(results, real, real_masked,
self.config.run_TCI)
results = normalize_array(results)
real = normalize_array(real)
real_masked = normalize_array(real_masked)
self.show_images(results, real, real_masked,
self.config.run_TCI)
index = index + 1
end_time = datetime.now()
time_ran = str(end_time - start_time)
print("It took " + str(time_ran))
print("stored at " + str(self.output_path))
return self.output_path, time_ran
def trainGAN(self):
gen = self.generator().to(self.device)
gen_opt = torch.optim.Adam(gen.parameters(),
lr=self.lr,
betas=(self.beta1, self.beta2))
critic = self.critic().to(self.device)
critic_opt = torch.optim.Adam(critic.parameters(),
lr=self.lr,
betas=(self.beta1, self.beta2))
cur_step = 0
loadAndAgumentMasks = makeMasks.MaskClass(self.config, rand_seed=None)
# måske nn.Conv2d med vægte ikke virker når vi bruger partconv2d, i så fald måske tilføje
# or isinstance(m,partConv2d) og læg partconv2d et sted hvor den er accessible.
def weights_init(m):
if isinstance(m, nn.Conv2d) or isinstance(
m, nn.ConvTranspose2d) or isinstance(m, PartialConv2d):
nn.init.kaiming_normal_(m.weight,
mode='fan_out',
nonlinearity='relu')
if m.bias is not None:
nn.init.constant_(m.bias, 0)
elif isinstance(m, nn.BatchNorm2d):
torch.nn.init.constant_(m.weight, 1)
torch.nn.init.constant_(m.bias, 0)
elif isinstance(m, nn.LayerNorm):
torch.nn.init.normal_(m.weight, 0.0, 0.02)
torch.nn.init.constant_(m.bias, 0)
elif isinstance(m, nn.Linear):
nn.init.normal_(m.weight, 0.0, 0.02)
gen = gen.apply(weights_init)
critic = critic.apply(weights_init)
print("Setup loss function...")
loss_func = CalculateLoss(config=self.config).to(self.device)
for epoch in range(self.epochs):
for real, target in tqdm(self.dataloader,
position=0,
leave=True,
disable=self.config.run_polyaxon):
masks = loadAndAgumentMasks.returnTensorMasks(self.batchSize)
masks = torch.from_numpy(masks)
masks = masks.type(torch.cuda.FloatTensor)
masks = 1 - masks
masks.to(self.device)
real = real.to(self.device)
# ---------------------
# Train critic
# ---------------------
critic.zero_grad()
# Real images
real_validity = critic(real)
d_real = real_validity.mean()
# Generate a batch of images with mask
Masked_fake_img = torch.mul(real, masks)
fake_imgs = gen(Masked_fake_img, masks)
# Fake images
fake_validity = critic(fake_imgs) # Detach or not?
d_fake = fake_validity.mean()
gradient_penalty = self.calc_gradient_penalty(
critic, real.data, fake_imgs.data)
d_loss = d_fake - d_real + gradient_penalty
d_loss.backward()
critic_opt.step()
# Values for txt / logging
critic_cost = d_fake - d_real + gradient_penalty
wasserstein_d = d_real - d_fake
critic_score = real_validity.mean().item()
gen_score = fake_validity.mean().item()
# Train the generator every n_critic steps
if cur_step % self.n_critic == 0:
# -----------------
# Train Generator
# -----------------
gen.zero_grad()
# Generate a batch of images
fake_noise = torch.mul(real, masks)
fake_imgs = gen(fake_noise, masks)
# Loss measures generator's ability to fool the critic
# Train on fake images
fake_validity1 = critic(fake_imgs)
loss_dict = loss_func(fake_noise, masks, fake_imgs, real)
loss = 0.0
# sums up each loss value
for key, value in loss_dict.items():
loss += value
loss.backward(retain_graph=True)
g_loss = fake_validity1.mean()
#g_lossMSE = criterionMSE(real, fake_imgs)
#g_lossMSE.backward(retain_graph=True)
g_loss = -g_loss
g_loss.backward() #mone
gen_opt.step()
gen_cost = g_loss
cur_step += 1
if self.config.run_polyaxon and epoch % 5 == 0:
metrics = {}
for key, value in loss_dict.items():
modelHelper.saveMetrics(metrics, key, value.item(),
self.config.polyaxon_experiment,
epoch)
modelHelper.saveMetrics(metrics, 'critic cost', critic_cost,
self.config.polyaxon_experiment, epoch)
modelHelper.saveMetrics(metrics, 'Wasserstein distance',
wasserstein_d,
self.config.polyaxon_experiment, epoch)
modelHelper.saveMetrics(metrics, 'Gen cost', gen_cost,
self.config.polyaxon_experiment, epoch)
if epoch % self.save_model_step == 0 and self.trainMode == True:
name = str(self.modelName) + '_' + str(epoch)
model_path = modelHelper.saveModel(name, self.modelOutputPath,
gen, self.modelName)
if self.config.nir_data:
modelHelper.save_tensor_batch_NIR(
real, Masked_fake_img, fake_imgs, self.batchSize,
Path.joinpath(self.ImageOutputPath,
'epoch_' + str(epoch)))
else:
modelHelper.save_tensor_batch_TCI(
real, Masked_fake_img, fake_imgs, self.batchSize,
Path.joinpath(self.ImageOutputPath,
'epoch_' + str(epoch)))
# Save loss from generator and critic to a file
filename = Path.joinpath(
self.modelOutputPath,
self.modelName + '_' + str(self.batchSize) + 'Errors.txt')
saveString = 'wasserStein Number: ' + str(
wasserstein_d) + ' Generator loss: ' + str(
g_loss.item()) + '\n' + 'critic loss: ' + str(
d_loss.item()
) + '\n' + 'critic guess on reals: ' + str(
critic_score) + ' critic guess on fakes: ' + str(
gen_score
) + ' Updated critic guess on fake: ' + str(
gen_cost) + '\n'
modelHelper.saveToTxt(filename, saveString)
if self.trainWithFreeze:
#trainFrozenModel = trainFrozenGan(self.dataloader,gen,critic,gen_opt,critic_opt, self.config)
#trainFrozenGan.trainGAN()
#Frys BN i encoder parts of the network
#Bruge affine? eller sætte weight og bias til module.eval
for name, module in gen.named_modules():
if isinstance(module, nn.BatchNorm2d) and 'down' in name:
module.eval()
for epoch in range(self.epochsFrozen):
for real, target in tqdm(self.dataloader,
position=0,
leave=True,
disable=self.config.run_polyaxon):
masks = loadAndAgumentMasks.returnTensorMasks(
self.batchSize)
masks = torch.from_numpy(masks)
masks = masks.type(torch.cuda.FloatTensor)
masks = 1 - masks
masks.to(self.device)
real = real.to(self.device)
# ---------------------
# Train critic
# ---------------------
critic.zero_grad()
# Real images
real_validity = critic(real)
d_real = real_validity.mean()
# Generate a batch of images with mask
Masked_fake_img = torch.mul(real, masks)
fake_imgs = gen(Masked_fake_img, masks)
# Fake images
fake_validity = critic(fake_imgs) # Detach or not?
d_fake = fake_validity.mean()
gradient_penalty = self.calc_gradient_penalty(
critic, real.data, fake_imgs.data)
d_loss = d_fake - d_real + gradient_penalty
d_loss.backward()
critic_opt.step()
# Values for txt / logging
critic_cost = d_fake - d_real + gradient_penalty
wasserstein_d = d_real - d_fake
critic_score = real_validity.mean().item()
gen_score = fake_validity.mean().item()
# Train the generator every n_critic steps
if cur_step % self.n_critic == 0:
# -----------------
# Train Generator
# -----------------
gen.zero_grad()
# Generate a batch of images
fake_noise = torch.mul(real, masks)
fake_imgs = gen(fake_noise, masks)
# Loss measures generator's ability to fool the critic
# Train on fake images
fake_validity1 = critic(fake_imgs)
loss_dict = loss_func(fake_noise, masks, fake_imgs,
real)
loss = 0.0
# sums up each loss value
for key, value in loss_dict.items():
loss += value
loss.backward(retain_graph=True)
g_loss = fake_validity1.mean()
# g_lossMSE = criterionMSE(real, fake_imgs)
# g_lossMSE.backward(retain_graph=True)
g_loss = -g_loss
g_loss.backward() # mone
gen_opt.step()
gen_cost = g_loss
cur_step += 1
if self.config.run_polyaxon and epoch % 5 == 0:
metrics = {}
for key, value in loss_dict.items():
modelHelper.saveMetrics(
metrics, key, value.item(),
self.config.polyaxon_experiment,
epoch + self.epochs)
modelHelper.saveMetrics(metrics, 'critic cost',
critic_cost,
self.config.polyaxon_experiment,
epoch + self.epochs)
modelHelper.saveMetrics(metrics, 'Wasserstein distance',
wasserstein_d,
self.config.polyaxon_experiment,
epoch + self.epochs)
modelHelper.saveMetrics(metrics, 'Gen cost', gen_cost,
self.config.polyaxon_experiment,
epoch + self.epochs)
if epoch % self.save_model_step == 0 and self.trainMode == True:
name = str(self.modelName) + '_' + str(epoch + self.epochs)
model_path = modelHelper.saveModel(name,
self.modelOutputPath,
gen, self.modelName)
if self.config.nir_data:
modelHelper.save_tensor_batch_NIR(
real, Masked_fake_img, fake_imgs, self.batchSize,
Path.joinpath(self.ImageOutputPath,
'epoch_' + str(epoch)))
else:
modelHelper.save_tensor_batch_TCI(
real, Masked_fake_img, fake_imgs, self.batchSize,
Path.joinpath(self.ImageOutputPath,
'epoch_' + str(epoch + self.epochs)))
# Save loss from generator and critic to a file
filename = Path.joinpath(
self.modelOutputPath, self.modelName + '_' +
str(self.batchSize) + 'Errors.txt')
saveString = 'wasserStein Number: ' + str(
wasserstein_d) + ' Generator loss: ' + str(g_loss.item(
)) + '\n' + 'critic loss: ' + str(d_loss.item(
)) + '\n' + 'critic guess on reals: ' + str(
critic_score) + ' critic guess on fakes: ' + str(
gen_score
) + ' Updated critic guess on fake: ' + str(
gen_cost) + '\n'
modelHelper.saveToTxt(filename, saveString)
return model_path
def trainGAN(self):
gen = self.generator().to(self.device)
gen_opt = torch.optim.Adam(gen.parameters(), lr=self.lr, betas=(self.beta1, self.beta2))
disc = self.discriminator().to(self.device)
disc_opt = torch.optim.Adam(disc.parameters(), lr=self.lr, betas=(self.beta1, self.beta2))
filename = Path.joinpath(self.modelOutputPath, self.modelName + '_Errors_' + str(self.batchSize) + '.txt')
criterionBCE = nn.BCELoss().cuda()
criterionMSE = nn.L1Loss().cuda()
# Loss function
# Moves vgg16 model to gpu, used for feature map in loss function
loss_func = CalculateLoss(self.config).to(self.device)
print("Setup loss function...")
cur_step = 0
discriminator_loss = []
generator_loss = []
generator_loss_BCE = []
loadAndAgumentMasks = makeMasks.MaskClass(self.config,rand_seed=None)
# måske nn.Conv2d med vægte ikke virker når vi bruger partconv2d, i så fald måske tilføje
# or isinstance(m,partConv2d) og læg partconv2d et sted hvor den er accessible.
def weights_init(m):
if isinstance(m, nn.Conv2d) or isinstance(m, nn.ConvTranspose2d) or isinstance(m,PartialConv2d):
nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
if m.bias is not None:
nn.init.constant_(m.bias, 0)
elif isinstance(m, nn.BatchNorm2d):
torch.nn.init.constant_(m.weight, 1)
torch.nn.init.constant_(m.bias, 0)
elif isinstance(m,nn.LayerNorm):
torch.nn.init.normal_(m.weight, 0.0, 0.02)
torch.nn.init.constant_(m.bias, 0)
def weights_initOld(m):
if isinstance(m, nn.Conv2d) or isinstance(m, nn.ConvTranspose2d) or isinstance(m,PartialConv2d):
torch.nn.init.normal_(m.weight, 0.0, 0.02)
elif isinstance(m, nn.BatchNorm2d):
torch.nn.init.normal_(m.weight, 0.0, 0.02)
torch.nn.init.constant_(m.bias, 0)
elif isinstance(m,nn.LayerNorm):
torch.nn.init.normal_(m.weight, 0.0, 0.02)
torch.nn.init.constant_(m.bias, 0)
gen = gen.apply(weights_init)
disc = disc.apply(weights_init)
for epoch in range(self.epochs):
# Dataloader returns the batches
for real in tqdm(self.dataloader,position=0,leave=True,disable=self.config.run_polyaxon):
masks = loadAndAgumentMasks.returnTensorMasks(self.batchSize)
masks = torch.from_numpy(masks)
masks = masks.type(torch.cuda.FloatTensor)
masks = 1 - masks
masks.to(self.device)
real = real.to(self.device)
## Update discriminator ##
disc.zero_grad()
fake_noise = torch.mul(real, masks)
fake = gen(fake_noise, masks)
disc_fake_pred = disc(fake.detach())
disc_fake_loss = criterionBCE(disc_fake_pred, torch.zeros_like(disc_fake_pred))
disc_real_pred = disc(real)
disc_real_loss = criterionBCE(disc_real_pred, torch.ones_like(disc_real_pred))
disc_loss = (disc_fake_loss + disc_real_loss)
gen_score_fakes = disc_fake_pred.mean().item()
disc_score_reals = disc_real_pred.mean().item()
# Keep track of the average discriminator loss
discriminator_loss.append(disc_loss.item())
# Update gradients
disc_loss.backward()
# Update optimizer
disc_opt.step()
## Update generator ##
gen.zero_grad()
fake_2 = gen(fake_noise, masks)
disc_fake_pred2 = disc(fake_2)
#Calculate loss
gen_lossMSE = criterionMSE(real, fake_2)
gen_loss_Adversarial = criterionBCE(disc_fake_pred2, torch.ones_like(disc_real_pred))
#Add heavy penalty to pixels underneath the mask, ie, try not to make it mode_collaps?
masks = 1-masks
real_masked_area = torch.mul(real,masks)
fake_masked_area = torch.mul(fake_2,masks)
#gen_loss_Inpainted_area = criterionMSE(real_masked_area,fake_masked_area)
gen_loss = gen_lossMSE + gen_loss_Adversarial #+ (gen_loss_Inpainted_area*5)
gen_score_fakes1 = disc_fake_pred2.mean().item()
# få lavet en loss function, der penalizer pixels ændret udenfor maske
# + regner MSE/L1 på alle pixels
loss_dict = loss_func(fake_noise, masks, fake_2, real)
loss = 0.0
# sums up each loss value
for key, value in loss_dict.items():
loss += value
loss.backward()
gen_opt.step()
# Keep track of the average generator loss
generator_loss.append(gen_loss.item())
generator_loss_BCE.append(gen_loss_Adversarial.item())
## Visualization code ##
if cur_step % self.save_model_step == 0 and cur_step > 0 and self.trainMode == False:
#if not training, it means we are messing around testing stuff, so no need to save model
#and losses
print(
f"Step {cur_step}: Generator loss: {gen_loss.item()}, discriminator loss: {disc_loss.item()}")
# Save loss from generator and discriminator to a file, and reset them, to avoid the list perpetually growing
# Name of file = model name + batch_size +
discriminator_loss = [sum(discriminator_loss) / len(discriminator_loss)]
generator_loss = [sum(generator_loss) / len(generator_loss)]
generator_loss_BCE = [sum(generator_loss_BCE) / len(generator_loss_BCE)]
self.show_tensor_images(fake_2, real, fake_noise)
#If in train mode, it should not display images at xx display steps, but only save the model and
#and losses during training
cur_step += 1
if self.config.run_polyaxon and epoch % 5 == 0:
metrics = {}
#modelHelper.saveMetrics(metrics,'G_loss',generator_loss[-1],self.config.polyaxon_experiment,epoch)
#modelHelper.saveMetrics(metrics, 'G_BCE_loss', generator_loss_BCE[-1],self.config.polyaxon_experiment,epoch)
#modelHelper.saveMetrics(metrics,'D_loss',discriminator_loss[-1],self.config.polyaxon_experiment,epoch)
for key,value in loss_dict.items():
modelHelper.saveMetrics(metrics, key, value.item(), self.config.polyaxon_experiment, epoch)
modelHelper.saveMetrics(metrics, 'Disc guess on reals', disc_score_reals, self.config.polyaxon_experiment,epoch)
modelHelper.saveMetrics(metrics, 'Disc guess on fakes', gen_score_fakes, self.config.polyaxon_experiment,
epoch)
modelHelper.saveMetrics(metrics, 'Updated disc guess on fakes', gen_score_fakes1, self.config.polyaxon_experiment,
epoch)
if epoch % self.save_model_step == 0 and self.trainMode == True:
saveString = 'Epoch Number: ' + str(epoch) + ' Generator loss: ' + str(generator_loss[-1]) + '\n' + 'Generator loss BCE: ' + str(generator_loss_BCE[-1]) + '\n' + 'Discriminator loss: ' + str(discriminator_loss[-1]) + '\n' + 'Disc guess on reals: ' + str(disc_score_reals) + ' Disc guess on fakes: ' + str(gen_score_fakes) + ' Updated disc guess on fakes: ' + str(gen_score_fakes1) + '\n'
modelHelper.saveToTxt(filename, saveString)
name = str(self.modelName) + '_' + str(epoch)
model_path = modelHelper.saveModel(name, self.modelOutputPath, gen, self.modelName)
modelHelper.save_tensor_batch(real, fake_noise, fake_2, self.batchSize,
Path.joinpath(self.ImageOutputPath, 'epoch_' + str(epoch)))
elif epoch % self.save_error_step == 0 and self.trainMode == True:
saveString = 'Epoch Number: ' + str(epoch) +'\n' + ' Generator loss: ' + str(
generator_loss[-1]) + '\n' + 'Generator loss BCE: ' + str(
generator_loss_BCE[-1]) + '\n' + 'Discriminator loss: ' + str(
discriminator_loss[-1]) + '\n' + 'Disc guess on reals: ' + str(disc_score_reals) + ' Disc guess on fakes: ' + str(
gen_score_fakes) + ' Updated disc guess on fakes: ' + str(gen_score_fakes1) + '\n'
modelHelper.saveToTxt(filename, saveString)
return model_path
def trainTemporalGAN(self):
gen = self.generator().to(self.device)
gen_opt = torch.optim.Adam(gen.parameters(), lr=self.lr, betas=(self.beta1, self.beta2))
disc = self.discriminator().to(self.device)
disc_opt = torch.optim.Adam(disc.parameters(), lr=self.lr, betas=(self.beta1, self.beta2))
#display_step = 4
criterionBCE = nn.BCELoss().cuda()
criterionMSE = nn.MSELoss().cuda()
#display_step = 5
cur_step = 0
discriminator_loss = []
generator_loss = []
generator_loss_BCE = []
loadAndAgumentMasks = makeMasks.MaskClass(rand_seed=None)
# måske nn.Conv2d med vægte ikke virker når vi bruger partconv2d, i så fald måske tilføje
# or isinstance(m,partConv2d) og læg partconv2d et sted hvor den er accessible.
def weights_init(m):
if isinstance(m, nn.Conv2d) or isinstance(m, nn.ConvTranspose2d):
torch.nn.init.normal_(m.weight, 0.0, 0.02)
if isinstance(m, nn.BatchNorm2d):
torch.nn.init.normal_(m.weight, 0.0, 0.02)
torch.nn.init.constant_(m.bias, 0)
gen = gen.apply(weights_init)
disc = disc.apply(weights_init)
for epoch in range(self.epochs):
# Dataloader returns the batches
for temp0,temp1,temp2,temp3,temp4 in tqdm(self.dataloader):
masks0 = loadAndAgumentMasks.returnTensorMasks(self.batchSize)
# masksInverted = 1-masks
# masksInverted = torch.from_numpy(masksInverted)
# masksInverted = masksInverted.type(torch.cuda.FloatTensor)
# masksInverted.to(self.device)
masks0 = torch.from_numpy(masks0)
masks0 = masks0.type(torch.cuda.FloatTensor)
masks0 = 1 - masks0
masks1 = loadAndAgumentMasks.returnTensorMasks(self.batchSize)
masks1 = torch.from_numpy(masks1)
masks1 = masks1.type(torch.cuda.FloatTensor)
masks1 = 1 - masks1
masks2 = loadAndAgumentMasks.returnTensorMasks(self.batchSize)
masks2 = torch.from_numpy(masks2)
masks2 = masks2.type(torch.cuda.FloatTensor)
masks2 = 1 - masks2
masks3 = loadAndAgumentMasks.returnTensorMasks(self.batchSize)
masks3 = torch.from_numpy(masks3)
masks3 = masks3.type(torch.cuda.FloatTensor)
masks3 = 1 - masks3
masks4 = loadAndAgumentMasks.returnTensorMasks(self.batchSize)
masks4 = torch.from_numpy(masks4)
masks4 = masks4.type(torch.cuda.FloatTensor)
masks4 = 1 - masks4
masks =torch.cat((masks0,masks1,masks2,masks3,masks4),1).to(self.device)
real = torch.cat((temp0[0],temp1[0],temp2[0],temp3[0],temp4[0],),1).to(self.device)
#real = real.to(self.device)
#t = torch.cuda.get_device_properties(0).total_memory
#c = torch.cuda.memory_cached(0)
#a = torch.cuda.memory_allocated(0)
#print(t)
#print(c)
#print(a)
## Update discriminator ##
disc_opt.zero_grad()
# lav om så den kører på masker
fake_noise = torch.mul(real, masks)
fake = gen(fake_noise, masks)
disc_fake_pred = disc(fake.detach())
disc_fake_loss = criterionBCE(disc_fake_pred, torch.zeros_like(disc_fake_pred))
disc_real_pred = disc(real)
disc_real_loss = criterionBCE(disc_real_pred, torch.ones_like(disc_real_pred))
disc_loss = (disc_fake_loss + disc_real_loss) / 2
# Keep track of the average discriminator loss
discriminator_loss.append(disc_loss.item())
# Update gradients
disc_loss.backward(retain_graph=True)
# Update optimizer
disc_opt.step()
## Update generator ##
gen_opt.zero_grad()
# fake_noise_2 = real*masksInverted
fake_2 = gen(fake_noise, masks)
disc_fake_pred = disc(fake_2)
gen_lossMSE = criterionMSE(real, fake_2)
gen_loss_Adversarial = criterionBCE(disc_fake_pred, torch.ones_like(disc_real_pred))
gen_loss = gen_lossMSE + gen_loss_Adversarial
# få lavet en loss function, der penalizer pixels ændret udenfor maske
# + regner MSE/L1 på alle pixels
gen_loss.backward()
gen_opt.step()
# Keep track of the average generator loss
generator_loss.append(gen_loss.item())
generator_loss_BCE.append(gen_loss_Adversarial.item())
## Visualization code ##
if cur_step % self.save_model_step == 0 and cur_step > 0 and self.trainMode == False:
#if not training, it means we are messing around testing stuff, so no need to save model
#and losses
print(
f"Step {cur_step}: Generator loss: {gen_loss.item()}, discriminator loss: {disc_loss.item()}")
# Save loss from generator and discriminator to a file, and reset them, to avoid the list perpetually growing
# Name of file = model name + batch_size +
discriminator_loss = [sum(discriminator_loss) / len(discriminator_loss)]
generator_loss = [sum(generator_loss) / len(generator_loss)]
generator_loss_BCE = [sum(generator_loss_BCE) / len(generator_loss_BCE)]
self.show_tensor_images(fake_2, real, fake_noise)
#If in train mode, it should not display images at xx display steps, but only save the model and
#and losses during training
elif cur_step % self.save_model_step == 0 and cur_step > 0 and self.trainMode == True:
#save model
torch.save(gen.state_dict(),
Path.joinpath(self.modelOutputPath, self.modelName + '_' + str(epoch) + '.pt'))
# Save loss from generator and discriminator to a file, and reset them, to avoid the list perpetually growing
# Name of file = model name + batch_size +
discriminator_loss = [sum(discriminator_loss) / len(discriminator_loss)]
generator_loss = [sum(generator_loss) / len(generator_loss)]
generator_loss_BCE = [sum(generator_loss_BCE)/len(generator_loss_BCE)]
self.saveToTxt(generator_loss_BCE, generator_loss_BCE, discriminator_loss)
cur_step += 1
def run_eval(self,
output_path,
store_path,
model_path=None,
test_dataloader=None):
curdatLayer = importData(self.config)
if test_dataloader is None:
train, test_dataloader = curdatLayer.getRGBDataLoader()
del train
if Path.exists(
Path.joinpath(
output_path,
self.config.model_name + '_' + str(self.config.epochs) +
'.pt')) and self.config.run_polyaxon == False:
##Hvis det er med wgan generator, altså layernorm, indsæt Wgangenerator istedet for generator()
if self.config.new_generator:
gen = Wgangenerator().to(self.config.device)
else:
gen = generator().to(self.config.device)
gen.load_state_dict(
torch.load(
Path.joinpath(
output_path, self.config.model_name + '_' +
str(self.config.epochs) +
'.pt'))) ## Use epochs to identify model number
elif Path.exists(Path(str(model_path))):
if self.config.nir_data:
gen = generatorNIR().to(self.config.device)
elif self.config.new_generator:
gen = Wgangenerator().to(self.config.device)
else:
gen = generator().to(self.config.device)
print("Just loaded model from path " + str(model_path))
gen.load_state_dict(
torch.load(model_path)) ## Use epochs to identify model number
else:
print("Unable to find path to model")
gen.eval()
loadAndAgumentMasks = makeMasks.MaskClass(self.config,
rand_seed=None,
evaluation=True)
names = []
# Find names of test images, in order to save the generated files with same name, for further reference
localImg = test_dataloader.dataset.image_list
# Slice string to only include the name of the file, ie after the last //
localNames = []
# if self.config.run_polyaxon:
# split_path = localImg[0].split('/') ##Linux
# else:
# split_path = localImg[0].split("\\")
# local_index= split_path.index('processed')
# local_country= split_path[local_index+1]
for i in localImg:
if self.config.run_polyaxon:
selected_image = i.split('/')[-1] ##Linux
else:
selected_image = i.split("\\")[-1]
localNames.append(selected_image)
names = names + localNames
print("Found this many names " + str(len(names)))
current_number = 0
if not os.path.exists(
Path.joinpath(output_path, self.config.model_name)):
os.makedirs(Path.joinpath(output_path, self.config.model_name))
now = datetime.now()
dt_string = now.strftime("%d_%m_%Y_%H_%M_%S")
local_test_path = output_path / self.config.model_name / dt_string / 'Data'
local_test_nir_path = output_path / self.config.model_name / dt_string / 'DataNir'
local_store_path = store_path / self.config.model_name / dt_string / 'stored_Data'
os.makedirs(local_test_path)
os.makedirs(local_store_path)
os.makedirs(local_test_nir_path)
start_time = datetime.now()
for real, sar in tqdm(test_dataloader,
disable=self.config.run_polyaxon):
masks = loadAndAgumentMasks.returnTensorMasks(
self.config.batch_size)
masks = torch.from_numpy(masks)
masks = masks.type(torch.cuda.FloatTensor)
masks = 1 - masks
masks.to(self.config.device)
real = real.to(self.config.device)
fake_masked_images = torch.mul(real, masks)
generated_images = gen(fake_masked_images, masks)
image_names = names[current_number:current_number +
self.config.batch_size]
current_number = current_number + self.config.batch_size ## Change naming to include all names
# modelHelper.save_tensor_batch(generated_images,fake_masked_images,config.batch_size,path)
for index, image in enumerate(generated_images):
namePath = Path.joinpath(local_test_path, image_names[index])
if self.config.nir_data:
modelHelper.save_tensor_single_NIR(
image,
Path.joinpath(local_test_path, image_names[index]),
Path.joinpath(local_test_nir_path, image_names[index]),
raw=True)
else:
modelHelper.save_tensor_single(image,
Path.joinpath(
local_test_path,
image_names[index]),
raw=True)
end_time = datetime.now()
time_ran = str(end_time - start_time)
# create dataloader with generated images
generated_images_dataloader = curdatLayer.getGeneratedImagesDataloader(
local_test_path)
#print("generated image 429 "+str(generated_images_dataloader.dataset.image_list[429]))
#print ("test image 429 "+str(test_dataloader.dataset.image_list[429]))
# calculate FID
if self.config.nir_data: # Loader test_dataloader in for NIR igen da den skal have 3 channels og ikke 4.
train, test_dataloader = curdatLayer.getRGBDataLoader()
del train
FID_Value = FIDCalculator(
test_dataloader, generated_images_dataloader,
len(test_dataloader) * self.config.batch_size,
self.config.batch_size, self.config).get_FID_scores()
# Calculate PSNR and SSIM
dataloader_iterator = iter(generated_images_dataloader)
#dataloader_iterator = iter(test_dataloader)
maeValues = []
sddValues = []
ssimscikitValues = []
SSIMValues = []
psnrValues = []
CCValues = []
rmseValues = []
# loop to calculate PSNR and SSIM for all test and generated images.
count = 0
for i, images_real in enumerate(test_dataloader):
try:
images_generated = next(dataloader_iterator)
except StopIteration:
dataloader_iterator = iter(generated_images_dataloader)
images_generated = next(dataloader_iterator)
for index2 in range(self.config.batch_size):
psnrValues.append(PSNR().__call__(images_real[index2],
images_generated[index2]))
if psnrValues[-1] < 3:
print(str(psnrValues[-1]))
modelHelper.save_tensor_single(
normalize_batch_tensor(images_real[index2]),
Path.joinpath(local_store_path,
str(i) + '_' + str(count) +
'_real.tiff'))
modelHelper.save_tensor_single(
normalize_batch_tensor(images_generated[index2]),
Path.joinpath(local_store_path,
str(i) + '_' + str(count) + 'gen.tiff'))
CCValues.append(CC().__call__(images_real[index2],
images_generated[index2]))
maeValues.append(MSE().__call__(images_real[index2],
images_generated[index2]))
sddValues.append(
SDD.__call__(images_real[index2],
images_generated[index2]))
ssimscikitValues.append(
SSIM_SKI.__call__(images_real[index2],
images_generated[index2]))
image1 = images_real[index2].unsqueeze(0)
image2 = images_generated[index2].unsqueeze(0)
SSIMValues.append(ssim(image1, image2))
rmseValues.append(
RMSE.__call__(images_real[index2],
images_generated[index2]))
count = count + 1
meanMAE = sum(maeValues) / len(maeValues)
minMAE = min(maeValues)
maxMAE = max(maeValues)
meanSDD = sum(sddValues) / len(sddValues)
minSDD = min(sddValues)
maxSDD = max(sddValues)
meanPSNR = sum(psnrValues) / len(psnrValues)
minPSNR = min(psnrValues)
maxPSNR = max(psnrValues)
meanSSIM = sum(SSIMValues) / len(SSIMValues)
minSSIM = min(SSIMValues)
maxSSIM = max(SSIMValues)
meanSCISSIM = sum(ssimscikitValues) / len(ssimscikitValues)
minSCISSIM = min(ssimscikitValues)
maxSCISSIM = max(ssimscikitValues)
meanCC = sum(CCValues) / len(CCValues)
minCC = min(CCValues)
maxCC = max(CCValues)
meanRMSE = sum(rmseValues) / len(rmseValues)
minRMSE = min(rmseValues)
maxRMSE = max(rmseValues)
# Save final results of evaluation metrics
saveEvalToTxt(self.config.model_name, meanMAE, minMAE, maxMAE, meanSDD,
minSDD, maxSDD, meanSSIM.item(), minSSIM.item(),
maxSSIM.item(), meanSCISSIM, minSCISSIM, maxSCISSIM,
meanPSNR, minPSNR, maxPSNR, meanCC, minCC, maxCC,
meanRMSE, minRMSE, maxRMSE, FID_Value, time_ran,
local_store_path)
# Clean
modelHelper.clearFolder(local_test_path.parent)
def main(args):
config = TrainingConfig()
config = update_config(args, config)
logger = logging.getLogger(__name__)
if config.run_polyaxon:
input_root_path = Path(get_data_paths()['data'])
output_root_path = Path(get_outputs_path())
inpainting_data_path = input_root_path / 'inpainting'
os.environ['TORCH_HOME'] = str(input_root_path / 'pytorch_cache')
config.data_path=inpainting_data_path
config.output_path=output_root_path
imageOutputPath = config.data_path /'data' /'generated'
model_path =inpainting_data_path /'models'
modelOutputPath = Path.joinpath(model_path, 'OutputModels')
stores_output_path = config.output_path /'data'/'storedData'
else:
imageOutputPath = Path().absolute().parent /'data' /'generated'
localdir = Path().absolute().parent
modelOutputPath = Path.joinpath(localdir, 'OutputModels')
stores_output_path = localdir /'data'/'storedData'
#Import test data
test = eval_model(config)
test.run_eval(modelOutputPath,stores_output_path)
curdatLayer = importData(config)
train, test_dataloader = curdatLayer.getRGBDataLoader()
del train
test = Path.joinpath(modelOutputPath, config.model_name + '_'+str(config.epochs) + '.pt')
print(Path.joinpath(modelOutputPath, config.model_name + '_'+str(config.epochs) + '.pt'))
if Path.exists(Path.joinpath(modelOutputPath, config.model_name + '_'+str(config.epochs) + '.pt')):
##Hvis det er med wgan generator, altså layernorm, indsæt Wgangenerator istedet for generator()
gen = generator().to(config.device)
gen.load_state_dict(torch.load(Path.joinpath(modelOutputPath, config.model_name + '_'+str(config.epochs) + '.pt'))) ## Use epochs to identify model number
else:
print("Unable to find path to model")
gen.eval()
loadAndAgumentMasks = makeMasks.MaskClass(config,rand_seed=None,evaluation=True)
names = []
for i in range(len(test_dataloader.dataset.datasets)):
# Find names of test images, in order to save the generated files with same name, for further reference
localImg = test_dataloader.dataset.datasets[i].image_list
# Slice string to only include the name of the file, ie after the last //
localNames = []
for i in localImg:
if config.run_polyaxon:
selected_image=i.split('/')[-1] ##Linux
else:
selected_image=i.split("\\")[-1]
localNames.append(selected_image)
names=names+localNames
print("Found this many names "+str(len(names)))
current_number = 0
if not os.path.exists(Path.joinpath(imageOutputPath, config.model_name)):
os.makedirs(Path.joinpath(imageOutputPath, config.model_name))
now = datetime.now()
dt_string = now.strftime("%d_%m_%Y_%H_%M_%S")
local_test_path= imageOutputPath / config.model_name / dt_string /'Data'
local_store_path = stores_output_path / config.model_name / dt_string /'stored_Data'
os.makedirs(local_test_path)
os.makedirs(local_store_path)
start_time = datetime.now()
testCount = 3
for real in tqdm(test_dataloader):
masks = loadAndAgumentMasks.returnTensorMasks(config.batch_size)
masks = torch.from_numpy(masks)
masks = masks.type(torch.cuda.FloatTensor)
masks = 1 - masks
masks.to(config.device)
real = real.to(config.device)
fake_masked_images = torch.mul(real, masks)
generated_images = gen(fake_masked_images, masks)
image_names = names[current_number:current_number+config.batch_size]
current_number = current_number + config.batch_size ## Change naming to include all names
#modelHelper.save_tensor_batch(generated_images,fake_masked_images,config.batch_size,path)
for index, image in enumerate(generated_images):
namePath= Path.joinpath(local_test_path,image_names[index])
modelHelper.save_tensor_single(image,Path.joinpath(local_test_path, image_names[index]),raw=True)
if testCount<0:
break
testCount = testCount-1
print("Saved image to " +str(local_test_path))
end_time = datetime.now()
time_ran = str(end_time - start_time)
#create dataloader with generated images
generated_images_dataloader = curdatLayer.getGeneratedImagesDataloader(local_test_path)
#calculate FID
FID_Value = FIDCalculator(test_dataloader,generated_images_dataloader, len(test_dataloader)*config.batch_size, config.batch_size,config).get_FID_scores()
#Calculate PSNR and SSIM
dataloader_iterator = iter(generated_images_dataloader)
psnrValues = []
maeValues = []
sddValues=[]
##ssimValues=[]
SSIMValues = []
CCValues = []
#loop to calculate PSNR and SSIM for all test and generated images.
for images_real in test_dataloader:
try:
images_generated = next(dataloader_iterator)
except StopIteration:
dataloader_iterator = iter(generated_images_dataloader)
images_generated = next(dataloader_iterator)
for index2 in range(config.batch_size):
psnrValues.append(PSNR().__call__(images_real[index2], images_generated[index2]))
##CCValues.append(CC().__call__(images_real[index2], images_generated[index2]))
maeValues.append(MSE().__call__(images_real[index2], images_generated[index2]))
sddValues.append(SDD.__call__(images_real[index2], images_generated[index2]))
##ssimValues.append(SSIM.__call__(images_real[index2], images_generated[index2]))
image1 = images_real[index2].unsqueeze(0)
image2 = images_generated[index2].unsqueeze(0)
SSIMValues.append(ssim(image1, image2))
break
meanMAE= sum(maeValues)/len(maeValues)
minMAE = min(maeValues)
maxMAE = max(maeValues)
meanSDD = sum(sddValues) / len(sddValues)
minSDD = min(sddValues)
maxSDD = max(sddValues)
meanPSNR = sum(psnrValues)/len(psnrValues)
minPSNR = min(psnrValues)
maxPSNR = max(psnrValues)
meanSSIM = sum(SSIMValues) / len(SSIMValues)
minSSIM = min(SSIMValues)
maxSSIM = max(SSIMValues)
meanCC = sum(CCValues) / len(CCValues)
minCC = min(CCValues)
maxCC = max(CCValues)
#Save final results of evaluation metrics
saveEvalToTxt(config.model_name,meanMAE,minMAE,maxMAE,meanSSIM,minSDD,maxSDD,meanSSIM,minSSIM,maxSSIM,PSNR.item(),minPSNR.item(),maxPSNR.item(),meanCC.item(),minCC.item(),maxCC.item(),FID_Value,time_ran, local_store_path)
#Clean
modelHelper.clearFolder(local_test_path.parent)
def main(args):
""" Runs dataLayer processing scripts to turn raw dataLayer from (../raw) into
cleaned dataLayer ready to be analyzed (saved in ../processed).
"""
os.environ["KMP_DUPLICATE_LIB_OK"] = "TRUE"
## Talk to Rune about how dataLayer is handle.
config = TrainingConfig()
config = update_config(args,config)
## For polyaxon
#if config.run_polyaxon:
localdir = Path().absolute().parent
dataPath = Path.joinpath(localdir, 'data\ImagesForVisualization')
logger = logging.getLogger(__name__)
logger.info('making final dataLayer set from raw dataLayer')
curdatLayer = importData(config)
## Original
pathToNIR = r"E:\Speciale\NDVIExperiment\Croatia\Original_Data\NIR"
nir_images = curdatLayer.open_Imagefiles_as_array(pathToNIR)
nir_image = nir_images[0]
pathtoRGB = r"E:\Speciale\NDVIExperiment\Croatia\Original_Data\RGB"
rgb_images=curdatLayer.open_Imagefiles_as_array(pathtoRGB)
rgb_image= rgb_images[0]
r,g,b = cv2.split(rgb_image)
org_ndvi = (nir_image - r) / (nir_image + r)
titles = ["Sentinel 2 - Normalized Difference Vegetation Index (NDVI) over Original"]
# https://earthpy.readthedocs.io/en/latest/gallery_vignettes/plot_calculate_classify_ndvi.html
# Turn off bytescale scaling due to float values for NDVI
ep.plot_bands(org_ndvi, cmap="RdYlGn", cols=1, title=titles, vmin=-1, vmax=1)
## Inpainted
pathToNIR = r"E:\Speciale\NDVIExperiment\Croatia\PartialConvolutions\big_mask\DataNir"
nir_images = curdatLayer.open_Imagefiles_as_array(pathToNIR)
nir_image = nir_images[0]
pathtoRGB = r"E:\Speciale\NDVIExperiment\Croatia\PartialConvolutions\big_mask\Data"
rgb_images = curdatLayer.open_Imagefiles_as_array(pathtoRGB)
rgb_image = rgb_images[0]
r, g, b = cv2.split(rgb_image)
gen_ndvi = (nir_image - r) / (nir_image + r)
titles = ["Sentinel 2- Normalized Difference Vegetation Index (NDVI) over generated"]
# https://earthpy.readthedocs.io/en/latest/gallery_vignettes/plot_calculate_classify_ndvi.html
# Turn off bytescale scaling due to float values for NDVI
ep.plot_bands(gen_ndvi, cmap="RdYlGn", cols=1, title=titles, vmin=-1, vmax=1)
diff_ndvi=org_ndvi-gen_ndvi
old= sum(gen_ndvi.flatten())
new = sum(org_ndvi.flatten())
diffSumsWithMaria = ((new-old)/old)
diff_percent_sum = sum((gen_ndvi.flatten()-org_ndvi.flatten())/org_ndvi.flatten()*100)
print("The NDVI have changed " +str(diffSumsWithMaria)+" %")
titles = ["Sentinel 2- Normalized Difference Vegetation Index (NDVI) difference"]
# https://earthpy.readthedocs.io/en/latest/gallery_vignettes/plot_calculate_classify_ndvi.html
# Turn off bytescale scaling due to float values for NDVI
ep.plot_bands(diff_ndvi, cmap="RdYlGn", cols=1, title=titles, vmin=-1, vmax=1)
loadAndAgumentMasks = makeMasks.MaskClass(config, rand_seed=None, evaluation=True)
mask = loadAndAgumentMasks.returnMask(787)
mask = mask[0, :, :]
# Get real and set to GPU
#Invert
mask = 1-mask
# Augment with masks
# Check if this applies to all three color channels?
gen_ndvi_masked = gen_ndvi.copy()
org_ndvi_masked = org_ndvi.copy()
for layer in range(gen_ndvi_masked.shape[-1]):
gen_ndvi_masked[np.where(mask)] = 0
for layer in range(org_ndvi_masked.shape[-1]):
org_ndvi_masked[np.where(mask)] = 0
ep.plot_bands(gen_ndvi_masked, cmap="RdYlGn", cols=1, title=titles, vmin=-1, vmax=1)
maeValues = []
sddValues = []
ssimscikitValues = []
psnrValues = []
CCValues = []
rmseValues = []
org_ndvi_masked=org_ndvi_masked[org_ndvi_masked!=0]
gen_ndvi_masked=gen_ndvi_masked[gen_ndvi_masked!=0]
psnrValues.append(PSNR().__call__(org_ndvi_masked, gen_ndvi_masked,tensor=False))
CCValues.append(CC().__call__(org_ndvi_masked, gen_ndvi_masked,tensor=False))
maeValues.append(MSE().__call__(org_ndvi_masked, gen_ndvi_masked,tensor=False))
sddValues.append(SDD.__call__(org_ndvi_masked, gen_ndvi_masked,tensor=False))
#ssimscikitValues.append(SSIM_SKI.__call__(org_ndvi_masked, gen_ndvi_masked,tensor=False))
rmseValues.append(RMSE.__call__(org_ndvi_masked, gen_ndvi_masked,tensor=False))
meanMAE = sum(maeValues) / len(maeValues)
minMAE = min(maeValues)
maxMAE = max(maeValues)
meanSDD = sum(sddValues) / len(sddValues)
minSDD = min(sddValues)
maxSDD = max(sddValues)
meanPSNR = sum(psnrValues) / len(psnrValues)
minPSNR = min(psnrValues)
maxPSNR = max(psnrValues)
meanSSIM = 0
minSSIM = 0
maxSSIM = 0
meanSCISSIM = 0
minSCISSIM = 0
maxSCISSIM = 0
# meanSCISSIM = sum(ssimscikitValues) / len(ssimscikitValues)
# minSCISSIM = min(ssimscikitValues)
# maxSCISSIM = max(ssimscikitValues)
meanCC = sum(CCValues) / len(CCValues)
minCC = min(CCValues)
maxCC = max(CCValues)
meanRMSE = sum(rmseValues) / len(rmseValues)
minRMSE = min(rmseValues)
maxRMSE = max(rmseValues)
FID_Value=0.0
time_ran=0.0
local_store_path=Path(r"E:\Speciale\NDVIExperiment\Croatia")
saveEvalToTxt(config.model_name, meanMAE, minMAE, maxMAE, meanSDD, minSDD, maxSDD, meanSSIM,
minSSIM,
maxSSIM, meanSCISSIM, minSCISSIM, maxSCISSIM, meanPSNR, minPSNR, maxPSNR, meanCC, minCC, maxCC,
meanRMSE, minRMSE, maxRMSE, FID_Value, time_ran, local_store_path)
# v_nir_image = np.concatenate((nir_images[0], nir_images[1]), axis=1)
# v_rgb_image = np.concatenate((rgb_images[0], rgb_images[1]), axis=1)
# v_r, v_g, v_b = cv2.split(v_rgb_image)
#
# v_ndvi = (v_nir_image - v_r) / (v_nir_image + v_r)
# titles = ["Sentinel 2- Normalized Difference Vegetation Index (NDVI) over two samples"]
# # https://earthpy.readthedocs.io/en/latest/gallery_vignettes/plot_calculate_classify_ndvi.html
# # Turn off bytescale scaling due to float values for NDVI
# ep.plot_bands(v_ndvi, cmap="RdYlGn", cols=1, title=titles, vmin=-1, vmax=1)
# nir_images= nir_images/10000
# r = r/10000
# ndvi = (nir_images - r) / (nir_images + r)
#
# fig = plt.figure(figsize=(10, 10))
# fig.set_facecolor('white')
# plt.imshow(ndvi, cmap='RdYlGn') # Typically the color map for NDVI maps are the Red to Yellow to Green
# plt.title('NDVI')
# plt.show()
#
lol = ""
def main(args):
""" Runs dataLayer processing scripts to turn raw dataLayer from (../raw) into
cleaned dataLayer ready to be analyzed (saved in ../processed).
"""
## Talk to Rune about how dataLayer is handle.
config = TrainingConfig()
config = update_config(args, config)
## For polyaxon
if config.run_polyaxon:
input_root_path = Path(get_data_paths()['data'])
output_root_path = Path(get_outputs_path())
inpainting_data_path = input_root_path / 'inpainting'
os.environ['TORCH_HOME'] = str(input_root_path / 'pytorch_cache')
config.data_path = inpainting_data_path
config.output_path = output_root_path
config.polyaxon_experiment = Experiment()
logger = logging.getLogger(__name__)
logger.info('making final dataLayer set from raw dataLayer')
curdatLayer = importData(config)
if config.nir_data:
train, test_dataloader = curdatLayer.getNIRDataLoader()
else:
train, test_dataloader = curdatLayer.getRGBDataLoader()
local_model_path = r"C:\Users\panda\PycharmProjects\Image_Inpainting_Sat\Master_Satelite_Image_Inpainting\OutputModels\PartialConvolutionsWgan_301.pt"
local_output_path = Path(r"E:\Speciale\final_model")
#gen = Wgangenerator().to(config.device)
if config.nir_data:
gen = generatorNIR().to(config.device)
else:
gen = generator().to(config.device)
gen.load_state_dict(
torch.load(local_model_path)) ## Use epochs to identify model number
gen.eval()
loadAndAgumentMasks = makeMasks.MaskClass(config,
rand_seed=None,
evaluation=True,
noFlip=True)
names = []
# Find names of test images, in order to save the generated files with same name, for further reference
localImg = test_dataloader.dataset.image_list
# Slice string to only include the name of the file, ie after the last //
localNames = []
# if self.config.run_polyaxon:
# split_path = localImg[0].split('/') ##Linux
# else:
# split_path = localImg[0].split("\\")
# local_index= split_path.index('processed')
# local_country= split_path[local_index+1]
for i in localImg:
if config.run_polyaxon:
selected_image = i.split('/')[-1] ##Linux
else:
selected_image = i.split("\\")[-1]
localNames.append(selected_image)
names = names + localNames
print("Found this many names " + str(len(names)))
current_number = 0
if not os.path.exists(Path.joinpath(local_output_path, config.model_name)):
os.makedirs(Path.joinpath(local_output_path, config.model_name))
now = datetime.now()
dt_string = now.strftime("%d_%m_%Y_%H_%M_%S")
local_test_path = local_output_path / config.model_name / dt_string / 'Data'
local_test_nir_path = local_output_path / config.model_name / dt_string / 'DataNir'
local_store_path = local_output_path / config.model_name / dt_string / 'stored_Data'
os.makedirs(local_test_path)
if config.nir_data:
os.makedirs(local_test_nir_path)
start_time = datetime.now()
for real in tqdm(test_dataloader, disable=config.run_polyaxon):
masks = loadAndAgumentMasks.returnTensorMasks(config.batch_size)
masks = torch.from_numpy(masks)
masks = masks.type(torch.cuda.FloatTensor)
masks = 1 - masks
masks.to(config.device)
real = real.to(config.device)
fake_masked_images = torch.mul(real, masks)
generated_images = gen(fake_masked_images, masks)
image_names = names[current_number:current_number + config.batch_size]
current_number = current_number + config.batch_size ## Change naming to include all names
# modelHelper.save_tensor_batch(generated_images,fake_masked_images,config.batch_size,path)
if config.nir_data:
for index, image in enumerate(generated_images):
namePath = Path.joinpath(local_test_path, image_names[index])
if config.nir_data:
modelHelper.save_tensor_single_NIR(
image,
Path.joinpath(local_test_path, image_names[index]),
Path.joinpath(local_test_nir_path, image_names[index]),
raw=True)
else:
modelHelper.save_tensor_batch(
real, fake_masked_images, generated_images, config.batch_size,
Path.joinpath(local_test_path,
"_final_model_" + str(current_number)))
current_number = current_number + 1
end_time = datetime.now()