Initialize a inference state.
(self, video_path, offload_video_to_cpu=False, offload_state_to_cpu=False, async_loading_frames=False, frame_names=None)
| 13 | |
| 14 | @torch.inference_mode() |
| 15 | def init_state(self, video_path, offload_video_to_cpu=False, offload_state_to_cpu=False, async_loading_frames=False, frame_names=None): |
| 16 | """Initialize a inference state.""" |
| 17 | images, video_height, video_width = load_video_frames( |
| 18 | video_path=video_path, image_size=self.image_size, offload_video_to_cpu=offload_video_to_cpu, async_loading_frames=async_loading_frames, frame_names=frame_names |
| 19 | ) |
| 20 | inference_state = {} |
| 21 | inference_state["images"] = images |
| 22 | inference_state["num_frames"] = len(images) |
| 23 | # whether to offload the video frames to CPU memory |
| 24 | # turning on this option saves the GPU memory with only a very small overhead |
| 25 | inference_state["offload_video_to_cpu"] = offload_video_to_cpu |
| 26 | # whether to offload the inference state to CPU memory |
| 27 | # turning on this option saves the GPU memory at the cost of a lower tracking fps |
| 28 | # (e.g. in a test case of 768x768 model, fps dropped from 27 to 24 when tracking one object |
| 29 | # and from 24 to 21 when tracking two objects) |
| 30 | inference_state["offload_state_to_cpu"] = offload_state_to_cpu |
| 31 | # the original video height and width, used for resizing final output scores |
| 32 | inference_state["video_height"] = video_height |
| 33 | inference_state["video_width"] = video_width |
| 34 | inference_state["device"] = torch.device("cuda") |
| 35 | if offload_state_to_cpu: |
| 36 | inference_state["storage_device"] = torch.device("cpu") |
| 37 | else: |
| 38 | inference_state["storage_device"] = torch.device("cuda") |
| 39 | # inputs on each frame |
| 40 | inference_state["point_inputs_per_obj"] = {} |
| 41 | inference_state["mask_inputs_per_obj"] = {} |
| 42 | # visual features on a small number of recently visited frames for quick interactions |
| 43 | inference_state["cached_features"] = {} |
| 44 | # values that don't change across frames (so we only need to hold one copy of them) |
| 45 | inference_state["constants"] = {} |
| 46 | # mapping between client-side object id and model-side object index |
| 47 | inference_state["obj_id_to_idx"] = OrderedDict() |
| 48 | inference_state["obj_idx_to_id"] = OrderedDict() |
| 49 | inference_state["obj_ids"] = [] |
| 50 | # A storage to hold the model's tracking results and states on each frame |
| 51 | inference_state["output_dict"] = { |
| 52 | "cond_frame_outputs": {}, # dict containing {frame_idx: <out>} |
| 53 | "non_cond_frame_outputs": {}, # dict containing {frame_idx: <out>} |
| 54 | } |
| 55 | # Slice (view) of each object tracking results, sharing the same memory with "output_dict" |
| 56 | inference_state["output_dict_per_obj"] = {} |
| 57 | # A temporary storage to hold new outputs when user interact with a frame |
| 58 | # to add clicks or mask (it's merged into "output_dict" before propagation starts) |
| 59 | inference_state["temp_output_dict_per_obj"] = {} |
| 60 | # Frames that already holds consolidated outputs from click or mask inputs |
| 61 | # (we directly use their consolidated outputs during tracking) |
| 62 | inference_state["consolidated_frame_inds"] = { |
| 63 | "cond_frame_outputs": set(), # set containing frame indices |
| 64 | "non_cond_frame_outputs": set(), # set containing frame indices |
| 65 | } |
| 66 | # metadata for each tracking frame (e.g. which direction it's tracked) |
| 67 | inference_state["tracking_has_started"] = False |
| 68 | inference_state["frames_already_tracked"] = {} |
| 69 | # Warm up the visual backbone and cache the image feature on frame 0 |
| 70 | self._get_image_feature(inference_state, frame_idx=0, batch_size=1) |
| 71 | return inference_state |
| 72 |
nothing calls this directly
no test coverage detected