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)