| """End-to-end Ref-AVS: SAM2 visual backbone + AuralFuser fusion + tracking head. |
| |
| Orchestration follows ``avs.code/v1m.code/model/mymodel.py``. |
| """ |
| import torch |
|
|
| from model.visual.sam2.build_sam import build_sam2_visual_predictor |
| from model.visual.sam2.utils.transforms import SAM2Transforms |
| from model.aural_fuser import AuralFuser |
| from transformers import AutoTokenizer, AutoModel |
|
|
|
|
| class AVmodel(torch.nn.Module): |
| """SAM2 + audio/text fusion (``aural_fuser``) + SAM2 tracking decoder.""" |
|
|
| def __init__(self, param, mask_threshold=0.0, max_hole_area=0.0, max_sprinkle_area=0.0): |
| super().__init__() |
| self.param = param |
| self.mask_threshold = mask_threshold |
| self._bb_feat_sizes = [ |
| (int(self.param.image_size / 4), int(self.param.image_size / 4)), |
| (int(self.param.image_size / 8), int(self.param.image_size / 8)), |
| (int(self.param.image_size / 16), int(self.param.image_size / 16)), |
| ] |
|
|
| self.v_model = build_sam2_visual_predictor( |
| self.param.sam_config_path, |
| self.param.backbone_weight, |
| apply_postprocessing=True, |
| mode='train', |
| hydra_overrides_extra=["++model.image_size={}".format(self.param.image_size)], |
| ) |
| self._transforms = SAM2Transforms( |
| resolution=self.v_model.image_size, |
| mask_threshold=mask_threshold, |
| max_hole_area=max_hole_area, |
| max_sprinkle_area=max_sprinkle_area, |
| ) |
| self.aural_fuser = AuralFuser(hyp_param=self.param) |
| self.text_tokenizer = AutoTokenizer.from_pretrained('distilbert/distilroberta-base') |
| self.t_model = AutoModel.from_pretrained('distilbert/distilroberta-base') |
|
|
| def _encode_text(self, prompts): |
| """RoBERTa embeddings for referring expressions (frozen at train time).""" |
| enc = self.text_tokenizer( |
| *prompts, |
| max_length=25, |
| padding="max_length", |
| truncation=True, |
| return_tensors="pt", |
| ) |
| enc['input_ids'] = enc['input_ids'].cuda(self.param.local_rank, non_blocking=True) |
| enc['attention_mask'] = enc['attention_mask'].cuda(self.param.local_rank, non_blocking=True) |
| with torch.no_grad(): |
| return self.t_model(**enc).last_hidden_state |
|
|
| def forward_frame(self, frame_): |
| """Single-frame SAM2 image encoder pass (same helper pattern as v1m).""" |
| frame = torch.nn.functional.interpolate( |
| frame_, (self.param.image_size, self.param.image_size), |
| antialias=True, align_corners=False, mode='bilinear', |
| ) |
| return self.v_model.image_encoder(frame) |
|
|
| def forward(self, frames, spect, prompts, sam_process=False): |
| """Fuse audio+text into FPN, then run SAM2 tracking without box/mask prompts.""" |
| text_feats = self._encode_text(prompts) |
| backbone_feats = self.v_model.forward_image(frames, pre_compute=False) |
| audio_residual_feats = self.aural_fuser(backbone_feats, spect, text_feats) |
| visual_resfeats, audio_resfeats, proj_feats = audio_residual_feats |
|
|
| map_res = visual_resfeats[::-1] |
| vec_res = audio_resfeats[::-1] |
| av_feats = (map_res, vec_res) |
|
|
| backbone_feats = self.v_model.precompute_high_res_features(backbone_feats) |
| backbone_feats = self.v_model.dont_prepare_prompt_inputs( |
| backbone_feats, |
| num_frames=frames.shape[0], |
| condition_frame=int(frames.shape[0] / 2), |
| ) |
| outputs = self.v_model.forward_tracking_wo_prompt(backbone_feats, audio_res=av_feats) |
| return outputs, proj_feats |
|
|
| @property |
| def device(self) -> torch.device: |
| return self.v_model.device |
|
|
| def freeze_sam_parameters(self): |
| """Freeze SAM2 and text backbone; only ``aural_fuser`` is trained.""" |
| self.v_model.eval() |
| self.t_model.eval() |
| for _, parameter in self.v_model.named_parameters(): |
| parameter.requires_grad = False |
|
|
