add video func support

GLEE/glee/models/glee_model.py

@@ -170,8 +170,6 @@ class GLEE_Model(nn.Module):
             features = self.backbone(images.tensor)
 
 
-
-
         if 'spatial' in prompts:
             ## setp 1,2,3
             key_images = [ images ]  #bz*[1,3,H,W]
@@ -240,12 +238,130 @@ class GLEE_Model(nn.Module):
             extra['visual_prompt_tokens'] = src_spatial_queries #[len,bz,C]
             extra['visual_prompt_nonzero_mask'] = src_spatial_maskings  # [bz,len]
         
- 
         outputs = self.predictor(multi_scale_features, mask_features, extra=extra, task=task, masks=None, targets=targets)
         return  outputs 
  
 
 
+    
+    def vos_step1(self, previous_image, prompts, task, targets=None, batch_name_list=None, is_train = False):
+        extra = {}
+        if isinstance(previous_image,torch.Tensor):
+            features = self.backbone(previous_image)
+        else:
+            features = self.backbone(previous_image.tensor)
+        # bz = len(images)//2
+
+        ## setp 1,2,3
+        key_images = [previous_image]  #bz*[1,3,H,W]
+        key_promptmasks = [m.unsqueeze(0) for m in prompts['spatial']] #bz*[1,1,H,W]
+        ref_feats, ref_masks = self.get_template(key_images, key_promptmasks) 
+
+        early_fusion = {"hidden":ref_feats,"masks":ref_masks} 
+
+        
+
+        mask_features, _, multi_scale_features, zero_loss = self.pixel_decoder.forward_features(features, masks=None, early_fusion = early_fusion)
+        
+
+        prompt_multi_scale_features = multi_scale_features+[mask_features]
+
+        if 'spatial' in prompts:
+            pos_masks = prompts['spatial']
+            # neg_masks = [~p for p in prompts['spatial']]
+            neg_masks = [p&False for p in prompts['spatial']]
+            
+            extra.update({'spatial_query_pos_mask': pos_masks, 'spatial_query_neg_mask': neg_masks})
+            # import pdb;pdb.set_trace()
+
+
+            _,h,w = extra['spatial_query_pos_mask'][0].shape
+            divisor = torch.tensor([h,w], device=mask_features.device)[None,]
+            # Get mean pos spatial query
+            non_zero_pos_point = [rand_sample((m.nonzero()[:,1:]/divisor).t(), self.max_spatial_len[-1]).t() for m in extra['spatial_query_pos_mask']]
+            # [:,1:]第一个维度是指示属于那个batch，原本这里的mshape是【num_inst,H,W】,得到的nonzero 是【num_point，3】，[:,1:]是xy坐标，
+            # #这里舍弃第一个维度是表示每张图片上prompt覆盖的物体的point混在一起采样，没有instance之间的区分. 因此每个图片都得到一个[512,2]的point set，是采样过后的正样本
+            non_zero_pos_point = nn.utils.rnn.pad_sequence(non_zero_pos_point, padding_value=-1).permute(1,0,2)  # 把list中的结果通过padding concat到一起,得到的是[bz,512,2]
+            
+            non_zero_pos_mask = (non_zero_pos_point.sum(dim=-1) < 0) # 把xy坐标相加小于0的找出来
+            spatial_query_pos = point_sample(mask_features, non_zero_pos_point.flip(dims=(2,)).type(mask_features.dtype), align_corners=True) #[(N, C, P)
+            spatial_query_pos = torch.stack([x[m].mean(dim=0, keepdim=True) for x, m in zip(spatial_query_pos.transpose(1,2), ~non_zero_pos_mask)]).transpose(0,1).nan_to_num() # [1,bz,C]
+            # import pdb;pdb.set_trace()
+            # Get mean neg spatial query
+            non_zero_neg_point = [rand_sample((m.nonzero()[:,1:]/divisor).t(), self.max_spatial_len[-1]).t() for m in extra['spatial_query_neg_mask']]
+            non_zero_neg_point = nn.utils.rnn.pad_sequence(non_zero_neg_point, padding_value=-1).permute(1,0,2)
+            non_zero_neg_mask = (non_zero_neg_point.sum(dim=-1) < 0)
+            spatial_query_neg = point_sample(mask_features, non_zero_neg_point.flip(dims=(2,)).type(mask_features.dtype), align_corners=True)
+            spatial_query_neg = torch.stack([x[m].mean(dim=0, keepdim=True) for x, m in zip(spatial_query_neg.transpose(1,2), ~non_zero_neg_mask)]).transpose(0,1).nan_to_num()
+
+            # Get layerwise spatial query
+            src_spatial_queries = []
+            src_spatial_maskings = []
+            for i in range(len(prompt_multi_scale_features)):
+                bs,dc,h,w = prompt_multi_scale_features[i].shape
+                # src_mask_features = multi_scale_features[i].view(h,w,bs,dc)
+                src_mask_features = prompt_multi_scale_features[i].permute(2,3,0,1)
+                # import pdb;pdb.set_trace()
+                src_mask_features = src_mask_features @ self.mask_sptial_embed[i]
+
+                non_zero_query_point_pos = [rand_sample((m.nonzero()[:,1:]/divisor).t(), self.max_spatial_len[i]).t() for m in extra['spatial_query_pos_mask']]
+                non_zero_query_point_neg = [rand_sample((m.nonzero()[:,1:]/divisor).t(), self.max_spatial_len[i]).t() for m in extra['spatial_query_neg_mask']]
+                non_zero_query_point = [torch.cat([x,y], dim=0) for x,y in zip(non_zero_query_point_pos, non_zero_query_point_neg)]
+                pos_neg_indicator = [torch.cat([torch.ones(x.shape[0], device=x.device), -torch.ones(y.shape[0], device=y.device)]) for x,y in zip(non_zero_query_point_pos, non_zero_query_point_neg)]
+                pos_neg_indicator = nn.utils.rnn.pad_sequence(pos_neg_indicator, padding_value=0)
+                # import pdb;pdb.set_trace()
+                non_zero_query_point = nn.utils.rnn.pad_sequence(non_zero_query_point, padding_value=-1).permute(1,0,2)
+                non_zero_query_mask = (non_zero_query_point.sum(dim=-1) < 0)
+                non_zero_query_point[non_zero_query_mask] = 0
+                # import pdb;pdb.set_trace()
+                spatial_tokens = point_sample(src_mask_features.permute(2,3,0,1), non_zero_query_point.flip(dims=(2,)).type(src_mask_features.dtype), align_corners=True).permute(2,0,1)
+                spatial_tokens[pos_neg_indicator==1] += self.pn_indicator.weight[0:1]
+                spatial_tokens[pos_neg_indicator==-1] += self.pn_indicator.weight[1:2]
+
+                src_spatial_queries += [spatial_tokens]
+                src_spatial_maskings += [non_zero_query_mask]
+            
+            extra['visual_prompt_tokens'] = src_spatial_queries #[len,bz,C]
+            extra['visual_prompt_nonzero_mask'] = src_spatial_maskings  # [bz,len]
+
+        return  early_fusion, extra
+ 
+
+
+    def vos_step2(self, images, task, language_dict_features, last_extra, targets=None, batch_name_list=None, is_train = False):
+        extra =  last_extra
+        dist_loss = None
+        if True:
+            if task not in ['grounding','rvos']:
+                assert batch_name_list
+                calsses_name_list = batch_name_list
+                tokenized = self.tokenizer.batch_encode_plus(calsses_name_list,
+                        max_length=self.cfg.MODEL.LANGUAGE_BACKBONE.MAX_QUERY_LEN, # 256
+                        padding='max_length' if self.cfg.MODEL.LANGUAGE_BACKBONE.PAD_MAX else "longest", # max_length
+                        return_special_tokens_mask=True,
+                        return_tensors='pt',
+                        truncation=True).to("cuda")
+
+                texts = (tokenized['input_ids'], tokenized['attention_mask'])
+                token_x = self.text_encoder(*texts)['last_hidden_state']
+                token_x = token_x @ self.lang_projection
+                lang_feat_pool = agg_lang_feat(token_x, tokenized['attention_mask'], pool_type="average")  # (bs, 768)
+                extra['class_embeddings'] = lang_feat_pool
+ 
+        if isinstance(images,torch.Tensor):
+            features = self.backbone(images)
+        else:
+            features = self.backbone(images.tensor)
+        # bz = len(images)//2
+        # import pdb;pdb.set_trace()
+        mask_features, _, multi_scale_features, zero_loss = self.pixel_decoder.forward_features(features, masks=None, early_fusion = language_dict_features)
+
+
+        outputs = self.predictor(multi_scale_features, mask_features, extra=extra, task=task, masks=None, targets=targets)
+        return  outputs
+ 
+
+