Spaces:
Running on Zero
Running on Zero
| # sam3/model/sam3_tracker_train.py | |
| """ | |
| SAM3 Tracker for VOS training (similar to SAM2Train). | |
| """ | |
| import logging | |
| import numpy as np | |
| import torch | |
| from sam3.model.sam3_tracker_base import Sam3TrackerBase | |
| from sam3.model.sam3_tracker_utils import get_next_point, sample_box_points | |
| # from sam3.train.data.collator import BatchedDatapoint | |
| from sam3.model.data_misc import BatchedDatapoint | |
| from sam3.model.utils.misc import concat_points | |
| logger = logging.getLogger(__name__) | |
| class Sam3TrackerTrain(Sam3TrackerBase): | |
| def __init__( | |
| self, | |
| # Point/mask input sampling probabilities | |
| prob_to_use_pt_input_for_train=0.0, | |
| prob_to_use_pt_input_for_eval=0.0, | |
| prob_to_use_box_input_for_train=0.0, | |
| prob_to_use_box_input_for_eval=0.0, | |
| # Number of frames to add correction points | |
| num_frames_to_correct_for_train=1, | |
| num_frames_to_correct_for_eval=1, | |
| rand_frames_to_correct_for_train=False, | |
| rand_frames_to_correct_for_eval=False, | |
| # Initial conditioning frames | |
| num_init_cond_frames_for_train=1, | |
| num_init_cond_frames_for_eval=1, | |
| rand_init_cond_frames_for_train=True, | |
| rand_init_cond_frames_for_eval=False, | |
| add_all_frames_to_correct_as_cond=False, | |
| # Correction point sampling | |
| num_correction_pt_per_frame=7, | |
| pt_sampling_for_eval="center", | |
| prob_to_sample_from_gt_for_train=0.0, | |
| use_act_ckpt_iterative_pt_sampling=False, | |
| # Whether to freeze image encoder | |
| freeze_image_encoder=False, | |
| prob_to_dropout_spatial_mem=0.0, | |
| teacher_force_obj_scores_for_mem=False, | |
| use_memory_selection=False, | |
| **kwargs, | |
| ): | |
| super().__init__(**kwargs) | |
| # Point sampler settings | |
| self.prob_to_use_pt_input_for_train = prob_to_use_pt_input_for_train | |
| self.prob_to_use_box_input_for_train = prob_to_use_box_input_for_train | |
| self.prob_to_use_pt_input_for_eval = prob_to_use_pt_input_for_eval | |
| self.prob_to_use_box_input_for_eval = prob_to_use_box_input_for_eval | |
| self.num_frames_to_correct_for_train = num_frames_to_correct_for_train | |
| self.num_frames_to_correct_for_eval = num_frames_to_correct_for_eval | |
| self.rand_frames_to_correct_for_train = rand_frames_to_correct_for_train | |
| self.rand_frames_to_correct_for_eval = rand_frames_to_correct_for_eval | |
| self.num_init_cond_frames_for_train = num_init_cond_frames_for_train | |
| self.num_init_cond_frames_for_eval = num_init_cond_frames_for_eval | |
| self.rand_init_cond_frames_for_train = rand_init_cond_frames_for_train | |
| self.rand_init_cond_frames_for_eval = rand_init_cond_frames_for_eval | |
| self.add_all_frames_to_correct_as_cond = add_all_frames_to_correct_as_cond | |
| self.num_correction_pt_per_frame = num_correction_pt_per_frame | |
| self.pt_sampling_for_eval = pt_sampling_for_eval | |
| self.prob_to_sample_from_gt_for_train = prob_to_sample_from_gt_for_train | |
| self.use_act_ckpt_iterative_pt_sampling = use_act_ckpt_iterative_pt_sampling | |
| # Random number generator with fixed seed | |
| self.rng = np.random.default_rng(seed=42) | |
| self.prob_to_dropout_spatial_mem = prob_to_dropout_spatial_mem | |
| self.teacher_force_obj_scores_for_mem = teacher_force_obj_scores_for_mem | |
| self.use_memory_selection = use_memory_selection | |
| if freeze_image_encoder: | |
| for p in self.backbone.parameters(): | |
| p.requires_grad = False | |
| ############################################### | |
| # SAM3Matting | |
| for name, p in self.named_parameters(): | |
| if 'alpha' not in name and 'unknown' not in name: | |
| p.requires_grad = False | |
| # trainable = [name for name, p in self.named_parameters() if p.requires_grad] | |
| # for name in trainable: | |
| # print(f" {name}") | |
| ############################################### | |
| def forward(self, input: BatchedDatapoint): | |
| if self.training or not self.forward_backbone_per_frame_for_eval: | |
| backbone_out = self.forward_image(input.img_batch) | |
| else: | |
| backbone_out = {"backbone_fpn": None, "vision_pos_enc": None} | |
| backbone_out = self.prepare_prompt_inputs(backbone_out, input) | |
| previous_stages_out = self.forward_tracking(backbone_out, input) | |
| return previous_stages_out | |
| def prepare_prompt_inputs(self, backbone_out, input, start_frame_idx=0): | |
| """ | |
| Prepare input mask, point or box prompts. | |
| """ | |
| gt_masks_per_frame = {} | |
| for stage_id in range(len(input.find_inputs)): | |
| segments = input.find_targets[stage_id].segments | |
| gt_masks_per_frame[stage_id] = segments.unsqueeze(1) | |
| backbone_out["gt_masks_per_frame"] = gt_masks_per_frame | |
| num_frames = len(input.find_inputs) | |
| backbone_out["num_frames"] = num_frames | |
| # Decide whether to use point or mask inputs | |
| if self.training: | |
| prob_to_use_pt_input = self.prob_to_use_pt_input_for_train | |
| prob_to_use_box_input = self.prob_to_use_box_input_for_train | |
| num_frames_to_correct = self.num_frames_to_correct_for_train | |
| rand_frames_to_correct = self.rand_frames_to_correct_for_train | |
| num_init_cond_frames = self.num_init_cond_frames_for_train | |
| rand_init_cond_frames = self.rand_init_cond_frames_for_train | |
| else: | |
| prob_to_use_pt_input = self.prob_to_use_pt_input_for_eval | |
| prob_to_use_box_input = self.prob_to_use_box_input_for_eval | |
| num_frames_to_correct = self.num_frames_to_correct_for_eval | |
| rand_frames_to_correct = self.rand_frames_to_correct_for_eval | |
| num_init_cond_frames = self.num_init_cond_frames_for_eval | |
| rand_init_cond_frames = self.rand_init_cond_frames_for_eval | |
| if num_frames == 1: | |
| # Special case: single image (force point input) | |
| prob_to_use_pt_input = 1.0 | |
| num_frames_to_correct = 1 | |
| num_init_cond_frames = 1 | |
| assert num_init_cond_frames >= 1 | |
| use_pt_input = self.rng.random() < prob_to_use_pt_input | |
| ##################################### | |
| # SAM3Matting | |
| use_pt_input = False | |
| ##################################### | |
| # Sample initial conditioning frames | |
| if num_init_cond_frames == 1: | |
| init_cond_frames = [start_frame_idx] | |
| else: | |
| if rand_init_cond_frames: | |
| num_init_cond_frames = self.rng.integers(1, num_init_cond_frames, endpoint=True) | |
| init_cond_frames = [start_frame_idx] + self.rng.choice( | |
| range(start_frame_idx + 1, num_frames), | |
| num_init_cond_frames - 1, | |
| replace=False, | |
| ).tolist() | |
| backbone_out["init_cond_frames"] = init_cond_frames | |
| backbone_out["frames_not_in_init_cond"] = [ | |
| t for t in range(start_frame_idx, num_frames) if t not in init_cond_frames | |
| ] | |
| backbone_out["use_pt_input"] = use_pt_input | |
| # Prepare mask or point inputs on initial conditioning frames | |
| backbone_out["mask_inputs_per_frame"] = {} | |
| backbone_out["point_inputs_per_frame"] = {} | |
| assert len(init_cond_frames) == 1 | |
| for t in init_cond_frames: | |
| if not use_pt_input: | |
| # Use mask inputs directly | |
| backbone_out["mask_inputs_per_frame"][t] = gt_masks_per_frame[t] | |
| else: | |
| assert False | |
| # Sample points from GT masks | |
| use_box_input = self.rng.random() < prob_to_use_box_input | |
| if use_box_input: | |
| points, labels = sample_box_points(gt_masks_per_frame[t]) | |
| else: | |
| points, labels = get_next_point( | |
| gt_masks=gt_masks_per_frame[t], | |
| pred_masks=None, | |
| method="uniform" if self.training else self.pt_sampling_for_eval, | |
| ) | |
| point_inputs = {"point_coords": points, "point_labels": labels} | |
| backbone_out["point_inputs_per_frame"][t] = point_inputs | |
| # Sample frames for correction clicks | |
| if not use_pt_input: | |
| frames_to_add_correction_pt = [] | |
| elif num_frames_to_correct == num_init_cond_frames: | |
| frames_to_add_correction_pt = init_cond_frames | |
| else: | |
| if rand_frames_to_correct and num_frames_to_correct > num_init_cond_frames: | |
| num_frames_to_correct = self.rng.integers( | |
| num_init_cond_frames, num_frames_to_correct, endpoint=True | |
| ) | |
| extra_num = num_frames_to_correct - num_init_cond_frames | |
| frames_to_add_correction_pt = ( | |
| init_cond_frames | |
| + self.rng.choice( | |
| backbone_out["frames_not_in_init_cond"], extra_num, replace=False | |
| ).tolist() | |
| ) | |
| backbone_out["frames_to_add_correction_pt"] = frames_to_add_correction_pt | |
| return backbone_out | |
| def back_convert(self, targets): | |
| """ | |
| Convert targets for loss computation. | |
| For VOS, we return targets as-is since Sam3LossVOS | |
| expects BatchedFindTarget format. | |
| """ | |
| return targets |