File size: 3,277 Bytes
872b0a0 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 | import numpy as np
import torch
# import torch.nn.functional as F
class ObjectCache:
def __init__(self, cache_size=500):
self.cache_size = cache_size
self._obj_mask_cache = None
self._img_cache = None
self._motion_cache = None
self.count = 0
# We initialize the cache when we receive the first push sample so that it adapts to the current img size automatically
def init_cache(self, img_size):
self._obj_mask_cache = torch.zeros(
(self.cache_size, 1, *img_size), dtype=torch.float32
)
self._img_cache = torch.zeros(
(self.cache_size, 3, *img_size), dtype=torch.float32
)
self._motion_cache = torch.zeros((self.cache_size, 2), dtype=torch.float32)
return
def pop(self, B=8, with_aug=True): # we do not remove objects after popping
if self.count < self.cache_size: # do not use it before it is full
return None
idx = np.random.choice(self.cache_size, B, replace=False)
obj_mask = self._obj_mask_cache[idx]
img = self._img_cache[idx]
motion = self._motion_cache[idx]
if with_aug:
rand_scale = (
torch.rand(B) * 0.7 + 0.8
) # randomly rescale motion by 0.8-1.5 times
rand_scale *= (-1) ** (
torch.rand(B) > 0.5
).float() # randomly reverse motion
motion = motion * rand_scale[:, None]
flip_flag = torch.rand(B) > 0.5 # randomly horitontal-flip obj mask
img[flip_flag] = img[flip_flag].flip(dims=[3])
obj_mask[flip_flag] = obj_mask[flip_flag].flip(dims=[3])
motion[flip_flag, 0] *= -1
return obj_mask, img, motion
def push(self, obj_mask, img, motion):
"""
obj_mask: [B, 1, H, W]
img: [B, 3, H, W]
motion: [B, 2]
"""
if self._obj_mask_cache is None:
self.init_cache(img_size=img.shape[-2:])
B = obj_mask.shape[0]
if self.count <= self.cache_size - B: # many spaces
self._obj_mask_cache[self.count : (self.count + B)] = obj_mask
self._img_cache[self.count : (self.count + B)] = img
self._motion_cache[self.count : (self.count + B)] = motion
self.count += B
return
elif self.count < self.cache_size: # partial space
space = self.cache_size - self.count
self._obj_mask_cache[self.count :] = obj_mask[:space]
self._img_cache[self.count :] = img[:space]
self._motion_cache[self.count :] = motion[:space]
overwrite_idx = np.random.choice(self.count, B - space, replace=False)
self._obj_mask_cache[overwrite_idx] = obj_mask[space:]
self._img_cache[overwrite_idx] = img[space:]
self._motion_cache[overwrite_idx] = motion[space:]
self.count += space
return
else: # no spaces; random overwrite
overwrite_idx = np.random.choice(self.cache_size, B, replace=False)
self._obj_mask_cache[overwrite_idx] = obj_mask
self._img_cache[overwrite_idx] = img
self._motion_cache[overwrite_idx] = motion
return
|