Spaces:
Running
Running
| # Copyright 2022 - Valeo Comfort and Driving Assistance - Oriane Siméoni @ valeo.ai | |
| # | |
| # Licensed under the Apache License, Version 2.0 (the "License"); | |
| # you may not use this file except in compliance with the License. | |
| # You may obtain a copy of the License at | |
| # | |
| # http://www.apache.org/licenses/LICENSE-2.0 | |
| # | |
| # Unless required by applicable law or agreed to in writing, software | |
| # distributed under the License is distributed on an "AS IS" BASIS, | |
| # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. | |
| # See the License for the specific language governing permissions and | |
| # limitations under the License. | |
| import torch | |
| import torch.nn.functional as F | |
| from typing import Tuple | |
| def compute_img_bkg_seg( | |
| attentions, | |
| feats, | |
| featmap_dims, | |
| th_bkg, | |
| dim=64, | |
| epsilon: float = 1e-10, | |
| apply_weights: bool = True, | |
| ) -> Tuple[torch.Tensor, float]: | |
| """ | |
| inputs | |
| - attentions [B, ] | |
| """ | |
| w_featmap, h_featmap = featmap_dims | |
| nb, nh, _ = attentions.shape[:3] | |
| # we keep only the output patch attention | |
| att = attentions[:, :, 0, 1:].reshape(nb, nh, -1) | |
| att = att.reshape(nb, nh, w_featmap, h_featmap) | |
| # ----------------------------------------------- | |
| # Inspired by CroW sparsity channel weighting of each head CroW, Kalantidis etal. | |
| threshold = torch.mean(att.reshape(nb, -1), dim=1) # Find threshold per image | |
| Q = torch.sum( | |
| att.reshape(nb, nh, w_featmap * h_featmap) > threshold[:, None, None], axis=2 | |
| ) / (w_featmap * h_featmap) | |
| beta = torch.log(torch.sum(Q + epsilon, dim=1)[:, None] / (Q + epsilon)) | |
| # Weight features based on attention sparsity | |
| descs = feats[ | |
| :, | |
| 1:, | |
| ] | |
| if apply_weights: | |
| descs = (descs.reshape(nb, -1, nh, dim) * beta[:, None, :, None]).reshape( | |
| nb, -1, nh * dim | |
| ) | |
| else: | |
| descs = (descs.reshape(nb, -1, nh, dim)).reshape(nb, -1, nh * dim) | |
| # ----------------------------------------------- | |
| # Compute cosine-similarities | |
| descs = F.normalize(descs, dim=-1, p=2) | |
| cos_sim = torch.bmm(descs, descs.permute(0, 2, 1)) | |
| # ----------------------------------------------- | |
| # Find pixel with least amount of attention | |
| if apply_weights: | |
| att = att.reshape(nb, nh, w_featmap, h_featmap) * beta[:, :, None, None] | |
| else: | |
| att = att.reshape(nb, nh, w_featmap, h_featmap) | |
| id_pixel_ref = torch.argmin(torch.sum(att, axis=1).reshape(nb, -1), dim=-1) | |
| # ----------------------------------------------- | |
| # Mask of definitely background pixels: 1 on the background | |
| cos_sim = cos_sim.reshape(nb, -1, w_featmap * h_featmap) | |
| bkg_mask = ( | |
| cos_sim[torch.arange(cos_sim.size(0)), id_pixel_ref, :].reshape( | |
| nb, w_featmap, h_featmap | |
| ) | |
| > th_bkg | |
| ) # mask to be used to remove background | |
| return bkg_mask.float() | |