diff options
Diffstat (limited to 'model/utils/slot_utils.py')
| -rw-r--r-- | model/utils/slot_utils.py | 326 |
1 files changed, 326 insertions, 0 deletions
diff --git a/model/utils/slot_utils.py b/model/utils/slot_utils.py new file mode 100644 index 0000000..936248b --- /dev/null +++ b/model/utils/slot_utils.py @@ -0,0 +1,326 @@ +import torch.nn as nn +import torch as th +import torchvision.transforms as transforms +import torch.nn.functional as F +from einops import rearrange, repeat, reduce +from typing import Tuple, Union, List +from model.utils.nn_utils import Gaus2D, LambdaModule, TanhAlpha + +class InitialLatentStates(nn.Module): + def __init__( + self, + gestalt_size: int, + num_objects: int, + bottleneck: str, + size: Tuple[int, int], + teacher_forcing: int + ): + super(InitialLatentStates, self).__init__() + self.bottleneck = bottleneck + + self.num_objects = num_objects + self.gestalt_mean = nn.Parameter(th.zeros(1, gestalt_size)) + self.gestalt_std = nn.Parameter(th.ones(1, gestalt_size)) + self.std = nn.Parameter(th.zeros(1)) + self.gestalt_strength = 2 + self.teacher_forcing = teacher_forcing + + self.init = TanhAlpha(start = -1) + self.register_buffer('priority', th.arange(num_objects).float() * 25, persistent=False) + self.register_buffer('threshold', th.ones(1) * 0.8) + self.last_mask = None + self.binarize_first = round(gestalt_size * 0.8) + + self.gaus2d = nn.Sequential( + Gaus2D((size[0] // 16, size[1] // 16)), + Gaus2D((size[0] // 4, size[1] // 4)), + Gaus2D(size) + ) + + self.level = 1 + self.t = 0 + + self.to_batch = LambdaModule(lambda x: rearrange(x, 'b (o c) -> (b o) c', o = num_objects)) + self.to_shared = LambdaModule(lambda x: rearrange(x, '(b o) c -> b (o c)', o = num_objects)) + + self.blur = transforms.GaussianBlur(13) + self.size = size + + def reset_state(self): + self.last_mask = None + self.t = 0 + self.to_next_spawn = 0 + + def set_level(self, level): + self.level = level + factor = int(4 / (level ** 2)) + self.to_position = ErrorToPosition((self.size[0] // factor, self.size[1] // factor)) + + def forward( + self, + error: th.Tensor, + mask: th.Tensor = None, + position: th.Tensor = None, + gestalt: th.Tensor = None, + priority: th.Tensor = None, + shuffleslots: bool = True, + slots_bounded_last: th.Tensor = None, + slots_occlusionfactor_last: th.Tensor = None, + allow_spawn: bool = True, + clean_slots: bool = False + ): + + batch_size = error.shape[0] + device = error.device + + if self.init.get() < 1: + self.gestalt_strength = self.init() + + if self.last_mask is None: + self.last_mask = th.zeros((batch_size * self.num_objects, 1), device = device) + if shuffleslots: + self.slots_assigned = th.ones((batch_size * self.num_objects, 1), device = device) + else: + self.slots_assigned = th.zeros((batch_size * self.num_objects, 1), device = device) + + if not allow_spawn: + unnassigned = self.slots_assigned - slots_bounded_last + self.slots_assigned = self.slots_assigned - unnassigned + + if clean_slots and (slots_occlusionfactor_last is not None): + occluded = self.slots_assigned * (self.to_batch(slots_occlusionfactor_last) > 0.1).float() + self.slots_assigned = self.slots_assigned - occluded + + if (slots_bounded_last is None) or (self.gestalt_strength < 1): + + if mask is not None: + # maximum berechnung --> slot gebunden c=o + mask2 = reduce(mask[:,:-1], 'b c h w -> (b c) 1' , 'max').detach() + + if self.gestalt_strength <= 0: + self.last_mask = mask2 + elif self.gestalt_strength < 1: + self.last_mask = th.maximum(self.last_mask, mask2) + self.last_mask = self.last_mask - th.relu(-1 * (mask2 - self.threshold) * (1 - self.gestalt_strength)) + else: + self.last_mask = th.maximum(self.last_mask, mask2) + + slots_bounded = (self.last_mask > self.threshold).float().detach() * self.slots_assigned + else: + slots_bounded = slots_bounded_last * self.slots_assigned + + if self.bottleneck == "binar": + gestalt_new = repeat(th.sigmoid(self.gestalt_mean), '1 c -> b c', b = batch_size * self.num_objects) + gestalt_new = gestalt_new + gestalt_new * (1 - gestalt_new) * th.randn_like(gestalt_new) + else: + gestalt_mean = repeat(self.gestalt_mean, '1 c -> b c', b = batch_size * self.num_objects) + gestalt_std = repeat(self.gestalt_std, '1 c -> b c', b = batch_size * self.num_objects) + gestalt_new = th.sigmoid(gestalt_mean + gestalt_std * th.randn_like(gestalt_std)) + + if gestalt is None: + gestalt = gestalt_new + else: + gestalt = self.to_batch(gestalt) * slots_bounded + gestalt_new * (1 - slots_bounded) + + if priority is None: + priority = repeat(self.priority, 'o -> (b o) 1', b = batch_size) + else: + priority = self.to_batch(priority) * slots_bounded + repeat(self.priority, 'o -> (b o) 1', b = batch_size) * (1 - slots_bounded) + + + if shuffleslots: + self.slots_assigned = th.ones_like(self.slots_assigned) + + xy_rand_new = th.rand((batch_size * self.num_objects * 10, 2), device = device) * 2 - 1 + std_new = th.zeros((batch_size * self.num_objects * 10, 1), device = device) + position_new = th.cat((xy_rand_new, std_new), dim=1) + + position2d = self.gaus2d[self.level](position_new) + position2d = rearrange(position2d, '(b o) 1 h w -> b o h w', b = batch_size) + + rand_error = reduce(position2d * error, 'b o h w -> (b o) 1', 'sum') + + xy_rand_new = rearrange(xy_rand_new, '(b r) c -> r b c', r = 10) + rand_error = rearrange(rand_error, '(b r) c -> r b c', r = 10) + + max_error = th.argmax(rand_error, dim=0, keepdim=True) + x, y = th.chunk(xy_rand_new, 2, dim=2) + x = th.gather(x, dim=0, index=max_error).detach().squeeze(dim=0) + y = th.gather(y, dim=0, index=max_error).detach().squeeze(dim=0) + std = repeat(self.std, '1 -> (b o) 1', b = batch_size, o=self.num_objects) + + if position is None: + position = th.cat((x, y, std), dim=1) + else: + position = self.to_batch(position) * slots_bounded + th.cat((x, y, std), dim=1) * (1 - slots_bounded) + + else: + + # set unassigned slots to empty position + empty_position = th.tensor([-1,-1,0]).to(device) + empty_position = repeat(empty_position, 'c -> (b o) c', b = batch_size, o=self.num_objects).detach() + + if position is None: + position = empty_position + else: + position = self.to_batch(position) * self.slots_assigned + empty_position * (1 - self.slots_assigned) + + + # blur errror, and set masked areas to zero + error = self.blur(error) + if mask is not None: + mask2 = mask[:,:-1] * rearrange(slots_bounded, '(b o) 1 -> b o 1 1', b = batch_size) + mask2 = th.sum(mask2, dim=1, keepdim=True) + error = error * (1-mask2) + max_error = reduce(error, 'b o h w -> (b o) 1', 'max') + + if self.to_next_spawn <= 0 and allow_spawn: + + self.to_next_spawn = 2 + + # calculate the position with the highest error + new_pos = self.to_position(error) + std = repeat(self.std, '1 -> b 1', b = batch_size) + new_pos = repeat(th.cat((new_pos, std), dim=1), 'b c -> (b o) c', o = self.num_objects) + + # calculate if an assigned slot is unbound (-->free) + n_slots_assigned = self.to_shared(self.slots_assigned).sum(dim=1, keepdim=True) + n_slots_bounded = self.to_shared(slots_bounded).sum(dim=1, keepdim=True) + free_slot_given = th.clip(n_slots_assigned - n_slots_bounded, 0, 1) + + # either spawn a new slot or use the one that is free + slots_new_index = n_slots_assigned * (1-free_slot_given) + n_slots_bounded * free_slot_given # reset the free slot each timespawn + + # new slot index + free_slot_required = (max_error > 0).float() + slots_new_index = F.one_hot(slots_new_index.long(), num_classes=self.num_objects+1).float().squeeze(dim=1)[:,:-1] + slots_new_index = self.to_batch(slots_new_index * free_slot_required) + + # place new free slot + position = new_pos * slots_new_index + position * (1 - slots_new_index) + self.slots_assigned = th.clip(self.slots_assigned + slots_new_index, 0, 1) + + self.to_next_spawn -= 1 + return self.to_shared(position), self.to_shared(gestalt), self.to_shared(priority), error + + def get_slots_unassigned(self): + return self.to_shared(1-self.slots_assigned) + + def get_slots_assigned(self): + return self.to_shared(self.slots_assigned) + + +class OcclusionTracker(nn.Module): + def __init__(self, batch_size, num_objects, device): + super(OcclusionTracker, self).__init__() + self.batch_size = batch_size + self.num_objects = num_objects + self.slots_bounded_all = th.zeros((batch_size * num_objects, 1)).to(device) + self.threshold = 0.8 + self.device = device + self.to_shared = LambdaModule(lambda x: rearrange(x, '(b o) c -> b (o c)', o = num_objects)) + self.slots_bounded_next_last = None + + def forward( + self, + mask: th.Tensor = None, + rawmask: th.Tensor = None, + reset_mask: bool = False, + update: bool = True + ): + + if mask is not None: + + # compute bounding mask + slots_bounded_smooth_cur = reduce(mask[:,:-1], 'b o h w -> (b o) 1' , 'max').detach() + slots_bounded_cur = (slots_bounded_smooth_cur > self.threshold).float().detach() + if reset_mask: + self.slots_bounded_next_last = slots_bounded_cur # allow immediate spawn + + if update: + slots_bounded_cur = slots_bounded_cur * th.clip(self.slots_bounded_next_last + self.slots_bounded_all, 0, 1) + else: + self.slots_bounded_next_last = slots_bounded_cur + + if reset_mask: + self.slots_bounded_smooth_all = slots_bounded_smooth_cur + self.slots_bounded_all = slots_bounded_cur + elif update: + self.slots_bounded_all = th.maximum(self.slots_bounded_all, slots_bounded_cur) + self.slots_bounded_smooth_all = th.maximum(self.slots_bounded_smooth_all, slots_bounded_smooth_cur) + + # compute occlusion mask + slots_occluded_cur = self.slots_bounded_all - slots_bounded_cur + + # compute partially occluded mask + mask = (mask[:,:-1] > self.threshold).float().detach() + rawmask = (rawmask[:,:-1] > self.threshold).float().detach() + masked = rawmask - mask + + masked = reduce(masked, 'b o h w -> (b o) 1' , 'sum') + rawmask = reduce(rawmask, 'b o h w -> (b o) 1' , 'sum') + + slots_occlusionfactor_cur = (masked / (rawmask + 1)) * (1-slots_occluded_cur) + slots_occluded_cur + slots_partially_occluded = (slots_occlusionfactor_cur > 0.1).float() #* slots_bounded_cur + slots_fully_visible = (slots_occlusionfactor_cur <= 0.1).float() * slots_bounded_cur + + if reset_mask: + self.slots_fully_visible_all = slots_fully_visible + elif update: + self.slots_fully_visible_all = th.maximum(self.slots_fully_visible_all, slots_fully_visible) + + return self.to_shared(self.slots_bounded_all), self.to_shared(self.slots_bounded_smooth_all), self.to_shared(slots_occluded_cur), self.to_shared(slots_partially_occluded), self.to_shared(slots_fully_visible), self.to_shared(slots_occlusionfactor_cur) + + def get_slots_fully_visible_all(self): + return self.to_shared(self.slots_fully_visible_all) + +class ErrorToPosition(nn.Module): + def __init__(self, size: Union[int, Tuple[int, int]]): + super(ErrorToPosition, self).__init__() + + self.register_buffer("grid_x", th.arange(size[0]), persistent=False) + self.register_buffer("grid_y", th.arange(size[1]), persistent=False) + + self.grid_x = (self.grid_x / (size[0]-1)) * 2 - 1 + self.grid_y = (self.grid_y / (size[1]-1)) * 2 - 1 + + self.grid_x = self.grid_x.view(1, 1, -1, 1).expand(1, 1, *size).clone() + self.grid_y = self.grid_y.view(1, 1, 1, -1).expand(1, 1, *size).clone() + + self.grid_x = self.grid_x.view(1, 1, -1) + self.grid_y = self.grid_y.view(1, 1, -1) + + self.size = size + + def forward(self, input: th.Tensor): + assert input.shape[1] == 1 + + input = rearrange(input, 'b c h w -> b c (h w)') + argmax = th.argmax(input, dim=2, keepdim=True) + + x = self.grid_x[0,0,argmax].squeeze(dim=2) + y = self.grid_y[0,0,argmax].squeeze(dim=2) + + return th.cat((x,y),dim=1) + + +def compute_rawmask(mask, bg_mask): + + num_objects = mask.shape[1] + + # d is a diagonal matrix which defines what to take the softmax over + d_mask = th.diag(th.ones(num_objects+1)).to(mask.device) + d_mask[:,-1] = 1 + d_mask[-1,-1] = 0 + + # take subset of rawmask with the diagonal matrix + rawmask = th.cat((mask, bg_mask), dim=1) + rawmask = repeat(rawmask, 'b o h w -> b r o h w', r = num_objects+1) + rawmask = rawmask[:,d_mask.bool()] + rawmask = rearrange(rawmask, 'b (o r) h w -> b o r h w', o = num_objects) + + # take softmax between each object mask and the background mask + rawmask = th.squeeze(th.softmax(rawmask, dim=2)[:,:,0], dim=2) + rawmask = th.cat((rawmask, bg_mask), dim=1) # add background mask + + return rawmask
\ No newline at end of file |