Spaces:
Runtime error
Runtime error
| 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) | |