def load_batch(self, start_frames, target_size): # pylint: disable=too-many-locals
assert self.contiguous_count == 1
x_batch = [[], []]
y_batch = []
self.skip_connections_pipe.send(
(start_frames, self.encoded_features_path))
for scene, frame in start_frames:
features_tm1 = np.load('{}_encoded.npz'.format(
encoded_feature_path(self.encoded_features_path, scene,
frame)))['arr_0']
features = np.load('{}_encoded.npz'.format(
encoded_feature_path(self.encoded_features_path, scene,
frame + self.contiguous_count)))['arr_0']
l, ab = dataset.read_frame_lab(scene,
frame + self.contiguous_count,
target_size)
l_tm1, _ = dataset.read_frame_lab(scene, frame, target_size)
warped_features = warp_features(l_tm1, l, features_tm1)
x_batch[0].append(warped_features)
x_batch[1].append(features)
y_batch.append(ab)
for skip_connection in self.skip_connections_pipe.recv():
x_batch.append(skip_connection)
return [np.array(model_input)
for model_input in x_batch], np.array(y_batch)
def load_batch(self, start_frames, target_size): # pylint: disable=too-many-locals
assert self.contiguous_count == 1
x_batch = [[], [], []]
y_batch = []
for scene, frame in start_frames:
l, ab = dataset.read_frame_lab(scene,
frame + self.contiguous_count,
target_size)
l_tm1, ab_tm1 = dataset.read_frame_lab(scene, frame, target_size)
# TODO implement augmentation
# TODO generate artificial flow data from ImageNet
features_tm1 = self.model.predict([
np.array([ab_and_mask_matrix(ab_tm1, .00016)]),
np.array([l_tm1]),
np.empty((1, 32, 32, 512)),
])[1]
warped_features = warp_features(l_tm1, l, features_tm1)
x_batch[0].append(ab_and_mask_matrix(ab, .00016))
x_batch[1].append(l)
x_batch[2].append(warped_features)
y_batch.append(ab)
return [np.array(model_input)
for model_input in x_batch], np.array(y_batch)
def main(args):
model = load_model(str(args.model))
scenes = dataset.get_all_scenes(args.dataset)
for scene, frames in scenes.items():
def write_image_lab(filename, l, ab):
image = np.round(255 * dataset.lab_to_bgr(l, ab)).astype('uint8')
if not cv2.imwrite(filename, image):
raise RuntimeError('Failed to write {}'.format(filename))
def predict(state, grayscale_input):
return model.predict(
{
'state_input': np.expand_dims(state, axis=0),
'grayscale_input': np.expand_dims(grayscale_input, axis=0)
},
verbose=1)
scene_number = scene.stem
state = np.dstack(dataset.read_frame_lab(scene, frames[0], (256, 256)))
for i, frame in enumerate(frames[1:]):
l, _ = dataset.read_frame_lab(scene, frame, (256, 256))
output = predict(state, l)
write_image_lab('output_{:04d}_{:04d}.png'.format(scene_number, i),
l, output[0])
state = np.dstack((l, output[0]))
def load_sample(self, scene, start_frame, target_size):
'''Load a sample to build a batch.'''
# y = expected colorization in last frame
# state = previous frames colorized and current frame in grayscale
grayscale, y = dataset.read_frame_lab(scene, start_frame + self.contiguous_count, target_size)
state = [
dataset.read_frame_lab(scene, start_frame + i, target_size)[0]
for i in range(self.contiguous_count)
]
return state + [grayscale, y]
def main(args): # pylint: disable=too-many-locals
scenes = dataset.get_all_scenes(args.dataset)
m = model()
load_weights(m, args.encoder, by_name=True)
load_weights(m, args.decoder, by_name=True)
for scene, frames in scenes.items():
l_tm1 = None
for frame in frames:
l, ab = dataset.read_frame_lab(scene, frame, (256, 256))
if l_tm1 is None:
# Set warped_features = encoded_features on the first frame
_, warped_features, _ = m.predict([
np.array([ab_and_mask_matrix(ab, .00016)]),
np.array([l]),
np.empty((1, 32, 32, 512))
],
verbose=1)
else:
warped_features = warp_features(
l_tm1, l, interpolated_features_tm1)[np.newaxis]
x, _, interpolated_features = m.predict([
np.array([ab_and_mask_matrix(ab, .00016)]),
np.array([l]), warped_features
],
verbose=1)
image = np.round(255 * dataset.lab_to_bgr(l, x[0])).astype('uint8')
cv2.imshow('Video', image)
cv2.waitKey(1)
interpolated_features_tm1 = interpolated_features[0]
l_tm1 = l
def load_batch(self, start_frames, target_size):
assert self.contiguous_count == 0
x_batch = [[], []]
y_batch = []
for scene, frame in start_frames:
l, ab = dataset.read_frame_lab(scene, frame, target_size)
if self.augmentation:
x = np.dstack((l, ab))
x = self.augment(x)
l, ab = x[:, :, :1], x[:, :, 1:]
x_batch[0].append(ab_and_mask_matrix(ab, .00016))
x_batch[1].append(l)
y_batch.append(ab)
return [np.array(model_input)
for model_input in x_batch], np.array(y_batch)
def main(args):
scenes = dataset.get_all_scenes(args.dataset)
encoded_total = 0
with InferenceConsumerPool(args.weights) as consume_pool:
try:
for scene, frames in scenes.items():
for frame in frames:
frame_path = dataset.get_frame_path(scene, frame)
encoded_path = args.encoded / '{}_encoded'.format(frame_path.relative_to(args.dataset))
encoded_path.parents[1].mkdir(exist_ok=True)
encoded_path.parent.mkdir(exist_ok=True)
consume_pool.put((
dataset.read_frame_lab(scene, frame, args.resolution),
encoded_path,
))
encoded_total += len(frames)
print('Total encoded: {}'.format(encoded_total))
finally:
consume_pool.put(None)
def skip_connections_eval_cpu(pipe, weights, target_size):
tf_allow_growth()
encoder = encoder_head_model()
load_weights(encoder, weights, by_name=True)
while True:
try:
start_frames, encoded_features_path = pipe.recv()
except EOFError:
break # end of data from parent process
l_batch, ab_and_mask_matrix_t_batch = [], []
for scene, frame in start_frames:
l, _ = dataset.read_frame_lab(scene, frame + 1, target_size)
l_batch.append(l)
ab_and_mask = np.load('{}_encoded_mask.npz'.format(
encoded_feature_path(encoded_features_path, scene,
frame + 1)))['arr_0']
ab_and_mask_matrix_t_batch.append(ab_and_mask)
pipe.send(
encoder.predict(
[np.array(x) for x in (l_batch, ab_and_mask_matrix_t_batch)]))
pipe.close()
def encoder_eval(queue, train_pipe, validation_pipe, weights, target_size): # pylint: disable=too-many-locals
tf_allow_growth()
image_data_generator = ImageDataGenerator()
# inputs: l_input, ab_and_mask_input
# outputs: encoded_features, conv1_2norm, conv2_2norm, conv3_3norm
encoder = encoder_model()
load_weights(encoder, weights, by_name=True)
while True:
data = queue.get()
if data is None:
# If end of data, return (will automatically GC the pipes)
return
start_frames, augment = data
# decoder inputs: warped_features, features, conv1_2norm, conv2_2norm, conv3_3norm
# decoder outputs: x
y_batch = []
l_batch = []
l_tm1_batch = []
ab_mask_batch = []
ab_mask_tm1_batch = []
for scene, frame in start_frames:
l_tm1, ab_tm1 = dataset.read_frame_lab(scene, frame, target_size)
if augment:
transform = random_augmentation()
l_tm1, ab_tm1 = augment_l_ab(image_data_generator,
(l_tm1, ab_tm1), transform)
if isinstance(frame, int):
l, ab = dataset.read_frame_lab(scene, frame + 1, target_size)
if augment:
l, ab = augment_l_ab(image_data_generator, (l, ab),
transform)
else:
# Augment artificially
l, ab = augment_l_ab(image_data_generator, (l_tm1, ab_tm1),
small_flow_transform())
l_tm1_batch.append(l_tm1)
ab_mask_tm1_batch.append(ab_and_mask_matrix(ab_tm1, .00016))
l_batch.append(l)
ab_mask_batch.append(ab_and_mask_matrix(ab, .00008))
y_batch.append(ab)
features_tm1 = encoder.predict([
np.array(l_tm1_batch),
np.array(ab_mask_tm1_batch),
])[0]
warped_features = [
warp_features(l_tm1, l,
feature_tm1) for l_tm1, l, feature_tm1 in zip(
l_tm1_batch, l_batch, features_tm1)
]
features, conv1_2norm, conv2_2norm, conv3_3norm = encoder.predict([
np.array(l_batch),
np.array(ab_mask_batch),
])
x_y_batch = (
[
np.array(warped_features), features, conv1_2norm, conv2_2norm,
conv3_3norm
],
np.array(y_batch),
)
if augment:
train_pipe.send(x_y_batch)
else:
validation_pipe.send(x_y_batch)
