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): # 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 main(args):
scenes = dataset.get_all_scenes(args.dataset)
m = model()
load_weights(m, args.model)
for scene, frames in scenes.items():
def read_image_lab(scene, frame_number):
bgr_image = dataset.read_image(dataset.get_frame_path(
scene, frame_number),
resolution=(256, 256))
return dataset.bgr_to_lab(bgr_image)
def predict(grayscale_input, ab_and_mask_input):
return m.predict(
{
'grayscale_input': np.expand_dims(grayscale_input, axis=0),
'ab_and_mask_input': np.expand_dims(ab_and_mask_input,
axis=0),
},
verbose=1)
for frame in frames:
l, ab = read_image_lab(scene, frame)
x = predict(l, ab_and_mask_matrix(ab, .1))
image = np.round(255 * dataset.lab_to_bgr(l, x[0])).astype('uint8')
cv2.imshow('Video', image)
cv2.waitKey(1)
def test_ab_and_mask_matrix(mock_apply_mask):
masked_data, mask = np.empty((64, 128, 2)), np.empty((64, 128, 1))
mock_apply_mask.return_value = masked_data, mask
data = np.empty((64, 128, 2))
m = ab_and_mask_matrix(data, .5)
mock_apply_mask.assert_called_once_with(data, .5)
assert m.shape == (64, 128, 3)
assert (m[:, :, :2] == masked_data).all()
assert (m[:, :, -1:] == mask).all()
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 thread(self): # pylint: disable=too-many-locals
encoder = encoder_model()
load_weights(encoder, self.weights, by_name=True)
with ConsumerPool(lambda args: np.savez_compressed(*args), queue_size=16) as save_consumer_pool:
while True:
l_batch, ab_and_mask_batch, filenames = [], [], []
for _ in range(self.batch_size):
job = self.queue.get()
if job is None:
self.queue.task_done()
break
lab, encoded_filename = job
l_batch.append(lab[0])
ab_and_mask_batch.append(ab_and_mask_matrix(lab[1], .00008))
filenames.append(encoded_filename)
self.queue.task_done()
if l_batch:
print('Encoding batch of size {}'.format(len(l_batch)))
start = time.time()
encoded_batch, _, _, _ = encoder.predict([np.array(x) for x in (l_batch, ab_and_mask_batch)])
print('Encoded in {}'.format(time.time() - start))
for filename, ab_and_mask, encoded_features in zip(filenames, ab_and_mask_batch, encoded_batch):
save_consumer_pool.put((filename, encoded_features))
save_consumer_pool.put(('{}_mask'.format(filename), ab_and_mask))
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)
def main(args): # pylint: disable=too-many-locals,too-many-statements
m = model()
load_weights(m, args.encoder, by_name=True)
load_weights(m, args.decoder, by_name=True)
video = args.video
capture = open_video(video)
writer = None
if args.save:
truth = random.choice(['L', 'R'])
fourcc = cv2.VideoWriter_fourcc(*'XVID')
if video.isdigit():
stem = video
else:
stem = Path(video).stem
filename = 'output_{}_{}.avi'.format(stem, truth)
print('Saving to {}'.format(filename))
writer = cv2.VideoWriter(filename, fourcc, 30.0, (512, 256))
truth = cv2.VideoWriter('truth_{}'.format(filename), fourcc, 30.0,
(256, 256))
colormotion = cv2.VideoWriter('colormotion_{}'.format(filename),
fourcc, 30.0, (256, 256))
def on_trackbar(value):
on_trackbar.value = value / 100000
on_trackbar(16)
if args.ui:
cv2.namedWindow('ColorMotion')
cv2.createTrackbar('Mask percentage * 0.1%', 'ColorMotion', 16, 100,
on_trackbar)
l_tm1 = None
prev = None
interpolated_features_tm1 = None
prev_mask = None
# while True:
for _, frame in zip(range(300), capture):
frame = cv2.resize(frame, (256, 256), interpolation=cv2.INTER_AREA)
l, ab = dataset.bgr_to_lab((frame / 255).astype(np.float32))
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, on_trackbar.value)]),
np.array([l]),
np.empty((1, 32, 32, 512))
])
else:
warped_features = warp_features(
l_tm1, l, interpolated_features_tm1)[np.newaxis]
mask = np.array([ab_and_mask_matrix(ab, on_trackbar.value)])
if prev_mask is not None:
prev_mask[:, :, 2] *= .8 # pylint: disable=unsupported-assignment-operation
# mask_valid = mask[:, :, :, 2:3]
# condition = np.stack((mask_valid, ) * 3, axis=-1)
# mask = np.where(condition, mask, prev_mask)
mask += prev_mask
x, _, interpolated_features = m.predict([
np.array([ab_and_mask_matrix(ab, on_trackbar.value)]),
np.array([l]), warped_features
])
prev_mask = mask
ab = x[0]
if prev is not None:
ab = (ab + 2 * prev) / 3
prev = ab
bgr = np.round(255 * dataset.lab_to_bgr(l, ab)).astype('uint8')
if writer:
output = np.concatenate((bgr, frame), axis=1)
writer.write(output)
truth.write(frame)
colormotion.write(bgr)
if args.ui:
cv2.imshow('Original stream', frame)
cv2.imshow('ColorMotion', bgr)
if (cv2.waitKey(1) & 255) == ord('q'):
break
interpolated_features_tm1 = interpolated_features[0]
l_tm1 = l