| import numpy as np |
|
|
| from utils.image import one_hot_mask |
|
|
| from networks.layers.basic import seq_to_2d |
| from networks.engines.aot_engine import AOTEngine, AOTInferEngine |
|
|
|
|
| class DeAOTEngine(AOTEngine): |
| def __init__(self, |
| aot_model, |
| gpu_id=0, |
| long_term_mem_gap=9999, |
| short_term_mem_skip=1, |
| layer_loss_scaling_ratio=2., |
| max_len_long_term=9999): |
| super().__init__(aot_model, gpu_id, long_term_mem_gap, |
| short_term_mem_skip, max_len_long_term) |
| self.layer_loss_scaling_ratio = layer_loss_scaling_ratio |
| def update_short_term_memory(self, curr_mask, curr_id_emb=None, skip_long_term_update=False): |
|
|
| if curr_id_emb is None: |
| if len(curr_mask.size()) == 3 or curr_mask.size()[0] == 1: |
| curr_one_hot_mask = one_hot_mask(curr_mask, self.max_obj_num) |
| else: |
| curr_one_hot_mask = curr_mask |
| curr_id_emb = self.assign_identity(curr_one_hot_mask) |
|
|
| lstt_curr_memories = self.curr_lstt_output[1] |
| lstt_curr_memories_2d = [] |
| for layer_idx in range(len(lstt_curr_memories)): |
| curr_k, curr_v, curr_id_k, curr_id_v = lstt_curr_memories[ |
| layer_idx] |
| curr_id_k, curr_id_v = self.AOT.LSTT.layers[ |
| layer_idx].fuse_key_value_id(curr_id_k, curr_id_v, curr_id_emb) |
| lstt_curr_memories[layer_idx][2], lstt_curr_memories[layer_idx][ |
| 3] = curr_id_k, curr_id_v |
| local_curr_id_k = seq_to_2d( |
| curr_id_k, self.enc_size_2d) if curr_id_k is not None else None |
| local_curr_id_v = seq_to_2d(curr_id_v, self.enc_size_2d) |
| lstt_curr_memories_2d.append([ |
| seq_to_2d(curr_k, self.enc_size_2d), |
| seq_to_2d(curr_v, self.enc_size_2d), local_curr_id_k, |
| local_curr_id_v |
| ]) |
|
|
| self.short_term_memories_list.append(lstt_curr_memories_2d) |
| self.short_term_memories_list = self.short_term_memories_list[ |
| -self.short_term_mem_skip:] |
| self.short_term_memories = self.short_term_memories_list[0] |
|
|
| if self.frame_step - self.last_mem_step >= self.long_term_mem_gap: |
| |
| if not skip_long_term_update: |
| self.update_long_term_memory(lstt_curr_memories) |
| self.last_mem_step = self.frame_step |
|
|
|
|
| class DeAOTInferEngine(AOTInferEngine): |
| def __init__(self, |
| aot_model, |
| gpu_id=0, |
| long_term_mem_gap=9999, |
| short_term_mem_skip=1, |
| max_aot_obj_num=None, |
| max_len_long_term=9999): |
| super().__init__(aot_model, gpu_id, long_term_mem_gap, |
| short_term_mem_skip, max_aot_obj_num, max_len_long_term) |
| def add_reference_frame(self, img, mask, obj_nums, frame_step=-1): |
| if isinstance(obj_nums, list): |
| obj_nums = obj_nums[0] |
| self.obj_nums = obj_nums |
| aot_num = max(np.ceil(obj_nums / self.max_aot_obj_num), 1) |
| while (aot_num > len(self.aot_engines)): |
| new_engine = DeAOTEngine(self.AOT, self.gpu_id, |
| self.long_term_mem_gap, |
| self.short_term_mem_skip, |
| max_len_long_term = self.max_len_long_term) |
| new_engine.eval() |
| self.aot_engines.append(new_engine) |
|
|
| separated_masks, separated_obj_nums = self.separate_mask( |
| mask, obj_nums) |
| img_embs = None |
| for aot_engine, separated_mask, separated_obj_num in zip( |
| self.aot_engines, separated_masks, separated_obj_nums): |
| if aot_engine.obj_nums is None or aot_engine.obj_nums[0] < separated_obj_num: |
| aot_engine.add_reference_frame(img, |
| separated_mask, |
| obj_nums=[separated_obj_num], |
| frame_step=frame_step, |
| img_embs=img_embs) |
| else: |
| aot_engine.update_short_term_memory(separated_mask) |
| if img_embs is None: |
| img_embs = aot_engine.curr_enc_embs |
|
|
| self.update_size() |
|
|
