def write_video_scenes(video, scene_boundaries, dataset_path, resolution): if not scene_boundaries: return scene_boundaries = iter(scene_boundaries) destination = dataset.create_video_destination_folder(video, dataset_path, exist_ok=True) video = cv2.VideoCapture(str(video)) scene_lower, scene_upper = next(scene_boundaries) def write(path, frame): if path.exists(): return frame = cv2.resize(frame, (resolution[0], resolution[1]), interpolation=cv2.INTER_AREA) cv2.imwrite(str(path), frame, [cv2.IMWRITE_PNG_COMPRESSION, 9]) with ConsumerPool(lambda args: write(*args)) as write_consumer_pool: for frame_count in count(): has_frame, frame = video.read() if not has_frame: return if scene_lower <= frame_count <= scene_upper: scene = dataset.get_scene_directory(destination, scene_lower) write_consumer_pool.put(( dataset.get_frame_path(scene, frame_count), frame, )) elif frame_count > scene_upper: scene_lower, scene_upper = next(scene_boundaries, (None, None)) if scene_lower is None: return
def consume_function(args): frame_data, previous_data, i = args frame, ready_event, _ = frame_data previous_frame, previous_ready, _ = previous_data or (None, ) * 3 if is_new_scene(frame, previous_frame): print('Frame {} is a new scene'.format(i)) scene_destination = dataset.get_scene_directory(destination, i) else: # If this frame is a continuation of the previous scene, save to the same folder previous_ready.wait() scene_destination = previous_data[2] frame_data[2] = scene_destination # Destination has been processed and folder created if needed, so it's ready to save the next frame ready_event.set() cv2.imwrite(str(dataset.get_frame_path(scene_destination, i)), frame, [cv2.IMWRITE_PNG_COMPRESSION, 9])
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 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 encoded_feature_path(root, scene, frame): return dataset.get_frame_path(root, scene.relative_to(scene.parents[1]), frame)
def read_image_lab(scene, frame_number): bgr_image = dataset.read_image(str(dataset.get_frame_path(scene, frame_number)), resolution=target_size) return dataset.to_lab(bgr_image)
def test_get_frame_path(scene_directory, frame_number, expected): assert dataset.get_frame_path(scene_directory, frame_number) == Path(expected)