diff --git a/sam3/model/sam3_video_base.py b/sam3/model/sam3_video_base.py index e3b40ca0a..e2777a588 100644 --- a/sam3/model/sam3_video_base.py +++ b/sam3/model/sam3_video_base.py @@ -722,7 +722,7 @@ def _recondition_masklets( for idx in reconditioned_states_idx: self.tracker.propagate_in_video_preflight( - tracker_states_local[idx], run_mem_encoder=True + tracker_states_local[idx], run_mem_encoder=False ) return tracker_states_local @@ -947,6 +947,21 @@ def run_tracker_update_planning_phase( tracker_metadata_prev, tracker_obj_scores_global, ) + # Replace tracker masks with detector masks for reconditioned objects, + # so that the memory encoder below encodes the correct (detector) masks. + # Without this, _tracker_update_memories would overwrite the reconditioned + # memory with the original tracker masks. + for trk_obj_id in reconditioned_obj_ids: + det_idx = trk_id_to_max_iou_high_conf_det.get(trk_obj_id) + if det_idx is None: + continue + trk_idx_arr = np.where( + tracker_metadata_prev["obj_ids_all_gpu"] == trk_obj_id + )[0] + if len(trk_idx_arr) == 0: + continue + trk_idx = trk_idx_arr.item() + tracker_low_res_masks_global[trk_idx] = det_out["mask"][det_idx] # Step 4: Run SAM2 memory encoder on the current frame's prediction masks # This is done on all GPUs