| 
						
						
							
								
							
						
						
					 | 
				
				 | 
				
					@ -41,12 +41,14 @@ class Omnivore(NNOperator): | 
				
			
			
		
	
		
			
				
					 | 
					 | 
				
				 | 
				
					    def __init__(self, | 
				
			
			
		
	
		
			
				
					 | 
					 | 
				
				 | 
				
					                 model_name: str = 'omnivore_swinT', | 
				
			
			
		
	
		
			
				
					 | 
					 | 
				
				 | 
				
					                 framework: str = 'pytorch', | 
				
			
			
		
	
		
			
				
					 | 
					 | 
				
				 | 
				
					                 input_type: str = 'video', | 
				
			
			
		
	
		
			
				
					 | 
					 | 
				
				 | 
				
					                 skip_preprocess: bool = False, | 
				
			
			
		
	
		
			
				
					 | 
					 | 
				
				 | 
				
					                 classmap: dict = None, | 
				
			
			
		
	
		
			
				
					 | 
					 | 
				
				 | 
				
					                 topk: int = 5, | 
				
			
			
		
	
		
			
				
					 | 
					 | 
				
				 | 
				
					                 ): | 
				
			
			
		
	
		
			
				
					 | 
					 | 
				
				 | 
				
					        super().__init__(framework=framework) | 
				
			
			
		
	
		
			
				
					 | 
					 | 
				
				 | 
				
					        self.model_name = model_name | 
				
			
			
		
	
		
			
				
					 | 
					 | 
				
				 | 
				
					        self.input_type = input_type | 
				
			
			
		
	
		
			
				
					 | 
					 | 
				
				 | 
				
					        self.skip_preprocess = skip_preprocess | 
				
			
			
		
	
		
			
				
					 | 
					 | 
				
				 | 
				
					        self.topk = topk | 
				
			
			
		
	
		
			
				
					 | 
					 | 
				
				 | 
				
					        self.dataset_name = 'kinetics_400' | 
				
			
			
		
	
	
		
			
				
					| 
						
							
								
							
						
						
							
								
							
						
						
					 | 
				
				 | 
				
					@ -99,11 +101,14 @@ class Omnivore(NNOperator): | 
				
			
			
		
	
		
			
				
					 | 
					 | 
				
				 | 
				
					                    ) | 
				
			
			
		
	
		
			
				
					 | 
					 | 
				
				 | 
				
					        inputs = data.to(self.device)[None, ...] | 
				
			
			
		
	
		
			
				
					 | 
					 | 
				
				 | 
				
					
 | 
				
			
			
		
	
		
			
				
					 | 
					 | 
				
				 | 
				
					        outs = self.model(inputs,input_type="video") | 
				
			
			
		
	
		
			
				
					 | 
					 | 
				
				 | 
				
					        feats = self.model.forward_features(inputs ,input_type = self.input_type) | 
				
			
			
		
	
		
			
				
					 | 
					 | 
				
				 | 
				
					        features = feats.to('cpu').squeeze(0).detach().numpy() | 
				
			
			
		
	
		
			
				
					 | 
					 | 
				
				 | 
				
					
 | 
				
			
			
		
	
		
			
				
					 | 
					 | 
				
				 | 
				
					        outs = self.model.head(feats) | 
				
			
			
		
	
		
			
				
					 | 
					 | 
				
				 | 
				
					        post_act = torch.nn.Softmax(dim=1) | 
				
			
			
		
	
		
			
				
					 | 
					 | 
				
				 | 
				
					        preds = post_act(outs) | 
				
			
			
		
	
		
			
				
					 | 
					 | 
				
				 | 
				
					        pred_scores, pred_classes = preds.topk(k=self.topk) | 
				
			
			
		
	
		
			
				
					 | 
					 | 
				
				 | 
				
					        labels = [self.classmap[int(i)] for i in pred_classes[0]] | 
				
			
			
		
	
		
			
				
					 | 
					 | 
				
				 | 
				
					        scores = [round(float(x), 5) for x in pred_scores[0]] | 
				
			
			
		
	
		
			
				
					 | 
					 | 
				
				 | 
				
					        print(labels,scores) | 
				
			
			
		
	
		
			
				
					 | 
					 | 
				
				 | 
				
					        return labels, scores | 
				
			
			
		
	
		
			
				
					 | 
					 | 
				
				 | 
				
					
 | 
				
			
			
		
	
		
			
				
					 | 
					 | 
				
				 | 
				
					        return labels, scores, features | 
				
			
			
		
	
	
		
			
				
					| 
						
						
						
					 | 
				
				 | 
				
					
  |