(cfg)
| 98 | |
| 99 | @torch.no_grad() |
| 100 | def run_preprocess(cfg): |
| 101 | Log.info(f"[Preprocess] Start!") |
| 102 | tic = Log.time() |
| 103 | video_path = cfg.video_path |
| 104 | paths = cfg.paths |
| 105 | static_cam = cfg.static_cam |
| 106 | verbose = cfg.verbose |
| 107 | |
| 108 | # Get bbx tracking result |
| 109 | if not Path(paths.bbx).exists(): |
| 110 | tracker = Tracker() |
| 111 | bbx_xyxy = tracker.get_one_track(video_path).float() # (L, 4) |
| 112 | bbx_xys = get_bbx_xys_from_xyxy(bbx_xyxy, base_enlarge=1.2).float() # (L, 3) apply aspect ratio and enlarge |
| 113 | torch.save({"bbx_xyxy": bbx_xyxy, "bbx_xys": bbx_xys}, paths.bbx) |
| 114 | del tracker |
| 115 | else: |
| 116 | bbx_xys = torch.load(paths.bbx)["bbx_xys"] |
| 117 | Log.info(f"[Preprocess] bbx (xyxy, xys) from {paths.bbx}") |
| 118 | if verbose: |
| 119 | video = read_video_np(video_path) |
| 120 | bbx_xyxy = torch.load(paths.bbx)["bbx_xyxy"] |
| 121 | video_overlay = draw_bbx_xyxy_on_image_batch(bbx_xyxy, video) |
| 122 | save_video(video_overlay, cfg.paths.bbx_xyxy_video_overlay) |
| 123 | |
| 124 | # Get VitPose |
| 125 | if not Path(paths.vitpose).exists(): |
| 126 | vitpose_extractor = VitPoseExtractor() |
| 127 | vitpose = vitpose_extractor.extract(video_path, bbx_xys) |
| 128 | torch.save(vitpose, paths.vitpose) |
| 129 | del vitpose_extractor |
| 130 | else: |
| 131 | vitpose = torch.load(paths.vitpose) |
| 132 | Log.info(f"[Preprocess] vitpose from {paths.vitpose}") |
| 133 | if verbose: |
| 134 | video = read_video_np(video_path) |
| 135 | video_overlay = draw_coco17_skeleton_batch(video, vitpose, 0.5) |
| 136 | save_video(video_overlay, paths.vitpose_video_overlay) |
| 137 | |
| 138 | # Get vit features |
| 139 | if not Path(paths.vit_features).exists(): |
| 140 | extractor = Extractor() |
| 141 | vit_features = extractor.extract_video_features(video_path, bbx_xys) |
| 142 | torch.save(vit_features, paths.vit_features) |
| 143 | del extractor |
| 144 | else: |
| 145 | Log.info(f"[Preprocess] vit_features from {paths.vit_features}") |
| 146 | |
| 147 | # Get visual odometry results |
| 148 | if not static_cam: # use slam to get cam rotation |
| 149 | if not Path(paths.slam).exists(): |
| 150 | if not cfg.use_dpvo: |
| 151 | simple_vo = SimpleVO(cfg.video_path, scale=0.5, step=8, method="sift", f_mm=cfg.f_mm) |
| 152 | vo_results = simple_vo.compute() # (L, 4, 4), numpy |
| 153 | torch.save(vo_results, paths.slam) |
| 154 | else: # DPVO |
| 155 | from hmr4d.utils.preproc.slam import SLAMModel |
| 156 | |
| 157 | length, width, height = get_video_lwh(cfg.video_path) |
no test coverage detected