ZiyuG/SAM2Point
16
1# Copyright (c) Meta Platforms, Inc. and affiliates.2# All rights reserved.3 4# This source code is licensed under the license found in the5# LICENSE file in the root directory of this source tree.6 7import warnings8from collections import OrderedDict9 10import torch11 12from tqdm import tqdm13 14from sam2.modeling.sam2_base import NO_OBJ_SCORE, SAM2Base15from sam2.utils.misc import concat_points, fill_holes_in_mask_scores, load_video_frames16 17 18class SAM2VideoPredictor(SAM2Base):19 """The predictor class to handle user interactions and manage inference states."""20 21 def __init__(22 self,23 fill_hole_area=0,24 # whether to apply non-overlapping constraints on the output object masks25 non_overlap_masks=False,26 # whether to clear non-conditioning memory of the surrounding frames (which may contain outdated information) after adding correction clicks;27 # note that this would only apply to *single-object tracking* unless `clear_non_cond_mem_for_multi_obj` is also set to True)28 clear_non_cond_mem_around_input=False,29 # whether to also clear non-conditioning memory of the surrounding frames (only effective when `clear_non_cond_mem_around_input` is True).30 clear_non_cond_mem_for_multi_obj=False,31 **kwargs,32 ):33 super().__init__(**kwargs)34 self.fill_hole_area = fill_hole_area35 self.non_overlap_masks = non_overlap_masks36 self.clear_non_cond_mem_around_input = clear_non_cond_mem_around_input37 self.clear_non_cond_mem_for_multi_obj = clear_non_cond_mem_for_multi_obj38 39 @torch.inference_mode()40 def init_state(41 self,42 frame_paths,43 offload_video_to_cpu=False,44 offload_state_to_cpu=False,45 async_loading_frames=False,46 ):47 """Initialize a inference state."""48 images, video_height, video_width = load_video_frames(49 img_paths=frame_paths,50 image_size=self.image_size,51 offload_video_to_cpu=offload_video_to_cpu,52 async_loading_frames=async_loading_frames,53 )54 inference_state = {}55 inference_state["images"] = images56 inference_state["num_frames"] = len(images)57 # whether to offload the video frames to CPU memory58 # turning on this option saves the GPU memory with only a very small overhead59 inference_state["offload_video_to_cpu"] = offload_video_to_cpu60 # whether to offload the inference state to CPU memory61 # turning on this option saves the GPU memory at the cost of a lower tracking fps62 # (e.g. in a test case of 768x768 model, fps dropped from 27 to 24 when tracking one object63 # and from 24 to 21 when tracking two objects)64 inference_state["offload_state_to_cpu"] = offload_state_to_cpu65 # the original video height and width, used for resizing final output scores66 inference_state["video_height"] = video_height67 inference_state["video_width"] = video_width68 inference_state["device"] = torch.device("cuda")69 if offload_state_to_cpu:70 inference_state["storage_device"] = torch.device("cpu")71 else:72 inference_state["storage_device"] = torch.device("cuda")73 # inputs on each frame74 inference_state["point_inputs_per_obj"] = {}75 inference_state["mask_inputs_per_obj"] = {}76 # visual features on a small number of recently visited frames for quick interactions77 inference_state["cached_features"] = {}78 # values that don't change across frames (so we only need to hold one copy of them)79 inference_state["constants"] = {}80 # mapping between client-side object id and model-side object index81 inference_state["obj_id_to_idx"] = OrderedDict()82 inference_state["obj_idx_to_id"] = OrderedDict()83 inference_state["obj_ids"] = []84 # A storage to hold the model's tracking results and states on each frame85 inference_state["output_dict"] = {86 "cond_frame_outputs": {}, # dict containing {frame_idx: <out>}87 "non_cond_frame_outputs": {}, # dict containing {frame_idx: <out>}88 }89 # Slice (view) of each object tracking results, sharing the same memory with "output_dict"90 inference_state["output_dict_per_obj"] = {}91 # A temporary storage to hold new outputs when user interact with a frame92 # to add clicks or mask (it's merged into "output_dict" before propagation starts)93 inference_state["temp_output_dict_per_obj"] = {}94 # Frames that already holds consolidated outputs from click or mask inputs95 # (we directly use their consolidated outputs during tracking)96 inference_state["consolidated_frame_inds"] = {97 "cond_frame_outputs": set(), # set containing frame indices98 "non_cond_frame_outputs": set(), # set containing frame indices99 }100 # metadata for each tracking frame (e.g. which direction it's tracked)101 inference_state["tracking_has_started"] = False102 inference_state["frames_already_tracked"] = {}103 # Warm up the visual backbone and cache the image feature on frame 0104 self._get_image_feature(inference_state, frame_idx=0, batch_size=1)105 return inference_state106 107 @classmethod108 def from_pretrained(cls, model_id: str, **kwargs) -> "SAM2VideoPredictor":109 """110 Load a pretrained model from the Hugging Face hub.111 112 Arguments:113 model_id (str): The Hugging Face repository ID.114 **kwargs: Additional arguments to pass to the model constructor.115 116 Returns:117 (SAM2VideoPredictor): The loaded model.118 """119 from sam2.build_sam import build_sam2_video_predictor_hf120 121 sam_model = build_sam2_video_predictor_hf(model_id, **kwargs)122 return cls(sam_model)123 124 def _obj_id_to_idx(self, inference_state, obj_id):125 """Map client-side object id to model-side object index."""126 obj_idx = inference_state["obj_id_to_idx"].get(obj_id, None)127 if obj_idx is not None:128 return obj_idx129 130 # This is a new object id not sent to the server before. We only allow adding131 # new objects *before* the tracking starts.132 allow_new_object = not inference_state["tracking_has_started"]133 if allow_new_object:134 # get the next object slot135 obj_idx = len(inference_state["obj_id_to_idx"])136 inference_state["obj_id_to_idx"][obj_id] = obj_idx137 inference_state["obj_idx_to_id"][obj_idx] = obj_id138 inference_state["obj_ids"] = list(inference_state["obj_id_to_idx"])139 # set up input and output structures for this object140 inference_state["point_inputs_per_obj"][obj_idx] = {}141 inference_state["mask_inputs_per_obj"][obj_idx] = {}142 inference_state["output_dict_per_obj"][obj_idx] = {143 "cond_frame_outputs": {}, # dict containing {frame_idx: <out>}144 "non_cond_frame_outputs": {}, # dict containing {frame_idx: <out>}145 }146 inference_state["temp_output_dict_per_obj"][obj_idx] = {147 "cond_frame_outputs": {}, # dict containing {frame_idx: <out>}148 "non_cond_frame_outputs": {}, # dict containing {frame_idx: <out>}149 }150 return obj_idx151 else:152 raise RuntimeError(153 f"Cannot add new object id {obj_id} after tracking starts. "154 f"All existing object ids: {inference_state['obj_ids']}. "155 f"Please call 'reset_state' to restart from scratch."156 )157 158 def _obj_idx_to_id(self, inference_state, obj_idx):159 """Map model-side object index to client-side object id."""160 return inference_state["obj_idx_to_id"][obj_idx]161 162 def _get_obj_num(self, inference_state):163 """Get the total number of unique object ids received so far in this session."""164 return len(inference_state["obj_idx_to_id"])165 166 @torch.inference_mode()167 def add_new_points_or_box(168 self,169 inference_state,170 frame_idx,171 obj_id,172 points=None,173 labels=None,174 clear_old_points=True,175 normalize_coords=True,176 box=None,177 ):178 """Add new points to a frame."""179 obj_idx = self._obj_id_to_idx(inference_state, obj_id)180 point_inputs_per_frame = inference_state["point_inputs_per_obj"][obj_idx]181 mask_inputs_per_frame = inference_state["mask_inputs_per_obj"][obj_idx]182 183 if (points is not None) != (labels is not None):184 raise ValueError("points and labels must be provided together")185 if points is None and box is None:186 raise ValueError("at least one of points or box must be provided as input")187 188 if points is None:189 points = torch.zeros(0, 2, dtype=torch.float32)190 elif not isinstance(points, torch.Tensor):191 points = torch.tensor(points, dtype=torch.float32)192 if labels is None:193 labels = torch.zeros(0, dtype=torch.int32)194 elif not isinstance(labels, torch.Tensor):195 labels = torch.tensor(labels, dtype=torch.int32)196 if points.dim() == 2:197 points = points.unsqueeze(0) # add batch dimension198 if labels.dim() == 1:199 labels = labels.unsqueeze(0) # add batch dimension200 201 # If `box` is provided, we add it as the first two points with labels 2 and 3202 # along with the user-provided points (consistent with how SAM 2 is trained).203 if box is not None:204 if not clear_old_points:205 raise ValueError(206 "cannot add box without clearing old points, since "207 "box prompt must be provided before any point prompt "208 "(please use clear_old_points=True instead)"209 )210 if inference_state["tracking_has_started"]:211 warnings.warn(212 "You are adding a box after tracking starts. SAM 2 may not always be "213 "able to incorporate a box prompt for *refinement*. If you intend to "214 "use box prompt as an *initial* input before tracking, please call "215 "'reset_state' on the inference state to restart from scratch.",216 category=UserWarning,217 stacklevel=2,218 )219 if not isinstance(box, torch.Tensor):220 box = torch.tensor(box, dtype=torch.float32, device=points.device)221 box_coords = box.reshape(1, 2, 2)222 box_labels = torch.tensor([2, 3], dtype=torch.int32, device=labels.device)223 box_labels = box_labels.reshape(1, 2)224 points = torch.cat([box_coords, points], dim=1)225 labels = torch.cat([box_labels, labels], dim=1)226 227 if normalize_coords:228 video_H = inference_state["video_height"]229 video_W = inference_state["video_width"]230 points = points / torch.tensor([video_W, video_H]).to(points.device)231 # scale the (normalized) coordinates by the model's internal image size232 points = points * self.image_size233 points = points.to(inference_state["device"])234 labels = labels.to(inference_state["device"])235 236 if not clear_old_points:237 point_inputs = point_inputs_per_frame.get(frame_idx, None)238 else:239 point_inputs = None240 point_inputs = concat_points(point_inputs, points, labels)241 242 point_inputs_per_frame[frame_idx] = point_inputs243 mask_inputs_per_frame.pop(frame_idx, None)244 # If this frame hasn't been tracked before, we treat it as an initial conditioning245 # frame, meaning that the inputs points are to generate segments on this frame without246 # using any memory from other frames, like in SAM. Otherwise (if it has been tracked),247 # the input points will be used to correct the already tracked masks.248 is_init_cond_frame = frame_idx not in inference_state["frames_already_tracked"]249 # whether to track in reverse time order250 if is_init_cond_frame:251 reverse = False252 else:253 reverse = inference_state["frames_already_tracked"][frame_idx]["reverse"]254 obj_output_dict = inference_state["output_dict_per_obj"][obj_idx]255 obj_temp_output_dict = inference_state["temp_output_dict_per_obj"][obj_idx]256 # Add a frame to conditioning output if it's an initial conditioning frame or257 # if the model sees all frames receiving clicks/mask as conditioning frames.258 is_cond = is_init_cond_frame or self.add_all_frames_to_correct_as_cond259 storage_key = "cond_frame_outputs" if is_cond else "non_cond_frame_outputs"260 261 # Get any previously predicted mask logits on this object and feed it along with262 # the new clicks into the SAM mask decoder.263 prev_sam_mask_logits = None264 # lookup temporary output dict first, which contains the most recent output265 # (if not found, then lookup conditioning and non-conditioning frame output)266 prev_out = obj_temp_output_dict[storage_key].get(frame_idx)267 if prev_out is None:268 prev_out = obj_output_dict["cond_frame_outputs"].get(frame_idx)269 if prev_out is None:270 prev_out = obj_output_dict["non_cond_frame_outputs"].get(frame_idx)271 272 if prev_out is not None and prev_out["pred_masks"] is not None:273 prev_sam_mask_logits = prev_out["pred_masks"].cuda(non_blocking=True)274 # Clamp the scale of prev_sam_mask_logits to avoid rare numerical issues.275 prev_sam_mask_logits = torch.clamp(prev_sam_mask_logits, -32.0, 32.0)276 current_out, _ = self._run_single_frame_inference(277 inference_state=inference_state,278 output_dict=obj_output_dict, # run on the slice of a single object279 frame_idx=frame_idx,280 batch_size=1, # run on the slice of a single object281 is_init_cond_frame=is_init_cond_frame,282 point_inputs=point_inputs,283 mask_inputs=None,284 reverse=reverse,285 # Skip the memory encoder when adding clicks or mask. We execute the memory encoder286 # at the beginning of `propagate_in_video` (after user finalize their clicks). This287 # allows us to enforce non-overlapping constraints on all objects before encoding288 # them into memory.289 run_mem_encoder=False,290 prev_sam_mask_logits=prev_sam_mask_logits,291 )292 # Add the output to the output dict (to be used as future memory)293 obj_temp_output_dict[storage_key][frame_idx] = current_out294 295 # Resize the output mask to the original video resolution296 obj_ids = inference_state["obj_ids"]297 consolidated_out = self._consolidate_temp_output_across_obj(298 inference_state,299 frame_idx,300 is_cond=is_cond,301 run_mem_encoder=False,302 consolidate_at_video_res=True,303 )304 _, video_res_masks = self._get_orig_video_res_output(305 inference_state, consolidated_out["pred_masks_video_res"]306 )307 return frame_idx, obj_ids, video_res_masks308 309 def add_new_points(self, *args, **kwargs):310 """Deprecated method. Please use `add_new_points_or_box` instead."""311 return self.add_new_points_or_box(*args, **kwargs)312 313 @torch.inference_mode()314 def add_new_mask(315 self,316 inference_state,317 frame_idx,318 obj_id,319 mask,320 ):321 """Add new mask to a frame."""322 obj_idx = self._obj_id_to_idx(inference_state, obj_id)323 point_inputs_per_frame = inference_state["point_inputs_per_obj"][obj_idx]324 mask_inputs_per_frame = inference_state["mask_inputs_per_obj"][obj_idx]325 326 if not isinstance(mask, torch.Tensor):327 mask = torch.tensor(mask, dtype=torch.bool)328 assert mask.dim() == 2329 mask_H, mask_W = mask.shape330 mask_inputs_orig = mask[None, None] # add batch and channel dimension331 mask_inputs_orig = mask_inputs_orig.float().to(inference_state["device"])332 333 # resize the mask if it doesn't match the model's image size334 if mask_H != self.image_size or mask_W != self.image_size:335 mask_inputs = torch.nn.functional.interpolate(336 mask_inputs_orig,337 size=(self.image_size, self.image_size),338 align_corners=False,339 mode="bilinear",340 antialias=True, # use antialias for downsampling341 )342 mask_inputs = (mask_inputs >= 0.5).float()343 else:344 mask_inputs = mask_inputs_orig345 346 mask_inputs_per_frame[frame_idx] = mask_inputs347 point_inputs_per_frame.pop(frame_idx, None)348 # If this frame hasn't been tracked before, we treat it as an initial conditioning349 # frame, meaning that the inputs points are to generate segments on this frame without350 # using any memory from other frames, like in SAM. Otherwise (if it has been tracked),351 # the input points will be used to correct the already tracked masks.352 is_init_cond_frame = frame_idx not in inference_state["frames_already_tracked"]353 # whether to track in reverse time order354 if is_init_cond_frame:355 reverse = False356 else:357 reverse = inference_state["frames_already_tracked"][frame_idx]["reverse"]358 obj_output_dict = inference_state["output_dict_per_obj"][obj_idx]359 obj_temp_output_dict = inference_state["temp_output_dict_per_obj"][obj_idx]360 # Add a frame to conditioning output if it's an initial conditioning frame or361 # if the model sees all frames receiving clicks/mask as conditioning frames.362 is_cond = is_init_cond_frame or self.add_all_frames_to_correct_as_cond363 storage_key = "cond_frame_outputs" if is_cond else "non_cond_frame_outputs"364 365 current_out, _ = self._run_single_frame_inference(366 inference_state=inference_state,367 output_dict=obj_output_dict, # run on the slice of a single object368 frame_idx=frame_idx,369 batch_size=1, # run on the slice of a single object370 is_init_cond_frame=is_init_cond_frame,371 point_inputs=None,372 mask_inputs=mask_inputs,373 reverse=reverse,374 # Skip the memory encoder when adding clicks or mask. We execute the memory encoder375 # at the beginning of `propagate_in_video` (after user finalize their clicks). This376 # allows us to enforce non-overlapping constraints on all objects before encoding377 # them into memory.378 run_mem_encoder=False,379 )380 # Add the output to the output dict (to be used as future memory)381 obj_temp_output_dict[storage_key][frame_idx] = current_out382 383 # Resize the output mask to the original video resolution384 obj_ids = inference_state["obj_ids"]385 consolidated_out = self._consolidate_temp_output_across_obj(386 inference_state,387 frame_idx,388 is_cond=is_cond,389 run_mem_encoder=False,390 consolidate_at_video_res=True,391 )392 _, video_res_masks = self._get_orig_video_res_output(393 inference_state, consolidated_out["pred_masks_video_res"]394 )395 return frame_idx, obj_ids, video_res_masks396 397 def _get_orig_video_res_output(self, inference_state, any_res_masks):398 """399 Resize the object scores to the original video resolution (video_res_masks)400 and apply non-overlapping constraints for final output.401 """402 device = inference_state["device"]403 video_H = inference_state["video_height"]404 video_W = inference_state["video_width"]405 any_res_masks = any_res_masks.to(device, non_blocking=True)406 if any_res_masks.shape[-2:] == (video_H, video_W):407 video_res_masks = any_res_masks408 else:409 video_res_masks = torch.nn.functional.interpolate(410 any_res_masks,411 size=(video_H, video_W),412 mode="bilinear",413 align_corners=False,414 )415 if self.non_overlap_masks:416 video_res_masks = self._apply_non_overlapping_constraints(video_res_masks)417 return any_res_masks, video_res_masks418 419 def _consolidate_temp_output_across_obj(420 self,421 inference_state,422 frame_idx,423 is_cond,424 run_mem_encoder,425 consolidate_at_video_res=False,426 ):427 """428 Consolidate the per-object temporary outputs in `temp_output_dict_per_obj` on429 a frame into a single output for all objects, including430 1) fill any missing objects either from `output_dict_per_obj` (if they exist in431 `output_dict_per_obj` for this frame) or leave them as placeholder values432 (if they don't exist in `output_dict_per_obj` for this frame);433 2) if specified, rerun memory encoder after apply non-overlapping constraints434 on the object scores.435 """436 batch_size = self._get_obj_num(inference_state)437 storage_key = "cond_frame_outputs" if is_cond else "non_cond_frame_outputs"438 # Optionally, we allow consolidating the temporary outputs at the original439 # video resolution (to provide a better editing experience for mask prompts).440 if consolidate_at_video_res:441 assert not run_mem_encoder, "memory encoder cannot run at video resolution"442 consolidated_H = inference_state["video_height"]443 consolidated_W = inference_state["video_width"]444 consolidated_mask_key = "pred_masks_video_res"445 else:446 consolidated_H = consolidated_W = self.image_size // 4447 consolidated_mask_key = "pred_masks"448 449 # Initialize `consolidated_out`. Its "maskmem_features" and "maskmem_pos_enc"450 # will be added when rerunning the memory encoder after applying non-overlapping451 # constraints to object scores. Its "pred_masks" are prefilled with a large452 # negative value (NO_OBJ_SCORE) to represent missing objects.453 consolidated_out = {454 "maskmem_features": None,455 "maskmem_pos_enc": None,456 consolidated_mask_key: torch.full(457 size=(batch_size, 1, consolidated_H, consolidated_W),458 fill_value=NO_OBJ_SCORE,459 dtype=torch.float32,460 device=inference_state["storage_device"],461 ),462 "obj_ptr": torch.full(463 size=(batch_size, self.hidden_dim),464 fill_value=NO_OBJ_SCORE,465 dtype=torch.float32,466 device=inference_state["device"],467 ),468 }469 empty_mask_ptr = None470 for obj_idx in range(batch_size):471 obj_temp_output_dict = inference_state["temp_output_dict_per_obj"][obj_idx]472 obj_output_dict = inference_state["output_dict_per_obj"][obj_idx]473 out = obj_temp_output_dict[storage_key].get(frame_idx, None)474 # If the object doesn't appear in "temp_output_dict_per_obj" on this frame,475 # we fall back and look up its previous output in "output_dict_per_obj".476 # We look up both "cond_frame_outputs" and "non_cond_frame_outputs" in477 # "output_dict_per_obj" to find a previous output for this object.478 if out is None:479 out = obj_output_dict["cond_frame_outputs"].get(frame_idx, None)480 if out is None:481 out = obj_output_dict["non_cond_frame_outputs"].get(frame_idx, None)482 # If the object doesn't appear in "output_dict_per_obj" either, we skip it483 # and leave its mask scores to the default scores (i.e. the NO_OBJ_SCORE484 # placeholder above) and set its object pointer to be a dummy pointer.485 if out is None:486 # Fill in dummy object pointers for those objects without any inputs or487 # tracking outcomes on this frame (only do it under `run_mem_encoder=True`,488 # i.e. when we need to build the memory for tracking).489 if run_mem_encoder:490 if empty_mask_ptr is None:491 empty_mask_ptr = self._get_empty_mask_ptr(492 inference_state, frame_idx493 )494 # fill object pointer with a dummy pointer (based on an empty mask)495 consolidated_out["obj_ptr"][obj_idx : obj_idx + 1] = empty_mask_ptr496 continue497 # Add the temporary object output mask to consolidated output mask498 obj_mask = out["pred_masks"]499 consolidated_pred_masks = consolidated_out[consolidated_mask_key]500 if obj_mask.shape[-2:] == consolidated_pred_masks.shape[-2:]:501 consolidated_pred_masks[obj_idx : obj_idx + 1] = obj_mask502 else:503 # Resize first if temporary object mask has a different resolution504 resized_obj_mask = torch.nn.functional.interpolate(505 obj_mask,506 size=consolidated_pred_masks.shape[-2:],507 mode="bilinear",508 align_corners=False,509 )510 consolidated_pred_masks[obj_idx : obj_idx + 1] = resized_obj_mask511 consolidated_out["obj_ptr"][obj_idx : obj_idx + 1] = out["obj_ptr"]512 513 # Optionally, apply non-overlapping constraints on the consolidated scores514 # and rerun the memory encoder515 if run_mem_encoder:516 device = inference_state["device"]517 high_res_masks = torch.nn.functional.interpolate(518 consolidated_out["pred_masks"].to(device, non_blocking=True),519 size=(self.image_size, self.image_size),520 mode="bilinear",521 align_corners=False,522 )523 if self.non_overlap_masks_for_mem_enc:524 high_res_masks = self._apply_non_overlapping_constraints(high_res_masks)525 maskmem_features, maskmem_pos_enc = self._run_memory_encoder(526 inference_state=inference_state,527 frame_idx=frame_idx,528 batch_size=batch_size,529 high_res_masks=high_res_masks,530 is_mask_from_pts=True, # these frames are what the user interacted with531 )532 consolidated_out["maskmem_features"] = maskmem_features533 consolidated_out["maskmem_pos_enc"] = maskmem_pos_enc534 535 return consolidated_out536 537 def _get_empty_mask_ptr(self, inference_state, frame_idx):538 """Get a dummy object pointer based on an empty mask on the current frame."""539 # A dummy (empty) mask with a single object540 batch_size = 1541 mask_inputs = torch.zeros(542 (batch_size, 1, self.image_size, self.image_size),543 dtype=torch.float32,544 device=inference_state["device"],545 )546 547 # Retrieve correct image features548 (549 _,550 _,551 current_vision_feats,552 current_vision_pos_embeds,553 feat_sizes,554 ) = self._get_image_feature(inference_state, frame_idx, batch_size)555 556 # Feed the empty mask and image feature above to get a dummy object pointer557 current_out = self.track_step(558 frame_idx=frame_idx,559 is_init_cond_frame=True,560 current_vision_feats=current_vision_feats,561 current_vision_pos_embeds=current_vision_pos_embeds,562 feat_sizes=feat_sizes,563 point_inputs=None,564 mask_inputs=mask_inputs,565 output_dict={},566 num_frames=inference_state["num_frames"],567 track_in_reverse=False,568 run_mem_encoder=False,569 prev_sam_mask_logits=None,570 )571 return current_out["obj_ptr"]572 573 @torch.inference_mode()574 def propagate_in_video_preflight(self, inference_state):575 """Prepare inference_state and consolidate temporary outputs before tracking."""576 # Tracking has started and we don't allow adding new objects until session is reset.577 inference_state["tracking_has_started"] = True578 batch_size = self._get_obj_num(inference_state)579 580 # Consolidate per-object temporary outputs in "temp_output_dict_per_obj" and581 # add them into "output_dict".582 temp_output_dict_per_obj = inference_state["temp_output_dict_per_obj"]583 output_dict = inference_state["output_dict"]584 # "consolidated_frame_inds" contains indices of those frames where consolidated585 # temporary outputs have been added (either in this call or any previous calls586 # to `propagate_in_video_preflight`).587 consolidated_frame_inds = inference_state["consolidated_frame_inds"]588 for is_cond in [False, True]:589 # Separately consolidate conditioning and non-conditioning temp outptus590 storage_key = "cond_frame_outputs" if is_cond else "non_cond_frame_outputs"591 # Find all the frames that contain temporary outputs for any objects592 # (these should be the frames that have just received clicks for mask inputs593 # via `add_new_points_or_box` or `add_new_mask`)594 temp_frame_inds = set()595 for obj_temp_output_dict in temp_output_dict_per_obj.values():596 temp_frame_inds.update(obj_temp_output_dict[storage_key].keys())597 consolidated_frame_inds[storage_key].update(temp_frame_inds)598 # consolidate the temprary output across all objects on this frame599 for frame_idx in temp_frame_inds:600 consolidated_out = self._consolidate_temp_output_across_obj(601 inference_state, frame_idx, is_cond=is_cond, run_mem_encoder=True602 )603 # merge them into "output_dict" and also create per-object slices604 output_dict[storage_key][frame_idx] = consolidated_out605 self._add_output_per_object(606 inference_state, frame_idx, consolidated_out, storage_key607 )608 clear_non_cond_mem = self.clear_non_cond_mem_around_input and (609 self.clear_non_cond_mem_for_multi_obj or batch_size <= 1610 )611 if clear_non_cond_mem:612 # clear non-conditioning memory of the surrounding frames613 self._clear_non_cond_mem_around_input(inference_state, frame_idx)614 615 # clear temporary outputs in `temp_output_dict_per_obj`616 for obj_temp_output_dict in temp_output_dict_per_obj.values():617 obj_temp_output_dict[storage_key].clear()618 619 # edge case: if an output is added to "cond_frame_outputs", we remove any prior620 # output on the same frame in "non_cond_frame_outputs"621 for frame_idx in output_dict["cond_frame_outputs"]:622 output_dict["non_cond_frame_outputs"].pop(frame_idx, None)623 for obj_output_dict in inference_state["output_dict_per_obj"].values():624 for frame_idx in obj_output_dict["cond_frame_outputs"]:625 obj_output_dict["non_cond_frame_outputs"].pop(frame_idx, None)626 for frame_idx in consolidated_frame_inds["cond_frame_outputs"]:627 assert frame_idx in output_dict["cond_frame_outputs"]628 consolidated_frame_inds["non_cond_frame_outputs"].discard(frame_idx)629 630 # Make sure that the frame indices in "consolidated_frame_inds" are exactly those frames631 # with either points or mask inputs (which should be true under a correct workflow).632 all_consolidated_frame_inds = (633 consolidated_frame_inds["cond_frame_outputs"]634 | consolidated_frame_inds["non_cond_frame_outputs"]635 )636 input_frames_inds = set()637 for point_inputs_per_frame in inference_state["point_inputs_per_obj"].values():638 input_frames_inds.update(point_inputs_per_frame.keys())639 for mask_inputs_per_frame in inference_state["mask_inputs_per_obj"].values():640 input_frames_inds.update(mask_inputs_per_frame.keys())641 assert all_consolidated_frame_inds == input_frames_inds642 643 @torch.inference_mode()644 def propagate_in_video(645 self,646 inference_state,647 start_frame_idx=None,648 max_frame_num_to_track=None,649 reverse=False,650 ):651 """Propagate the input points across frames to track in the entire video."""652 self.propagate_in_video_preflight(inference_state)653 654 output_dict = inference_state["output_dict"]655 consolidated_frame_inds = inference_state["consolidated_frame_inds"]656 obj_ids = inference_state["obj_ids"]657 num_frames = inference_state["num_frames"]658 batch_size = self._get_obj_num(inference_state)659 if len(output_dict["cond_frame_outputs"]) == 0:660 raise RuntimeError("No points are provided; please add points first")661 clear_non_cond_mem = self.clear_non_cond_mem_around_input and (662 self.clear_non_cond_mem_for_multi_obj or batch_size <= 1663 )664 665 # set start index, end index, and processing order666 if start_frame_idx is None:667 # default: start from the earliest frame with input points668 start_frame_idx = min(output_dict["cond_frame_outputs"])669 if max_frame_num_to_track is None:670 # default: track all the frames in the video671 max_frame_num_to_track = num_frames672 if reverse:673 end_frame_idx = max(start_frame_idx - max_frame_num_to_track, 0)674 if start_frame_idx > 0:675 processing_order = range(start_frame_idx, end_frame_idx - 1, -1)676 else:677 processing_order = [] # skip reverse tracking if starting from frame 0678 else:679 end_frame_idx = min(680 start_frame_idx + max_frame_num_to_track, num_frames - 1681 )682 processing_order = range(start_frame_idx, end_frame_idx + 1)683 684 for frame_idx in tqdm(processing_order, desc="propagate in video"):685 # We skip those frames already in consolidated outputs (these are frames686 # that received input clicks or mask). Note that we cannot directly run687 # batched forward on them via `_run_single_frame_inference` because the688 # number of clicks on each object might be different.689 if frame_idx in consolidated_frame_inds["cond_frame_outputs"]:690 storage_key = "cond_frame_outputs"691 current_out = output_dict[storage_key][frame_idx]692 pred_masks = current_out["pred_masks"]693 if clear_non_cond_mem:694 # clear non-conditioning memory of the surrounding frames695 self._clear_non_cond_mem_around_input(inference_state, frame_idx)696 elif frame_idx in consolidated_frame_inds["non_cond_frame_outputs"]:697 storage_key = "non_cond_frame_outputs"698 current_out = output_dict[storage_key][frame_idx]699 pred_masks = current_out["pred_masks"]700 else:701 storage_key = "non_cond_frame_outputs"702 current_out, pred_masks = self._run_single_frame_inference(703 inference_state=inference_state,704 output_dict=output_dict,705 frame_idx=frame_idx,706 batch_size=batch_size,707 is_init_cond_frame=False,708 point_inputs=None,709 mask_inputs=None,710 reverse=reverse,711 run_mem_encoder=True,712 )713 output_dict[storage_key][frame_idx] = current_out714 # Create slices of per-object outputs for subsequent interaction with each715 # individual object after tracking.716 self._add_output_per_object(717 inference_state, frame_idx, current_out, storage_key718 )719 inference_state["frames_already_tracked"][frame_idx] = {"reverse": reverse}720 721 # Resize the output mask to the original video resolution (we directly use722 # the mask scores on GPU for output to avoid any CPU conversion in between)723 _, video_res_masks = self._get_orig_video_res_output(724 inference_state, pred_masks725 )726 yield frame_idx, obj_ids, video_res_masks727 728 def _add_output_per_object(729 self, inference_state, frame_idx, current_out, storage_key730 ):731 """732 Split a multi-object output into per-object output slices and add them into733 `output_dict_per_obj`. The resulting slices share the same tensor storage.734 """735 maskmem_features = current_out["maskmem_features"]736 assert maskmem_features is None or isinstance(maskmem_features, torch.Tensor)737 738 maskmem_pos_enc = current_out["maskmem_pos_enc"]739 assert maskmem_pos_enc is None or isinstance(maskmem_pos_enc, list)740 741 output_dict_per_obj = inference_state["output_dict_per_obj"]742 for obj_idx, obj_output_dict in output_dict_per_obj.items():743 obj_slice = slice(obj_idx, obj_idx + 1)744 obj_out = {745 "maskmem_features": None,746 "maskmem_pos_enc": None,747 "pred_masks": current_out["pred_masks"][obj_slice],748 "obj_ptr": current_out["obj_ptr"][obj_slice],749 }750 if maskmem_features is not None:751 obj_out["maskmem_features"] = maskmem_features[obj_slice]752 if maskmem_pos_enc is not None:753 obj_out["maskmem_pos_enc"] = [x[obj_slice] for x in maskmem_pos_enc]754 obj_output_dict[storage_key][frame_idx] = obj_out755 756 @torch.inference_mode()757 def reset_state(self, inference_state):758 """Remove all input points or mask in all frames throughout the video."""759 self._reset_tracking_results(inference_state)760 # Remove all object ids761 inference_state["obj_id_to_idx"].clear()762 inference_state["obj_idx_to_id"].clear()763 inference_state["obj_ids"].clear()764 inference_state["point_inputs_per_obj"].clear()765 inference_state["mask_inputs_per_obj"].clear()766 inference_state["output_dict_per_obj"].clear()767 inference_state["temp_output_dict_per_obj"].clear()768 769 def _reset_tracking_results(self, inference_state):770 """Reset all tracking inputs and results across the videos."""771 for v in inference_state["point_inputs_per_obj"].values():772 v.clear()773 for v in inference_state["mask_inputs_per_obj"].values():774 v.clear()775 for v in inference_state["output_dict_per_obj"].values():776 v["cond_frame_outputs"].clear()777 v["non_cond_frame_outputs"].clear()778 for v in inference_state["temp_output_dict_per_obj"].values():779 v["cond_frame_outputs"].clear()780 v["non_cond_frame_outputs"].clear()781 inference_state["output_dict"]["cond_frame_outputs"].clear()782 inference_state["output_dict"]["non_cond_frame_outputs"].clear()783 inference_state["consolidated_frame_inds"]["cond_frame_outputs"].clear()784 inference_state["consolidated_frame_inds"]["non_cond_frame_outputs"].clear()785 inference_state["tracking_has_started"] = False786 inference_state["frames_already_tracked"].clear()787 788 def _get_image_feature(self, inference_state, frame_idx, batch_size):789 """Compute the image features on a given frame."""790 # Look up in the cache first791 image, backbone_out = inference_state["cached_features"].get(792 frame_idx, (None, None)793 )794 if backbone_out is None:795 # Cache miss -- we will run inference on a single image796 image = inference_state["images"][frame_idx].cuda().float().unsqueeze(0)797 backbone_out = self.forward_image(image)798 # Cache the most recent frame's feature (for repeated interactions with799 # a frame; we can use an LRU cache for more frames in the future).800 inference_state["cached_features"] = {frame_idx: (image, backbone_out)}801 802 # expand the features to have the same dimension as the number of objects803 expanded_image = image.expand(batch_size, -1, -1, -1)804 expanded_backbone_out = {805 "backbone_fpn": backbone_out["backbone_fpn"].copy(),806 "vision_pos_enc": backbone_out["vision_pos_enc"].copy(),807 }808 for i, feat in enumerate(expanded_backbone_out["backbone_fpn"]):809 expanded_backbone_out["backbone_fpn"][i] = feat.expand(810 batch_size, -1, -1, -1811 )812 for i, pos in enumerate(expanded_backbone_out["vision_pos_enc"]):813 pos = pos.expand(batch_size, -1, -1, -1)814 expanded_backbone_out["vision_pos_enc"][i] = pos815 816 features = self._prepare_backbone_features(expanded_backbone_out)817 features = (expanded_image,) + features818 return features819 820 def _run_single_frame_inference(821 self,822 inference_state,823 output_dict,824 frame_idx,825 batch_size,826 is_init_cond_frame,827 point_inputs,828 mask_inputs,829 reverse,830 run_mem_encoder,831 prev_sam_mask_logits=None,832 ):833 """Run tracking on a single frame based on current inputs and previous memory."""834 # Retrieve correct image features835 (836 _,837 _,838 current_vision_feats,839 current_vision_pos_embeds,840 feat_sizes,841 ) = self._get_image_feature(inference_state, frame_idx, batch_size)842 843 # point and mask should not appear as input simultaneously on the same frame844 assert point_inputs is None or mask_inputs is None845 current_out = self.track_step(846 frame_idx=frame_idx,847 is_init_cond_frame=is_init_cond_frame,848 current_vision_feats=current_vision_feats,849 current_vision_pos_embeds=current_vision_pos_embeds,850 feat_sizes=feat_sizes,851 point_inputs=point_inputs,852 mask_inputs=mask_inputs,853 output_dict=output_dict,854 num_frames=inference_state["num_frames"],855 track_in_reverse=reverse,856 run_mem_encoder=run_mem_encoder,857 prev_sam_mask_logits=prev_sam_mask_logits,858 )859 860 # optionally offload the output to CPU memory to save GPU space861 storage_device = inference_state["storage_device"]862 maskmem_features = current_out["maskmem_features"]863 if maskmem_features is not None:864 maskmem_features = maskmem_features.to(torch.bfloat16)865 maskmem_features = maskmem_features.to(storage_device, non_blocking=True)866 pred_masks_gpu = current_out["pred_masks"]867 # potentially fill holes in the predicted masks868 if self.fill_hole_area > 0:869 pred_masks_gpu = fill_holes_in_mask_scores(870 pred_masks_gpu, self.fill_hole_area871 )872 pred_masks = pred_masks_gpu.to(storage_device, non_blocking=True)873 # "maskmem_pos_enc" is the same across frames, so we only need to store one copy of it874 maskmem_pos_enc = self._get_maskmem_pos_enc(inference_state, current_out)875 # object pointer is a small tensor, so we always keep it on GPU memory for fast access876 obj_ptr = current_out["obj_ptr"]877 # make a compact version of this frame's output to reduce the state size878 compact_current_out = {879 "maskmem_features": maskmem_features,880 "maskmem_pos_enc": maskmem_pos_enc,881 "pred_masks": pred_masks,882 "obj_ptr": obj_ptr,883 }884 return compact_current_out, pred_masks_gpu885 886 def _run_memory_encoder(887 self, inference_state, frame_idx, batch_size, high_res_masks, is_mask_from_pts888 ):889 """890 Run the memory encoder on `high_res_masks`. This is usually after applying891 non-overlapping constraints to object scores. Since their scores changed, their892 memory also need to be computed again with the memory encoder.893 """894 # Retrieve correct image features895 _, _, current_vision_feats, _, feat_sizes = self._get_image_feature(896 inference_state, frame_idx, batch_size897 )898 maskmem_features, maskmem_pos_enc = self._encode_new_memory(899 current_vision_feats=current_vision_feats,900 feat_sizes=feat_sizes,901 pred_masks_high_res=high_res_masks,902 is_mask_from_pts=is_mask_from_pts,903 )904 905 # optionally offload the output to CPU memory to save GPU space906 storage_device = inference_state["storage_device"]907 maskmem_features = maskmem_features.to(torch.bfloat16)908 maskmem_features = maskmem_features.to(storage_device, non_blocking=True)909 # "maskmem_pos_enc" is the same across frames, so we only need to store one copy of it910 maskmem_pos_enc = self._get_maskmem_pos_enc(911 inference_state, {"maskmem_pos_enc": maskmem_pos_enc}912 )913 return maskmem_features, maskmem_pos_enc914 915 def _get_maskmem_pos_enc(self, inference_state, current_out):916 """917 `maskmem_pos_enc` is the same across frames and objects, so we cache it as918 a constant in the inference session to reduce session storage size.919 """920 model_constants = inference_state["constants"]921 # "out_maskmem_pos_enc" should be either a list of tensors or None922 out_maskmem_pos_enc = current_out["maskmem_pos_enc"]923 if out_maskmem_pos_enc is not None:924 if "maskmem_pos_enc" not in model_constants:925 assert isinstance(out_maskmem_pos_enc, list)926 # only take the slice for one object, since it's same across objects927 maskmem_pos_enc = [x[0:1].clone() for x in out_maskmem_pos_enc]928 model_constants["maskmem_pos_enc"] = maskmem_pos_enc929 else:930 maskmem_pos_enc = model_constants["maskmem_pos_enc"]931 # expand the cached maskmem_pos_enc to the actual batch size932 batch_size = out_maskmem_pos_enc[0].size(0)933 expanded_maskmem_pos_enc = [934 x.expand(batch_size, -1, -1, -1) for x in maskmem_pos_enc935 ]936 else:937 expanded_maskmem_pos_enc = None938 return expanded_maskmem_pos_enc939 940 def _clear_non_cond_mem_around_input(self, inference_state, frame_idx):941 """942 Remove the non-conditioning memory around the input frame. When users provide943 correction clicks, the surrounding frames' non-conditioning memories can still944 contain outdated object appearance information and could confuse the model.945 946 This method clears those non-conditioning memories surrounding the interacted947 frame to avoid giving the model both old and new information about the object.948 """949 r = self.memory_temporal_stride_for_eval950 frame_idx_begin = frame_idx - r * self.num_maskmem951 frame_idx_end = frame_idx + r * self.num_maskmem952 output_dict = inference_state["output_dict"]953 non_cond_frame_outputs = output_dict["non_cond_frame_outputs"]954 for t in range(frame_idx_begin, frame_idx_end + 1):955 non_cond_frame_outputs.pop(t, None)956 for obj_output_dict in inference_state["output_dict_per_obj"].values():957 obj_output_dict["non_cond_frame_outputs"].pop(t, None)958