Initialize a inference state.
(self, frames, offload_video_to_cpu=False, offload_state_to_cpu=False, async_loading_frames=False, frame_names=None)
| 72 | |
| 73 | @torch.inference_mode() |
| 74 | def init_state_v2(self, frames, offload_video_to_cpu=False, offload_state_to_cpu=False, async_loading_frames=False, frame_names=None): |
| 75 | """Initialize a inference state.""" |
| 76 | images, video_height, video_width = load_video_frames_v2( |
| 77 | frames=frames, image_size=self.image_size, offload_video_to_cpu=offload_video_to_cpu, async_loading_frames=async_loading_frames, frame_names=frame_names |
| 78 | ) |
| 79 | inference_state = {} |
| 80 | inference_state["images"] = images |
| 81 | inference_state["num_frames"] = len(images) |
| 82 | # whether to offload the video frames to CPU memory |
| 83 | # turning on this option saves the GPU memory with only a very small overhead |
| 84 | inference_state["offload_video_to_cpu"] = offload_video_to_cpu |
| 85 | # whether to offload the inference state to CPU memory |
| 86 | # turning on this option saves the GPU memory at the cost of a lower tracking fps |
| 87 | # (e.g. in a test case of 768x768 model, fps dropped from 27 to 24 when tracking one object |
| 88 | # and from 24 to 21 when tracking two objects) |
| 89 | inference_state["offload_state_to_cpu"] = offload_state_to_cpu |
| 90 | # the original video height and width, used for resizing final output scores |
| 91 | inference_state["video_height"] = video_height |
| 92 | inference_state["video_width"] = video_width |
| 93 | inference_state["device"] = torch.device("cuda") |
| 94 | if offload_state_to_cpu: |
| 95 | inference_state["storage_device"] = torch.device("cpu") |
| 96 | else: |
| 97 | inference_state["storage_device"] = torch.device("cuda") |
| 98 | # inputs on each frame |
| 99 | inference_state["point_inputs_per_obj"] = {} |
| 100 | inference_state["mask_inputs_per_obj"] = {} |
| 101 | # visual features on a small number of recently visited frames for quick interactions |
| 102 | inference_state["cached_features"] = {} |
| 103 | # values that don't change across frames (so we only need to hold one copy of them) |
| 104 | inference_state["constants"] = {} |
| 105 | # mapping between client-side object id and model-side object index |
| 106 | inference_state["obj_id_to_idx"] = OrderedDict() |
| 107 | inference_state["obj_idx_to_id"] = OrderedDict() |
| 108 | inference_state["obj_ids"] = [] |
| 109 | # A storage to hold the model's tracking results and states on each frame |
| 110 | inference_state["output_dict"] = { |
| 111 | "cond_frame_outputs": {}, # dict containing {frame_idx: <out>} |
| 112 | "non_cond_frame_outputs": {}, # dict containing {frame_idx: <out>} |
| 113 | } |
| 114 | # Slice (view) of each object tracking results, sharing the same memory with "output_dict" |
| 115 | inference_state["output_dict_per_obj"] = {} |
| 116 | # A temporary storage to hold new outputs when user interact with a frame |
| 117 | # to add clicks or mask (it's merged into "output_dict" before propagation starts) |
| 118 | inference_state["temp_output_dict_per_obj"] = {} |
| 119 | # Frames that already holds consolidated outputs from click or mask inputs |
| 120 | # (we directly use their consolidated outputs during tracking) |
| 121 | inference_state["consolidated_frame_inds"] = { |
| 122 | "cond_frame_outputs": set(), # set containing frame indices |
| 123 | "non_cond_frame_outputs": set(), # set containing frame indices |
| 124 | } |
| 125 | # metadata for each tracking frame (e.g. which direction it's tracked) |
| 126 | inference_state["tracking_has_started"] = False |
| 127 | inference_state["frames_already_tracked"] = {} |
| 128 | inference_state["frames_tracked_per_obj"] = {} |
| 129 | # Warm up the visual backbone and cache the image feature on frame 0 |
| 130 | self._get_image_feature(inference_state, frame_idx=0, batch_size=1) |
| 131 | return inference_state |
no test coverage detected