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)
