diff --git a/compute_dists_dirs.py b/compute_dists_dirs.py index 528ee5f7..38c09e79 100644 --- a/compute_dists_dirs.py +++ b/compute_dists_dirs.py @@ -7,12 +7,13 @@ parser.add_argument('-d0','--dir0', type=str, default='./imgs/ex_dir0') parser.add_argument('-d1','--dir1', type=str, default='./imgs/ex_dir1') parser.add_argument('-o','--out', type=str, default='./imgs/example_dists.txt') +parser.add_argument('-v','--version', type=str, default='0.1') parser.add_argument('--use_gpu', action='store_true', help='turn on flag to use GPU') opt = parser.parse_args() ## Initializing the model -model = models.PerceptualLoss(model='net-lin',net='alex',use_gpu=opt.use_gpu) +model = models.PerceptualLoss(model='net-lin',net='alex',use_gpu=opt.use_gpu,version=opt.version) # crawl directories f = open(opt.out,'w') diff --git a/compute_dists_pair.py b/compute_dists_pair.py index b284adf3..9000d369 100644 --- a/compute_dists_pair.py +++ b/compute_dists_pair.py @@ -1,42 +1,58 @@ import argparse import os import models -import numpy as np from util import util +import numpy as np +from IPython import embed parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter) -parser.add_argument('-d','--dir', type=str, default='./imgs/ex_dir_pair') +parser.add_argument('-d','--dir', type=str, default='./imgs/ex_dir0') parser.add_argument('-o','--out', type=str, default='./imgs/example_dists.txt') +parser.add_argument('-v','--version', type=str, default='0.1') +parser.add_argument('--all-pairs', action='store_true', help='turn on to test all N(N-1)/2 pairs, leave off to just do consecutive pairs (N-1)') +parser.add_argument('-N', type=int, default=None) parser.add_argument('--use_gpu', action='store_true', help='turn on flag to use GPU') opt = parser.parse_args() ## Initializing the model -model = models.PerceptualLoss(model='net-lin',net='alex',use_gpu=opt.use_gpu) +model = models.PerceptualLoss(model='net-lin',net='alex',use_gpu=opt.use_gpu,version=opt.version) # crawl directories f = open(opt.out,'w') files = os.listdir(opt.dir) +if(opt.N is not None): + files = files[:opt.N] +F = len(files) dists = [] -for (ff,file0) in enumerate(files[:-1]): - img0 = util.im2tensor(util.load_image(os.path.join(opt.dir,file0))) # RGB image from [-1,1] +for (ff,file) in enumerate(files[:-1]): + img0 = util.im2tensor(util.load_image(os.path.join(opt.dir,file))) # RGB image from [-1,1] if(opt.use_gpu): img0 = img0.cuda() - for (gg,file1) in enumerate(files[ff+1:]): + if(opt.all_pairs): + files1 = files[ff+1:] + else: + files1 = [files[ff+1],] + + for file1 in files1: img1 = util.im2tensor(util.load_image(os.path.join(opt.dir,file1))) + if(opt.use_gpu): img1 = img1.cuda() # Compute distance - dist01 = model.forward(img0,img1).item() - dists.append(dist01) - print('(%s, %s): %.3f'%(file0,file1,dist01)) - f.writelines('(%s, %s): %.3f'%(file0,file1,dist01)) - -dist_mean = np.mean(np.array(dists)) -print('Mean: %.3f'%dist_mean) -f.writelines('Mean: %.3f'%dist_mean) + dist01 = model.forward(img0,img1) + print('(%s,%s): %.3f'%(file,file1,dist01)) + f.writelines('(%s,%s): %.6f\n'%(file,file1,dist01)) + + dists.append(dist01.item()) + +avg_dist = np.mean(np.array(dists)) +stderr_dist = np.std(np.array(dists))/np.sqrt(len(dists)) + +print('Avg: %.5f +/- %.5f'%(avg_dist,stderr_dist)) +f.writelines('Avg: %.6f +/- %.6f'%(avg_dist,stderr_dist)) f.close() diff --git a/models/__init__.py b/models/__init__.py index 83ae8107..9056aa9e 100755 --- a/models/__init__.py +++ b/models/__init__.py @@ -11,7 +11,7 @@ from models import dist_model class PerceptualLoss(torch.nn.Module): - def __init__(self, model='net-lin', net='alex', colorspace='rgb', spatial=False, use_gpu=True, gpu_ids=[0]): # VGG using our perceptually-learned weights (LPIPS metric) + def __init__(self, model='net-lin', net='alex', colorspace='rgb', spatial=False, use_gpu=True, gpu_ids=[0], version='0.1'): # VGG using our perceptually-learned weights (LPIPS metric) # def __init__(self, model='net', net='vgg', use_gpu=True): # "default" way of using VGG as a perceptual loss super(PerceptualLoss, self).__init__() print('Setting up Perceptual loss...') @@ -19,7 +19,7 @@ def __init__(self, model='net-lin', net='alex', colorspace='rgb', spatial=False, self.spatial = spatial self.gpu_ids = gpu_ids self.model = dist_model.DistModel() - self.model.initialize(model=model, net=net, use_gpu=use_gpu, colorspace=colorspace, spatial=self.spatial, gpu_ids=gpu_ids) + self.model.initialize(model=model, net=net, use_gpu=use_gpu, colorspace=colorspace, spatial=self.spatial, gpu_ids=gpu_ids, version=version) print('...[%s] initialized'%self.model.name()) print('...Done')