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 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):
child_pipe, parent_pipe = Pipe()
p = Process(target=skip_connections_eval_cpu,
args=(child_pipe, args.encoder_weights, (256, 256)))
p.start()
tf_allow_growth()
train_generator, _ = data_generators(args.dataset,
args.encoded_features_path,
parent_pipe)
checkpoint = ModelCheckpoint('epoch-{epoch:03d}-{loss:.3f}.h5',
verbose=1,
period=1)
decoder = interpolate_and_decode()
load_weights(decoder, args.weights, by_name=True)
fit = decoder.fit_generator(
train_generator,
steps_per_epoch=args.steps_per_epoch,
epochs=args.epochs,
# validation_data=test_generator,
# validation_steps=args.validation_steps,
callbacks=[checkpoint])
decoder.save('optical_flow_decoder.h5')
parent_pipe.close()
p.join()
print(fit.history)
def main(args):
train_generator, test_generator = data_generators(args.dataset)
m = model()
if args.weights:
load_weights(m, args.weights)
checkpoint = ModelCheckpoint('epoch-{epoch:03d}-{val_loss:.3f}.hdf5', verbose=1, period=5)
fit = m.fit_generator(
train_generator,
steps_per_epoch=2000,
epochs=80,
validation_data=test_generator,
validation_steps=10,
callbacks=[checkpoint])
print(fit.history)
m.save('Colorful_model.h5')
def main(args):
m = model()
train_generator, _ = data_generators(m, args.dataset)
checkpoint = ModelCheckpoint('epoch-{epoch:03d}-{loss:.3f}.h5',
verbose=1,
period=1)
load_weights(m, args.weights, by_name=True)
fit = m.fit_generator(
train_generator,
steps_per_epoch=args.steps_per_epoch,
epochs=args.epochs,
# validation_data=test_generator,
# validation_steps=args.validation_steps,
callbacks=[checkpoint])
m.save('optical_flow_decoder.h5')
print(fit.history)
def main(args):
train_generator, test_generator = data_generators(args.dataset)
m = model()
if args.weights:
load_weights(m, args.weights)
checkpoint = ModelCheckpoint('epoch-{epoch:03d}-{val_loss:.3f}.hdf5',
verbose=1,
period=5)
fit = m.fit_generator(train_generator,
steps_per_epoch=args.steps_per_epoch,
epochs=args.epochs,
validation_data=test_generator,
validation_steps=args.validation_steps,
callbacks=[checkpoint])
print(fit.history)
if not args.dry:
m.save('user_guided_model.h5')
def train_decoder(queue, train_recv_pipe, validation_recv_pipe, args):
tf_allow_growth(.5)
train_generator, validation_generator = data_generators(
args.dataset, queue, train_recv_pipe, validation_recv_pipe)
checkpoint = ModelCheckpoint(
'epoch-{epoch:03d}-{loss:.3f}-{val_loss:.3f}.h5', verbose=1, period=5)
decoder = interpolate_and_decode()
if args.weights:
load_weights(decoder, args.weights, by_name=True)
fit = decoder.fit_generator(train_generator,
steps_per_epoch=args.steps_per_epoch,
epochs=args.epochs,
validation_data=validation_generator,
validation_steps=args.validation_steps,
callbacks=[checkpoint])
decoder.save('optical_flow_decoder.h5')
print(fit.history)
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 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 prune(args):
_, validation_generator = data_generators(args.dataset)
validation_generator = GeneratorToObject(validation_generator)
encoder = encoder_model()
load_weights(encoder, args.weights, by_name=True)
percent_pruned = 0
percent_pruning = 80
# while percent_pruned <= .3:
total_channels = get_total_channels(encoder)
n_channels_delete = int(percent_pruning // 100 * total_channels)
# Prune the model
apoz_df = get_model_apoz(encoder, validation_generator)
# percent_pruned += percent_pruning
print('pruning up to {}% of the original model weights'.format(
percent_pruned))
model = prune_model(encoder, apoz_df, n_channels_delete)
# Clean up tensorflow session after pruning and re-load model
output = str(
args.weights.parent /
'compressed_{}%_{}'.format(percent_pruning, args.weights.name))
print('Saving to {}'.format(output))
model.compile(loss='mean_squared_error', optimizer='adam')
model.save(output)
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