def _make_video(video_path, imgs):
"""Code used to generate a video using cv2.
Parameters:
video_path: a path ending with .mp4, for instance: "/results/pose2d.mp4"
imgs: an iterable or generator with the images to turn into a video
"""
first_frame = next(imgs)
imgs = itertools.chain([first_frame], imgs)
shape = int(first_frame.shape[1]), int(first_frame.shape[0])
logger.debug('Saving video to: ' + video_path)
fourcc = cv2.VideoWriter_fourcc(*'mp4v')
fps = 30
output_shape = _resize(current_shape=shape, new_width=video_width)
logger.debug('Video size is: {}'.format(output_shape))
video_writer = cv2.VideoWriter(video_path, fourcc, fps, output_shape)
progress_bar = tqdm if logger.info_enabled() else lambda x: x
for img in progress_bar(imgs):
resized = cv2.resize(img, output_shape)
rgb = cv2.cvtColor(resized, cv2.COLOR_BGR2RGB)
video_writer.write(rgb)
video_writer.release()
logger.info('Video created at {}\n'.format(video_path))
def process_folder(model, loader, unlabeled, output_folder, overwrite,
num_classes, acc_joints):
save_path_pred, save_path_heatmap = (
get_save_path_pred(unlabeled, output_folder),
get_save_path_heatmap(unlabeled, output_folder),
)
if os.path.isfile(save_path_pred) and not overwrite:
logger.info("Prediction file exists, skipping pose estimation")
return None, None
elif os.path.isfile(save_path_pred) and overwrite:
logger.info("Overwriting existing predictions")
save_path_heatmap.parent.mkdir(exist_ok=True, parents=True)
save_path_pred.parent.mkdir(exist_ok=True, parents=True)
logger.debug(f"creaint heatmap path: {save_path_heatmap}")
heatmap = np.memmap(
filename=save_path_heatmap,
dtype="float32",
mode="w+",
shape=(
config["num_cameras"] + 1,
loader.dataset.greatest_image_id() + 1,
config["num_predict"],
config["heatmap_shape"][0],
config["heatmap_shape"][1],
),
) # num_cameras+1 for the mirrored camera 3
logger.debug(f"creating heatmap shape: {heatmap.shape}")
pred, heatmap, _, _, _ = step(
loader=loader,
model=model,
optimizer=None,
mode=Mode.test,
heatmap=heatmap,
epoch=0,
num_classes=num_classes,
acc_joints=acc_joints,
)
_, cid2cidread = read_camera_order(
get_output_path(unlabeled, output_folder))
cid_to_reverse = config["flip_cameras"]
cid_read_to_reverse = [cid2cidread[cid] for cid in cid_to_reverse]
pred = flip_pred(pred, cid_read_to_reverse)
logger.debug("Flipping heatmaps")
heatmap = flip_heatmap(heatmap, cid_read_to_reverse)
logger.debug("Flipping heatmaps")
save_dict(pred, save_path_pred)
if type(heatmap) != np.memmap:
save_dict(heatmap, save_path_heatmap)
print(pred.shape)
return pred, heatmap
def run_recursive(args):
"""Processes every subfolder named 'images' in the args.input_folder folder.
Parameters:
args: the parsed command-line arguments (see: parse_cli_args())
"""
subfolder_name = "images"
msg = f"{Style.BRIGHT}Recursively looking for subfolders named `{subfolder_name}` inside `{args.input_folder}`{Style.RESET_ALL}"
logger.info(msg)
subfolders = find_subfolders(args.input_folder, "images")
s = "s" if len(subfolders) > 1 else ""
folders_str = "\n-".join(subfolders)
logger.info(f"Found {len(subfolders)} subfolder{s}:\n-{folders_str}")
args.recursive = False
run_in_folders(args, subfolders)
def run(args):
"""Processes the image folder args.input_folder.
Parameters:
args: the parsed command-line arguments (see: parse_cli_args())
"""
nothing_to_do = args.skip_estimation and (not args.video_2d) and (
not args.video_3d)
if nothing_to_do:
logger.info(
f"{Style.BRIGHT}Nothing to do. Check your command-line arguments.{Style.RESET_ALL}"
)
return 0
logger.info(
f"{Style.BRIGHT}\nWorking in {args.input_folder}{Style.RESET_ALL}")
core = Core(args.input_folder, args.output_folder, args.num_images_max)
core.overwrite = args.overwrite # monkey-patch: save it for later
if args.camera_ids is not None:
core.update_camera_ordering(args.camera_ids)
if not args.skip_estimation:
core.pose2d_estimation(core.overwrite)
core.calibrate_calc(0, core.max_img_id)
core.save_pose()
else:
core.calibrate_calc(0, core.max_img_id)
if args.video_2d:
video.make_pose2d_video(core.plot_2d, core.num_images,
core.input_folder, core.output_folder)
if args.video_3d:
video.make_pose3d_video(
core.get_points3d(),
core.plot_2d,
core.num_images,
core.input_folder,
core.output_folder,
)
return 0
def run_from_file(args):
"""Processes every folder listed in the args.input_folder text file.
Parameters:
args: the parsed command-line arguments (see: parse_cli_args())
"""
logger.info(
f"{Style.BRIGHT}Looking for folders listed in {args.input_folder}{Style.RESET_ALL}"
)
try:
with open(args.input_folder, "r") as f:
folders = list(line.strip() for line in f)
except FileNotFoundError:
logger.error(f"Unable to find the file {args.input_folder}")
return 1
except IsADirectoryError:
logger.error(
f"{args.input_folder} is a directory, please provide a file instead."
)
return 1
folders = list(dict.fromkeys(folders)) # removes duplicate entries
folders = [f for f in folders if f.strip()] # remove blank lines
folders = [Path(f) for f in folders] # convert to path objects
bad = [f for f in folders if not f.is_dir()]
for f in bad:
logger.error(f"[Error] Not a directory or does not exist: {str(f)}")
if bad:
return 1
s = "s" if len(folders) > 1 else ""
folders_str = "\n-".join(folders)
logger.info(f"Folder{s} found:\n-{folders_str}")
args.from_file = False
run_in_folders(args, folders)