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