Spaces:
Running
Running
| from inference.memory_manager import MemoryManager | |
| from model.network import XMem | |
| from model.aggregate import aggregate | |
| from tracker.util.tensor_util import pad_divide_by, unpad | |
| class InferenceCore: | |
| def __init__(self, network: XMem, config): | |
| self.config = config | |
| self.network = network | |
| self.mem_every = config["mem_every"] | |
| self.deep_update_every = config["deep_update_every"] | |
| self.enable_long_term = config["enable_long_term"] | |
| # if deep_update_every < 0, synchronize deep update with memory frame | |
| self.deep_update_sync = self.deep_update_every < 0 | |
| self.clear_memory() | |
| self.all_labels = None | |
| def clear_memory(self): | |
| self.curr_ti = -1 | |
| self.last_mem_ti = 0 | |
| if not self.deep_update_sync: | |
| self.last_deep_update_ti = -self.deep_update_every | |
| self.memory = MemoryManager(config=self.config) | |
| def update_config(self, config): | |
| self.mem_every = config["mem_every"] | |
| self.deep_update_every = config["deep_update_every"] | |
| self.enable_long_term = config["enable_long_term"] | |
| # if deep_update_every < 0, synchronize deep update with memory frame | |
| self.deep_update_sync = self.deep_update_every < 0 | |
| self.memory.update_config(config) | |
| def set_all_labels(self, all_labels): | |
| # self.all_labels = [l.item() for l in all_labels] | |
| self.all_labels = all_labels | |
| def step(self, image, mask=None, valid_labels=None, end=False): | |
| # image: 3*H*W | |
| # mask: num_objects*H*W or None | |
| self.curr_ti += 1 | |
| image, self.pad = pad_divide_by(image, 16) | |
| image = image.unsqueeze(0) # add the batch dimension | |
| is_mem_frame = ( | |
| (self.curr_ti - self.last_mem_ti >= self.mem_every) or (mask is not None) | |
| ) and (not end) | |
| need_segment = (self.curr_ti > 0) and ( | |
| (valid_labels is None) or (len(self.all_labels) != len(valid_labels)) | |
| ) | |
| is_deep_update = ( | |
| (self.deep_update_sync and is_mem_frame) | |
| or ( # synchronized | |
| not self.deep_update_sync | |
| and self.curr_ti - self.last_deep_update_ti >= self.deep_update_every | |
| ) # no-sync | |
| ) and (not end) | |
| is_normal_update = (not self.deep_update_sync or not is_deep_update) and ( | |
| not end | |
| ) | |
| key, shrinkage, selection, f16, f8, f4 = self.network.encode_key( | |
| image, need_ek=(self.enable_long_term or need_segment), need_sk=is_mem_frame | |
| ) | |
| multi_scale_features = (f16, f8, f4) | |
| # segment the current frame is needed | |
| if need_segment: | |
| memory_readout = self.memory.match_memory(key, selection).unsqueeze(0) | |
| hidden, pred_logits_with_bg, pred_prob_with_bg = self.network.segment( | |
| multi_scale_features, | |
| memory_readout, | |
| self.memory.get_hidden(), | |
| h_out=is_normal_update, | |
| strip_bg=False, | |
| ) | |
| # remove batch dim | |
| pred_prob_with_bg = pred_prob_with_bg[0] | |
| pred_prob_no_bg = pred_prob_with_bg[1:] | |
| pred_logits_with_bg = pred_logits_with_bg[0] | |
| pred_logits_no_bg = pred_logits_with_bg[1:] | |
| if is_normal_update: | |
| self.memory.set_hidden(hidden) | |
| else: | |
| pred_prob_no_bg = ( | |
| pred_prob_with_bg | |
| ) = pred_logits_with_bg = pred_logits_no_bg = None | |
| # use the input mask if any | |
| if mask is not None: | |
| mask, _ = pad_divide_by(mask, 16) | |
| if pred_prob_no_bg is not None: | |
| # if we have a predicted mask, we work on it | |
| # make pred_prob_no_bg consistent with the input mask | |
| mask_regions = mask.sum(0) > 0.5 | |
| pred_prob_no_bg[:, mask_regions] = 0 | |
| # shift by 1 because mask/pred_prob_no_bg do not contain background | |
| mask = mask.type_as(pred_prob_no_bg) | |
| if valid_labels is not None: | |
| shift_by_one_non_labels = [ | |
| i | |
| for i in range(pred_prob_no_bg.shape[0]) | |
| if (i + 1) not in valid_labels | |
| ] | |
| # non-labelled objects are copied from the predicted mask | |
| mask[shift_by_one_non_labels] = pred_prob_no_bg[ | |
| shift_by_one_non_labels | |
| ] | |
| pred_prob_with_bg = aggregate(mask, dim=0) | |
| # also create new hidden states | |
| self.memory.create_hidden_state(len(self.all_labels), key) | |
| # save as memory if needed | |
| if is_mem_frame: | |
| value, hidden = self.network.encode_value( | |
| image, | |
| f16, | |
| self.memory.get_hidden(), | |
| pred_prob_with_bg[1:].unsqueeze(0), | |
| is_deep_update=is_deep_update, | |
| ) | |
| self.memory.add_memory( | |
| key, | |
| shrinkage, | |
| value, | |
| self.all_labels, | |
| selection=selection if self.enable_long_term else None, | |
| ) | |
| self.last_mem_ti = self.curr_ti | |
| if is_deep_update: | |
| self.memory.set_hidden(hidden) | |
| self.last_deep_update_ti = self.curr_ti | |
| if pred_logits_with_bg is None: | |
| return unpad(pred_prob_with_bg, self.pad), None | |
| else: | |
| return unpad(pred_prob_with_bg, self.pad), unpad( | |
| pred_logits_with_bg, self.pad | |
| ) | |