initiaö commit

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