def create_ckpt_fpaths(num_items, opt_ckpt_size):
ckpts = list(range(0, num_items, opt_ckpt_size))
if np.max(np.array(ckpts)) < num_items:
ckpts.append(num_items)
for i, ckpt in enumerate(ckpts[:-1]):
n_start = file_utils.zpad(ckpt, num_zeros=cfg.CKPT_ZERO_PADDING)
n_end = file_utils.zpad(ckpts[i+1], num_zeros=cfg.CKPT_ZERO_PADDING)
ckpt_suffix = 'ckpt_{}_{}{}'.format(n_start, n_end, fpp_out.suffix) # 0_10.pkl
fp = join(str(fpp_out.parent), '{}_{}'.format(fpp_out.stem, ckpt_suffix))
ckpt_fpaths.append(fp)
return ckpt_fpaths
def load_images(self,
dir_media,
opt_size,
opt_density,
opt_drawframes=False):
sha256_tree = file_utils.sha256_tree(self._sha256)
dir_sha256 = join(dir_media, sha256_tree, self._sha256)
opt_size_label = cfg.IMAGE_SIZE_LABELS[opt_size]
# get the keyframe status data to check if images available
try:
keyframe_status = self.get_metadata(types.Metadata.KEYFRAME_STATUS)
except Exception as ex:
self.log.error(
'no keyframe metadata. Try: "append -t keyframe_status"')
return
keyframes = {}
# if keyframe images were generated and exist locally
if keyframe_status and keyframe_status.get_status(opt_size):
keyframe_metadata = self.get_metadata(types.Metadata.KEYFRAME)
if not keyframe_metadata:
self.log.error(
'no keyframe metadata. Try: "append -t keyframe"')
return
# get keyframe indices
frame_idxs = keyframe_metadata.get_keyframes(opt_density)
for frame_idx in frame_idxs:
# get keyframe filepath
fp_keyframe = join(dir_sha256, file_utils.zpad(frame_idx),
opt_size_label, 'index.jpg')
try:
im = cv.imread(fp_keyframe)
im.shape # used to invoke error if file didn't load correctly
except:
self.log.warn('file not found: {}'.format(fp_keyframe))
# don't add to keyframe dict
continue
keyframes[frame_idx] = im
# append metadata to chair_item's mapping item
self.set_keyframes(keyframes, opt_drawframes)
def cli(ctx, sink, opt_dir_media):
"""Saves keyframes for still-frame-video"""
# -------------------------------------------------
# imports
from os.path import join
import cv2 as cv
from vframe.utils import file_utils, logger_utils
from vframe.settings.paths import Paths
# -------------------------------------------------
# initialize
log = logger_utils.Logger.getLogger()
log.debug('init saves images')
file_utils.mkdirs(opt_dir_media)
# -------------------------------------------------
# process
frame_count = 0
while True:
chair_item = yield
for frame_idx, frame in chair_item.drawframes.items():
# save frame to the output folder
fp_im = join(opt_dir_media, 'frame_{}.png'.format(file_utils.zpad(frame_count)))
cv.imwrite(fp_im, frame)
frame_count += 1
# ------------------------------------------------------------
# send back to generator
sink.send(chair_item)
def thread_processor(task_obj):
tl = threading.local()
tl.fp_video = task_obj['fp_video']
tl.idxs = task_obj['idxs']
tl.dir_out = task_obj['dir_out']
tl.sha256_tree = task_obj['sha256_tree']
tl.sha256 = task_obj['sha256']
try:
tl.frame_ims = im_utils.vid2frames(tl.fp_video, idxs=tl.idxs)
except Exception as ex:
logging.getLogger().error('Could not read video file')
logging.getLogger().error('file: {}'.format(tl.fp_video))
logging.getLogger().error('sha256: {}'.format(tl.sha256))
return
tl.labels = cfg.IMAGE_SIZE_LABELS
tl.sizes = cfg.IMAGE_SIZES
for tl.k_label, tl.k_width in zip(reversed(tl.labels), reversed(tl.sizes)):
tl.label = tl.labels[tl.k_label]
tl.width = tl.sizes[tl.k_width]
# pyramid down frame sizes 1280, 640, 320, 160
try:
tl.frame_ims = [im_utils.resize(tl.im, width=tl.width) for tl.im in tl.frame_ims]
except:
logging.getLogger().error('')
logging.getLogger().error('Could not resize. Bad video or missing file')
logging.getLogger().error(tl.sha256)
logging.getLogger().error('')
return
for tl.idx, tl.im in zip(tl.idxs, tl.frame_ims):
# ensure path exists
tl.zpad = file_utils.zpad(tl.idx)
tl.fp_dst = join(tl.dir_out, tl.sha256_tree, tl.sha256, tl.zpad, tl.label, 'index.jpg')
# convert to PIL
tl.im_pil = im_utils.ensure_pil(tl.im, bgr2rgb=True)
file_utils.ensure_path_exists(tl.fp_dst)
tl.im_pil.save(tl.fp_dst, quality=cfg.JPG_SAVE_QUALITY)
def extract(items, dir_out, dir_videos, keyframe_type, threads=1):
"""Extracts keyframes from images"""
task_queue = Queue()
print_lock = threading.Lock()
log = logging.getLogger()
if threads > 1:
def thread_processor(task_obj):
tl = threading.local()
tl.fp_video = task_obj['fp_video']
tl.idxs = task_obj['idxs']
tl.dir_out = task_obj['dir_out']
tl.sha256_tree = task_obj['sha256_tree']
tl.sha256 = task_obj['sha256']
try:
tl.frame_ims = im_utils.vid2frames(tl.fp_video, idxs=tl.idxs)
except Exception as ex:
logging.getLogger().error('Could not read video file')
logging.getLogger().error('file: {}'.format(tl.fp_video))
logging.getLogger().error('sha256: {}'.format(tl.sha256))
return
tl.labels = cfg.IMAGE_SIZE_LABELS
tl.sizes = cfg.IMAGE_SIZES
for tl.k_label, tl.k_width in zip(reversed(tl.labels), reversed(tl.sizes)):
tl.label = tl.labels[tl.k_label]
tl.width = tl.sizes[tl.k_width]
# pyramid down frame sizes 1280, 640, 320, 160
try:
tl.frame_ims = [im_utils.resize(tl.im, width=tl.width) for tl.im in tl.frame_ims]
except:
logging.getLogger().error('')
logging.getLogger().error('Could not resize. Bad video or missing file')
logging.getLogger().error(tl.sha256)
logging.getLogger().error('')
return
for tl.idx, tl.im in zip(tl.idxs, tl.frame_ims):
# ensure path exists
tl.zpad = file_utils.zpad(tl.idx)
tl.fp_dst = join(tl.dir_out, tl.sha256_tree, tl.sha256, tl.zpad, tl.label, 'index.jpg')
# convert to PIL
tl.im_pil = im_utils.ensure_pil(tl.im, bgr2rgb=True)
file_utils.ensure_path_exists(tl.fp_dst)
tl.im_pil.save(tl.fp_dst, quality=cfg.JPG_SAVE_QUALITY)
def process_queue(num_items):
# TODO: progress bar
while True:
task_obj = task_queue.get()
thread_processor(task_obj)
logging.getLogger().info('process: {:.2f}% {:,}/{:,}'.format(
(task_queue.qsize() / num_items)*100, num_items-task_queue.qsize(), num_items))
task_queue.task_done()
# avoid race conditions by creating dir structure here
log.info('create directory structure first to avoid race conditions')
log.info('TODO: this needs to be fixed, thread lock maybe')
for sha256, item in tqdm(items.items()):
item_metadata = item.metadata.get(MetadataType.KEYFRAME, {})
sha256_tree = file_utils.sha256_tree(sha256)
fp_dst = join(dir_out, sha256_tree)
file_utils.ensure_path_exists(fp_dst)
# init threads
num_items = len(items)
for i in range(threads):
t = threading.Thread(target=process_queue, args=(num_items,))
t.daemon = True
t.start()
# process threads
start = time.time()
for sha256, item in items.items():
sha256_tree = file_utils.sha256_tree(sha256)
item_metadata = item.metadata.get(MetadataType.KEYFRAME, {})
if not item_metadata:
continue
keyframe_data = item_metadata.metadata
idxs = keyframe_data.get(keyframe_type)
fp_video = join(dir_videos, sha256_tree, '{}.{}'.format(sha256, item.ext))
task_obj = {
'fp_video': fp_video,
'idxs': idxs,
'dir_out':dir_out,
'sha256': sha256,
'sha256_tree': sha256_tree
}
task_queue.put(task_obj)
task_queue.join()
else:
for sha256, item in tqdm(items.items()):
item_metadata = item.metadata.get(MetadataType.KEYFRAME, {})
if not item_metadata:
continue
sha256_tree = file_utils.sha256_tree(sha256)
keyframe_data = item_metadata.metadata
#idxs_basic = keyframe_data.get(KeyframeMetadataType.BASIC)
#idxs_dense = keyframe_data.get(KeyframeMetadataType.DENSE)
#idxs_expanded = keyframe_data.get(KeyframeMetadataType.EXPANDED)
# fetches the metadata by the enum type from the custom click param
idxs = keyframe_data.get(keyframe_type)
# get frames from video
fp_video = join(dir_videos, sha256_tree, '{}.{}'.format(sha256, item.ext))
frame_ims = im_utils.vid2frames(fp_video, idxs=idxs)
labels = cfg.IMAGE_SIZE_LABELS
sizes = cfg.IMAGE_SIZES
for k_label, k_width in zip(reversed(labels), reversed(sizes)):
label = labels[k_label]
width = sizes[k_width]
# pyramid down frame sizes 1280, 640, 320, 160
frame_ims = [im_utils.resize(im, width=width) for im in frame_ims]
for idx, im in zip(idxs, frame_ims):
# ensure path exists
zpad = file_utils.zpad(idx)
fp_dst = join(dir_out, sha256_tree, sha256, zpad, label, 'index.jpg')
# conver to PIL
im_pil = im_utils.ensure_pil(im, bgr2rgb=True)
file_utils.ensure_path_exists(fp_dst)
im_pil.save(fp_dst, quality=cfg.JPG_SAVE_QUALITY)
def cli(ctx, sink, opt_disk, opt_density):
"""Generates KeyframeStatus metadata"""
# Recommended: Use Expanded density to check for all keyframes
# -------------------------------------------------
# imports
import os
from os.path import join
from pathlib import Path
from vframe.settings.paths import Paths
from vframe.settings import vframe_cfg as cfg
from vframe.utils import file_utils, logger_utils
from vframe.models.metadata_item import KeyframeStatusMetadataItem
# -------------------------------------------------
# process
log = logger_utils.Logger.getLogger()
# set paths
media_type = types.Metadata.KEYFRAME
metadata_type = types.Metadata.KEYFRAME_STATUS
dir_keyframes = Paths.media_dir(media_type,
data_store=opt_disk,
verified=ctx.opts['verified'])
# iterate sink
while True:
chair_item = yield
sha256 = chair_item.sha256
sha256_tree = file_utils.sha256_tree(sha256)
dir_parent = join(dir_keyframes, sha256_tree, sha256)
# check if keyframe metadata exists
keyframe_metadata_item = chair_item.item.get_metadata(
types.Metadata.KEYFRAME)
if not keyframe_metadata_item:
log.error(
'no keyframe metadata. try "append -t keyframe", {}'.format(
keyframe_metadata_item))
chair_item.item.set_metadata(metadata_type, {})
else:
# check if the keyframes images exist
status = {k: False for k in cfg.IMAGE_SIZE_LABELS}
if Path(dir_parent).exists():
# get keyframe numbers
idxs = keyframe_metadata_item.get_keyframes(opt_density)
for idx in idxs:
labels = [v for k, v in cfg.IMAGE_SIZE_LABELS.items()]
for k, label in cfg.IMAGE_SIZE_LABELS.items():
fpp_im = Path(dir_parent, file_utils.zpad(idx), label,
'index.jpg')
if fpp_im.exists():
status[k] = True
# append metadata to chair_item's mapping item
chair_item.item.set_metadata(
metadata_type, KeyframeStatusMetadataItem(status))
# -------------------------------------------------
# continue processing other items
sink.send(chair_item)