release infer and demo
This commit is contained in:
pq-yang
@@ -0,0 +1,56 @@
|
||||
import warnings
|
||||
from typing import Iterable
|
||||
import torch
|
||||
from matanyone2.model.matanyone2 import MatAnyone2
|
||||
|
||||
|
||||
class ImageFeatureStore:
|
||||
"""
|
||||
A cache for image features.
|
||||
These features might be reused at different parts of the inference pipeline.
|
||||
This class provide an interface for reusing these features.
|
||||
It is the user's responsibility to delete redundant features.
|
||||
|
||||
Feature of a frame should be associated with a unique index -- typically the frame id.
|
||||
"""
|
||||
def __init__(self, network: MatAnyone2, no_warning: bool = False):
|
||||
self.network = network
|
||||
self._store = {}
|
||||
self.no_warning = no_warning
|
||||
|
||||
def _encode_feature(self, index: int, image: torch.Tensor, last_feats=None) -> None:
|
||||
ms_features, pix_feat = self.network.encode_image(image, last_feats=last_feats)
|
||||
key, shrinkage, selection = self.network.transform_key(ms_features[0])
|
||||
self._store[index] = (ms_features, pix_feat, key, shrinkage, selection)
|
||||
|
||||
def get_all_features(self, images: torch.Tensor) -> (Iterable[torch.Tensor], torch.Tensor):
|
||||
seq_length = images.shape[0]
|
||||
ms_features, pix_feat = self.network.encode_image(images, seq_length)
|
||||
key, shrinkage, selection = self.network.transform_key(ms_features[0])
|
||||
for index in range(seq_length):
|
||||
self._store[index] = ([f[index].unsqueeze(0) for f in ms_features], pix_feat[index].unsqueeze(0), key[index].unsqueeze(0), shrinkage[index].unsqueeze(0), selection[index].unsqueeze(0))
|
||||
|
||||
def get_features(self, index: int,
|
||||
image: torch.Tensor, last_feats=None) -> (Iterable[torch.Tensor], torch.Tensor):
|
||||
if index not in self._store:
|
||||
self._encode_feature(index, image, last_feats)
|
||||
|
||||
return self._store[index][:2]
|
||||
|
||||
def get_key(self, index: int,
|
||||
image: torch.Tensor, last_feats=None) -> (torch.Tensor, torch.Tensor, torch.Tensor):
|
||||
if index not in self._store:
|
||||
self._encode_feature(index, image, last_feats)
|
||||
|
||||
return self._store[index][2:]
|
||||
|
||||
def delete(self, index: int) -> None:
|
||||
if index in self._store:
|
||||
del self._store[index]
|
||||
|
||||
def __len__(self):
|
||||
return len(self._store)
|
||||
|
||||
def __del__(self):
|
||||
if len(self._store) > 0 and not self.no_warning:
|
||||
warnings.warn(f'Leaking {self._store.keys()} in the image feature store')
|
||||
@@ -0,0 +1,550 @@
|
||||
import logging
|
||||
from omegaconf import DictConfig
|
||||
from typing import List, Optional, Iterable, Union,Tuple
|
||||
|
||||
import os
|
||||
import cv2
|
||||
import torch
|
||||
import imageio
|
||||
import tempfile
|
||||
import numpy as np
|
||||
from tqdm import tqdm
|
||||
from PIL import Image
|
||||
import torch.nn.functional as F
|
||||
|
||||
from matanyone2.inference.memory_manager import MemoryManager
|
||||
from matanyone2.inference.object_manager import ObjectManager
|
||||
from matanyone2.inference.image_feature_store import ImageFeatureStore
|
||||
from matanyone2.model.matanyone2 import MatAnyone2
|
||||
from matanyone2.utils.tensor_utils import pad_divide_by, unpad, aggregate
|
||||
from matanyone2.utils.inference_utils import gen_dilate, gen_erosion, read_frame_from_videos
|
||||
from matanyone2.utils.device import get_default_device, safe_autocast
|
||||
|
||||
|
||||
log = logging.getLogger()
|
||||
|
||||
|
||||
class InferenceCore:
|
||||
|
||||
def __init__(self,
|
||||
network: Union[MatAnyone2,str],
|
||||
cfg: DictConfig = None,
|
||||
*,
|
||||
image_feature_store: ImageFeatureStore = None,
|
||||
device: Optional[Union[str, torch.device]] = None
|
||||
):
|
||||
if device is None:
|
||||
device = get_default_device()
|
||||
self.device = device
|
||||
if isinstance(network, str):
|
||||
network = MatAnyone2.from_pretrained(network)
|
||||
network.to(device)
|
||||
network.eval()
|
||||
self.network = network
|
||||
cfg = cfg if cfg is not None else network.cfg
|
||||
self.cfg = cfg
|
||||
self.mem_every = cfg.mem_every
|
||||
stagger_updates = cfg.stagger_updates
|
||||
self.chunk_size = cfg.chunk_size
|
||||
self.save_aux = cfg.save_aux
|
||||
self.max_internal_size = cfg.max_internal_size
|
||||
self.flip_aug = cfg.flip_aug
|
||||
|
||||
self.curr_ti = -1
|
||||
self.last_mem_ti = 0
|
||||
# at which time indices should we update the sensory memory
|
||||
if stagger_updates >= self.mem_every:
|
||||
self.stagger_ti = set(range(1, self.mem_every + 1))
|
||||
else:
|
||||
self.stagger_ti = set(
|
||||
np.round(np.linspace(1, self.mem_every, stagger_updates)).astype(int))
|
||||
self.object_manager = ObjectManager()
|
||||
self.memory = MemoryManager(cfg=cfg, object_manager=self.object_manager)
|
||||
|
||||
if image_feature_store is None:
|
||||
self.image_feature_store = ImageFeatureStore(self.network)
|
||||
else:
|
||||
self.image_feature_store = image_feature_store
|
||||
|
||||
self.last_mask = None
|
||||
self.last_pix_feat = None
|
||||
self.last_msk_value = None
|
||||
|
||||
def clear_memory(self):
|
||||
self.curr_ti = -1
|
||||
self.last_mem_ti = 0
|
||||
self.memory = MemoryManager(cfg=self.cfg, object_manager=self.object_manager)
|
||||
|
||||
def clear_non_permanent_memory(self):
|
||||
self.curr_ti = -1
|
||||
self.last_mem_ti = 0
|
||||
self.memory.clear_non_permanent_memory()
|
||||
|
||||
def clear_sensory_memory(self):
|
||||
self.curr_ti = -1
|
||||
self.last_mem_ti = 0
|
||||
self.memory.clear_sensory_memory()
|
||||
|
||||
def update_config(self, cfg):
|
||||
self.mem_every = cfg['mem_every']
|
||||
self.memory.update_config(cfg)
|
||||
|
||||
def clear_temp_mem(self):
|
||||
self.memory.clear_work_mem()
|
||||
# self.object_manager = ObjectManager()
|
||||
self.memory.clear_obj_mem()
|
||||
# self.memory.clear_sensory_memory()
|
||||
|
||||
def _add_memory(self,
|
||||
image: torch.Tensor,
|
||||
pix_feat: torch.Tensor,
|
||||
prob: torch.Tensor,
|
||||
key: torch.Tensor,
|
||||
shrinkage: torch.Tensor,
|
||||
selection: torch.Tensor,
|
||||
*,
|
||||
is_deep_update: bool = True,
|
||||
force_permanent: bool = False) -> None:
|
||||
"""
|
||||
Memorize the given segmentation in all memory stores.
|
||||
|
||||
The batch dimension is 1 if flip augmentation is not used.
|
||||
image: RGB image, (1/2)*3*H*W
|
||||
pix_feat: from the key encoder, (1/2)*_*H*W
|
||||
prob: (1/2)*num_objects*H*W, in [0, 1]
|
||||
key/shrinkage/selection: for anisotropic l2, (1/2)*_*H*W
|
||||
selection can be None if not using long-term memory
|
||||
is_deep_update: whether to use deep update (e.g. with the mask encoder)
|
||||
force_permanent: whether to force the memory to be permanent
|
||||
"""
|
||||
if prob.shape[1] == 0:
|
||||
# nothing to add
|
||||
log.warn('Trying to add an empty object mask to memory!')
|
||||
return
|
||||
|
||||
if force_permanent:
|
||||
as_permanent = 'all'
|
||||
else:
|
||||
as_permanent = 'first'
|
||||
|
||||
self.memory.initialize_sensory_if_needed(key, self.object_manager.all_obj_ids)
|
||||
msk_value, sensory, obj_value, _ = self.network.encode_mask(
|
||||
image,
|
||||
pix_feat,
|
||||
self.memory.get_sensory(self.object_manager.all_obj_ids),
|
||||
prob,
|
||||
deep_update=is_deep_update,
|
||||
chunk_size=self.chunk_size,
|
||||
need_weights=self.save_aux)
|
||||
self.memory.add_memory(key,
|
||||
shrinkage,
|
||||
msk_value,
|
||||
obj_value,
|
||||
self.object_manager.all_obj_ids,
|
||||
selection=selection,
|
||||
as_permanent=as_permanent)
|
||||
self.last_mem_ti = self.curr_ti
|
||||
if is_deep_update:
|
||||
self.memory.update_sensory(sensory, self.object_manager.all_obj_ids)
|
||||
self.last_msk_value = msk_value
|
||||
|
||||
def _segment(self,
|
||||
key: torch.Tensor,
|
||||
selection: torch.Tensor,
|
||||
pix_feat: torch.Tensor,
|
||||
ms_features: Iterable[torch.Tensor],
|
||||
update_sensory: bool = True) -> torch.Tensor:
|
||||
"""
|
||||
Produce a segmentation using the given features and the memory
|
||||
|
||||
The batch dimension is 1 if flip augmentation is not used.
|
||||
key/selection: for anisotropic l2: (1/2) * _ * H * W
|
||||
pix_feat: from the key encoder, (1/2) * _ * H * W
|
||||
ms_features: an iterable of multiscale features from the encoder, each is (1/2)*_*H*W
|
||||
with strides 16, 8, and 4 respectively
|
||||
update_sensory: whether to update the sensory memory
|
||||
|
||||
Returns: (num_objects+1)*H*W normalized probability; the first channel is the background
|
||||
"""
|
||||
bs = key.shape[0]
|
||||
if self.flip_aug:
|
||||
assert bs == 2
|
||||
else:
|
||||
assert bs == 1
|
||||
|
||||
if not self.memory.engaged:
|
||||
log.warn('Trying to segment without any memory!')
|
||||
return torch.zeros((1, key.shape[-2] * 16, key.shape[-1] * 16),
|
||||
device=key.device,
|
||||
dtype=key.dtype)
|
||||
|
||||
uncert_output = None
|
||||
|
||||
if self.curr_ti == 0: # ONLY for the first frame for prediction
|
||||
memory_readout = self.memory.read_first_frame(self.last_msk_value, pix_feat, self.last_mask, self.network, uncert_output=uncert_output)
|
||||
else:
|
||||
memory_readout = self.memory.read(pix_feat, key, selection, self.last_mask, self.network, uncert_output=uncert_output, last_msk_value=self.last_msk_value, ti=self.curr_ti,
|
||||
last_pix_feat=self.last_pix_feat, last_pred_mask=self.last_mask)
|
||||
memory_readout = self.object_manager.realize_dict(memory_readout)
|
||||
|
||||
sensory, _, pred_prob_with_bg = self.network.segment(ms_features,
|
||||
memory_readout,
|
||||
self.memory.get_sensory(
|
||||
self.object_manager.all_obj_ids),
|
||||
chunk_size=self.chunk_size,
|
||||
update_sensory=update_sensory)
|
||||
# remove batch dim
|
||||
if self.flip_aug:
|
||||
# average predictions of the non-flipped and flipped version
|
||||
pred_prob_with_bg = (pred_prob_with_bg[0] +
|
||||
torch.flip(pred_prob_with_bg[1], dims=[-1])) / 2
|
||||
else:
|
||||
pred_prob_with_bg = pred_prob_with_bg[0]
|
||||
if update_sensory:
|
||||
self.memory.update_sensory(sensory, self.object_manager.all_obj_ids)
|
||||
return pred_prob_with_bg
|
||||
|
||||
def pred_all_flow(self, images):
|
||||
self.total_len = images.shape[0]
|
||||
images, self.pad = pad_divide_by(images, 16)
|
||||
images = images.unsqueeze(0) # add the batch dimension: (1,t,c,h,w)
|
||||
|
||||
self.flows_forward, self.flows_backward = self.network.pred_forward_backward_flow(images)
|
||||
|
||||
def encode_all_images(self, images):
|
||||
images, self.pad = pad_divide_by(images, 16)
|
||||
self.image_feature_store.get_all_features(images) # t c h w
|
||||
return images
|
||||
|
||||
def step(self,
|
||||
image: torch.Tensor,
|
||||
mask: Optional[torch.Tensor] = None,
|
||||
objects: Optional[List[int]] = None,
|
||||
*,
|
||||
idx_mask: bool = False,
|
||||
end: bool = False,
|
||||
delete_buffer: bool = True,
|
||||
force_permanent: bool = False,
|
||||
matting: bool = True,
|
||||
first_frame_pred: bool = False) -> torch.Tensor:
|
||||
"""
|
||||
Take a step with a new incoming image.
|
||||
If there is an incoming mask with new objects, we will memorize them.
|
||||
If there is no incoming mask, we will segment the image using the memory.
|
||||
In both cases, we will update the memory and return a segmentation.
|
||||
|
||||
image: 3*H*W
|
||||
mask: H*W (if idx mask) or len(objects)*H*W or None
|
||||
objects: list of object ids that are valid in the mask Tensor.
|
||||
The ids themselves do not need to be consecutive/in order, but they need to be
|
||||
in the same position in the list as the corresponding mask
|
||||
in the tensor in non-idx-mask mode.
|
||||
objects is ignored if the mask is None.
|
||||
If idx_mask is False and objects is None, we sequentially infer the object ids.
|
||||
idx_mask: if True, mask is expected to contain an object id at every pixel.
|
||||
If False, mask should have multiple channels with each channel representing one object.
|
||||
end: if we are at the end of the sequence, we do not need to update memory
|
||||
if unsure just set it to False
|
||||
delete_buffer: whether to delete the image feature buffer after this step
|
||||
force_permanent: the memory recorded this frame will be added to the permanent memory
|
||||
"""
|
||||
if objects is None and mask is not None:
|
||||
assert not idx_mask
|
||||
objects = list(range(1, mask.shape[0] + 1))
|
||||
|
||||
# resize input if needed -- currently only used for the GUI
|
||||
resize_needed = False
|
||||
if self.max_internal_size > 0:
|
||||
h, w = image.shape[-2:]
|
||||
min_side = min(h, w)
|
||||
if min_side > self.max_internal_size:
|
||||
resize_needed = True
|
||||
new_h = int(h / min_side * self.max_internal_size)
|
||||
new_w = int(w / min_side * self.max_internal_size)
|
||||
image = F.interpolate(image.unsqueeze(0),
|
||||
size=(new_h, new_w),
|
||||
mode='bilinear',
|
||||
align_corners=False)[0]
|
||||
if mask is not None:
|
||||
if idx_mask:
|
||||
mask = F.interpolate(mask.unsqueeze(0).unsqueeze(0).float(),
|
||||
size=(new_h, new_w),
|
||||
mode='nearest-exact',
|
||||
align_corners=False)[0, 0].round().long()
|
||||
else:
|
||||
mask = F.interpolate(mask.unsqueeze(0),
|
||||
size=(new_h, new_w),
|
||||
mode='bilinear',
|
||||
align_corners=False)[0]
|
||||
|
||||
self.curr_ti += 1
|
||||
|
||||
image, self.pad = pad_divide_by(image, 16) # DONE alreay for 3DCNN!!
|
||||
image = image.unsqueeze(0) # add the batch dimension
|
||||
if self.flip_aug:
|
||||
image = torch.cat([image, torch.flip(image, dims=[-1])], dim=0)
|
||||
|
||||
# whether to update the working memory
|
||||
is_mem_frame = ((self.curr_ti - self.last_mem_ti >= self.mem_every) or
|
||||
(mask is not None)) and (not end)
|
||||
# segment when there is no input mask or when the input mask is incomplete
|
||||
need_segment = (mask is None) or (self.object_manager.num_obj > 0
|
||||
and not self.object_manager.has_all(objects))
|
||||
update_sensory = ((self.curr_ti - self.last_mem_ti) in self.stagger_ti) and (not end)
|
||||
|
||||
# reinit if it is the first frame for prediction
|
||||
if first_frame_pred:
|
||||
self.curr_ti = 0
|
||||
self.last_mem_ti = 0
|
||||
is_mem_frame = True
|
||||
need_segment = True
|
||||
update_sensory = True
|
||||
|
||||
# encoding the image
|
||||
ms_feat, pix_feat = self.image_feature_store.get_features(self.curr_ti, image)
|
||||
key, shrinkage, selection = self.image_feature_store.get_key(self.curr_ti, image)
|
||||
|
||||
# segmentation from memory if needed
|
||||
if need_segment:
|
||||
pred_prob_with_bg = self._segment(key,
|
||||
selection,
|
||||
pix_feat,
|
||||
ms_feat,
|
||||
update_sensory=update_sensory)
|
||||
|
||||
# use the input mask if provided
|
||||
if mask is not None:
|
||||
# inform the manager of the new objects, and get a list of temporary id
|
||||
# temporary ids -- indicates the position of objects in the tensor
|
||||
# (starts with 1 due to the background channel)
|
||||
corresponding_tmp_ids, _ = self.object_manager.add_new_objects(objects)
|
||||
|
||||
mask, _ = pad_divide_by(mask, 16)
|
||||
if need_segment:
|
||||
# merge predicted mask with the incomplete input mask
|
||||
pred_prob_no_bg = pred_prob_with_bg[1:]
|
||||
# use the mutual exclusivity of segmentation
|
||||
if idx_mask:
|
||||
pred_prob_no_bg[:, mask > 0] = 0
|
||||
else:
|
||||
pred_prob_no_bg[:, mask.max(0) > 0.5] = 0
|
||||
|
||||
new_masks = []
|
||||
for mask_id, tmp_id in enumerate(corresponding_tmp_ids):
|
||||
if idx_mask:
|
||||
this_mask = (mask == objects[mask_id]).type_as(pred_prob_no_bg)
|
||||
else:
|
||||
this_mask = mask[tmp_id]
|
||||
if tmp_id > pred_prob_no_bg.shape[0]:
|
||||
new_masks.append(this_mask.unsqueeze(0))
|
||||
else:
|
||||
# +1 for padding the background channel
|
||||
pred_prob_no_bg[tmp_id - 1] = this_mask
|
||||
# new_masks are always in the order of tmp_id
|
||||
mask = torch.cat([pred_prob_no_bg, *new_masks], dim=0)
|
||||
elif idx_mask:
|
||||
# simply convert cls to one-hot representation
|
||||
if len(objects) == 0:
|
||||
if delete_buffer:
|
||||
self.image_feature_store.delete(self.curr_ti)
|
||||
log.warn('Trying to insert an empty mask as memory!')
|
||||
return torch.zeros((1, key.shape[-2] * 16, key.shape[-1] * 16),
|
||||
device=key.device,
|
||||
dtype=key.dtype)
|
||||
mask = torch.stack(
|
||||
[mask == objects[mask_id] for mask_id, _ in enumerate(corresponding_tmp_ids)],
|
||||
dim=0)
|
||||
if matting:
|
||||
mask = mask.unsqueeze(0).float() / 255.
|
||||
pred_prob_with_bg = torch.cat([1-mask, mask], 0)
|
||||
else:
|
||||
pred_prob_with_bg = aggregate(mask, dim=0)
|
||||
pred_prob_with_bg = torch.softmax(pred_prob_with_bg, dim=0)
|
||||
|
||||
self.last_mask = pred_prob_with_bg[1:].unsqueeze(0)
|
||||
if self.flip_aug:
|
||||
self.last_mask = torch.cat(
|
||||
[self.last_mask, torch.flip(self.last_mask, dims=[-1])], dim=0)
|
||||
self.last_pix_feat = pix_feat
|
||||
|
||||
# save as memory if needed
|
||||
if is_mem_frame or force_permanent:
|
||||
# clear the memory for given mask and add the first predicted mask
|
||||
if first_frame_pred:
|
||||
self.clear_temp_mem()
|
||||
self._add_memory(image,
|
||||
pix_feat,
|
||||
self.last_mask,
|
||||
key,
|
||||
shrinkage,
|
||||
selection,
|
||||
force_permanent=force_permanent,
|
||||
is_deep_update=True)
|
||||
else: # compute self.last_msk_value for non-memory frame
|
||||
msk_value, _, _, _ = self.network.encode_mask(
|
||||
image,
|
||||
pix_feat,
|
||||
self.memory.get_sensory(self.object_manager.all_obj_ids),
|
||||
self.last_mask,
|
||||
deep_update=False,
|
||||
chunk_size=self.chunk_size,
|
||||
need_weights=self.save_aux)
|
||||
self.last_msk_value = msk_value
|
||||
|
||||
if delete_buffer:
|
||||
self.image_feature_store.delete(self.curr_ti)
|
||||
|
||||
output_prob = unpad(pred_prob_with_bg, self.pad)
|
||||
if resize_needed:
|
||||
# restore output to the original size
|
||||
output_prob = F.interpolate(output_prob.unsqueeze(0),
|
||||
size=(h, w),
|
||||
mode='bilinear',
|
||||
align_corners=False)[0]
|
||||
|
||||
return output_prob
|
||||
|
||||
def delete_objects(self, objects: List[int]) -> None:
|
||||
"""
|
||||
Delete the given objects from the memory.
|
||||
"""
|
||||
self.object_manager.delete_objects(objects)
|
||||
self.memory.purge_except(self.object_manager.all_obj_ids)
|
||||
|
||||
def output_prob_to_mask(self, output_prob: torch.Tensor, matting: bool = True) -> torch.Tensor:
|
||||
if matting:
|
||||
new_mask = output_prob[1:].squeeze(0)
|
||||
else:
|
||||
mask = torch.argmax(output_prob, dim=0)
|
||||
|
||||
# index in tensor != object id -- remap the ids here
|
||||
new_mask = torch.zeros_like(mask)
|
||||
for tmp_id, obj in self.object_manager.tmp_id_to_obj.items():
|
||||
new_mask[mask == tmp_id] = obj.id
|
||||
|
||||
return new_mask
|
||||
|
||||
@torch.inference_mode()
|
||||
@safe_autocast()
|
||||
def process_video(
|
||||
self,
|
||||
input_path: str,
|
||||
mask_path: str,
|
||||
output_path: str = None,
|
||||
n_warmup: int = 10,
|
||||
r_erode: int = 10,
|
||||
r_dilate: int = 10,
|
||||
suffix: str = "",
|
||||
save_image: bool = False,
|
||||
max_size: int = -1,
|
||||
) -> Tuple:
|
||||
"""
|
||||
Process a video for object segmentation and matting.
|
||||
This method processes a video file by performing object segmentation and matting on each frame.
|
||||
It supports warmup frames, mask erosion/dilation, and various output options.
|
||||
Args:
|
||||
input_path (str): Path to the input video file
|
||||
mask_path (str): Path to the mask image file used for initial segmentation
|
||||
output_path (str, optional): Directory path where output files will be saved. Defaults to a temporary directory
|
||||
n_warmup (int, optional): Number of warmup frames to use. Defaults to 10
|
||||
r_erode (int, optional): Erosion radius for mask processing. Defaults to 10
|
||||
r_dilate (int, optional): Dilation radius for mask processing. Defaults to 10
|
||||
suffix (str, optional): Suffix to append to output filename. Defaults to ""
|
||||
save_image (bool, optional): Whether to save individual frames. Defaults to False
|
||||
max_size (int, optional): Maximum size for frame dimension. Use -1 for no limit. Defaults to -1
|
||||
Returns:
|
||||
Tuple[str, str]: A tuple containing:
|
||||
- Path to the output foreground video file (str)
|
||||
- Path to the output alpha matte video file (str)
|
||||
Output:
|
||||
- Saves processed video files with foreground (_fgr) and alpha matte (_pha)
|
||||
- If save_image=True, saves individual frames in separate directories
|
||||
"""
|
||||
output_path = output_path if output_path is not None else tempfile.TemporaryDirectory().name
|
||||
r_erode = int(r_erode)
|
||||
r_dilate = int(r_dilate)
|
||||
n_warmup = int(n_warmup)
|
||||
max_size = int(max_size)
|
||||
|
||||
vframes, fps, length, video_name = read_frame_from_videos(input_path)
|
||||
repeated_frames = vframes[0].unsqueeze(0).repeat(n_warmup, 1, 1, 1)
|
||||
vframes = torch.cat([repeated_frames, vframes], dim=0).float()
|
||||
length += n_warmup
|
||||
|
||||
new_h, new_w = vframes.shape[-2:]
|
||||
if max_size > 0:
|
||||
h, w = new_h, new_w
|
||||
min_side = min(h, w)
|
||||
if min_side > max_size:
|
||||
new_h = int(h / min_side * max_size)
|
||||
new_w = int(w / min_side * max_size)
|
||||
vframes = F.interpolate(vframes, size=(new_h, new_w), mode="area")
|
||||
|
||||
os.makedirs(output_path, exist_ok=True)
|
||||
if suffix:
|
||||
video_name = f"{video_name}_{suffix}"
|
||||
if save_image:
|
||||
os.makedirs(f"{output_path}/{video_name}", exist_ok=True)
|
||||
os.makedirs(f"{output_path}/{video_name}/pha", exist_ok=True)
|
||||
os.makedirs(f"{output_path}/{video_name}/fgr", exist_ok=True)
|
||||
|
||||
mask = np.array(Image.open(mask_path).convert("L"))
|
||||
if r_dilate > 0:
|
||||
mask = gen_dilate(mask, r_dilate, r_dilate)
|
||||
if r_erode > 0:
|
||||
mask = gen_erosion(mask, r_erode, r_erode)
|
||||
|
||||
mask = torch.from_numpy(mask).float().to(self.device)
|
||||
if max_size > 0:
|
||||
mask = F.interpolate(
|
||||
mask.unsqueeze(0).unsqueeze(0), size=(new_h, new_w), mode="nearest"
|
||||
)[0, 0]
|
||||
|
||||
bgr = (np.array([120, 255, 155], dtype=np.float32) / 255).reshape((1, 1, 3))
|
||||
objects = [1]
|
||||
|
||||
phas = []
|
||||
fgrs = []
|
||||
for ti in tqdm(range(length)):
|
||||
image = vframes[ti]
|
||||
image_np = np.array(image.permute(1, 2, 0))
|
||||
image = (image / 255.0).float().to(self.device)
|
||||
|
||||
if ti == 0:
|
||||
output_prob = self.step(image, mask, objects=objects)
|
||||
output_prob = self.step(image, first_frame_pred=True)
|
||||
else:
|
||||
if ti <= n_warmup:
|
||||
output_prob = self.step(image, first_frame_pred=True)
|
||||
else:
|
||||
output_prob = self.step(image)
|
||||
|
||||
mask = self.output_prob_to_mask(output_prob)
|
||||
pha = mask.unsqueeze(2).cpu().numpy()
|
||||
com_np = image_np / 255.0 * pha + bgr * (1 - pha)
|
||||
|
||||
if ti > (n_warmup - 1):
|
||||
com_np = (com_np * 255).astype(np.uint8)
|
||||
pha = (pha * 255).astype(np.uint8)
|
||||
fgrs.append(com_np)
|
||||
phas.append(pha)
|
||||
if save_image:
|
||||
cv2.imwrite(
|
||||
f"{output_path}/{video_name}/pha/{str(ti - n_warmup).zfill(5)}.png",
|
||||
pha,
|
||||
)
|
||||
cv2.imwrite(
|
||||
f"{output_path}/{video_name}/fgr/{str(ti - n_warmup).zfill(5)}.png",
|
||||
com_np[..., [2, 1, 0]],
|
||||
)
|
||||
|
||||
fgrs = np.array(fgrs)
|
||||
phas = np.array(phas)
|
||||
|
||||
fgr_filename = f"{output_path}/{video_name}_fgr.mp4"
|
||||
alpha_filename = f"{output_path}/{video_name}_pha.mp4"
|
||||
|
||||
imageio.mimwrite(fgr_filename, fgrs, fps=fps, quality=7)
|
||||
imageio.mimwrite(alpha_filename, phas, fps=fps, quality=7)
|
||||
|
||||
return (fgr_filename,alpha_filename)
|
||||
@@ -0,0 +1,348 @@
|
||||
from typing import Dict, List, Optional, Literal
|
||||
from collections import defaultdict
|
||||
import torch
|
||||
|
||||
|
||||
def _add_last_dim(dictionary, key, new_value, prepend=False):
|
||||
# append/prepend a new value to the last dimension of a tensor in a dictionary
|
||||
# if the key does not exist, put the new value in
|
||||
# append by default
|
||||
if key in dictionary:
|
||||
dictionary[key] = torch.cat([dictionary[key], new_value], -1)
|
||||
else:
|
||||
dictionary[key] = new_value
|
||||
|
||||
|
||||
class KeyValueMemoryStore:
|
||||
"""
|
||||
Works for key/value pairs type storage
|
||||
e.g., working and long-term memory
|
||||
"""
|
||||
def __init__(self, save_selection: bool = False, save_usage: bool = False):
|
||||
"""
|
||||
We store keys and values of objects that first appear in the same frame in a bucket.
|
||||
Each bucket contains a set of object ids.
|
||||
Each bucket is associated with a single key tensor
|
||||
and a dictionary of value tensors indexed by object id.
|
||||
|
||||
The keys and values are stored as the concatenation of a permanent part and a temporary part.
|
||||
"""
|
||||
self.save_selection = save_selection
|
||||
self.save_usage = save_usage
|
||||
|
||||
self.global_bucket_id = 0 # does not reduce even if buckets are removed
|
||||
self.buckets: Dict[int, List[int]] = {} # indexed by bucket id
|
||||
self.k: Dict[int, torch.Tensor] = {} # indexed by bucket id
|
||||
self.v: Dict[int, torch.Tensor] = {} # indexed by object id
|
||||
|
||||
# indexed by bucket id; the end point of permanent memory
|
||||
self.perm_end_pt: Dict[int, int] = defaultdict(int)
|
||||
|
||||
# shrinkage and selection are just like the keys
|
||||
self.s = {}
|
||||
if self.save_selection:
|
||||
self.e = {} # does not contain the permanent memory part
|
||||
|
||||
# usage
|
||||
if self.save_usage:
|
||||
self.use_cnt = {} # indexed by bucket id, does not contain the permanent memory part
|
||||
self.life_cnt = {} # indexed by bucket id, does not contain the permanent memory part
|
||||
|
||||
def add(self,
|
||||
key: torch.Tensor,
|
||||
values: Dict[int, torch.Tensor],
|
||||
shrinkage: torch.Tensor,
|
||||
selection: torch.Tensor,
|
||||
supposed_bucket_id: int = -1,
|
||||
as_permanent: Literal['no', 'first', 'all'] = 'no') -> None:
|
||||
"""
|
||||
key: (1/2)*C*N
|
||||
values: dict of values ((1/2)*C*N), object ids are used as keys
|
||||
shrinkage: (1/2)*1*N
|
||||
selection: (1/2)*C*N
|
||||
|
||||
supposed_bucket_id: used to sync the bucket id between working and long-term memory
|
||||
if provided, the input should all be in a single bucket indexed by this id
|
||||
as_permanent: whether to store the input as permanent memory
|
||||
'no': don't
|
||||
'first': only store it as permanent memory if the bucket is empty
|
||||
'all': always store it as permanent memory
|
||||
"""
|
||||
bs = key.shape[0]
|
||||
ne = key.shape[-1]
|
||||
assert len(key.shape) == 3
|
||||
assert len(shrinkage.shape) == 3
|
||||
assert not self.save_selection or len(selection.shape) == 3
|
||||
assert as_permanent in ['no', 'first', 'all']
|
||||
|
||||
# add the value and create new buckets if necessary
|
||||
if supposed_bucket_id >= 0:
|
||||
enabled_buckets = [supposed_bucket_id]
|
||||
bucket_exist = supposed_bucket_id in self.buckets
|
||||
for obj, value in values.items():
|
||||
if bucket_exist:
|
||||
assert obj in self.v
|
||||
assert obj in self.buckets[supposed_bucket_id]
|
||||
_add_last_dim(self.v, obj, value, prepend=(as_permanent == 'all'))
|
||||
else:
|
||||
assert obj not in self.v
|
||||
self.v[obj] = value
|
||||
self.buckets[supposed_bucket_id] = list(values.keys())
|
||||
else:
|
||||
new_bucket_id = None
|
||||
enabled_buckets = set()
|
||||
for obj, value in values.items():
|
||||
assert len(value.shape) == 3
|
||||
if obj in self.v:
|
||||
_add_last_dim(self.v, obj, value, prepend=(as_permanent == 'all'))
|
||||
bucket_used = [
|
||||
bucket_id for bucket_id, object_ids in self.buckets.items()
|
||||
if obj in object_ids
|
||||
]
|
||||
assert len(bucket_used) == 1 # each object should only be in one bucket
|
||||
enabled_buckets.add(bucket_used[0])
|
||||
else:
|
||||
self.v[obj] = value
|
||||
if new_bucket_id is None:
|
||||
# create new bucket
|
||||
new_bucket_id = self.global_bucket_id
|
||||
self.global_bucket_id += 1
|
||||
self.buckets[new_bucket_id] = []
|
||||
# put the new object into the corresponding bucket
|
||||
self.buckets[new_bucket_id].append(obj)
|
||||
enabled_buckets.add(new_bucket_id)
|
||||
|
||||
# increment the permanent size if necessary
|
||||
add_as_permanent = {} # indexed by bucket id
|
||||
for bucket_id in enabled_buckets:
|
||||
add_as_permanent[bucket_id] = False
|
||||
if as_permanent == 'all':
|
||||
self.perm_end_pt[bucket_id] += ne
|
||||
add_as_permanent[bucket_id] = True
|
||||
elif as_permanent == 'first':
|
||||
if self.perm_end_pt[bucket_id] == 0:
|
||||
self.perm_end_pt[bucket_id] = ne
|
||||
add_as_permanent[bucket_id] = True
|
||||
|
||||
# create new counters for usage if necessary
|
||||
if self.save_usage and as_permanent != 'all':
|
||||
new_count = torch.zeros((bs, ne), device=key.device, dtype=torch.float32)
|
||||
new_life = torch.zeros((bs, ne), device=key.device, dtype=torch.float32) + 1e-7
|
||||
|
||||
# add the key to every bucket
|
||||
for bucket_id in self.buckets:
|
||||
if bucket_id not in enabled_buckets:
|
||||
# if we are not adding new values to a bucket, we should skip it
|
||||
continue
|
||||
|
||||
_add_last_dim(self.k, bucket_id, key, prepend=add_as_permanent[bucket_id])
|
||||
_add_last_dim(self.s, bucket_id, shrinkage, prepend=add_as_permanent[bucket_id])
|
||||
if not add_as_permanent[bucket_id]:
|
||||
if self.save_selection:
|
||||
_add_last_dim(self.e, bucket_id, selection)
|
||||
if self.save_usage:
|
||||
_add_last_dim(self.use_cnt, bucket_id, new_count)
|
||||
_add_last_dim(self.life_cnt, bucket_id, new_life)
|
||||
|
||||
def update_bucket_usage(self, bucket_id: int, usage: torch.Tensor) -> None:
|
||||
# increase all life count by 1
|
||||
# increase use of indexed elements
|
||||
if not self.save_usage:
|
||||
return
|
||||
|
||||
usage = usage[:, self.perm_end_pt[bucket_id]:]
|
||||
if usage.shape[-1] == 0:
|
||||
# if there is no temporary memory, we don't need to update
|
||||
return
|
||||
self.use_cnt[bucket_id] += usage.view_as(self.use_cnt[bucket_id])
|
||||
self.life_cnt[bucket_id] += 1
|
||||
|
||||
def sieve_by_range(self, bucket_id: int, start: int, end: int, min_size: int) -> None:
|
||||
# keep only the temporary elements *outside* of this range (with some boundary conditions)
|
||||
# the permanent elements are ignored in this computation
|
||||
# i.e., concat (a[:start], a[end:])
|
||||
# bucket with size <= min_size are not modified
|
||||
|
||||
assert start >= 0
|
||||
assert end <= 0
|
||||
|
||||
object_ids = self.buckets[bucket_id]
|
||||
bucket_num_elements = self.k[bucket_id].shape[-1] - self.perm_end_pt[bucket_id]
|
||||
if bucket_num_elements <= min_size:
|
||||
return
|
||||
|
||||
if end == 0:
|
||||
# negative 0 would not work as the end index!
|
||||
# effectively make the second part an empty slice
|
||||
end = self.k[bucket_id].shape[-1] + 1
|
||||
|
||||
p_size = self.perm_end_pt[bucket_id]
|
||||
start = start + p_size
|
||||
|
||||
k = self.k[bucket_id]
|
||||
s = self.s[bucket_id]
|
||||
if self.save_selection:
|
||||
e = self.e[bucket_id]
|
||||
if self.save_usage:
|
||||
use_cnt = self.use_cnt[bucket_id]
|
||||
life_cnt = self.life_cnt[bucket_id]
|
||||
|
||||
self.k[bucket_id] = torch.cat([k[:, :, :start], k[:, :, end:]], -1)
|
||||
self.s[bucket_id] = torch.cat([s[:, :, :start], s[:, :, end:]], -1)
|
||||
if self.save_selection:
|
||||
self.e[bucket_id] = torch.cat([e[:, :, :start - p_size], e[:, :, end:]], -1)
|
||||
if self.save_usage:
|
||||
self.use_cnt[bucket_id] = torch.cat([use_cnt[:, :start - p_size], use_cnt[:, end:]], -1)
|
||||
self.life_cnt[bucket_id] = torch.cat([life_cnt[:, :start - p_size], life_cnt[:, end:]],
|
||||
-1)
|
||||
for obj_id in object_ids:
|
||||
v = self.v[obj_id]
|
||||
self.v[obj_id] = torch.cat([v[:, :, :start], v[:, :, end:]], -1)
|
||||
|
||||
def remove_old_memory(self, bucket_id: int, max_len: int) -> None:
|
||||
self.sieve_by_range(bucket_id, 0, -max_len, max_len)
|
||||
|
||||
def remove_obsolete_features(self, bucket_id: int, max_size: int) -> None:
|
||||
# for long-term memory only
|
||||
object_ids = self.buckets[bucket_id]
|
||||
|
||||
assert self.perm_end_pt[bucket_id] == 0 # permanent memory should be empty in LT memory
|
||||
|
||||
# normalize with life duration
|
||||
usage = self.get_usage(bucket_id)
|
||||
bs = usage.shape[0]
|
||||
|
||||
survivals = []
|
||||
|
||||
for bi in range(bs):
|
||||
_, survived = torch.topk(usage[bi], k=max_size)
|
||||
survivals.append(survived.flatten())
|
||||
assert survived.shape[-1] == survivals[0].shape[-1]
|
||||
|
||||
self.k[bucket_id] = torch.stack(
|
||||
[self.k[bucket_id][bi, :, survived] for bi, survived in enumerate(survivals)], 0)
|
||||
self.s[bucket_id] = torch.stack(
|
||||
[self.s[bucket_id][bi, :, survived] for bi, survived in enumerate(survivals)], 0)
|
||||
|
||||
if self.save_selection:
|
||||
# Long-term memory does not store selection so this should not be needed
|
||||
self.e[bucket_id] = torch.stack(
|
||||
[self.e[bucket_id][bi, :, survived] for bi, survived in enumerate(survivals)], 0)
|
||||
for obj_id in object_ids:
|
||||
self.v[obj_id] = torch.stack(
|
||||
[self.v[obj_id][bi, :, survived] for bi, survived in enumerate(survivals)], 0)
|
||||
|
||||
self.use_cnt[bucket_id] = torch.stack(
|
||||
[self.use_cnt[bucket_id][bi, survived] for bi, survived in enumerate(survivals)], 0)
|
||||
self.life_cnt[bucket_id] = torch.stack(
|
||||
[self.life_cnt[bucket_id][bi, survived] for bi, survived in enumerate(survivals)], 0)
|
||||
|
||||
def get_usage(self, bucket_id: int) -> torch.Tensor:
|
||||
# return normalized usage
|
||||
if not self.save_usage:
|
||||
raise RuntimeError('I did not count usage!')
|
||||
else:
|
||||
usage = self.use_cnt[bucket_id] / self.life_cnt[bucket_id]
|
||||
return usage
|
||||
|
||||
def get_all_sliced(
|
||||
self, bucket_id: int, start: int, end: int
|
||||
) -> (torch.Tensor, torch.Tensor, torch.Tensor, Dict[int, torch.Tensor], torch.Tensor):
|
||||
# return k, sk, ek, value, normalized usage in order, sliced by start and end
|
||||
# this only queries the temporary memory
|
||||
|
||||
assert start >= 0
|
||||
assert end <= 0
|
||||
|
||||
p_size = self.perm_end_pt[bucket_id]
|
||||
start = start + p_size
|
||||
|
||||
if end == 0:
|
||||
# negative 0 would not work as the end index!
|
||||
k = self.k[bucket_id][:, :, start:]
|
||||
sk = self.s[bucket_id][:, :, start:]
|
||||
ek = self.e[bucket_id][:, :, start - p_size:] if self.save_selection else None
|
||||
value = {obj_id: self.v[obj_id][:, :, start:] for obj_id in self.buckets[bucket_id]}
|
||||
usage = self.get_usage(bucket_id)[:, start - p_size:] if self.save_usage else None
|
||||
else:
|
||||
k = self.k[bucket_id][:, :, start:end]
|
||||
sk = self.s[bucket_id][:, :, start:end]
|
||||
ek = self.e[bucket_id][:, :, start - p_size:end] if self.save_selection else None
|
||||
value = {obj_id: self.v[obj_id][:, :, start:end] for obj_id in self.buckets[bucket_id]}
|
||||
usage = self.get_usage(bucket_id)[:, start - p_size:end] if self.save_usage else None
|
||||
|
||||
return k, sk, ek, value, usage
|
||||
|
||||
def purge_except(self, obj_keep_idx: List[int]):
|
||||
# purge certain objects from the memory except the one listed
|
||||
obj_keep_idx = set(obj_keep_idx)
|
||||
|
||||
# remove objects that are not in the keep list from the buckets
|
||||
buckets_to_remove = []
|
||||
for bucket_id, object_ids in self.buckets.items():
|
||||
self.buckets[bucket_id] = [obj_id for obj_id in object_ids if obj_id in obj_keep_idx]
|
||||
if len(self.buckets[bucket_id]) == 0:
|
||||
buckets_to_remove.append(bucket_id)
|
||||
|
||||
# remove object values that are not in the keep list
|
||||
self.v = {k: v for k, v in self.v.items() if k in obj_keep_idx}
|
||||
|
||||
# remove buckets that are empty
|
||||
for bucket_id in buckets_to_remove:
|
||||
del self.buckets[bucket_id]
|
||||
del self.k[bucket_id]
|
||||
del self.s[bucket_id]
|
||||
if self.save_selection:
|
||||
del self.e[bucket_id]
|
||||
if self.save_usage:
|
||||
del self.use_cnt[bucket_id]
|
||||
del self.life_cnt[bucket_id]
|
||||
|
||||
def clear_non_permanent_memory(self):
|
||||
# clear all non-permanent memory
|
||||
for bucket_id in self.buckets:
|
||||
self.sieve_by_range(bucket_id, 0, 0, 0)
|
||||
|
||||
def get_v_size(self, obj_id: int) -> int:
|
||||
return self.v[obj_id].shape[-1]
|
||||
|
||||
def size(self, bucket_id: int) -> int:
|
||||
if bucket_id not in self.k:
|
||||
return 0
|
||||
else:
|
||||
return self.k[bucket_id].shape[-1]
|
||||
|
||||
def perm_size(self, bucket_id: int) -> int:
|
||||
return self.perm_end_pt[bucket_id]
|
||||
|
||||
def non_perm_size(self, bucket_id: int) -> int:
|
||||
return self.size(bucket_id) - self.perm_size(bucket_id)
|
||||
|
||||
def engaged(self, bucket_id: Optional[int] = None) -> bool:
|
||||
if bucket_id is None:
|
||||
return len(self.buckets) > 0
|
||||
else:
|
||||
return bucket_id in self.buckets
|
||||
|
||||
@property
|
||||
def num_objects(self) -> int:
|
||||
return len(self.v)
|
||||
|
||||
@property
|
||||
def key(self) -> Dict[int, torch.Tensor]:
|
||||
return self.k
|
||||
|
||||
@property
|
||||
def value(self) -> Dict[int, torch.Tensor]:
|
||||
return self.v
|
||||
|
||||
@property
|
||||
def shrinkage(self) -> Dict[int, torch.Tensor]:
|
||||
return self.s
|
||||
|
||||
@property
|
||||
def selection(self) -> Dict[int, torch.Tensor]:
|
||||
return self.e
|
||||
|
||||
def __contains__(self, key):
|
||||
return key in self.v
|
||||
@@ -0,0 +1,453 @@
|
||||
import logging
|
||||
from omegaconf import DictConfig
|
||||
from typing import List, Dict
|
||||
import torch
|
||||
|
||||
from matanyone2.inference.object_manager import ObjectManager
|
||||
from matanyone2.inference.kv_memory_store import KeyValueMemoryStore
|
||||
from matanyone2.model.matanyone2 import MatAnyone2
|
||||
from matanyone2.model.utils.memory_utils import get_similarity, do_softmax
|
||||
|
||||
log = logging.getLogger()
|
||||
|
||||
|
||||
class MemoryManager:
|
||||
"""
|
||||
Manages all three memory stores and the transition between working/long-term memory
|
||||
"""
|
||||
def __init__(self, cfg: DictConfig, object_manager: ObjectManager):
|
||||
self.object_manager = object_manager
|
||||
self.sensory_dim = cfg.model.sensory_dim
|
||||
self.top_k = cfg.top_k
|
||||
self.chunk_size = cfg.chunk_size
|
||||
|
||||
self.save_aux = cfg.save_aux
|
||||
|
||||
self.use_long_term = cfg.use_long_term
|
||||
self.count_long_term_usage = cfg.long_term.count_usage
|
||||
# subtract 1 because the first-frame is now counted as "permanent memory"
|
||||
# and is not counted towards max_mem_frames
|
||||
# but we want to keep the hyperparameters consistent as before for the same behavior
|
||||
if self.use_long_term:
|
||||
self.max_mem_frames = cfg.long_term.max_mem_frames - 1
|
||||
self.min_mem_frames = cfg.long_term.min_mem_frames - 1
|
||||
self.num_prototypes = cfg.long_term.num_prototypes
|
||||
self.max_long_tokens = cfg.long_term.max_num_tokens
|
||||
self.buffer_tokens = cfg.long_term.buffer_tokens
|
||||
else:
|
||||
self.max_mem_frames = cfg.max_mem_frames - 1
|
||||
|
||||
# dimensions will be inferred from input later
|
||||
self.CK = self.CV = None
|
||||
self.H = self.W = None
|
||||
|
||||
# The sensory memory is stored as a dictionary indexed by object ids
|
||||
# each of shape bs * C^h * H * W
|
||||
self.sensory = {}
|
||||
|
||||
# a dictionary indexed by object ids, each of shape bs * T * Q * C
|
||||
self.obj_v = {}
|
||||
|
||||
self.work_mem = KeyValueMemoryStore(save_selection=self.use_long_term,
|
||||
save_usage=self.use_long_term)
|
||||
if self.use_long_term:
|
||||
self.long_mem = KeyValueMemoryStore(save_usage=self.count_long_term_usage)
|
||||
|
||||
self.config_stale = True
|
||||
self.engaged = False
|
||||
|
||||
def update_config(self, cfg: DictConfig) -> None:
|
||||
self.config_stale = True
|
||||
self.top_k = cfg['top_k']
|
||||
|
||||
assert self.use_long_term == cfg.use_long_term, 'cannot update this'
|
||||
assert self.count_long_term_usage == cfg.long_term.count_usage, 'cannot update this'
|
||||
|
||||
self.use_long_term = cfg.use_long_term
|
||||
self.count_long_term_usage = cfg.long_term.count_usage
|
||||
if self.use_long_term:
|
||||
self.max_mem_frames = cfg.long_term.max_mem_frames - 1
|
||||
self.min_mem_frames = cfg.long_term.min_mem_frames - 1
|
||||
self.num_prototypes = cfg.long_term.num_prototypes
|
||||
self.max_long_tokens = cfg.long_term.max_num_tokens
|
||||
self.buffer_tokens = cfg.long_term.buffer_tokens
|
||||
else:
|
||||
self.max_mem_frames = cfg.max_mem_frames - 1
|
||||
|
||||
def _readout(self, affinity, v, uncert_mask=None) -> torch.Tensor:
|
||||
# affinity: bs*N*HW
|
||||
# v: bs*C*N or bs*num_objects*C*N
|
||||
# returns bs*C*HW or bs*num_objects*C*HW
|
||||
if len(v.shape) == 3:
|
||||
# single object
|
||||
if uncert_mask is not None:
|
||||
return v @ affinity * uncert_mask
|
||||
else:
|
||||
return v @ affinity
|
||||
else:
|
||||
bs, num_objects, C, N = v.shape
|
||||
v = v.view(bs, num_objects * C, N)
|
||||
out = v @ affinity
|
||||
if uncert_mask is not None:
|
||||
uncert_mask = uncert_mask.flatten(start_dim=2).expand(-1, C, -1)
|
||||
out = out * uncert_mask
|
||||
return out.view(bs, num_objects, C, -1)
|
||||
|
||||
def _get_mask_by_ids(self, mask: torch.Tensor, obj_ids: List[int]) -> torch.Tensor:
|
||||
# -1 because the mask does not contain the background channel
|
||||
return mask[:, [self.object_manager.find_tmp_by_id(obj) - 1 for obj in obj_ids]]
|
||||
|
||||
def _get_sensory_by_ids(self, obj_ids: List[int]) -> torch.Tensor:
|
||||
return torch.stack([self.sensory[obj] for obj in obj_ids], dim=1)
|
||||
|
||||
def _get_object_mem_by_ids(self, obj_ids: List[int]) -> torch.Tensor:
|
||||
return torch.stack([self.obj_v[obj] for obj in obj_ids], dim=1)
|
||||
|
||||
def _get_visual_values_by_ids(self, obj_ids: List[int]) -> torch.Tensor:
|
||||
# All the values that the object ids refer to should have the same shape
|
||||
value = torch.stack([self.work_mem.value[obj] for obj in obj_ids], dim=1)
|
||||
if self.use_long_term and obj_ids[0] in self.long_mem.value:
|
||||
lt_value = torch.stack([self.long_mem.value[obj] for obj in obj_ids], dim=1)
|
||||
value = torch.cat([lt_value, value], dim=-1)
|
||||
|
||||
return value
|
||||
|
||||
def read_first_frame(self, last_msk_value, pix_feat: torch.Tensor,
|
||||
last_mask: torch.Tensor, network: MatAnyone2, uncert_output=None) -> Dict[int, torch.Tensor]:
|
||||
"""
|
||||
Read from all memory stores and returns a single memory readout tensor for each object
|
||||
|
||||
pix_feat: (1/2) x C x H x W
|
||||
query_key: (1/2) x C^k x H x W
|
||||
selection: (1/2) x C^k x H x W
|
||||
last_mask: (1/2) x num_objects x H x W (at stride 16)
|
||||
return a dict of memory readouts, indexed by object indices. Each readout is C*H*W
|
||||
"""
|
||||
h, w = pix_feat.shape[-2:]
|
||||
bs = pix_feat.shape[0]
|
||||
assert last_mask.shape[0] == bs
|
||||
|
||||
"""
|
||||
Compute affinity and perform readout
|
||||
"""
|
||||
all_readout_mem = {}
|
||||
buckets = self.work_mem.buckets
|
||||
for bucket_id, bucket in buckets.items():
|
||||
|
||||
if self.chunk_size < 1:
|
||||
object_chunks = [bucket]
|
||||
else:
|
||||
object_chunks = [
|
||||
bucket[i:i + self.chunk_size] for i in range(0, len(bucket), self.chunk_size)
|
||||
]
|
||||
|
||||
for objects in object_chunks:
|
||||
this_sensory = self._get_sensory_by_ids(objects)
|
||||
this_last_mask = self._get_mask_by_ids(last_mask, objects)
|
||||
this_msk_value = self._get_visual_values_by_ids(objects) # (1/2)*num_objects*C*N
|
||||
pixel_readout = network.pixel_fusion(pix_feat, last_msk_value, this_sensory,
|
||||
this_last_mask)
|
||||
this_obj_mem = self._get_object_mem_by_ids(objects).unsqueeze(2)
|
||||
readout_memory, aux_features = network.readout_query(pixel_readout, this_obj_mem)
|
||||
for i, obj in enumerate(objects):
|
||||
all_readout_mem[obj] = readout_memory[:, i]
|
||||
|
||||
if self.save_aux:
|
||||
aux_output = {
|
||||
# 'sensory': this_sensory,
|
||||
# 'pixel_readout': pixel_readout,
|
||||
'q_logits': aux_features['logits'] if aux_features else None,
|
||||
# 'q_weights': aux_features['q_weights'] if aux_features else None,
|
||||
# 'p_weights': aux_features['p_weights'] if aux_features else None,
|
||||
# 'attn_mask': aux_features['attn_mask'].float() if aux_features else None,
|
||||
}
|
||||
self.aux = aux_output
|
||||
|
||||
return all_readout_mem
|
||||
|
||||
def read(self, pix_feat: torch.Tensor, query_key: torch.Tensor, selection: torch.Tensor,
|
||||
last_mask: torch.Tensor, network: MatAnyone2, uncert_output=None, last_msk_value=None, ti=None,
|
||||
last_pix_feat=None, last_pred_mask=None) -> Dict[int, torch.Tensor]:
|
||||
"""
|
||||
Read from all memory stores and returns a single memory readout tensor for each object
|
||||
|
||||
pix_feat: (1/2) x C x H x W
|
||||
query_key: (1/2) x C^k x H x W
|
||||
selection: (1/2) x C^k x H x W
|
||||
last_mask: (1/2) x num_objects x H x W (at stride 16)
|
||||
return a dict of memory readouts, indexed by object indices. Each readout is C*H*W
|
||||
"""
|
||||
h, w = pix_feat.shape[-2:]
|
||||
bs = pix_feat.shape[0]
|
||||
assert query_key.shape[0] == bs
|
||||
assert selection.shape[0] == bs
|
||||
assert last_mask.shape[0] == bs
|
||||
|
||||
uncert_mask = uncert_output["mask"] if uncert_output is not None else None
|
||||
|
||||
query_key = query_key.flatten(start_dim=2) # bs*C^k*HW
|
||||
selection = selection.flatten(start_dim=2) # bs*C^k*HW
|
||||
"""
|
||||
Compute affinity and perform readout
|
||||
"""
|
||||
all_readout_mem = {}
|
||||
buckets = self.work_mem.buckets
|
||||
for bucket_id, bucket in buckets.items():
|
||||
if self.use_long_term and self.long_mem.engaged(bucket_id):
|
||||
# Use long-term memory
|
||||
long_mem_size = self.long_mem.size(bucket_id)
|
||||
memory_key = torch.cat([self.long_mem.key[bucket_id], self.work_mem.key[bucket_id]],
|
||||
-1)
|
||||
shrinkage = torch.cat(
|
||||
[self.long_mem.shrinkage[bucket_id], self.work_mem.shrinkage[bucket_id]], -1)
|
||||
|
||||
similarity = get_similarity(memory_key, shrinkage, query_key, selection)
|
||||
affinity, usage = do_softmax(similarity,
|
||||
top_k=self.top_k,
|
||||
inplace=True,
|
||||
return_usage=True)
|
||||
"""
|
||||
Record memory usage for working and long-term memory
|
||||
"""
|
||||
# ignore the index return for long-term memory
|
||||
work_usage = usage[:, long_mem_size:]
|
||||
self.work_mem.update_bucket_usage(bucket_id, work_usage)
|
||||
|
||||
if self.count_long_term_usage:
|
||||
# ignore the index return for working memory
|
||||
long_usage = usage[:, :long_mem_size]
|
||||
self.long_mem.update_bucket_usage(bucket_id, long_usage)
|
||||
else:
|
||||
# no long-term memory
|
||||
memory_key = self.work_mem.key[bucket_id]
|
||||
shrinkage = self.work_mem.shrinkage[bucket_id]
|
||||
similarity = get_similarity(memory_key, shrinkage, query_key, selection, uncert_mask=uncert_mask)
|
||||
|
||||
if self.use_long_term:
|
||||
affinity, usage = do_softmax(similarity,
|
||||
top_k=self.top_k,
|
||||
inplace=True,
|
||||
return_usage=True)
|
||||
self.work_mem.update_bucket_usage(bucket_id, usage)
|
||||
else:
|
||||
affinity = do_softmax(similarity, top_k=self.top_k, inplace=True)
|
||||
|
||||
if self.chunk_size < 1:
|
||||
object_chunks = [bucket]
|
||||
else:
|
||||
object_chunks = [
|
||||
bucket[i:i + self.chunk_size] for i in range(0, len(bucket), self.chunk_size)
|
||||
]
|
||||
|
||||
for objects in object_chunks:
|
||||
this_sensory = self._get_sensory_by_ids(objects)
|
||||
this_last_mask = self._get_mask_by_ids(last_mask, objects)
|
||||
this_msk_value = self._get_visual_values_by_ids(objects) # (1/2)*num_objects*C*N
|
||||
visual_readout = self._readout(affinity,
|
||||
this_msk_value, uncert_mask).view(bs, len(objects), self.CV, h, w)
|
||||
|
||||
uncert_output = network.pred_uncertainty(last_pix_feat, pix_feat, last_pred_mask, visual_readout[:,0]-last_msk_value[:,0])
|
||||
|
||||
if uncert_output is not None:
|
||||
uncert_prob = uncert_output["prob"].unsqueeze(1) # b n 1 h w
|
||||
visual_readout = visual_readout*uncert_prob + last_msk_value*(1-uncert_prob)
|
||||
|
||||
pixel_readout = network.pixel_fusion(pix_feat, visual_readout, this_sensory,
|
||||
this_last_mask)
|
||||
this_obj_mem = self._get_object_mem_by_ids(objects).unsqueeze(2)
|
||||
readout_memory, aux_features = network.readout_query(pixel_readout, this_obj_mem)
|
||||
for i, obj in enumerate(objects):
|
||||
all_readout_mem[obj] = readout_memory[:, i]
|
||||
|
||||
if self.save_aux:
|
||||
aux_output = {
|
||||
# 'sensory': this_sensory,
|
||||
# 'pixel_readout': pixel_readout,
|
||||
'q_logits': aux_features['logits'] if aux_features else None,
|
||||
# 'q_weights': aux_features['q_weights'] if aux_features else None,
|
||||
# 'p_weights': aux_features['p_weights'] if aux_features else None,
|
||||
# 'attn_mask': aux_features['attn_mask'].float() if aux_features else None,
|
||||
}
|
||||
self.aux = aux_output
|
||||
|
||||
return all_readout_mem
|
||||
|
||||
def add_memory(self,
|
||||
key: torch.Tensor,
|
||||
shrinkage: torch.Tensor,
|
||||
msk_value: torch.Tensor,
|
||||
obj_value: torch.Tensor,
|
||||
objects: List[int],
|
||||
selection: torch.Tensor = None,
|
||||
*,
|
||||
as_permanent: bool = False) -> None:
|
||||
# key: (1/2)*C*H*W
|
||||
# msk_value: (1/2)*num_objects*C*H*W
|
||||
# obj_value: (1/2)*num_objects*Q*C
|
||||
# objects contains a list of object ids corresponding to the objects in msk_value/obj_value
|
||||
bs = key.shape[0]
|
||||
assert shrinkage.shape[0] == bs
|
||||
assert msk_value.shape[0] == bs
|
||||
assert obj_value.shape[0] == bs
|
||||
|
||||
self.engaged = True
|
||||
if self.H is None or self.config_stale:
|
||||
self.config_stale = False
|
||||
self.H, self.W = msk_value.shape[-2:]
|
||||
self.HW = self.H * self.W
|
||||
# convert from num. frames to num. tokens
|
||||
self.max_work_tokens = self.max_mem_frames * self.HW
|
||||
if self.use_long_term:
|
||||
self.min_work_tokens = self.min_mem_frames * self.HW
|
||||
|
||||
# key: bs*C*N
|
||||
# value: bs*num_objects*C*N
|
||||
key = key.flatten(start_dim=2)
|
||||
shrinkage = shrinkage.flatten(start_dim=2)
|
||||
self.CK = key.shape[1]
|
||||
|
||||
msk_value = msk_value.flatten(start_dim=3)
|
||||
self.CV = msk_value.shape[2]
|
||||
|
||||
if selection is not None:
|
||||
# not used in non-long-term mode
|
||||
selection = selection.flatten(start_dim=2)
|
||||
|
||||
# insert object values into object memory
|
||||
for obj_id, obj in enumerate(objects):
|
||||
if obj in self.obj_v:
|
||||
"""streaming average
|
||||
each self.obj_v[obj] is (1/2)*num_summaries*(embed_dim+1)
|
||||
first embed_dim keeps track of the sum of embeddings
|
||||
the last dim keeps the total count
|
||||
averaging in done inside the object transformer
|
||||
|
||||
incoming obj_value is (1/2)*num_objects*num_summaries*(embed_dim+1)
|
||||
self.obj_v[obj] = torch.cat([self.obj_v[obj], obj_value[:, obj_id]], dim=0)
|
||||
"""
|
||||
last_acc = self.obj_v[obj][:, :, -1]
|
||||
new_acc = last_acc + obj_value[:, obj_id, :, -1]
|
||||
|
||||
self.obj_v[obj][:, :, :-1] = (self.obj_v[obj][:, :, :-1] +
|
||||
obj_value[:, obj_id, :, :-1])
|
||||
self.obj_v[obj][:, :, -1] = new_acc
|
||||
else:
|
||||
self.obj_v[obj] = obj_value[:, obj_id]
|
||||
|
||||
# convert mask value tensor into a dict for insertion
|
||||
msk_values = {obj: msk_value[:, obj_id] for obj_id, obj in enumerate(objects)}
|
||||
self.work_mem.add(key,
|
||||
msk_values,
|
||||
shrinkage,
|
||||
selection=selection,
|
||||
as_permanent=as_permanent)
|
||||
|
||||
for bucket_id in self.work_mem.buckets.keys():
|
||||
# long-term memory cleanup
|
||||
if self.use_long_term:
|
||||
# Do memory compressed if needed
|
||||
if self.work_mem.non_perm_size(bucket_id) >= self.max_work_tokens:
|
||||
# Remove obsolete features if needed
|
||||
if self.long_mem.non_perm_size(bucket_id) >= (self.max_long_tokens -
|
||||
self.num_prototypes):
|
||||
self.long_mem.remove_obsolete_features(
|
||||
bucket_id,
|
||||
self.max_long_tokens - self.num_prototypes - self.buffer_tokens)
|
||||
|
||||
self.compress_features(bucket_id)
|
||||
else:
|
||||
# FIFO
|
||||
self.work_mem.remove_old_memory(bucket_id, self.max_work_tokens)
|
||||
|
||||
def purge_except(self, obj_keep_idx: List[int]) -> None:
|
||||
# purge certain objects from the memory except the one listed
|
||||
self.work_mem.purge_except(obj_keep_idx)
|
||||
if self.use_long_term and self.long_mem.engaged():
|
||||
self.long_mem.purge_except(obj_keep_idx)
|
||||
self.sensory = {k: v for k, v in self.sensory.items() if k in obj_keep_idx}
|
||||
|
||||
if not self.work_mem.engaged():
|
||||
# everything is removed!
|
||||
self.engaged = False
|
||||
|
||||
def compress_features(self, bucket_id: int) -> None:
|
||||
|
||||
# perform memory consolidation
|
||||
prototype_key, prototype_value, prototype_shrinkage = self.consolidation(
|
||||
*self.work_mem.get_all_sliced(bucket_id, 0, -self.min_work_tokens))
|
||||
|
||||
# remove consolidated working memory
|
||||
self.work_mem.sieve_by_range(bucket_id,
|
||||
0,
|
||||
-self.min_work_tokens,
|
||||
min_size=self.min_work_tokens)
|
||||
|
||||
# add to long-term memory
|
||||
self.long_mem.add(prototype_key,
|
||||
prototype_value,
|
||||
prototype_shrinkage,
|
||||
selection=None,
|
||||
supposed_bucket_id=bucket_id)
|
||||
|
||||
def consolidation(self, candidate_key: torch.Tensor, candidate_shrinkage: torch.Tensor,
|
||||
candidate_selection: torch.Tensor, candidate_value: Dict[int, torch.Tensor],
|
||||
usage: torch.Tensor) -> (torch.Tensor, Dict[int, torch.Tensor], torch.Tensor):
|
||||
# find the indices with max usage
|
||||
bs = candidate_key.shape[0]
|
||||
assert bs in [1, 2]
|
||||
|
||||
prototype_key = []
|
||||
prototype_selection = []
|
||||
for bi in range(bs):
|
||||
_, max_usage_indices = torch.topk(usage[bi], k=self.num_prototypes, dim=-1, sorted=True)
|
||||
prototype_indices = max_usage_indices.flatten()
|
||||
prototype_key.append(candidate_key[bi, :, prototype_indices])
|
||||
prototype_selection.append(candidate_selection[bi, :, prototype_indices])
|
||||
prototype_key = torch.stack(prototype_key, dim=0)
|
||||
prototype_selection = torch.stack(prototype_selection, dim=0)
|
||||
"""
|
||||
Potentiation step
|
||||
"""
|
||||
similarity = get_similarity(candidate_key, candidate_shrinkage, prototype_key,
|
||||
prototype_selection)
|
||||
affinity = do_softmax(similarity)
|
||||
|
||||
# readout the values
|
||||
prototype_value = {k: self._readout(affinity, v) for k, v in candidate_value.items()}
|
||||
|
||||
# readout the shrinkage term
|
||||
prototype_shrinkage = self._readout(affinity, candidate_shrinkage)
|
||||
|
||||
return prototype_key, prototype_value, prototype_shrinkage
|
||||
|
||||
def initialize_sensory_if_needed(self, sample_key: torch.Tensor, ids: List[int]):
|
||||
for obj in ids:
|
||||
if obj not in self.sensory:
|
||||
# also initializes the sensory memory
|
||||
bs, _, h, w = sample_key.shape
|
||||
self.sensory[obj] = torch.zeros((bs, self.sensory_dim, h, w),
|
||||
device=sample_key.device)
|
||||
|
||||
def update_sensory(self, sensory: torch.Tensor, ids: List[int]):
|
||||
# sensory: 1*num_objects*C*H*W
|
||||
for obj_id, obj in enumerate(ids):
|
||||
self.sensory[obj] = sensory[:, obj_id]
|
||||
|
||||
def get_sensory(self, ids: List[int]):
|
||||
# returns (1/2)*num_objects*C*H*W
|
||||
return self._get_sensory_by_ids(ids)
|
||||
|
||||
def clear_non_permanent_memory(self):
|
||||
self.work_mem.clear_non_permanent_memory()
|
||||
if self.use_long_term:
|
||||
self.long_mem.clear_non_permanent_memory()
|
||||
|
||||
def clear_sensory_memory(self):
|
||||
self.sensory = {}
|
||||
|
||||
def clear_work_mem(self):
|
||||
self.work_mem = KeyValueMemoryStore(save_selection=self.use_long_term,
|
||||
save_usage=self.use_long_term)
|
||||
|
||||
def clear_obj_mem(self):
|
||||
self.obj_v = {}
|
||||
@@ -0,0 +1,24 @@
|
||||
class ObjectInfo:
|
||||
"""
|
||||
Store meta information for an object
|
||||
"""
|
||||
def __init__(self, id: int):
|
||||
self.id = id
|
||||
self.poke_count = 0 # count number of detections missed
|
||||
|
||||
def poke(self) -> None:
|
||||
self.poke_count += 1
|
||||
|
||||
def unpoke(self) -> None:
|
||||
self.poke_count = 0
|
||||
|
||||
def __hash__(self):
|
||||
return hash(self.id)
|
||||
|
||||
def __eq__(self, other):
|
||||
if type(other) == int:
|
||||
return self.id == other
|
||||
return self.id == other.id
|
||||
|
||||
def __repr__(self):
|
||||
return f'(ID: {self.id})'
|
||||
@@ -0,0 +1,149 @@
|
||||
from typing import Union, List, Dict
|
||||
|
||||
import torch
|
||||
from matanyone2.inference.object_info import ObjectInfo
|
||||
|
||||
|
||||
class ObjectManager:
|
||||
"""
|
||||
Object IDs are immutable. The same ID always represent the same object.
|
||||
Temporary IDs are the positions of each object in the tensor. It changes as objects get removed.
|
||||
Temporary IDs start from 1.
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
self.obj_to_tmp_id: Dict[ObjectInfo, int] = {}
|
||||
self.tmp_id_to_obj: Dict[int, ObjectInfo] = {}
|
||||
self.obj_id_to_obj: Dict[int, ObjectInfo] = {}
|
||||
|
||||
self.all_historical_object_ids: List[int] = []
|
||||
|
||||
def _recompute_obj_id_to_obj_mapping(self) -> None:
|
||||
self.obj_id_to_obj = {obj.id: obj for obj in self.obj_to_tmp_id}
|
||||
|
||||
def add_new_objects(
|
||||
self, objects: Union[List[ObjectInfo], ObjectInfo,
|
||||
List[int]]) -> (List[int], List[int]):
|
||||
if not isinstance(objects, list):
|
||||
objects = [objects]
|
||||
|
||||
corresponding_tmp_ids = []
|
||||
corresponding_obj_ids = []
|
||||
for obj in objects:
|
||||
if isinstance(obj, int):
|
||||
obj = ObjectInfo(id=obj)
|
||||
|
||||
if obj in self.obj_to_tmp_id:
|
||||
# old object
|
||||
corresponding_tmp_ids.append(self.obj_to_tmp_id[obj])
|
||||
corresponding_obj_ids.append(obj.id)
|
||||
else:
|
||||
# new object
|
||||
new_obj = ObjectInfo(id=obj.id)
|
||||
|
||||
# new object
|
||||
new_tmp_id = len(self.obj_to_tmp_id) + 1
|
||||
self.obj_to_tmp_id[new_obj] = new_tmp_id
|
||||
self.tmp_id_to_obj[new_tmp_id] = new_obj
|
||||
self.all_historical_object_ids.append(new_obj.id)
|
||||
corresponding_tmp_ids.append(new_tmp_id)
|
||||
corresponding_obj_ids.append(new_obj.id)
|
||||
|
||||
self._recompute_obj_id_to_obj_mapping()
|
||||
assert corresponding_tmp_ids == sorted(corresponding_tmp_ids)
|
||||
return corresponding_tmp_ids, corresponding_obj_ids
|
||||
|
||||
def delete_objects(self, obj_ids_to_remove: Union[int, List[int]]) -> None:
|
||||
# delete an object or a list of objects
|
||||
# re-sort the tmp ids
|
||||
if isinstance(obj_ids_to_remove, int):
|
||||
obj_ids_to_remove = [obj_ids_to_remove]
|
||||
|
||||
new_tmp_id = 1
|
||||
total_num_id = len(self.obj_to_tmp_id)
|
||||
|
||||
local_obj_to_tmp_id = {}
|
||||
local_tmp_to_obj_id = {}
|
||||
|
||||
for tmp_iter in range(1, total_num_id + 1):
|
||||
obj = self.tmp_id_to_obj[tmp_iter]
|
||||
if obj.id not in obj_ids_to_remove:
|
||||
local_obj_to_tmp_id[obj] = new_tmp_id
|
||||
local_tmp_to_obj_id[new_tmp_id] = obj
|
||||
new_tmp_id += 1
|
||||
|
||||
self.obj_to_tmp_id = local_obj_to_tmp_id
|
||||
self.tmp_id_to_obj = local_tmp_to_obj_id
|
||||
self._recompute_obj_id_to_obj_mapping()
|
||||
|
||||
def purge_inactive_objects(self,
|
||||
max_missed_detection_count: int) -> (bool, List[int], List[int]):
|
||||
# remove tmp ids of objects that are removed
|
||||
obj_id_to_be_deleted = []
|
||||
tmp_id_to_be_deleted = []
|
||||
tmp_id_to_keep = []
|
||||
obj_id_to_keep = []
|
||||
|
||||
for obj in self.obj_to_tmp_id:
|
||||
if obj.poke_count > max_missed_detection_count:
|
||||
obj_id_to_be_deleted.append(obj.id)
|
||||
tmp_id_to_be_deleted.append(self.obj_to_tmp_id[obj])
|
||||
else:
|
||||
tmp_id_to_keep.append(self.obj_to_tmp_id[obj])
|
||||
obj_id_to_keep.append(obj.id)
|
||||
|
||||
purge_activated = len(obj_id_to_be_deleted) > 0
|
||||
if purge_activated:
|
||||
self.delete_objects(obj_id_to_be_deleted)
|
||||
return purge_activated, tmp_id_to_keep, obj_id_to_keep
|
||||
|
||||
def tmp_to_obj_cls(self, mask) -> torch.Tensor:
|
||||
# remap tmp id cls representation to the true object id representation
|
||||
new_mask = torch.zeros_like(mask)
|
||||
for tmp_id, obj in self.tmp_id_to_obj.items():
|
||||
new_mask[mask == tmp_id] = obj.id
|
||||
return new_mask
|
||||
|
||||
def get_tmp_to_obj_mapping(self) -> Dict[int, ObjectInfo]:
|
||||
# returns the mapping in a dict format for saving it with pickle
|
||||
return {obj.id: tmp_id for obj, tmp_id in self.tmp_id_to_obj.items()}
|
||||
|
||||
def realize_dict(self, obj_dict, dim=1) -> torch.Tensor:
|
||||
# turns a dict indexed by obj id into a tensor, ordered by tmp IDs
|
||||
output = []
|
||||
for _, obj in self.tmp_id_to_obj.items():
|
||||
if obj.id not in obj_dict:
|
||||
raise NotImplementedError
|
||||
output.append(obj_dict[obj.id])
|
||||
output = torch.stack(output, dim=dim)
|
||||
return output
|
||||
|
||||
def make_one_hot(self, cls_mask) -> torch.Tensor:
|
||||
output = []
|
||||
for _, obj in self.tmp_id_to_obj.items():
|
||||
output.append(cls_mask == obj.id)
|
||||
if len(output) == 0:
|
||||
output = torch.zeros((0, *cls_mask.shape), dtype=torch.bool, device=cls_mask.device)
|
||||
else:
|
||||
output = torch.stack(output, dim=0)
|
||||
return output
|
||||
|
||||
@property
|
||||
def all_obj_ids(self) -> List[int]:
|
||||
return [k.id for k in self.obj_to_tmp_id]
|
||||
|
||||
@property
|
||||
def num_obj(self) -> int:
|
||||
return len(self.obj_to_tmp_id)
|
||||
|
||||
def has_all(self, objects: List[int]) -> bool:
|
||||
for obj in objects:
|
||||
if obj not in self.obj_to_tmp_id:
|
||||
return False
|
||||
return True
|
||||
|
||||
def find_object_by_id(self, obj_id) -> ObjectInfo:
|
||||
return self.obj_id_to_obj[obj_id]
|
||||
|
||||
def find_tmp_by_id(self, obj_id) -> int:
|
||||
return self.obj_to_tmp_id[self.obj_id_to_obj[obj_id]]
|
||||
@@ -0,0 +1,30 @@
|
||||
import logging
|
||||
from omegaconf import DictConfig
|
||||
|
||||
log = logging.getLogger()
|
||||
|
||||
|
||||
def get_dataset_cfg(cfg: DictConfig):
|
||||
dataset_name = cfg.dataset
|
||||
data_cfg = cfg.datasets[dataset_name]
|
||||
|
||||
potential_overrides = [
|
||||
'image_directory',
|
||||
'mask_directory',
|
||||
'json_directory',
|
||||
'size',
|
||||
'save_all',
|
||||
'use_all_masks',
|
||||
'use_long_term',
|
||||
'mem_every',
|
||||
]
|
||||
|
||||
for override in potential_overrides:
|
||||
if cfg[override] is not None:
|
||||
log.info(f'Overriding config {override} from {data_cfg[override]} to {cfg[override]}')
|
||||
data_cfg[override] = cfg[override]
|
||||
# escalte all potential overrides to the top-level config
|
||||
if override in data_cfg:
|
||||
cfg[override] = data_cfg[override]
|
||||
|
||||
return data_cfg
|
||||
Reference in New Issue
Block a user