| import logging |
|
|
| from typing import List, Optional, Tuple, Union |
|
|
| import numpy |
| import numpy as np |
| import torch |
| from PIL.Image import Image |
|
|
| from model.visual.sam2.modeling.sam2_base import SAM2Base |
|
|
| from model.visual.sam2.modeling.backbones.hieradet import Hiera |
| from model.visual.sam2.modeling.backbones.image_encoder import FpnNeck |
| from model.visual.sam2.modeling.backbones.image_encoder import ImageEncoder |
| from model.visual.sam2.modeling.position_encoding import PositionEmbeddingSine |
|
|
| from model.visual.sam2.modeling.memory_attention import MemoryAttention |
| from model.visual.sam2.modeling.memory_attention import MemoryAttentionLayer |
| from model.visual.sam2.modeling.sam.transformer import RoPEAttention |
| from model.visual.sam2.modeling.memory_encoder import MemoryEncoder |
| from model.visual.sam2.modeling.memory_encoder import MaskDownSampler |
| from model.visual.sam2.modeling.memory_encoder import Fuser |
| from model.visual.sam2.modeling.memory_encoder import CXBlock |
|
|
| from model.visual.sam2.utils.transforms import SAM2Transforms |
| from model.visual.sam2.modeling.backbones.hieradet import do_pool |
| from model.visual.sam2.modeling.backbones.utils import ( |
| PatchEmbed, |
| window_partition, |
| window_unpartition, |
| ) |
|
|
|
|
| class AVmodel(torch.nn.Module): |
| """End-to-end AV segmentation: SAM2 visual backbone + AuralFuser audio-visual fusion + tracking head.""" |
|
|
| 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))] |
|
|
| from model.visual.sam2.build_sam import build_sam2_visual_predictor |
| self.v_model = build_sam2_visual_predictor(self.param.sam_config_path, self.param.backbone_weight, |
| apply_postprocessing=True, mode='train') |
| 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, |
| ) |
| from model.aural_fuser import AuralFuser |
| self.aural_fuser = AuralFuser(hyp_param=self.param) |
|
|
|
|
|
|
| def _prepare_backbone_features(self, backbone_out): |
| """Prepare and flatten visual features.""" |
| backbone_out = backbone_out.copy() |
| assert len(backbone_out["backbone_fpn"]) == len(backbone_out["vision_pos_enc"]) |
| assert len(backbone_out["backbone_fpn"]) >= self.v_model.num_feature_levels |
|
|
| feature_maps = backbone_out["backbone_fpn"][-self.v_model.num_feature_levels:] |
| vision_pos_embeds = backbone_out["vision_pos_enc"][-self.v_model.num_feature_levels:] |
|
|
| feat_sizes = [(x.shape[-2], x.shape[-1]) for x in vision_pos_embeds] |
| vision_feats = [x.flatten(2).permute(2, 0, 1) for x in feature_maps] |
| vision_pos_embeds = [x.flatten(2).permute(2, 0, 1) for x in vision_pos_embeds] |
|
|
| return backbone_out, vision_feats, vision_pos_embeds, feat_sizes |
|
|
| def forward_frame(self, frame_): |
| 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 into FPN features, then run SAM2 tracking. `sam_process` is reserved for prompt path.""" |
| backbone_feats = self.v_model.forward_image(frames, pre_compute=False) |
| audio_residual_feats = self.aural_fuser(backbone_feats, spect) |
| 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], |
| cond_frame=int(frames.shape[0]/2) if self.training else 0) |
| 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): |
| self.v_model.eval() |
| for name, parameter in self.v_model.named_parameters(): |
| parameter.requires_grad = False |
|
|
