def attack(image):
temps = []
start = time.time()
for count in range(20): #可修改的最少遍历次数
succ = 0
new_image = random_attack(image, 500) #可修改的遍历深度
if (new_image is not None):
# 攻击成功
succ = 1
else:
print("fail")
ran = random.randint(0, 10000)
global ccc
ccc += 1
return train_images[ran] / 255.0
# 失败 条件为500*20次无法生成相似度>0.75的标签不同的新图片 参数可修改
temps.append((new_image, SSIM(new_image, image)))
if (SSIM(new_image, image) > 0.9):
break
index = 0
max_ssim = 0
for i in range(len(temps)):
ssim = temps[i][1]
if ssim > max_ssim:
index = i
max_ssim = ssim
#print(time.time()-start)
#print(max_ssim)
res = temps[index]
return res[0]
def test():
oSSIM = 0
mSSIM = 0
oPSNR = 0
mPSNR = 0
tf.get_default_graph()
saver = tf.train.import_meta_graph("../temp/ckpt/model.ckpt-50.meta")
with tf.Session() as sess:
saver.restore(sess, tf.train.latest_checkpoint("../temp/ckpt/"))
for i in range(1, test_size + 1):
print("Iteration {0}".format(i))
original_image = "../dataset/test_orig/orig_{0}.png".format(i)
mf_image = "../dataset/test_mf/mf_{0}.png".format(i)
original_image_object = Image.open(original_image)
mf_image_object = Image.open(mf_image)
test_size_s = 16
Z = np.array(mf_image_object)
Z = Z.reshape((1, 128, 128, 1))
image = sess.run(rest, feed_dict={z: Z})
image = image.reshape((128, 128))
test_size_p = 160
gen_image_object = Image.fromarray(image, mode='L')
original_width, original_height, original_npix = get_image_data_from_file(
original_image)
gen_width, gen_height, gen_npix = get_image_data(gen_image_object)
mf_width, mf_height, mf_npix = get_image_data_from_file(mf_image)
if original_width != gen_width or original_height != gen_height or original_npix != gen_npix:
print("ERROR: Images should have same dimensions. \n")
exit(1)
if mf_width != gen_width or mf_height != gen_height or mf_npix != gen_npix:
print("ERROR: Images should have same dimensions. \n")
exit(1)
original_y, original_cb, original_cr = get_yuv(
original_image_object)
gen_y, gen_cb, gen_cr = get_yuv(gen_image_object)
mf_y, mf_cb, mf_cr = get_yuv(mf_image_object)
psnr = calculate_psnr(original_y, gen_y, original_cb, gen_cb,
original_cr, gen_cr, original_npix)
oPSNR = oPSNR + psnr
psnr = calculate_psnr(mf_y, gen_y, mf_cb, gen_cb, mf_cr, gen_cr,
mf_npix)
mPSNR = mPSNR + psnr
size = (256, 256)
original_image_object = original_image_object.resize(
size, Image.ANTIALIAS)
gen_image_object = gen_image_object.resize(size, Image.ANTIALIAS)
mf_image_object = mf_image_object.resize(size, Image.ANTIALIAS)
ssim = SSIM(original_image_object,
gaussian_kernel_1d).ssim_value(gen_image_object)
oSSIM = oSSIM + ssim
ssim = SSIM(mf_image_object,
gaussian_kernel_1d).ssim_value(gen_image_object)
mSSIM = mSSIM + ssim
print("Average: ")
print("oPSNR={0}".format(oPSNR / test_size_p))
print("mPSNR={0}".format(mPSNR / test_size_p))
print("oSSIM={0}".format(oSSIM / test_size_s))
print("mSSIM={0}".format(mSSIM / test_size_s))
def SSIM_Array(arr1, arr2):
img1 = Image.fromarray(arr1.astype('uint8'), 'RGB')
img2 = Image.fromarray(arr2.astype('uint8'),'RGB')
gaussian_kernel_sigma = 1.5
gaussian_kernel_width = 11
gaussian_kernel_1d = get_gaussian_kernel(gaussian_kernel_width, gaussian_kernel_sigma)
return SSIM(img1,gaussian_kernel_1d).ssim_value(img2)
def ssim_loss(self, hr, sr):
margin = (tf.shape(hr)[1] - tf.shape(sr)[1]) // 2
hr_crop = tf.cond(tf.equal(margin, 0), lambda: hr,
lambda: hr[:, margin:-margin, margin:-margin, :])
hr = K.in_train_phase(hr_crop, hr)
hr.uses_learning_phase = True
return -SSIM(hr, sr)
def train(config=None):
# Initialize a new wandb run
with wandb.init(config=config):
# If called by wandb.agent, as below,
# this config will be set by Sweep Controller
config = wandb.config
train_loader, val_loader = build_dataset(config)
model = DnCNN(channels=3)
model.apply(weights_init_kaiming)
lr = config['lr']
wd = config['wd']
# Loss function
criterion = SSIM()
# Create the contiguous parameters.
model_params = [p for p in model.parameters() if p.requires_grad]
optimizer = holocron.optim.RAdam(model_params,
lr,
betas=(0.95, 0.99),
eps=1e-6,
weight_decay=wd)
#Trainer
trainer = AutoencoderTrainer(model,
train_loader,
val_loader,
criterion,
optimizer,
0,
output_file=config['checkpoint'],
configwb=True)
trainer.fit_n_epochs(config['epochs'], config['lr'], config['freeze'])
def test_image(model, save_path, image_path):
img_transforms = transforms.Compose(
[transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])])
size_transform = Compose([PadIfNeeded(736, 1280)])
crop = CenterCrop(720, 1280)
img = cv2.imread(image_path)
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
img_s = size_transform(image=img)['image']
img_tensor = torch.from_numpy(
np.transpose(img_s / 255, (2, 0, 1)).astype('float32'))
img_tensor = img_transforms(img_tensor)
with torch.no_grad():
img_tensor = Variable(img_tensor.unsqueeze(0).cuda())
result_image = model(img_tensor)
result_image = result_image[0].cpu().float().numpy()
result_image = (np.transpose(result_image, (1, 2, 0)) + 1) / 2.0 * 255.0
result_image = crop(image=result_image)['image']
result_image = result_image.astype('uint8')
gt_image = get_gt_image(image_path)
lap = estimate_blur(result_image)
lap_sharp = estimate_blur(gt_image)
lap_blur = estimate_blur(img)
_, filename = os.path.split(image_path)
psnr = PSNR(result_image, gt_image)
pilFake = Image.fromarray(result_image)
pilReal = Image.fromarray(gt_image)
ssim = SSIM(pilFake).cw_ssim_value(pilReal)
#result_image = np.hstack((img_s, result_image, gt_image))
#cv2.imwrite(os.path.join(save_path, filename), cv2.cvtColor(result_image, cv2.COLOR_RGB2BGR))
return psnr, ssim, lap, lap_blur, lap_sharp
def __init__(self, smoothl1 = True, l1 = False, mse = False, instance_ssim = True, perceptual_loss = True):
super(loss_function, self).__init__()
print("=========> Building criterion")
self.loss_function = nn.ModuleDict()
self.weight = dict()
self.loss_function['loss_smooth_l1'] = nn.SmoothL1Loss() if smoothl1 else None
self.loss_function['loss_l1'] = nn.L1Loss() if l1 else None
self.loss_function['loss_mse'] = torch.nn.MSELoss() if mse else None
self.loss_function['instance_ssim'] = SSIM(reduction = 'mean', window_size = 7, asloss = True) if instance_ssim else None
self.loss_function['loss_perceptual'] = perceptual() if perceptual_loss else None
self.weight['loss_smooth_l1'] = 1
self.weight['loss_l1'] = 1
self.weight['loss_mse'] = 1
self.weight['instance_ssim'] = 1
self.weight['loss_perceptual'] = 1
if opt.cuda:
if opt.parallel:
for key in self.loss_function.keys():
if self.loss_function[key] is not None:
self.loss_function[key] = nn.DataParallel(self.loss_function[key], [0, 1, 2, 3]).cuda()
else:
for key in self.loss_function.keys():
if self.loss_function[key] is not None:
self.loss_function[key] = self.loss_function[key].cuda()
else:
for key in self.loss_function.keys():
if self.loss_function[key] is not None:
self.loss_function[key] = self.loss_function[key].cpu()
def __init__(self,
data,
trn_avails,
logger,
writer,
batch_size,
transform,
content_font,
language,
meta,
val_loaders,
n_workers=2):
self.data = data
self.logger = logger
self.writer = writer
self.batch_size = batch_size
self.transform = transform
self.n_workers = n_workers
self.unify_resize_method = True
self.trn_avails = trn_avails
self.val_loaders = val_loaders
self.content_font = content_font
self.language = language
if self.language == 'kor':
self.n_comp_types = 3
elif self.language == 'thai':
self.n_comp_types = 4
else:
raise ValueError()
# setup cross-validation
self.SSIM = SSIM().cuda()
weights = [0.25, 0.3, 0.3, 0.15]
self.MSSSIM = MSSSIM(weights=weights).cuda()
n_batches = [len(loader) for loader in self.val_loaders.values()]
self.n_cv_batches = min(n_batches)
self.logger.info("# of cross-validation batches = {}".format(
self.n_cv_batches))
# the number of chars/fonts for CV visualization
n_chars = 16
n_fonts = 16
seen_chars = uniform_sample(meta['train']['chars'], n_chars // 2)
unseen_chars = uniform_sample(meta['valid']['chars'], n_chars // 2)
unseen_fonts = uniform_sample(meta['valid']['fonts'], n_fonts)
self.cv_comparable_fonts = unseen_fonts
self.cv_comparable_chars = seen_chars + unseen_chars
allchars = meta['train']['chars'] + meta['valid']['chars']
self.cv_comparable_avails = {
font: allchars
for font in self.cv_comparable_fonts
}
def compare(path_Groundimage, path_Generate_image):
gaussian_kernel_sigma = 1.5
gaussian_kernel_width = 11
gaussian_kernel_1d = get_gaussian_kernel(gaussian_kernel_width,
gaussian_kernel_sigma)
size = (64, 64)
im = Image.open(path_Groundimage)
im = im.resize(size, Image.ANTIALIAS)
# slightly rotated image
im_rot = Image.open(path_Generate_image)
im_rot = im_rot.resize(size, Image.ANTIALIAS)
# print("========== SSIM results ==========")
ssim_rot = SSIM(im, gaussian_kernel_1d).ssim_value(im_rot)
# print("========== CV-SSIM results ==========")
cw_ssim_rot = SSIM(im).cw_ssim_value(im_rot)
# print("CW-SSIM of generated image and ground truth with view 0 %.4f" % cw_ssim_rot)
return ssim_rot, cw_ssim_rot
def __init__(self, device, fidelity=None):
super().__init__()
self.smoothing = GaussianSmoothing(3, 10, 5)
self.smoothing = self.smoothing.to(device)
self.ssim = SSIM()
self.w1 = 1
if fidelity == 'l1':
self.fidelity = nn.L1Loss().to(device)
else:
self.w1 = 0.00001
self.fidelity = self.color_loss
def random_attack(image, max_time):
label = predict(image)
loop = 1
new_image = random_line_attack(image, 1)
i = 0
while (predict(new_image) == label
or SSIM(new_image, image) < 0.75) and i < max_time: #结构相似度可调整
i += 1
new_image = random_line_attack(image, r=loop + 1)
if i < max_time:
return new_image
return None
def get_diff(im1, im2, mode='color'):
im1 = cv2.resize(im1, (200, 200))
im2 = cv2.resize(im2, (200, 200))
saveData('x_1920x1080.RGB', im1)
saveData('y_1920x1080.RGB', im2)
im1 = im1.astype('float32') / 255
im2 = im2.astype('float32') / 255
if mode == 'gray':
im1 = cv2.cvtColor(im1, cv2.COLOR_BGR2GRAY)
im2 = cv2.cvtColor(im2, cv2.COLOR_BGR2GRAY)
score = SSIM(im1, im2, max_value=1)
return score
def ssim_compare(img1, img2):
im1 = Image.open(img1)
im1 = im1.resize(size, Image.ANTIALIAS)
im2 = Image.open(img2)
im2 = im2.resize(size, Image.ANTIALIAS)
similar = SSIM(im1).cw_ssim_value(im2) * 100
if similar > 90:
return ('Same Image', similar)
else:
return ('Different Image', similar)
def __init__(self, model, layer_ids, layer_wgts, base_loss=F.l1_loss):
super().__init__()
self.model = nn.Sequential(*list(model.children())[:layer_ids[-1] + 1])
self.loss_features = [self.model[i] for i in layer_ids]
self.hooks = hook_outputs(self.loss_features, detach=False)
self.wgts = layer_wgts
self.metric_names = [
'pixel',
] + [f'feat_{i}' for i in range(len(layer_ids))
] + [f'gram_{i}'
for i in range(len(layer_ids))] + ['PSNR', 'SSIM', 'pEPs']
self.base_loss = base_loss
self.mse = nn.MSELoss()
self.ssim = SSIM(window_size=11)
def run(y):
LR = Variable(ToTensor()(y).unsqueeze(0).pin_memory(), volatile=True)
if opt.cuda:
LR = LR.cuda()
HR_2 = timeit(model.__call__)(LR).cpu()
if opt.reference:
img = Image.open(opt.reference).convert('YCbCr')
y, _, _ = img.split()
# y = y.resize((y.size[0] // 2, y.size[1] // 2), Image.BILINEAR)
y = transform(y)
y.save('ref.png')
HR_ref = Variable(ToTensor()(y).unsqueeze(0), volatile=True)
print('SSIM:', 1 - SSIM()(HR_ref, HR_2).data[0])
print('MS-SSIM:', 1 - MSSSIM()(HR_ref, HR_2).data[0])
print('Charbonnier:', CharbonnierLoss(HR_ref, HR_2).data[0])
return process(HR_2.cpu())
def get_ssim():
image_dir = '/home/marcovaldo/Data'
cover_dir = 'ILSVRC2012_img_val'
stego_dir = 'UNet_ImageNet_stego'
ssim_criterion = SSIM(window_size=11).cuda()
total_loss, idx = list(), 0
dataloader2 = DatasetFromFolder2(image_dir=image_dir, cover_dir=cover_dir,
stego_dir=stego_dir, transform=transform)
dataloader2 = torch.utils.data.DataLoader(dataset=dataloader2, batch_size=params.batch_size, shuffle=False)
for idx, (covers, stegos) in enumerate(dataloader2, 1):
covers = Variable(covers.cuda())
stegos = Variable(stegos.cuda())
loss = ssim_criterion(covers, stegos)
total_loss.append(loss.data[0])
if idx % 1 == 0:
print('{}/{} loss: {:.6f}'.format(idx*params.batch_size, len(dataloader2.dataset), loss.data[0]))
print(sum(total_loss)/len(total_loss))
def __init__(self, args):
super(Model, self).__init__()
self.fusion = Decomposition()
self.D = MultiscaleDiscriminator(input_nc=1)
self.MSE_fun = nn.MSELoss()
self.L1_loss = nn.L1Loss()
self.SSIM_fun = SSIM()
if args.contiguousparams == True:
print("ContiguousParams---")
parametersF = ContiguousParams(self.fusion.parameters())
parametersD = ContiguousParams(self.D.parameters())
self.optimizer_G = optim.Adam(parametersF.contiguous(), lr=args.lr)
self.optimizer_D = optim.Adam(parametersD.contiguous(), lr=args.lr)
else:
self.optimizer_G = optim.Adam(self.fusion.parameters(), lr=args.lr)
self.optimizer_D = optim.Adam(self.D.parameters(), lr=args.lr)
self.g1 = self.g2 = self.g3 = self.s = self.img_re = None
self.loss = torch.zeros(1)
self.scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(
self.optimizer_G,
mode='min',
factor=0.5,
patience=2,
verbose=False,
threshold=0.0001,
threshold_mode='rel',
cooldown=0,
min_lr=0,
eps=1e-10)
self.min_loss = 1000
self.args = args
self.downsample = downsample()
self.criterionGAN = torch.nn.MSELoss()
if args.multiGPU:
self.mulgpus()
self.load()
self.load_D()
def __init__(self, *args, **kwargs):
self.ssim = SSIM(*args, **kwargs)
sprite_gray = ImageOps.grayscale(Image.open(filename))
sprites[filename] = sprite_gray
gaussian_kernel_sigma = 1.5
gaussian_kernel_width = 11
gaussian_kernel_1d = get_gaussian_kernel(gaussian_kernel_width,
gaussian_kernel_sigma)
matching = {}
maxVal = float('-inf')
for sprite1 in sprites:
matching[sprite1] = {}
for sprite2 in sprites:
matching[sprite1][sprite2] = SSIM(
sprites[sprite1], gaussian_kernel_1d
).ssim_value(
sprites[sprite2]
) #np.max(sp.signal.correlate2d(sprites[sprite1],sprites[sprite2]))
if matching[sprite1][sprite2] > maxVal:
maxVal = matching[sprite1][sprite2]
for s1 in matching:
for s2 in matching[s1]:
matching[s1][s2] /= maxVal
#print s1, s2, matching[s1][s2]
frames = []
with open(frameInformation) as frameData:
categories = frameData.readline().rstrip().split(',')
curFrame = -1
frame = None
if (h_s + img_s >= img_height):
break
h_s += step
img[:, :, 0] = img[:, :, 0] / area_count
img[:, :, 1] = img[:, :, 1] / area_count
img[:, :, 2] = img[:, :, 2] / area_count
fake_B = img.astype(np.uint8)
if opt.dataset_mode != 'single':
real_B = util.tensor2im(data['B'])
avgPSNR += PSNR(fake_B, real_B)
pilFake = Image.fromarray(fake_B)
pilReal = Image.fromarray(real_B)
avgSSIM += SSIM(pilFake).cw_ssim_value(pilReal)
img_path = model.get_image_paths()
print('process image... %s' % img_path)
visualizer.save_images(
webpage,
OrderedDict([('real_A', util.tensor2im(data['A'])),
('fake_B', fake_B)]), img_path)
if opt.dataset_mode != 'single':
avgPSNR /= counter
avgSSIM /= counter
with open(
os.path.join(opt.results_dir, opt.name, 'test_latest',
'result.txt'), 'w') as f:
f.write('PSNR = %f\n' % avgPSNR)
def main():
global opt, name, logger, netG, netD, vgg, curriculum_ssim, loss_mse, rgb2yuv, instance_ssim, loss_bce
opt = parser.parse_args()
name = "ShadowRemoval"
print(opt)
# Tag_ResidualBlocks_BatchSize
cuda = opt.cuda
if cuda and not torch.cuda.is_available():
raise Exception("No GPU found, please run without --cuda")
seed = 1334
torch.manual_seed(seed)
if 'WORLD_SIZE' in os.environ:
opt.distributed = int(os.environ['WORLD_SIZE']) > 1
if cuda and not torch.cuda.is_available():
raise Exception("No GPU found, please run without --cuda")
if opt.parallel:
opt.gpu = opt.local_rank
torch.cuda.set_device(opt.gpu)
torch.distributed.init_process_group(backend='nccl',
init_method='env://')
opt.world_size = torch.distributed.get_world_size()
if cuda:
torch.cuda.manual_seed(seed)
cudnn.benchmark = True
print("==========> Loading datasets")
test_dataset = DatasetFromFolder(opt.test,
transform=Compose([ToTensor()]),
training=False,
experiments="ShadowRemoval")
data_loader = DataLoader(dataset=test_dataset,
num_workers=4,
batch_size=opt.batchSize,
pin_memory=True,
shuffle=False)
print("==========> Building model")
netG = ShadowRemoval(channels=64)
print("=========> Building criterion")
loss_smooth_l1 = nn.SmoothL1Loss()
loss_l1 = nn.L1Loss()
loss_mse = torch.nn.MSELoss()
loss_bce = torch.nn.BCELoss()
loss_perceptual = perceptual()
instance_ssim = SSIM(reduction='mean', window_size=7)
rgb2yuv = rgb2yuv()
curriculum_ssim_mask = CLBase(lossfunc=nn.BCELoss(reduce=False))
curriculum_ssim_clean = CLBase()
# optionally copy weights from a checkpoint
if opt.pretrained and opt.continue_training:
if os.path.isfile(opt.pretrained):
print("=> loading model '{}'".format(opt.pretrained))
weights = torch.load(opt.pretrained)
netG.load_state_dict(weights['state_dict'])
else:
print("=> no model found at '{}'".format(opt.pretrained))
print("==========> Setting Optimizer")
optimizerG = optim.Adam(filter(lambda p: p.requires_grad,
netG.parameters()),
lr=opt.lr_g,
betas=(0.9, 0.999))
print("==========> Setting GPU")
if cuda:
netG = netG.cuda()
instance_ssim = instance_ssim.cuda()
loss_smooth_l1 = loss_smooth_l1.cuda()
loss_mse = loss_mse.cuda()
loss_l1 = loss_l1.cuda()
loss_bce = loss_bce.cuda()
curriculum_ssim_mask = curriculum_ssim_mask.cuda()
curriculum_ssim_clean = curriculum_ssim_clean.cuda()
loss_perceptual = loss_perceptual.cuda()
rgb2yuv = rgb2yuv.cuda()
if opt.acceleration:
print("FP 16 Trianing")
amp.register_float_function(torch, 'sigmoid')
netG, optimizerG = amp.initialize(netG,
optimizerG,
opt_level=opt.opt_level)
else:
netG = netG.cpu()
instance_ssim = instance_ssim.cpu()
loss_smooth_l1 = loss_smooth_l1.cpu()
loss_mse = loss_mse.cpu()
loss_l1 = loss_l1.cpu()
loss_bce = loss_bce.cpu()
curriculum_ssim = curriculum_ssim.cpu()
loss_perceptual = loss_perceptual.cpu()
rgb2yuv = rgb2yuv.cpu()
test(data_loader, netG)
print_network(encoder)
print_network(decoder)
print_network(discriminator)
print_network(discriminator2)
if params.use_cuda:
encoder.cuda()
decoder.cuda()
discriminator.cuda()
discriminator2.cuda()
encoder.apply(weights_init)
decoder.apply(weights_init)
# loss function
BCE_loss = nn.BCELoss().cuda()
MSE_loss = nn.MSELoss().cuda()
SSIM_loss = SSIM(window_size=11).cuda()
MSSIM_loss = MSSSIM().cuda()
# Data pre-processing
transform = transforms.Compose([
transforms.Scale(params.input_size),
transforms.ToTensor(),
])
# transforms.Normalize(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225))])
# transforms.Normalize(mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5))])
# Train data
# train_data = DatasetFromFolder(data_dir, subfolder=subfolder, transform=transform, resize_scale=params.input_size,
# name=False)
train_data = DatasetFromFolder2(data_dir,
subfolder=subfolder,
transform=transform,
def calculate(path_true,
path_pred,
output_path,
log_file,
image_size=256,
ssim_winsize=11,
gray_mode=False,
debug_mode=False):
psnr = []
ssim = []
mae = []
names = []
index = 1
if not log_file:
log_file = os.path.join(output_path, 'metrics.txt')
psnr_torch = []
psnr_metric = PSNR(1.)
ssim_torch_class = []
ssim_torch_fun = []
ssim_metric_class = SSIM(window_size=11)
ssim_metric_fun = get_ssim
# files = list(glob(path_true + '/*.jpg')) + list(glob(path_true + '/*.png'))
files = load_flist(path_true)
for fn in sorted(files):
name = basename(str(fn))
names.append(name)
img_gt = (imread(str(fn)) / 255.0).astype(np.float32)
img_pred = (imread(path_pred + '/' + basename(str(fn))) /
255.0).astype(np.float32)
if img_gt.ndim == 2:
img_gt = gray2rgb(img_gt)
img_gt = resize(img_gt, image_size, image_size).astype(
np.float32) / 255.0
img_pred = resize(img_pred, image_size, image_size).astype(
np.float32) / 255.0
if gray_mode:
img_gt = rgb2gray(img_gt)
img_pred = rgb2gray(img_pred)
if debug_mode:
plt.subplot('121')
plt.imshow(img_gt)
plt.title('Groud truth')
plt.subplot('122')
plt.imshow(img_pred)
plt.title('Output')
plt.show()
psnr.append(compare_psnr(img_gt, img_pred, data_range=1))
psnr_torch.append(psnr_metric(img_gt, img_pred).item())
ssim.append(
compare_ssim(img_gt,
img_pred,
data_range=1,
win_size=11,
multichannel=True,
sigma=1.5,
use_sample_covariance=False,
gaussian_weights=True))
ssim_torch_class.append(
ssim_metric_class(to_tensor(img_gt), to_tensor(img_pred)).item())
ssim_torch_fun.append(
ssim_metric_fun(to_tensor(img_gt), to_tensor(img_pred)).item())
mae.append(compare_mae(img_gt, img_pred))
if np.mod(index, 1) == 0:
print(
str(index) + ' images processed',
"PSNR: %.4f" % round(np.mean(psnr), 4),
"PSNR_torch: %.4f" % round(np.mean(psnr_torch), 4),
"SSIM: %.4f" % round(np.mean(ssim), 4),
"SSIM_torch_class: %.4f" % round(np.mean(ssim_torch_class), 4),
"SSIM_torch_fun: %.4f" % round(np.mean(ssim_torch_fun), 4),
"MAE: %.4f" % round(np.mean(mae), 4),
)
index += 1
np.savez(output_path + '/metrics.npz',
psnr=psnr,
ssim=ssim,
mae=mae,
names=names)
print_fun("PSNR: %.4f" % round(np.mean(psnr), 4),
"PSNR Variance: %.4f" % round(np.var(psnr), 4),
"PSNR_torch: %.4f" % round(np.mean(psnr_torch), 4),
"SSIM: %.4f" % round(np.mean(ssim), 4),
"SSIM Variance: %.4f" % round(np.var(ssim), 4),
"MAE: %.4f" % round(np.mean(mae), 4),
"MAE Variance: %.4f" % round(np.var(mae), 4),
log_file=log_file)
print()
for i in range(6):
start = i * 2000
end = (i + 1) * 2000
print_fun("mask ratio: [%2d%% - %2d%%]" % (i * 10, (i + 1) * 10),
log_file=log_file)
print_fun("PSNR: %.4f" % round(np.mean(psnr[start:end]), 4),
"PSNR Variance: %.4f" % round(np.var(psnr[start:end]), 4),
"SSIM: %.4f" % round(np.mean(ssim[start:end]), 4),
"SSIM Variance: %.4f" % round(np.var(ssim[start:end]), 4),
"MAE: %.4f" % round(np.mean(mae[start:end]), 4),
"MAE Variance: %.4f" % round(np.var(mae[start:end]), 4),
log_file=log_file)
def ssimLoss(x,y):
ssim_loss = SSIM()
return -ssim_loss(x,y)
# plt.savefig('figs/' + speaker + '_' + basefile + '_mean.png')
# plt.close()
ult_all[speaker + '-' + basefile] = ult_data_3d_mean
# now all ult files are loaded
CW_SSIM_all = np.zeros((len(ult_all), len(ult_all)))
# calculate CW_SSIM utterance by utterance
ult_keys = list(sorted(ult_all.keys()))
for i in range(len(ult_all)):
for j in range(len(ult_all)):
print('calc CW_SSIM', speaker, i, j, ' ', end='\r')
CW_SSIM_all[i, j] = SSIM( \
Image.fromarray(ult_all[ult_keys[i]].reshape(NumVectors, PixPerVector), 'L')).cw_ssim_value( \
Image.fromarray(ult_all[ult_keys[j]].reshape(NumVectors, PixPerVector), 'L'))
print('calc CW_SSIM', speaker, 'done')
# serialize results to file
pickle.dump(CW_SSIM_all,
open('measures/' + speaker + '_CW_SSIM_avg.pkl', 'wb'))
else:
CW_SSIM_all = pickle.load(
open('measures/' + speaker + '_CW_SSIM_avg.pkl', 'rb'))
print(speaker, ', len:', len(CW_SSIM_all), ', CW_SSIM min/max: ',
np.min(CW_SSIM_all), np.max(CW_SSIM_all))
# visualize results
# vmax set according to manual checking
def __init__(self):
super(SSIMLoss, self).__init__()
self.ssim = SSIM(data_range=1.)
from ssim import SSIM
from PIL import Image, ImageOps
im = Image.open('1.jpg')
im2 = Image.open('2.jpg')
cw_ssim_rot = SSIM(im).cw_ssim_value(im2)
print cw_ssim_rot
import torch
import torch.nn.functional as F
from torch import nn
from ssim import SSIM
from geometrics import flow_warp, inverse_warp, pose2flow, mask_gen
from collections import defaultdict
import numpy as np
import cv2
get_ssim = SSIM().cuda()
eps = 1e-3
MIN_DEPTH = 1e-3
MAX_DEPTH = 80
def gradient(pred):
dy = torch.abs(pred[..., :-1, :] - pred[..., 1:, :])
dx = torch.abs(pred[..., :, -1:] - pred[..., :, 1:])
return dx, dy
def upsample(src, shape):
return F.interpolate(src, shape, mode="bilinear", align_corners=False)
def compute_errors(gt, pred):
"""Computation of error metrics between predicted and ground truth depths
"""
error_dict = {}
thresh = np.maximum((gt / pred), (pred / gt))
error_dict['a1'] = (thresh < 1.25).mean()
print(args)
if not os.path.exists(args.checkpoint_path):
os.mkdir(args.checkpoint_path)
print('===>make directory', args.checkpoint_path)
if not os.path.exists(args.result_path):
os.mkdir(args.result_path)
print('===>make directory', args.result_path)
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
dataset = RGB_NIR_Dataset(mode='train')
dataloader = DataLoader(dataset,
batch_size=args.batch_size,
shuffle=True,
pin_memory=True)
generator = RGB_Generator().to(device)
discriminator = Discriminator().to(device)
g_optim = optim.Adam(generator.parameters(), lr=args.lr, betas=(0., 0.99))
d_optim = optim.Adam(discriminator.parameters(),
lr=args.lr,
betas=(0., 0.99))
MSE_Loss = nn.MSELoss().to(device)
ssim_loss = SSIM()
for i in range(args.epochs):
train(i, dataloader, generator, discriminator, g_optim, d_optim,
MSE_Loss, ssim_loss)
print('epoch %3d is done' % i)
def get_loss_function(self, model, colorspace='RGB', reduction='sum', deterministic=False, pretrained=True, image_size=None):
"""
returns a trained loss class.
model: one of the values returned by self.loss_functions
colorspace: 'LA' or 'RGB'
deterministic: bool, if false (default) uses shifting of image blocks for watson-fft
image_size: tuple, size of input images. Only required for adaptive loss. Eg: [3, 64, 64]
"""
is_greyscale = colorspace in ['grey', 'Grey', 'LA', 'greyscale', 'grey-scale']
if model.lower() in ['l2']:
loss = nn.MSELoss(reduction=reduction)
elif model.lower() in ['l1']:
loss = nn.L1Loss(reduction=reduction)
elif model.lower() in ['ssim']:
loss = SSIM(size_average=(reduction in ['sum', 'mean']))
elif model.lower() in ['watson', 'watson-dct']:
if is_greyscale:
if deterministic:
loss = WatsonDistance(reduction=reduction)
if pretrained:
loss.load_state_dict(self.load_state_dict('gray_watson_dct_trial0.pth'))
else:
loss = ShiftWrapper(WatsonDistance, (), {'reduction': reduction})
if pretrained:
loss.loss.load_state_dict(self.load_state_dict('gray_watson_dct_trial0.pth'))
else:
if deterministic:
loss = ColorWrapper(WatsonDistance, (), {'reduction': reduction})
if pretrained:
loss.load_state_dict(self.load_state_dict('rgb_watson_dct_trial0.pth'))
else:
loss = ShiftWrapper(ColorWrapper, (WatsonDistance, (), {'reduction': reduction}), {})
if pretrained:
loss.loss.load_state_dict(self.load_state_dict('rgb_watson_dct_trial0.pth'))
elif model.lower() in ['watson-fft', 'watson-dft']:
if is_greyscale:
if deterministic:
loss = WatsonDistanceFft(reduction=reduction)
if pretrained:
loss.load_state_dict(self.load_state_dict('gray_watson_fft_trial0.pth'))
else:
loss = ShiftWrapper(WatsonDistanceFft, (), {'reduction': reduction})
if pretrained:
loss.loss.load_state_dict(self.load_state_dict('gray_watson_fft_trial0.pth'))
else:
if deterministic:
loss = ColorWrapper(WatsonDistanceFft, (), {'reduction': reduction})
if pretrained:
loss.load_state_dict(self.load_state_dict('rgb_watson_fft_trial0.pth'))
else:
loss = ShiftWrapper(ColorWrapper, (WatsonDistanceFft, (), {'reduction': reduction}), {})
if pretrained:
loss.loss.load_state_dict(self.load_state_dict('rgb_watson_fft_trial0.pth'))
elif model.lower() in ['watson-vgg', 'watson-deep']:
if is_greyscale:
loss = GreyscaleWrapper(WatsonDistanceVgg, (), {'reduction': reduction})
if pretrained:
loss.loss.load_state_dict(self.load_state_dict('gray_watson_vgg_trial0.pth'))
else:
loss = WatsonDistanceVgg(reduction=reduction)
if pretrained:
loss.load_state_dict(self.load_state_dict('rgb_watson_vgg_trial0.pth'))
elif model.lower() in ['deeploss-vgg']:
if is_greyscale:
loss = GreyscaleWrapper(PNetLin, (), {'pnet_type': 'vgg', 'reduction': reduction, 'use_dropout': False})
if pretrained:
loss.loss.load_state_dict(self.load_state_dict('gray_pnet_lin_vgg_trial0.pth'))
else:
loss = PNetLin(pnet_type='vgg', reduction=reduction, use_dropout=False)
if pretrained:
loss.load_state_dict(self.load_state_dict('rgb_pnet_lin_vgg_trial0.pth'))
elif model.lower() in ['deeploss-squeeze']:
if is_greyscale:
loss = GreyscaleWrapper(PNetLin, (), {'pnet_type': 'squeeze', 'reduction': reduction, 'use_dropout': False})
if pretrained:
loss.loss.load_state_dict(self.load_state_dict('gray_pnet_lin_squeeze_trial0.pth'))
else:
loss = PNetLin(pnet_type='squeeze', reduction=reduction, use_dropout=False)
if pretrained:
loss.load_state_dict(self.load_state_dict('rgb_pnet_lin_squeeze_trial0.pth'))
elif model.lower() in ['adaptive']:
def map_weights(states):
return OrderedDict([(k[1:], v) for k, v in states.items()])
assert image_size is not None, 'Adaptive loss requires image size input'
if is_greyscale:
loss = GreyscaleWrapper(RobustLoss, (), {'image_size':image_size, 'use_gpu':False, 'trainable':False, 'reduction':reduction})
if pretrained:
loss.loss.load_state_dict(map_weights(self.load_state_dict('gray_adaptive_trial0.pth')))
else:
loss = RobustLoss(image_size=image_size, use_gpu=False, trainable=False, reduction=reduction)
if pretrained:
loss.load_state_dict(map_weights(self.load_state_dict('rgb_adaptive_trial0.pth')))
else:
raise Exception('Metric "{}" not implemented'.format(model))
# freeze all training of the loss functions
if pretrained:
for param in loss.parameters():
param.requires_grad = False
return loss