def __call__(self, tarinfos_pair):
basename, tarinfos = tarinfos_pair
json_source = ShardedJsonDumpSource(
self.mode, self.tar_path, tarinfos, self.suppress_end_fail
)
path = pjoin(self.out, basename + ".untracked.h5")
stats: Counter = Counter()
if self.skip_existing and os.path.exists(path):
return stats
os.makedirs(os.path.dirname(path), exist_ok=True)
h5ctx = h5out(path)
h5f = h5ctx.__enter__()
try:
try:
write_conv(h5f, self.mode, basename, json_source, "monolithic-tar")
except Exception:
print(f"Exception while dumping to {path}", file=sys.stderr)
h5f.attrs["fatal_exception"] = True
traceback.print_exc(file=sys.stderr)
stats["total_fatal_exceptions"] += 1
else:
stats["total_dumps"] += 1
stats["total_frames"] += json_source.num_frames
stats["total_corrupt_frames"] += json_source.corrupt_frames
stats["total_corrupt_shards"] += json_source.corrupt_shards
stats["total_remaining_heaps"] += json_source.remaining_heaps
if json_source.end_fail:
stats["total_end_fail"] += 1
finally:
try:
h5f.__exit__(None, None, None)
except Exception:
print(f"Exception while trying to close {path}", file=sys.stderr)
traceback.print_exc(file=sys.stderr)
return stats
def maybe_h5out(h5fn, dry_run):
from skelshop.utils.h5py import h5out
if dry_run:
yield None
else:
with h5out(h5fn) as h5f:
yield h5f
def embedall(
face_extractor,
video,
h5fn,
from_skels,
start_frame,
skel_thresh_pool,
skel_thresh_val,
batch_size,
write_bboxes,
write_chip,
):
"""
Create a HDF5 face dump from a video using dlib.
"""
from skelshop.face.pipe import all_faces_from_skel_batched, iter_faces_from_dlib
check_extractor(face_extractor, from_skels)
extractor_info = EXTRACTORS[face_extractor]
num_frames = count_frames(video) - start_frame
with ExitStack() as stack:
h5f = stack.enter_context(h5out(h5fn))
add_basic_metadata(h5f, video, num_frames)
has_fod = extractor_info["type"] == "dlib"
writer = FaceWriter(
h5f,
write_fod_bbox=write_bboxes and has_fod,
write_chip_bbox=write_bboxes,
write_chip=write_chip,
)
kwargs = {
"include_chip": write_chip,
"include_bboxes": write_bboxes,
}
if batch_size is not None:
kwargs["batch_size"] = batch_size
if extractor_info["type"] == "dlib":
vid_read = stack.enter_context(load_video_rgb(video))
face_iter = iter_faces_from_dlib(
vid_read,
detector=extractor_info["detector"],
keypoints=extractor_info["keypoints"],
**kwargs,
)
else:
vid_read = decord_video_reader(video)
skel_read = open_skels(from_skels)
mode = mode_of_extractor_info(extractor_info)
face_iter = all_faces_from_skel_batched(
vid_read,
skel_read,
thresh_pool=skel_thresh_pool,
thresh_val=skel_thresh_val,
mode=mode,
**kwargs,
)
write_faces(face_iter, writer)
def conv(input_fmt, legacy_dump, out, mode, cores, suppress_end_fail, skip_existing):
"""
Convert a exiting dump from another format into HDF5 format.
LEGACY_DUMP is the dump in the old format. OUT is a file path when run with
single-zip, otherwise it is the base of a directory tree which will be
created during processing.
"""
if input_fmt != "monolithic-tar":
if cores != 1:
raise click.UsageError(
"--cores must be 1 for INPUT_FMT other than monolithic-tar"
)
if out is None:
raise click.UsageError(
"OUT required for INPUT_FMT other than monolithic-tar"
)
if input_fmt == "monolithic-tar":
set_start_method("forkserver")
if out is None:
out = ""
stats: Counter = Counter()
processor = TarInfosProcessor(
mode,
legacy_dump,
suppress_end_fail,
skip_existing,
out,
iter_tarinfos(legacy_dump),
processes=cores,
)
for new_stats in processor:
stats += new_stats
print("Stats", stats)
elif input_fmt in ("single-zip", "ordered-tar"):
if input_fmt == "single-zip":
json_source_ctx = zip_json_source(mode, legacy_dump)
else:
json_source_ctx = ordered_tar_source(mode, legacy_dump)
with h5out(out) as h5f, json_source_ctx as json_source:
write_conv(h5f, mode, None, json_source, input_fmt)
else:
assert False
def filter(h5infn, h5outfn, pipeline, compression, start_frame, end_frame):
"""
Apply tracking to an untracked HDF5 pose dump.
"""
with h5py.File(h5infn, "r") as h5in, h5out(h5outfn) as h5fout:
for attr, val in h5in.attrs.items():
h5fout.attrs[attr] = val
pipeline.apply_metadata(h5fout)
add_fmt_metadata(h5fout, "trackshots")
limbs = h5in.attrs["limbs"]
frame_iter = UnsegmentedReader(h5in).iter_from(start_frame)
if end_frame is not None:
frame_iter = take(end_frame - start_frame, frame_iter)
stage = IterStage(frame_iter)
frame_iter = pipeline(stage)
lossless_kwargs, lossy_kwargs = COMPRESSIONS[compression]
write_shots(
h5fout,
limbs,
frame_iter,
start_frame=start_frame,
lossless_kwargs=lossless_kwargs,
lossy_kwargs=lossy_kwargs,
)
def embedselect(
video,
selection,
h5fn,
from_skels,
batch_size: Optional[str],
write_bboxes,
write_chip,
):
"""
Embed faces into a sparse face dump according to a predetermined selection
of frame-person pairs.
"""
from skelshop.face.pipe import select_faces_from_skel_batched
batch_size_num = process_batch_size(batch_size)
vid_read = decord_video_reader(video)
next(selection)
# First group by extractor, then by frame
grouped: Dict[str, Dict[int, Tuple[int, int, List[int]]]] = {}
for line in selection:
seg, pers_id, seg_frame_num, abs_frame_num, extractor = line.strip(
).split(",")
grouped.setdefault(extractor, {}).setdefault(
int(abs_frame_num),
(int(seg), int(seg_frame_num), []))[2].append(int(pers_id))
for extractor in grouped.keys():
check_extractor(extractor, from_skels)
skels_h5 = h5py.File(from_skels, "r")
skel_read = ShotSegmentedReader(skels_h5, infinite=False)
merged = ResultMerger()
for extractor, frames in grouped.items():
extractor_info = EXTRACTORS[extractor]
mode = mode_of_extractor_info(extractor_info)
targets = [(abs_frame_num, *frame_info)
for abs_frame_num, frame_info in frames.items()]
merged.add_results(
((frame_num, pers_ids) for frame_num, _, _, pers_ids in targets),
select_faces_from_skel_batched(
iter(targets),
vid_read,
skel_read,
batch_size=batch_size_num,
mode=mode,
include_bboxes=write_bboxes,
include_chip=write_chip,
),
)
num_frames = count_frames(video)
with h5out(h5fn) as h5f:
add_basic_metadata(h5f, video, num_frames)
writer = FaceWriter(
h5f,
sparse=True,
write_chip_bbox=write_bboxes,
write_chip=write_chip,
)
for (abs_frame, pers_ids), face in merged:
writer.write_frame_faces(
face["embeddings"],
[(abs_frame, pers_id) for pers_id in pers_ids],
chip_bboxes=face.get("chip_bboxes"),
face_chips=face.get("chips"),
)