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 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_get_all_scenes():
dataset_directory = base_dir / 'test_dataset'
frames = dataset.get_all_scenes(dataset_directory)
assert frames == OrderedDict([
(dataset_directory / 'movie/000000', [0, 1]),
(dataset_directory / 'movie/000002', [2]),
])
def flow_from_directory(self, root, batch_size=32, target_size=None, seed=None):
contiguous_frames = self.get_contiguous_frames(dataset.get_all_scenes(root))
print('Dataset {} has {} contiguous subscenes'.format(root, len(contiguous_frames)))
random = Random(seed)
while True:
yield self.load_batch(random.choices(contiguous_frames, k=batch_size),
target_size=target_size)
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 flow_from_directory(self, root, batch_size=32, target_size=None): # pylint: disable=arguments-differ
contiguous_frames = self.get_contiguous_frames(
dataset.get_all_scenes(root, names_as_int=False))
print('Dataset {} has {} contiguous subscenes'.format(
root, len(contiguous_frames)))
while True:
yield self.load_batch(random.choices(contiguous_frames,
k=batch_size),
target_size=target_size)
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 main(args):
model = load_model(str(args.model))
scenes = dataset.get_all_scenes(args.dataset)
for scene, frames in scenes.items():
def read_image_l(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)[0]
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(grayscale_input):
return model.predict({
'grayscale_input': np.expand_dims(grayscale_input, axis=0)
}, verbose=1)
scene_number = scene.stem
for i, frame in enumerate(frames[1:]):
l = read_image_l(scene, frame)
output = predict(l)
write_image_lab('output_{:04d}_{:04d}.png'.format(scene_number, i), l, output[0])
def main(args):
dataset = get_all_scenes(args.dataset)
metadata = build_metadata(dataset)
with open(args.metadata, 'w') as f:
json.dump(metadata, f)
