def prepare():
# Create the archives.
os.makedirs(OUT_DIR, exist_ok=True)
with concurrent.futures.ThreadPoolExecutor(max_workers=20) as executor:
with darc.DataArchive(darc_path, 'w') as archive:
archive_elements = []
def handle_crop_index(crop_index):
images = []
for image_index in range(images_per_crop):
images.append(combineNoisyData(crop_index, image_index))
images.append(combineCleanData(crop_index))
image_stack = np.stack(images)
return image_stack.astype(np.float16)
for crop_index in range(crop_count):
archive_elements.append(executor.submit(handle_crop_index, crop_index))
for i, archive_element in enumerate(archive_elements):
archive.append(archive_element.result(), chunks=(1, 128, 128, 3))
print("Adding crop {}/{} to archive.".format(i+1, crop_count))
def __init__(self, directory, dataset_mode, load_size):
self.directory = directory
if not os.path.isdir('dataset'):
raise Exception(
'Could not find the dataset directory \'dataset\'. Download the dataset first.'
)
self.full_path = os.path.join('dataset', '2afc', directory)
self.dataset_mode = dataset_mode
self.load_size = load_size
self.cache = {
'judge': [],
'p0': [],
'p1': [],
'ref': [],
'judge_path': [],
'p0_path': [],
'p1_path': [],
'ref_path': []
}
darc_path = self.getDarcPath()
if not os.path.exists(darc_path):
self.createDarc()
self.darc = darc.DataArchive(darc_path)
def handle_task(indices):
'''Constructs a minibatch with the given indices.'''
thread_id = threading.get_ident()
if not thread_id in _darcs:
_darcs[thread_id] = {}
judges, p0s, p1s, refs = [], [], [], []
for i in indices:
dataset, index = self._getDatasetByIndex(i)
dataset_path = dataset.getDarcPath()
if not dataset_path in _darcs[thread_id]:
# One Darc per thread.
_darcs[thread_id][dataset_path] = darc.DataArchive(
dataset_path)
db = _darcs[thread_id][dataset_path]
p_data = db["{}_p".format(index)]
judge_data = db["{}_judge".format(index)]
judges.append(judge_data.data())
p0s.append(p_data[0, :, :, :])
p1s.append(p_data[1, :, :, :])
refs.append(p_data[2, :, :, :])
return {
'judge': np.concatenate(judges),
'p0': np.stack(p0s),
'p1': np.stack(p1s),
'ref': np.stack(refs)
}
def __init__(self, path, validation=False): self.path = path self.darc = darc.DataArchive(path, 'r') assert len(self.darc) > 0 # Empty archive? self.crop_count = self.darc[0].shape[0] - 1 self.image_count = len(self.darc) * self.crop_count self.validation = validation self.dimensions = config.TOTAL_COUNT
def run_metrics(): # TODO HOX.
import tensorflow as tf
import elpips
import darc
import csv
# Build graph.
tf_X_input = tf.placeholder(tf.float32)
tf_Y_input = tf.placeholder(tf.float32)
tf_X = tf.expand_dims(tf_X_input, axis=0)
tf_Y = tf.expand_dims(tf_Y_input, axis=0)
tf_Y_grayscale = tf.reduce_mean(tf_Y, axis=3, keepdims=True)
tf_l2 = tf.reduce_mean(tf.square(tf_X - tf_Y))
tf_l1 = tf.reduce_mean(tf.abs(tf_X - tf_Y))
tf_relmse = tf.reduce_mean(
tf.square(tf_X - tf_Y) / (0.001 + tf.square(tf_Y_grayscale)))
# Note: It would be somewhat faster to just use n=args.elpips_sample_count but TF has
# problems with n > 1 on some GPUs.
elpips_vgg_model = elpips.Metric(elpips.elpips_vgg(n=1), back_prop=False)
tf_elpips_vgg = elpips_vgg_model.forward(tf_X, tf_Y)[0]
print("Creating Tensorflow session.")
tf_config = tf.ConfigProto(allow_soft_placement=True,
log_device_placement=False)
tf_config.gpu_options.per_process_gpu_memory_fraction = 0.8
with tf.Session(config=tf_config) as sess:
# Initialize model.
sess.run([
tf.global_variables_initializer(),
tf.local_variables_initializer()
])
# Iterate over the archives.
tasks = collections.deque()
with concurrent.futures.ThreadPoolExecutor(max_workers=5) as executor:
for archive_path in DATASETS:
# Reconstruct scene.
scene_name = get_scene_name(archive_path)
current_darc = darc.DataArchive(archive_path)
crop_count = current_darc[0].shape[0] - 1
image_count = len(current_darc)
# Read minibatches.
for image_index in range(image_count):
# Execute previous tasks.
while tasks:
task = tasks[0].result()
tasks.popleft()
print("Loading reference... ", end="")
sys.stdout.flush()
stime = time.time()
reference = current_darc[image_index][-1, :, :, 0:3]
etime = time.time()
print("Done in {:.01f}s.".format(etime - stime))
directory = os.path.join(OUT_DIRECTORY, config.model.name,
scene_name)
if not os.path.exists(directory):
print(
"Skipping directory '{}': Directory does not exist."
.format(directory))
continue
for crop_index in range(crop_count):
if args.crops and crop_index not in args.crops:
print("(Skipping scene {}, image {}, crop {}).".
format(scene_name, image_index, crop_index))
continue
crop_path = os.path.join(
directory, "img{:04d}_crop{:02d}".format(
image_index, crop_index))
if not os.path.exists(crop_path + ".npz"):
print("Skipping: Not found.")
continue
with open(
os.path.join(
directory,
"img{:04d}_results.{}.csv".format(
image_index, crop_index)),
'w') as csvfile:
fields = [
'crop_index', 'l1', 'l2', 'relmse',
'elpips-vgg', 'elpips-vgg-stdev'
]
csv_writer = csv.DictWriter(csvfile,
fieldnames=fields)
csv_writer.writeheader()
print(
"Handling scene {}, image {}, crop {}.".format(
scene_name, image_index, crop_index))
# Load image.
print("Loading image... ", end="")
sys.stdout.flush()
stime = time.time()
current_image = image.load_npz(crop_path + ".npz")
etime = time.time()
print("Done in {:.01f}s.".format(etime - stime))
# Run metrics.
print("Running metrics... ", end="")
sys.stdout.flush()
stime = time.time()
err_l1, err_l2, err_relmse = sess.run(
[tf_l1, tf_l2, tf_relmse],
feed_dict={
tf_X_input: current_image,
tf_Y_input: reference
})
print_dot()
err_elpips_vgg = []
for i in range(args.elpips_sample_count):
if i > 0 and i % 10 == 0:
print_dot()
err_elpips_vgg_single = sess.run(
tf_elpips_vgg,
feed_dict={
tf_X_input: current_image,
tf_Y_input: reference
})
err_elpips_vgg.append(err_elpips_vgg_single)
err_elpips_vgg_mean = np.mean(err_elpips_vgg)
err_elpips_vgg_std = np.std(
err_elpips_vgg, ddof=1) / np.sqrt(
args.elpips_sample_count)
etime = time.time()
print("Done in {:.01f}s.".format(etime - stime))
# Save results.
csv_writer.writerow({
'crop_index':
crop_index,
'l1':
err_l1,
'l2':
err_l2,
'relmse':
err_relmse,
'elpips-vgg':
err_elpips_vgg_mean,
'elpips-vgg-stdev':
err_elpips_vgg_std
})
def createDarc(self):
'''Loads a dataset of images from disk and stores it in the Darc format.
Darc is similar in idea to HDF5 but much simpler and should handle multi-process parallelization issues better.'''
print("Reading dataset '{}' from disk.".format(self.getCacheKey()))
archive = darc.DataArchive(self.getDarcPath(), 'w')
judge_files = cached_listdir(os.path.join(self.full_path, 'judge'))
judge_files = [
os.path.join(self.full_path, 'judge', file) for file in judge_files
if os.path.splitext(file)[1] == '.npy'
]
judge_files = sorted(judge_files)
p0_files = cached_listdir(os.path.join(self.full_path, 'p0'))
p0_files = [
os.path.join(self.full_path, 'p0', file) for file in p0_files
if os.path.splitext(file)[1] == '.png'
]
p0_files = sorted(p0_files)
p1_files = cached_listdir(os.path.join(self.full_path, 'p1'))
p1_files = [
os.path.join(self.full_path, 'p1', file) for file in p1_files
if os.path.splitext(file)[1] == '.png'
]
p1_files = sorted(p1_files)
ref_files = cached_listdir(os.path.join(self.full_path, 'ref'))
ref_files = [
os.path.join(self.full_path, 'ref', file) for file in ref_files
if os.path.splitext(file)[1] == '.png'
]
ref_files = sorted(ref_files)
def handle_task(judge_path, p0_path, p1_path, ref_path):
judge = np.load(judge_path)[0]
# Modify to use this code to use exactly the same data as the original LPIPS paper (as far as we know).
#
# This applies a bilinear 4x downscaling to the input images of resolution 256x256, which in reality
# are just 64x64 images upscaled to 256x256 with nearest neighbor. We do not know why this is the case.
#
# A couple of the images are, for some reason, 252x252, and there's a slight difference in the images in that case.
#####
#import PIL
#
#p0 = PIL.Image.open(p0_path).convert('RGB')
#p0 = p0.resize((self.load_size, self.load_size), PIL.Image.BILINEAR)
#p0 = np.array(p0) / 255.0
#
#p1 = PIL.Image.open(p1_path).convert('RGB')
#p1 = p1.resize((self.load_size, self.load_size), PIL.Image.BILINEAR)
#p1 = np.array(p1) / 255.0
#
#ref = PIL.Image.open(ref_path).convert('RGB')
#ref = ref.resize((self.load_size, self.load_size), PIL.Image.BILINEAR)
#ref = np.array(ref) / 255.0
#
#p0 = p0[:,:,0:3]
#p1 = p1[:,:,0:3]
#ref = ref[:,:,0:3]
#####
p0 = load_image_uint(p0_path)
p1 = load_image_uint(p1_path)
ref = load_image_uint(ref_path)
assert len(p0.shape) >= 3 and p0.shape[2] >= 3
assert len(p1.shape) >= 3 and p1.shape[2] >= 3
assert len(ref.shape) >= 3 and ref.shape[2] >= 3
p0 = p0[:, :, 0:3]
p1 = p1[:, :, 0:3]
ref = ref[:, :, 0:3]
if self.dataset_mode == '2afc' and p0.shape[0:2] != (
self.load_size, self.load_size):
p0 = skimage.transform.resize(
p0, [self.load_size, self.load_size, 3],
mode='reflect',
anti_aliasing=False)
p1 = skimage.transform.resize(
p1, [self.load_size, self.load_size, 3],
mode='reflect',
anti_aliasing=False)
ref = skimage.transform.resize(
ref, [self.load_size, self.load_size, 3],
mode='reflect',
anti_aliasing=False)
# Note: Images now in [0, 1].
return judge, p0, p1, ref, judge_path, p0_path, p1_path, ref_path
# Check that each dataset sample has judge, p0, p1 and reference.
if not all(
len(collection) == len(judge_files)
for collection in (p0_files, p1_files, ref_files)):
raise Exception('Dataset files missing!')
with concurrent.futures.ThreadPoolExecutor(max_workers=5) as executor:
postponed_image_sets = []
for i in range(len(judge_files)):
judge_path, p0_path, p1_path, ref_path = judge_files[
i], p0_files[i], p1_files[i], ref_files[i]
postponed_image_sets.append(
(judge_path, p0_path, p1_path, ref_path))
queued_image_sets = []
i = 0
while queued_image_sets or postponed_image_sets:
while len(queued_image_sets) < 20 and postponed_image_sets:
judge_path, p0_path, p1_path, ref_path = postponed_image_sets[
0]
queued_image_sets.append(
executor.submit(handle_task, judge_path, p0_path,
p1_path, ref_path))
postponed_image_sets.pop(0)
judge, p0, p1, ref, judge_path, p0_path, p1_path, ref_path = queued_image_sets[
0].result()
queued_image_sets.pop(0)
print("{}/{}".format(i, len(judge_files)))
p_tensor = np.stack([p0, p1, ref])
archive.append(p_tensor,
chunks=[1, -1, -1, -1],
name="{}_p".format(i))
judge_tensor = np.asarray(judge).reshape([1])
archive.append(judge_tensor, name="{}_judge".format(i))
i += 1
archive.close()
def extract_dataset():
import darc
if config.model.input_enabled('gradients'):
out_dir_input = "input-gpt"
out_dir_reference = "references"
else:
out_dir_input = "input-pt-2.5x"
out_dir_reference = "references"
# Iterate over the archives.
tasks = collections.deque()
with concurrent.futures.ThreadPoolExecutor(max_workers=10) as executor:
for archive_path in DATASETS:
# Get scene inputs.
scene_name = get_scene_name(archive_path)
current_darc = darc.DataArchive(archive_path)
crop_count = current_darc[0].shape[0] - 1
image_count = len(current_darc)
# Read minibatches.
for image_index in range(image_count):
# Execute previous tasks.
while tasks:
task = tasks[0].result()
tasks.popleft()
print("Loading metadata... ", end="")
stime = time.time()
sys.stdout.flush()
minibatch = current_darc[image_index]
etime = time.time()
print("Done in {:.01f}s.".format(etime - stime))
# Extract inputs.
directory = os.path.join(OUT_DIRECTORY, out_dir_input,
scene_name)
os.makedirs(directory, exist_ok=True)
for crop_index in range(crop_count):
print("Handling scene {}, image {}, crop {}.".format(
scene_name, image_index, crop_index))
print("Extracting inputs... ", end="")
stime = time.time()
def getSlice(which):
return minibatch[crop_index, :, :, which[0]:which[1]]
def saveSlice(data, slice_name):
tasks.append(
executor.submit(
image.save_png, data,
os.path.join(
directory,
"img{:04d}_crop{:02d}_{}.png".format(
image_index, crop_index, slice_name))))
tasks.append(
executor.submit(
image.save_npz, data,
os.path.join(
directory,
"img{:04d}_crop{:02d}_{}.npz".format(
image_index, crop_index, slice_name))))
saveSlice(getSlice(config.IND_PRIMAL_COLOR), 'primal')
print_dot()
saveSlice(getSlice(config.IND_DX_COLOR), 'dx')
print_dot()
saveSlice(getSlice(config.IND_DY_COLOR), 'dy')
print_dot()
saveSlice(getSlice(config.IND_ALBEDO), 'albedo')
print_dot()
saveSlice(getSlice(config.IND_DEPTH), 'depth')
print_dot()
saveSlice(getSlice(config.IND_NORMAL), 'normal')
print_dot()
saveSlice(getSlice(config.IND_DIFFUSE), 'diffuse')
print_dot()
saveSlice(getSlice(config.IND_SPECULAR), 'specular')
print_dot()
saveSlice(getSlice(config.IND_VAR_PRIMAL), 'var_primal')
print_dot()
saveSlice(getSlice(config.IND_VAR_DX), 'var_dx')
print_dot()
saveSlice(getSlice(config.IND_VAR_DY), 'var_dy')
print_dot()
saveSlice(getSlice(config.IND_VAR_ALBEDO), 'var_albedo')
print_dot()
saveSlice(getSlice(config.IND_VAR_DEPTH), 'var_depth')
print_dot()
saveSlice(getSlice(config.IND_VAR_NORMAL), 'var_normal')
print_dot()
saveSlice(getSlice(config.IND_VAR_DIFFUSE), 'var_diffuse')
print_dot()
saveSlice(getSlice(config.IND_VAR_SPECULAR),
'var_specular')
etime = time.time()
print("Done in {:.01f}s.".format(etime - stime))
# Extract references.
print("Extracting reference... ", end="")
stime = time.time()
directory = os.path.join(OUT_DIRECTORY, out_dir_reference,
scene_name)
os.makedirs(directory, exist_ok=True)
crop_name = os.path.join(directory,
"img{:04d}_ref".format(image_index))
tasks.append(
executor.submit(image.save_npz, minibatch[crop_count, :, :,
0:3],
crop_name + '.npz'))
tasks.append(
executor.submit(image.save_png, minibatch[crop_count, :, :,
0:3],
crop_name + ".png"))
etime = time.time()
print("Done in {:.01f}s.".format(etime - stime))